@grest-ts/websocket 0.0.5

package/src/server/GGSocketServer.ts

@@ -0,0 +1,217 @@
+/**
+ * Manages WebSocket connections for a single endpoint
+ * Handles connection acceptance, effect parsing via handshake, and socket lifecycle
+ */
+
+import {WebSocket, WebSocketServer} from "ws";
+import {ERROR, GGValidator} from "@grest-ts/schema";
+import {NodeSocketAdapter} from "../adapter/NodeSocketAdapter";
+import * as http from "http";
+import * as url from "url";
+import {GGLog} from "@grest-ts/logger";
+import {GG_WS_CONNECTION} from "./GG_WS_CONNECTION";
+import {GGWebSocketMetrics} from "./GGWebSocketMetrics";
+import {Message, MessageType} from "../socket/SocketMessage";
+import {GG_TRACE} from "@grest-ts/trace";
+import {GGSocket} from "../socket/GGSocket";
+import {GGLocator, GGLocatorScope} from "@grest-ts/locator";
+import {GG_METRICS} from "@grest-ts/metrics";
+import {withTimeout} from "@grest-ts/common";
+import {GGContext} from "@grest-ts/context";
+import {GG_DISCOVERY} from "@grest-ts/discovery";
+import {GGWebSocketHandshakeContext, GGWebSocketMiddleware} from "../schema/GGWebSocketMiddleware";
+import {GGHttpServer} from "@grest-ts/http";
+
+export interface GGSocketServerConfig<TContext, Query> {
+    path: string;
+    apiName: string;
+    queryValidator?: GGValidator<Query>;
+    middlewares: readonly GGWebSocketMiddleware[];
+}
+
+export class GGSocketServer<TContext, Query> {
+
+    private readonly wss: WebSocketServer;
+    private readonly http: GGHttpServer;
+    public readonly path: string;
+    private readonly apiName: string;
+    private readonly middlewares: readonly GGWebSocketMiddleware[];
+    private readonly queryValidator: GGValidator<Query>;
+
+    private readonly activeSockets: Set<GGSocket> = new Set();
+    private readonly onConnectionHandlers: Array<(socket: GGSocket, query: Query) => void> = [];
+
+    // Capture context at construction - WebSocket events lose AsyncLocalStorage context
+    private readonly scope: GGLocatorScope;
+
+    constructor(http: GGHttpServer, config: GGSocketServerConfig<TContext, Query>) {
+        this.scope = GGLocator.getScope();
+        this.path = config.path;
+        this.apiName = config.apiName;
+        this.middlewares = config.middlewares;
+        this.queryValidator = config.queryValidator;
+        this.wss = new WebSocketServer({server: http.httpServer, path: this.path});
+        this.wss.on('connection', this.scope.wrapWithEnter(this._onConnection));
+        this.http = http
+            .onStart(() => {
+                GGLog.info(this, "WebSocket server started", {api: config.apiName, path: config.path});
+                if (GG_DISCOVERY.has()) {
+                    GG_DISCOVERY.get().registerRoutes([{
+                        runtime: this.scope.serviceName,
+                        api: this.apiName,
+                        pathPrefix: this.path,
+                        protocol: "ws",
+                        port: http.port
+                    }]);
+                }
+            })
+            .onTeardown(async () => {
+                GGLog.info(this, "WebSocket server closing");
+                try {
+                    this.wss.close(); // We use external http, so this does not close existing connections.
+                } catch (error) {
+                    GGLog.error(this, error);
+                }
+                await Promise.allSettled(Array.from(this.activeSockets).map(socket => socket.teardown()));
+                GGLog.info(this, "WebSocket server stopped");
+            });
+    }
+
+    public onConnection(handler: (socket: GGSocket, query: Query) => void): void {
+        this.onConnectionHandlers.push(handler);
+    }
+
+    private _onConnection = async (ws: WebSocket, req: http.IncomingMessage) => {
+        const connectionLabels = {api: this.apiName, path: this.path};
+
+        const context = new GGContext("ws-connection");
+        await context.run(async () => {
+            GG_TRACE.init()
+            GG_WS_CONNECTION.set({
+                port: this.http.port,
+                path: this.path
+            })
+            try {
+                // Parse and validate query parameters
+                const parsedUrl = url.parse(req.url || '', true);
+                let queryArgs = parsedUrl.query as Query;
+
+                if (this.queryValidator) {
+                    const result = this.queryValidator.safeParse(queryArgs, true)
+                    if (result.success === false) {
+                        ws.close(4000, "Invalid query parameters");
+                        GGLog.warn(this, "REJECTED - bad query", result.issues);
+                        if (GG_METRICS.has()) GGWebSocketMetrics.connections.inc(1, {...connectionLabels, result: 'QUERY_INVALID'});
+                        return;
+                    }
+                    queryArgs = result.value as Query
+                }
+
+                const adapter = new NodeSocketAdapter(ws);
+
+                // Wait for handshake message with headers
+                const handshakeResult = await this.handleHandshake(context, adapter, queryArgs);
+
+                if (!handshakeResult.success) {
+                    GGLog.warn(this, "REJECTED - handshake failed", (handshakeResult as { success: false; error: any }).error);
+                    if (GG_METRICS.has()) GGWebSocketMetrics.connections.inc(1, {...connectionLabels, result: 'HANDSHAKE_FAILED'});
+                    adapter.send(Message.create(MessageType.HANDSHAKE_ERR, "", "", (handshakeResult as { success: false; error: any }).error));
+                    ws.close(4001, "Handshake failed");
+                    return;
+                }
+
+                // Send handshake success
+                adapter.send(Message.create(MessageType.HANDSHAKE_OK, "", "", null));
+
+                GGLog.debug(this, "New websocket connection", queryArgs);
+                const socket = new GGSocket(adapter, {
+                    apiName: this.apiName,
+                    socketPath: this.path,
+                    connectionContext: context
+                });
+                this.activeSockets.add(socket);
+
+                // Track connection metrics
+                if (GG_METRICS.has()) {
+                    GGWebSocketMetrics.connections.inc(1, {...connectionLabels, result: 'OK'});
+                    GGWebSocketMetrics.connectionsActive.inc(1, connectionLabels);
+                }
+
+                socket.onClose(() => {
+                    this.activeSockets.delete(socket);
+                    if (GG_METRICS.has()) GGWebSocketMetrics.connectionsActive.dec(1, connectionLabels);
+                });
+
+                // Run connection handlers inside the connectionScope so they can access context
+                for (const handler of this.onConnectionHandlers) {
+                    try {
+                        handler(socket, queryArgs);
+                    } catch (error) {
+                        GGLog.error(this, error);
+                    }
+                }
+            } catch (error) {
+                GGLog.error(this, error);
+                if (GG_METRICS.has()) GGWebSocketMetrics.connections.inc(1, {...connectionLabels, result: 'ERROR'});
+                // Only close if socket is still open (fix #5 - avoid double close)
+                if (ws.readyState === ws.OPEN || ws.readyState === ws.CONNECTING) {
+                    ws.close(4001, "Connection rejected");
+                }
+            }
+        });
+    }
+
+    /**
+     * Handle handshake message from client.
+     * Waits for HANDSHAKE message with headers, runs effects, returns context.
+     */
+    private async handleHandshake(
+        context: GGContext,
+        adapter: NodeSocketAdapter,
+        queryArgs: Query
+    ): Promise<{ success: true } | { success: false; error: any }> {
+        type HandshakeResult = { success: true } | { success: false; error: any };
+        return withTimeout<HandshakeResult>(
+            new Promise<HandshakeResult>((resolve) => {
+                // Wrap with scope.wrapWithEnter to preserve locator scope when callback fires
+                const onMessage = this.scope.wrapWithEnter((data: string) => {
+                    const msg = Message.parse(data);
+                    if (!msg) return;
+
+                    if (msg.type === MessageType.HANDSHAKE) {
+                        adapter.offMessage(onMessage);
+                        context.run(async () => {
+                            try {
+                                // Build handshake context from headers
+                                const headers = msg.data || {};
+                                const handshakeContext: GGWebSocketHandshakeContext = {
+                                    headers,
+                                    queryArgs: queryArgs as Record<string, string>
+                                };
+
+                                // Run middlewares
+                                for (const middleware of this.middlewares) {
+                                    middleware.parseHandshake?.(handshakeContext);
+                                    await middleware.process?.();
+                                }
+                                resolve({success: true});
+                            } catch (error: any) {
+                                const errorJson = error instanceof ERROR
+                                    ? error.toJSON()
+                                    : {message: String(error)};
+                                resolve({success: false, error: errorJson});
+                            }
+                        });
+                    }
+                });
+
+                adapter.onMessage(onMessage);
+            }),
+            5000,
+            'Handshake timeout'
+        ).catch((error): { success: false; error: any } => ({
+            success: false,
+            error: {message: error.message || 'Handshake timeout'}
+        }));
+    }
+}

