@grest-ts/websocket 0.0.5

package/src/client/GGSocketPool.ts

@@ -0,0 +1,244 @@
+import {GGSocket} from '../socket/GGSocket';
+import {GGWebSocketHandshakeContext, GGWebSocketMiddleware} from "../schema/GGWebSocketMiddleware";
+import {SocketAdapter} from "../socket/SocketAdapter";
+import {GG_WS_CONNECTION} from "../server/GG_WS_CONNECTION";
+import {Message, MessageType} from "../socket/SocketMessage";
+import {GGValidator, SERVER_ERROR} from "@grest-ts/schema";
+import {withTimeout} from "@grest-ts/common";
+import {GGContext} from "@grest-ts/context";
+import {GG_TRACE} from "@grest-ts/trace";
+import {getDefaultAdapter} from "../adapter/getDefaultAdapter";
+
+export interface GGSocketPoolConfig<Query> {
+    domain: string,
+    path: string,
+    query?: Query
+    queryValidator?: GGValidator<Query>
+    middlewares?: readonly GGWebSocketMiddleware[]
+}
+
+/**
+ * Connection pool for WebSocket connections
+ * Reuses existing connections when the same URL + headers combination is requested
+ */
+export class GGSocketPool {
+    private static sockets = new Map<string, GGSocket>();
+    private static pendingSockets = new Map<string, Promise<GGSocket>>();
+    private static adapter: any = null;
+    private static adapterPromise: Promise<any> | null = null;
+
+    /**
+     * Get the current number of active connections in the pool
+     */
+    public static get size(): number {
+        return this.sockets.size;
+    }
+
+    /**
+     * Get the current number of pending connections being established
+     */
+    public static get pendingSize(): number {
+        return this.pendingSockets.size;
+    }
+
+    /**
+     * Close all pooled connections gracefully.
+     * Waits for pending connections to establish before closing them.
+     * @param graceful - If true, calls teardown() on each socket allowing pending requests to complete.
+     *                   If false, calls close() for immediate termination.
+     */
+    public static async closeAll(graceful: boolean = true): Promise<void> {
+        // Wait for any pending connections to establish first
+        const pendingPromises = Array.from(this.pendingSockets.values());
+        if (pendingPromises.length > 0) {
+            await Promise.allSettled(pendingPromises);
+        }
+
+        // Close all active connections
+        const sockets = Array.from(this.sockets.values());
+        if (graceful) {
+            await Promise.allSettled(sockets.map(socket => socket.teardown()));
+        } else {
+            sockets.forEach(socket => socket.close());
+        }
+
+        // Clear the maps (onClose handlers should have already removed them,
+        // but clear explicitly to handle any edge cases)
+        this.sockets.clear();
+        this.pendingSockets.clear();
+    }
+
+    /**
+     * Clear all maps without closing connections.
+     * WARNING: This will cause memory leaks if connections are still active.
+     * Only use for testing purposes.
+     */
+    public static __clearForTesting(): void {
+        this.sockets.clear();
+        this.pendingSockets.clear();
+        this.adapter = null;
+        this.adapterPromise = null;
+    }
+
+    /**
+     * Remove a specific connection from the pool by its key.
+     * Does NOT close the socket - just removes it from the pool.
+     * @returns true if the connection was found and removed, false otherwise
+     */
+    public static removeFromPool(key: string): boolean {
+        return this.sockets.delete(key);
+    }
+
+    /**
+     * Get all connection keys currently in the pool (for debugging/monitoring)
+     */
+    public static getConnectionKeys(): string[] {
+        return Array.from(this.sockets.keys());
+    }
+
+    /**
+     * Ensure adapter is loaded (lazy initialization)
+     */
+    private static async ensureAdapter(): Promise<any> {
+        if (this.adapter) {
+            return this.adapter;
+        }
+
+        if (!this.adapterPromise) {
+            this.adapterPromise = getDefaultAdapter();
+        }
+
+        // Always await the promise to handle race conditions
+        this.adapter = await this.adapterPromise;
+        return this.adapter;
+    }
+
+    public static setAdapter(adapter: new(args: any, options?: any) => SocketAdapter) {
+        this.adapter = adapter;
+        this.adapterPromise = Promise.resolve(adapter);
+    }
+
+    /**
+     * Build headers from middlewares' updateHandshake()
+     */
+    private static buildHeaders(config: GGSocketPoolConfig<any>): Record<string, string> {
+        if (!config.middlewares) {
+            return {};
+        }
+
+        const handshakeContext: GGWebSocketHandshakeContext = {
+            headers: {},
+            queryArgs: (config.query as Record<string, string>) ?? {}
+        };
+
+        for (const middleware of config.middlewares) {
+            middleware.updateHandshake?.(handshakeContext);
+        }
+
+        return handshakeContext.headers;
+    }
+
+    static async getOrConnect<Query>(
+        config: GGSocketPoolConfig<Query>
+    ): Promise<GGSocket> {
+        // Build headers from middlewares
+        const headers = this.buildHeaders(config);
+
+        // Build full URL with query string if provided
+        let fullUrl = config.domain + config.path;
+        if (config.query) {
+            const queryEntries: [string, string][] = Object.entries(config.query).map(([key, value]) => [key, String(value)]);
+            fullUrl += '?' + new URLSearchParams(queryEntries).toString();
+        }
+
+        // Create connection key based on URL + headers
+        const headerKey = Object.entries(headers).sort().map(([k, v]) => `${k}=${v}`).join('&');
+        const key = fullUrl + "::" + headerKey;
+
+        // Check for existing connection first
+        if (this.sockets.has(key)) {
+            return this.sockets.get(key);
+        }
+        if (this.pendingSockets.has(key)) {
+            return this.pendingSockets.get(key);
+        }
+
+        // Create the connection promise BEFORE any async operations to prevent race conditions
+        // This ensures that concurrent calls will see the pending promise
+        const connectionPromise = (async () => {
+            // Ensure adapter is loaded (this is async but safely inside the promise)
+            const adapterClass = await this.ensureAdapter();
+
+            return new Promise<GGSocket>((resolve, reject) => {
+                const adapter = new adapterClass(fullUrl);
+                adapter.onOpen(async () => {
+                    try {
+                        const context = new GGContext("ws-client-connection");
+                        await context.run(async () => {
+                            GG_TRACE.init();
+                            GG_WS_CONNECTION.set({
+                                port: undefined,
+                                path: config.domain
+                            });
+
+                            // Send handshake with headers
+                            adapter.send(Message.create(MessageType.HANDSHAKE, "", "", headers));
+
+                            // Wait for handshake response
+                            await withTimeout(
+                                new Promise<void>((handshakeResolve, handshakeReject) => {
+                                    const onMessage = (data: string) => {
+                                        const msg = Message.parse(data);
+                                        if (!msg) return;
+
+                                        if (msg.type === MessageType.HANDSHAKE_OK) {
+                                            adapter.offMessage(onMessage);
+                                            handshakeResolve();
+                                        } else if (msg.type === MessageType.HANDSHAKE_ERR) {
+                                            adapter.offMessage(onMessage);
+                                            handshakeReject(new SERVER_ERROR({
+                                                displayMessage: 'WebSocket handshake failed',
+                                                originalError: msg.data
+                                            }));
+                                        }
+                                    };
+                                    adapter.onMessage(onMessage);
+                                }),
+                                5000,
+                                'Handshake timeout'
+                            );
+                            resolve(new GGSocket(adapter, {connectionContext: context}));
+                        });
+                    } catch (error) {
+                        reject(error);
+                    }
+                });
+                adapter.onError((error: Error) => {
+                    reject(error);
+                });
+            });
+        })();
+
+        // Store the pending promise IMMEDIATELY (before awaiting)
+        this.pendingSockets.set(key, connectionPromise);
+
+        try {
+            const socket = await connectionPromise;
+
+            // Store the connection
+            this.sockets.set(key, socket);
+            this.pendingSockets.delete(key);
+
+            // Clean up on close
+            socket.onClose(() => {
+                this.sockets.delete(key);
+            });
+
+            return socket;
+        } catch (error) {
+            // Clean up failed connection attempt
+            this.pendingSockets.delete(key);
+            throw error;
+        }
+    }
+}

