Extract the session SocketRegistry port

packages/control-plane/src/session/disconnect-handler.ts

@@ -4,12 +4,12 @@ import type {
   ConnectedClient,
   SandboxDisconnectMonitor,
   SessionBroadcaster,
-  SocketRegistry,
+  RoutedSocketRegistry,
 } from "./ports";
 
 export interface SessionDisconnectHandlerDeps<Connection, Client extends ConnectedClient> {
   getLogger: () => Logger;
-  sockets: SocketRegistry<Connection, Client>;
+  sockets: RoutedSocketRegistry<Connection, Client>;
   sandbox: SandboxDisconnectMonitor;
   broadcaster: SessionBroadcaster;
 }

packages/control-plane/src/session/durable-object-session-connections.test.ts

@@ -1,16 +1,6 @@
 import { describe, expect, it, vi } from "vitest";
 import { DurableObjectSessionConnections } from "./durable-object-session-connections";
-import type { SessionWebSocketManager } from "./websocket-manager";
-
-vi.stubGlobal(
-  "WebSocketRequestResponsePair",
-  class WebSocketRequestResponsePair {
-    constructor(
-      readonly request: string,
-      readonly response: string
-    ) {}
-  }
-);
+import type { DurableObjectSocketRegistry } from "./durable-object-socket-registry";
 
 function harness() {
   const browser = { readyState: WebSocket.OPEN } as WebSocket;
@@ -40,22 +30,22 @@ function harness() {
     detachSandboxSocket: vi.fn(),
     getAuthenticatedClients: vi.fn(() => [clientInfo].values()),
     recoverClientMapping: vi.fn(() => null),
-  } as unknown as SessionWebSocketManager;
-  const state = { setWebSocketAutoResponse: vi.fn() } as unknown as DurableObjectState;
+    configureAutoPing: vi.fn(),
+    createUpgradeSockets: vi.fn(),
+  } as unknown as DurableObjectSocketRegistry;
   return {
-    connections: new DurableObjectSessionConnections(state, manager),
+    connections: new DurableObjectSessionConnections(manager),
     manager,
-    state,
     browser,
     sandbox,
   };
 }
 
 describe("DurableObjectSessionConnections", () => {
   it("owns Cloudflare auto-response configuration", () => {
-    const { state } = harness();
+    const { manager } = harness();
 
-    expect(state.setWebSocketAutoResponse).toHaveBeenCalledTimes(1);
+    expect(manager.configureAutoPing).toHaveBeenCalledTimes(1);
   });
 
   it("registers and broadcasts to a browser through the WebSocket registry", async () => {

packages/control-plane/src/session/durable-object-session-connections.ts

@@ -8,32 +8,25 @@ import type {
 } from "./connections";
 import { projectConnectedParticipants, SandboxDeliveryUnavailableError } from "./connections";
 import type { SandboxCommand } from "./types";
-import type { SessionWebSocketManager } from "./websocket-manager";
+import type { DurableObjectSocketRegistry } from "./durable-object-socket-registry";
 
 /** Cloudflare Durable Object implementation of the session connection port. */
 export class DurableObjectSessionConnections implements SessionConnections {
-  constructor(
-    private readonly ctx: DurableObjectState,
-    private readonly wsManager: SessionWebSocketManager
-  ) {
-    this.ctx.setWebSocketAutoResponse(
-      new WebSocketRequestResponsePair(
-        JSON.stringify({ type: "ping" }),
-        JSON.stringify({ type: "pong", timestamp: Date.now() })
-      )
+  constructor(private readonly sockets: DurableObjectSocketRegistry) {
+    this.sockets.configureAutoPing(
+      JSON.stringify({ type: "ping" }),
+      JSON.stringify({ type: "pong", timestamp: Date.now() })
     );
   }
 
   createUpgradeSockets(): { client: WebSocket; server: WebSocket } {
-    const pair = new WebSocketPair();
-    const [client, server] = Object.values(pair);
-    return { client, server };
+    return this.sockets.createUpgradeSockets();
   }
 
   registerBrowser(input: BrowserConnection): Promise<void> {
     let matchedSocket: WebSocket | null = null;
-    this.wsManager.forEachClientSocket("all_clients", (ws) => {
-      const connection = this.wsManager.classify(ws);
+    this.sockets.forEachClientSocket("all_clients", (ws) => {
+      const connection = this.sockets.classify(ws);
       if (connection.kind === "client" && connection.wsId === input.connectionId) {
         matchedSocket = ws;
       }
@@ -42,12 +35,12 @@ export class DurableObjectSessionConnections implements SessionConnections {
       return Promise.reject(new Error(`Browser connection ${input.connectionId} not found`));
     }
 
-    this.wsManager.setClient(matchedSocket, {
+    this.sockets.setClient(matchedSocket, {
       ...input.participant,
       clientId: input.clientId,
       ws: matchedSocket,
     });
-    this.wsManager.persistClientMapping(
+    this.sockets.persistClientMapping(
       input.connectionId,
       input.participant.participantId,
       input.clientId
@@ -56,9 +49,9 @@ export class DurableObjectSessionConnections implements SessionConnections {
   }
 
   registerSandbox(input: SandboxConnection): Promise<void> {
-    const ws = this.wsManager.getSandboxSocket();
+    const ws = this.sockets.getSandboxSocket();
     if (!ws) return Promise.reject(new Error(`Sandbox connection ${input.connectionId} not found`));
-    const connection = this.wsManager.classify(ws);
+    const connection = this.sockets.classify(ws);
     if (
       connection.kind !== "sandbox" ||
       (input.sandboxId !== undefined && connection.sandboxId !== input.sandboxId)
@@ -69,26 +62,26 @@ export class DurableObjectSessionConnections implements SessionConnections {
   }
 
   sendToSandbox(message: SandboxCommand): Promise<void> {
-    const ws = this.wsManager.getSandboxSocket();
+    const ws = this.sockets.getSandboxSocket();
     if (!ws) return Promise.reject(new SandboxDeliveryUnavailableError());
-    return this.wsManager.send(ws, message)
+    return this.sockets.send(ws, message)
       ? Promise.resolve()
       : Promise.reject(new SandboxDeliveryUnavailableError("Failed to send message to sandbox"));
   }
 
   broadcastToBrowsers(message: ServerMessage): Promise<void> {
-    this.wsManager.forEachClientSocket("authenticated_only", (ws) => {
-      this.wsManager.send(ws, message);
+    this.sockets.forEachClientSocket("authenticated_only", (ws) => {
+      this.sockets.send(ws, message);
     });
     return Promise.resolve();
   }
 
   disconnectSandbox(reason: DisconnectReason): Promise<void> {
-    this.wsManager.detachSandboxSocket(reason.code, reason.reason);
+    this.sockets.detachSandboxSocket(reason.code, reason.reason);
     return Promise.resolve();
   }
 
   listParticipants(): Promise<ConnectedParticipant[]> {
-    return Promise.resolve(projectConnectedParticipants(this.wsManager.getAuthenticatedClients()));
+    return Promise.resolve(projectConnectedParticipants(this.sockets.getAuthenticatedClients()));
   }
 }

packages/control-plane/src/session/websocket-manager.test.ts → packages/control-plane/src/session/durable-object-socket-registry.test.ts

@@ -1,13 +1,15 @@
 /**
- * Unit tests for SessionWebSocketManagerImpl.
+ * Unit tests for DurableObjectSocketRegistry.
  *
  * Uses fake DurableObjectState and mock repositories to test
  * all WebSocket mechanics in isolation from the full DO.
  */
 
 import { describe, it, expect, vi } from "vitest";
-import { SessionWebSocketManagerImpl } from "./websocket-manager";
-import type { WebSocketManagerConfig } from "./websocket-manager";
+import {
+  DurableObjectSocketRegistry,
+  type SocketRegistryConfig,
+} from "./durable-object-socket-registry";
 import type { Logger } from "../logger";
 import type { ClientInfo } from "../types";
 import type { SandboxRepository } from "./sandbox-repository";
@@ -70,8 +72,10 @@ function createFakeCtx(): FakeCtx {
     getTags(ws: WebSocket): string[] {
       return sockets.get(ws) ?? [];
     },
-    getWebSockets(): WebSocket[] {
-      return Array.from(sockets.keys());
+    getWebSockets(tag?: string): WebSocket[] {
+      return Array.from(sockets, ([ws, tags]) => ({ ws, tags }))
+        .filter(({ tags }) => !tag || tags.includes(tag))
+        .map(({ ws }) => ws);
     },
     setWebSocketAutoResponse: vi.fn(),
     storage: { setAlarm: vi.fn() },
@@ -180,15 +184,15 @@ function createSandboxRow(modalSandboxId: string): SandboxRow {
   };
 }
 
-const TEST_CONFIG: WebSocketManagerConfig = { authTimeoutMs: 100 };
+const TEST_CONFIG: SocketRegistryConfig = { authTimeoutMs: 100 };
 
 /** Create a fresh manager with all dependencies. */
 function createManager() {
   const fakeCtx = createFakeCtx();
   const mockRepo = createMockRepository();
   const log = createMockLogger();
 
-  const manager = new SessionWebSocketManagerImpl(
+  const manager = new DurableObjectSocketRegistry(
     fakeCtx.state,
     mockRepo.repo,
     mockRepo.repo as unknown as WsClientMappingRepository,
@@ -203,7 +207,7 @@ function createManager() {
 // Tests
 // ---------------------------------------------------------------------------
 
-describe("SessionWebSocketManagerImpl", () => {
+describe("DurableObjectSocketRegistry", () => {
   describe("classify", () => {
     it("classifies sandbox socket with sandbox ID", () => {
       const { manager, sockets } = createManager();

packages/control-plane/src/session/websocket-manager.ts → packages/control-plane/src/session/durable-object-socket-registry.ts

@@ -1,6 +1,5 @@
 /**
- * SessionWebSocketManager — centralizes all Cloudflare WebSocket API usage
- * into a single, testable module.
+ * DurableObjectSocketRegistry - the Cloudflare implementation of the session socket port.
  *
  * The manager is a registry for ClientInfo, not a factory. The DO builds
  * ClientInfo and stores it here via setClient/getClient.
@@ -10,83 +9,21 @@ import type { Logger } from "../logger";
 import type { ClientInfo } from "../types";
 import type { ConnectionClassification } from "./ports";
 import type { SandboxRepository } from "./sandbox-repository";
+import type { SocketRegistry } from "./socket-registry";
 import type {
   WsClientMappingRepository,
   WsClientMappingResult,
 } from "./ws-client-mapping-repository";
 
-/** Configuration for the WebSocket manager. */
-export interface WebSocketManagerConfig {
+export interface SocketRegistryConfig {
   authTimeoutMs: number;
 }
 
-// ---------------------------------------------------------------------------
-// Interface
-// ---------------------------------------------------------------------------
-
-export interface SessionWebSocketManager {
-  /** Accept a client WebSocket with a wsId tag for hibernation recovery. */
-  acceptClientSocket(ws: WebSocket, wsId: string): void;
-
-  /**
-   * Accept a sandbox WebSocket, close any existing sandbox socket, and set
-   * as the active sandbox connection.
-   */
-  acceptAndSetSandboxSocket(ws: WebSocket, sandboxId?: string): { replaced: boolean };
-
-  /** Parse a WebSocket's tags to determine its kind and identity. */
-  classify(ws: WebSocket): ConnectionClassification;
-
-  /**
-   * Get the active sandbox socket, recovering from hibernation if needed.
-   * Validates sandbox ID against the repository during hibernation recovery.
-   */
-  getSandboxSocket(): WebSocket | null;
-
-  /** Clear the in-memory sandbox socket reference. */
-  clearSandboxSocket(): void;
-
-  /** Clear and close all active sandbox sockets without consulting persisted dispatch status. */
-  detachSandboxSocket(code: number, reason: string): void;
-
-  /** Clear sandbox socket only if ws matches current reference. Returns true if it was the active socket. */
-  clearSandboxSocketIfMatch(ws: WebSocket): boolean;
-
-  setClient(ws: WebSocket, info: ClientInfo): void;
-  getClient(ws: WebSocket): ClientInfo | null;
-  removeClient(ws: WebSocket): ClientInfo | null;
-
-  /** Returns raw DB mapping for hibernation recovery. The DO builds ClientInfo from this. */
-  recoverClientMapping(ws: WebSocket): WsClientMappingResult | null;
-
-  /** Persist ws-to-participant mapping for hibernation survival. */
-  persistClientMapping(wsId: string, participantId: string, clientId: string): void;
-
-  setClientSynchronizing(ws: WebSocket, synchronizing: boolean): void;
-  isClientSynchronizing(ws: WebSocket): boolean;
-  isClientAuthenticated(ws: WebSocket): boolean;
-
-  /** Check if a wsId has a persisted mapping (used by auth timeout). */
-  hasPersistedMapping(wsId: string): boolean;
-
-  send(ws: WebSocket, message: string | object): boolean;
-  close(ws: WebSocket, code: number, reason: string): void;
-
-  forEachClientSocket(
-    mode: "all_clients" | "authenticated_only",
-    fn: (ws: WebSocket) => void
-  ): void;
-
-  enforceAuthTimeout(ws: WebSocket, wsId: string): Promise<void>;
-  getAuthenticatedClients(): IterableIterator<ClientInfo>;
-  getConnectedClientCount(): number;
-}
-
 // ---------------------------------------------------------------------------
 // Implementation
 // ---------------------------------------------------------------------------
 
-export class SessionWebSocketManagerImpl implements SessionWebSocketManager {
+export class DurableObjectSocketRegistry implements SocketRegistry<WebSocket> {
   private clients = new Map<WebSocket, ClientInfo>();
   private synchronizingClients = new Set<WebSocket>();
   private sandboxWs: WebSocket | null = null;
@@ -96,9 +33,19 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager {
     private readonly sandboxRepository: SandboxRepository,
     private readonly wsClientMappingRepository: WsClientMappingRepository,
     private readonly log: Logger,
-    private readonly config: WebSocketManagerConfig
+    private readonly config: SocketRegistryConfig
   ) {}
 
+  createUpgradeSockets(): { client: WebSocket; server: WebSocket } {
+    const pair = new WebSocketPair();
+    const [client, server] = Object.values(pair);
+    return { client, server };
+  }
+
+  configureAutoPing(request: string, response: string): void {
+    this.ctx.setWebSocketAutoResponse(new WebSocketRequestResponsePair(request, response));
+  }
+
   // -------------------------------------------------------------------------
   // Accept
   // -------------------------------------------------------------------------
@@ -157,11 +104,8 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager {
     if (sandbox && terminalStatuses.includes(sandbox.status)) {
       this.sandboxWs = null;
       // Close any lingering sandbox WebSockets so they don't persist
-      for (const ws of this.ctx.getWebSockets()) {
-        const parsed = this.classify(ws);
-        if (parsed.kind === "sandbox") {
-          this.close(ws, 1000, "Sandbox terminated");
-        }
+      for (const ws of this.ctx.getWebSockets("sandbox")) {
+        this.close(ws, 1000, "Sandbox terminated");
       }
       return null;
     }
@@ -172,7 +116,7 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager {
 
     // Hibernation recovery: scan all WebSockets, validate sandbox identity
 
-    for (const ws of this.ctx.getWebSockets()) {
+    for (const ws of this.ctx.getWebSockets("sandbox")) {
       const parsed = this.classify(ws);
       if (parsed.kind !== "sandbox" || ws.readyState !== WebSocket.OPEN) continue;
 
@@ -199,9 +143,7 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager {
   detachSandboxSocket(code: number, reason: string): void {
     const sockets = new Set<WebSocket>();
     if (this.sandboxWs) sockets.add(this.sandboxWs);
-    for (const ws of this.ctx.getWebSockets()) {
-      if (this.classify(ws).kind === "sandbox") sockets.add(ws);
-    }
+    for (const ws of this.ctx.getWebSockets("sandbox")) sockets.add(ws);
     this.sandboxWs = null;
     for (const ws of sockets) this.close(ws, code, reason);
   }