From a1e24d4a71f3f79b48a91efb97cc9dbb02150774 Mon Sep 17 00:00:00 2001 From: Cole Murray Date: Thu, 3 Sep 2026 21:18:16 -0700 Subject: [PATCH 1/2] Require the active sandbox socket identity at message dispatch (COL-128) The message router treated any socket tagged `sandbox` as authoritative. Replacing a bridge closed the previous sockets, but close is cleanup, not a fence: a replaced socket keeps its tags until its close completes, and a frame already queued on it (or delivered to a restored instance) still mutated session state after the new bridge was selected. Recovery after hibernation had the mirror problem: it re-adopted the first open sandbox socket, which could be the closing one. Every accepted bridge socket now carries a fresh `socket:` tag, and accepting persists that id on the sandbox row (migration 48) before the replaced sockets are closed. Dispatch, recovery, and close handling compare a socket's tag against the row: the router refuses frames from any other sandbox socket and closes it again; recovery re-adopts only the socket the row names; a replaced socket's close no longer counts as losing the bridge. A spawn reservation clears the id with the credentials it invalidates. Sockets accepted before this change (no tag, NULL column) stay authoritative until the next bridge connects. --- .../src/sandbox/lifecycle/manager.test.ts | 1 + .../control-plane/src/session/components.ts | 1 + .../handlers/child-summary.handler.test.ts | 1 + .../session-lifecycle.handler.test.ts | 1 + .../src/session/message-router.ts | 18 +- packages/control-plane/src/session/ports.ts | 15 +- .../src/session/sandbox-repository.test.ts | 28 +++ .../src/session/sandbox-repository.ts | 21 ++- .../control-plane/src/session/schema.test.ts | 7 + packages/control-plane/src/session/schema.ts | 6 + .../control-plane/src/session/server.test.ts | 24 +++ packages/control-plane/src/session/types.ts | 2 + .../src/session/websocket-manager.test.ts | 160 ++++++++++++++++-- .../src/session/websocket-manager.ts | 61 +++++-- .../durable-object-eviction.test.ts | 68 +++++++- .../integration/websocket-sandbox.test.ts | 60 +++++++ 16 files changed, 438 insertions(+), 36 deletions(-) diff --git a/packages/control-plane/src/sandbox/lifecycle/manager.test.ts b/packages/control-plane/src/sandbox/lifecycle/manager.test.ts index b5b2c57528..ca1d9ec4c1 100644 --- a/packages/control-plane/src/sandbox/lifecycle/manager.test.ts +++ b/packages/control-plane/src/sandbox/lifecycle/manager.test.ts @@ -121,6 +121,7 @@ function createMockSandbox( tunnel_urls: null, ttyd_url: null, ttyd_token: null, + active_socket_id: null, created_at: Date.now() - 60000, spawn_failure_count: 0, last_spawn_failure: 0, diff --git a/packages/control-plane/src/session/components.ts b/packages/control-plane/src/session/components.ts index 73e7f5cb22..6bbbb04a2f 100644 --- a/packages/control-plane/src/session/components.ts +++ b/packages/control-plane/src/session/components.ts @@ -775,6 +775,7 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi send: (ws, message) => wsManager.send(ws, message), getClient: (ws) => connectionAuthenticator.getClientInfo(ws), close: (ws, code, reason) => wsManager.close(ws, code, reason), + isActiveSandbox: (ws) => wsManager.isActiveSandboxSocket(ws), clearSandboxIfMatch: (ws) => wsManager.clearSandboxSocketIfMatch(ws), removeClient: (ws) => wsManager.removeClient(ws), hasParticipant: (participantId) => diff --git a/packages/control-plane/src/session/http/handlers/child-summary.handler.test.ts b/packages/control-plane/src/session/http/handlers/child-summary.handler.test.ts index 2abe3725b3..e511a3928a 100644 --- a/packages/control-plane/src/session/http/handlers/child-summary.handler.test.ts +++ b/packages/control-plane/src/session/http/handlers/child-summary.handler.test.ts @@ -67,6 +67,7 @@ function createSandbox(overrides: Partial = {}): SandboxRow { tunnel_urls: null, ttyd_url: null, ttyd_token: null, + active_socket_id: null, created_at: 1, ...overrides, }; diff --git a/packages/control-plane/src/session/http/handlers/session-lifecycle.handler.test.ts b/packages/control-plane/src/session/http/handlers/session-lifecycle.handler.test.ts index 34a69575c3..7102245f6c 100644 --- a/packages/control-plane/src/session/http/handlers/session-lifecycle.handler.test.ts +++ b/packages/control-plane/src/session/http/handlers/session-lifecycle.handler.test.ts @@ -62,6 +62,7 @@ function createSandbox(overrides: Partial = {}): SandboxRow { tunnel_urls: null, ttyd_url: null, ttyd_token: null, + active_socket_id: null, created_at: 1, ...overrides, }; diff --git a/packages/control-plane/src/session/message-router.ts b/packages/control-plane/src/session/message-router.ts index d3118a04b1..be9c2fb351 100644 --- a/packages/control-plane/src/session/message-router.ts +++ b/packages/control-plane/src/session/message-router.ts @@ -56,11 +56,23 @@ export class SessionMessageRouter { // The wire protocol is JSON text; binary frames have always been ignored. if (typeof message !== "string") return; - if (this.deps.sockets.classify(connection).kind === "sandbox") { - await this.handleSandboxMessage(message); - } else { + const classified = this.deps.sockets.classify(connection); + if (classified.kind !== "sandbox") { await this.handleClientMessage(connection, message); + return; + } + if (!this.deps.sockets.isActiveSandbox(connection)) { + // A replaced bridge keeps its tags until its close completes. A frame + // from it proves it is still open, so close it again instead of + // letting it mutate the session. + this.deps.log.debug("Ignoring frame from a replaced sandbox socket", { + sandbox_id: classified.sandboxId, + socket_id: classified.socketId, + }); + this.deps.sockets.close(connection, 1000, "Sandbox socket replaced"); + return; } + await this.handleSandboxMessage(message); } private async handleSandboxMessage(message: string): Promise { diff --git a/packages/control-plane/src/session/ports.ts b/packages/control-plane/src/session/ports.ts index fcca2e58cb..2266e17c2c 100644 --- a/packages/control-plane/src/session/ports.ts +++ b/packages/control-plane/src/session/ports.ts @@ -16,9 +16,14 @@ export interface ConnectedClient { lastFetchHistoryAtMs?: number; } -/** Result of classifying an opaque runtime connection. */ +/** + * Result of classifying an opaque runtime connection. A sandbox socket + * carries the sandbox it authenticated as and its own accept-time identity; + * only the socket whose identity the session persisted as active is + * authoritative (see `SocketRegistry.isActiveSandbox`). + */ export type ConnectionClassification = - | { kind: "sandbox"; sandboxId?: string } + | { kind: "sandbox"; sandboxId?: string; socketId?: string } | { kind: "client"; wsId?: string }; /** Wall and monotonic time sources used by session application code. */ @@ -33,6 +38,12 @@ export interface SocketRegistry { send(connection: Connection, message: ServerMessage): boolean; getClient(connection: Connection): Client | null; close(connection: Connection, code: number, reason: string): void; + /** + * Whether `connection` is the sandbox socket the session currently + * dispatches to. A replaced bridge keeps its tags until its close + * completes; its frames carry no authority. + */ + isActiveSandbox(connection: Connection): boolean; clearSandboxIfMatch(connection: Connection): boolean; removeClient(connection: Connection): Client | null; hasParticipant(participantId: string): boolean; diff --git a/packages/control-plane/src/session/sandbox-repository.test.ts b/packages/control-plane/src/session/sandbox-repository.test.ts index eadccfb20f..77e435b57c 100644 --- a/packages/control-plane/src/session/sandbox-repository.test.ts +++ b/packages/control-plane/src/session/sandbox-repository.test.ts @@ -131,6 +131,8 @@ describe("SandboxRepository", () => { expect(mock.calls[0].query).toContain("vnc_password = NULL"); // A replacement sandbox must not inherit the predecessor's runtime. expect(mock.calls[0].query).toContain("runtime_version = NULL"); + // ...nor its bridge: the predecessor's socket loses dispatch authority here. + expect(mock.calls[0].query).toContain("active_socket_id = NULL"); expect(mock.calls[0].params).toEqual(["spawning", 1000, "modal-sb-1"]); }); @@ -146,6 +148,32 @@ describe("SandboxRepository", () => { }); }); + describe("active socket id", () => { + const query = `SELECT active_socket_id FROM sandbox LIMIT 1`; + + it("reads null before any bridge has connected", () => { + mock.setData(query, [{ active_socket_id: null }]); + expect(repository.getActiveSocketId()).toBeNull(); + }); + + it("reads null without a sandbox row", () => { + expect(repository.getActiveSocketId()).toBeNull(); + }); + + it("reads the persisted identity", () => { + mock.setData(query, [{ active_socket_id: "sbws-1" }]); + expect(repository.getActiveSocketId()).toBe("sbws-1"); + }); + + it("writes the identity to the session's one sandbox row", () => { + repository.setActiveSocketId("sbws-2"); + + expect(mock.calls.length).toBe(1); + expect(mock.calls[0].query).toContain("UPDATE sandbox SET active_socket_id = ?"); + expect(mock.calls[0].params).toEqual(["sbws-2"]); + }); + }); + describe("updateSandboxAuthTokenHash", () => { const query = `UPDATE sandbox SET auth_token_hash = ? WHERE modal_sandbox_id = ?`; diff --git a/packages/control-plane/src/session/sandbox-repository.ts b/packages/control-plane/src/session/sandbox-repository.ts index 112b501fdf..0e4bf6a2c0 100644 --- a/packages/control-plane/src/session/sandbox-repository.ts +++ b/packages/control-plane/src/session/sandbox-repository.ts @@ -137,7 +137,8 @@ export class SandboxRepository { tunnel_urls = NULL, ttyd_url = NULL, ttyd_token = NULL, - runtime_version = NULL + runtime_version = NULL, + active_socket_id = NULL WHERE id = (SELECT id FROM sandbox LIMIT 1)`, data.status, data.createdAt, @@ -145,6 +146,24 @@ export class SandboxRepository { ); } + /** + * The bridge socket the session dispatches to, by its `socket:` tag. + * Null once a spawn reserves a new identity or before any bridge connects. + */ + getActiveSocketId(): string | null { + const result = this.sql.exec(`SELECT active_socket_id FROM sandbox LIMIT 1`); + const rows = this.rows<{ active_socket_id: string | null }>(result); + return rows[0]?.active_socket_id ?? null; + } + + /** Make `socketId` the socket the session dispatches to; every earlier socket loses authority. */ + setActiveSocketId(socketId: string): void { + this.sql.exec( + `UPDATE sandbox SET active_socket_id = ? WHERE id = (SELECT id FROM sandbox LIMIT 1)`, + socketId + ); + } + /** * Phase 2 of the two-phase spawn write (#1589): publish the reserved * identity's hash. Scoped to that identity so a delayed publisher cannot diff --git a/packages/control-plane/src/session/schema.test.ts b/packages/control-plane/src/session/schema.test.ts index bce73b9ce1..bad520a8b1 100644 --- a/packages/control-plane/src/session/schema.test.ts +++ b/packages/control-plane/src/session/schema.test.ts @@ -273,6 +273,13 @@ describe("applyMigrations", () => { ]); }); + it("adds sandbox.active_socket_id for fresh and migrated DOs", () => { + expect(SCHEMA_SQL).toContain("active_socket_id TEXT"); + + const migration = MIGRATIONS.find((entry) => entry.id === 48); + expect(migration?.run).toBe("ALTER TABLE sandbox ADD COLUMN active_socket_id TEXT"); + }); + it("keeps repository context consistent at the session table boundary", () => { expect(SCHEMA_SQL).toContain("(repo_owner IS NULL) = (repo_name IS NULL)"); expect(SCHEMA_SQL).toContain("repo_owner IS NOT NULL"); diff --git a/packages/control-plane/src/session/schema.ts b/packages/control-plane/src/session/schema.ts index 6e8e892364..09072e59ae 100644 --- a/packages/control-plane/src/session/schema.ts +++ b/packages/control-plane/src/session/schema.ts @@ -189,6 +189,7 @@ CREATE TABLE IF NOT EXISTS sandbox ( tunnel_urls TEXT, -- JSON mapping of port -> tunnel URL for extra ports ttyd_url TEXT, -- ttyd proxy tunnel URL ttyd_token TEXT, -- Encrypted JWT token for ttyd auth + active_socket_id TEXT, -- Bridge socket the session dispatches to (socket: tag) created_at INTEGER NOT NULL ); @@ -648,6 +649,11 @@ export const MIGRATIONS: readonly SchemaMigration[] = [ description: "Persist terminal message projections awaiting retry", run: TERMINAL_MESSAGE_PROJECTION_TABLE_SQL, }, + { + id: 48, + description: "Add active_socket_id to sandbox", + run: `ALTER TABLE sandbox ADD COLUMN active_socket_id TEXT`, + }, ]; /** diff --git a/packages/control-plane/src/session/server.test.ts b/packages/control-plane/src/session/server.test.ts index eed3b9a873..ff9cfaeaff 100644 --- a/packages/control-plane/src/session/server.test.ts +++ b/packages/control-plane/src/session/server.test.ts @@ -44,6 +44,7 @@ function createHarness() { send: vi.fn(() => true), getClient: vi.fn(() => currentClient), close: vi.fn(), + isActiveSandbox: vi.fn(() => true), clearSandboxIfMatch: vi.fn(() => true), removeClient: vi.fn(() => client), hasParticipant: vi.fn(() => false), @@ -320,6 +321,29 @@ describe("SessionServer", () => { }); }); + it("refuses frames from a replaced sandbox socket and closes it again", async () => { + const { server, messageDeps, sockets, log, setConnectionKind } = createHarness(); + setConnectionKind("sandbox"); + vi.mocked(sockets.isActiveSandbox).mockReturnValue(false); + + await server.onMessage( + "sandbox", + JSON.stringify({ + type: "heartbeat", + sandboxId: "sandbox-1", + timestamp: 1000, + status: "ready", + }) + ); + + expect(messageDeps.processSandboxEvent).not.toHaveBeenCalled(); + expect(sockets.close).toHaveBeenCalledWith("sandbox", 1000, "Sandbox socket replaced"); + expect(log.debug).toHaveBeenCalledWith( + "Ignoring frame from a replaced sandbox socket", + expect.objectContaining({ sandbox_id: "sandbox-1" }) + ); + }); + it("schedules sandbox reconnect checks and always reciprocates close", async () => { const { server, sockets, sandbox, setConnectionKind } = createHarness(); setConnectionKind("sandbox"); diff --git a/packages/control-plane/src/session/types.ts b/packages/control-plane/src/session/types.ts index 626b003953..fdf82acacc 100644 --- a/packages/control-plane/src/session/types.ts +++ b/packages/control-plane/src/session/types.ts @@ -170,6 +170,8 @@ export interface SandboxRow { tunnel_urls: string | null; // JSON mapping of port -> tunnel URL ttyd_url: string | null; ttyd_token: string | null; + /** The `socket:` tag of the bridge socket the session dispatches to. */ + active_socket_id: string | null; created_at: number; } diff --git a/packages/control-plane/src/session/websocket-manager.test.ts b/packages/control-plane/src/session/websocket-manager.test.ts index 0a56ff0754..0598f4242a 100644 --- a/packages/control-plane/src/session/websocket-manager.test.ts +++ b/packages/control-plane/src/session/websocket-manager.test.ts @@ -108,6 +108,11 @@ function createMockRepository() { const repo = { getSandbox: () => sandboxRow, + getActiveSocketId: () => sandboxRow?.active_socket_id ?? null, + setActiveSocketId: (socketId: string) => { + // Like the UPDATE it stands in for: nothing to write without a row. + if (sandboxRow) sandboxRow.active_socket_id = socketId; + }, getWsClientMapping: (wsId: string) => mappings.get(wsId) ?? null, hasWsClientMapping: (wsId: string) => mappings.has(wsId), upsertWsClientMapping: (data: { @@ -195,6 +200,7 @@ function createSandboxRow(modalSandboxId: string): SandboxRow { tunnel_urls: null, ttyd_url: null, ttyd_token: null, + active_socket_id: null, created_at: Date.now(), }; } @@ -304,7 +310,34 @@ describe("SessionWebSocketManagerImpl", () => { manager.acceptAndSetSandboxSocket(ws); - expect(sockets.get(ws)).toEqual(["sandbox"]); + expect(sockets.get(ws)).toEqual(["sandbox", expect.stringMatching(/^socket:sbws-/)]); + }); + + it("persists the new socket's identity before closing the socket it replaces", () => { + const { manager, mockRepo } = createManager(); + mockRepo.setSandbox(createSandboxRow("sb-1")); + const order: string[] = []; + const oldWs = createFakeWebSocket(); + vi.mocked(oldWs.close).mockImplementation(() => { + order.push(`close:${mockRepo.repo.getActiveSocketId()}`); + }); + const newWs = createFakeWebSocket(); + + manager.acceptAndSetSandboxSocket(oldWs, "sb-1"); + const oldId = mockRepo.repo.getActiveSocketId(); + manager.acceptAndSetSandboxSocket(newWs, "sb-1"); + const newId = mockRepo.repo.getActiveSocketId(); + + expect(oldId).toMatch(/^sbws-/); + expect(newId).toMatch(/^sbws-/); + expect(newId).not.toBe(oldId); + // The row already named the replacement when the old socket was closed. + expect(order).toEqual([`close:${newId}`]); + expect(manager.classify(newWs)).toEqual({ + kind: "sandbox", + sandboxId: "sb-1", + socketId: newId, + }); }); it("closes existing sandbox socket and returns replaced=true", () => { @@ -345,7 +378,8 @@ describe("SessionWebSocketManagerImpl", () => { }); it("sets new socket as active sandbox", () => { - const { manager } = createManager(); + const { manager, mockRepo } = createManager(); + mockRepo.setSandbox(createSandboxRow("sb-1")); const ws = createFakeWebSocket(); manager.acceptAndSetSandboxSocket(ws, "sb-1"); @@ -354,9 +388,63 @@ describe("SessionWebSocketManagerImpl", () => { }); }); + describe("isActiveSandboxSocket", () => { + it("is true only for the most recently accepted sandbox socket", () => { + const { manager, mockRepo } = createManager(); + mockRepo.setSandbox(createSandboxRow("sb-1")); + const oldWs = createFakeWebSocket(); + const newWs = createFakeWebSocket(); + + manager.acceptAndSetSandboxSocket(oldWs, "sb-1"); + expect(manager.isActiveSandboxSocket(oldWs)).toBe(true); + + // Same sandbox reconnecting: the replaced socket is still OPEN while + // its close completes, and still tagged, but no longer authoritative. + manager.acceptAndSetSandboxSocket(newWs, "sb-1"); + expect(manager.isActiveSandboxSocket(oldWs)).toBe(false); + expect(manager.isActiveSandboxSocket(newWs)).toBe(true); + }); + + it("is false for client sockets", () => { + const { manager, mockRepo } = createManager(); + mockRepo.setSandbox(createSandboxRow("sb-1")); + const ws = createFakeWebSocket(); + manager.acceptClientSocket(ws, "ws-1"); + + expect(manager.isActiveSandboxSocket(ws)).toBe(false); + }); + + it("is false once a spawn reservation has cleared the persisted identity", () => { + const { manager, mockRepo } = createManager(); + const row = createSandboxRow("sb-1"); + mockRepo.setSandbox(row); + const ws = createFakeWebSocket(); + manager.acceptAndSetSandboxSocket(ws, "sb-1"); + + row.active_socket_id = null; + row.modal_sandbox_id = "sb-2"; + + expect(manager.isActiveSandboxSocket(ws)).toBe(false); + }); + + it("keeps a socket accepted before identities were persisted authoritative", () => { + const { manager, sockets, mockRepo } = createManager(); + mockRepo.setSandbox(createSandboxRow("sb-1")); + const legacyWs = createFakeWebSocket(); + sockets.set(legacyWs, ["sandbox", "sid:sb-1"]); + + expect(manager.isActiveSandboxSocket(legacyWs)).toBe(true); + + const newWs = createFakeWebSocket(); + manager.acceptAndSetSandboxSocket(newWs, "sb-1"); + expect(manager.isActiveSandboxSocket(legacyWs)).toBe(false); + }); + }); + describe("getSandboxSocket", () => { it("returns cached socket if open", () => { - const { manager } = createManager(); + const { manager, mockRepo } = createManager(); + mockRepo.setSandbox(createSandboxRow("sb-1")); const ws = createFakeWebSocket(); manager.acceptAndSetSandboxSocket(ws, "sb-1"); @@ -364,6 +452,34 @@ describe("SessionWebSocketManagerImpl", () => { expect(manager.getSandboxSocket()).toBe(ws); }); + it("recovers the socket whose identity the row names, not the first open one", () => { + const { manager, sockets, mockRepo } = createManager(); + const row = createSandboxRow("sb-1"); + row.active_socket_id = "sbws-active"; + mockRepo.setSandbox(row); + const replacedWs = createFakeWebSocket(); + const activeWs = createFakeWebSocket(); + + // Both sockets belong to the same sandbox and both still look OPEN + // after a restart; the replaced one is enumerated first. + sockets.set(replacedWs, ["sandbox", "sid:sb-1", "socket:sbws-replaced"]); + sockets.set(activeWs, ["sandbox", "sid:sb-1", "socket:sbws-active"]); + + expect(manager.getSandboxSocket()).toBe(activeWs); + expect(replacedWs.close).not.toHaveBeenCalled(); + }); + + it("returns null when only replaced sockets survive a restart", () => { + const { manager, sockets, mockRepo } = createManager(); + const row = createSandboxRow("sb-1"); + row.active_socket_id = "sbws-active"; + mockRepo.setSandbox(row); + const replacedWs = createFakeWebSocket(); + sockets.set(replacedWs, ["sandbox", "sid:sb-1", "socket:sbws-replaced"]); + + expect(manager.getSandboxSocket()).toBeNull(); + }); + it("returns null when no sandbox socket exists", () => { const { manager } = createManager(); expect(manager.getSandboxSocket()).toBeNull(); @@ -508,7 +624,8 @@ describe("SessionWebSocketManagerImpl", () => { describe("clearSandboxSocketIfMatch", () => { it("clears and returns true when ws matches", () => { - const { manager } = createManager(); + const { manager, mockRepo } = createManager(); + mockRepo.setSandbox(createSandboxRow("sb-1")); const ws = createFakeWebSocket(); manager.acceptAndSetSandboxSocket(ws, "sb-1"); @@ -521,12 +638,13 @@ describe("SessionWebSocketManagerImpl", () => { }); it("returns false and does not clear when ws does not match", () => { - const { manager } = createManager(); + const { manager, mockRepo } = createManager(); + mockRepo.setSandbox(createSandboxRow("sb-1")); const oldWs = createFakeWebSocket(); const newWs = createFakeWebSocket(); manager.acceptAndSetSandboxSocket(oldWs, "sb-1"); - manager.acceptAndSetSandboxSocket(newWs, "sb-2"); + manager.acceptAndSetSandboxSocket(newWs, "sb-1"); // Try to clear with old socket — should not affect new socket const result = manager.clearSandboxSocketIfMatch(oldWs); @@ -535,13 +653,33 @@ describe("SessionWebSocketManagerImpl", () => { expect(manager.getSandboxSocket()).toBe(newWs); }); - it("returns true when no sandbox socket is set (post-hibernation)", () => { - const { manager } = createManager(); + it("recognizes the active socket after a restart by its persisted identity", () => { + const { manager, sockets, mockRepo } = createManager(); + const row = createSandboxRow("sb-1"); + row.active_socket_id = "sbws-active"; + mockRepo.setSandbox(row); + const activeWs = createFakeWebSocket(); + const replacedWs = createFakeWebSocket(); + sockets.set(activeWs, ["sandbox", "sid:sb-1", "socket:sbws-active"]); + sockets.set(replacedWs, ["sandbox", "sid:sb-1", "socket:sbws-replaced"]); + + expect(manager.clearSandboxSocketIfMatch(replacedWs)).toBe(false); + expect(manager.clearSandboxSocketIfMatch(activeWs)).toBe(true); + }); + + it("treats the close of a socket a spawn reservation displaced as a replacement", () => { + const { manager, mockRepo } = createManager(); + const row = createSandboxRow("sb-1"); + mockRepo.setSandbox(row); const ws = createFakeWebSocket(); + manager.acceptAndSetSandboxSocket(ws, "sb-1"); + row.active_socket_id = null; + row.modal_sandbox_id = "sb-2"; - // When sandboxWs is null (e.g., post-hibernation), the closing socket - // is treated as active since there's no replacement to compare against. - expect(manager.clearSandboxSocketIfMatch(ws)).toBe(true); + expect(manager.clearSandboxSocketIfMatch(ws)).toBe(false); + // The pointer is still dropped: nothing may keep sending into it. + Object.defineProperty(ws, "readyState", { value: WebSocket.CLOSED }); + expect(manager.getSandboxSocket()).toBeNull(); }); }); diff --git a/packages/control-plane/src/session/websocket-manager.ts b/packages/control-plane/src/session/websocket-manager.ts index fdf4fa8100..ecec51bd7d 100644 --- a/packages/control-plane/src/session/websocket-manager.ts +++ b/packages/control-plane/src/session/websocket-manager.ts @@ -5,6 +5,12 @@ * The manager owns socket identity, persistence, and authorization leases. * The connection authenticator builds ClientInfo and stores it here after * snapshot synchronization. + * + * Exactly one sandbox socket is authoritative at a time. Accepting a bridge + * tags its socket with a fresh `socket:` and persists that id on the + * sandbox row before the socket is published; dispatch, recovery after a + * restart, and close handling all compare a socket's tag against the row. + * Closing a replaced socket is cleanup — the persisted identity is the fence. */ import type { Logger } from "../logger"; @@ -37,11 +43,14 @@ export interface SessionWebSocketManager { acceptClientSocket(ws: WebSocket, wsId: string): void; /** - * Accept a sandbox WebSocket, close any existing sandbox socket, and set - * as the active sandbox connection. + * Accept a sandbox WebSocket as the session's active bridge: persist its + * identity, then close every other sandbox socket. */ acceptAndSetSandboxSocket(ws: WebSocket, sandboxId?: string): { replaced: boolean }; + /** Whether `ws` is the sandbox socket the session currently dispatches to. */ + isActiveSandboxSocket(ws: WebSocket): boolean; + /** Parse a WebSocket's tags to determine its kind and identity. */ classify(ws: WebSocket): ConnectionClassification; @@ -57,7 +66,11 @@ export interface SessionWebSocketManager { /** 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. */ + /** + * Drop the in-memory pointer for a closing sandbox socket. Returns whether + * it was the active socket — whether its close is the loss of the session's + * bridge rather than the tail of a replacement. + */ clearSandboxSocketIfMatch(ws: WebSocket): boolean; setClient(ws: WebSocket, info: ClientInfo): void; @@ -130,8 +143,14 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { } acceptAndSetSandboxSocket(ws: WebSocket, sandboxId?: string): { replaced: boolean } { - const tags = ["sandbox", ...(sandboxId ? [`sid:${sandboxId}`] : [])]; + const socketId = `sbws-${Date.now()}-${Math.random().toString(36).slice(2, 10)}`; + const tags = ["sandbox", ...(sandboxId ? [`sid:${sandboxId}`] : []), `socket:${socketId}`]; this.host.accept(ws, tags); + // Advance the persisted identity before anything else observes the new + // socket: from here on every earlier bridge socket is refused at dispatch + // whether or not its close below completes, and recovery after a restart + // re-selects this socket by its tag. + this.sandboxRepository.setActiveSocketId(socketId); let replaced = false; // Close every other live sandbox socket, not only the cached one: after @@ -161,7 +180,8 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { const tags = this.host.tags(ws); if (tags.includes("sandbox")) { const sidTag = tags.find((t) => t.startsWith("sid:")); - return { kind: "sandbox", sandboxId: sidTag?.slice(4) }; + const socketTag = tags.find((t) => t.startsWith("socket:")); + return { kind: "sandbox", sandboxId: sidTag?.slice(4), socketId: socketTag?.slice(7) }; } const wsIdTag = tags.find((t) => t.startsWith("wsid:")); return { kind: "client", wsId: wsIdTag?.slice(5) }; @@ -171,9 +191,18 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { // Sandbox socket // ------------------------------------------------------------------------- + isActiveSandboxSocket(ws: WebSocket): boolean { + const parsed = this.classify(ws); + if (parsed.kind !== "sandbox") return false; + // A socket accepted before identities were persisted (no tag, no row + // value) stays authoritative until the next bridge connects. + return (parsed.socketId ?? null) === this.sandboxRepository.getActiveSocketId(); + } + getSandboxSocket(): WebSocket | null { const sandbox = this.sandboxRepository.getSandbox(); const expectedSandboxId = sandbox?.modal_sandbox_id; + const activeSocketId = sandbox?.active_socket_id ?? null; // If the sandbox is in a terminal state, don't re-adopt stale WebSockets. // After inactivity timeout or heartbeat stale, the DO closes the WS and sets @@ -195,8 +224,9 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { if (this.sandboxWs?.readyState === WebSocket.OPEN) { const cached = this.classify(this.sandboxWs); if ( - !expectedSandboxId || - (cached.kind === "sandbox" && cached.sandboxId === expectedSandboxId) + cached.kind === "sandbox" && + (!expectedSandboxId || cached.sandboxId === expectedSandboxId) && + (cached.socketId ?? null) === activeSocketId ) { return this.sandboxWs; } @@ -204,8 +234,10 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { this.sandboxWs = null; } - // Hibernation recovery: scan all WebSockets, validate sandbox identity - + // Recovery after a restart: the pointer is gone but the accepted sockets + // and their tags survive. Only the socket carrying the persisted active + // identity is re-adopted — a same-sandbox socket it replaced may still be + // open while its close completes, and must not win for coming first. for (const ws of this.host.sockets()) { const parsed = this.classify(ws); if (parsed.kind !== "sandbox" || ws.readyState !== WebSocket.OPEN) continue; @@ -218,6 +250,7 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { this.close(ws, 1000, "Sandbox identity changed"); continue; } + if ((parsed.socketId ?? null) !== activeSocketId) continue; this.log.info("Recovered sandbox WebSocket from hibernation"); this.sandboxWs = ws; @@ -242,13 +275,9 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { } clearSandboxSocketIfMatch(ws: WebSocket): boolean { - if (this.sandboxWs === ws) { - this.sandboxWs = null; - return true; - } - // sandboxWs is null (post-hibernation or already cleared) — treat as active. - // The only definitive "replaced" signal is sandboxWs pointing to a different socket. - return this.sandboxWs === null; + const active = this.isActiveSandboxSocket(ws); + if (active || this.sandboxWs === ws) this.sandboxWs = null; + return active; } // ------------------------------------------------------------------------- diff --git a/packages/control-plane/test/integration/durable-object-eviction.test.ts b/packages/control-plane/test/integration/durable-object-eviction.test.ts index 593db74a2d..b9e25f19ce 100644 --- a/packages/control-plane/test/integration/durable-object-eviction.test.ts +++ b/packages/control-plane/test/integration/durable-object-eviction.test.ts @@ -1,5 +1,6 @@ import { beforeEach, describe, expect, it } from "vitest"; import { env, runDurableObjectAlarm } from "cloudflare:test"; +import type { SandboxStatus } from "@open-inspect/shared/types/sessions"; import { runInSessionDO } from "./session-do-access"; import type { SessionDO } from "../../src/session/durable-object"; import { cleanD1Tables } from "./cleanup"; @@ -7,8 +8,10 @@ import { INTEGRATION_WEBSOCKET_TIMEOUT_MS, initNamedSession, openClientWs, + openSandboxWs, queryDO, seedMessage, + seedSandboxAuth, waitForSandboxStatus, } from "./helpers"; @@ -16,10 +19,17 @@ const INSTANCE_MARKER = "pre-eviction-instance"; type MarkedSessionDO = SessionDO & { __evictionMarker?: string }; -/** Tear down the running instance and return a stub bound to its replacement. */ -async function evictSessionDO(sessionName: string): Promise { +/** + * Tear down the running instance and return a stub bound to its replacement. + * Waits for the sandbox to settle in `settledStatus` first so no in-flight + * status write from the (always-failing) test spawn lands on the replacement. + */ +async function evictSessionDO( + sessionName: string, + settledStatus: SandboxStatus = "failed" +): Promise { const stub = env.SESSION.get(env.SESSION.idFromName(sessionName)); - await waitForSandboxStatus(stub, "failed"); + await waitForSandboxStatus(stub, settledStatus); await expect( runInSessionDO(stub, (instance: MarkedSessionDO) => { instance.__evictionMarker = INSTANCE_MARKER; @@ -157,6 +167,58 @@ describe("SessionDO eviction and hibernation restore", () => { ]); }); + it("dispatches sandbox frames by the persisted socket identity after a restore", async () => { + const sessionName = `do-evict-sandbox-${Date.now()}`; + const sandboxId = "sb-evict"; + const { stub } = await initNamedSession(sessionName); + await seedSandboxAuth(stub, { authToken: "sandbox-token-evict", sandboxId }); + const { ws } = await openSandboxWs(sessionName, { + authToken: "sandbox-token-evict", + sandboxId, + }); + expect(ws).not.toBeNull(); + ws!.accept(); + await waitForSandboxStatus(stub, "ready"); + const [{ active_socket_id: activeSocketId }] = await queryDO<{ + active_socket_id: string | null; + }>(stub, "SELECT active_socket_id FROM sandbox"); + expect(activeSocketId).toMatch(/^sbws-/); + + const restored = await evictSessionDO(sessionName, "ready"); + const toolCall = (callId: string) => + JSON.stringify({ + type: "tool_call", + tool: "read_file", + args: { path: "/src/main.ts" }, + callId, + messageId: "msg-evict", + sandboxId, + timestamp: Date.now() / 1000, + }); + // Two same-sandbox sockets survive only as their tags: the one the row + // names is dispatched from, the one it replaced is not. + await runInSessionDO(restored, async (instance: SessionDO, state) => { + const replaced = new WebSocketPair(); + state.acceptWebSocket(replaced[1], ["sandbox", `sid:${sandboxId}`, "socket:sbws-replaced"]); + replaced[0].accept(); + const active = new WebSocketPair(); + state.acceptWebSocket(active[1], ["sandbox", `sid:${sandboxId}`, `socket:${activeSocketId}`]); + active[0].accept(); + + await instance.webSocketMessage(replaced[1], toolCall("call-replaced")); + await instance.webSocketMessage(active[1], toolCall("call-active")); + }); + + const events = await queryDO<{ data: string }>( + restored, + "SELECT data FROM events WHERE type = ?", + "tool_call" + ); + expect(events.map((event) => (JSON.parse(event.data) as { callId: string }).callId)).toEqual([ + "call-active", + ]); + }); + it("rebuilds client identity from ws_client_mapping when the in-memory cache is gone", async () => { const sessionName = `do-evict-identity-${Date.now()}`; await initNamedSession(sessionName); diff --git a/packages/control-plane/test/integration/websocket-sandbox.test.ts b/packages/control-plane/test/integration/websocket-sandbox.test.ts index c713218113..c7ad1274be 100644 --- a/packages/control-plane/test/integration/websocket-sandbox.test.ts +++ b/packages/control-plane/test/integration/websocket-sandbox.test.ts @@ -317,6 +317,66 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { replacementWs!.close(); }); + it("dispatches only from the sandbox socket whose identity the session persisted", async () => { + const name = `ws-sandbox-authority-${Date.now()}`; + const { stub } = await initNamedSession(name); + await seedSandboxAuth(stub, { authToken: SANDBOX_TOKEN, sandboxId: SANDBOX_ID }); + + const { ws: activeWs } = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + expect(activeWs).not.toBeNull(); + activeWs!.accept(); + + const [{ active_socket_id: activeSocketId }] = await queryDO<{ + active_socket_id: string | null; + }>(stub, "SELECT active_socket_id FROM sandbox"); + expect(activeSocketId).toMatch(/^sbws-/); + + const toolCall = (callId: string) => + JSON.stringify({ + type: "tool_call", + tool: "read_file", + args: { path: "/src/main.ts" }, + callId, + messageId: "msg-authority", + sandboxId: SANDBOX_ID, + timestamp: Date.now() / 1000, + }); + + // A bridge the active one replaced, in the shape a restart leaves it: + // still accepted at the platform level under the same sandbox id, still + // open, carrying an identity the row no longer names. + await runInSessionDO(stub, async (instance: SessionDO, state) => { + const pair = new WebSocketPair(); + state.acceptWebSocket(pair[1], ["sandbox", `sid:${SANDBOX_ID}`, "socket:sbws-replaced"]); + pair[0].accept(); + expect(openSandboxSockets(state)).toHaveLength(2); + + await instance.webSocketMessage(pair[1], toolCall("call-replaced")); + + // Refused, and closed again rather than left open. + const live = openSandboxSockets(state); + expect(live).toHaveLength(1); + expect(state.getTags(live[0])).toContain(`socket:${activeSocketId}`); + }); + + activeWs!.send(toolCall("call-active")); + await new Promise((r) => setTimeout(r, 200)); + + const events = await queryDO<{ data: string }>( + stub, + "SELECT data FROM events WHERE type = ?", + "tool_call" + ); + const callIds = events.map((event) => (JSON.parse(event.data) as { callId: string }).callId); + expect(callIds).toContain("call-active"); + expect(callIds).not.toContain("call-replaced"); + + activeWs!.close(); + }); + it("sandbox connect sets status to ready", async () => { const name = `ws-sandbox-ready-${Date.now()}`; const { stub } = await initNamedSession(name); From bbdbb42e6bfe986bdfa9d77d91a696fa33832a3b Mon Sep 17 00:00:00 2001 From: Cole Murray Date: Thu, 3 Sep 2026 21:35:27 -0700 Subject: [PATCH 2/2] Revoke sandbox socket authority on detach and spawn, distinct from legacy NULL The persisted identity was only advanced on accept. Detaching a socket (heartbeat stale, inactivity, unresponsive sends, fatal runtime errors) left the row naming it, so a trailing frame still passed the dispatch check and a restart could re-adopt the still-open socket. The spawn reservation cleared the id to NULL, which the migration-compatibility branch reads as "untagged sockets are authoritative", reopening the fence for a displaced legacy socket during rollout. `active_socket_id` is now three-valued: a tag id, '' for revoked, and NULL only on rows that predate persisted identities. Detach revokes before closing; the spawn reservation revokes instead of clearing; the legacy branch additionally requires the socket's sandbox id to match the row. One predicate now answers dispatch, recovery, and close handling. --- .../control-plane/src/session/components.ts | 2 + .../src/session/sandbox-repository.test.ts | 29 +++---- .../src/session/sandbox-repository.ts | 23 ++--- packages/control-plane/src/session/types.ts | 5 +- .../src/session/websocket-manager.test.ts | 84 +++++++++++++++++-- .../src/session/websocket-manager.ts | 44 ++++++---- .../integration/websocket-sandbox.test.ts | 58 +++++++++++++ 7 files changed, 194 insertions(+), 51 deletions(-) diff --git a/packages/control-plane/src/session/components.ts b/packages/control-plane/src/session/components.ts index 6bbbb04a2f..16936862dd 100644 --- a/packages/control-plane/src/session/components.ts +++ b/packages/control-plane/src/session/components.ts @@ -169,6 +169,7 @@ export interface SessionRuntime { */ export interface SessionComponents { sandboxRepository: SandboxRepository; + wsManager: SessionWebSocketManager; /** * Assignable — the setter swaps the underlying cell for tests. Substitution * swaps operations only: the provider NAME was captured at construction and @@ -826,6 +827,7 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi const components: SessionComponents = { sandboxRepository, + wsManager, // Accessor pair over the local cell: production reads never go through // this property; the setter is the live-DO integration seam. get sourceControlProvider() { diff --git a/packages/control-plane/src/session/sandbox-repository.test.ts b/packages/control-plane/src/session/sandbox-repository.test.ts index 77e435b57c..7a2e071c1d 100644 --- a/packages/control-plane/src/session/sandbox-repository.test.ts +++ b/packages/control-plane/src/session/sandbox-repository.test.ts @@ -131,8 +131,9 @@ describe("SandboxRepository", () => { expect(mock.calls[0].query).toContain("vnc_password = NULL"); // A replacement sandbox must not inherit the predecessor's runtime. expect(mock.calls[0].query).toContain("runtime_version = NULL"); - // ...nor its bridge: the predecessor's socket loses dispatch authority here. - expect(mock.calls[0].query).toContain("active_socket_id = NULL"); + // ...nor its bridge: the predecessor's socket loses dispatch authority + // here. Revoked is '' — NULL is reserved for rows that predate identities. + expect(mock.calls[0].query).toContain("active_socket_id = ''"); expect(mock.calls[0].params).toEqual(["spawning", 1000, "modal-sb-1"]); }); @@ -149,22 +150,6 @@ describe("SandboxRepository", () => { }); describe("active socket id", () => { - const query = `SELECT active_socket_id FROM sandbox LIMIT 1`; - - it("reads null before any bridge has connected", () => { - mock.setData(query, [{ active_socket_id: null }]); - expect(repository.getActiveSocketId()).toBeNull(); - }); - - it("reads null without a sandbox row", () => { - expect(repository.getActiveSocketId()).toBeNull(); - }); - - it("reads the persisted identity", () => { - mock.setData(query, [{ active_socket_id: "sbws-1" }]); - expect(repository.getActiveSocketId()).toBe("sbws-1"); - }); - it("writes the identity to the session's one sandbox row", () => { repository.setActiveSocketId("sbws-2"); @@ -172,6 +157,14 @@ describe("SandboxRepository", () => { expect(mock.calls[0].query).toContain("UPDATE sandbox SET active_socket_id = ?"); expect(mock.calls[0].params).toEqual(["sbws-2"]); }); + + it("revokes with the empty sentinel rather than NULL", () => { + repository.revokeActiveSocketId(); + + expect(mock.calls.length).toBe(1); + expect(mock.calls[0].query).toContain("UPDATE sandbox SET active_socket_id = ''"); + expect(mock.calls[0].params).toEqual([]); + }); }); describe("updateSandboxAuthTokenHash", () => { diff --git a/packages/control-plane/src/session/sandbox-repository.ts b/packages/control-plane/src/session/sandbox-repository.ts index 0e4bf6a2c0..531ba9ee5a 100644 --- a/packages/control-plane/src/session/sandbox-repository.ts +++ b/packages/control-plane/src/session/sandbox-repository.ts @@ -138,7 +138,7 @@ export class SandboxRepository { ttyd_url = NULL, ttyd_token = NULL, runtime_version = NULL, - active_socket_id = NULL + active_socket_id = '' WHERE id = (SELECT id FROM sandbox LIMIT 1)`, data.status, data.createdAt, @@ -147,16 +147,12 @@ export class SandboxRepository { } /** - * The bridge socket the session dispatches to, by its `socket:` tag. - * Null once a spawn reserves a new identity or before any bridge connects. + * Make `socketId` the socket the session dispatches to; every earlier + * socket loses authority. `active_socket_id` is three-valued: a tag id, + * `''` for revoked (no socket matches, see `revokeActiveSocketId` and the + * spawn reservation above), and NULL only on rows that predate persisted + * identities. */ - getActiveSocketId(): string | null { - const result = this.sql.exec(`SELECT active_socket_id FROM sandbox LIMIT 1`); - const rows = this.rows<{ active_socket_id: string | null }>(result); - return rows[0]?.active_socket_id ?? null; - } - - /** Make `socketId` the socket the session dispatches to; every earlier socket loses authority. */ setActiveSocketId(socketId: string): void { this.sql.exec( `UPDATE sandbox SET active_socket_id = ? WHERE id = (SELECT id FROM sandbox LIMIT 1)`, @@ -164,6 +160,13 @@ export class SandboxRepository { ); } + /** Leave the session with no authoritative bridge socket until the next accept. */ + revokeActiveSocketId(): void { + this.sql.exec( + `UPDATE sandbox SET active_socket_id = '' WHERE id = (SELECT id FROM sandbox LIMIT 1)` + ); + } + /** * Phase 2 of the two-phase spawn write (#1589): publish the reserved * identity's hash. Scoped to that identity so a delayed publisher cannot diff --git a/packages/control-plane/src/session/types.ts b/packages/control-plane/src/session/types.ts index fdf82acacc..9cfc5374c3 100644 --- a/packages/control-plane/src/session/types.ts +++ b/packages/control-plane/src/session/types.ts @@ -170,7 +170,10 @@ export interface SandboxRow { tunnel_urls: string | null; // JSON mapping of port -> tunnel URL ttyd_url: string | null; ttyd_token: string | null; - /** The `socket:` tag of the bridge socket the session dispatches to. */ + /** + * The `socket:` tag of the bridge socket the session dispatches to; + * `''` once revoked, NULL only on rows that predate persisted identities. + */ active_socket_id: string | null; created_at: number; } diff --git a/packages/control-plane/src/session/websocket-manager.test.ts b/packages/control-plane/src/session/websocket-manager.test.ts index 0598f4242a..ea8636c54c 100644 --- a/packages/control-plane/src/session/websocket-manager.test.ts +++ b/packages/control-plane/src/session/websocket-manager.test.ts @@ -108,11 +108,13 @@ function createMockRepository() { const repo = { getSandbox: () => sandboxRow, - getActiveSocketId: () => sandboxRow?.active_socket_id ?? null, setActiveSocketId: (socketId: string) => { // Like the UPDATE it stands in for: nothing to write without a row. if (sandboxRow) sandboxRow.active_socket_id = socketId; }, + revokeActiveSocketId: () => { + if (sandboxRow) sandboxRow.active_socket_id = ""; + }, getWsClientMapping: (wsId: string) => mappings.get(wsId) ?? null, hasWsClientMapping: (wsId: string) => mappings.has(wsId), upsertWsClientMapping: (data: { @@ -315,18 +317,19 @@ describe("SessionWebSocketManagerImpl", () => { it("persists the new socket's identity before closing the socket it replaces", () => { const { manager, mockRepo } = createManager(); - mockRepo.setSandbox(createSandboxRow("sb-1")); + const row = createSandboxRow("sb-1"); + mockRepo.setSandbox(row); const order: string[] = []; const oldWs = createFakeWebSocket(); vi.mocked(oldWs.close).mockImplementation(() => { - order.push(`close:${mockRepo.repo.getActiveSocketId()}`); + order.push(`close:${row.active_socket_id}`); }); const newWs = createFakeWebSocket(); manager.acceptAndSetSandboxSocket(oldWs, "sb-1"); - const oldId = mockRepo.repo.getActiveSocketId(); + const oldId = row.active_socket_id; manager.acceptAndSetSandboxSocket(newWs, "sb-1"); - const newId = mockRepo.repo.getActiveSocketId(); + const newId = row.active_socket_id; expect(oldId).toMatch(/^sbws-/); expect(newId).toMatch(/^sbws-/); @@ -414,19 +417,56 @@ describe("SessionWebSocketManagerImpl", () => { expect(manager.isActiveSandboxSocket(ws)).toBe(false); }); - it("is false once a spawn reservation has cleared the persisted identity", () => { + it("is false once a spawn reservation has revoked the persisted identity", () => { const { manager, mockRepo } = createManager(); const row = createSandboxRow("sb-1"); mockRepo.setSandbox(row); const ws = createFakeWebSocket(); manager.acceptAndSetSandboxSocket(ws, "sb-1"); - row.active_socket_id = null; + // What updateSandboxForSpawn writes. + row.active_socket_id = ""; row.modal_sandbox_id = "sb-2"; expect(manager.isActiveSandboxSocket(ws)).toBe(false); }); + it("is false for every socket once detach has revoked authority", () => { + const { manager, sockets, mockRepo } = createManager(); + const row = createSandboxRow("sb-1"); + mockRepo.setSandbox(row); + const ws = createFakeWebSocket(); + manager.acceptAndSetSandboxSocket(ws, "sb-1"); + + manager.detachSandboxSocket(1000, "Heartbeat stale"); + + expect(row.active_socket_id).toBe(""); + expect(ws.close).toHaveBeenCalledWith(1000, "Heartbeat stale"); + // The close is cleanup; the row is the fence, so a trailing frame from + // the still-tagged socket is refused whether or not the close landed. + expect(manager.isActiveSandboxSocket(ws)).toBe(false); + expect(manager.clearSandboxSocketIfMatch(ws)).toBe(false); + // Nor does a restart re-adopt it, even after an in-place resume. + row.status = "connecting"; + expect(sockets.get(ws)).toBeDefined(); + expect(manager.getSandboxSocket()).toBeNull(); + }); + + it("revokes authority before closing on detach", () => { + const { manager, mockRepo } = createManager(); + const row = createSandboxRow("sb-1"); + mockRepo.setSandbox(row); + const ws = createFakeWebSocket(); + vi.mocked(ws.close).mockImplementation(() => { + expect(row.active_socket_id).toBe(""); + }); + manager.acceptAndSetSandboxSocket(ws, "sb-1"); + + manager.detachSandboxSocket(1011, "Fatal sandbox runtime error"); + + expect(ws.close).toHaveBeenCalledOnce(); + }); + it("keeps a socket accepted before identities were persisted authoritative", () => { const { manager, sockets, mockRepo } = createManager(); mockRepo.setSandbox(createSandboxRow("sb-1")); @@ -439,6 +479,34 @@ describe("SessionWebSocketManagerImpl", () => { manager.acceptAndSetSandboxSocket(newWs, "sb-1"); expect(manager.isActiveSandboxSocket(legacyWs)).toBe(false); }); + + it("requires a pre-identity socket to belong to the row's sandbox", () => { + const { manager, sockets, mockRepo } = createManager(); + mockRepo.setSandbox(createSandboxRow("sb-2")); + const staleWs = createFakeWebSocket(); + sockets.set(staleWs, ["sandbox", "sid:sb-1"]); + + expect(manager.isActiveSandboxSocket(staleWs)).toBe(false); + }); + + it("refuses a pre-identity socket once a spawn reservation has revoked authority", () => { + const { manager, sockets, mockRepo } = createManager(); + const row = createSandboxRow("sb-1"); + mockRepo.setSandbox(row); + const legacyWs = createFakeWebSocket(); + sockets.set(legacyWs, ["sandbox", "sid:sb-1"]); + expect(manager.isActiveSandboxSocket(legacyWs)).toBe(true); + + // The reservation revokes rather than clears, so the migration + // compatibility branch never reopens for a displaced sandbox. + row.active_socket_id = ""; + row.modal_sandbox_id = "sb-2"; + + expect(manager.isActiveSandboxSocket(legacyWs)).toBe(false); + expect(manager.clearSandboxSocketIfMatch(legacyWs)).toBe(false); + expect(manager.getSandboxSocket()).toBeNull(); + expect(legacyWs.close).toHaveBeenCalledWith(1000, "Sandbox identity changed"); + }); }); describe("getSandboxSocket", () => { @@ -673,7 +741,7 @@ describe("SessionWebSocketManagerImpl", () => { mockRepo.setSandbox(row); const ws = createFakeWebSocket(); manager.acceptAndSetSandboxSocket(ws, "sb-1"); - row.active_socket_id = null; + row.active_socket_id = ""; row.modal_sandbox_id = "sb-2"; expect(manager.clearSandboxSocketIfMatch(ws)).toBe(false); diff --git a/packages/control-plane/src/session/websocket-manager.ts b/packages/control-plane/src/session/websocket-manager.ts index ecec51bd7d..949dbe9626 100644 --- a/packages/control-plane/src/session/websocket-manager.ts +++ b/packages/control-plane/src/session/websocket-manager.ts @@ -19,6 +19,7 @@ import type { ClientInfo } from "../types"; import type { SocketHost } from "./platform"; import type { ConnectionClassification } from "./ports"; import type { SandboxRepository } from "./sandbox-repository"; +import type { SandboxRow } from "./types"; import type { WsClientMappingRepository, WsClientMappingResult, @@ -63,7 +64,11 @@ export interface SessionWebSocketManager { /** Clear the in-memory sandbox socket reference. */ clearSandboxSocket(): void; - /** Clear and close all active sandbox sockets without consulting persisted dispatch status. */ + /** + * Revoke the persisted dispatch authority, then close every sandbox socket + * without consulting the sandbox status. Nothing dispatches from, or is + * recovered as, the bridge until the next accept. + */ detachSandboxSocket(code: number, reason: string): void; /** @@ -192,17 +197,30 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { // ------------------------------------------------------------------------- isActiveSandboxSocket(ws: WebSocket): boolean { - const parsed = this.classify(ws); - if (parsed.kind !== "sandbox") return false; - // A socket accepted before identities were persisted (no tag, no row - // value) stays authoritative until the next bridge connects. - return (parsed.socketId ?? null) === this.sandboxRepository.getActiveSocketId(); + return this.isAuthoritative(this.classify(ws), this.sandboxRepository.getSandbox()); + } + + /** + * Whether a sandbox socket is the one the row names. The row's + * `active_socket_id` is three-valued: a tag id matches exactly that socket; + * `''` (revoked by detach or a spawn reservation) matches nothing; NULL + * means the row predates persisted identities, and then an untagged socket + * for the row's sandbox stays authoritative until the next accept. + */ + private isAuthoritative(parsed: ConnectionClassification, sandbox: SandboxRow | null): boolean { + if (parsed.kind !== "sandbox" || !sandbox) return false; + if (sandbox.active_socket_id === null) { + return ( + parsed.socketId === undefined && + (!sandbox.modal_sandbox_id || parsed.sandboxId === sandbox.modal_sandbox_id) + ); + } + return parsed.socketId !== undefined && parsed.socketId === sandbox.active_socket_id; } getSandboxSocket(): WebSocket | null { const sandbox = this.sandboxRepository.getSandbox(); const expectedSandboxId = sandbox?.modal_sandbox_id; - const activeSocketId = sandbox?.active_socket_id ?? null; // If the sandbox is in a terminal state, don't re-adopt stale WebSockets. // After inactivity timeout or heartbeat stale, the DO closes the WS and sets @@ -222,12 +240,7 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { } if (this.sandboxWs?.readyState === WebSocket.OPEN) { - const cached = this.classify(this.sandboxWs); - if ( - cached.kind === "sandbox" && - (!expectedSandboxId || cached.sandboxId === expectedSandboxId) && - (cached.socketId ?? null) === activeSocketId - ) { + if (this.isAuthoritative(this.classify(this.sandboxWs), sandbox)) { return this.sandboxWs; } this.close(this.sandboxWs, 1000, "Sandbox identity changed"); @@ -250,7 +263,7 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { this.close(ws, 1000, "Sandbox identity changed"); continue; } - if ((parsed.socketId ?? null) !== activeSocketId) continue; + if (!this.isAuthoritative(parsed, sandbox)) continue; this.log.info("Recovered sandbox WebSocket from hibernation"); this.sandboxWs = ws; @@ -270,6 +283,9 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { for (const ws of this.host.sockets()) { if (this.classify(ws).kind === "sandbox") sockets.add(ws); } + // Revoke before closing: a trailing frame from a detached socket, or a + // restart that still finds it open, must not find the row naming it. + this.sandboxRepository.revokeActiveSocketId(); this.sandboxWs = null; for (const ws of sockets) this.close(ws, code, reason); } diff --git a/packages/control-plane/test/integration/websocket-sandbox.test.ts b/packages/control-plane/test/integration/websocket-sandbox.test.ts index c7ad1274be..94f2854cf5 100644 --- a/packages/control-plane/test/integration/websocket-sandbox.test.ts +++ b/packages/control-plane/test/integration/websocket-sandbox.test.ts @@ -377,6 +377,64 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { activeWs!.close(); }); + it("revokes dispatch authority on detach before the close completes", async () => { + const name = `ws-sandbox-detach-${Date.now()}`; + const { stub } = await initNamedSession(name); + await seedSandboxAuth(stub, { authToken: SANDBOX_TOKEN, sandboxId: SANDBOX_ID }); + + const { ws } = await openSandboxWs(name, { authToken: SANDBOX_TOKEN, sandboxId: SANDBOX_ID }); + expect(ws).not.toBeNull(); + ws!.accept(); + const [{ active_socket_id: activeSocketId }] = await queryDO<{ + active_socket_id: string | null; + }>(stub, "SELECT active_socket_id FROM sandbox"); + expect(activeSocketId).toMatch(/^sbws-/); + + const toolCall = (callId: string) => + JSON.stringify({ + type: "tool_call", + tool: "read_file", + args: { path: "/src/main.ts" }, + callId, + messageId: "msg-detach", + sandboxId: SANDBOX_ID, + timestamp: Date.now() / 1000, + }); + + await runInSessionDO(stub, async (instance: SessionDO, state) => { + const [serverWs] = openSandboxSockets(state); + expect(serverWs).toBeDefined(); + const { wsManager } = componentsOf(instance); + + wsManager.detachSandboxSocket(1000, "Heartbeat stale"); + + // The row was revoked, so a frame that was already in flight from the + // detached socket is refused even though the socket is still tagged. + await instance.webSocketMessage(serverWs, toolCall("call-detached")); + expect(wsManager.getSandboxSocket()).toBeNull(); + + // A restart that still finds the detached socket open, tagged with the + // identity the row named before, must not re-adopt it. + const pair = new WebSocketPair(); + state.acceptWebSocket(pair[1], ["sandbox", `sid:${SANDBOX_ID}`, `socket:${activeSocketId}`]); + pair[0].accept(); + expect(wsManager.getSandboxSocket()).toBeNull(); + await instance.webSocketMessage(pair[1], toolCall("call-restored-detached")); + }); + + const [{ active_socket_id: revoked }] = await queryDO<{ active_socket_id: string | null }>( + stub, + "SELECT active_socket_id FROM sandbox" + ); + expect(revoked).toBe(""); + const events = await queryDO<{ data: string }>( + stub, + "SELECT data FROM events WHERE type = ?", + "tool_call" + ); + expect(events).toHaveLength(0); + }); + it("sandbox connect sets status to ready", async () => { const name = `ws-sandbox-ready-${Date.now()}`; const { stub } = await initNamedSession(name);