Add WebSocket transport

src/app.ts

@@ -3,6 +3,7 @@ import { Database, MigrationService } from './services/storage/index.ts';
 import { Broadcaster } from './services/broadcaster.ts';
 import { ApplicationRouter } from './services/router.ts';
 import { HttpTransportRouter } from './services/transport/http-transport.ts';
+import { WsTransportRouter } from './services/transport/ws-transport.ts';
 import { ServerHost } from './services/server-host.ts';
 import { Logger } from './utils/logger.ts';
 
@@ -31,7 +32,8 @@ export class App {
         const router = await ApplicationRouter.create(routes);
 
         const http = new HttpTransportRouter(router, debug);
-        const host = new ServerHost(config, debug, [http]);
+        const ws = new WsTransportRouter(router, debug, config.server.maxRequestBodyBytes);
+        const host = new ServerHost(config, debug, [http, ws]);
 
         return new App(host, database);
     }

src/services/stream/ws-stream.ts

@@ -0,0 +1,94 @@
+import { randomUUID } from 'node:crypto';
+import type { WSContext } from 'hono/ws';
+import { toExtendedJson } from '@xo-cash/utils';
+
+import { BaseStream, type StreamMessage } from './base-stream.ts';
+
+/** WebSocket readyState value indicating an open socket. */
+const OPEN = 1;
+
+/** Adapts one WebSocket connection to the shared connection-stream contract. */
+export class WSStream extends BaseStream {
+    readonly streaming = true;
+    readonly bidirectional = true;
+
+    /** Stable server-side identity used only for diagnostics. */
+    readonly id: string;
+
+    /** Observers notified when the socket closes. */
+    private closeCallbacks: Array<() => void> = [];
+
+    /** Guards against sends after local or remote closure. */
+    private closed = false;
+
+    /**
+     * @param ws - Minimal WebSocket surface required for send/close operations.
+     */
+    constructor(private readonly ws: Pick<WSContext, 'send' | 'close' | 'readyState'>) {
+        super();
+        this.id = randomUUID();
+    }
+
+    /**
+     * Encode and send one Extended JSON message frame.
+     *
+     * @param message - Neutral response or event envelope to send.
+     */
+    async send(message: StreamMessage): Promise<void> {
+        if (this.closed || this.ws.readyState !== OPEN) {
+            throw new Error('Cannot send to a closed WebSocket stream');
+        }
+
+        // WebSocket carries the complete neutral response or event envelope.
+        this.ws.send(toExtendedJson(message));
+    }
+
+    /** Close the socket from the server side and notify observers. */
+    close(): void {
+        if (this.closed) {
+            return;
+        }
+
+        this.closed = true;
+        this.ws.close();
+        this.emitClose();
+    }
+
+    /**
+     * Register a callback invoked when the socket closes.
+     *
+     * @param callback - Called immediately if the stream is already closed.
+     */
+    onClose(callback: () => void): void {
+        if (this.closed) {
+            callback();
+
+            return;
+        }
+
+        this.closeCallbacks.push(callback);
+    }
+
+    /**
+     * Record closure reported by WebSocket callbacks without closing the socket
+     * again. This shares the same observer notification path as local closure.
+     */
+    markClosed(): void {
+        if (this.closed) {
+            return;
+        }
+
+        this.closed = true;
+        this.emitClose();
+    }
+
+    /** Drain and invoke all registered close observers. */
+    private emitClose(): void {
+        const callbacks = this.closeCallbacks;
+        this.closeCallbacks = [];
+
+        for (const callback of callbacks) {
+            callback();
+        }
+    }
+}

src/services/transport/ws-transport.ts