package/src/index-browser.ts

@@ -0,0 +1,21 @@
+// Type helpers
+export * from "./socket/WebSocketTypes";
+
+// Metrics
+export * from "./server/GGWebSocketMetrics";
+
+// Context
+export * from "./server/GG_WS_CONNECTION";
+export * from "./server/GG_WS_MESSAGE";
+
+// Core
+export * from "./socket/GGSocket";
+
+// API Schema
+export * from "./schema/webSocketSchema";
+export * from "./schema/GGWebSocketSchema";
+export * from "./schema/GGWebSocketMiddleware";
+
+// Client
+export * from "./client/GGSocketClient";
+export * from "./client/GGSocketPool";

package/src/index-node.ts

@@ -0,0 +1,28 @@
+// Type helpers
+export * from "./socket/WebSocketTypes";
+
+// Metrics
+export * from "./server/GGWebSocketMetrics";
+
+// Context
+export * from "./server/GG_WS_CONNECTION";
+export * from "./server/GG_WS_MESSAGE";
+
+// Core
+export * from "./socket/GGSocket";
+
+// API Schema
+export * from "./schema/webSocketSchema";
+export * from "./schema/GGWebSocketSchema";
+export * from "./schema/GGWebSocketMiddleware";
+
+// Server
+export * from "./server/GGSocketServer";
+export * from "./server/GGWebSocketSchema.startServer";
+
+// Client
+export * from "./client/GGSocketClient";
+export * from "./client/GGSocketPool";
+
+// Extensions
+import "./server/GGWebSocketSchema.startServer";

