Add WebSocket client SDK and split out shared methods into ws-utils, add tests

Author: federicociner

package-lock.json

@@ -33,6 +33,10 @@
       "types": "./dist/http-stream.d.ts",
       "default": "./dist/http-stream.js"
     },
+    "./ws-client": {
+      "types": "./dist/ws-stream.d.ts",
+      "default": "./dist/ws-stream.js"
+    },
     "./server": {
       "types": "./dist/server.d.ts",
       "default": "./dist/server.js"
@@ -69,6 +73,7 @@
     "@eslint/js": "^10.0.1",
     "@hey-api/openapi-ts": "^0.97.0",
     "@types/node": "^25.5.0",
+    "@types/ws": "^8.5.13",
     "@typescript-eslint/eslint-plugin": "^8.57.1",
     "@typescript-eslint/parser": "^8.57.1",
     "concurrently": "^9.2.1",
@@ -82,6 +87,7 @@
     "typedoc-github-theme": "^0.4.0",
     "typescript": "^6.0.2",
     "vitest": "^4.1.0",
+    "ws": "^8.18.0",
     "zod": "^3.25.0 || ^4.0.0"
   }
 }

package.json

@@ -1,4 +1,5 @@
 import http from "node:http";
+import { WebSocketServer } from "ws";
 
 import { AcpServer } from "../server.js";
 import { createNodeHttpHandler } from "../node-adapter.js";
@@ -9,6 +10,7 @@ import type { Agent, AgentSideConnection } from "../acp.js";
 
 export interface TestHttpServer {
   readonly url: string;
+  readonly wsUrl: string;
   readonly close: () => Promise<void>;
 }
 
@@ -19,6 +21,13 @@ export async function startTestServer(
 ): Promise<TestHttpServer> {
   const acpServer = new AcpServer({ createAgent: agentFactory });
   const httpServer = http.createServer(createNodeHttpHandler(acpServer));
+  const webSocketServer = new WebSocketServer({ noServer: true });
+
+  httpServer.on("upgrade", (req, socket, head) => {
+    webSocketServer.handleUpgrade(req, socket, head, (webSocket) => {
+      acpServer.handleWebSocket(webSocket);
+    });
+  });
 
   await listen(httpServer, options.port ?? 0);
 
@@ -30,8 +39,14 @@ export async function startTestServer(
 
   return {
     url: `http://127.0.0.1:${address.port}`,
+    wsUrl: `ws://127.0.0.1:${address.port}`,
     close: async () => {
-      await Promise.all([acpServer.close(), closeHttpServer(httpServer)]);
+      terminateWebSockets(webSocketServer);
+      await Promise.all([
+        acpServer.close(),
+        closeWebSocketServer(webSocketServer),
+        closeHttpServer(httpServer),
+      ]);
     },
   };
 }
@@ -54,6 +69,25 @@ function listen(server: http.Server, port: number): Promise<void> {
   });
 }
 