@@ -0,0 +1,230 @@
+import { upgradeWebSocket } from '@hono/node-server';
+import type { WebSocketServerLike } from '@hono/node-server';
+import type { Hono } from 'hono';
+import type { WSContext, WSMessageReceive } from 'hono/ws';
+import { WebSocketServer } from 'ws';
+import { toExtendedJson, fromExtendedJson } from '@xo-cash/utils';
+import { z } from 'zod';
+
+import type { Logger } from '../../utils/logger.ts';
+import { ApplicationError, normalizePublicError } from '../../errors/index.ts';
+import type { ApplicationRequest, ApplicationRouter } from '../router.ts';
+import { WSStream } from '../stream/ws-stream.ts';
+import type { AppEnv, UpgradeTransportRouter } from './transport-router.ts';
+
+/** Default WebSocket upgrade path for application messages. */
+const WS_ROUTE = '/ws';
+
+// Strict validation prevents legacy or protocol-specific fields reaching routes.
+const wsRequestSchema = z
+    .object({
+        id: z.string().min(1).optional(),
+        path: z.string().min(1),
+        body: z.unknown().optional(),
+    })
+    .strict();
+
+/**
+ * WebSocket protocol adapter.
+ *
+ * One WSStream is shared by every message on a socket, allowing subscribe and
+ * unsubscribe requests to operate on the same broadcaster registration.
+ */
+export class WsTransportRouter implements UpgradeTransportRouter {
+    /**
+     * Native Node WebSocket server used by Hono's Node adapter.
+     *
+     * Hono handles the HTTP upgrade and lifecycle callbacks, while the `ws`
+     * server owns frame parsing and therefore enforces the payload limit.
+     */
+    private readonly wsServer: WebSocketServer;
+
+    /** Upgrade server wired into the Node HTTP listener by ServerHost. */
+    readonly websocketServer: WebSocketServerLike;
+
+    private readonly debug: Logger;
+
+    /**
+     * @param router - Shared application router for route dispatch.
+     * @param debug - Root logger extended with a ws-transport namespace.
+     * @param maxRequestBodyBytes - Maximum encoded WebSocket message size.
+     * @param url - WebSocket upgrade path exposed by this adapter.
+     */
+    constructor(
+        private readonly router: ApplicationRouter,
+        debug: Logger,
+        private readonly maxRequestBodyBytes: number,
+        private readonly url: string = WS_ROUTE,
+    ) {
+        this.debug = debug.extend('ws-transport');
+
+        // Hono's Node WebSocket helper delegates frame handling to `ws`; unlike
+        // HTTP requests, WebSocket messages do not pass through Hono middleware.
+        // maxPayload rejects oversized messages before onMessage receives them.
+        this.wsServer = new WebSocketServer({
+            noServer: true,
+            maxPayload: this.maxRequestBodyBytes,
+        });
+        this.websocketServer = this.wsServer as unknown as WebSocketServerLike;
+    }
+
+    /**
+     * Register the connection endpoint and translate socket lifecycle events.
+     *
+     * @param server - Hono instance to attach the upgrade handler to.
+     */
+    register(server: Hono<AppEnv>): void {
+        server.get(
+            this.url,
+            upgradeWebSocket(
+                () => {
+                    // One WSStream per socket, not per message. Broadcaster topic
+                    // subscriptions live on the connection stream so a later
+                    // /data/unsubscribe message can remove topics registered by an
+                    // earlier /data/subscribe on the same socket.
+                    let stream: WSStream | undefined;
+
+                    return {
+                        onOpen: (_event, ws): void => {
+                            stream = new WSStream(ws);
+                            this.debug('socket opened: %s', stream.id);
+                        },
+                        onMessage: (event, ws): void => {
+                            // onOpen may not run before the first message on some stacks.
+                            stream ??= new WSStream(ws);
+
+                            // Each WebSocket frame is an independent application request
+                            // (path + body + optional id). Dispatch is fire-and-forget so
+                            // multiple in-flight requests can share one connection.
+                            this.handleMessage(ws, stream, event.data).catch((error: unknown) => {
+                                this.debug('unexpected message failure: %O', error);
+                            });
+                        },
+                        onClose: (): void => {
+                            this.debug('socket closed: %s', stream?.id ?? 'unknown');
+                            // markClosed, not close — the socket is already gone. This
+                            // notifies the broadcaster's onClose hook so all topic
+                            // registrations for this connection are torn down.
+                            stream?.markClosed();
+                        },
+                        onError: (event): void => {
+                            this.debug('socket error: %O', event);
+                            stream?.markClosed();
+                        },
+                    };
+                },
+                { onError: (error) => this.debug('upgrade failed: %O', error) },
+            ),
+        );
+    }
+
+    /** Stop accepting upgrades and close active sockets during shutdown. */
+    async stop(): Promise<void> {
+        this.debug('closing %d active WebSocket(s)', this.wsServer.clients.size);
+        for (const client of this.wsServer.clients) {
+            client.close(1001, 'Server shutting down');
+        }
+
+        await new Promise<void>((resolve, reject) => {
+            this.wsServer.close((error) => (error ? reject(error) : resolve()));
+        });
+    }
+
+    /**
+     * Dispatch one independently correlated message.
+     *
+     * Route failures become error envelopes and intentionally leave the socket open.
+     *
+     * @param ws - Active WebSocket context for error replies.
+     * @param stream - Shared connection stream for this socket.
+     * @param message - Raw WebSocket payload from the client.
+     */
+    private async handleMessage(ws: WSContext, stream: WSStream, message: WSMessageReceive): Promise<void> {
+        let request: ApplicationRequest;
+
+        try {
+            // Decode the transport envelope { id?, path, body? } and revive any
+            // Extended JSON values before the message reaches application routes.
+            request = await WsTransportRouter.decodeWebSocketRequest(message);
+        } catch (error) {
+            this.debug('invalid message: %O', error);
+            // Malformed envelopes have no request id to echo back.
+            this.sendError(ws, undefined, error);
+
+            return;
+        }
+
+        try {
+            // ApplicationRouter wraps the shared connection stream in a per-request
+            // RouteStream (body + requestId). A subscription dispatch intentionally
+            // remains pending until a later request unsubscribes or the socket closes.
+            await this.router.dispatch(request, stream);
+        } catch (error) {
+            this.debug('dispatch failed for %s on socket %s: %O', request.path, stream.id, error);
+            // Reply with a correlated error envelope but keep the socket open.
+            // The client may have other in-flight requests or active subscriptions
+            // on this connection that must survive a single route failure.
+            this.sendError(ws, request.requestId, error);
+        }
+    }
+
+    /**
+     * Serialize the shared public error contract without exposing internals.
+     *
+     * @param ws - Active WebSocket context to send the error on.
+     * @param requestId - Optional correlation ID echoed back to the client.
+     * @param error - Failure to normalize into the public error shape.
+     */
+    private sendError(ws: WSContext, requestId: string | undefined, error: unknown): void {
+        ws.send(
+            toExtendedJson({
+                ...(requestId === undefined ? {} : { id: requestId }),
+                type: 'error',
+                ...normalizePublicError(error),
+            }),
+        );
+    }
+
+    /**
+     * Decode the WebSocket message into an application request.
+     *
+     * @param message - Raw WebSocket payload from the client.
+     * @returns The decoded WebSocket message as an application request.
+     */
+    static async decodeWebSocketRequest(message: WSMessageReceive): Promise<ApplicationRequest> {
+        const payload = await WsTransportRouter.decodeWebSocketMessage(message);
+        let decoded: unknown;
+
+        try {
+            decoded = fromExtendedJson(payload);
+        } catch {
+            throw new ApplicationError(400, 'Invalid JSON in WebSocket message');
+        }
+
+        const envelope = wsRequestSchema.parse(decoded);
+
+        return {
+            path: envelope.path,
+            ...(envelope.id === undefined ? {} : { requestId: envelope.id }),
+            ...(envelope.body === undefined ? {} : { body: envelope.body }),
+        };
+    }
+
+    /**
+     * Decode the WebSocket message into a string.
+     *
+     * @param message - Raw WebSocket payload from the client.
+     * @returns The decoded WebSocket message as a string.
+     */
+    static async decodeWebSocketMessage(message: WSMessageReceive): Promise<string> {
+        if (typeof message === 'string') {
+            return message;
+        }
+
+        if (message instanceof Blob) {
+            return Buffer.from(await message.arrayBuffer()).toString('utf8');
+        }
+
+        return Buffer.from(message).toString('utf8');
+    }
+}