package/src/server/GGWebSocketMetrics.ts

@@ -0,0 +1,58 @@
+import {GGCounterKey, GGGaugeKey, GGHistogramKey, GGMetrics} from "@grest-ts/metrics";
+import {EXISTS, FORBIDDEN, NOT_AUTHORIZED, NOT_FOUND, OK, ROUTE_NOT_FOUND, SERVER_ERROR, VALIDATION_ERROR} from "@grest-ts/schema";
+
+/**
+ * Result type for WebSocket operations.
+ * Known types listed for documentation, but accepts any string for custom error types.
+ */
+type ResultType =
+    | typeof OK.TYPE
+    | typeof VALIDATION_ERROR.TYPE
+    | typeof NOT_AUTHORIZED.TYPE
+    | typeof FORBIDDEN.TYPE
+    | typeof NOT_FOUND.TYPE
+    | typeof ROUTE_NOT_FOUND.TYPE
+    | typeof EXISTS.TYPE
+    | typeof SERVER_ERROR.TYPE
+    | string;
+
+type SocketConnectionResult = 'OK' | 'AUTH_FAILED' | 'QUERY_INVALID' | string;
+
+export const GGWebSocketMetrics = GGMetrics.define('/websocket/', () => ({
+    connectionsActive: new GGGaugeKey<{ api: string, path: string }>('connections_active', {
+        help: 'Active WebSocket connections',
+        labelNames: ['api', 'path'],
+        groupBy: {labels: ["api"]}
+    }),
+    connections: new GGCounterKey<{ api: string, path: string, result: SocketConnectionResult }>('connections_total', {
+        help: 'Total WebSocket connection attempts',
+        labelNames: ['api', 'path', 'result'],
+        groupBy: {labels: ["api"]}
+    }),
+
+    // Incoming (server handling REQ + MSG)
+    requests: new GGCounterKey<{ api: string, path: string, method: string, result: ResultType }>('requests_total', {
+        help: 'Total incoming WebSocket commands (REQ + MSG)',
+        labelNames: ['api', 'path', 'method', 'result'],
+        groupBy: {labels: ["api", "method"], template: "{api}.{method}"}
+    }),
+    requestDuration: new GGHistogramKey<{ api: string, path: string, method: string }>('request_duration_ms', {
+        help: 'Incoming WebSocket command processing duration in milliseconds',
+        labelNames: ['api', 'path', 'method'],
+        buckets: [1, 5, 10, 25, 50, 100, 250, 500, 1000],
+        groupBy: {labels: ["api", "method"], template: "{api}.{method}"}
+    }),
+
+    // Outgoing (client calls via socket)
+    outRequests: new GGCounterKey<{ api: string, path: string, method: string, result: ResultType }>('out_requests_total', {
+        help: 'Total outgoing WebSocket commands (REQ + MSG)',
+        labelNames: ['api', 'path', 'method', 'result'],
+        groupBy: {labels: ["api", "method"], template: "{api}.{method}"}
+    }),
+    outRequestDuration: new GGHistogramKey<{ api: string, path: string, method: string }>('out_request_duration_ms', {
+        help: 'Outgoing WebSocket request round-trip duration in milliseconds (REQ only)',
+        labelNames: ['api', 'path', 'method'],
+        buckets: [1, 5, 10, 25, 50, 100, 250, 500, 1000, 2500],
+        groupBy: {labels: ["api", "method"], template: "{api}.{method}"}
+    }),
+}));

