diff --git a/packages/control-plane/src/cloudflare/session-platform.test.ts b/packages/control-plane/src/cloudflare/session-platform.test.ts new file mode 100644 index 0000000000..a16c89e54d --- /dev/null +++ b/packages/control-plane/src/cloudflare/session-platform.test.ts @@ -0,0 +1,100 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import type { SqlDatabase } from "../db/sql-database"; +import type { Logger } from "../logger"; +import { createDurableObjectSessionPlatform } from "./session-platform"; + +/** Stand-in for the Workers runtime's request/response pair. */ +class FakeRequestResponsePair { + constructor( + readonly request: string, + readonly response: string + ) {} +} + +function createFakeState() { + const storage = { + sql: { exec: vi.fn() }, + transactionSync: vi.fn((closure: () => T): T => closure()), + getAlarm: vi.fn(async () => null), + setAlarm: vi.fn(async () => {}), + deleteAlarm: vi.fn(async () => {}), + }; + const calls = { + id: { toString: () => "do-id" }, + storage, + acceptWebSocket: vi.fn(), + getTags: vi.fn(() => ["sandbox", "sid:sb-1"]), + getWebSockets: vi.fn(() => []), + setWebSocketAutoResponse: vi.fn(), + waitUntil: vi.fn(), + }; + const db = {} as SqlDatabase; + return { state: calls as unknown as DurableObjectState, storage, calls, db }; +} + +describe("createDurableObjectSessionPlatform", () => { + beforeEach(() => { + vi.stubGlobal("WebSocketRequestResponsePair", FakeRequestResponsePair); + }); + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it("exposes the object's id, storage, alarm store, and the global store", () => { + const { state, storage, db } = createFakeState(); + + const platform = createDurableObjectSessionPlatform(state, db); + + expect(platform.id).toBe("do-id"); + expect(platform.storage).toBe(storage); + expect(platform.db).toBe(db); + expect(platform.alarmStore).toBe(storage); + expect(platform.storage.transactionSync(() => 42)).toBe(42); + expect(storage.transactionSync).toHaveBeenCalledTimes(1); + }); + + it("delegates socket acceptance, tags, and enumeration, passing the tag filter through", () => { + const { state, calls, db } = createFakeState(); + const ws = {} as WebSocket; + + const { sockets: host } = createDurableObjectSessionPlatform(state, db); + host.accept(ws, ["sandbox", "sid:sb-1"]); + host.sockets(); + host.sockets("sandbox"); + + expect(calls.acceptWebSocket).toHaveBeenCalledWith(ws, ["sandbox", "sid:sb-1"]); + expect(host.tags(ws)).toEqual(["sandbox", "sid:sb-1"]); + expect(calls.getTags).toHaveBeenCalledWith(ws); + expect(calls.getWebSockets.mock.calls).toEqual([[undefined], ["sandbox"]]); + }); + + it("installs the auto-response as a request/response pair", () => { + const { state, calls, db } = createFakeState(); + + const platform = createDurableObjectSessionPlatform(state, db); + platform.sockets.setAutoResponse('{"type":"ping"}', '{"type":"pong"}'); + + expect(calls.setWebSocketAutoResponse).toHaveBeenCalledTimes(1); + const pair = calls.setWebSocketAutoResponse.mock.calls[0][0] as FakeRequestResponsePair; + expect(pair).toBeInstanceOf(FakeRequestResponsePair); + expect(pair.request).toBe('{"type":"ping"}'); + expect(pair.response).toBe('{"type":"pong"}'); + }); + + it("builds background tasks over the object's event lifetime that report to the given logger", async () => { + const { state, calls, db } = createFakeState(); + const logger = { error: vi.fn() } as unknown as Logger; + + const platform = createDurableObjectSessionPlatform(state, db); + platform.createBackgroundTasks(logger).submit(() => Promise.reject(new Error("boom")), { + name: "session.task", + }); + + expect(calls.waitUntil).toHaveBeenCalledTimes(1); + await calls.waitUntil.mock.calls[0][0]; + expect(logger.error).toHaveBeenCalledWith( + "background_task.failed", + expect.objectContaining({ task_name: "session.task" }) + ); + }); +}); diff --git a/packages/control-plane/src/cloudflare/session-platform.ts b/packages/control-plane/src/cloudflare/session-platform.ts new file mode 100644 index 0000000000..bc8e2dcbb0 --- /dev/null +++ b/packages/control-plane/src/cloudflare/session-platform.ts @@ -0,0 +1,29 @@ +import type { SqlDatabase } from "../db/sql-database"; +import type { SessionPlatform } from "../session/platform"; +import { createCloudflareBackgroundTasks } from "./background-tasks"; + +/** + * A Durable Object's storage, hibernatable sockets, alarm, and event lifetime + * as the session platform, over the deployment's global store. + */ +export function createDurableObjectSessionPlatform( + ctx: DurableObjectState, + db: SqlDatabase +): SessionPlatform { + return { + id: ctx.id.toString(), + storage: ctx.storage, + db, + alarmStore: ctx.storage, + sockets: { + accept: (ws, tags) => ctx.acceptWebSocket(ws, tags), + tags: (ws) => ctx.getTags(ws), + sockets: (tag) => ctx.getWebSockets(tag), + // Hibernation-level auto-response: matched by the runtime without + // waking the object. + setAutoResponse: (request, response) => + ctx.setWebSocketAutoResponse(new WebSocketRequestResponsePair(request, response)), + }, + createBackgroundTasks: (log) => createCloudflareBackgroundTasks(ctx, log), + }; +} diff --git a/packages/control-plane/src/session/components.ts b/packages/control-plane/src/session/components.ts index 68e4cfaf81..9313ed115f 100644 --- a/packages/control-plane/src/session/components.ts +++ b/packages/control-plane/src/session/components.ts @@ -50,6 +50,7 @@ import { requireRepoSecretsEncryptionKey, requireTokenEncryptionKey } from "../e import type { Env, ClientInfo } from "../types"; import type { SessionRow } from "./types"; import type { SqlDatabase } from "../db/sql-database"; +import type { SessionPlatform } from "./platform"; import { SessionCoreRepository } from "./session-core-repository"; import { SandboxRepository } from "./sandbox-repository"; import { SessionAttachmentRepository } from "./session-attachment-repository"; @@ -80,7 +81,6 @@ import { CallbackNotificationService } from "./callback-notification-service"; import { UserEnvResolver } from "./user-env-resolver"; import { resolveSessionRepoId } from "./repo-id-resolution"; import { Scheduler } from "../scheduler/scheduler"; -import { createCloudflareBackgroundTasks } from "../cloudflare/background-tasks"; import { PresenceService } from "./presence-service"; import { SessionMessageQueue } from "./message-queue"; import { SandboxArtifactEventHandler } from "./sandbox-events/artifact.handler"; @@ -138,13 +138,6 @@ import { AuthorizationError, AuthorizationService } from "../authorization/servi */ const WS_AUTH_TIMEOUT_MS = 30000; // 30 seconds -/** The platform surface the session graph is built over. */ -export interface SessionPlatform { - ctx: DurableObjectState; - sql: SqlStorage; - db: SqlDatabase | null; -} - /** * What the platform adapter (SessionDO) is allowed to touch. Everything else * stays inside the factory; `internals` exists for integration tests that @@ -211,9 +204,16 @@ function resolveExecutionTimeoutMs( /** Build the session runtime, including authorization verification and lease expiry handling. */ export function createSessionRuntime(platform: SessionPlatform, env: Env): SessionRuntime { - const { ctx, sql, db } = platform; - const durableObjectId = ctx.id.toString(); - const transaction = (closure: () => T): T => ctx.storage.transactionSync(closure); + const { + id: durableObjectId, + storage, + db, + alarmStore, + sockets: socketHost, + createBackgroundTasks, + } = platform; + const { sql } = storage; + const transaction = (closure: () => T): T => storage.transactionSync(closure); // Tier 1 — repositories and alarm persistence (leaves over SqlStorage). const attachmentRepository = new SessionAttachmentRepository(sql); @@ -249,29 +249,27 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi createLogger("session-do", {}, parseLogLevel(env.LOG_LEVEL)), getPublicSessionId ); - const backgroundTasks = createCloudflareBackgroundTasks(ctx, log); + const backgroundTasks = createBackgroundTasks(log); // The sandbox repository validates the status it reads and warns on anything // unmodelled, so it needs the session logger — and it owns encrypt-at-rest // for access secrets, so it takes the key. const sandboxRepository = new SandboxRepository(sql, log, repoSecretsEncryptionKey); // Tier 2 — sockets and alarm scheduling. - const alarmScheduler = createEarliestAlarmScheduler(ctx.storage, alarmDeadlines); + const alarmScheduler = createEarliestAlarmScheduler(alarmStore, alarmDeadlines); const wsManager: SessionWebSocketManager = new SessionWebSocketManagerImpl( - ctx, + socketHost, sandboxRepository, wsClientMappingRepository, alarmScheduler, log, { authTimeoutMs: WS_AUTH_TIMEOUT_MS } ); - // Hibernation-level ping/pong: the runtime answers keepalives without - // waking the Durable Object. Platform-global wiring, so it lives here. - ctx.setWebSocketAutoResponse( - new WebSocketRequestResponsePair( - JSON.stringify({ type: "ping" }), - JSON.stringify({ type: "pong", timestamp: Date.now() }) - ) + // Platform-level ping/pong: keepalives are answered without waking the + // runtime. Session-wide wiring, so it lives here. + socketHost.setAutoResponse( + JSON.stringify({ type: "ping" }), + JSON.stringify({ type: "pong", timestamp: Date.now() }) ); // Tier 3 — outbound delivery over the socket registry. @@ -288,8 +286,8 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi // Shared single instances/closures — every consumer below takes these // rather than re-deriving its own copy. - const sessionIndexStore = db ? new SessionIndexStore(db) : null; - const sessionPullRequestStore = db ? new SessionPullRequestStore(db) : null; + const sessionIndexStore = new SessionIndexStore(db); + const sessionPullRequestStore = new SessionPullRequestStore(db); const resolveRepoId = (sessionRow: SessionRow) => resolveSessionRepoId(sessionRow, sessionCoreRepository, sourceControlProvider); @@ -332,7 +330,7 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi terminalMessageCompletedAt: completedAt, }); - const userScmTokenStore = db ? new UserScmTokenStore(db, tokenEncryptionKey) : null; + const userScmTokenStore = new UserScmTokenStore(db, tokenEncryptionKey); const participantService = new ParticipantService({ repository: participantRepository, getProcessingMessageAuthor: () => messageRepository.getProcessingMessageAuthor(), @@ -342,14 +340,12 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi userScmTokenStore, }); - const scheduler = db ? new Scheduler(db, env, backgroundTasks) : undefined; + const scheduler = new Scheduler(db, env, backgroundTasks); const callbackService = new CallbackNotificationService({ repository: sessionCoreRepository, messageRepository, env, - completeAutomationRun: scheduler - ? (completion) => scheduler.runComplete(completion) - : undefined, + completeAutomationRun: (completion) => scheduler.runComplete(completion), log, getSessionId: () => resolvePublicSessionId(sessionCoreRepository.getSession(), durableObjectId), }); @@ -558,7 +554,7 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi // service around the request-scoped log, so these stay functions. const refreshOpenAIToken = async (sessionRow: SessionRow, requestLog: Logger) => { const service = new OpenAITokenRefreshService( - db!, + db, repoSecretsEncryptionKey, resolveRepoId, requestLog @@ -567,7 +563,7 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi }; const refreshXaiToken = async (sessionRow: SessionRow, requestLog: Logger) => { const service = new XaiTokenRefreshService( - db!, + db, repoSecretsEncryptionKey, resolveRepoId, requestLog @@ -585,7 +581,6 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi sandboxRepository, sandboxEventProcessor, messenger, - Boolean(db), refreshOpenAIToken, refreshXaiToken, getScmCredentials, @@ -646,7 +641,7 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi pushBranchToRemote: (pushSpec) => pushService.pushBranchToRemote(pushSpec), messenger, appName: resolveAppName(env), - sessionPullRequests: sessionPullRequestStore ?? undefined, + sessionPullRequests: sessionPullRequestStore, resolveScmSettings: (repo) => resolveScmSettings(db, repo), }); @@ -694,7 +689,6 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi schedulePullRequestRefresh, scmProviderName, resolveAuthorization: async (userId) => { - if (!db) return { kind: "unavailable" }; try { const authorization = await new AuthorizationService(db).getEffectiveAuthorization(userId); return authorization.suspendedAt === null @@ -868,7 +862,7 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi interface LifecycleManagerDeps { env: Env; - db: SqlDatabase | null; + db: SqlDatabase; /** The latched public-session-id resolver shared with the session logger. */ getSessionId: () => string; /** The repository, satisfying the manager's storage port structurally. */ @@ -913,34 +907,27 @@ function createLifecycleManager(deps: LifecycleManagerDeps): SandboxLifecycleMan env.WORKER_URL || `https://open-inspect-control-plane.${env.CF_ACCOUNT_ID || "workers"}.workers.dev`; - // Create D1-backed lookups if database is available - let mcpServerLookup: McpServerLookup | undefined; - if (db) { - const mcpStore = new McpServerStore(db, repoSecretsEncryptionKey); - mcpServerLookup = { - getDecryptedForSession: (repositories) => mcpStore.getDecryptedForSession(repositories), - }; - } + const mcpStore = new McpServerStore(db, repoSecretsEncryptionKey); + const mcpServerLookup: McpServerLookup = { + getDecryptedForSession: (repositories) => mcpStore.getDecryptedForSession(repositories), + }; // Session-scoped gate: resolved from the primary member (the scalar mirror // this lookup is called with) — see resolveSessionScopedSettings for the // per-feature scope rules. Token absence short-circuits to false so a // misconfigured deployment never installs a tool that would 503 on every call. - let slackAgentNotifyLookup: SlackAgentNotifyLookup | undefined; - if (db) { - const tokenPresent = !!env.SLACK_BOT_TOKEN; - const settingsStore = new IntegrationSettingsStore(db); - slackAgentNotifyLookup = { - isEnabledForRepo: async (repoOwner, repoName) => { - if (!tokenPresent) return false; - const settings = - repoOwner && repoName - ? (await settingsStore.getResolvedConfig("slack", `${repoOwner}/${repoName}`)).settings - : ((await settingsStore.getGlobal("slack"))?.defaults ?? {}); - return resolveSlackSettings(settings).agentNotificationsEnabled; - }, - }; - } + const tokenPresent = !!env.SLACK_BOT_TOKEN; + const settingsStore = new IntegrationSettingsStore(db); + const slackAgentNotifyLookup: SlackAgentNotifyLookup = { + isEnabledForRepo: async (repoOwner, repoName) => { + if (!tokenPresent) return false; + const settings = + repoOwner && repoName + ? (await settingsStore.getResolvedConfig("slack", `${repoOwner}/${repoName}`)).settings + : ((await settingsStore.getGlobal("slack"))?.defaults ?? {}); + return resolveSlackSettings(settings).agentNotificationsEnabled; + }, + }; const sandboxDashboardUrlBuilder = sandboxBackend === "modal" @@ -966,13 +953,11 @@ function createLifecycleManager(deps: LifecycleManagerDeps): SandboxLifecycleMan sandboxDashboardUrlBuilder, }; - // Create the image lookup if D1 is available and the provider supports - // prebuilt images. - let imageBuildLookup: ImageBuildLookup | undefined; + // The image lookup exists only for providers that support prebuilt images. const imageBuildProvider = resolveImageBuildProvider(sandboxBackend); - if (db && imageBuildProvider) { - imageBuildLookup = createImageBuildLookup(db, imageBuildProvider); - } + const imageBuildLookup: ImageBuildLookup | undefined = imageBuildProvider + ? createImageBuildLookup(db, imageBuildProvider) + : undefined; return new SandboxLifecycleManager( provider, diff --git a/packages/control-plane/src/session/durable-object.ts b/packages/control-plane/src/session/durable-object.ts index 0a606fffc1..46417582f8 100644 --- a/packages/control-plane/src/session/durable-object.ts +++ b/packages/control-plane/src/session/durable-object.ts @@ -10,25 +10,31 @@ import { DurableObject } from "cloudflare:workers"; import { initSchema } from "./schema"; import type { Env } from "../types"; -import type { SqlDatabase } from "../db/sql-database"; +import { createDurableObjectSessionPlatform } from "../cloudflare/session-platform"; +import type { SessionPlatform } from "./platform"; import { createSessionRuntime, type SessionRuntime } from "./components"; export class SessionDO extends DurableObject { - private sql: SqlStorage; /** - * The DO's global-database handle — the single point where env.DB is read. - * Nullable to preserve the existing defensive guards against a missing - * binding at runtime. Distinct from `this.sql`, the DO-embedded SQLite. + * This object's storage, sockets, alarm, and event lifetime as the ports + * the runtime is built over, with the deployment's global store. The + * constructor is the single point where env.DB is read; a missing binding + * fails construction instead of running a degraded session. */ - private readonly db: SqlDatabase | null; + private readonly platform: SessionPlatform; // The per-activation runtime; null until ensureInitialized() builds it. private _runtime: SessionRuntime | null = null; constructor(ctx: DurableObjectState, env: Env) { super(ctx, env); // eslint-disable-next-line no-restricted-syntax -- composition root input: the DO's one env.DB read - this.db = env.DB ?? null; - this.sql = ctx.storage.sql; + const db = env.DB; + if (!db) { + throw new Error( + "SessionDO requires the DB binding; sessions cannot run without the global store" + ); + } + this.platform = createDurableObjectSessionPlatform(ctx, db); } /** The runtime, (re)built on first touch after construction or eviction. */ @@ -44,8 +50,8 @@ export class SessionDO extends DurableObject { private ensureInitialized(rehydrateAlarm = true): void { if (this._runtime) return; const initStart = performance.now(); - initSchema(this.sql); - const runtime = createSessionRuntime({ ctx: this.ctx, sql: this.sql, db: this.db }, this.env); + initSchema(this.platform.storage.sql); + const runtime = createSessionRuntime(this.platform, this.env); // Publish only after the graph is fully built: a throw above leaves the // activation uninitialized, so the next event retries initialization // instead of dereferencing an undefined runtime. diff --git a/packages/control-plane/src/session/http/handlers/sandbox.handler.test.ts b/packages/control-plane/src/session/http/handlers/sandbox.handler.test.ts index 35488b95d6..7d6cbf757b 100644 --- a/packages/control-plane/src/session/http/handlers/sandbox.handler.test.ts +++ b/packages/control-plane/src/session/http/handlers/sandbox.handler.test.ts @@ -15,7 +15,7 @@ import type { SessionCoreRepository } from "../../session-core-repository"; import type { SandboxRepository } from "../../sandbox-repository"; import type { SessionSandboxEventProcessor } from "../../sandbox-events/processor"; -function createHandler({ managedSecretsConfigured = true } = {}) { +function createHandler() { const repository = { createEvent: vi.fn(), getProcessingMessage: vi.fn(), @@ -50,7 +50,6 @@ function createHandler({ managedSecretsConfigured = true } = {}) { { getSandbox } as unknown as SandboxRepository, { processSandboxEvent } as unknown as SessionSandboxEventProcessor, messenger, - managedSecretsConfigured, refreshOpenAIToken, refreshXaiToken, getScmCredentials, @@ -540,16 +539,6 @@ describe("SandboxHandler", () => { expect(await response.json()).toEqual({ error: "No session" }); }); - it("returns 500 when openai secrets are not configured", async () => { - const { handler, getSession } = createHandler({ managedSecretsConfigured: false }); - getSession.mockReturnValue({} as SessionRow); - - const response = await handler.openaiTokenRefresh(); - - expect(response.status).toBe(500); - expect(await response.json()).toEqual({ error: "Secrets not configured" }); - }); - it.each([ [OpenAITokenNotConfiguredError, 404, "OPENAI_OAUTH_REFRESH_TOKEN not configured"], [OpenAITokenUnauthorizedError, 401, "OpenAI token refresh failed: unauthorized"], @@ -626,16 +615,6 @@ describe("SandboxHandler", () => { expect(await response.json()).toEqual({ error: "No session" }); }); - it("returns 500 when managed secrets are not configured for xAI", async () => { - const { handler, getSession } = createHandler({ managedSecretsConfigured: false }); - getSession.mockReturnValue({} as SessionRow); - - const response = await handler.xaiTokenRefresh(); - - expect(response.status).toBe(500); - expect(await response.json()).toEqual({ error: "Secrets not configured" }); - }); - it("returns mapped service error from xAI token refresh", async () => { const { handler, getSession, refreshXaiToken } = createHandler(); getSession.mockReturnValue({ id: "session-1" } as SessionRow); diff --git a/packages/control-plane/src/session/http/handlers/sandbox.handler.ts b/packages/control-plane/src/session/http/handlers/sandbox.handler.ts index 3572570b8c..4fb7f0af45 100644 --- a/packages/control-plane/src/session/http/handlers/sandbox.handler.ts +++ b/packages/control-plane/src/session/http/handlers/sandbox.handler.ts @@ -46,8 +46,6 @@ export class SandboxHandler { private readonly sandboxRepository: SandboxRepository, private readonly sandboxEventProcessor: SessionSandboxEventProcessor, private readonly messenger: SessionMessenger, - /** Fixed at composition time: managed secrets exist only when D1 is bound. */ - private readonly managedSecretsConfigured: boolean, private readonly refreshOpenAIToken: (session: SessionRow, log: Logger) => Promise, private readonly refreshXaiToken: ( session: SessionRow, @@ -244,10 +242,6 @@ export class SandboxHandler { return Response.json({ error: "No session" }, { status: 404 }); } - if (!this.managedSecretsConfigured) { - return Response.json({ error: "Secrets not configured" }, { status: 500 }); - } - let token: OpenAIToken; try { token = await this.refreshOpenAIToken(session, log); @@ -282,9 +276,6 @@ export class SandboxHandler { if (!session) { return Response.json({ error: "No session" }, { status: 404 }); } - if (!this.managedSecretsConfigured) { - return Response.json({ error: "Secrets not configured" }, { status: 500 }); - } const result = await this.refreshXaiToken(session, log); if (!result.ok) { return Response.json({ error: result.error }, { status: result.status }); diff --git a/packages/control-plane/src/session/platform.ts b/packages/control-plane/src/session/platform.ts new file mode 100644 index 0000000000..200cbc84a1 --- /dev/null +++ b/packages/control-plane/src/session/platform.ts @@ -0,0 +1,56 @@ +/** + * The platform surface one session runtime is built over: what a host must + * supply for `createSessionRuntime` to assemble the collaborator graph. The + * Cloudflare adapter is `createDurableObjectSessionPlatform` + * (cloudflare/session-platform.ts); a Node host supplies the same record from + * its own storage, sockets, and process facilities. + */ + +import type { Logger } from "../logger"; +import type { BackgroundTasks } from "../platform-ports"; +import type { SqlDatabase } from "../db/sql-database"; +import type { AlarmScheduleStore } from "./alarm/scheduler"; +import type { SqlStorage, TransactionSync } from "./sql-storage"; + +/** The host that owns the session's accepted sockets. */ +export interface SocketHost { + /** Adopt `ws` into the runtime, tagged so its identity survives a restart. */ + accept(ws: WebSocket, tags: string[]): void; + /** The tags `ws` was accepted with. */ + tags(ws: WebSocket): string[]; + /** Every accepted socket, or only those carrying `tag`. */ + sockets(tag?: string): WebSocket[]; + /** + * Answer `request` frames with `response` at the platform level, without + * waking the runtime. + */ + setAutoResponse(request: string, response: string): void; +} + +/** + * The session's own SQLite store. The statements and the transaction + * primitive are bound to the same connection. + */ +export interface SessionStorage { + readonly sql: SqlStorage; + transactionSync: TransactionSync; +} + +export interface SessionPlatform { + /** + * The host's identity for this runtime. It stands in for the session id + * until `init` writes the session row. + */ + id: string; + storage: SessionStorage; + /** The global store. A host that cannot supply one cannot run sessions. */ + db: SqlDatabase; + /** The runtime's single scheduled wake-up. */ + alarmStore: AlarmScheduleStore; + sockets: SocketHost; + /** + * Build the deferred-work port for this runtime. Takes the session-scoped + * logger so failures of background work are attributed to the session. + */ + createBackgroundTasks(log: Logger): BackgroundTasks; +} diff --git a/packages/control-plane/src/session/websocket-manager.test.ts b/packages/control-plane/src/session/websocket-manager.test.ts index 95bca24fd3..5130a4e499 100644 --- a/packages/control-plane/src/session/websocket-manager.test.ts +++ b/packages/control-plane/src/session/websocket-manager.test.ts @@ -1,14 +1,15 @@ /** * Unit tests for SessionWebSocketManagerImpl. * - * Uses fake DurableObjectState and mock repositories to test - * all WebSocket mechanics in isolation from the full DO. + * Uses a fake SocketHost and mock repositories to test + * all WebSocket mechanics in isolation from the host. */ import { describe, it, expect, vi } from "vitest"; import { SessionWebSocketManagerImpl } from "./websocket-manager"; import type { WebSocketManagerConfig } from "./websocket-manager"; import type { Logger } from "../logger"; +import type { SocketHost } from "./platform"; import type { ClientInfo } from "../types"; import type { SandboxRepository } from "./sandbox-repository"; import type { @@ -52,34 +53,34 @@ function createFakeWebSocket(readyState = WebSocket.OPEN): WebSocket { } /** Type for the fake DurableObjectState with test helpers. */ -interface FakeCtx { +interface FakeSocketHost { sockets: Map; - state: DurableObjectState; + host: SocketHost; } /** - * Fake DurableObjectState that tracks accepted WebSockets and their tags. + * Fake SocketHost that tracks accepted WebSockets and their tags. */ -function createFakeCtx(): FakeCtx { +function createFakeSocketHost(): FakeSocketHost { const sockets = new Map(); - const state = { - acceptWebSocket(ws: WebSocket, tags: string[]) { + const host: SocketHost = { + accept(ws, tags) { sockets.set(ws, tags); }, - getTags(ws: WebSocket): string[] { + tags(ws) { return sockets.get(ws) ?? []; }, - getWebSockets(): WebSocket[] { - return Array.from(sockets.keys()); + sockets(tag) { + const accepted = Array.from(sockets.keys()); + return tag === undefined + ? accepted + : accepted.filter((ws) => (sockets.get(ws) ?? []).includes(tag)); }, - setWebSocketAutoResponse: vi.fn(), - storage: { setAlarm: vi.fn() }, - id: { toString: () => "test-do-id" }, - waitUntil: vi.fn(), - } as unknown as DurableObjectState; + setAutoResponse: vi.fn(), + }; - return { sockets, state }; + return { sockets, host }; } /** Create a minimal mock Logger. */ @@ -202,7 +203,7 @@ const TEST_CONFIG: WebSocketManagerConfig = { authTimeoutMs: 100 }; /** Create a fresh manager with all dependencies. */ function createManager() { - const fakeCtx = createFakeCtx(); + const fakeHost = createFakeSocketHost(); const mockRepo = createMockRepository(); const alarmScheduler = { schedule: vi.fn(async () => {}), @@ -212,7 +213,7 @@ function createManager() { const log = createMockLogger(); const manager = new SessionWebSocketManagerImpl( - fakeCtx.state, + fakeHost.host, mockRepo.repo, mockRepo.repo as unknown as WsClientMappingRepository, alarmScheduler, @@ -222,8 +223,8 @@ function createManager() { return { manager, - sockets: fakeCtx.sockets, - state: fakeCtx.state, + sockets: fakeHost.sockets, + host: fakeHost.host, mockRepo, alarmScheduler, log, diff --git a/packages/control-plane/src/session/websocket-manager.ts b/packages/control-plane/src/session/websocket-manager.ts index a4090ffbf4..6ca4a16f3c 100644 --- a/packages/control-plane/src/session/websocket-manager.ts +++ b/packages/control-plane/src/session/websocket-manager.ts @@ -1,14 +1,16 @@ /** - * SessionWebSocketManager — centralizes all Cloudflare WebSocket API usage - * into a single, testable module. + * SessionWebSocketManager — the session's socket registry over its + * `SocketHost`. * * The manager owns socket identity, persistence, and authorization leases. - * The DO builds ClientInfo and stores it here after snapshot synchronization. + * The connection authenticator builds ClientInfo and stores it here after + * snapshot synchronization. */ import type { Logger } from "../logger"; import type { AlarmScheduler } from "../platform-ports"; import type { ClientInfo } from "../types"; +import type { SocketHost } from "./platform"; import type { ConnectionClassification } from "./ports"; import type { SandboxRepository } from "./sandbox-repository"; import type { @@ -106,15 +108,15 @@ export type ClientLookup = // Implementation // --------------------------------------------------------------------------- -/** Durable Object WebSocket manager with persisted authorization leases. */ +/** Session WebSocket manager with persisted authorization leases. */ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { private clients = new Map(); private synchronizingClients = new Set(); private sandboxWs: WebSocket | null = null; - /** Create a WebSocket manager backed by Durable Object state and persisted client mappings. */ + /** Create a WebSocket manager over the host's sockets and persisted client mappings. */ constructor( - private readonly ctx: DurableObjectState, + private readonly host: SocketHost, private readonly sandboxRepository: SandboxRepository, private readonly wsClientMappingRepository: WsClientMappingRepository, private readonly alarmScheduler: AlarmScheduler, @@ -133,12 +135,12 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { } acceptClientSocket(ws: WebSocket, wsId: string): void { - this.ctx.acceptWebSocket(ws, [`wsid:${wsId}`]); + this.host.accept(ws, [`wsid:${wsId}`]); } acceptAndSetSandboxSocket(ws: WebSocket, sandboxId?: string): { replaced: boolean } { const tags = ["sandbox", ...(sandboxId ? [`sid:${sandboxId}`] : [])]; - this.ctx.acceptWebSocket(ws, tags); + this.host.accept(ws, tags); let replaced = false; if (this.sandboxWs && this.sandboxWs !== ws) { @@ -161,7 +163,7 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { // ------------------------------------------------------------------------- classify(ws: WebSocket): ConnectionClassification { - const tags = this.ctx.getTags(ws); + 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) }; @@ -186,7 +188,7 @@ 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()) { + for (const ws of this.host.sockets()) { const parsed = this.classify(ws); if (parsed.kind === "sandbox") { this.close(ws, 1000, "Sandbox terminated"); @@ -201,7 +203,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.host.sockets()) { const parsed = this.classify(ws); if (parsed.kind !== "sandbox" || ws.readyState !== WebSocket.OPEN) continue; @@ -229,7 +231,7 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { detachSandboxSocket(code: number, reason: string): void { const sockets = new Set(); if (this.sandboxWs) sockets.add(this.sandboxWs); - for (const ws of this.ctx.getWebSockets()) { + for (const ws of this.host.sockets()) { if (this.classify(ws).kind === "sandbox") sockets.add(ws); } this.sandboxWs = null; @@ -319,7 +321,7 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { /** Close and remove expired client leases, then schedule the next deadline. */ async expireAuthorizationLeases(now: number): Promise { - for (const ws of this.ctx.getWebSockets()) { + for (const ws of this.host.sockets()) { const parsed = this.classify(ws); if (parsed.kind !== "client") continue; const expiresAt = this.authorizationExpiry(ws, parsed); @@ -386,7 +388,7 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { mode: "all_clients" | "authenticated_only", fn: (ws: WebSocket) => void ): void { - for (const ws of this.ctx.getWebSockets()) { + for (const ws of this.host.sockets()) { const parsed = this.classify(ws); if (parsed.kind === "sandbox") continue; @@ -468,7 +470,7 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { getConnectedClientCount(): number { let count = 0; - for (const ws of this.ctx.getWebSockets()) { + for (const ws of this.host.sockets()) { const parsed = this.classify(ws); if (parsed.kind !== "sandbox" && ws.readyState === WebSocket.OPEN) { count++; diff --git a/packages/control-plane/test/integration/session-components.test.ts b/packages/control-plane/test/integration/session-components.test.ts index bb41af2753..4a3655a5c4 100644 --- a/packages/control-plane/test/integration/session-components.test.ts +++ b/packages/control-plane/test/integration/session-components.test.ts @@ -3,6 +3,7 @@ import { env } from "cloudflare:test"; import type { SessionDO } from "../../src/session/durable-object"; import type { Env } from "../../src/types"; import { createSessionRuntime } from "../../src/session/components"; +import { createDurableObjectSessionPlatform } from "../../src/cloudflare/session-platform"; import { componentsOf, runInSessionDO } from "./session-do-access"; /** @@ -26,14 +27,7 @@ describe("createSessionRuntime", () => { let error: string | null = null; try { - createSessionRuntime( - { - ctx: state, - sql: state.storage.sql, - db: null, - }, - doctored - ); + createSessionRuntime(createDurableObjectSessionPlatform(state, env.DB), doctored); } catch (caught) { error = caught instanceof Error ? caught.message : String(caught); }