diff --git a/docs/plans/early-sandbox-single-websocket.md b/docs/plans/early-sandbox-single-websocket.md new file mode 100644 index 0000000000..7dd0313659 --- /dev/null +++ b/docs/plans/early-sandbox-single-websocket.md @@ -0,0 +1,544 @@ +# Early Sandbox Single-WebSocket Plan + +## Core Decision + +Use one existing sandbox WebSocket for both boot-time control and execution. + +- WebSocket connection proves runtime liveness. +- The existing `ready` event grants execution capability. +- Heartbeats run over the WebSocket during boot and execution. +- No boot-progress HTTP endpoint. +- No database migration. +- No new sandbox status or event type. +- No boot phases or failure-ack protocol. + +```text +Supervisor starts + | + v +Bridge connects to existing sandbox WebSocket + | + +---- heartbeat(status="booting") + | + v +Repository boot and setup + | + v +OpenCode becomes healthy + | + v +Bridge completes signing/session initialization + | + v +Existing ready event + | + v +Control plane grants execution capability + | + v +Queued prompt dispatches +``` + +## 1. Required Invariants + +1. A connected socket does not imply prompt readiness. +2. Every new or replacement socket starts execution-unready. +3. Only `ready` from the current authenticated socket grants execution. +4. Pending prompts remain `pending` until execution readiness. +5. Heartbeats from the current socket renew startup liveness. +6. A socket that stops heartbeating eventually fails. +7. Old runtimes continue working because they already send `ready` immediately after their late + connection. +8. Persisted `connecting`/`ready` status keeps readiness safe across hibernation. +9. Stale or replaced sockets cannot grant readiness. +10. Shutdown remains deliverable while the runtime is booting. + +## 2. Connection State + +Use the existing statuses: + +| Socket state | Persisted status | Meaning | +| ------------------------ | ------------------------------------ | --------------------------------------------- | +| No socket | `spawning` or `connecting` | Waiting for runtime | +| Socket attached, unready | `connecting` | Supervisor alive, startup incomplete | +| Socket attached, ready | `ready` | Prompts and repository commands allowed | +| Socket disconnected | Existing status temporarily retained | Existing disconnect watchdog applies | +| Recoverable failure | `failed` | A current runtime may reconnect and self-heal | +| Terminal sandbox | `stopped` or `stale` | Connection rejected or fenced | + +No `booting` status is required. + +## 3. Readiness State + +Use the existing persisted sandbox status as the only readiness state: + +```text +connecting = control only +ready = execution capable +``` + +The socket keeps its existing `sid:` identity tag. The current sender and persisted sandbox identity +fence `ready`; no serialized attachment, protocol version, or readiness column is required. + +On every successful admission, persist `connecting` before accepting the socket. On an authoritative +`ready`, persist `ready`. Because status survives Durable Object hibernation, the recovered current +socket retains the correct capability without a second representation of readiness. + +## 4. WebSocket Admission + +Change `packages/control-plane/src/session/connection-authenticator.ts`. + +Current behavior on admission: + +- Sets status `ready`. +- Broadcasts readiness. +- Starts inactivity monitoring. +- Processes the queue. + +New behavior: + +1. Authenticate the sandbox as today. +2. Validate sandbox identity and lifecycle state. +3. Reject reconnect with retryable `503` while a snapshot is in progress. +4. Persist sandbox status `connecting` and a server-received heartbeat. +5. Accept the new socket and replace any previous or hibernated sandbox socket. +6. Schedule startup liveness monitoring. +7. Do not start inactivity monitoring. +8. Do not broadcast `ready`. +9. Do not process the prompt queue. + +Old runtimes remain compatible because their bridge sends `ready` immediately after connecting. + +## 5. Authoritative Ready Event + +Change: + +- `packages/control-plane/src/session/message-router.ts` +- `packages/control-plane/src/session/sandbox-events/processor.ts` +- `packages/control-plane/src/session/sandbox-events/context.ts` +- `packages/control-plane/src/session/sandbox-events/runtime.handler.ts` +- `packages/control-plane/src/session/components.ts` + +The message router must retain the actual sending socket when dispatching sandbox events. + +Ready processing: + +1. Verify the sending socket's `sid:` tag matches payload `sandboxId`. +2. Verify the socket is still the active socket. +3. Verify the persisted sandbox identity still matches. +4. Transition `spawning`, `connecting`, or current-attempt `failed` to `ready`. +5. Record runtime version and repository baselines as today. +6. Initialize server-received heartbeat and activity timestamps. +7. Schedule heartbeat and inactivity monitoring. +8. Broadcast ready status and access changes. +9. Process the pending prompt queue last. + +Repeated `ready` from the same current socket must be idempotent. + +A late `ready` from a replaced socket must do nothing. + +## 6. Socket Accessors + +Change `packages/control-plane/src/session/websocket-manager.ts`. + +Keep the existing control accessor and add one execution accessor: + +```ts +getSandboxSocket(): WebSocket | null; +getExecutionSocket(): WebSocket | null; +``` + +`getSandboxSocket()` requires: + +- Current authenticated socket. +- Current sandbox identity. +- Non-terminal sandbox state. + +`getExecutionSocket()` returns that current socket only when persisted sandbox status is `ready`. + +## 7. Command Classification + +| Command | Socket required | +| ------------------------ | --------------------- | +| Heartbeat | Control | +| Shutdown | Control | +| Stop | Control | +| ACK | Exact captured sender | +| Prompt | Execution | +| Push | Execution | +| Diff refresh | Execution | +| Runtime snapshot command | Execution | + +Affected paths: + +- `packages/control-plane/src/session/message-queue.ts` +- `packages/control-plane/src/session/messenger.ts` +- `packages/control-plane/src/session/sandbox-push-service.ts` +- `packages/control-plane/src/session/diffs/service.ts` +- Relevant lifecycle and stop handlers + +Prefer explicit messenger methods: + +```ts +sendControlCommand(...); +sendExecutionCommand(...); +``` + +Avoid a generic readiness framework or command registry. + +## 8. Queue Behavior + +Change `packages/control-plane/src/session/message-queue.ts`. + +The queue needs three cases: + +```text +No control socket + -> retain existing spawn/resume behavior + +Control socket but no execution socket + -> leave message pending + -> do not claim it + -> do not start execution timeout + -> do not spawn another sandbox + +Execution socket available + -> claim and send using existing behavior +``` + +This is the central behavioral change. + +Warm-on-typing should treat an attached control socket as an existing sandbox so it does not create +a duplicate. + +## 9. Runtime Startup + +Change: + +- `packages/sandbox-runtime/src/sandbox_runtime/supervisor.py` +- `packages/sandbox-runtime/src/sandbox_runtime/bridge.py` +- Possibly `packages/sandbox-runtime/src/sandbox_runtime/agent_bridge_process.py` +- Possibly `packages/sandbox-runtime/src/sandbox_runtime/opencode_client.py` + +New supervisor order: + +1. Exclude image-build mode as today. +2. Start the bridge process immediately. +3. Start desktop best-effort. +4. Run repository boot. +5. Materialize managed skills. +6. Start code-server and terminal best-effort. +7. Start OpenCode and wait for health. +8. Continue normal process monitoring. + +Remove the HTTP boot-progress task entirely. + +## 10. Deferred Bridge Initialization + +The bridge currently performs OpenCode and repository-dependent work before connecting: + +- Restored-session validation +- Commit-signing initialization +- Repository manifest reads +- Ready-event construction + +Split bridge startup into two concurrent responsibilities. + +### Transport Initialization + +Runs immediately: + +- Connect WebSocket. +- Start heartbeat. +- Handle reconnects. +- Receive shutdown and stop. +- Buffer outbound events. +- Reject execution commands defensively before ready. + +### Execution Initialization + +Runs locally while the transport is connected: + +1. Poll OpenCode `/global/health`. +2. Wait until health succeeds. +3. Load or validate the OpenCode session. +4. Read the final repository manifest. +5. Initialize commit signing. +6. Build the ready payload. +7. Mark local execution initialization complete. +8. Send `ready` on the current connection. + +OpenCode health is sufficient as the readiness boundary because the supervisor starts OpenCode only +after repository boot and manifest creation complete. + +This avoids supervisor-to-bridge IPC. + +## 11. Runtime Reconnect + +Maintain one process-level flag: + +```py +execution_initialized: bool +``` + +Before initialization: + +- Connect. +- Heartbeat with `status="booting"`. +- Do not send `ready`. + +After initialization: + +- Send `ready`. +- Heartbeat with `status="ready"`. + +On WebSocket reconnect after initialization: + +- The new control-plane socket starts unready. +- Bridge immediately re-sends `ready`. +- Control plane marks the replacement socket execution-ready. + +The existing heartbeat schema already allows arbitrary status strings, so no shared contract change +is required. + +## 12. Runtime Command Defense + +The control plane is the primary gate, but the bridge should also reject execution commands before +readiness. + +Before ready, permit: + +- `shutdown` +- `stop` +- `ack` + +Before ready, reject or return unavailable for: + +- `prompt` +- `push` +- `refresh_diff` +- `snapshot` + +This protects against rollout mistakes and stale control-plane versions. + +## 13. Heartbeat And Timeout + +Retain the renewable lease model. + +Before socket connection: + +```text +created_at + existing startup timeout +``` + +After attachment: + +```text +latest server-received heartbeat + heartbeat timeout +``` + +Recommended timing: + +- Heartbeat every 30 seconds. +- Failure after 90-120 seconds without heartbeat. +- No total boot-duration limit. + +The same persisted `last_heartbeat` field works for booting and ready states. + +Timeout transition must remain fenced by: + +- Sandbox ID. +- Spawn attempt timestamp. +- Startup status. +- Observed heartbeat timestamp. + +No new timestamp column is required. + +## 14. Failure Handling + +Do not introduce `boot_failed` events or acknowledgements. + +On startup failure: + +1. Supervisor uses the existing fatal-error reporting path. +2. Supervisor shuts down the bridge. +3. Heartbeats stop. +4. Existing heartbeat timeout is the fallback if fatal reporting fails. +5. Control plane marks the attempt failed and re-drives pending work according to existing policy. + +This preserves current behavior without another delivery protocol. + +## 15. Hibernation + +Hibernation recovery must: + +1. Recover the tagged sandbox socket. +2. Verify the `sid:` tag matches the current sandbox row. +3. Use persisted `connecting` or `ready` status as the capability gate. +4. Recover connecting sockets as control-only. +5. Close identity-mismatched sockets. + +No additional database or WebSocket attachment persistence is needed. + +## 16. Snapshot Interaction + +If a runtime reconnects while persisted status is `snapshotting`, reject the upgrade with retryable +`503` and `Retry-After`. The bridge reconnect loop retries after snapshot completion. Its next +accepted socket is persisted as `connecting` and must reannounce `ready` before execution resumes. + +## 17. Compatibility And Rollout + +A same-route rollout is possible without protocol V2. + +| Control plane | Runtime | Result | +| ------------------- | ------------------------- | ------------------------------------------------- | +| New readiness-aware | Old late-connect runtime | Safe; old runtime sends `ready` immediately | +| New readiness-aware | New early-connect runtime | Safe; waits for explicit `ready` | +| Old control plane | Old runtime | Existing behavior | +| Old control plane | New runtime in late mode | Existing behavior | +| Old control plane | New runtime in early mode | Unsafe; old control plane dispatches on admission | + +Use one temporary activation environment variable: + +```text +EARLY_SANDBOX_CONNECTION=1 +``` + +New runtime defaults to late mode when absent. + +Deployment sequence: + +1. Deploy readiness-aware control plane. +2. Verify old runtimes still connect and dispatch. +3. Deploy early-capable runtime with activation off. +4. Enable early mode for a canary environment/provider. +5. Expand activation after metrics are stable. +6. Remove the HTTP boot-progress endpoint and polling loop. +7. Remove the temporary flag after rollback no longer requires late mode. + +No protocol-version schema or dedicated route is required under this controlled deployment order. + +Rollback: + +1. Disable early-mode activation. +2. Let or force active early sessions drain. +3. Roll runtime artifacts back if needed. +4. Roll the control plane back only after no early sessions remain. + +## 18. PR Structure + +To keep reviews manageable, split the work. + +### PR 1: Explicit Readiness In Control Plane + +Scope: + +- Persisted connecting/ready capability +- Admission no longer grants readiness +- Current-sender ready handling +- Queue gating +- Command classification +- Reconnect and hibernation tests +- Legacy late runtime compatibility + +Expected scope: approximately 15-20 files. + +This PR is independently deployable because existing runtimes already emit `ready`. + +### PR 2: Early Runtime Connection + +Scope: + +- Start bridge before repository boot +- Deferred bridge execution initialization +- Booting heartbeat +- Runtime command defense +- Temporary activation variable +- Remove HTTP boot-progress loop and endpoint +- Runtime and cross-version tests + +Expected scope: approximately 8-12 files. + +## 19. Test Plan + +### Control-Plane Unit Tests + +- Admission persists connecting before accepting the socket. +- Admission does not mark status ready. +- Admission does not process the queue. +- Ready from current socket grants execution. +- Ready from replaced socket is ignored. +- Identity mismatch closes or rejects the sender. +- Duplicate ready is idempotent. +- Attached-unready message stays pending. +- Attached-unready state does not spawn a duplicate. +- Stop and shutdown use control socket. +- Prompt, push, and diff use execution socket. +- ACK returns to the captured sender. +- Hibernation preserves connecting and ready capability through persisted status. +- Reconnect is retryably rejected during snapshots. + +### Control-Plane Integration Tests + +- Pending prompt survives early attachment. +- Prompt dispatches only after ready. +- Reconnect before ready remains unready. +- Reconnect after ready requires reannouncement. +- Replaced socket cannot grant readiness. +- Heartbeating boot exceeds the former startup duration. +- Missing heartbeat fails the attempt. +- Durable Object eviction preserves readiness correctly. +- Legacy runtime admission plus immediate ready still works. +- Snapshot reconnect is rejected until snapshot completion, then requires ready. + +### Runtime Tests + +- Bridge connects before repository boot completes. +- Heartbeats are sent while OpenCode is unavailable. +- No ready event before OpenCode health. +- No signing/session initialization before health. +- Ready follows health and initialization. +- Reconnect after initialization re-sends ready. +- Shutdown works while booting. +- Prompt and push are rejected before ready. +- Repository or OpenCode startup failure stops bridge. +- Image-build mode never starts bridge. +- Activation off retains existing late-start order. + +### Cross-Version Tests + +- New control plane with old runtime. +- New control plane with new runtime, activation off. +- New control plane with new runtime, activation on. +- Old control plane with new runtime, activation off. +- Activation cannot be enabled against an old control plane. + +## 20. Acceptance Criteria + +1. Runtime WebSocket attaches before repository boot. +2. Attached socket immediately provides heartbeat liveness. +3. No prompt is claimed before explicit ready. +4. Existing ready event is the only execution grant. +5. Every replacement socket starts unready. +6. Ready is accepted only from the active authenticated socket. +7. A heartbeating boot can run indefinitely. +8. Missing heartbeats eventually fail the current attempt. +9. Shutdown remains available while booting. +10. Legacy runtimes continue working. +11. Durable Object hibernation preserves connection readiness safely. +12. No database migration is introduced. +13. No protocol V2, boot phases, or failure-ack system is introduced. +14. HTTP boot-progress polling is removed. +15. Image-build behavior is unchanged. +16. The implementation remains split into reviewable PRs. + +## Estimated Scope + +- Production files: approximately 16-20 +- Test files: approximately 10-14 +- Total files across two PRs: approximately 24-30 +- No shared schema, web UI, provider-specific behavior, or database migration expected + +The unavoidable complexity is limited to connection-scoped readiness, sender fencing, queue gating, +and deferred bridge initialization. Everything else from the earlier 99-file approach remains +excluded. diff --git a/packages/control-plane/src/sandbox/lifecycle/decisions.test.ts b/packages/control-plane/src/sandbox/lifecycle/decisions.test.ts index 911fe57e13..d5e20d6be8 100644 --- a/packages/control-plane/src/sandbox/lifecycle/decisions.test.ts +++ b/packages/control-plane/src/sandbox/lifecycle/decisions.test.ts @@ -288,7 +288,7 @@ describe("evaluateSpawnDecision", () => { expect(decision.action).toBe("skip"); }); - it('returns "spawn" when stuck in "spawning" past the spawning timeout (recovers interrupted spawn)', () => { + it('returns "spawn" when stuck in "spawning" past the spawning timeout', () => { const now = Date.now(); const state: SandboxState = { status: "spawning", @@ -847,7 +847,7 @@ describe("evaluateConnectingTimeout", () => { it("returns not timed out for non-connecting status", () => { const now = Date.now(); - const result = evaluateConnectingTimeout("ready", now - 200_000, config, now); + const result = evaluateConnectingTimeout("ready", now - 200_000, null, config, now); expect(result.isTimedOut).toBe(false); expect(result.elapsedMs).toBe(0); @@ -857,7 +857,7 @@ describe("evaluateConnectingTimeout", () => { const now = Date.now(); const createdAt = now - 60_000; // 60s ago, well within 120s timeout - const result = evaluateConnectingTimeout("connecting", createdAt, config, now); + const result = evaluateConnectingTimeout("connecting", createdAt, null, config, now); expect(result.isTimedOut).toBe(false); expect(result.elapsedMs).toBe(60_000); @@ -867,7 +867,7 @@ describe("evaluateConnectingTimeout", () => { const now = Date.now(); const createdAt = now - 130_000; // 130s ago, past 120s timeout - const result = evaluateConnectingTimeout("connecting", createdAt, config, now); + const result = evaluateConnectingTimeout("connecting", createdAt, null, config, now); expect(result.isTimedOut).toBe(true); expect(result.elapsedMs).toBe(130_000); @@ -877,27 +877,48 @@ describe("evaluateConnectingTimeout", () => { const now = Date.now(); const createdAt = now - config.timeoutMs; // Exactly at timeout - const result = evaluateConnectingTimeout("connecting", createdAt, config, now); + const result = evaluateConnectingTimeout("connecting", createdAt, null, config, now); expect(result.isTimedOut).toBe(true); expect(result.elapsedMs).toBe(config.timeoutMs); }); - it("returns timed out when stuck in spawning past timeout (interrupted spawn)", () => { + it("uses the existing timeout when no heartbeat has arrived", () => { const now = Date.now(); - const createdAt = now - 130_000; // 130s ago, past 120s timeout + const createdAt = now - config.timeoutMs; - const result = evaluateConnectingTimeout("spawning", createdAt, config, now); + const result = evaluateConnectingTimeout("spawning", createdAt, null, config, now); expect(result.isTimedOut).toBe(true); - expect(result.elapsedMs).toBe(130_000); + expect(result.livenessAt).toBe(createdAt); }); - it("returns not timed out for spawning within timeout window", () => { + it("extends boot indefinitely while authenticated heartbeats remain recent", () => { const now = Date.now(); - const result = evaluateConnectingTimeout("spawning", now - 60_000, config, now); + const createdAt = now - 60 * 60_000; + const lastHeartbeat = now - 30_000; + + const result = evaluateConnectingTimeout("connecting", createdAt, lastHeartbeat, config, now); expect(result.isTimedOut).toBe(false); + expect(result.livenessAt).toBe(lastHeartbeat); + expect(result.deadlineAt).toBe(lastHeartbeat + config.timeoutMs); + }); + + it("times out when boot heartbeats become stale", () => { + const now = Date.now(); + const lastHeartbeat = now - config.timeoutMs; + + const result = evaluateConnectingTimeout( + "connecting", + now - 600_000, + lastHeartbeat, + config, + now + ); + + expect(result.isTimedOut).toBe(true); + expect(result.livenessAt).toBe(lastHeartbeat); }); it("ignores all non-spawning/connecting statuses", () => { @@ -905,7 +926,7 @@ describe("evaluateConnectingTimeout", () => { const old = now - 999_999; for (const status of ["pending", "ready", "stopped", "failed", "stale"] as const) { - const result = evaluateConnectingTimeout(status, old, config, now); + const result = evaluateConnectingTimeout(status, old, null, config, now); expect(result.isTimedOut).toBe(false); } }); diff --git a/packages/control-plane/src/sandbox/lifecycle/decisions.ts b/packages/control-plane/src/sandbox/lifecycle/decisions.ts index 2bdd772066..3663bc12b4 100644 --- a/packages/control-plane/src/sandbox/lifecycle/decisions.ts +++ b/packages/control-plane/src/sandbox/lifecycle/decisions.ts @@ -146,6 +146,8 @@ export interface SandboxState { status: SandboxStatus; /** When the sandbox was created/spawned */ createdAt: number; + /** Last server-received runtime liveness report during startup. */ + lastHeartbeat?: number | null; /** Provider object ID if the sandbox exists remotely */ providerObjectId?: string | null; /** Snapshot image ID if available for restore */ @@ -168,16 +170,7 @@ export interface SpawnConfig { cooldownMs: number; /** Time to wait for WebSocket after spawn (default: 60s) */ readyWaitMs: number; - /** - * Max time a sandbox may remain in "spawning"/"connecting" before it is - * treated as dead and a fresh spawn is allowed (default: 120s). - * - * Guards against spawns interrupted before the sandbox connects (provider - * crash, redeploy, cancelled provider call). Such a spawn can leave the - * persisted status pinned at "spawning"/"connecting" indefinitely — the - * connecting-timeout alarm may never have been scheduled — which otherwise - * makes every later spawn attempt skip with "already spawning" forever. - */ + /** Max time without startup liveness before a fresh spawn is allowed. */ spawningTimeoutMs: number; } @@ -312,11 +305,17 @@ export function evaluateSpawnDecision( // "connecting" forever — the connecting-timeout alarm may never have been // scheduled. Treat a stale spawn/connect as dead so a fresh spawn can recover // the session, instead of skipping indefinitely. - if ( - (state.status === "spawning" || state.status === "connecting") && - timeSinceLastSpawn < config.spawningTimeoutMs - ) { - return { action: "skip", reason: `already ${state.status}` }; + if (state.status === "spawning" || state.status === "connecting") { + const startup = evaluateConnectingTimeout( + state.status, + state.createdAt, + state.lastHeartbeat ?? null, + { timeoutMs: config.spawningTimeoutMs }, + now + ); + if (!startup.isTimedOut) { + return { action: "skip", reason: `already ${state.status}` }; + } } // Don't spawn if status is "ready" and we have an active WebSocket @@ -542,14 +541,12 @@ export function evaluateHeartbeatHealth( * Configuration for the initial-connect watchdog. */ export interface ConnectingTimeoutConfig { - /** Maximum time in ms a sandbox can stay in "connecting" before being failed */ + /** Maximum gap without startup liveness. */ timeoutMs: number; } /** - * Default connecting timeout: 2 minutes. - * Boot sequence (git clone → setup.sh → start.sh → opencode → bridge connect) typically - * takes 30–90 seconds. Two minutes provides margin without leaving users waiting too long. + * Default startup liveness timeout: 2 minutes. */ export const DEFAULT_CONNECTING_TIMEOUT_CONFIG: ConnectingTimeoutConfig = { timeoutMs: 120_000, @@ -561,8 +558,12 @@ export const DEFAULT_CONNECTING_TIMEOUT_CONFIG: ConnectingTimeoutConfig = { export interface ConnectingTimeoutResult { /** Whether the sandbox has exceeded the connecting timeout */ isTimedOut: boolean; - /** Time elapsed since sandbox was created (ms) */ + /** Time elapsed since the latest startup liveness signal. */ elapsedMs: number; + /** Timestamp from which the active timeout window is measured. */ + livenessAt: number; + /** Next timestamp at which startup should be evaluated. */ + deadlineAt: number; } /** @@ -582,24 +583,29 @@ export interface ConnectingTimeoutResult { * * @param status - Current sandbox status * @param createdAt - Timestamp (ms) when the sandbox was spawned + * @param lastHeartbeat - Latest authenticated startup heartbeat receipt time * @param config - Connecting timeout configuration * @param now - Current timestamp (ms) - * @returns Whether the sandbox has timed out and how long it's been spawning/connecting + * @returns Current timeout state and deadline */ export function evaluateConnectingTimeout( status: SandboxStatus, createdAt: number, + lastHeartbeat: number | null, config: ConnectingTimeoutConfig, now: number ): ConnectingTimeoutResult { if (status !== "connecting" && status !== "spawning") { - return { isTimedOut: false, elapsedMs: 0 }; + return { isTimedOut: false, elapsedMs: 0, livenessAt: createdAt, deadlineAt: createdAt }; } - const elapsedMs = now - createdAt; + const livenessAt = Math.max(createdAt, lastHeartbeat ?? 0); + const elapsedMs = now - livenessAt; return { isTimedOut: elapsedMs >= config.timeoutMs, elapsedMs, + livenessAt, + deadlineAt: livenessAt + config.timeoutMs, }; } diff --git a/packages/control-plane/src/sandbox/lifecycle/manager.test.ts b/packages/control-plane/src/sandbox/lifecycle/manager.test.ts index 4ceebdcc5b..4aa3a2a636 100644 --- a/packages/control-plane/src/sandbox/lifecycle/manager.test.ts +++ b/packages/control-plane/src/sandbox/lifecycle/manager.test.ts @@ -179,6 +179,7 @@ function createMockStorage( sandbox.auth_token = null; sandbox.modal_sandbox_id = data.modalSandboxId; sandbox.runtime_version = null; + sandbox.last_heartbeat = null; if (!data.preserveProviderObjectId) sandbox.modal_object_id = null; } }), @@ -193,6 +194,7 @@ function createMockStorage( if (sandbox) { sandbox.status = data.status; sandbox.created_at = data.createdAt; + sandbox.last_heartbeat = null; } }), updateSandboxModalObjectId: vi.fn((id: string | null) => { @@ -216,6 +218,35 @@ function createMockStorage( calls.push("updateSandboxLastActivity"); if (sandbox) sandbox.last_activity = timestamp; }), + recordStartupHeartbeat: vi.fn((sandboxId: string, timestamp: number) => { + calls.push(`recordStartupHeartbeat:${sandboxId}`); + if ( + !sandbox || + sandbox.modal_sandbox_id !== sandboxId || + sandbox.status === "snapshotting" || + sandbox.status === "stopped" || + sandbox.status === "stale" + ) { + return false; + } + sandbox.status = "connecting"; + sandbox.last_heartbeat = timestamp; + return true; + }), + failStartupIfUnchanged: vi.fn((sandboxId: string, createdAt: number, livenessAt: number) => { + calls.push(`failStartupIfUnchanged:${sandboxId}`); + if ( + !sandbox || + sandbox.modal_sandbox_id !== sandboxId || + sandbox.created_at !== createdAt || + Math.max(sandbox.created_at, sandbox.last_heartbeat ?? 0) !== livenessAt || + (sandbox.status !== "spawning" && sandbox.status !== "connecting") + ) { + return false; + } + sandbox.status = "failed"; + return true; + }), incrementCircuitBreakerFailure: vi.fn((timestamp: number) => { calls.push("incrementCircuitBreakerFailure"); if (sandbox) { @@ -927,6 +958,36 @@ describe("SandboxLifecycleManager", () => { expect(provider.createSandbox).toHaveBeenCalledWith(expect.objectContaining({ userEnvVars })); }); + it.each([ + [false, {}], + [true, { EARLY_SANDBOX_CONNECTION: "1" }], + ])( + "sets the early connection runtime variable only when activated", + async (enabled, expected) => { + const sandbox = createMockSandbox({ status: "pending", created_at: Date.now() - 60000 }); + const storage = createMockStorage(createMockSession(), sandbox, { + EARLY_SANDBOX_CONNECTION: "user-value", + }); + const provider = createMockProvider(); + const manager = new SandboxLifecycleManager( + provider, + storage, + storage, + createMockBroadcaster(), + createMockWebSocketManager(false), + createMockAlarmScheduler(), + createMockIdGenerator(), + { ...createTestConfig(), earlySandboxConnection: enabled } + ); + + await manager.spawnSandbox(); + + expect(provider.createSandbox).toHaveBeenCalledWith( + expect.objectContaining({ userEnvVars: expected }) + ); + } + ); + it("spawns no-repository sessions without repo-only sandbox features", async () => { const sandbox = createMockSandbox({ status: "pending", created_at: Date.now() - 60000 }); const storage = createMockStorage( @@ -1666,10 +1727,8 @@ describe("SandboxLifecycleManager", () => { await manager.spawnSandbox(); - // Should go: pending -> spawning -> connecting - const statusCalls = storage.calls.filter((c) => c.startsWith("updateSandbox")); - expect(statusCalls).toContain("updateSandboxForSpawn"); - expect(statusCalls).toContain("updateSandboxStatus:connecting"); + expect(storage.calls).toContain("updateSandboxForSpawn"); + expect(storage.calls).toContain("updateSandboxStatus:connecting"); }); it("handles provider errors and increments failure count for permanent errors", async () => { @@ -1786,7 +1845,7 @@ describe("SandboxLifecycleManager", () => { const storage = createMockStorage(createMockSession(), sandbox); const broadcaster = createMockBroadcaster(); const provider = createMockProvider(); - + const processMessageQueue = vi.fn(async () => {}); const manager = new SandboxLifecycleManager( provider, storage, @@ -1798,7 +1857,7 @@ describe("SandboxLifecycleManager", () => { createTestConfig() ); - await manager.triggerSnapshot("test_reason"); + await manager.triggerSnapshot("test_reason", processMessageQueue); expect(provider.takeSnapshot).toHaveBeenCalled(); expect(storage.calls).toContain( @@ -1811,6 +1870,27 @@ describe("SandboxLifecycleManager", () => { { type: "sandbox_status", status: "ready" }, { type: "sandbox_access_changed" }, ]); + expect(processMessageQueue).not.toHaveBeenCalled(); + }); + + it("drains the queue after a ready sandbox snapshots with its control socket connected", async () => { + const sandbox = createMockSandbox({ status: "ready" }); + const storage = createMockStorage(createMockSession(), sandbox); + const processMessageQueue = vi.fn(async () => {}); + const manager = new SandboxLifecycleManager( + createMockProvider(), + storage, + storage, + createMockBroadcaster(), + createMockWebSocketManager(true), + createMockAlarmScheduler(), + createMockIdGenerator(), + createTestConfig() + ); + + await manager.triggerSnapshot("test_reason", processMessageQueue); + + expect(processMessageQueue).toHaveBeenCalledOnce(); }); it("skips when provider does not support snapshots", async () => { @@ -1904,6 +1984,32 @@ describe("SandboxLifecycleManager", () => { }); describe("handleAlarm", () => { + it("re-arms a healthy sandbox at its renewable heartbeat deadline", async () => { + const now = Date.now(); + const sandbox = createMockSandbox({ + status: "ready", + last_heartbeat: now - 10000, + last_activity: now, + }); + const storage = createMockStorage(createMockSession(), sandbox); + const alarmScheduler = createMockAlarmScheduler(); + const config = createTestConfig(); + const manager = new SandboxLifecycleManager( + createMockProvider(), + storage, + storage, + createMockBroadcaster(), + createMockWebSocketManager(), + alarmScheduler, + createMockIdGenerator(), + config + ); + + await manager.handleAlarm(); + + expect(alarmScheduler.alarms[0]).toBe(sandbox.last_heartbeat! + config.heartbeat.timeoutMs); + }); + it("detects heartbeat timeout and sets stale", async () => { const now = Date.now(); const sandbox = createMockSandbox({ @@ -1995,7 +2101,7 @@ describe("SandboxLifecycleManager", () => { // Should extend, not timeout expect(storage.calls).not.toContain("updateSandboxStatus:stopped"); - expect(alarmScheduler.alarms.length).toBe(1); + expect(alarmScheduler.alarms.length).toBe(2); expect( broadcaster.messages.some((m) => (m as { type: string }).type === "sandbox_warning") ).toBe(true); @@ -2028,7 +2134,7 @@ describe("SandboxLifecycleManager", () => { await manager.handleAlarm(); expect(storage.calls).not.toContain("updateSandboxStatus:stopped"); - expect(alarmScheduler.alarms.length).toBe(1); + expect(alarmScheduler.alarms.length).toBe(2); }); it("triggers snapshot before stopping", async () => { @@ -2246,7 +2352,7 @@ describe("SandboxLifecycleManager", () => { const result = await manager.handleAlarm(); expect(result).toBe("sandbox_failed"); - expect(storage.calls).toContain("updateSandboxStatus:failed"); + expect(storage.calls).toContain(`failStartupIfUnchanged:${sandbox.modal_sandbox_id}`); expect(storage.calls).toContain("clearSandboxAccess:codeServer"); expect(broadcaster.messages.some((m) => (m as { status?: string }).status === "failed")).toBe( true @@ -2261,6 +2367,58 @@ describe("SandboxLifecycleManager", () => { expect(provider.takeSnapshot).not.toHaveBeenCalled(); }); + it("keeps an arbitrarily old boot alive while progress is recent", async () => { + const now = Date.now(); + const sandbox = createMockSandbox({ + status: "connecting", + created_at: now - 60 * 60_000, + last_heartbeat: now - 10_000, + }); + const storage = createMockStorage(createMockSession(), sandbox); + const alarmScheduler = createMockAlarmScheduler(); + const manager = new SandboxLifecycleManager( + createMockProvider(), + storage, + storage, + createMockBroadcaster(), + createMockWebSocketManager(), + alarmScheduler, + createMockIdGenerator(), + createTestConfig() + ); + + expect(await manager.handleAlarm()).toBe("no_action"); + expect(storage.calls).not.toContain("updateSandboxStatus:failed"); + expect(alarmScheduler.alarms).toEqual([ + sandbox.last_heartbeat! + createTestConfig().connectingTimeout.timeoutMs, + ]); + }); + + it("does not fail a replacement when timeout state changed", async () => { + const now = Date.now(); + const sandbox = createMockSandbox({ + status: "connecting", + created_at: now - 130_000, + last_heartbeat: null, + }); + const storage = createMockStorage(createMockSession(), sandbox); + vi.mocked(storage.failStartupIfUnchanged).mockReturnValue(false); + const broadcaster = createMockBroadcaster(); + const manager = new SandboxLifecycleManager( + createMockProvider(), + storage, + storage, + broadcaster, + createMockWebSocketManager(), + createMockAlarmScheduler(), + createMockIdGenerator(), + createTestConfig() + ); + + expect(await manager.handleAlarm()).toBe("no_action"); + expect(broadcaster.messages).toEqual([]); + }); + it("does not timeout connecting sandbox within timeout window", async () => { const now = Date.now(); const sandbox = createMockSandbox({ @@ -2290,6 +2448,53 @@ describe("SandboxLifecycleManager", () => { }); }); + describe("recordStartupHeartbeat", () => { + it("records current startup liveness and rearms the watchdog", async () => { + const now = Date.now(); + const sandbox = createMockSandbox({ status: "connecting", last_heartbeat: null }); + const storage = createMockStorage(createMockSession(), sandbox); + const alarmScheduler = createMockAlarmScheduler(); + const config = createTestConfig(); + const manager = new SandboxLifecycleManager( + createMockProvider(), + storage, + storage, + createMockBroadcaster(), + createMockWebSocketManager(), + alarmScheduler, + createMockIdGenerator(), + config + ); + + await expect(manager.recordStartupHeartbeat(sandbox.modal_sandbox_id!, now)).resolves.toBe( + true + ); + expect(sandbox.last_heartbeat).toBe(now); + expect(alarmScheduler.alarms).toEqual([now + config.connectingTimeout.timeoutMs]); + }); + + it("rejects progress from a replaced sandbox", async () => { + const sandbox = createMockSandbox({ status: "connecting" }); + const storage = createMockStorage(createMockSession(), sandbox); + const alarmScheduler = createMockAlarmScheduler(); + const manager = new SandboxLifecycleManager( + createMockProvider(), + storage, + storage, + createMockBroadcaster(), + createMockWebSocketManager(), + alarmScheduler, + createMockIdGenerator(), + createTestConfig() + ); + + await expect(manager.recordStartupHeartbeat("sandbox-replaced", Date.now())).resolves.toBe( + false + ); + expect(alarmScheduler.alarms).toEqual([]); + }); + }); + describe("terminateUnresponsiveSandbox", () => { it.each([ ["prompt_dispatch_send_failed", "Prompt dispatch send failed"], diff --git a/packages/control-plane/src/sandbox/lifecycle/manager.ts b/packages/control-plane/src/sandbox/lifecycle/manager.ts index f697b27115..c6683c1abe 100644 --- a/packages/control-plane/src/sandbox/lifecycle/manager.ts +++ b/packages/control-plane/src/sandbox/lifecycle/manager.ts @@ -76,6 +76,7 @@ const PROVIDER_REPLACEMENT_STOP_TIMEOUT_MS = 10_000; interface SandboxCircuitBreakerInfo { status: SandboxStatus; created_at: number; + last_heartbeat: number | null; modal_object_id: string | null; snapshot_image_id: string | null; snapshot_runtime_version: string | null; @@ -152,6 +153,10 @@ export interface SandboxStorage { ): void; /** Update last activity timestamp */ updateSandboxLastActivity(timestamp: number): void; + /** Record authenticated startup liveness for the current logical sandbox. */ + recordStartupHeartbeat(sandboxId: string, timestamp: number): boolean; + /** Fail startup only if identity, attempt, status, and liveness are unchanged. */ + failStartupIfUnchanged(sandboxId: string, createdAt: number, livenessAt: number): boolean; /** Increment circuit breaker failure count */ incrementCircuitBreakerFailure(timestamp: number): void; /** Reset circuit breaker failure count */ @@ -213,6 +218,7 @@ export interface SandboxLifecycleConfig { heartbeat: HeartbeatConfig; connectingTimeout: ConnectingTimeoutConfig; controlPlaneUrl: string; + earlySandboxConnection?: boolean; /** Default model ID used when the session has no model override. */ model: string; /** @@ -241,6 +247,17 @@ export const DEFAULT_LIFECYCLE_CONFIG: Omit | undefined, + earlySandboxConnection: boolean | undefined +): Record | undefined { + if (!userEnvVars && !earlySandboxConnection) return undefined; + const resolved = { ...userEnvVars }; + delete resolved.EARLY_SANDBOX_CONNECTION; + if (earlySandboxConnection) resolved.EARLY_SANDBOX_CONNECTION = "1"; + return resolved; +} + function buildSandboxIdForSession(session: SessionRow, now: number): string { const sandboxName = sessionHasRepository(session) ? `${session.repo_owner}-${session.repo_name}` @@ -298,6 +315,7 @@ export interface SlackAgentNotifyLookup { */ export interface SandboxLifecycle { spawnSandbox(): Promise; + isSnapshotting(): boolean; updateLastActivity(timestamp: number): void; terminateUnresponsiveSandbox(trigger: UnresponsiveSandboxTrigger): Promise; reportSandboxError(reason: string): void; @@ -413,6 +431,7 @@ export class SandboxLifecycleManager implements SandboxLifecycle { const spawnState = { status: sandboxState?.status ?? DEFAULT_SANDBOX_STATUS, createdAt: sandboxState?.created_at || 0, + lastHeartbeat: sandboxState?.last_heartbeat ?? null, providerObjectId: sandboxState?.modal_object_id || null, snapshotImageId: sandboxState?.snapshot_image_id || null, snapshotRuntimeVersion: sandboxState?.snapshot_runtime_version || null, @@ -577,7 +596,7 @@ export class SandboxLifecycleManager implements SandboxLifecycle { sandboxAuthToken, provider, model: modelId, - userEnvVars, + userEnvVars: runtimeUserEnvVars(userEnvVars, this.config.earlySandboxConnection), prebuiltImageId, prebuiltImageSha, timeoutSeconds, @@ -900,7 +919,7 @@ export class SandboxLifecycleManager implements SandboxLifecycle { repoName: session.repo_name, provider, model: modelId, - userEnvVars, + userEnvVars: runtimeUserEnvVars(userEnvVars, this.config.earlySandboxConnection), timeoutSeconds, branch: session.base_branch, codeServerEnabled, @@ -1073,7 +1092,7 @@ export class SandboxLifecycleManager implements SandboxLifecycle { /** * Trigger a filesystem snapshot of the sandbox. */ - async triggerSnapshot(reason: string): Promise { + async triggerSnapshot(reason: string, processMessageQueue?: () => Promise): Promise { if (!this.provider.takeSnapshot) { this.log.debug("Provider does not support snapshots"); return; @@ -1151,6 +1170,9 @@ export class SandboxLifecycleManager implements SandboxLifecycle { this.broadcaster.broadcast({ type: "sandbox_status", status: previousStatus }); if (previousStatus === "ready") { this.broadcaster.broadcast({ type: "sandbox_access_changed" }); + if (processMessageQueue && this.wsManager.getSandboxWebSocket()) { + await processMessageQueue(); + } } } } @@ -1288,6 +1310,7 @@ export class SandboxLifecycleManager implements SandboxLifecycle { const connectingResult = evaluateConnectingTimeout( sandbox.status, sandbox.created_at, + sandbox.last_heartbeat, this.config.connectingTimeout, now ); @@ -1296,9 +1319,20 @@ export class SandboxLifecycleManager implements SandboxLifecycle { this.log.warn("Connecting timeout", { event: "sandbox.connecting_timeout", elapsed_ms: connectingResult.elapsedMs, - timeout_ms: this.config.connectingTimeout.timeoutMs, + timeout_ms: connectingResult.deadlineAt - connectingResult.livenessAt, }); - this.storage.updateSandboxStatus("failed"); + const sandboxId = sandbox.modal_sandbox_id; + const failed = + sandboxId && + this.storage.failStartupIfUnchanged( + sandboxId, + sandbox.created_at, + connectingResult.livenessAt + ); + if (!failed) { + this.log.info("Connecting timeout superseded by newer startup liveness"); + return "no_action"; + } this.clearSandboxAccessState(); if (this.canStopProviderSandbox()) { try { @@ -1316,6 +1350,11 @@ export class SandboxLifecycleManager implements SandboxLifecycle { return "sandbox_failed"; } + if (sandbox.status === "spawning" || sandbox.status === "connecting") { + await this.alarmScheduler.schedule(connectingResult.deadlineAt); + return "no_action"; + } + // Check heartbeat health const heartbeatHealth = evaluateHeartbeatHealth( sandbox.last_heartbeat, @@ -1366,6 +1405,10 @@ export class SandboxLifecycleManager implements SandboxLifecycle { return "sandbox_terminated"; } + if (sandbox.last_heartbeat !== null) { + await this.alarmScheduler.schedule(sandbox.last_heartbeat + this.config.heartbeat.timeoutMs); + } + // Evaluate inactivity timeout const connectedClients = this.getConnectedClientCount(); const inactivityState = { @@ -1509,6 +1552,12 @@ export class SandboxLifecycleManager implements SandboxLifecycle { this.storage.updateSandboxLastActivity(timestamp); } + async recordStartupHeartbeat(sandboxId: string, timestamp: number): Promise { + if (!this.storage.recordStartupHeartbeat(sandboxId, timestamp)) return false; + await this.alarmScheduler.schedule(timestamp + this.config.connectingTimeout.timeoutMs); + return true; + } + /** * Schedule an inactivity check alarm. */ @@ -1652,6 +1701,13 @@ export class SandboxLifecycleManager implements SandboxLifecycle { await this.alarmScheduler.schedule(createdAt + this.config.connectingTimeout.timeoutMs); } + /** + * Check whether queue work must wait for the current snapshot to finish. + */ + isSnapshotting(): boolean { + return this.storage.getSandbox()?.status === "snapshotting"; + } + /** * Check if a sandbox spawn is currently in progress. * Used by SessionDO to coordinate spawn decisions. diff --git a/packages/control-plane/src/session/components.ts b/packages/control-plane/src/session/components.ts index c0eebd985e..aa65e473c0 100644 --- a/packages/control-plane/src/session/components.ts +++ b/packages/control-plane/src/session/components.ts @@ -457,7 +457,7 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi messenger, recordTerminalMessage, statusService, - (reason) => lifecycleManager.triggerSnapshot(reason), + (reason) => lifecycleManager.triggerSnapshot(reason, () => messageQueue.processMessageQueue()), updateLastActivity, () => lifecycleManager.scheduleInactivityCheck(), () => messageQueue.processMessageQueue(), @@ -469,8 +469,13 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi eventRepository, messenger, diffService, + wsManager, (title, options) => titleService.applySessionTitleUpdate(title, options), - updateLastActivity + updateLastActivity, + () => lifecycleManager.scheduleInactivityCheck(), + () => messageQueue.processMessageQueue(), + () => lifecycleManager.isProviderStartupPending(), + () => lifecycleManager.scheduleDisconnectCheck() ); const pushService = new SandboxPushService(log, wsManager); const sandboxEventProcessor = new SessionSandboxEventProcessor( @@ -676,9 +681,7 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi sessionCoreRepository, sandboxRepository, lifecycleManager, - messenger, backgroundTasks, - messageQueue, participantService, presenceService, snapshotReader, @@ -782,7 +785,8 @@ export function createSessionRuntime(platform: SessionPlatform, env: Env): Sessi log, sockets, clientCommands, - processSandboxEvent: (event) => sandboxEventProcessor.processSandboxEvent(event), + processSandboxEvent: (event, sender) => + sandboxEventProcessor.processSandboxEvent(event, sender), clock, }), disconnects: new SessionDisconnectHandler({ @@ -916,6 +920,7 @@ function createLifecycleManager(deps: LifecycleManagerDeps): SandboxLifecycleMan const config = { ...DEFAULT_LIFECYCLE_CONFIG, controlPlaneUrl, + earlySandboxConnection: env.EARLY_SANDBOX_CONNECTION === "1", model: DEFAULT_MODEL, // Re-derived per use until the session row exists: on the first-ever // activation the manager is built during the init request, before the row diff --git a/packages/control-plane/src/session/connection-authenticator.test.ts b/packages/control-plane/src/session/connection-authenticator.test.ts new file mode 100644 index 0000000000..fadc113224 --- /dev/null +++ b/packages/control-plane/src/session/connection-authenticator.test.ts @@ -0,0 +1,56 @@ +import { describe, expect, it, vi } from "vitest"; +import type { Logger } from "../logger"; +import { + SessionConnectionAuthenticator, + type SessionConnectionAuthenticatorDeps, +} from "./connection-authenticator"; +import type { SandboxRow } from "./types"; + +describe("SessionConnectionAuthenticator", () => { + it("persists admission state before accepting a sandbox socket", async () => { + const calls: string[] = []; + const sandbox = { + modal_sandbox_id: "sandbox-id", + auth_token: "sandbox-token", + auth_token_hash: null, + status: "ready", + } as SandboxRow; + const deps = { + wsManager: { + createUpgradeSockets: vi.fn(() => ({ client: {} as WebSocket, server: {} as WebSocket })), + acceptAndSetSandboxSocket: vi.fn(() => { + calls.push("accept"); + throw new Error("accept failed"); + }), + }, + sessionCoreRepository: { getSession: vi.fn(() => null) }, + sandboxRepository: { + getSandbox: vi.fn(() => sandbox), + }, + lifecycleManager: { + recordStartupHeartbeat: vi.fn(async () => { + calls.push("persist"); + return true; + }), + }, + } as unknown as SessionConnectionAuthenticatorDeps; + const log = { + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + } as unknown as Logger; + const authenticator = new SessionConnectionAuthenticator(deps); + const request = new Request("https://example.test?type=sandbox", { + headers: { + Authorization: "Bearer sandbox-token", + "X-Sandbox-ID": "sandbox-id", + }, + }); + + const response = await authenticator.handleWebSocketUpgrade(request, new URL(request.url), log); + + expect(response.status).toBe(500); + expect(calls).toEqual(["persist", "accept"]); + }); +}); diff --git a/packages/control-plane/src/session/connection-authenticator.ts b/packages/control-plane/src/session/connection-authenticator.ts index e66e6a1f71..53dfaf563e 100644 --- a/packages/control-plane/src/session/connection-authenticator.ts +++ b/packages/control-plane/src/session/connection-authenticator.ts @@ -11,8 +11,6 @@ import { isValidSandboxToken } from "./sandbox-access"; import { resolveParticipantName } from "./participant-name"; import { getAvatarUrl, type ParticipantService } from "./participant-service"; import type { PresenceService } from "./presence-service"; -import type { SessionMessageQueue } from "./message-queue"; -import type { SessionMessenger } from "./messenger"; import type { SandboxRepository } from "./sandbox-repository"; import type { SessionCoreRepository } from "./session-core-repository"; import type { SessionSnapshotReader } from "./snapshot-reader"; @@ -30,9 +28,7 @@ export interface SessionConnectionAuthenticatorDeps { sessionCoreRepository: SessionCoreRepository; sandboxRepository: SandboxRepository; lifecycleManager: SandboxLifecycleManager; - messenger: SessionMessenger; backgroundTasks: BackgroundTasks; - messageQueue: Pick; participantService: ParticipantService; presenceService: PresenceService; snapshotReader: SessionSnapshotReader; @@ -60,9 +56,7 @@ export class SessionConnectionAuthenticator { sessionCoreRepository, sandboxRepository, lifecycleManager, - messenger, backgroundTasks, - messageQueue, } = this.deps; log.debug("WebSocket upgrade requested"); const isSandbox = url.searchParams.get("type") === "sandbox"; @@ -81,7 +75,7 @@ export class SessionConnectionAuthenticator { const expectedSandboxId = sandbox?.modal_sandbox_id; // Validate sandbox ID first (catches stale sandboxes reconnecting after restore) - if (expectedSandboxId && sandboxId !== expectedSandboxId) { + if (!expectedSandboxId || sandboxId !== expectedSandboxId) { log.warn("ws.connect", { event: "ws.connect", ws_type: "sandbox", @@ -142,6 +136,19 @@ export class SessionConnectionAuthenticator { }); return new Response("Sandbox is stopped", { status: 410 }); } + if (currentSandbox?.status === "snapshotting") { + log.info("ws.connect", { + event: "ws.connect", + ws_type: "sandbox", + outcome: "rejected", + reject_reason: "sandbox_snapshotting", + duration_ms: Date.now() - wsStartTime, + }); + return new Response("Sandbox snapshot is in progress", { + status: 503, + headers: { "Retry-After": "1" }, + }); + } if ( currentSandbox?.modal_sandbox_id !== expectedSandboxId || currentSandbox?.auth_token_hash !== sandbox?.auth_token_hash || @@ -160,25 +167,13 @@ export class SessionConnectionAuthenticator { const sandboxId = request.headers.get("X-Sandbox-ID"); if (isSandbox) { - // The lifecycle manager publishes access after any pending provider - // startup has persisted its URLs and credentials. - const accessIsPersisted = !lifecycleManager.isProviderStartupPending(); - const { replaced } = wsManager.acceptAndSetSandboxSocket(server, sandboxId ?? undefined); + const now = Date.now(); + if (!(await lifecycleManager.recordStartupHeartbeat(sandboxId!, now))) { + return new Response("Forbidden: Sandbox credentials changed", { status: 403 }); + } + const { replaced } = wsManager.acceptAndSetSandboxSocket(server, sandboxId!); // Notify manager that sandbox connected so it can reset the spawning flag lifecycleManager.onSandboxConnected(); - sandboxRepository.updateSandboxStatus("ready"); - messenger.broadcast({ type: "sandbox_status", status: "ready" }); - if (accessIsPersisted) { - messenger.broadcast({ type: "sandbox_access_changed" }); - } - - // Set initial activity timestamp and schedule inactivity check - // IMPORTANT: Must await to ensure alarm is scheduled before returning - const now = Date.now(); - lifecycleManager.updateLastActivity(now); - sandboxRepository.updateSandboxHeartbeat(now); - await lifecycleManager.scheduleInactivityCheck(); - log.info("ws.connect", { event: "ws.connect", ws_type: "sandbox", @@ -187,11 +182,6 @@ export class SessionConnectionAuthenticator { replaced_existing: replaced, duration_ms: Date.now() - now, }); - - // Process any pending messages now that sandbox is connected - backgroundTasks.submit(() => messageQueue.processMessageQueue(), { - name: "message_queue.process", - }); } else { const wsId = `ws-${Date.now()}-${Math.random().toString(36).slice(2, 10)}`; wsManager.acceptClientSocket(server, wsId); diff --git a/packages/control-plane/src/session/message-queue.test.ts b/packages/control-plane/src/session/message-queue.test.ts index 016ad8f78c..0c83a6fc6d 100644 --- a/packages/control-plane/src/session/message-queue.test.ts +++ b/packages/control-plane/src/session/message-queue.test.ts @@ -184,8 +184,10 @@ function buildQueue() { getUnreferenced: vi.fn((): SessionAttachmentRow[] => []), }; + const getSandboxSocket = vi.fn(() => null as WebSocket | null); const wsManager = { - getSandboxSocket: vi.fn(() => null as WebSocket | null), + getSandboxSocket, + getExecutionSocket: vi.fn(() => getSandboxSocket()), send: vi.fn((_ws: WebSocket, _message: ServerMessage) => true), }; @@ -208,6 +210,7 @@ function buildQueue() { }; const sandboxLifecycle = { spawnSandbox: vi.fn(async () => {}), + isSnapshotting: vi.fn(() => false), updateLastActivity: vi.fn((_timestamp: number) => {}), terminateUnresponsiveSandbox: vi.fn(async () => {}), reportSandboxError: vi.fn((_reason: string) => {}), @@ -477,6 +480,75 @@ describe("SessionMessageQueue", () => { expect(h.callbackService.notifyStarted).not.toHaveBeenCalled(); }); + it("keeps a prompt pending while a socketless sandbox is snapshotting", async () => { + const h = buildQueue(); + h.repository.getNextPendingMessage.mockReturnValue(createMessage()); + h.sandboxLifecycle.isSnapshotting.mockReturnValue(true); + + await h.queue.processMessageQueue(); + + expect(h.sandboxLifecycle.spawnSandbox).not.toHaveBeenCalled(); + expect(h.repository.startMessageProcessing).not.toHaveBeenCalled(); + }); + + it("leaves a prompt pending without spawning when only a control socket is attached", async () => { + const h = buildQueue(); + h.repository.getNextPendingMessage.mockReturnValue(createMessage()); + h.wsManager.getSandboxSocket.mockReturnValue({ + readyState: WebSocket.OPEN, + } as WebSocket); + h.wsManager.getExecutionSocket.mockReturnValue(null); + + await h.queue.processMessageQueue(); + + expect(h.sandboxLifecycle.spawnSandbox).not.toHaveBeenCalled(); + expect(h.repository.startMessageProcessing).not.toHaveBeenCalled(); + expect(h.wsManager.send).not.toHaveBeenCalled(); + expect(h.setAlarm).not.toHaveBeenCalled(); + expect(h.getProviderAuthenticationError).not.toHaveBeenCalled(); + }); + + it("recovers an expired stop confirmation before gating a control-only socket", async () => { + const h = buildQueue(); + h.wsManager.getSandboxSocket.mockReturnValue({ readyState: WebSocket.OPEN } as WebSocket); + h.wsManager.getExecutionSocket.mockReturnValue(null); + h.repository.getMessageAwaitingStopConfirmation + .mockReturnValueOnce({ id: "msg-stopped", deadline: Date.now() - 1 }) + .mockReturnValueOnce({ id: "msg-stopped", deadline: Date.now() - 1 }) + .mockReturnValue(null); + + await h.queue.processMessageQueue(); + + expect(h.sandboxLifecycle.terminateUnresponsiveSandbox).toHaveBeenCalledWith( + "stop_confirmation_timeout" + ); + }); + + it("re-reads sandbox sockets after provider auth validation", async () => { + const h = buildQueue(); + const staleSocket = { readyState: WebSocket.OPEN } as WebSocket; + let resolveAuth!: () => void; + h.repository.getNextPendingMessage.mockReturnValue(createMessage()); + h.wsManager.getSandboxSocket.mockReturnValue(staleSocket); + h.wsManager.getExecutionSocket.mockReturnValue(staleSocket); + h.getProviderAuthenticationError.mockImplementation( + () => + new Promise((resolve) => { + resolveAuth = () => resolve(null); + }) + ); + + const processing = h.queue.processMessageQueue(); + await vi.waitFor(() => expect(h.getProviderAuthenticationError).toHaveBeenCalledOnce()); + h.wsManager.getSandboxSocket.mockReturnValue({ readyState: WebSocket.OPEN } as WebSocket); + h.wsManager.getExecutionSocket.mockReturnValue(null); + resolveAuth(); + await processing; + + expect(h.repository.startMessageProcessing).not.toHaveBeenCalled(); + expect(h.wsManager.send).not.toHaveBeenCalled(); + }); + it.each(["cancelled", "archived"] as const)( "does not dispatch queued work for a %s session", async (status) => { @@ -809,6 +881,7 @@ describe("SessionMessageQueue", () => { const sandboxWs = { readyState: 1 } as WebSocket; h.repository.getNextPendingMessage.mockReturnValue(createMessage()); h.wsManager.getSandboxSocket.mockReturnValue(sandboxWs); + h.wsManager.getExecutionSocket.mockReturnValue(sandboxWs); await h.queue.processMessageQueue(); @@ -840,7 +913,9 @@ describe("SessionMessageQueue", () => { h.repository.getNextPendingMessage.mockReturnValue( createMessage({ source: "github", origin_context: JSON.stringify(origin) }) ); - h.wsManager.getSandboxSocket.mockReturnValue({ readyState: 1 } as WebSocket); + const sandboxWs = { readyState: 1 } as WebSocket; + h.wsManager.getSandboxSocket.mockReturnValue(sandboxWs); + h.wsManager.getExecutionSocket.mockReturnValue(sandboxWs); await h.queue.processMessageQueue(); @@ -895,6 +970,7 @@ describe("SessionMessageQueue", () => { model === "xai/grok-4.5" ? "No xAI authentication is configured" : null ); h.wsManager.getSandboxSocket.mockReturnValue(sandboxWs); + h.wsManager.getExecutionSocket.mockReturnValue(sandboxWs); await h.queue.processMessageQueue(); diff --git a/packages/control-plane/src/session/message-queue.ts b/packages/control-plane/src/session/message-queue.ts index d2a15a5d7b..59476d76ad 100644 --- a/packages/control-plane/src/session/message-queue.ts +++ b/packages/control-plane/src/session/message-queue.ts @@ -371,6 +371,12 @@ export class SessionMessageQueue { return; } + const initialControlSocket = this.wsManager.getSandboxSocket(); + if (initialControlSocket && !this.wsManager.getExecutionSocket()) { + this.log.debug("processMessageQueue: sandbox execution not ready"); + return; + } + const message = this.messageRepository.getNextPendingMessage(); if (!message) { return; @@ -392,8 +398,12 @@ export class SessionMessageQueue { return; } - const sandboxWs = this.wsManager.getSandboxSocket(); - if (!sandboxWs) { + const controlSocket = this.wsManager.getSandboxSocket(); + if (!controlSocket) { + if (this.sandboxLifecycle.isSnapshotting()) { + this.log.debug("processMessageQueue: sandbox snapshot in progress"); + return; + } this.log.info("prompt.dispatch", { event: "prompt.dispatch", message_id: message.id, @@ -424,6 +434,11 @@ export class SessionMessageQueue { ); return; } + const sandboxWs = this.wsManager.getExecutionSocket(); + if (!sandboxWs) { + this.log.debug("processMessageQueue: sandbox execution not ready"); + return; + } const author = this.participantRepository.getParticipantById(message.author_id); if (!author) { diff --git a/packages/control-plane/src/session/message-router.ts b/packages/control-plane/src/session/message-router.ts index 8cab30ee86..e97267787f 100644 --- a/packages/control-plane/src/session/message-router.ts +++ b/packages/control-plane/src/session/message-router.ts @@ -39,7 +39,7 @@ export interface SessionMessageRouterDeps; clientCommands: SessionClientCommands; - processSandboxEvent: (event: SandboxEvent) => Promise; + processSandboxEvent: (event: SandboxEvent, sender: Connection) => Promise; clock: Clock; } @@ -52,18 +52,18 @@ export class SessionMessageRouter { if (typeof message !== "string") return; if (this.deps.sockets.classify(connection).kind === "sandbox") { - await this.handleSandboxMessage(message); + await this.handleSandboxMessage(connection, message); } else { await this.handleClientMessage(connection, message); } } - private async handleSandboxMessage(message: string): Promise { + private async handleSandboxMessage(connection: Connection, message: string): Promise { const parsed = this.parseMessage(message, "sandbox", sandboxEventSchema); if (!parsed.valid) return; try { - await this.deps.processSandboxEvent(parsed.data); + await this.deps.processSandboxEvent(parsed.data, connection); } catch (error) { this.deps.log.error("Error processing sandbox message", { error: error instanceof Error ? error : String(error), diff --git a/packages/control-plane/src/session/messenger.test.ts b/packages/control-plane/src/session/messenger.test.ts index 61ebc0afd4..51247eb6d0 100644 --- a/packages/control-plane/src/session/messenger.test.ts +++ b/packages/control-plane/src/session/messenger.test.ts @@ -15,7 +15,7 @@ function harness(overrides: { sandboxSocket?: WebSocket | null; sendResult?: boo fn(clientB); } ), - getSandboxSocket: vi.fn(() => sandbox), + getExecutionSocket: vi.fn(() => sandbox), send: vi.fn(() => overrides.sendResult ?? true), }; return { messenger: new SessionMessengerImpl(wsManager), wsManager, clientA, clientB, sandbox }; diff --git a/packages/control-plane/src/session/messenger.ts b/packages/control-plane/src/session/messenger.ts index ea7226d09b..d5b6bc18a1 100644 --- a/packages/control-plane/src/session/messenger.ts +++ b/packages/control-plane/src/session/messenger.ts @@ -20,7 +20,7 @@ import type { SessionWebSocketManager } from "./websocket-manager"; */ type DeliverySockets = Pick< SessionWebSocketManager, - "forEachClientSocket" | "getSandboxSocket" | "send" + "forEachClientSocket" | "getExecutionSocket" | "send" >; export class SandboxDeliveryUnavailableError extends Error { @@ -49,7 +49,7 @@ export class SessionMessengerImpl implements SessionMessenger { } sendToSandbox(command: SandboxCommand): Promise { - const ws = this.wsManager.getSandboxSocket(); + const ws = this.wsManager.getExecutionSocket(); if (!ws) return Promise.reject(new SandboxDeliveryUnavailableError()); return this.wsManager.send(ws, command) ? Promise.resolve() diff --git a/packages/control-plane/src/session/sandbox-events/context.ts b/packages/control-plane/src/session/sandbox-events/context.ts index 98430dba66..eaf623e5a4 100644 --- a/packages/control-plane/src/session/sandbox-events/context.ts +++ b/packages/control-plane/src/session/sandbox-events/context.ts @@ -15,6 +15,7 @@ export interface SandboxEventContext { messageId: string | null; /** The processing message as of event arrival (single DO turn — stable). */ processingMessage: { id: string } | null; + sender: WebSocket | null; } /** diff --git a/packages/control-plane/src/session/sandbox-events/processor.test.ts b/packages/control-plane/src/session/sandbox-events/processor.test.ts index 9faa897616..80d054a416 100644 --- a/packages/control-plane/src/session/sandbox-events/processor.test.ts +++ b/packages/control-plane/src/session/sandbox-events/processor.test.ts @@ -35,6 +35,8 @@ function createProcessor() { const getProcessingMessage = vi.fn(() => null as { id: string } | null); const repository = { updateSandboxHeartbeat: vi.fn(), + updateSandboxStatus: vi.fn(), + getSandbox: vi.fn(() => ({ status: "connecting" })), recordReportedSandboxRuntimeVersion: vi.fn(), getProcessingMessage, addSessionCost: vi.fn(), @@ -65,8 +67,12 @@ function createProcessor() { notifyComplete: vi.fn(async () => {}), }; + const defaultSender = { readyState: WebSocket.OPEN } as WebSocket; + const getSandboxSocket = vi.fn(() => null as WebSocket | null); const wsManager = { - getSandboxSocket: vi.fn(() => null as WebSocket | null), + getSandboxSocket, + getExecutionSocket: vi.fn(() => getSandboxSocket()), + isCurrentSandboxSocket: vi.fn(() => true), send: vi.fn(() => true), }; @@ -77,6 +83,7 @@ function createProcessor() { const projectTerminalMessage = vi.fn(async () => {}); const statusService = { reconcileAfterExecution: vi.fn(async (_success: boolean) => {}) }; const scheduleInactivityCheck = vi.fn(async () => {}); + const scheduleDisconnectCheck = vi.fn(async () => {}); const processMessageQueue = vi.fn(async () => {}); const broadcastPromptQueue = vi.fn(); const updateLastActivity = vi.fn(); @@ -131,14 +138,25 @@ function createProcessor() { eventRepository, messenger, diffService as unknown as SessionDiffService, + wsManager as unknown as SessionWebSocketManager, applySessionTitleUpdate, - updateLastActivity + updateLastActivity, + scheduleInactivityCheck, + processMessageQueue, + () => false, + scheduleDisconnectCheck ), pushService ); return { - processor, + processor: { + processSandboxEvent: (event: SandboxEvent & { ackId?: string }, sender?: WebSocket | null) => + processor.processSandboxEvent( + event, + sender === undefined ? (getSandboxSocket() ?? defaultSender) : sender + ), + }, pushService, artifactRepository, repository, @@ -151,6 +169,7 @@ function createProcessor() { projectTerminalMessage, statusService, scheduleInactivityCheck, + scheduleDisconnectCheck, processMessageQueue, broadcastPromptQueue, updateLastActivity, @@ -210,6 +229,17 @@ describe("SessionSandboxEventProcessor", () => { expect(h.broadcast).not.toHaveBeenCalled(); }); + it("accepts a senderless heartbeat from authenticated HTTP ingestion", async () => { + const h = createProcessor(); + + await h.processor.processSandboxEvent( + { type: "heartbeat", sandboxId: "sb-1", status: "booting", timestamp: 1000 }, + null + ); + + expect(h.repository.updateSandboxHeartbeat).toHaveBeenCalledWith(expect.any(Number)); + }); + it("applies session_title without storing a timeline event", async () => { const h = createProcessor(); const event: SandboxEvent = { @@ -242,6 +272,82 @@ describe("SessionSandboxEventProcessor", () => { expect(h.diffService.pinBaselines).toHaveBeenCalledWith(event); }); + it("grants readiness, starts inactivity, and drains the queue for the captured sender", async () => { + const h = createProcessor(); + const sender = { readyState: WebSocket.OPEN } as WebSocket; + + await h.processor.processSandboxEvent( + { type: "ready", sandboxId: "sb-1", timestamp: 1000 }, + sender + ); + + expect(h.wsManager.isCurrentSandboxSocket).toHaveBeenCalledWith(sender, "sb-1"); + expect(h.repository.updateSandboxStatus).toHaveBeenCalledWith("ready"); + expect(h.repository.updateSandboxHeartbeat).toHaveBeenCalledWith(expect.any(Number)); + expect(h.updateLastActivity).toHaveBeenCalledWith(expect.any(Number)); + expect(h.scheduleDisconnectCheck).toHaveBeenCalledOnce(); + expect(h.scheduleInactivityCheck).toHaveBeenCalledOnce(); + expect(h.processMessageQueue).toHaveBeenCalledOnce(); + }); + + it("ignores ready from a stale sender", async () => { + const h = createProcessor(); + h.wsManager.isCurrentSandboxSocket.mockReturnValue(false); + + await h.processor.processSandboxEvent( + { type: "ready", sandboxId: "sb-1", timestamp: 1000 }, + {} as WebSocket + ); + + expect(h.repository.recordReportedSandboxRuntimeVersion).not.toHaveBeenCalled(); + expect(h.processMessageQueue).not.toHaveBeenCalled(); + }); + + it("retries post-status ready work after a scheduling failure", async () => { + const h = createProcessor(); + h.scheduleInactivityCheck.mockRejectedValueOnce(new Error("alarm unavailable")); + const sender = {} as WebSocket; + const event = { type: "ready", sandboxId: "sb-1", timestamp: 1000 } as const; + + await expect(h.processor.processSandboxEvent(event, sender)).rejects.toThrow( + "alarm unavailable" + ); + h.repository.getSandbox.mockReturnValue({ status: "ready" }); + await h.processor.processSandboxEvent(event, sender); + + expect(h.repository.updateSandboxStatus).toHaveBeenCalledWith("ready"); + expect(h.scheduleInactivityCheck).toHaveBeenCalledTimes(2); + expect(h.processMessageQueue).toHaveBeenCalledOnce(); + }); + + it("preserves snapshotting and defers queue drain when a reconnect announces ready", async () => { + const h = createProcessor(); + h.repository.getSandbox.mockReturnValue({ status: "snapshotting" }); + + await h.processor.processSandboxEvent( + { type: "ready", sandboxId: "sb-1", timestamp: 1000 }, + {} as WebSocket + ); + + expect(h.repository.updateSandboxStatus).not.toHaveBeenCalled(); + expect(h.processMessageQueue).not.toHaveBeenCalled(); + expect(h.scheduleInactivityCheck).not.toHaveBeenCalled(); + expect(h.scheduleDisconnectCheck).toHaveBeenCalledOnce(); + }); + + it("treats senderless ready as metadata-only", async () => { + const h = createProcessor(); + + await h.processor.processSandboxEvent( + { type: "ready", sandboxId: "sb-1", timestamp: 1000 }, + null + ); + + expect(h.diffService.pinBaselines).toHaveBeenCalledOnce(); + expect(h.repository.updateSandboxStatus).not.toHaveBeenCalled(); + expect(h.processMessageQueue).not.toHaveBeenCalled(); + }); + it("records the reported runtime version on ready", async () => { const h = createProcessor(); @@ -703,6 +809,29 @@ describe("SessionSandboxEventProcessor", () => { }); describe("ACK mechanism", () => { + it("returns a critical ACK to the captured sender after replacement", async () => { + const h = createProcessor(); + const sender = {} as WebSocket; + h.wsManager.getSandboxSocket.mockReturnValue({} as WebSocket); + + await h.processor.processSandboxEvent( + { + type: "error", + messageId: "msg-stale", + error: "failed", + sandboxId: "sb-1", + timestamp: 1000, + ackId: "error:stale", + }, + sender + ); + + expect(h.wsManager.send).toHaveBeenCalledWith(sender, { + type: "ack", + ackId: "error:stale", + }); + }); + it("sends ACK after execution_complete when ackId is present", async () => { const h = createProcessor(); const sandboxWs = {} as WebSocket; diff --git a/packages/control-plane/src/session/sandbox-events/processor.ts b/packages/control-plane/src/session/sandbox-events/processor.ts index 2a600d233b..b920fae1e7 100644 --- a/packages/control-plane/src/session/sandbox-events/processor.ts +++ b/packages/control-plane/src/session/sandbox-events/processor.ts @@ -39,7 +39,10 @@ export class SessionSandboxEventProcessor { private readonly pushService: SandboxPushService ) {} - async processSandboxEvent(event: SandboxEventWithAck): Promise { + async processSandboxEvent( + event: SandboxEventWithAck, + sender: WebSocket | null = null + ): Promise { if (event.type === "heartbeat" || event.type === "token") { this.log.debug("Sandbox event", { event_type: event.type }); } else if (event.type !== "execution_complete") { @@ -53,25 +56,26 @@ export class SessionSandboxEventProcessor { now, messageId: eventMessageId ?? processingMessage?.id ?? null, processingMessage, + sender, }; await this.dispatch(event, context); if (CRITICAL_EVENT_TYPES.has(event.type)) { - this.sendAck(event.ackId); + if (sender) this.sendAck(sender, event.ackId); } } private async dispatch(event: SandboxEvent, context: SandboxEventContext): Promise { switch (event.type) { case "heartbeat": - this.runtime.handleHeartbeat(context); + this.runtime.handleHeartbeat(event, context); return; case "session_title": this.runtime.handleSessionTitle(event); return; case "ready": - this.runtime.handleReady(event, context); + await this.runtime.handleReady(event, context); return; case "git_sync": this.runtime.handleGitSync(event, context); @@ -117,13 +121,10 @@ export class SessionSandboxEventProcessor { } } - private sendAck(ackId: string | undefined): void { + private sendAck(sender: WebSocket, ackId: string | undefined): void { if (!ackId) return; - const sandboxWs = this.wsManager.getSandboxSocket(); - if (sandboxWs) { - this.wsManager.send(sandboxWs, { type: "ack", ackId }); - } else { - this.log.debug("Cannot send ACK: no sandbox socket", { ack_id: ackId }); + if (!this.wsManager.send(sender, { type: "ack", ackId })) { + this.log.debug("Cannot send ACK: sender unavailable", { ack_id: ackId }); } } } diff --git a/packages/control-plane/src/session/sandbox-events/runtime.handler.ts b/packages/control-plane/src/session/sandbox-events/runtime.handler.ts index 758a20e11d..d49711a574 100644 --- a/packages/control-plane/src/session/sandbox-events/runtime.handler.ts +++ b/packages/control-plane/src/session/sandbox-events/runtime.handler.ts @@ -5,6 +5,7 @@ import type { SessionMessenger } from "../messenger"; import type { SandboxRepository } from "../sandbox-repository"; import type { SessionCoreRepository } from "../session-core-repository"; import type { SessionTitleUpdateOptions, SessionTitleUpdateResult } from "../title"; +import type { SessionWebSocketManager } from "../websocket-manager"; import { persistSandboxEvent, type SandboxEventContext } from "./context"; /** @@ -21,14 +22,25 @@ export class SandboxRuntimeEventHandler { private readonly eventRepository: EventRepository, private readonly messenger: SessionMessenger, private readonly diffService: SessionDiffService, + private readonly wsManager: SessionWebSocketManager, private readonly applySessionTitleUpdate: ( title: string, options?: SessionTitleUpdateOptions ) => SessionTitleUpdateResult, - private readonly updateLastActivity: (timestamp: number) => void + private readonly updateLastActivity: (timestamp: number) => void, + private readonly scheduleInactivityCheck: () => Promise, + private readonly processMessageQueue: () => Promise, + private readonly isProviderStartupPending: () => boolean, + private readonly scheduleDisconnectCheck: () => Promise ) {} - handleHeartbeat(context: SandboxEventContext): void { + handleHeartbeat( + event: Extract, + context: SandboxEventContext + ): void { + if (context.sender && !this.wsManager.isCurrentSandboxSocket(context.sender, event.sandboxId)) { + return; + } this.sandboxRepository.updateSandboxHeartbeat(context.now); // A quiet tool call may emit no events for longer than the inactivity // timeout. While its message is processing, the bridge heartbeat proves @@ -42,7 +54,46 @@ export class SandboxRuntimeEventHandler { this.applySessionTitleUpdate(event.title, { onlyIfUnset: true }); } - handleReady(event: Extract, context: SandboxEventContext): void { + async handleReady( + event: Extract, + context: SandboxEventContext + ): Promise { + const status = this.sandboxRepository.getSandbox()?.status; + if (!context.sender) { + this.recordReadyMetadata(event, context); + return; + } + if ( + !status || + !["spawning", "connecting", "failed", "ready", "snapshotting"].includes(status) || + !this.wsManager.isCurrentSandboxSocket(context.sender, event.sandboxId) + ) { + return; + } + + if (status !== "ready") this.recordReadyMetadata(event, context); + if (status === "snapshotting") { + this.sandboxRepository.updateSandboxHeartbeat(context.now); + await this.scheduleDisconnectCheck(); + return; + } + + this.sandboxRepository.updateSandboxStatus("ready"); + this.sandboxRepository.updateSandboxHeartbeat(context.now); + this.updateLastActivity(context.now); + await this.scheduleDisconnectCheck(); + await this.scheduleInactivityCheck(); + this.messenger.broadcast({ type: "sandbox_status", status: "ready" }); + if (!this.isProviderStartupPending()) { + this.messenger.broadcast({ type: "sandbox_access_changed" }); + } + await this.processMessageQueue(); + } + + private recordReadyMetadata( + event: Extract, + context: SandboxEventContext + ): void { this.diffService.pinBaselines(event); // Fills the column a fresh spawn cleared; a restore has already seeded // the snapshot's version, which outranks whatever this sandbox reports. diff --git a/packages/control-plane/src/session/sandbox-push-service.test.ts b/packages/control-plane/src/session/sandbox-push-service.test.ts index a6099abcff..3cedd1e9f2 100644 --- a/packages/control-plane/src/session/sandbox-push-service.test.ts +++ b/packages/control-plane/src/session/sandbox-push-service.test.ts @@ -18,7 +18,8 @@ function createPushSpec(repoOwner: string, repoName: string, targetBranch: strin function createService() { const sandboxWs = { readyState: WebSocket.OPEN } as WebSocket; const wsManager = { - getSandboxSocket: vi.fn(() => sandboxWs), + getSandboxSocket: vi.fn<() => WebSocket | null>(() => sandboxWs), + getExecutionSocket: vi.fn<() => WebSocket | null>(() => sandboxWs), send: vi.fn(() => true), }; const log = { @@ -33,6 +34,19 @@ function createService() { } describe("SandboxPushService", () => { + it("rejects a push while the sandbox is connected but not execution-ready", async () => { + const h = createService(); + h.wsManager.getExecutionSocket.mockReturnValue(null); + + await expect( + h.service.pushBranchToRemote(createPushSpec("acme", "web", "feature/test")) + ).resolves.toEqual({ + success: false, + error: "Sandbox is not ready to push branches", + }); + expect(h.wsManager.send).not.toHaveBeenCalled(); + }); + it("fails a push immediately when the command cannot be delivered", async () => { vi.useFakeTimers(); try { diff --git a/packages/control-plane/src/session/sandbox-push-service.ts b/packages/control-plane/src/session/sandbox-push-service.ts index a4c7bebf14..7d97dff9f1 100644 --- a/packages/control-plane/src/session/sandbox-push-service.ts +++ b/packages/control-plane/src/session/sandbox-push-service.ts @@ -41,9 +41,12 @@ export class SandboxPushService { async pushBranchToRemote( pushSpec: GitPushSpec ): Promise<{ success: true } | { success: false; error: string }> { - const sandboxWs = this.wsManager.getSandboxSocket(); + const sandboxWs = this.wsManager.getExecutionSocket(); if (!sandboxWs) { + if (this.wsManager.getSandboxSocket()) { + return { success: false, error: "Sandbox is not ready to push branches" }; + } this.log.info("No sandbox connected, assuming branch was pushed manually"); return { success: true }; } diff --git a/packages/control-plane/src/session/sandbox-repository.test.ts b/packages/control-plane/src/session/sandbox-repository.test.ts index eadccfb20f..3d8c00558a 100644 --- a/packages/control-plane/src/session/sandbox-repository.test.ts +++ b/packages/control-plane/src/session/sandbox-repository.test.ts @@ -131,6 +131,7 @@ 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"); + expect(mock.calls[0].query).toContain("last_heartbeat = NULL"); expect(mock.calls[0].params).toEqual(["spawning", 1000, "modal-sb-1"]); }); @@ -229,6 +230,28 @@ describe("SandboxRepository", () => { }); }); + describe("recordStartupHeartbeat", () => { + it("atomically records connecting liveness for an admissible sandbox identity", () => { + repository.recordStartupHeartbeat("sandbox-current", 5000); + + expect(mock.calls[0].query).toContain("SET status = 'connecting', last_heartbeat = ?"); + expect(mock.calls[0].query).toContain("status NOT IN ('snapshotting', 'stopped', 'stale')"); + expect(mock.calls[0].query).toContain("modal_sandbox_id = ?"); + expect(mock.calls[0].params).toEqual([5000, "sandbox-current"]); + }); + }); + + describe("failStartupIfUnchanged", () => { + it("compares identity, attempt, status, and observed liveness", () => { + repository.failStartupIfUnchanged("sandbox-current", 1000, 3000); + + expect(mock.calls[0].query).toContain("status = 'failed'"); + expect(mock.calls[0].query).toContain("status IN ('spawning', 'connecting')"); + expect(mock.calls[0].query).toContain("MAX(created_at, COALESCE(last_heartbeat, 0)) = ?"); + expect(mock.calls[0].params).toEqual(["sandbox-current", 1000, 3000]); + }); + }); + describe("updateSandboxLastActivity", () => { it("updates activity timestamp", () => { repository.updateSandboxLastActivity(6000); diff --git a/packages/control-plane/src/session/sandbox-repository.ts b/packages/control-plane/src/session/sandbox-repository.ts index 112b501fdf..a371ce9b52 100644 --- a/packages/control-plane/src/session/sandbox-repository.ts +++ b/packages/control-plane/src/session/sandbox-repository.ts @@ -23,6 +23,7 @@ const ACCESS_ARTIFACT_COLUMNS: Record< export interface SandboxCircuitBreakerState { status: SandboxStatus; created_at: number; + last_heartbeat: number | null; modal_object_id: string | null; snapshot_image_id: string | null; snapshot_runtime_version: string | null; @@ -91,7 +92,7 @@ export class SandboxRepository { getSandboxWithCircuitBreaker(): SandboxCircuitBreakerState | null { const result = this.sql.exec( - `SELECT status, created_at, modal_object_id, snapshot_image_id, snapshot_runtime_version, spawn_failure_count, last_spawn_failure FROM sandbox LIMIT 1` + `SELECT status, created_at, last_heartbeat, modal_object_id, snapshot_image_id, snapshot_runtime_version, spawn_failure_count, last_spawn_failure FROM sandbox LIMIT 1` ); const rows = this.rows & { status: string }>(result); const row = rows[0]; @@ -137,7 +138,8 @@ export class SandboxRepository { tunnel_urls = NULL, ttyd_url = NULL, ttyd_token = NULL, - runtime_version = NULL + runtime_version = NULL, + last_heartbeat = NULL WHERE id = (SELECT id FROM sandbox LIMIT 1)`, data.status, data.createdAt, @@ -232,6 +234,32 @@ export class SandboxRepository { ); } + recordStartupHeartbeat(sandboxId: string, timestamp: number): boolean { + const result = this.sql.exec( + `UPDATE sandbox SET status = 'connecting', last_heartbeat = ? + WHERE modal_sandbox_id = ? AND status NOT IN ('snapshotting', 'stopped', 'stale')`, + timestamp, + sandboxId + ); + result.toArray(); + return (result.rowsWritten ?? 0) > 0; + } + + failStartupIfUnchanged(sandboxId: string, createdAt: number, livenessAt: number): boolean { + const result = this.sql.exec( + `UPDATE sandbox SET status = 'failed' + WHERE modal_sandbox_id = ? + AND created_at = ? + AND status IN ('spawning', 'connecting') + AND MAX(created_at, COALESCE(last_heartbeat, 0)) = ?`, + sandboxId, + createdAt, + livenessAt + ); + result.toArray(); + return (result.rowsWritten ?? 0) > 0; + } + updateSandboxLastActivity(timestamp: number): void { this.sql.exec( `UPDATE sandbox SET last_activity = ? WHERE id = (SELECT id FROM sandbox LIMIT 1)`, diff --git a/packages/control-plane/src/session/server.test.ts b/packages/control-plane/src/session/server.test.ts index 738de431d1..71c177e170 100644 --- a/packages/control-plane/src/session/server.test.ts +++ b/packages/control-plane/src/session/server.test.ts @@ -274,7 +274,7 @@ describe("SessionServer", () => { }); }); - it("parses and routes sandbox events without exposing a socket type", async () => { + it("routes sandbox events with the actual sending socket", async () => { const { server, messageDeps, setConnectionKind } = createHarness(); setConnectionKind("sandbox"); @@ -288,12 +288,15 @@ describe("SessionServer", () => { }) ); - expect(messageDeps.processSandboxEvent).toHaveBeenCalledWith({ - type: "heartbeat", - sandboxId: "sandbox-1", - timestamp: 1000, - status: "ready", - }); + expect(messageDeps.processSandboxEvent).toHaveBeenCalledWith( + { + type: "heartbeat", + sandboxId: "sandbox-1", + timestamp: 1000, + status: "ready", + }, + "sandbox" + ); }); it("schedules sandbox reconnect checks and always reciprocates close", async () => { diff --git a/packages/control-plane/src/session/websocket-manager.test.ts b/packages/control-plane/src/session/websocket-manager.test.ts index 9934c724ec..e0ca51d651 100644 --- a/packages/control-plane/src/session/websocket-manager.test.ts +++ b/packages/control-plane/src/session/websocket-manager.test.ts @@ -256,6 +256,19 @@ describe("SessionWebSocketManagerImpl", () => { }); describe("acceptAndSetSandboxSocket", () => { + it("gates a newly accepted socket with persisted connecting status", () => { + const { manager, mockRepo } = createManager(); + const ws = createFakeWebSocket(); + const row = createSandboxRow("sb-1"); + row.status = "connecting"; + mockRepo.setSandbox(row); + + manager.acceptAndSetSandboxSocket(ws, "sb-1"); + + expect(manager.getSandboxSocket()).toBe(ws); + expect(manager.getExecutionSocket()).toBeNull(); + }); + it("accepts with sandbox + sid tags", () => { const { manager, sockets } = createManager(); const ws = createFakeWebSocket(); @@ -268,13 +281,13 @@ describe("SessionWebSocketManagerImpl", () => { expect(tags).toContain("sid:sandbox-abc"); }); - it("accepts with only sandbox tag when no sandboxId", () => { + it("accepts with sandbox identity tags", () => { const { manager, sockets } = createManager(); const ws = createFakeWebSocket(); - manager.acceptAndSetSandboxSocket(ws); + manager.acceptAndSetSandboxSocket(ws, "sb-1"); - expect(sockets.get(ws)).toEqual(["sandbox"]); + expect(sockets.get(ws)).toEqual(["sandbox", "sid:sb-1"]); }); it("closes existing sandbox socket and returns replaced=true", () => { @@ -289,21 +302,36 @@ describe("SessionWebSocketManagerImpl", () => { expect(oldWs.close).toHaveBeenCalledWith(1000, "New sandbox connecting"); }); + it("closes a hibernated sandbox socket when the in-memory cache is empty", () => { + const { manager, sockets, mockRepo } = createManager(); + const oldWs = createFakeWebSocket(); + const newWs = createFakeWebSocket(); + sockets.set(oldWs, ["sandbox", "sid:sb-1"]); + mockRepo.setSandbox(createSandboxRow("sb-1")); + + const result = manager.acceptAndSetSandboxSocket(newWs, "sb-1"); + + expect(result.replaced).toBe(true); + expect(oldWs.close).toHaveBeenCalledWith(1000, "New sandbox connecting"); + expect(manager.getSandboxSocket()).toBe(newWs); + }); + it("does not try to close an already-closed sandbox socket", () => { const { manager } = createManager(); const oldWs = createFakeWebSocket(WebSocket.CLOSED); const newWs = createFakeWebSocket(); - manager.acceptAndSetSandboxSocket(oldWs); - const result = manager.acceptAndSetSandboxSocket(newWs); + manager.acceptAndSetSandboxSocket(oldWs, "sb-1"); + const result = manager.acceptAndSetSandboxSocket(newWs, "sb-1"); expect(result.replaced).toBe(false); expect(oldWs.close).not.toHaveBeenCalled(); }); it("sets new socket as active sandbox", () => { - const { manager } = createManager(); + const { manager, mockRepo } = createManager(); const ws = createFakeWebSocket(); + mockRepo.setSandbox(createSandboxRow("sb-1")); manager.acceptAndSetSandboxSocket(ws, "sb-1"); @@ -313,8 +341,9 @@ describe("SessionWebSocketManagerImpl", () => { describe("getSandboxSocket", () => { it("returns cached socket if open", () => { - const { manager } = createManager(); + const { manager, mockRepo } = createManager(); const ws = createFakeWebSocket(); + mockRepo.setSandbox(createSandboxRow("sb-1")); manager.acceptAndSetSandboxSocket(ws, "sb-1"); @@ -408,7 +437,7 @@ describe("SessionWebSocketManagerImpl", () => { expect(ws.close).toHaveBeenCalledWith(1000, "Sandbox terminated"); }); - it("returns null and closes zombie WS when sandbox status is failed", () => { + it("keeps a current-attempt failed socket available for self-healing", () => { const { manager, sockets, mockRepo } = createManager(); const ws = createFakeWebSocket(); @@ -417,8 +446,47 @@ describe("SessionWebSocketManagerImpl", () => { row.status = "failed"; mockRepo.setSandbox(row); - expect(manager.getSandboxSocket()).toBeNull(); - expect(ws.close).toHaveBeenCalledWith(1000, "Sandbox terminated"); + expect(manager.getSandboxSocket()).toBe(ws); + expect(ws.close).not.toHaveBeenCalled(); + }); + }); + + describe("execution readiness", () => { + it("uses persisted ready status as the only execution gate", () => { + const { manager, mockRepo } = createManager(); + const ws = createFakeWebSocket(); + const row = createSandboxRow("sb-1"); + row.status = "connecting"; + mockRepo.setSandbox(row); + manager.acceptAndSetSandboxSocket(ws, "sb-1"); + + expect(manager.getExecutionSocket()).toBeNull(); + row.status = "ready"; + expect(manager.getExecutionSocket()).toBe(ws); + }); + + it("fences sender identity against the active sid-tagged socket", () => { + const { manager, mockRepo } = createManager(); + const oldWs = createFakeWebSocket(); + const newWs = createFakeWebSocket(); + mockRepo.setSandbox(createSandboxRow("sb-1")); + manager.acceptAndSetSandboxSocket(oldWs, "sb-1"); + manager.acceptAndSetSandboxSocket(newWs, "sb-1"); + + expect(manager.isCurrentSandboxSocket(oldWs, "sb-1")).toBe(false); + expect(manager.isCurrentSandboxSocket(newWs, "sb-1")).toBe(true); + expect(manager.isCurrentSandboxSocket(newWs, "wrong-id")).toBe(false); + }); + + it("recovers a hibernated socket for execution when persisted status is ready", () => { + const { manager, sockets, mockRepo } = createManager(); + const ws = createFakeWebSocket(); + sockets.set(ws, ["sandbox", "sid:sb-1"]); + const row = createSandboxRow("sb-1"); + row.status = "ready"; + mockRepo.setSandbox(row); + + expect(manager.getExecutionSocket()).toBe(ws); }); }); @@ -468,9 +536,10 @@ describe("SessionWebSocketManagerImpl", () => { }); it("returns false and does not clear when ws does not match", () => { - const { manager } = createManager(); + const { manager, mockRepo } = createManager(); const oldWs = createFakeWebSocket(); const newWs = createFakeWebSocket(); + mockRepo.setSandbox(createSandboxRow("sb-2")); manager.acceptAndSetSandboxSocket(oldWs, "sb-1"); manager.acceptAndSetSandboxSocket(newWs, "sb-2"); diff --git a/packages/control-plane/src/session/websocket-manager.ts b/packages/control-plane/src/session/websocket-manager.ts index 36c0e8b423..338fbeaf2a 100644 --- a/packages/control-plane/src/session/websocket-manager.ts +++ b/packages/control-plane/src/session/websocket-manager.ts @@ -35,7 +35,7 @@ export interface SessionWebSocketManager { * Accept a sandbox WebSocket, close any existing sandbox socket, and set * as the active sandbox connection. */ - acceptAndSetSandboxSocket(ws: WebSocket, sandboxId?: string): { replaced: boolean }; + acceptAndSetSandboxSocket(ws: WebSocket, sandboxId: string): { replaced: boolean }; /** Parse a WebSocket's tags to determine its kind and identity. */ classify(ws: WebSocket): ConnectionClassification; @@ -46,6 +46,9 @@ export interface SessionWebSocketManager { */ getSandboxSocket(): WebSocket | null; + getExecutionSocket(): WebSocket | null; + isCurrentSandboxSocket(ws: WebSocket, sandboxId: string): boolean; + /** Clear the in-memory sandbox socket reference. */ clearSandboxSocket(): void; @@ -116,22 +119,19 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { this.ctx.acceptWebSocket(ws, [`wsid:${wsId}`]); } - acceptAndSetSandboxSocket(ws: WebSocket, sandboxId?: string): { replaced: boolean } { - const tags = ["sandbox", ...(sandboxId ? [`sid:${sandboxId}`] : [])]; - this.ctx.acceptWebSocket(ws, tags); - + acceptAndSetSandboxSocket(ws: WebSocket, sandboxId: string): { replaced: boolean } { let replaced = false; - if (this.sandboxWs && this.sandboxWs !== ws) { - try { - if (this.sandboxWs.readyState === WebSocket.OPEN) { - this.sandboxWs.close(1000, "New sandbox connecting"); - replaced = true; - } - } catch { - // Ignore errors closing old WebSocket + const existing = new Set(this.ctx.getWebSockets()); + if (this.sandboxWs) existing.add(this.sandboxWs); + for (const candidate of existing) { + if (candidate === ws || this.classify(candidate).kind !== "sandbox") continue; + if (candidate.readyState === WebSocket.OPEN) { + this.close(candidate, 1000, "New sandbox connecting"); + replaced = true; } } + this.ctx.acceptWebSocket(ws, ["sandbox", `sid:${sandboxId}`]); this.sandboxWs = ws; return { replaced }; } @@ -162,7 +162,7 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { // After inactivity timeout or heartbeat stale, the DO closes the WS and sets // status to stopped/stale, but the close handshake may not complete before // hibernation. On wake, the zombie WS still appears OPEN — skip it. - const terminalStatuses = ["stopped", "failed", "stale"]; + const terminalStatuses = ["stopped", "stale"]; if (sandbox && terminalStatuses.includes(sandbox.status)) { this.sandboxWs = null; // Close any lingering sandbox WebSockets so they don't persist @@ -176,7 +176,12 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { } if (this.sandboxWs?.readyState === WebSocket.OPEN) { - return this.sandboxWs; + const parsed = this.classify(this.sandboxWs); + if (parsed.kind === "sandbox" && parsed.sandboxId === expectedSandboxId) { + return this.sandboxWs; + } + this.close(this.sandboxWs, 1000, "Sandbox identity changed"); + this.sandboxWs = null; } // Hibernation recovery: scan all WebSockets, validate sandbox identity @@ -202,6 +207,21 @@ export class SessionWebSocketManagerImpl implements SessionWebSocketManager { return null; } + getExecutionSocket(): WebSocket | null { + return this.sandboxRepository.getSandbox()?.status === "ready" ? this.getSandboxSocket() : null; + } + + isCurrentSandboxSocket(ws: WebSocket, sandboxId: string): boolean { + const sandbox = this.sandboxRepository.getSandbox(); + const parsed = this.classify(ws); + return ( + this.getSandboxSocket() === ws && + parsed.kind === "sandbox" && + parsed.sandboxId === sandboxId && + sandbox?.modal_sandbox_id === sandboxId + ); + } + clearSandboxSocket(): void { this.sandboxWs = null; } diff --git a/packages/control-plane/src/types.ts b/packages/control-plane/src/types.ts index 3dc468f9cb..784059f6ce 100644 --- a/packages/control-plane/src/types.ts +++ b/packages/control-plane/src/types.ts @@ -78,6 +78,7 @@ export interface Env { UNSAFE_ALLOW_ALL_USERS?: string; CF_ACCOUNT_ID?: string; // Cloudflare account ID SANDBOX_PROVIDER?: string; // "modal" (default), "daytona", "vercel", "opencomputer", or "e2b" + EARLY_SANDBOX_CONNECTION?: string; // Temporary default-off runtime activation flag MODAL_WORKSPACE?: string; // Modal workspace name MODAL_ENVIRONMENT?: string; // Modal environment name for dashboard URLs MODAL_ENVIRONMENT_WEB_SUFFIX?: string; // Modal environment web suffix for endpoint URLs diff --git a/packages/control-plane/test/integration/prompt-enqueue.test.ts b/packages/control-plane/test/integration/prompt-enqueue.test.ts index 314ae5d147..76d62a41eb 100644 --- a/packages/control-plane/test/integration/prompt-enqueue.test.ts +++ b/packages/control-plane/test/integration/prompt-enqueue.test.ts @@ -6,6 +6,7 @@ import { openSandboxWs, queryDO, seedSandboxAuth, + waitForSandboxStatus, } from "./helpers"; const SANDBOX_TOKEN = "prompt-order-sandbox-token"; @@ -132,6 +133,14 @@ describe("POST /internal/prompt", () => { }); expect(sandboxWs).not.toBeNull(); sandboxWs!.accept(); + sandboxWs!.send( + JSON.stringify({ + type: "ready", + sandboxId: SANDBOX_ID, + timestamp: Date.now() / 1000, + }) + ); + await waitForSandboxStatus(stub, "ready"); const enqueue = async (content: string) => { const response = await stub.fetch("http://internal/internal/prompt", { @@ -195,6 +204,14 @@ describe("POST /internal/prompt", () => { }); expect(sandboxWs).not.toBeNull(); sandboxWs!.accept(); + sandboxWs!.send( + JSON.stringify({ + type: "ready", + sandboxId: SANDBOX_ID, + timestamp: Date.now() / 1000, + }) + ); + await waitForSandboxStatus(stub, "ready"); const sandboxMessages = collectMessages(sandboxWs!, { timeoutMs: 500 }); const enqueue = (content: string) => diff --git a/packages/control-plane/test/integration/session-diffs.test.ts b/packages/control-plane/test/integration/session-diffs.test.ts index 84d5c22690..ca75d9e61e 100644 --- a/packages/control-plane/test/integration/session-diffs.test.ts +++ b/packages/control-plane/test/integration/session-diffs.test.ts @@ -9,10 +9,13 @@ import { queryDO, seedSandboxAuth, serviceFetch, + waitForSandboxStatus, } from "./helpers"; async function reportReady( stub: DurableObjectStub, + sessionName: string, + auth: { authToken: string; sandboxId: string }, repositories: Array<{ position: number; repoOwner: string; @@ -20,17 +23,19 @@ async function reportReady( baseSha: string; }> ): Promise { - const response = await stub.fetch("http://internal/internal/sandbox-event", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ + const { ws } = await openSandboxWs(sessionName, auth); + expect(ws).not.toBeNull(); + ws!.accept(); + ws!.send( + JSON.stringify({ type: "ready", - sandboxId: "sandbox-diff", + sandboxId: auth.sandboxId, timestamp: 100, repositories, - }), - }); - expect(response.status).toBe(200); + }) + ); + await waitForSandboxStatus(stub, "ready"); + ws!.close(); } function bundle(baseSha: string, patch = "diff --git a/src/app.ts b/src/app.ts\n") { @@ -98,12 +103,15 @@ describe("session diff routes", () => { it("requires the current sandbox token for writes", async () => { const sessionName = `diff-auth-${Date.now()}`; const { stub } = await initNamedSession(sessionName); - await seedSandboxAuth(stub, { + const auth = { authToken: "current-diff-token", sandboxId: "current-diff-sandbox", - }); + }; + await seedSandboxAuth(stub, auth); const baseSha = "a".repeat(40); - await reportReady(stub, [{ position: 0, repoOwner: "acme", repoName: "web-app", baseSha }]); + await reportReady(stub, sessionName, auth, [ + { position: 0, repoOwner: "acme", repoName: "web-app", baseSha }, + ]); const response = await SELF.fetch(`https://test.local/sessions/${sessionName}/diff`, { method: "PUT", @@ -166,7 +174,9 @@ describe("session diff routes", () => { await seedSandboxAuth(stub, auth); const baseSha = "a".repeat(40); const patch = "diff --git a/src/app.ts b/src/app.ts\n@@ -1 +1 @@\n-old\n+new\n"; - await reportReady(stub, [{ position: 0, repoOwner: "acme", repoName: "web-app", baseSha }]); + await reportReady(stub, sessionName, auth, [ + { position: 0, repoOwner: "acme", repoName: "web-app", baseSha }, + ]); const upload = await SELF.fetch(`https://test.local/sessions/${sessionName}/diff`, { method: "PUT", @@ -224,7 +234,9 @@ describe("session diff routes", () => { const auth = { authToken: "diff-failure-token", sandboxId: "sandbox-failure" }; await seedSandboxAuth(stub, auth); const baseSha = "c".repeat(40); - await reportReady(stub, [{ position: 0, repoOwner: "acme", repoName: "web-app", baseSha }]); + await reportReady(stub, sessionName, auth, [ + { position: 0, repoOwner: "acme", repoName: "web-app", baseSha }, + ]); const headers = { Authorization: `Bearer ${auth.authToken}`, "Content-Type": "application/json", @@ -259,7 +271,9 @@ describe("session diff routes", () => { const auth = { authToken: "diff-membership-token", sandboxId: "sandbox-membership" }; await seedSandboxAuth(stub, auth); const baseSha = "d".repeat(40); - await reportReady(stub, [{ position: 0, repoOwner: "acme", repoName: "web-app", baseSha }]); + await reportReady(stub, sessionName, auth, [ + { position: 0, repoOwner: "acme", repoName: "web-app", baseSha }, + ]); const headers = { Authorization: `Bearer ${auth.authToken}`, "Content-Type": "application/json", @@ -290,7 +304,7 @@ describe("session diff routes", () => { await seedSandboxAuth(stub, auth); const firstSha = "1".repeat(40); const secondSha = "2".repeat(40); - await reportReady(stub, [ + await reportReady(stub, sessionName, auth, [ { position: 0, repoOwner: "acme", repoName: "web-app", baseSha: firstSha }, { position: 1, repoOwner: "group/subgroup", repoName: "api", baseSha: secondSha }, ]); @@ -333,6 +347,14 @@ describe("session diff routes", () => { const { ws } = await openSandboxWs(sessionName, auth); expect(ws).not.toBeNull(); ws!.accept(); + ws!.send( + JSON.stringify({ + type: "ready", + sandboxId: auth.sandboxId, + timestamp: Date.now() / 1000, + }) + ); + await waitForSandboxStatus(stub, "ready"); const messages = collectMessages(ws!, { until: (message) => message.type === "refresh_diff", timeoutMs: 2_000, diff --git a/packages/control-plane/test/integration/session-lifecycle-alarm-recovery.test.ts b/packages/control-plane/test/integration/session-lifecycle-alarm-recovery.test.ts index 185b042d3a..af9710d55c 100644 --- a/packages/control-plane/test/integration/session-lifecycle-alarm-recovery.test.ts +++ b/packages/control-plane/test/integration/session-lifecycle-alarm-recovery.test.ts @@ -16,11 +16,14 @@ const CONNECTING_TIMEOUT_BUFFER_MS = 1_000; async function parkSandboxPastConnectingTimeout(stub: DurableObjectStub): Promise { await waitForSandboxStatus(stub, "failed"); await runInSessionDO(stub, (instance: SessionDO, state) => { + const createdAt = + Date.now() - + (DEFAULT_LIFECYCLE_CONFIG.connectingTimeout.timeoutMs + CONNECTING_TIMEOUT_BUFFER_MS); state.storage.sql.exec( // modal_object_id stays null, so terminating never calls the provider. - "UPDATE sandbox SET status = 'connecting', modal_object_id = NULL, created_at = ?", - Date.now() - - (DEFAULT_LIFECYCLE_CONFIG.connectingTimeout.timeoutMs + CONNECTING_TIMEOUT_BUFFER_MS) + `UPDATE sandbox SET status = 'connecting', modal_object_id = NULL, + created_at = ?, last_heartbeat = NULL`, + createdAt ); }); } diff --git a/packages/control-plane/test/integration/stop-execution.test.ts b/packages/control-plane/test/integration/stop-execution.test.ts index 214b44a933..3696200db3 100644 --- a/packages/control-plane/test/integration/stop-execution.test.ts +++ b/packages/control-plane/test/integration/stop-execution.test.ts @@ -8,6 +8,7 @@ import { openSandboxWs, seedSandboxAuth, collectMessages, + waitForSandboxStatus, } from "./helpers"; describe("POST /internal/stop", () => { @@ -298,7 +299,17 @@ describe("POST /internal/stop", () => { // Connect sandbox WS so queue drain can dispatch const { ws: sandboxWs } = await openSandboxWs(name, sandboxAuth); - if (sandboxWs) sandboxWs.accept(); + if (sandboxWs) { + sandboxWs.accept(); + sandboxWs.send( + JSON.stringify({ + type: "ready", + sandboxId: sandboxAuth.sandboxId, + timestamp: Date.now() / 1000, + }) + ); + await waitForSandboxStatus(stub, "ready"); + } // Stop execution - marks A as failed await stub.fetch("http://internal/internal/stop", { method: "POST" }); diff --git a/packages/control-plane/test/integration/websocket-sandbox.test.ts b/packages/control-plane/test/integration/websocket-sandbox.test.ts index f22080f601..2fb85ee549 100644 --- a/packages/control-plane/test/integration/websocket-sandbox.test.ts +++ b/packages/control-plane/test/integration/websocket-sandbox.test.ts @@ -3,6 +3,10 @@ import { env } from "cloudflare:test"; import type { SessionDO } from "../../src/session/durable-object"; import { componentsOf, runInSessionDO } from "./session-do-access"; import { encryptToken } from "../../src/auth/crypto"; +import type { Logger } from "../../src/logger"; +import { SessionWebSocketManagerImpl } from "../../src/session/websocket-manager"; +import { WsClientMappingRepository } from "../../src/session/ws-client-mapping-repository"; +import { DEFAULT_HEARTBEAT_CONFIG } from "../../src/sandbox/lifecycle/decisions"; import { collectMessages, initNamedSession, @@ -16,6 +20,16 @@ import { const SANDBOX_TOKEN = "test-sandbox-auth-token-abc123"; const SANDBOX_ID = "sb-integration-test"; +function announceReady(ws: WebSocket): void { + ws.send( + JSON.stringify({ + type: "ready", + sandboxId: SANDBOX_ID, + timestamp: Date.now() / 1000, + }) + ); +} + describe("Sandbox WebSocket (via SELF.fetch)", () => { it("upgrade with valid auth returns 101", async () => { const name = `ws-sandbox-ok-${Date.now()}`; @@ -80,6 +94,116 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { expect(ws).toBeNull(); }); + it("rejects reconnects while snapshotting and requires ready after completion", async () => { + const name = `ws-sandbox-snapshot-reconnect-${Date.now()}`; + const { stub } = await initNamedSession(name); + await seedSandboxAuth(stub, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + status: "snapshotting", + }); + + const rejected = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + expect(rejected.response.status).toBe(503); + expect(rejected.ws).toBeNull(); + expect(await queryDO<{ status: string }>(stub, "SELECT status FROM sandbox")).toEqual([ + { status: "snapshotting" }, + ]); + + await runInSessionDO(stub, (instance: SessionDO) => { + componentsOf(instance).lifecycleManager.spawnSandbox = vi.fn(async () => {}); + }); + const promptResponse = await stub.fetch("http://internal/internal/prompt", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + content: "wait for snapshot reconnect", + authorId: "user-1", + source: "web", + }), + }); + const { messageId } = await promptResponse.json<{ messageId: string }>(); + expect( + await queryDO<{ status: string }>(stub, "SELECT status FROM messages WHERE id = ?", messageId) + ).toEqual([{ status: "pending" }]); + expect( + await runInSessionDO( + stub, + (instance: SessionDO) => + vi.mocked(componentsOf(instance).lifecycleManager.spawnSandbox).mock.calls.length + ) + ).toBe(0); + + await runInSessionDO(stub, async (instance: SessionDO, state) => { + state.storage.sql.exec( + "UPDATE sandbox SET status = 'ready', modal_object_id = 'snapshot-object'" + ); + const components = componentsOf(instance); + await components.lifecycleManager.triggerSnapshot("integration_test", () => + components.messageQueue.processMessageQueue() + ); + }); + expect( + await runInSessionDO( + stub, + (instance: SessionDO) => + vi.mocked(componentsOf(instance).lifecycleManager.spawnSandbox).mock.calls.length + ) + ).toBe(0); + const admitted = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + expect(admitted.response.status).toBe(101); + admitted.ws!.accept(); + expect(await queryDO<{ status: string }>(stub, "SELECT status FROM sandbox")).toEqual([ + { status: "connecting" }, + ]); + + const prompt = collectMessages(admitted.ws!, { + until: (message) => message.type === "prompt", + }); + announceReady(admitted.ws!); + expect(await prompt).toContainEqual(expect.objectContaining({ type: "prompt", messageId })); + await waitForSandboxStatus(stub, "ready"); + admitted.ws!.close(); + }); + + it("drains snapshot completion when the existing control socket remains connected", async () => { + const name = `ws-sandbox-connected-snapshot-${Date.now()}`; + const { stub } = await initNamedSession(name); + await seedSandboxAuth(stub, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + status: "ready", + }); + const { ws } = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + ws!.accept(); + announceReady(ws!); + await waitForSandboxStatus(stub, "ready"); + await runInSessionDO(stub, (_instance: SessionDO, state) => { + state.storage.sql.exec("UPDATE sandbox SET modal_object_id = 'snapshot-object'"); + }); + + const processCalls = await runInSessionDO(stub, async (instance: SessionDO) => { + const processMessageQueue = vi.fn(async () => {}); + await componentsOf(instance).lifecycleManager.triggerSnapshot( + "integration_test", + processMessageQueue + ); + return processMessageQueue.mock.calls.length; + }); + + expect(processCalls).toBe(1); + ws!.close(); + }); + it.each(["archived", "cancelled"] as const)( "upgrade for %s session returns 410", async (status) => { @@ -126,6 +250,7 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { expect(response.status).toBe(101); expect(ws).not.toBeNull(); ws!.accept(); + announceReady(ws!); await waitForSandboxStatus(stub, "ready"); ws!.close(); } @@ -268,11 +393,9 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { expect(ws).toBeNull(); }); - it("sandbox connect sets status to ready", async () => { + it("ready initializes persisted liveness, activity, and the heartbeat deadline", async () => { const name = `ws-sandbox-ready-${Date.now()}`; const { stub } = await initNamedSession(name); - // Model the production boot sequence: the sandbox connects while the - // lifecycle is still in "connecting", and the WS accept flips it to ready. await seedSandboxAuth(stub, { authToken: SANDBOX_TOKEN, sandboxId: SANDBOX_ID, @@ -285,15 +408,272 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { }); expect(ws).not.toBeNull(); ws!.accept(); - await waitForSandboxStatus(stub, "ready"); const stateRes = await stub.fetch("http://internal/internal/state"); const state = await stateRes.json<{ sandbox: { status: string } }>(); - expect(state.sandbox.status).toBe("ready"); + expect(state.sandbox.status).toBe("connecting"); + await runInSessionDO(stub, (instance: SessionDO, durableState) => + durableState.storage.deleteAlarm() + ); + + const readyAt = Date.now(); + announceReady(ws!); + await waitForSandboxStatus(stub, "ready"); + const [sandbox] = await queryDO<{ last_heartbeat: number; last_activity: number }>( + stub, + "SELECT last_heartbeat, last_activity FROM sandbox" + ); + expect(sandbox.last_heartbeat).toBeGreaterThanOrEqual(readyAt); + expect(sandbox.last_activity).toBeGreaterThanOrEqual(readyAt); + const alarm = await runInSessionDO(stub, (instance: SessionDO, durableState) => + durableState.storage.getAlarm() + ); + expect(alarm).not.toBeNull(); + expect(alarm!).toBeGreaterThan(readyAt); + expect(alarm!).toBeLessThan(readyAt + 5 * 60 * 1000); + + ws!.close(); + }); + + it("keeps an early prompt pending and dispatches it only after ready", async () => { + const name = `ws-sandbox-early-prompt-${Date.now()}`; + const { stub } = await initNamedSession(name); + await seedSandboxAuth(stub, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + status: "connecting", + }); + const { ws } = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + expect(ws).not.toBeNull(); + ws!.accept(); + + const response = await stub.fetch("http://internal/internal/prompt", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ content: "wait for ready", authorId: "user-1", source: "web" }), + }); + const { messageId } = await response.json<{ messageId: string }>(); + await new Promise((resolve) => setTimeout(resolve, 50)); + expect( + await queryDO<{ status: string }>(stub, "SELECT status FROM messages WHERE id = ?", messageId) + ).toEqual([{ status: "pending" }]); + + const prompt = collectMessages(ws!, { + until: (message) => message.type === "prompt", + }); + announceReady(ws!); + + expect((await prompt).find((message) => message.type === "prompt")).toEqual( + expect.objectContaining({ type: "prompt", messageId }) + ); + expect( + await queryDO<{ status: string }>(stub, "SELECT status FROM messages WHERE id = ?", messageId) + ).toEqual([{ status: "processing" }]); + ws!.close(); + }); + + it("does not accept ready from a replaced sender", async () => { + const name = `ws-sandbox-stale-ready-${Date.now()}`; + const { stub } = await initNamedSession(name); + await seedSandboxAuth(stub, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + status: "connecting", + }); + const { ws: firstWs } = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + expect(firstWs).not.toBeNull(); + firstWs!.accept(); + const firstServer = await runInSessionDO( + stub, + (instance: SessionDO, state) => state.getWebSockets("sandbox")[0] + ); + + const { ws: replacementWs } = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + expect(replacementWs).not.toBeNull(); + replacementWs!.accept(); + await runInSessionDO(stub, (instance: SessionDO) => + instance.webSocketMessage( + firstServer, + JSON.stringify({ + type: "ready", + sandboxId: SANDBOX_ID, + timestamp: Date.now() / 1000, + }) + ) + ); + expect(await queryDO<{ status: string }>(stub, "SELECT status FROM sandbox")).toEqual([ + { status: "connecting" }, + ]); + + announceReady(replacementWs!); + await waitForSandboxStatus(stub, "ready"); + replacementWs!.close(); + }); + + it("closes a hibernated sandbox socket when replacement is admitted after eviction", async () => { + const name = `ws-sandbox-evicted-replacement-${Date.now()}`; + const { stub } = await initNamedSession(name); + await seedSandboxAuth(stub, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + status: "connecting", + }); + const { ws: firstWs } = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + expect(firstWs).not.toBeNull(); + firstWs!.accept(); + const firstClosed = new Promise((resolve) => { + firstWs!.addEventListener("close", () => resolve(), { once: true }); + }); + + await runInSessionDO(stub, async (instance: SessionDO, state) => { + const pair = new WebSocketPair(); + // A fresh manager models a rehydrated DO whose socket cache is empty. + const manager = new SessionWebSocketManagerImpl( + state, + componentsOf(instance).sandboxRepository, + new WsClientMappingRepository(state.storage.sql), + { debug() {}, info() {}, warn() {}, error() {}, child() {} } as unknown as Logger, + { authTimeoutMs: 1000 } + ); + manager.acceptAndSetSandboxSocket(pair[1], SANDBOX_ID); + pair[0].accept(); + await instance.webSocketMessage( + pair[1], + JSON.stringify({ + type: "ready", + sandboxId: SANDBOX_ID, + timestamp: Date.now() / 1000, + }) + ); + pair[0].close(); + }); + + await expect(firstClosed).resolves.toBeUndefined(); + await waitForSandboxStatus(stub, "ready"); + }); + + it("keeps reconnect execution gated until ready is reannounced", async () => { + const name = `ws-sandbox-reconnect-gating-${Date.now()}`; + const { stub } = await initNamedSession(name); + await seedSandboxAuth(stub, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + status: "connecting", + }); + const { ws: firstWs } = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + expect(firstWs).not.toBeNull(); + firstWs!.accept(); + announceReady(firstWs!); + await waitForSandboxStatus(stub, "ready"); + firstWs!.close(1001, "reconnect"); + const { ws: replacementWs } = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + expect(replacementWs).not.toBeNull(); + replacementWs!.accept(); + const response = await stub.fetch("http://internal/internal/prompt", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ content: "after reconnect", authorId: "user-1", source: "web" }), + }); + const { messageId } = await response.json<{ messageId: string }>(); + expect( + await queryDO<{ status: string }>(stub, "SELECT status FROM messages WHERE id = ?", messageId) + ).toEqual([{ status: "pending" }]); + + const prompt = collectMessages(replacementWs!, { + until: (message) => message.type === "prompt", + }); + announceReady(replacementWs!); + expect(await prompt).toContainEqual(expect.objectContaining({ type: "prompt", messageId })); + replacementWs!.close(); + }); + + it("keeps a control-only prompt pending before provider-auth validation", async () => { + const name = `ws-sandbox-auth-order-${Date.now()}`; + const { stub } = await initNamedSession(name); + await seedSandboxAuth(stub, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + status: "connecting", + }); + const { ws } = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + expect(ws).not.toBeNull(); + ws!.accept(); + + const response = await stub.fetch("http://internal/internal/prompt", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + content: "wait before auth", + authorId: "user-1", + source: "web", + model: "xai/grok-4.5", + }), + }); + const { messageId } = await response.json<{ messageId: string }>(); + + expect( + await queryDO<{ status: string; error_message: string | null }>( + stub, + "SELECT status, error_message FROM messages WHERE id = ?", + messageId + ) + ).toEqual([{ status: "pending", error_message: null }]); ws!.close(); }); + it("preserves senderless authenticated HTTP heartbeat ingestion", async () => { + const name = `ws-sandbox-http-heartbeat-${Date.now()}`; + const { stub } = await initNamedSession(name); + await seedSandboxAuth(stub, { authToken: SANDBOX_TOKEN, sandboxId: SANDBOX_ID }); + const oldHeartbeat = Date.now() - 60_000; + await runInSessionDO(stub, (instance: SessionDO, state) => { + state.storage.sql.exec("UPDATE sandbox SET last_heartbeat = ?", oldHeartbeat); + }); + + const response = await stub.fetch("http://internal/internal/sandbox-event", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + type: "heartbeat", + sandboxId: SANDBOX_ID, + status: "booting", + timestamp: Date.now() / 1000, + }), + }); + + expect(response.status).toBe(200); + expect( + await queryDO<{ last_heartbeat: number }>(stub, "SELECT last_heartbeat FROM sandbox") + ).toEqual([{ last_heartbeat: expect.any(Number) }]); + const [sandbox] = await queryDO<{ last_heartbeat: number }>( + stub, + "SELECT last_heartbeat FROM sandbox" + ); + expect(sandbox.last_heartbeat).toBeGreaterThan(oldHeartbeat); + }); + it("publishes sandbox access only after it becomes readable", async () => { const name = `ws-sandbox-access-ready-${Date.now()}`; const { stub } = await initNamedSession(name); @@ -331,6 +711,7 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { }); expect(sandboxWs).not.toBeNull(); sandboxWs!.accept(); + announceReady(sandboxWs!); const messages = await collector; expect(messages.slice(-2).map((message) => message.type)).toEqual([ @@ -364,7 +745,6 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { lifecycleManager.providerStartupPending = true; }); const { ws: clientWs } = await openClientWs(name, { subscribe: true }); - const collector = collectMessages(clientWs, { timeoutMs: 100 }); const { ws: firstSandboxWs } = await openSandboxWs(name, { authToken: SANDBOX_TOKEN, @@ -372,6 +752,12 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { }); expect(firstSandboxWs).not.toBeNull(); firstSandboxWs!.accept(); + const firstReady = collectMessages(clientWs, { + until: (message) => message.type === "sandbox_status" && message.status === "ready", + }); + announceReady(firstSandboxWs!); + const firstMessages = await firstReady; + expect(firstMessages).toContainEqual({ type: "sandbox_status", status: "ready" }); const { ws: replacementSandboxWs } = await openSandboxWs(name, { authToken: SANDBOX_TOKEN, @@ -379,11 +765,17 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { }); expect(replacementSandboxWs).not.toBeNull(); replacementSandboxWs!.accept(); + expect(await queryDO<{ status: string }>(stub, "SELECT status FROM sandbox")).toEqual([ + { status: "connecting" }, + ]); + const replacementReady = collectMessages(clientWs, { + until: (message) => message.type === "sandbox_status" && message.status === "ready", + }); + announceReady(replacementSandboxWs!); - const messages = await collector; - expect( - messages.filter((message) => message.type === "sandbox_status" && message.status === "ready") - ).toHaveLength(2); + const replacementMessages = await replacementReady; + expect(replacementMessages).toContainEqual({ type: "sandbox_status", status: "ready" }); + const messages = [...firstMessages, ...replacementMessages]; expect(messages).not.toContainEqual({ type: "sandbox_access_changed" }); replacementSandboxWs!.close(); @@ -407,6 +799,8 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { }); expect(firstWs).not.toBeNull(); firstWs!.accept(); + announceReady(firstWs!); + await waitForSandboxStatus(stub, "ready"); const closed = new Promise((resolve) => { firstWs!.addEventListener("close", () => resolve()); @@ -444,6 +838,8 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { }); expect(firstWs).not.toBeNull(); firstWs!.accept(); + announceReady(firstWs!); + await waitForSandboxStatus(stub, "ready"); const closed = new Promise((resolve) => { firstWs!.addEventListener("close", () => resolve()); @@ -473,11 +869,70 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { await runInSessionDO(stub, (instance: SessionDO) => instance.alarm()); const sandboxAfterAlarm = await queryDO<{ status: string }>(stub, "SELECT status FROM sandbox"); - expect(sandboxAfterAlarm[0].status).toBe("ready"); + expect(sandboxAfterAlarm[0].status).toBe("connecting"); + announceReady(reconnectedWs!); + await waitForSandboxStatus(stub, "ready"); reconnectedWs!.close(); }); + it("renews the heartbeat lease when a heartbeat arrives before alarm delivery", async () => { + const name = `ws-sandbox-renewable-heartbeat-${Date.now()}`; + const { stub } = await initNamedSession(name); + await seedSandboxAuth(stub, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + status: "ready", + }); + const { ws } = await openSandboxWs(name, { + authToken: SANDBOX_TOKEN, + sandboxId: SANDBOX_ID, + }); + ws!.accept(); + announceReady(ws!); + await waitForSandboxStatus(stub, "ready"); + + const beforeHeartbeat = Date.now(); + ws!.send( + JSON.stringify({ + type: "heartbeat", + sandboxId: SANDBOX_ID, + status: "running", + timestamp: beforeHeartbeat / 1000, + }) + ); + await vi.waitFor(async () => { + const [sandbox] = await queryDO<{ last_heartbeat: number }>( + stub, + "SELECT last_heartbeat FROM sandbox" + ); + expect(sandbox.last_heartbeat).toBeGreaterThanOrEqual(beforeHeartbeat); + }); + const [renewed] = await queryDO<{ last_heartbeat: number }>( + stub, + "SELECT last_heartbeat FROM sandbox" + ); + + await runInSessionDO(stub, async (instance: SessionDO, state) => { + await state.storage.deleteAlarm(); + await instance.alarm(); + }); + const nextAlarm = await runInSessionDO(stub, (_instance: SessionDO, state) => + state.storage.getAlarm() + ); + expect(nextAlarm).toBe(renewed.last_heartbeat + DEFAULT_HEARTBEAT_CONFIG.timeoutMs); + + await runInSessionDO(stub, (_instance: SessionDO, state) => { + state.storage.sql.exec( + "UPDATE sandbox SET last_heartbeat = ?", + Date.now() - DEFAULT_HEARTBEAT_CONFIG.timeoutMs - 1 + ); + }); + await runInSessionDO(stub, (instance: SessionDO) => instance.alarm()); + await waitForSandboxStatus(stub, "stale"); + ws!.close(); + }); + it("failed sandbox can reconnect and self-heal to ready", async () => { const name = `ws-sandbox-selfheal-${Date.now()}`; const { stub } = await initNamedSession(name); @@ -496,6 +951,7 @@ describe("Sandbox WebSocket (via SELF.fetch)", () => { }); expect(ws).not.toBeNull(); ws!.accept(); + announceReady(ws!); await waitForSandboxStatus(stub, "ready"); ws!.close(); }); diff --git a/packages/sandbox-runtime/src/sandbox_runtime/bridge.py b/packages/sandbox-runtime/src/sandbox_runtime/bridge.py index fdb3f511b3..b5666bf99a 100644 --- a/packages/sandbox-runtime/src/sandbox_runtime/bridge.py +++ b/packages/sandbox-runtime/src/sandbox_runtime/bridge.py @@ -173,6 +173,8 @@ class AgentBridge: GIT_PUSH_TIMEOUT_SECONDS = 300.0 GIT_PUSH_TERMINATE_GRACE_SECONDS = 5.0 DIFF_REFRESH_SHUTDOWN_TIMEOUT_SECONDS = 5.0 + EXECUTION_HEALTH_POLL_INTERVAL_SECONDS = 0.5 + EXECUTION_UNAVAILABLE_MESSAGE = "Sandbox execution is unavailable while booting" def __init__( self, @@ -231,6 +233,9 @@ def __init__( self.ws: ClientConnection | None = None self.shutdown_event = asyncio.Event() self.git_sync_complete = asyncio.Event() + self._ready_event_payload: dict[str, Any] | None = None + self._connection_ready_event = asyncio.Event() + self._send_lock = asyncio.Lock() # Session state self.opencode_session_id: str | None = None @@ -285,8 +290,9 @@ def ws_url(self) -> str: url = self.control_plane_url.replace("https://", "wss://").replace("http://", "ws://") return f"{url}/sessions/{self.session_id}/ws?type=sandbox" - def _build_ready_event(self) -> dict[str, Any]: - repositories = load_repo_manifest(self.repo_manifest_path) + def _build_ready_event(self, repositories: list[Any] | None = None) -> dict[str, Any]: + if repositories is None: + repositories = load_repo_manifest(self.repo_manifest_path) # The image bakes SANDBOX_VERSION; reporting it lets the control plane # stamp snapshots with the runtime that produced them and retire the # ones a later compatibility floor rules out. @@ -325,19 +331,25 @@ async def run(self) -> None: """ self.log.info("bridge.run_start") - await self._load_session_id() reconnect_attempts = 0 run_outcome = "shutdown" - signing_initialized = False + early_connection = os.environ.get("EARLY_SANDBOX_CONNECTION") == "1" + execution_task: asyncio.Task[None] | None = None try: + if early_connection: + execution_task = asyncio.create_task(self._initialize_execution()) + else: + try: + await self._initialize_execution() + except Exception: + run_outcome = "fatal_error" + raise + while not self.shutdown_event.is_set(): run_outcome = "shutdown" try: - if not signing_initialized: - await self.git_signing.initialize(None) - signing_initialized = True - await self._connect_and_run() + await self._connect_and_run(execution_task) if not self.shutdown_event.is_set(): run_outcome = "connection_closed" reconnect_attempts = 0 @@ -348,6 +360,11 @@ async def run(self) -> None: except websockets.ConnectionClosed: run_outcome = "connection_closed" except Exception as e: + if execution_task is not None and execution_task.done(): + initialization_error = execution_task.exception() + if initialization_error is not None: + run_outcome = "fatal_error" + raise initialization_error error_str = str(e) # Check for fatal HTTP errors that shouldn't trigger retry if ( @@ -380,6 +397,13 @@ async def run(self) -> None: await asyncio.sleep(delay) finally: + if execution_task is not None: + if not execution_task.done(): + execution_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await execution_task + elif not execution_task.cancelled(): + execution_task.exception() # Cancel any in-flight prompt task before closing resources if self._current_prompt_task and not self._current_prompt_task.done(): self._current_prompt_task.cancel() @@ -455,7 +479,7 @@ def _is_fatal_connection_error(self, error_str: str) -> bool: ] return any(pattern in error_str for pattern in fatal_patterns) - async def _connect_and_run(self) -> None: + async def _connect_and_run(self, execution_task: asyncio.Task[None] | None = None) -> None: """Connect to control plane and handle messages. Raises: @@ -475,8 +499,11 @@ async def _connect_and_run(self) -> None: ping_timeout=10, ) as ws: self.ws = ws + self._connection_ready_event.clear() self._mark_connected() heartbeat_task: asyncio.Task[None] | None = None + readiness_task: asyncio.Task[None] | None = None + receiver_task: asyncio.Task[None] | None = None background_tasks: set[asyncio.Task[None]] = set() try: @@ -488,24 +515,31 @@ async def _connect_and_run(self) -> None: reconnect_attempt_count=self._reconnect_attempt_count, ) await self.event_forwarder.bind(ws) - await self._send_event(self._build_ready_event()) - await self._drain_boot_warnings() + if execution_task is None: + await self._announce_ready(ws, None) + else: + readiness_task = asyncio.create_task( + self._announce_ready(ws, execution_task) + ) heartbeat_task = asyncio.create_task(self._heartbeat_loop()) - async for message in ws: - if self.shutdown_event.is_set(): - break - - try: - cmd = json.loads(message) - task = await self._handle_command(cmd) - if task: - background_tasks.add(task) - task.add_done_callback(background_tasks.discard) - except json.JSONDecodeError as e: - self.log.warn("bridge.invalid_message", exc=e) - except Exception as e: - self.log.error("bridge.command_error", exc=e) + receiver_task = asyncio.create_task( + self._receive_commands(ws, background_tasks) + ) + + if readiness_task is None: + await receiver_task + else: + done, _pending = await asyncio.wait( + {readiness_task, receiver_task}, + return_when=asyncio.FIRST_COMPLETED, + ) + if readiness_task in done: + readiness_task.result() + readiness_task = None + await receiver_task + else: + receiver_task.result() except websockets.ConnectionClosed as e: self._log_disconnect( @@ -518,8 +552,22 @@ async def _connect_and_run(self) -> None: finally: if heartbeat_task is not None: heartbeat_task.cancel() + if readiness_task is not None: + readiness_task.cancel() + if receiver_task is not None: + receiver_task.cancel() for task in background_tasks: task.cancel() + await asyncio.gather( + *( + task + for task in (heartbeat_task, readiness_task, receiver_task) + if task is not None + ), + *background_tasks, + return_exceptions=True, + ) + self._connection_ready_event.clear() self.ws = None self.event_forwarder.unbind() if self._connected_at_monotonic is not None: @@ -543,20 +591,81 @@ async def _connect_and_run(self) -> None: ) from e raise + async def _receive_commands( + self, + ws: ClientConnection, + background_tasks: set[asyncio.Task[None]], + ) -> None: + async for message in ws: + if self.shutdown_event.is_set(): + break + + try: + cmd = json.loads(message) + task = await self._handle_command(cmd) + if task: + background_tasks.add(task) + task.add_done_callback(background_tasks.discard) + except json.JSONDecodeError as e: + self.log.warn("bridge.invalid_message", exc=e) + except Exception as e: + self.log.error("bridge.command_error", exc=e) + async def _heartbeat_loop(self) -> None: """Send periodic heartbeat events.""" while not self.shutdown_event.is_set(): - await asyncio.sleep(self.HEARTBEAT_INTERVAL) - if self.ws and self.ws.state == State.OPEN: await self._send_event( { "type": "heartbeat", "sandboxId": self.sandbox_id, - "status": "ready", + "status": "ready" if self._connection_ready_event.is_set() else "booting", "timestamp": time.time(), } ) + await asyncio.sleep(self.HEARTBEAT_INTERVAL) + + async def _initialize_execution(self) -> None: + while not self.shutdown_event.is_set(): + if await self.opencode_client.is_healthy(): + break + await asyncio.sleep(self.EXECUTION_HEALTH_POLL_INTERVAL_SECONDS) + if self.shutdown_event.is_set(): + return + + await self._load_session_id() + repositories = load_repo_manifest(self.repo_manifest_path) + while True: + try: + await self.git_signing.initialize(None) + break + except GitSigningError as error: + if not error.retryable: + raise + await asyncio.sleep(self.RECONNECT_BACKOFF_BASE) + + self._ready_event_payload = self._build_ready_event(repositories) + + async def _announce_ready( + self, + ws: ClientConnection, + execution_task: asyncio.Task[None] | None, + ) -> None: + if execution_task is not None: + await asyncio.shield(execution_task) + if self.shutdown_event.is_set() or self._ready_event_payload is None: + raise RuntimeError("Execution initialization completed without a ready payload") + ready_event = { + **self._ready_event_payload, + "sandboxId": self.sandbox_id, + "timestamp": time.time(), + } + async with self._send_lock: + if self.ws is not ws or ws.state != State.OPEN: + raise RuntimeError("Connection closed before ready could be sent") + await ws.send(json.dumps(ready_event)) + self._connection_ready_event.set() + await self._drain_boot_warnings() async def _drain_boot_warnings(self) -> None: """Forward supervisor boot warnings queued before the bridge existed. @@ -570,7 +679,6 @@ async def _drain_boot_warnings(self) -> None: return try: lines = path.read_text().splitlines() - path.unlink(missing_ok=True) except Exception as e: self.log.warn("bridge.boot_warnings_read_failed", exc=e) return @@ -586,6 +694,7 @@ async def _drain_boot_warnings(self) -> None: if not isinstance(entry, dict) or not entry.get("message"): continue await self._send_event({"type": "warning", **entry}) + path.unlink(missing_ok=True) async def _send_media_warning(self, message: str) -> None: """Surface non-fatal media handling failures to the user timeline.""" @@ -593,7 +702,8 @@ async def _send_media_warning(self, message: str) -> None: async def _send_event(self, event: dict[str, Any]) -> None: """Send event to control plane, buffering if WS is unavailable.""" - await self.event_forwarder.send(event) + async with self._send_lock: + await self.event_forwarder.send(event) async def _handle_command(self, cmd: dict[str, Any]) -> asyncio.Task[None] | None: """Handle command from control plane. @@ -606,6 +716,22 @@ async def _handle_command(self, cmd: dict[str, Any]) -> asyncio.Task[None] | Non cmd_type = cmd.get("type") self.log.debug("bridge.command_received", cmd_type=cmd_type) + if not self._connection_ready_event.is_set() and cmd_type in { + "prompt", + "push", + "refresh_diff", + "snapshot", + }: + await self._send_event( + { + "type": "error", + "error": self.EXECUTION_UNAVAILABLE_MESSAGE, + "commandType": cmd_type, + **({"messageId": cmd["messageId"]} if cmd.get("messageId") else {}), + } + ) + return None + if cmd_type == "prompt": message_id = cmd.get("messageId") or cmd.get("message_id", "unknown") self.diff_refresh.prompt_started() @@ -840,6 +966,8 @@ async def _handle_shutdown(self) -> None: if self._current_prompt_task and not self._current_prompt_task.done(): self._current_prompt_task.cancel() self.shutdown_event.set() + if self.ws is not None: + await self.ws.close() async def _handle_push(self, cmd: dict[str, Any]) -> None: """Handle push command using provider-generated push spec. diff --git a/packages/sandbox-runtime/src/sandbox_runtime/opencode_client.py b/packages/sandbox-runtime/src/sandbox_runtime/opencode_client.py index bd25fc6d0d..8ff986837b 100644 --- a/packages/sandbox-runtime/src/sandbox_runtime/opencode_client.py +++ b/packages/sandbox-runtime/src/sandbox_runtime/opencode_client.py @@ -94,6 +94,17 @@ async def session_exists(self, opencode_session_id: str) -> bool: ) return response.status_code == 200 + async def is_healthy(self) -> bool: + """Whether the local OpenCode server is accepting requests.""" + try: + response = await self._client().get( + f"{self._base_url}/global/health", + timeout=2.0, + ) + return response.status_code == 200 + except httpx.HTTPError: + return False + @asynccontextmanager async def events( self, *, inactivity_timeout_seconds: float diff --git a/packages/sandbox-runtime/src/sandbox_runtime/supervisor.py b/packages/sandbox-runtime/src/sandbox_runtime/supervisor.py index 4bc8329983..a4294c0a72 100644 --- a/packages/sandbox-runtime/src/sandbox_runtime/supervisor.py +++ b/packages/sandbox-runtime/src/sandbox_runtime/supervisor.py @@ -33,12 +33,21 @@ class ImageBuildExecutionCancelled(Exception): """A handled process signal interrupted image-build work.""" +class InteractiveStartupCancelled(Exception): + """A clean bridge exit interrupted pre-ready startup.""" + + +class InteractiveStartupBridgeFailed(Exception): + """The bridge exhausted its restart budget during pre-ready startup.""" + + class SandboxSupervisor: """Apply lifecycle policy to the composed runtime services.""" MAX_RESTARTS = 5 BACKOFF_BASE = 2.0 BACKOFF_MAX = 60.0 + EARLY_BRIDGE_MONITOR_INTERVAL_SECONDS = 1.0 def __init__( self, @@ -141,8 +150,13 @@ async def _handle_opencode_exit(self, restart_count: int) -> int: ) return restart_count - async def _handle_bridge_exit(self, restart_count: int) -> int: - exit_code = self.agent_bridge.exit_code() + async def _handle_bridge_exit( + self, + restart_count: int, + exit_code: int | None = None, + ) -> int: + if exit_code is None: + exit_code = self.agent_bridge.exit_code() if exit_code is None: return restart_count if exit_code == 0: @@ -251,10 +265,9 @@ async def _handle_desktop_crash(self, restart_count: int) -> int: self.log.warn("vnc.max_restarts", restart_count=restart_count) return restart_count - async def monitor_processes(self) -> None: + async def monitor_processes(self, *, bridge_restarts: int = 0) -> None: """Monitor each concrete process owner with its explicit restart policy.""" opencode_restarts = 0 - bridge_restarts = 0 code_server_restarts = 0 terminal_restarts = 0 desktop_restarts = 0 @@ -325,6 +338,70 @@ async def _run_image_build_execution( f"image build exceeded its {timeout_seconds}-second execution timeout" ) from error + async def _start_interactive_services( + self, + expected_tunnel_ports: list[int], + ) -> RepositoryBootResult: + try: + await self.browser_desktop.start() + except Exception as error: + self.log.warn("vnc.start_failed", exc=error) + await self.browser_desktop.stop() + + boot_result = await self.repository_boot.boot(self.boot_mode, expected_tunnel_ports) + self._repository_boot_result = boot_result + + # Materialization is sandbox-boot work; OpenCode process restarts + # reuse this tree and must not depend on control-plane availability. + if self.managed_skills is not None: + await self.managed_skills.materialize(boot_result.repositories, boot_result.workdir) + + try: + await self.code_server.start(boot_result.workdir) + except Exception as error: + self.log.warn("code_server.start_failed", exc=error) + await self.code_server.stop() + try: + await self.web_terminal.start(boot_result.workdir) + except Exception as error: + self.log.warn("web_terminal.start_failed", exc=error) + await self.web_terminal.stop() + + await self.opencode_server.start(boot_result.repositories, boot_result.workdir) + return boot_result + + async def _monitor_interactive_startup( + self, + startup_task: asyncio.Task[RepositoryBootResult], + ) -> tuple[RepositoryBootResult, int]: + bridge_restarts = 0 + try: + while not startup_task.done(): + await asyncio.wait( + {startup_task}, + timeout=self.EARLY_BRIDGE_MONITOR_INTERVAL_SECONDS, + ) + if startup_task.done(): + break + if self.shutdown_event.is_set(): + raise InteractiveStartupCancelled + exit_code = self.agent_bridge.exit_code() + if exit_code is None: + continue + bridge_restarts = await self._handle_bridge_exit( + bridge_restarts, + exit_code, + ) + if self.shutdown_event.is_set(): + if exit_code == 0: + raise InteractiveStartupCancelled + raise InteractiveStartupBridgeFailed + return await startup_task, bridge_restarts + finally: + if not startup_task.done(): + startup_task.cancel() + await asyncio.gather(startup_task, return_exceptions=True) + async def run(self, repo_image_callback: RepoImageBuildCallback | None = None) -> bool: startup_start = time.time() self.boot_mode = BootMode.from_env(os.environ) @@ -354,6 +431,8 @@ async def run(self, repo_image_callback: RepoImageBuildCallback | None = None) - Path(BOOT_WARNINGS_FILE_PATH).unlink(missing_ok=True) opencode_ready = False + early_connection = os.environ.get("EARLY_SANDBOX_CONNECTION") == "1" + bridge_restarts = 0 try: if self.boot_mode is BootMode.BUILD: boot_result = await self._run_image_build_execution(expected_tunnel_ports) @@ -378,34 +457,17 @@ async def run(self, repo_image_callback: RepoImageBuildCallback | None = None) - await self.shutdown_event.wait() return True - try: - await self.browser_desktop.start() - except Exception as error: - self.log.warn("vnc.start_failed", exc=error) - await self.browser_desktop.stop() - - boot_result = await self.repository_boot.boot(self.boot_mode, expected_tunnel_ports) - self._repository_boot_result = boot_result - - # Materialization is sandbox-boot work; OpenCode process restarts - # reuse this tree and must not depend on control-plane availability. - if self.managed_skills is not None: - await self.managed_skills.materialize(boot_result.repositories, boot_result.workdir) - - try: - await self.code_server.start(boot_result.workdir) - except Exception as error: - self.log.warn("code_server.start_failed", exc=error) - await self.code_server.stop() - try: - await self.web_terminal.start(boot_result.workdir) - except Exception as error: - self.log.warn("web_terminal.start_failed", exc=error) - await self.web_terminal.stop() - - await self.opencode_server.start(boot_result.repositories, boot_result.workdir) + if early_connection: + await self.agent_bridge.start() + startup_task = asyncio.create_task( + self._start_interactive_services(expected_tunnel_ports) + ) + boot_result, bridge_restarts = await self._monitor_interactive_startup(startup_task) + else: + boot_result = await self._start_interactive_services(expected_tunnel_ports) opencode_ready = True - await self.agent_bridge.start() + if not early_connection: + await self.agent_bridge.start() self.log.info( "sandbox.startup", repo_owner=self.config.repo_owner, @@ -420,10 +482,15 @@ async def run(self, repo_image_callback: RepoImageBuildCallback | None = None) - duration_ms=int((time.time() - startup_start) * 1000), outcome="success", ) - await self.monitor_processes() + await self.monitor_processes(bridge_restarts=bridge_restarts) except ImageBuildExecutionCancelled: self.log.info("image_build.cancelled", reason="shutdown_requested") return True + except InteractiveStartupCancelled: + self.log.info("supervisor.startup_cancelled", reason="bridge_graceful_exit") + return True + except InteractiveStartupBridgeFailed: + return False except Exception as error: self.log.error("supervisor.error", exc=error) if self.boot_mode is BootMode.BUILD and self.shutdown_event.is_set(): diff --git a/packages/sandbox-runtime/tests/test_bridge_boot_warnings.py b/packages/sandbox-runtime/tests/test_bridge_boot_warnings.py index d2eb6cc4df..4271ce738e 100644 --- a/packages/sandbox-runtime/tests/test_bridge_boot_warnings.py +++ b/packages/sandbox-runtime/tests/test_bridge_boot_warnings.py @@ -5,6 +5,7 @@ consumes the file exactly once. """ +import asyncio import json from pathlib import Path from unittest.mock import AsyncMock, patch @@ -66,3 +67,19 @@ async def test_drain_is_a_noop_without_file(tmp_path: Path): await bridge._drain_boot_warnings() bridge._send_event.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_cancelled_warning_send_preserves_warning_file(tmp_path: Path): + warnings_file = tmp_path / "warnings.jsonl" + warnings_file.write_text(json.dumps({"scope": "setup", "message": "setup failed"}) + "\n") + bridge = _create_bridge() + bridge._send_event = AsyncMock(side_effect=asyncio.CancelledError) + + with ( + patch("sandbox_runtime.bridge.BOOT_WARNINGS_FILE_PATH", str(warnings_file)), + pytest.raises(asyncio.CancelledError), + ): + await bridge._drain_boot_warnings() + + assert warnings_file.exists() diff --git a/packages/sandbox-runtime/tests/test_bridge_diff_capture.py b/packages/sandbox-runtime/tests/test_bridge_diff_capture.py index a3f359e921..7e51a186f7 100644 --- a/packages/sandbox-runtime/tests/test_bridge_diff_capture.py +++ b/packages/sandbox-runtime/tests/test_bridge_diff_capture.py @@ -18,12 +18,14 @@ def _bridge() -> AgentBridge: - return AgentBridge( + bridge = AgentBridge( sandbox_id="sandbox-1", session_id="session-1", control_plane_url="https://control.example.com", auth_token="sandbox-token", ) + bridge._connection_ready_event.set() + return bridge def _manifest(tmp_path: Path) -> Path: diff --git a/packages/sandbox-runtime/tests/test_bridge_early_connection.py b/packages/sandbox-runtime/tests/test_bridge_early_connection.py new file mode 100644 index 0000000000..f19a0fd452 --- /dev/null +++ b/packages/sandbox-runtime/tests/test_bridge_early_connection.py @@ -0,0 +1,282 @@ +import asyncio +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest +from websockets import State + +from sandbox_runtime.bridge import AgentBridge + + +def _bridge() -> AgentBridge: + bridge = AgentBridge( + sandbox_id="sandbox-1", + session_id="session-1", + control_plane_url="https://control.example", + auth_token="token", + ) + bridge.log = MagicMock() + return bridge + + +async def test_execution_initialization_waits_for_health_before_session_and_signing(): + bridge = _bridge() + first_health_check = asyncio.Event() + allow_health = asyncio.Event() + + async def is_healthy() -> bool: + first_health_check.set() + await allow_health.wait() + return True + + bridge.opencode_client.is_healthy = is_healthy + bridge._load_session_id = AsyncMock() + bridge.git_signing.initialize = AsyncMock() + bridge._build_ready_event = MagicMock(return_value={"type": "ready"}) + + task = asyncio.create_task(bridge._initialize_execution()) + await first_health_check.wait() + + bridge._load_session_id.assert_not_awaited() + bridge.git_signing.initialize.assert_not_awaited() + bridge._build_ready_event.assert_not_called() + + allow_health.set() + await task + + bridge._load_session_id.assert_awaited_once() + bridge.git_signing.initialize.assert_awaited_once_with(None) + assert bridge._ready_event_payload == {"type": "ready"} + + +async def test_heartbeat_reports_booting_until_ready_is_announced(monkeypatch): + bridge = _bridge() + bridge.ws = MagicMock(state=State.OPEN) + sent: list[dict] = [] + bridge._send_event = AsyncMock(side_effect=lambda event: sent.append(event.copy())) + + async def one_iteration(_delay: float) -> None: + if sent: + bridge.shutdown_event.set() + + monkeypatch.setattr("sandbox_runtime.bridge.asyncio.sleep", one_iteration) + await bridge._heartbeat_loop() + + assert sent[0]["status"] == "booting" + + bridge.shutdown_event.clear() + bridge._connection_ready_event.set() + sent.clear() + await bridge._heartbeat_loop() + + assert sent[0]["status"] == "ready" + + +class _FakeWs: + close_code = 1000 + + def __init__(self, messages: list[dict] | None = None, *, idle: bool = False) -> None: + self.state = State.OPEN + self.sent: list[str] = [] + self.messages = [json.dumps(message) for message in messages or []] + self.idle = idle + self.closed = asyncio.Event() + + async def send(self, data: str) -> None: + self.sent.append(data) + + async def close(self) -> None: + self.state = State.CLOSED + self.closed.set() + + def __aiter__(self): + return self + + async def __anext__(self): + if self.messages: + return self.messages.pop(0) + if self.idle and not self.closed.is_set(): + await self.closed.wait() + await asyncio.sleep(0) + raise StopAsyncIteration + + +class _Connection: + def __init__(self, ws: _FakeWs) -> None: + self.ws = ws + + async def __aenter__(self) -> _FakeWs: + return self.ws + + async def __aexit__(self, *_args) -> bool: + return False + + +async def test_late_mode_initializes_before_connecting(monkeypatch): + bridge = _bridge() + order: list[str] = [] + + async def initialize() -> None: + order.append("initialize") + bridge._ready_event_payload = {"type": "ready"} + + async def connect(*_args) -> None: + order.append("connect") + bridge.shutdown_event.set() + + bridge._initialize_execution = initialize + bridge._connect_and_run = connect + monkeypatch.delenv("EARLY_SANDBOX_CONNECTION", raising=False) + + await bridge.run() + + assert order == ["initialize", "connect"] + + +async def test_late_mode_sends_ready_before_immediate_execution_command(monkeypatch): + bridge = _bridge() + ws = _FakeWs([{"type": "refresh_diff"}, {"type": "shutdown"}], idle=True) + bridge.diff_refresh.request = MagicMock() + bridge._drain_boot_warnings = AsyncMock() + + async def initialize() -> None: + bridge._ready_event_payload = {"type": "ready", "sandboxId": bridge.sandbox_id} + + bridge._initialize_execution = initialize + monkeypatch.delenv("EARLY_SANDBOX_CONNECTION", raising=False) + monkeypatch.setattr( + "sandbox_runtime.bridge.websockets.connect", + lambda *_args, **_kwargs: _Connection(ws), + ) + + await bridge.run() + + assert json.loads(ws.sent[0])["type"] == "ready" + bridge.diff_refresh.request.assert_called_once_with(None) + + +@pytest.mark.parametrize("early", [False, True]) +async def test_execution_initialization_failure_exits_nonzero(monkeypatch, early: bool): + bridge = _bridge() + ws = _FakeWs(idle=True) + bridge._initialize_execution = AsyncMock(side_effect=RuntimeError("init failed")) + monkeypatch.setenv("EARLY_SANDBOX_CONNECTION", "1" if early else "0") + monkeypatch.setattr( + "sandbox_runtime.bridge.websockets.connect", + lambda *_args, **_kwargs: _Connection(ws), + ) + + with pytest.raises(RuntimeError, match="init failed"): + await asyncio.wait_for(bridge.run(), timeout=1) + + +async def test_shutdown_command_closes_idle_connection(monkeypatch): + bridge = _bridge() + bridge._ready_event_payload = {"type": "ready"} + ws = _FakeWs([{"type": "shutdown"}], idle=True) + monkeypatch.setattr( + "sandbox_runtime.bridge.websockets.connect", + lambda *_args, **_kwargs: _Connection(ws), + ) + + await asyncio.wait_for(bridge._connect_and_run(None), timeout=1) + + assert bridge.shutdown_event.is_set() + assert ws.closed.is_set() + + +async def test_initialization_completes_on_same_open_socket(monkeypatch): + bridge = _bridge() + allow_initialization = asyncio.Event() + ws = _FakeWs(idle=True) + + async def initialize() -> None: + await allow_initialization.wait() + bridge._ready_event_payload = {"type": "ready", "sandboxId": bridge.sandbox_id} + + initialization_task = asyncio.create_task(initialize()) + monkeypatch.setattr( + "sandbox_runtime.bridge.websockets.connect", + lambda *_args, **_kwargs: _Connection(ws), + ) + connection_task = asyncio.create_task(bridge._connect_and_run(initialization_task)) + await asyncio.sleep(0) + + assert not [event for event in map(json.loads, ws.sent) if event["type"] == "ready"] + + allow_initialization.set() + for _ in range(10): + await asyncio.sleep(0) + if any(event["type"] == "ready" for event in map(json.loads, ws.sent)): + break + + ready_events = [event for event in map(json.loads, ws.sent) if event["type"] == "ready"] + assert len(ready_events) == 1 + await ws.close() + await connection_task + + +async def test_ready_is_resent_on_reconnect(monkeypatch): + bridge = _bridge() + bridge._ready_event_payload = {"type": "ready", "sandboxId": bridge.sandbox_id} + first = _FakeWs() + second = _FakeWs() + sockets = iter([first, second]) + monkeypatch.setattr( + "sandbox_runtime.bridge.websockets.connect", + lambda *_args, **_kwargs: _Connection(next(sockets)), + ) + + await bridge._connect_and_run(None) + await bridge._connect_and_run(None) + + assert [json.loads(event)["type"] for event in first.sent].count("ready") == 1 + assert [json.loads(event)["type"] for event in second.sent].count("ready") == 1 + + +async def test_cancelled_ready_send_keeps_connection_execution_unready(): + bridge = _bridge() + bridge._ready_event_payload = {"type": "ready"} + ws = _FakeWs() + ws.send = AsyncMock(side_effect=asyncio.CancelledError) + bridge.ws = ws + + with pytest.raises(asyncio.CancelledError): + await bridge._announce_ready(ws, None) + + assert not bridge._connection_ready_event.is_set() + + +@pytest.mark.parametrize("command_type", ["prompt", "push", "refresh_diff", "snapshot"]) +async def test_execution_commands_are_rejected_before_ready(command_type: str): + bridge = _bridge() + bridge._handle_prompt = AsyncMock() + bridge._handle_push = AsyncMock() + bridge._handle_snapshot = AsyncMock() + bridge.diff_refresh.request = MagicMock() + bridge._send_event = AsyncMock() + + await bridge._handle_command( + {"type": command_type, "messageId": "message-1", "pushSpec": {"targetBranch": "main"}} + ) + + bridge._handle_prompt.assert_not_awaited() + bridge._handle_push.assert_not_awaited() + bridge._handle_snapshot.assert_not_awaited() + bridge.diff_refresh.request.assert_not_called() + bridge._send_event.assert_awaited_once() + + +async def test_control_commands_remain_available_before_ready(): + bridge = _bridge() + bridge._handle_stop = AsyncMock() + bridge._handle_shutdown = AsyncMock() + bridge.event_forwarder.acknowledge = MagicMock(return_value=True) + + await bridge._handle_command({"type": "stop"}) + await bridge._handle_command({"type": "shutdown"}) + await bridge._handle_command({"type": "ack", "ackId": "ack-1"}) + + bridge._handle_stop.assert_awaited_once() + bridge._handle_shutdown.assert_awaited_once() + bridge.event_forwarder.acknowledge.assert_called_once_with("ack-1") diff --git a/packages/sandbox-runtime/tests/test_bridge_event_buffer.py b/packages/sandbox-runtime/tests/test_bridge_event_buffer.py index f82d54af05..f14534c872 100644 --- a/packages/sandbox-runtime/tests/test_bridge_event_buffer.py +++ b/packages/sandbox-runtime/tests/test_bridge_event_buffer.py @@ -84,6 +84,8 @@ def bridge() -> AgentBridge: auth_token="test-token", ) bridge.opencode_session_id = "oc-session-123" + bridge._ready_event_payload = {"type": "ready", "sandboxId": bridge.sandbox_id} + bridge._connection_ready_event.set() wire_opencode_transport(bridge, MockHttpClient()) return bridge diff --git a/packages/sandbox-runtime/tests/test_bridge_reconnection.py b/packages/sandbox-runtime/tests/test_bridge_reconnection.py index 2ed804c181..d839b4d627 100644 --- a/packages/sandbox-runtime/tests/test_bridge_reconnection.py +++ b/packages/sandbox-runtime/tests/test_bridge_reconnection.py @@ -97,7 +97,7 @@ async def __aexit__(self, *_args): lambda *_args, **_kwargs: ConnectionContext(ws), ) bridge.log = MagicMock() - bridge._send_event = AsyncMock(side_effect=asyncio.CancelledError) + bridge.event_forwarder.bind = AsyncMock(side_effect=asyncio.CancelledError) with pytest.raises(asyncio.CancelledError): await bridge._connect_and_run() @@ -118,7 +118,7 @@ async def __aexit__(self, *_args): async def test_run_complete_does_not_retain_transient_outcome(self, bridge, monkeypatch): attempts = 0 - async def connect_and_run(): + async def connect_and_run(*_args): nonlocal attempts attempts += 1 if attempts == 1: @@ -126,8 +126,7 @@ async def connect_and_run(): bridge.shutdown_event.set() bridge.log = MagicMock() - bridge.git_signing.initialize = AsyncMock() - bridge._load_session_id = AsyncMock() + bridge._initialize_execution = AsyncMock() bridge._connect_and_run = connect_and_run monkeypatch.setattr("sandbox_runtime.bridge.asyncio.sleep", AsyncMock()) @@ -143,11 +142,11 @@ async def connect_and_run(): ) @pytest.mark.asyncio - async def test_run_retries_signing_initialization_before_connecting(self, bridge, monkeypatch): - async def connect_and_run(): - bridge.shutdown_event.set() - + async def test_execution_initialization_retries_retryable_signing_failure( + self, bridge, monkeypatch + ): bridge.log = MagicMock() + bridge.opencode_client.is_healthy = AsyncMock(return_value=True) bridge.git_signing.initialize = AsyncMock( side_effect=[ GitSigningError("Commit signing configuration unavailable", retryable=True), @@ -155,60 +154,54 @@ async def connect_and_run(): ] ) bridge._load_session_id = AsyncMock() - bridge._connect_and_run = AsyncMock(side_effect=connect_and_run) + bridge._build_ready_event = MagicMock(return_value={"type": "ready"}) sleep = AsyncMock() monkeypatch.setattr("sandbox_runtime.bridge.asyncio.sleep", sleep) - await bridge.run() + await bridge._initialize_execution() assert bridge.git_signing.initialize.await_count == 2 - bridge._connect_and_run.assert_awaited_once() sleep.assert_awaited_once_with(bridge.RECONNECT_BACKOFF_BASE) + assert bridge._ready_event_payload is not None @pytest.mark.asyncio @pytest.mark.parametrize("status", [401, 403, 404, 410]) - async def test_run_exits_on_terminal_signing_configuration_status( + async def test_execution_initialization_exits_on_terminal_signing_configuration_status( self, bridge, monkeypatch, status ): bridge.log = MagicMock() + bridge.opencode_client.is_healthy = AsyncMock(return_value=True) bridge.git_signing.initialize = AsyncMock( side_effect=GitSigningError( "Commit signing configuration unavailable", status_code=status ) ) bridge._load_session_id = AsyncMock() - bridge._connect_and_run = AsyncMock() sleep = AsyncMock() monkeypatch.setattr("sandbox_runtime.bridge.asyncio.sleep", sleep) - await bridge.run() + with pytest.raises(GitSigningError): + await bridge._initialize_execution() - bridge._connect_and_run.assert_not_awaited() sleep.assert_not_awaited() - assert bridge.shutdown_event.is_set() - bridge.log.info.assert_any_call( - "bridge.run_complete", - outcome="fatal_error", - connection_count=0, - reconnect_count=0, - reconnect_attempt_count=0, - total_connected_duration_seconds=0.0, - ) + assert bridge._ready_event_payload is None @pytest.mark.asyncio - async def test_run_exits_on_nonretryable_payload_failure(self, bridge, monkeypatch): + async def test_execution_initialization_exits_on_nonretryable_payload_failure( + self, bridge, monkeypatch + ): bridge.log = MagicMock() + bridge.opencode_client.is_healthy = AsyncMock(return_value=True) bridge.git_signing.initialize = AsyncMock( side_effect=GitSigningError("Invalid commit signing configuration") ) bridge._load_session_id = AsyncMock() - bridge._connect_and_run = AsyncMock() sleep = AsyncMock() monkeypatch.setattr("sandbox_runtime.bridge.asyncio.sleep", sleep) - await bridge.run() + with pytest.raises(GitSigningError): + await bridge._initialize_execution() - bridge._connect_and_run.assert_not_awaited() sleep.assert_not_awaited() diff --git a/packages/sandbox-runtime/tests/test_bridge_stop.py b/packages/sandbox-runtime/tests/test_bridge_stop.py index 1f0f962a32..0c665c2cf1 100644 --- a/packages/sandbox-runtime/tests/test_bridge_stop.py +++ b/packages/sandbox-runtime/tests/test_bridge_stop.py @@ -80,6 +80,7 @@ def bridge() -> AgentBridge: auth_token="test-token", ) bridge.opencode_session_id = "oc-session-123" + bridge._connection_ready_event.set() wire_opencode_transport(bridge, MockHttpClient()) return bridge diff --git a/packages/sandbox-runtime/tests/test_supervisor_lifecycle.py b/packages/sandbox-runtime/tests/test_supervisor_lifecycle.py index 3460f41c7f..5657a3aadd 100644 --- a/packages/sandbox-runtime/tests/test_supervisor_lifecycle.py +++ b/packages/sandbox-runtime/tests/test_supervisor_lifecycle.py @@ -2,6 +2,8 @@ from pathlib import Path from unittest.mock import AsyncMock, MagicMock +import pytest + from sandbox_runtime.repository_boot import RepositoryBootResult from sandbox_runtime.runtime_config import BootMode, RuntimeConfig from sandbox_runtime.supervisor import SandboxSupervisor @@ -66,6 +68,7 @@ async def test_regular_boot_phase_order(tmp_path, monkeypatch): monkeypatch.delenv("IMAGE_BUILD_MODE", raising=False) monkeypatch.delenv("RESTORED_FROM_SNAPSHOT", raising=False) monkeypatch.delenv("FROM_REPO_IMAGE", raising=False) + monkeypatch.delenv("EARLY_SANDBOX_CONNECTION", raising=False) assert await supervisor.run() is True supervisor.repository_boot.prepare_tunnel_environment.assert_called_once_with(BootMode.FRESH) @@ -80,6 +83,197 @@ async def test_regular_boot_phase_order(tmp_path, monkeypatch): ] +async def test_early_connection_starts_bridge_before_repository_boot(tmp_path, monkeypatch): + events = [] + supervisor, *_ = _supervisor(tmp_path, events) + monkeypatch.setenv("EARLY_SANDBOX_CONNECTION", "1") + monkeypatch.delenv("IMAGE_BUILD_MODE", raising=False) + monkeypatch.delenv("RESTORED_FROM_SNAPSHOT", raising=False) + monkeypatch.delenv("FROM_REPO_IMAGE", raising=False) + + assert await supervisor.run() is True + assert events == [ + "bridge", + "desktop", + "repository:fresh", + "skills", + "code_server", + "terminal", + "opencode", + ] + + +async def test_early_connection_repository_failure_stops_bridge(tmp_path, monkeypatch): + supervisor, repository, _opencode_server, agent_bridge, *_ = _supervisor(tmp_path, []) + repository.boot.side_effect = RuntimeError("clone failed") + supervisor._report_fatal_error = AsyncMock() + monkeypatch.setenv("EARLY_SANDBOX_CONNECTION", "1") + monkeypatch.delenv("IMAGE_BUILD_MODE", raising=False) + + assert await supervisor.run() is False + + agent_bridge.start.assert_awaited_once() + agent_bridge.stop.assert_awaited_once() + supervisor._report_fatal_error.assert_awaited_once_with("clone failed") + + +async def test_early_connection_restarts_bridge_while_repository_boot_is_blocked( + tmp_path, monkeypatch +): + supervisor, repository, _opencode_server, agent_bridge, *_ = _supervisor(tmp_path, []) + release_boot = asyncio.Event() + boot_started = asyncio.Event() + + async def blocked_boot(_mode, _ports): + boot_started.set() + await release_boot.wait() + return RepositoryBootResult(True, [], True, True, (), tmp_path) + + repository.boot.side_effect = blocked_boot + agent_bridge.exit_code.side_effect = [1, None] + + async def restart_bridge(): + if agent_bridge.start.await_count == 2: + release_boot.set() + + agent_bridge.start.side_effect = restart_bridge + monkeypatch.setattr(supervisor, "_wait_for_shutdown", AsyncMock(return_value=False)) + monkeypatch.setattr(supervisor, "EARLY_BRIDGE_MONITOR_INTERVAL_SECONDS", 0.01, raising=False) + monkeypatch.setenv("EARLY_SANDBOX_CONNECTION", "1") + monkeypatch.delenv("IMAGE_BUILD_MODE", raising=False) + + run_task = asyncio.create_task(supervisor.run()) + await boot_started.wait() + + assert await asyncio.wait_for(run_task, timeout=1) is True + assert agent_bridge.start.await_count == 2 + + +@pytest.mark.parametrize("blocked_stage", ["repository", "skills", "opencode"]) +async def test_clean_early_bridge_exit_cancels_pre_ready_startup( + tmp_path, monkeypatch, blocked_stage +): + supervisor, repository, opencode_server, agent_bridge, *_ = _supervisor(tmp_path, []) + stage_started = asyncio.Event() + stage_cancelled = asyncio.Event() + + async def block_stage(*_args): + stage_started.set() + try: + await asyncio.Event().wait() + finally: + stage_cancelled.set() + + if blocked_stage == "repository": + repository.boot.side_effect = block_stage + elif blocked_stage == "skills": + supervisor.managed_skills.materialize.side_effect = block_stage + else: + opencode_server.start.side_effect = block_stage + + agent_bridge.exit_code.return_value = 0 + supervisor._report_fatal_error = AsyncMock() + monkeypatch.setattr(supervisor, "EARLY_BRIDGE_MONITOR_INTERVAL_SECONDS", 0.01) + monkeypatch.setenv("EARLY_SANDBOX_CONNECTION", "1") + monkeypatch.delenv("IMAGE_BUILD_MODE", raising=False) + + run_task = asyncio.create_task(supervisor.run()) + await stage_started.wait() + + assert await asyncio.wait_for(run_task, timeout=1) is True + assert stage_cancelled.is_set() + supervisor._report_fatal_error.assert_not_awaited() + opencode_server.stop.assert_awaited_once() + + +async def test_early_bridge_crash_restarts_during_post_repository_startup(tmp_path, monkeypatch): + supervisor, _repository, _opencode_server, agent_bridge, *_ = _supervisor(tmp_path, []) + release_skills = asyncio.Event() + skills_started = asyncio.Event() + + async def blocked_skills(*_args): + skills_started.set() + await release_skills.wait() + + async def start_bridge(): + if agent_bridge.start.await_count == 2: + release_skills.set() + + supervisor.managed_skills.materialize.side_effect = blocked_skills + agent_bridge.start.side_effect = start_bridge + agent_bridge.exit_code.side_effect = [1, None] + monkeypatch.setattr(supervisor, "_wait_for_shutdown", AsyncMock(return_value=False)) + monkeypatch.setattr(supervisor, "EARLY_BRIDGE_MONITOR_INTERVAL_SECONDS", 0.01) + monkeypatch.setenv("EARLY_SANDBOX_CONNECTION", "1") + monkeypatch.delenv("IMAGE_BUILD_MODE", raising=False) + + run_task = asyncio.create_task(supervisor.run()) + await skills_started.wait() + + assert await asyncio.wait_for(run_task, timeout=1) is True + assert agent_bridge.start.await_count == 2 + + +async def test_early_bridge_restart_budget_carries_into_steady_monitoring(tmp_path, monkeypatch): + supervisor, _repository, _opencode_server, agent_bridge, *_ = _supervisor(tmp_path, []) + release_terminal = asyncio.Event() + terminal_started = asyncio.Event() + + async def blocked_terminal(*_args): + terminal_started.set() + await release_terminal.wait() + + async def start_bridge(): + if agent_bridge.start.await_count == 2: + release_terminal.set() + + supervisor.web_terminal.start.side_effect = blocked_terminal + agent_bridge.start.side_effect = start_bridge + agent_bridge.exit_code.side_effect = [1, None] + monkeypatch.setattr(supervisor, "_wait_for_shutdown", AsyncMock(return_value=False)) + monkeypatch.setattr(supervisor, "EARLY_BRIDGE_MONITOR_INTERVAL_SECONDS", 0.01) + monkeypatch.setenv("EARLY_SANDBOX_CONNECTION", "1") + monkeypatch.delenv("IMAGE_BUILD_MODE", raising=False) + + run_task = asyncio.create_task(supervisor.run()) + await terminal_started.wait() + + assert await asyncio.wait_for(run_task, timeout=1) is True + supervisor.monitor_processes.assert_awaited_once_with(bridge_restarts=1) + + +async def test_early_bridge_restart_exhaustion_cancels_startup_and_reports_fatal( + tmp_path, monkeypatch +): + supervisor, _repository, opencode_server, agent_bridge, *_ = _supervisor(tmp_path, []) + opencode_started = asyncio.Event() + opencode_cancelled = asyncio.Event() + + async def blocked_opencode(*_args): + opencode_started.set() + try: + await asyncio.Event().wait() + finally: + opencode_cancelled.set() + + opencode_server.start.side_effect = blocked_opencode + agent_bridge.exit_code.return_value = 1 + supervisor.MAX_RESTARTS = 1 + supervisor._report_fatal_error = AsyncMock() + monkeypatch.setattr(supervisor, "_wait_for_shutdown", AsyncMock(return_value=False)) + monkeypatch.setattr(supervisor, "EARLY_BRIDGE_MONITOR_INTERVAL_SECONDS", 0.01) + monkeypatch.setenv("EARLY_SANDBOX_CONNECTION", "1") + monkeypatch.delenv("IMAGE_BUILD_MODE", raising=False) + + run_task = asyncio.create_task(supervisor.run()) + await opencode_started.wait() + + assert await asyncio.wait_for(run_task, timeout=1) is False + assert opencode_cancelled.is_set() + assert agent_bridge.start.await_count == 2 + supervisor._report_fatal_error.assert_awaited_once_with("Bridge crashed 2 times, giving up") + + async def test_regular_boot_passes_repository_workspace_to_services(tmp_path, monkeypatch): supervisor, repository, opencode_server, _agent_bridge, code_server, terminal, _desktop = ( _supervisor(tmp_path, []) @@ -105,6 +299,7 @@ async def test_build_boot_excludes_runtime_services(tmp_path, monkeypatch): _supervisor(tmp_path, []) ) monkeypatch.setenv("IMAGE_BUILD_MODE", "true") + monkeypatch.setenv("EARLY_SANDBOX_CONNECTION", "1") callback = MagicMock() async def report_success(**_kwargs): diff --git a/terraform/environments/production/variables.tf b/terraform/environments/production/variables.tf index e04071fe5c..4651e4aab0 100644 --- a/terraform/environments/production/variables.tf +++ b/terraform/environments/production/variables.tf @@ -591,6 +591,12 @@ variable "sandbox_inactivity_timeout_ms" { default = 600000 } +variable "early_sandbox_connection" { + description = "Temporarily connect sandbox runtimes before execution initialization completes." + type = bool + default = false +} + variable "web_platform" { description = "Platform for the web app deployment: 'vercel' or 'cloudflare' (OpenNext)" type = string diff --git a/terraform/environments/production/workers-control-plane.tf b/terraform/environments/production/workers-control-plane.tf index 0bb48918ae..15032b5fea 100644 --- a/terraform/environments/production/workers-control-plane.tf +++ b/terraform/environments/production/workers-control-plane.tf @@ -107,6 +107,7 @@ module "control_plane_worker" { { name = "GITHUB_BOT_USERNAME", value = var.github_bot_username }, { name = "SANDBOX_PROVIDER", value = var.sandbox_provider }, { name = "SANDBOX_INACTIVITY_TIMEOUT_MS", value = tostring(var.sandbox_inactivity_timeout_ms) }, + { name = "EARLY_SANDBOX_CONNECTION", value = var.early_sandbox_connection ? "1" : "0" }, ], local.github_oauth_enabled ? [ { name = "GITHUB_CLIENT_ID", value = trimspace(var.github_client_id) },