tests/services/stream/stream.test.ts

@@ -0,0 +1,62 @@
+import { describe, expect, it, vi } from "vitest";
+
+import { HonoSSEStream } from "../../../src/services/stream/hono-sse-stream.js";
+import { HttpRequestStream } from "../../../src/services/stream/http-request-stream.js";
+import { WSStream } from "../../../src/services/stream/ws-stream.js";
+
+describe("stream lifecycle observers", () => {
+  it("buffers exactly one normal HTTP response", async () => {
+    const stream = new HttpRequestStream();
+
+    await stream.send({
+      type: "response",
+      statusCode: 200,
+      body: { ok: true },
+    });
+
+    expect(stream.getResponse()).toEqual({
+      type: "response",
+      statusCode: 200,
+      body: { ok: true },
+    });
+    await expect(
+      stream.send({
+        type: "response",
+        statusCode: 200,
+        body: { second: true },
+      }),
+    ).rejects.toThrow("only send one response");
+  });
+
+  it("notifies WebSocket observers registered after remote closure", () => {
+    const stream = new WSStream({
+      send: vi.fn(),
+      close: vi.fn(),
+      readyState: 1,
+    });
+    const onClose = vi.fn();
+
+    stream.markClosed();
+    stream.onClose(onClose);
+
+    expect(onClose).toHaveBeenCalledOnce();
+  });
+
+  it("notifies SSE observers registered after local closure", () => {
+    const streamApi = {
+      writeSSE: vi.fn(),
+      close: vi.fn(),
+    };
+    const stream = new HonoSSEStream(
+      streamApi as unknown as ConstructorParameters<typeof HonoSSEStream>[0],
+    );
+    const onClose = vi.fn();
+
+    stream.close();
+    stream.onClose(onClose);
+    stream.close();
+
+    expect(streamApi.close).toHaveBeenCalledOnce();
+    expect(onClose).toHaveBeenCalledOnce();
+  });
+});