package/src/schema/GGWebSocketMiddleware.ts

@@ -0,0 +1,43 @@
+/**
+ * WebSocket-specific middleware interface.
+ * Independent from @grest-ts/http - WebSocket has its own middleware system.
+ */
+
+/**
+ * Context available during WebSocket handshake.
+ */
+export interface GGWebSocketHandshakeContext {
+    /**
+     * Headers from the handshake message (client-sent or server-received).
+     */
+    headers: Record<string, string>;
+    /**
+     * Query parameters from the WebSocket URL.
+     */
+    queryArgs: Record<string, string>;
+}
+
+/**
+ * WebSocket middleware interface for both client and server.
+ * Unlike HTTP middleware, WebSocket middleware works with handshake data.
+ */
+export interface GGWebSocketMiddleware {
+    /**
+     * Client-side: Add headers to the handshake message before sending.
+     * Called by GGSocketPool when establishing a connection.
+     */
+    updateHandshake?(context: GGWebSocketHandshakeContext): void;
+
+    /**
+     * Server-side: Parse headers from the received handshake message.
+     * Called by GGSocketServer when receiving a connection.
+     */
+    parseHandshake?(context: GGWebSocketHandshakeContext): void;
+
+    /**
+     * Server-side: Async processing after parseHandshake.
+     * Can be used for authentication, authorization, etc.
+     * Throwing an error will reject the connection.
+     */
+    process?(): Promise<void>;
+}

package/src/schema/GGWebSocketSchema.ts

@@ -0,0 +1,57 @@
+import {GGWebSocketMiddleware} from "./GGWebSocketMiddleware";
+import {GGContractApiDefinition, GGContractClass} from "@grest-ts/schema";
+
+/**
+ * WebSocket API Schema - pure data definition with typed context
+ *
+ * Type parameters:
+ * - TClientToServer: Client-facing type for clientToServer methods (returns GGPromise)
+ * - TServerToClient: Client-facing type for serverToClient methods (returns GGPromise)
+ * - TContext: Accumulated context type (parseHandshake return types)
+ * - TQuery: Query parameters on connect
+ * - TClientToServerImpl: Server implementation type for clientToServer handlers (returns Promise)
+ */
+export class GGWebSocketSchema<
+    TClientToServer,
+    TServerToClient,
+    TContext = {},
+    TQuery = undefined,
+    TClientToServerImpl = TClientToServer
+> {
+    public readonly name: string
+    public readonly path: string
+    public readonly middlewares: readonly GGWebSocketMiddleware[]
+    private readonly contractFactory: () => GGWebSocketContractRuntime
+    private contractCache: GGWebSocketContractRuntime | null = null
+
+    constructor(
+        name: string,
+        path: string,
+        contractFactory: () => GGWebSocketContractRuntime,
+        middlewares: readonly GGWebSocketMiddleware[] = []
+    ) {
+        this.name = name
+        this.path = path
+        this.middlewares = middlewares
+        this.contractFactory = contractFactory
+        Object.freeze(this.middlewares)
+    }
+
+    get contract(): GGWebSocketContractRuntime {
+        if (!this.contractCache) {
+            this.contractCache = this.contractFactory()
+        }
+        Object.freeze(this)
+        return this.contractCache
+    }
+}
+
+/**
+ * Runtime contract structure for WebSocket APIs.
+ * Contains GGContractClass instances for proper validation via implement().
+ */
+export interface GGWebSocketContractRuntime {
+    apiName: string
+    clientToServer: GGContractClass<GGContractApiDefinition>
+    serverToClient: GGContractClass<GGContractApiDefinition>
+}