package/src/server/GGWebSocketSchema.startServer.ts

@@ -0,0 +1,136 @@
+/**
+ * Server extension for WebSocketSchema - adds startServer and register methods
+ * This file should only be imported in server (Node.js) context
+ */
+
+import type {GGSocket} from "../socket/GGSocket"
+import {GGWebSocketSchema} from "../schema/GGWebSocketSchema";
+import {GGWebSocketMiddleware} from "../schema/GGWebSocketMiddleware";
+import {GGSocketServer} from "./GGSocketServer";
+import {GGLocator} from "@grest-ts/locator";
+import {WebSocketIncoming, WebSocketOutgoing} from "../socket/WebSocketTypes";
+import {GG_HTTP_SERVER, GGHttpServer} from "@grest-ts/http";
+
+export interface WebSocketSchemaConfig {
+    /**
+     * The HTTP server adapter to attach the WebSocket server to.
+     * If not provided, will look up from locator.
+     */
+    http?: GGHttpServer;
+    /**
+     * Additional middlewares to apply to all connections.
+     */
+    middlewares?: GGWebSocketMiddleware[];
+}
+
+declare module "../schema/GGWebSocketSchema" {
+    interface GGWebSocketSchema<TClientToServer, TServerToClient, TContext = {}, TQuery = undefined, TClientToServerImpl = TClientToServer> {
+        /**
+         * Start the WebSocket server for this API.
+         */
+        startServer(
+            onConnection: (incoming: WebSocketIncoming<TClientToServerImpl>, outgoing: WebSocketOutgoing<TServerToClient>) => void,
+            config: WebSocketSchemaConfig
+        ): GGSocketServer<TContext, TQuery>
+
+        /**
+         * Register this WebSocket API with the default HTTP server.
+         * Uses GGHttpServerAdapter from locator if not explicitly provided.
+         */
+        register(
+            onConnection: (incoming: WebSocketIncoming<TClientToServerImpl>, outgoing: WebSocketOutgoing<TServerToClient>) => void,
+            config?: WebSocketSchemaConfig
+        ): void
+    }
+}
+
+GGWebSocketSchema.prototype.startServer = function (
+    this: GGWebSocketSchema<any, any, any, any>,
+    onConnection: any,
+    config: WebSocketSchemaConfig
+): GGSocketServer<any, any> {
+    const contract = this.contract
+    if (!contract) {
+        throw new Error(`WebSocketSchema "${this.name}" has no contract.`)
+    }
+
+    const normalizedPath = this.path.startsWith('/') ? this.path : '/' + this.path
+    const schemaName = this.name
+    const http = config.http ?? GGLocator.getScope().get(GG_HTTP_SERVER);
+
+    // @TODO We might want some lookup here based on path/middlewares etc. If I use same socket for multiple paths, we need to reuse also same GGSocketServer.
+    const socketServer = new GGSocketServer(http, {
+        apiName: schemaName,
+        path: normalizedPath,
+        middlewares: [...this.middlewares, ...(config?.middlewares ?? [])]
+    });
+
+    socketServer.onConnection((socket: GGSocket) => {
+        const clientToServerContract = contract.clientToServer
+        const serverToClientContract = contract.serverToClient
+
+        const incoming: any = {
+            on(handlers: any) {
+                const impl: Record<string, any> = {};
+                for (const methodName of Object.keys(clientToServerContract.methods)) {
+                    const methodDef = clientToServerContract.methods[methodName] as any;
+                    const params = methodDef.params;
+                    impl[methodName] = (data: any) => {
+                        // If method has params info, unpack data object to positional args
+                        if (params && params.length > 0 && data && typeof data === 'object') {
+                            const args = params.map((p: any) => data[p.name]);
+                            return handlers[methodName](...args);
+                        }
+                        // Single or no parameter - pass directly
+                        return handlers[methodName](data);
+                    };
+                }
+
+                const incomingInstance = clientToServerContract.implement(impl);
+
+                for (const methodName of Object.keys(clientToServerContract.methods)) {
+                    socket.registerHandler({
+                        path: `${schemaName}.${methodName}`,
+                        handler: (incomingInstance as any)[methodName]
+                    });
+                }
+            }
+        }
+
+        const impl: Record<string, any> = {};
+        for (const methodName of Object.keys(serverToClientContract.methods)) {
+            const method = serverToClientContract.methods[methodName];
+            const expectsResponse = 'success' in method;
+            impl[methodName] = (data: any) => {
+                return socket.send(`${schemaName}.${methodName}`, data, expectsResponse);
+            };
+        }
+
+        const outgoingInstance = serverToClientContract.implement(impl);
+        (outgoingInstance as any).onClose = (callback: () => void) => {
+            socket.onClose(callback)
+        }
+        onConnection(incoming, outgoingInstance)
+    });
+
+    return socketServer;
+}
+
+GGWebSocketSchema.prototype.register = function (
+    this: GGWebSocketSchema<any, any, any, any>,
+    onConnection: any,
+    config?: WebSocketSchemaConfig
+): void {
+    let httpServer = config?.http;
+    if (!httpServer) {
+        httpServer = GGLocator.getScope().get(GG_HTTP_SERVER);
+    }
+    if (!httpServer) {
+        throw new Error(`No HTTP server found. Make sure to register GGHttpServerAdapter in the scope or pass it via config`)
+    }
+
+    this.startServer(onConnection, {
+        http: httpServer,
+        middlewares: config?.middlewares
+    });
+}

package/src/server/GG_WS_CONNECTION.ts

@@ -0,0 +1,10 @@
+import {GGContextKey} from "@grest-ts/context";
+import {IsNumber, IsObject, IsString} from "@grest-ts/schema";
+
+const IsWsConnectionContext = IsObject({
+    port: IsNumber.orUndefined,
+    path: IsString
+});
+export type WsConnectionContext = typeof IsWsConnectionContext.infer;
+
+export const GG_WS_CONNECTION = new GGContextKey<WsConnectionContext>('ws-connection', IsWsConnectionContext);

package/src/server/GG_WS_MESSAGE.ts

@@ -0,0 +1,9 @@
+import {GGContextKey} from "@grest-ts/context";
+import {IsObject, IsString} from "@grest-ts/schema";
+
+const IsWsMessageContext = IsObject({
+    path: IsString
+});
+export type WsMessageContext = typeof IsWsMessageContext.infer;
+
+export const GG_WS_MESSAGE = new GGContextKey<WsMessageContext>('ws-message', IsWsMessageContext);