tests/services/transports/ws-transport.test.ts

@@ -0,0 +1,58 @@
+import { describe, expect, it } from "vitest";
+import { z } from "zod";
+import { WebSocketServer } from "ws";
+
+import { ApplicationRouter } from "../../../src/services/router.js";
+import {
+  WsTransportRouter,
+} from "../../../src/services/transport/ws-transport.js";
+import { Logger } from "../../../src/utils/logger.js";
+import { toExtendedJson } from "@xo-cash/utils";
+
+describe("WebSocket request decoding", () => {
+  it("decodes the minimal route-agnostic envelope and Extended JSON body", async () => {
+    await expect(
+      WsTransportRouter.decodeWebSocketRequest(
+        toExtendedJson({
+          id: "request-1",
+          path: "/data/write",
+          body: { value: new Uint8Array([1, 2, 3]) },
+        }),
+      ),
+    ).resolves.toEqual({
+      requestId: "request-1",
+      path: "/data/write",
+      body: { value: new Uint8Array([1, 2, 3]) },
+    });
+  });
+
+  it.each([
+    "{}",
+    '{"path":42}',
+    '{"path":"/data/get","id":1}',
+    '{"path":"/data/get","method":"POST"}',
+  ])("rejects an invalid envelope: %s", async (payload) => {
+    await expect(WsTransportRouter.decodeWebSocketRequest(payload)).rejects.toBeInstanceOf(
+      z.ZodError,
+    );
+  });
+
+  it("rejects malformed JSON", async () => {
+    await expect(WsTransportRouter.decodeWebSocketRequest("{")).rejects.toMatchObject({
+      statusCode: 400,
+      message: "Invalid JSON in WebSocket message",
+    });
+  });
+});
+
+describe("WsTransportRouter payload limits", () => {
+  it("configures Hono's ws server with the requested maxPayload", async () => {
+    const debug = new Logger("ws-transport-test");
+    const router = await ApplicationRouter.create([]);
+    const transport = new WsTransportRouter(router, debug, 1024);
+    const wsServer = transport.websocketServer as unknown as WebSocketServer;
+
+    expect(wsServer.options.maxPayload).toBe(1024);
+    await transport.stop();
+  });
+});