@procedures/client 0.0.0-stage → 0.2.0

package/src/adapters/websocket/server/transport.ts

@@ -0,0 +1,705 @@
+/**
+ * WebSocket Server Transport Implementation
+ *
+ * Adapts WebSocket connections to unified RPC format.
+ * Supports bidirectional streaming and persistent connections.
+ */
+
+import type { IncomingMessage } from "http";
+import { WebSocketServer, WebSocket } from "ws";
+import type { ServerTransport, ServerRequest, ServerResponse } from "../../../server/types.js";
+import type { Server } from "../../../server/server.js";
+import type {
+  WebSocketServerTransportOptions,
+  WebSocketMessage,
+  TrackedConnection,
+  ConnectionEventHandler,
+} from "./types.js";
+import { withoutInternalKeys } from "../../metadata.js";
+import { checkBrowserRequest, type BrowserGuardOptions } from "../../../server/browser-guard.js";
+
+/**
+ * WebSocket server transport adapter.
+ *
+ * Provides persistent, bidirectional RPC over WebSocket.
+ * Perfect for real-time updates and streaming responses.
+ *
+ * @example
+ * ```typescript
+ * import { createServer } from "http";
+ * const httpServer = createServer();
+ *
+ * const server = new Server();
+ * const wsTransport = new WebSocketServerTransport(server, {
+ *   server: httpServer,
+ *   path: "/ws",
+ *   authenticate: async (req) => {
+ *     // Verify token from query string or headers
+ *     return true;
+ *   }
+ * });
+ *
+ * await wsTransport.start();
+ * httpServer.listen(3000);
+ * ```
+ */
+/**
+ * The credit of one stream: the number of items the client can take now. The client sends
+ * more credit as its reader takes items (deep dive TRN-3).
+ */
+class StreamCredit {
+  private wake: (() => void) | undefined;
+
+  constructor(private credit: number) {}
+
+  grant(n: number): void {
+    this.credit += n;
+    const wake = this.wake;
+    this.wake = undefined;
+    wake?.();
+  }
+
+  /** Take one item of credit. Returns false when the signal aborted the wait. */
+  async take(signal: AbortSignal): Promise<boolean> {
+    while (this.credit <= 0) {
+      if (signal.aborted) return false;
+      await new Promise<void>((resolve) => {
+        const onAbort = (): void => resolve();
+        signal.addEventListener("abort", onAbort, { once: true });
+        this.wake = () => {
+          signal.removeEventListener("abort", onAbort);
+          resolve();
+        };
+      });
+    }
+    if (signal.aborted) return false;
+    this.credit--;
+    return true;
+  }
+}
+
+/** The error of a failed stream or handler: its code (a string) and its retryable flag. */
+function errorStatus(error: unknown): { type: "error"; code: string; message: string; retryable: boolean } {
+  const code = (error as { code?: unknown } | undefined)?.code;
+  const retryable = (error as { retryable?: unknown } | undefined)?.retryable;
+  return {
+    type: "error",
+    code: typeof code === "string" ? code : "HANDLER_ERROR",
+    message: error instanceof Error ? error.message : String(error),
+    // Before, every stream error lost its retryable flag (deep dive TRN-7)
+    retryable: typeof retryable === "boolean" ? retryable : false,
+  };
+}
+
+interface PendingRequest {
+  id: string;
+  /** The connection the request was sent to: only it may answer, and its close ends the request */
+  connectionId: string;
+  resolve: (value: unknown) => void;
+  reject: (error: Error) => void;
+  timeout: NodeJS.Timeout;
+}
+
+export class WebSocketServerTransport implements ServerTransport {
+  readonly name = "websocket";
+
+  private wss: WebSocketServer | null = null;
+  private options: Required<
+    Omit<WebSocketServerTransportOptions, "authenticate" | "onConnection" | "allowedOrigins" | "allowedHosts">
+  > & {
+    authenticate?: WebSocketServerTransportOptions["authenticate"];
+    onConnection?: WebSocketServerTransportOptions["onConnection"];
+  };
+  private server: Server;
+  private connections: Set<WebSocket> = new Set();
+
+  // Tracked connections with IDs for bidirectional RPC
+  private trackedConnections = new Map<string, TrackedConnection>();
+  private pendingRequests = new Map<string, PendingRequest>();
+  private connectionCounter = 0;
+
+  // The requests in progress, by "connectionId:requestId": a cancel message or a closed
+  // connection aborts their signal, so a stream stops and its handler gets return()
+  private inflight = new Map<string, AbortController>();
+
+  // The flow control of the streams in progress, by the same key
+  private credits = new Map<string, StreamCredit>();
+
+  /** The browser checks of each connection (see `server/browser-guard.ts`) */
+  private readonly guard: BrowserGuardOptions;
+
+  // Connection lifecycle event handlers
+  private connectHandlers: ConnectionEventHandler[] = [];
+  private disconnectHandlers: ConnectionEventHandler[] = [];
+
+  constructor(server: Server, options: WebSocketServerTransportOptions) {
+    this.server = server;
+    this.options = {
+      server: options.server,
+      path: options.path ?? "/ws",
+      authenticate: options.authenticate,
+      onConnection: options.onConnection,
+      perMessageDeflate: options.perMessageDeflate ?? false,
+      maxPayload: options.maxPayload ?? 1024 * 1024, // 1MB default
+      clientTracking: options.clientTracking ?? true,
+      highWaterMark: options.highWaterMark ?? 1024 * 1024,
+    };
+    this.guard = { allowedOrigins: options.allowedOrigins, allowedHosts: options.allowedHosts };
+  }
+
+  async start(): Promise<void> {
+    // Create WebSocket server
+    this.wss = new WebSocketServer({
+      server: this.options.server,
+      path: this.options.path,
+      perMessageDeflate: this.options.perMessageDeflate,
+      maxPayload: this.options.maxPayload,
+      clientTracking: this.options.clientTracking,
+      // The browser checks run for every connection: before, a connection was checked only
+      // when the host gave an authenticate handler, so any web page could open one and call
+      // procedures (cross-site WebSocket hijacking)
+      verifyClient: async (
+        { req }: { req: IncomingMessage },
+        callback: (result: boolean, code?: number, message?: string) => void
+      ) => {
+        const guard = checkBrowserRequest(req, this.guard);
+        if (!guard.ok) {
+          callback(false, guard.status, guard.message);
+          return;
+        }
+        if (!this.options.authenticate) {
+          callback(true);
+          return;
+        }
+        try {
+          callback(await this.options.authenticate(req));
+        } catch (error) {
+          console.error("[WebSocket] Authentication error:", error);
+          callback(false);
+        }
+      },
+    });
+
+    // Handle connections
+    this.wss.on("connection", (ws: WebSocket, req: IncomingMessage) => {
+      this.handleConnection(ws, req);
+    });
+
+    console.log(`[${this.name}] WebSocket server listening on ${this.options.path}`);
+  }
+
+  async stop(): Promise<void> {
+    if (this.wss) {
+      // Stop the requests in progress, so their handlers end (deep dive TRN-15)
+      for (const controller of this.inflight.values()) controller.abort();
+      this.inflight.clear();
+      for (const connectionId of this.trackedConnections.keys()) {
+        this.rejectPendingFor(connectionId, "Server stopped");
+      }
+
+      // Close all connections
+      for (const ws of this.connections) {
+        ws.close();
+      }
+      this.connections.clear();
+
+      // Close server
+      return new Promise((resolve, reject) => {
+        this.wss!.close((err?: Error) => {
+          if (err) {
+            reject(err);
+          } else {
+            console.log(`[${this.name}] WebSocket server stopped`);
+            this.wss = null;
+            resolve();
+          }
+        });
+      });
+    }
+  }
+
+  isRunning(): boolean {
+    return this.wss !== null;
+  }
+
+  /**
+   * Handle new WebSocket connection.
+   */
+  private handleConnection(ws: WebSocket, req: IncomingMessage): void {
+    // Track connection (legacy)
+    this.connections.add(ws);
+
+    // Create tracked connection with ID
+    const connectionId = `conn_${++this.connectionCounter}_${Date.now()}`;
+    const trackedConnection: TrackedConnection = {
+      id: connectionId,
+      socket: ws,
+      connectedAt: new Date(),
+      metadata: {
+        remoteAddress: req.socket.remoteAddress,
+        headers: req.headers,
+      },
+    };
+    this.trackedConnections.set(connectionId, trackedConnection);
+
+    // Call user connection handler (legacy)
+    if (this.options.onConnection) {
+      this.options.onConnection(ws, req);
+    }
+
+    // Emit connect event
+    for (const handler of this.connectHandlers) {
+      try {
+        handler(trackedConnection);
+      } catch (error) {
+        console.error("[WebSocket] Connect handler error:", error);
+      }
+    }
+
+    // Handle messages
+    ws.on("message", async (data: Buffer) => {
+      try {
+        await this.handleMessage(ws, data, req, connectionId);
+      } catch (error) {
+        console.error("[WebSocket] Message handling error:", error);
+        this.sendError(ws, "unknown", "Internal server error");
+      }
+    });
+
+    // Handle close
+    ws.on("close", () => {
+      this.connections.delete(ws);
+      this.rejectPendingFor(connectionId, "Connection closed");
+      this.abortInflightFor(connectionId);
+      const conn = this.trackedConnections.get(connectionId);
+      if (conn) {
+        this.trackedConnections.delete(connectionId);
+        // Emit disconnect event
+        for (const handler of this.disconnectHandlers) {
+          try {
+            handler(conn);
+          } catch (error) {
+            console.error("[WebSocket] Disconnect handler error:", error);
+          }
+        }
+      }
+    });
+
+    // Handle errors. The socket emits "close" after "error": the close handler does the
+    // cleanup and the disconnect event (deep dive TRN-16: before, this handler removed the
+    // connection first, so the close handler found nothing and sent no disconnect event).
+    ws.on("error", (error: Error) => {
+      console.error("[WebSocket] Connection error:", error);
+    });
+  }
+
+  /**
+   * Reject the server-to-client requests that wait for a connection that went away,
+   * instead of letting them wait for their timeout (BUGS-2026-07 L29).
+   */
+  private abortInflightFor(connectionId: string): void {
+    const prefix = `${connectionId}:`;
+    for (const [key, controller] of this.inflight) {
+      if (key.startsWith(prefix)) {
+        this.inflight.delete(key);
+        controller.abort();
+      }
+    }
+  }
+
+  private rejectPendingFor(connectionId: string, reason: string): void {
+    for (const [requestId, pending] of this.pendingRequests) {
+      if (pending.connectionId === connectionId) {
+        clearTimeout(pending.timeout);
+        this.pendingRequests.delete(requestId);
+        pending.reject(new Error(`${reason}: ${connectionId}`));
+      }
+    }
+  }
+
+  /**
+   * Handle WebSocket message.
+   */
+  private async handleMessage(
+    ws: WebSocket,
+    data: Buffer,
+    req: IncomingMessage,
+    connectionId: string
+  ): Promise<void> {
+    // Parse message
+    let message: WebSocketMessage;
+    try {
+      message = JSON.parse(data.toString());
+    } catch (error) {
+      this.sendError(ws, "unknown", "Invalid JSON message");
+      return;
+    }
+
+    // Handle ping (client heartbeat)
+    if (message.type === "ping") {
+      const pong: WebSocketMessage = {
+        id: message.id,
+        type: "pong",
+      };
+      this.send(ws, pong);
+      return;
+    }
+
+    // Handle pong (response to our ping)
+    if (message.type === "pong") {
+      // Just acknowledge, no action needed
+      return;
+    }
+
+    // Handle server-response (response to server-initiated request)
+    if (message.type === "server-response") {
+      const pending = this.pendingRequests.get(message.id);
+      // Only the connection the request was sent to may answer it
+      if (pending && pending.connectionId === connectionId) {
+        clearTimeout(pending.timeout);
+        this.pendingRequests.delete(message.id);
+        if (message.error) {
+          pending.reject(new Error(message.error.message));
+        } else {
+          pending.resolve(message.result);
+        }
+      }
+      return;
+    }
+
+    // Handle cancel: the client stopped reading a response (for example, it left a stream early)
+    if (message.type === "cancel") {
+      const key = `${connectionId}:${message.id}`;
+      this.inflight.get(key)?.abort();
+      this.inflight.delete(key);
+      return;
+    }
+
+    // Handle credit: the client's reader took items of a stream, so the stream can send more
+    if (message.type === "credit") {
+      if (typeof message.credit === "number" && message.credit > 0) {
+        this.credits.get(`${connectionId}:${message.id}`)?.grant(message.credit);
+      }
+      return;
+    }
+
+    // Validate request message
+    if (!message.id || message.type !== "request" || !message.method) {
+      this.sendError(ws, message.id || "unknown", "Invalid message format");
+      return;
+    }
+
+    const key = `${connectionId}:${message.id}`;
+    // A request id names one request of the connection while it runs. A second request with
+    // the id would take the first one's controller: its stream could not be cancelled, and the
+    // two outputs would mix (deep dive TRN-9).
+    if (this.inflight.has(key)) {
+      this.sendError(ws, message.id, "A request with this id is still running", "DUPLICATE_REQUEST");
+      return;
+    }
+    const controller = new AbortController();
+    this.inflight.set(key, controller);
+
+    // Convert to ServerRequest. Only an absent payload becomes {}: before, 0, false and ""
+    // became {} too (deep dive TRN-10).
+    const serverRequest: ServerRequest = {
+      id: message.id,
+      method: message.method,
+      payload: message.payload === undefined ? {} : message.payload,
+      metadata: {
+        ...withoutInternalKeys(message.metadata),
+        // Add connection info
+        connectionId,
+        remoteAddress: req.socket.remoteAddress,
+        headers: req.headers,
+      },
+      signal: controller.signal,
+    };
+
+    try {
+      // Process through universal server
+      const serverResponse = await this.server.handle(serverRequest);
+
+      if (serverResponse.stream) {
+        const credit =
+          typeof message.credit === "number" && message.credit > 0 ? new StreamCredit(message.credit) : undefined;
+        if (credit) this.credits.set(key, credit);
+        await this.sendStream(ws, serverResponse.id, serverResponse.stream, controller.signal, credit);
+      } else {
+        this.sendResponse(ws, serverResponse);
+      }
+    } catch (error) {
+      console.error("[WebSocket] Request handling error:", error);
+      this.sendError(
+        ws,
+        message.id,
+        error instanceof Error ? error.message : "Internal server error"
+      );
+    } finally {
+      // Delete only this request's entries (deep dive TRN-9)
+      if (this.inflight.get(key) === controller) this.inflight.delete(key);
+      this.credits.delete(key);
+    }
+  }
+
+  /**
+   * Send a stream response (BUGS-2026-07 H5): one "stream" frame for each item, then a frame
+   * with `stream.done`. An error of the stream ends it with an "error" frame. When the signal
+   * aborts (a cancel message, a closed connection), the loop stops, and the stream's handler
+   * gets return().
+   *
+   * Flow control (deep dive TRN-3): the stream takes its next item only when the client has
+   * credit for it (when the request carried credit), and waits while the connection's send
+   * buffer holds more than `highWaterMark` bytes.
+   */
+  private async sendStream(
+    ws: WebSocket,
+    id: string,
+    stream: AsyncIterable<unknown>,
+    signal: AbortSignal,
+    credit: StreamCredit | undefined
+  ): Promise<void> {
+    const iterator = stream[Symbol.asyncIterator]();
+    let done = false;
+    try {
+      for (;;) {
+        if (credit && !(await credit.take(signal))) return;
+        if (signal.aborted || ws.readyState !== WebSocket.OPEN) return;
+        const result = await iterator.next();
+        if (result.done) {
+          done = true;
+          break;
+        }
+        if (signal.aborted || ws.readyState !== WebSocket.OPEN) return;
+        await this.sendWithBackpressure(ws, { id, type: "stream", payload: result.value, stream: { done: false } }, signal);
+      }
+      this.send(ws, { id, type: "stream", stream: { done: true } });
+    } catch (error) {
+      done = true;
+      if (signal.aborted) return;
+      const status = errorStatus(error);
+      this.send(ws, {
+        id,
+        type: "error",
+        status,
+        error: { code: status.code, message: status.message, retryable: status.retryable },
+        stream: { done: true },
+      });
+    } finally {
+      // The loop stopped before the end: the handler gets return()
+      if (!done) await iterator.return?.();
+    }
+  }
+
+  /**
+   * Send a frame. When the connection's send buffer is past the high-water mark, wait until
+   * the socket has written this frame, or the signal aborts.
+   */
+  private sendWithBackpressure(ws: WebSocket, message: WebSocketMessage, signal: AbortSignal): Promise<void> {
+    if (ws.readyState !== WebSocket.OPEN) return Promise.resolve();
+    return new Promise<void>((resolve) => {
+      let waiting = false;
+      const onAbort = (): void => resolve();
+      ws.send(JSON.stringify(message), () => {
+        if (!waiting) return;
+        signal.removeEventListener("abort", onAbort);
+        resolve();
+      });
+      if (ws.bufferedAmount > this.options.highWaterMark && !signal.aborted) {
+        waiting = true;
+        signal.addEventListener("abort", onAbort, { once: true });
+      } else {
+        resolve();
+      }
+    });
+  }
+
+  private sendResponse(ws: WebSocket, response: ServerResponse): void {
+    const message: WebSocketMessage = {
+      id: response.id,
+      type: response.status.type === "success" ? "response" : "error",
+      payload: response.payload,
+      metadata: response.metadata,
+      status: response.status,
+    };
+
+    if (response.status.type === "error") {
+      message.error = {
+        code: response.status.code,
+        message: response.status.message || "Unknown error",
+        retryable: response.status.retryable,
+      };
+    }
+
+    this.send(ws, message);
+  }
+
+  /**
+   * Send error to client.
+   */
+  private sendError(ws: WebSocket, id: string, message: string, code = "INTERNAL_ERROR"): void {
+    const errorMessage: WebSocketMessage = {
+      id,
+      type: "error",
+      status: {
+        type: "error",
+        code,
+        message,
+        retryable: false,
+      },
+      error: {
+        code,
+        message,
+        retryable: false,
+      },
+    };
+
+    this.send(ws, errorMessage);
+  }
+
+  /**
+   * Send message to client.
+   */
+  private send(ws: WebSocket, message: WebSocketMessage): void {
+    if (ws.readyState === WebSocket.OPEN) {
+      ws.send(JSON.stringify(message));
+    }
+  }
+
+  /**
+   * Broadcast message to all connected clients.
+   */
+  broadcast(message: WebSocketMessage): void {
+    for (const ws of this.connections) {
+      this.send(ws, message);
+    }
+  }
+
+  /**
+   * Get number of connected clients.
+   */
+  getConnectionCount(): number {
+    return this.connections.size;
+  }
+
+  // ===========================================================================
+  // Tracked Connection Methods (for bidirectional RPC)
+  // ===========================================================================
+
+  /**
+   * Get all tracked connections.
+   */
+  getTrackedConnections(): TrackedConnection[] {
+    return Array.from(this.trackedConnections.values());
+  }
+
+  /**
+   * Get a specific tracked connection by ID.
+   */
+  getTrackedConnection(connectionId: string): TrackedConnection | undefined {
+    return this.trackedConnections.get(connectionId);
+  }
+
+  /**
+   * Send a message to a specific connection.
+   */
+  sendToConnection(connectionId: string, message: WebSocketMessage): void {
+    const conn = this.trackedConnections.get(connectionId);
+    if (conn) {
+      this.send(conn.socket, message);
+    }
+  }
+
+  /**
+   * Call a procedure on a specific client (server-to-client RPC).
+   * Returns a promise that resolves with the client's response.
+   */
+  callClient(
+    connectionId: string,
+    path: string[],
+    input: unknown,
+    timeout = 30000
+  ): Promise<unknown> {
+    const conn = this.trackedConnections.get(connectionId);
+    if (!conn) {
+      return Promise.reject(new Error(`Connection not found: ${connectionId}`));
+    }
+
+    const requestId = `sreq_${Date.now()}_${Math.random().toString(36).slice(2)}`;
+
+    return new Promise((resolve, reject) => {
+      const timeoutHandle = setTimeout(() => {
+        this.pendingRequests.delete(requestId);
+        reject(new Error(`Request timeout after ${timeout}ms`));
+      }, timeout);
+
+      this.pendingRequests.set(requestId, {
+        id: requestId,
+        connectionId,
+        resolve,
+        reject,
+        timeout: timeoutHandle,
+      });
+
+      const message: WebSocketMessage = {
+        type: "server-request",
+        id: requestId,
+        path,
+        input,
+      };
+
+      this.send(conn.socket, message);
+    });
+  }
+
+  /**
+   * Register a handler for new connections.
+   * Returns a function to unregister the handler.
+   */
+  onConnect(handler: ConnectionEventHandler): () => void {
+    this.connectHandlers.push(handler);
+    return () => {
+      const index = this.connectHandlers.indexOf(handler);
+      if (index >= 0) {
+        this.connectHandlers.splice(index, 1);
+      }
+    };
+  }
+
+  /**
+   * Register a handler for disconnections.
+   * Returns a function to unregister the handler.
+   */
+  onDisconnect(handler: ConnectionEventHandler): () => void {
+    this.disconnectHandlers.push(handler);
+    return () => {
+      const index = this.disconnectHandlers.indexOf(handler);
+      if (index >= 0) {
+        this.disconnectHandlers.splice(index, 1);
+      }
+    };
+  }
+
+  /**
+   * Update metadata for a tracked connection.
+   */
+  updateConnectionMetadata(
+    connectionId: string,
+    metadata: Record<string, unknown>
+  ): void {
+    const conn = this.trackedConnections.get(connectionId);
+    if (conn) {
+      conn.metadata = { ...conn.metadata, ...metadata };
+    }
+  }
+
+  /**
+   * Set discovered procedures for a tracked connection.
+   */
+  setConnectionProcedures(connectionId: string, procedures: string[]): void {
+    const conn = this.trackedConnections.get(connectionId);
+    if (conn) {
+      conn.procedures = procedures;
+    }
+  }
+}