package/src/schema/webSocketSchema.ts

@@ -0,0 +1,109 @@
+import {GGWebSocketSchema, GGWebSocketContractRuntime} from "./GGWebSocketSchema";
+import {GGWebSocketMiddleware} from "./GGWebSocketMiddleware";
+import {GGContractClass, GGContractClient, GGContractImplementation, GGContractMethod} from "@grest-ts/schema";
+
+/**
+ * Bidirectional websocket contract methods
+ */
+export interface GGSocketContractMethods {
+    clientToServer: Record<string, GGContractMethod>
+    serverToClient: Record<string, GGContractMethod>
+}
+
+/**
+ * WebSocket contract definition
+ */
+export interface GGSocketContract<TDef extends GGSocketContractMethods = GGSocketContractMethods> {
+    name: string
+    methods: TDef
+}
+
+/**
+ * Create a websocket contract
+ */
+export function defineSocketContract<TDef extends GGSocketContractMethods>(
+    name: string,
+    methods: TDef
+): GGSocketContract<TDef> {
+    return {name, methods}
+}
+
+/**
+ * Create a WebSocket API schema builder from a contract.
+ *
+ * @example
+ * export const ChatContract = defineSocketContract("Chat", {
+ *     clientToServer: {
+ *         sendMessage: { input: IsMessage, success: IsVoid, errors: [SERVER_ERROR] }
+ *     },
+ *     serverToClient: {
+ *         onMessage: { input: IsMessage }
+ *     }
+ * })
+ *
+ * export const ChatApi = webSocketSchema(ChatContract)
+ *     .path("/chat")
+ *     .use(AuthMiddleware)
+ *     .done()
+ */
+export function webSocketSchema<TDef extends GGSocketContractMethods>(
+    contract: GGSocketContract<TDef>
+): GGWebSocketSchemaBuilder<
+    GGContractClient<TDef["clientToServer"]>,
+    GGContractClient<TDef["serverToClient"]>,
+    undefined,
+    undefined,
+    GGContractImplementation<TDef["clientToServer"]>
+> {
+    return new GGWebSocketSchemaBuilder(contract)
+}
+
+class GGWebSocketSchemaBuilder<
+    TClientToServer,
+    TServerToClient,
+    TContext = undefined,
+    TQuery = undefined,
+    TClientToServerImpl = TClientToServer
+> {
+    private _path: string = ""
+    private _middlewares: GGWebSocketMiddleware[] = []
+
+    constructor(
+        private readonly _contract: GGSocketContract
+    ) {
+    }
+
+    path(path: string): this {
+        this._path = path
+        return this
+    }
+
+    use<M extends GGWebSocketMiddleware>(middleware: M): GGWebSocketSchemaBuilder<TClientToServer, TServerToClient, TContext | M, TQuery, TClientToServerImpl> {
+        this._middlewares.push(middleware)
+        return this as any
+    }
+
+    queryOnConnect<TNewQuery>(): GGWebSocketSchemaBuilder<TClientToServer, TServerToClient, TContext, TNewQuery, TClientToServerImpl> {
+        return this as any
+    }
+
+    done(): GGWebSocketSchema<TClientToServer, TServerToClient, TContext, TQuery, TClientToServerImpl> {
+        const contract = this._contract;
+        const contractFactory = (): GGWebSocketContractRuntime => {
+            const methods = contract.methods;
+            const name = contract.name;
+            return {
+                apiName: name,
+                clientToServer: new GGContractClass(name + ".clientToServer", methods.clientToServer),
+                serverToClient: new GGContractClass(name + ".serverToClient", methods.serverToClient)
+            };
+        };
+
+        return new GGWebSocketSchema<TClientToServer, TServerToClient, TContext, TQuery, TClientToServerImpl>(
+            contract.name,
+            this._path,
+            contractFactory,
+            this._middlewares
+        )
+    }
+}