+function terminateWebSockets(server: WebSocketServer): void {
+  for (const client of server.clients) {
+    client.terminate();
+  }
+}
+
+function closeWebSocketServer(server: WebSocketServer): Promise<void> {
+  return new Promise((resolve, reject) => {
+    server.close((error) => {
+      if (error) {
+        reject(error);
+        return;
+      }
+
+      resolve();
+    });
+  });
+}
+
 function closeHttpServer(server: http.Server): Promise<void> {
   return new Promise((resolve, reject) => {
     server.close((error) => {

src/test-support/test-http-server.ts

@@ -8,14 +8,17 @@ import {
   messageIdKey,
   sessionIdFromParams,
 } from "./protocol.js";
-
+import { onWebSocket, webSocketMessageToString } from "./ws-utils.js";
 import type { Agent, AgentSideConnection } from "./acp.js";
 import type {
   ConnectionRegistry,
   ConnectionState,
   ResponseRoute,
 } from "./connection.js";
 import type { AnyMessage, AnyRequest } from "./jsonrpc.js";
+import type { WebSocketLike } from "./ws-utils.js";
+
+export type WebSocketServerSocket = WebSocketLike;
 
 type ForwardResult =
   | {
@@ -26,27 +29,13 @@ type ForwardResult =
       message: string;
     };
 
-export interface WebSocketServerSocket {
-  readonly readyState?: number;
-  send(data: string): void;
-  close(code?: number, reason?: string): void;
-  addEventListener?(type: string, listener: (event: unknown) => void): void;
-  removeEventListener?(type: string, listener: (event: unknown) => void): void;
-  on?(type: string, listener: (...args: unknown[]) => void): unknown;
-  off?(type: string, listener: (...args: unknown[]) => void): unknown;
-  removeListener?(
-    type: string,
-    listener: (...args: unknown[]) => void,
-  ): unknown;
-}
-
 export interface WebSocketConnectionOptions {
   readonly registry: ConnectionRegistry;
   readonly createAgent: (conn: AgentSideConnection) => Agent;
 }
 
 export function handleWebSocketConnection(
-  socket: WebSocketServerSocket,
+  socket: WebSocketLike,
   options: WebSocketConnectionOptions,
 ): void {
   const session = new WebSocketServerSession(socket, options);
@@ -56,29 +45,30 @@ export function handleWebSocketConnection(
 class WebSocketServerSession {
   private connection: ConnectionState | undefined;
   private outboundReader: ReadableStreamDefaultReader<AnyMessage> | undefined;
+  private inboundWriteChain: Promise<void> = Promise.resolve();
   private isClosed = false;
   private readonly detachListeners: Array<() => void> = [];
 
   constructor(
-    private readonly socket: WebSocketServerSocket,
+    private readonly socket: WebSocketLike,
     private readonly options: WebSocketConnectionOptions,
   ) {}
 
   start(): void {
     this.detachListeners.push(
-      onSocket(this.socket, "message", (...args) => {
+      onWebSocket(this.socket, "message", (...args) => {
         void this.handleSocketMessage(args);
       }),
     );
 
     this.detachListeners.push(
-      onSocket(this.socket, "close", () => {
+      onWebSocket(this.socket, "close", () => {
         void this.closeSession();
       }),
     );
 
     this.detachListeners.push(
-      onSocket(this.socket, "error", () => {
+      onWebSocket(this.socket, "error", () => {
         void this.shutdown(1011, "WebSocket error");
       }),
     );
@@ -89,7 +79,7 @@ class WebSocketServerSession {
       return;
     }
 
-    const text = socketMessageToString(args);
+    const text = webSocketMessageToString(args);
     if (text === undefined) {
       console.warn("Ignoring non-text ACP WebSocket frame");
       return;
@@ -202,19 +192,33 @@ class WebSocketServerSession {
         connection.pendingRoutes.set(key, route);
       }
 
-      await writeInbound(connection, message);
+      await this.writeInbound(message);
       return { ok: true };
     }
 
     if (isResponseMessage(message)) {
-      await writeInbound(connection, message);
+      await this.writeInbound(message);
       return { ok: true };
     }
 
-    await writeInbound(connection, message);
+    await this.writeInbound(message);
     return { ok: true };
   }
 
+  private async writeInbound(message: AnyMessage): Promise<void> {
+    const connection = this.connection;
+
+    if (!connection) {
+      throw new Error("ACP WebSocket connection is not initialized");
+    }
+
+    const write = this.inboundWriteChain.then(() =>
+      writeInbound(connection, message),
+    );
+    this.inboundWriteChain = write.catch(() => undefined);
+    await write;
+  }
+
   private startOutboundPump(connection: ConnectionState): void {
     const subscription = connection.allOutbound.subscribe();
     const reader = subscription.stream.getReader();
@@ -345,81 +349,3 @@ function determineWebSocketRoute(message: AnyRequest): ResponseRoute {
 
   return "connection";
 }
-
-function onSocket(
-  socket: WebSocketServerSocket,
-  type: string,
-  listener: (...args: unknown[]) => void,
-): () => void {
-  if (socket.addEventListener) {
-    const eventListener = (event: unknown) => listener(event);
-    socket.addEventListener(type, eventListener);
-
-    return () => {
-      socket.removeEventListener?.(type, eventListener);
-    };
-  }
-
-  if (socket.on) {
-    socket.on(type, listener);
-
-    return () => {
-      if (socket.off) {
-        socket.off(type, listener);
-        return;
-      }
-
-      socket.removeListener?.(type, listener);
-    };
-  }
-
-  throw new Error("WebSocket object does not support event listeners");
-}
-
-function socketMessageToString(args: unknown[]): string | undefined {
-  const data = extractMessageData(args);
-
-  if (typeof data === "string") {
-    return data;
-  }
-
-  if (data instanceof ArrayBuffer || ArrayBuffer.isView(data)) {
-    return new TextDecoder().decode(data);
-  }
-
-  if (Array.isArray(data) && data.every(ArrayBuffer.isView)) {
-    return decodeArrayBufferViews(data);
-  }
-
-  return undefined;
-}
-
-function extractMessageData(args: unknown[]): unknown {
-  const [first] = args;
-
-  if (isMessageEventLike(first)) {
-    return first.data;
-  }
-
-  return first;
-}
-
-function isMessageEventLike(value: unknown): value is { data: unknown } {
-  return typeof value === "object" && value !== null && "data" in value;
-}
-
-function decodeArrayBufferViews(views: ArrayBufferView[]): string {
-  const totalLength = views.reduce((sum, view) => sum + view.byteLength, 0);
-  const combined = new Uint8Array(totalLength);
-  let offset = 0;
-
-  for (const view of views) {
-    combined.set(
-      new Uint8Array(view.buffer, view.byteOffset, view.byteLength),
-      offset,
-    );
-    offset += view.byteLength;
-  }
-
-  return new TextDecoder().decode(combined);
-}