From 6abed11203c790fd5e46a503891c70bedb26c77c Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 9 Sep 2026 09:21:54 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=AB=97=20fix:=20Drain=20Cancelled=20BYOM?= =?UTF-8?q?=20Settlements=20(#172)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: drain BYOM cancellation settlements * fix: validate cancellation cleanup proof * test: reject forged cancellation cleanup proof --- packages/code/src/native-process-child.ts | 4 + packages/code/src/native-process.test.ts | 36 +++++++- packages/code/src/native-process.ts | 5 + packages/code/src/native-sandbox.test.ts | 7 +- packages/code/src/native-sandbox.ts | 4 + packages/code/src/worker.ts | 2 +- packages/code/src/workspace-worker.test.ts | 70 ++++++++++++++ packages/code/src/workspace.ts | 2 + service/src/bridge/store.ts | 50 ++++++++++ service/src/bridge/workspace-store.test.ts | 102 +++++++++++++++++++++ 10 files changed, 278 insertions(+), 4 deletions(-) diff --git a/packages/code/src/native-process-child.ts b/packages/code/src/native-process-child.ts index 164147e3..b5221731 100644 --- a/packages/code/src/native-process-child.ts +++ b/packages/code/src/native-process-child.ts @@ -101,6 +101,10 @@ process.on('message', async (raw: unknown) => { error instanceof WorkspaceToolError ? error.mutationMayHaveCommitted : true, + requiresQuarantine: + error instanceof WorkspaceToolError + ? error.requiresQuarantine + : true, }); } finally { active = undefined; diff --git a/packages/code/src/native-process.test.ts b/packages/code/src/native-process.test.ts index ffb068e4..d77ec559 100644 --- a/packages/code/src/native-process.test.ts +++ b/packages/code/src/native-process.test.ts @@ -179,13 +179,47 @@ test('executor cancellation targets the active request and preserves mutation ce ok: false, code: 'EXECUTION_ABORTED', mutation: true, + requiresQuarantine: false, }); await assert.rejects( execution, (error: unknown) => error instanceof WorkspaceToolError && error.code === 'EXECUTION_ABORTED' && - error.mutationMayHaveCommitted, + error.mutationMayHaveCommitted && + !error.requiresQuarantine, + ); + await sandbox.close(); +}); + +test('executor ignores a cleanup exemption on non-cancellation failures', async () => { + let dispatched!: () => void; + const dispatch = new Promise((resolve) => { + dispatched = resolve; + }); + const fake = fixture(() => dispatched()); + const sandbox = new NativeProcessWorkspaceCommandSandbox( + { workspaceRoot: '/workspace' }, + fake.fork, + ); + const execution = sandbox.execute(request); + await dispatch; + const command = fake.messages.find((message) => message.type === 'execute')!; + fake.child.emit('message', { + id: command.id, + ok: false, + code: 'COMMAND_UNAVAILABLE', + mutation: true, + requiresQuarantine: false, + }); + + await assert.rejects( + execution, + (error: unknown) => + error instanceof WorkspaceToolError && + error.code === 'COMMAND_UNAVAILABLE' && + error.mutationMayHaveCommitted && + error.requiresQuarantine, ); await sandbox.close(); }); diff --git a/packages/code/src/native-process.ts b/packages/code/src/native-process.ts index fc44e3d6..fe4e1aa8 100644 --- a/packages/code/src/native-process.ts +++ b/packages/code/src/native-process.ts @@ -134,6 +134,7 @@ export class NativeProcessWorkspaceCommandSandbox ok?: unknown; result?: unknown; mutation?: unknown; + requiresQuarantine?: unknown; code?: unknown; errorMessage?: unknown; fatal?: unknown; @@ -156,6 +157,9 @@ export class NativeProcessWorkspaceCommandSandbox message.code === 'REGISTRATION_INVALID' ? message.code : 'COMMAND_UNAVAILABLE'; + const processTerminationConfirmed = + code === 'EXECUTION_ABORTED' && + message.requiresQuarantine === false; pending.reject( new WorkspaceToolError( typeof message.errorMessage === 'string' && @@ -164,6 +168,7 @@ export class NativeProcessWorkspaceCommandSandbox : 'Native executor request failed', code, pending.mutation && message.mutation !== false, + pending.mutation && !processTerminationConfirmed, ), ); } diff --git a/packages/code/src/native-sandbox.test.ts b/packages/code/src/native-sandbox.test.ts index a74f3727..e2dc4bda 100644 --- a/packages/code/src/native-sandbox.test.ts +++ b/packages/code/src/native-sandbox.test.ts @@ -1030,7 +1030,8 @@ test('reports cancellation after command start as a potentially committed mutati (error: unknown) => error instanceof WorkspaceToolError && error.code === 'EXECUTION_ABORTED' && - error.mutationMayHaveCommitted === true, + error.mutationMayHaveCommitted === true && + error.requiresQuarantine === false, ); }); @@ -1157,7 +1158,9 @@ test('cleans allocated command state exactly once on every execution exit', asyn error.code === (outcome.startsWith('abort') ? 'EXECUTION_ABORTED' : 'COMMAND_UNAVAILABLE') && - error.mutationMayHaveCommitted === (outcome === 'abort-after-spawn'), + error.mutationMayHaveCommitted === (outcome === 'abort-after-spawn') && + error.requiresQuarantine === + (outcome === 'abort-after-spawn' && process.platform === 'win32'), ); } assert.equal( diff --git a/packages/code/src/native-sandbox.ts b/packages/code/src/native-sandbox.ts index 49d269c0..139a10d6 100644 --- a/packages/code/src/native-sandbox.ts +++ b/packages/code/src/native-sandbox.ts @@ -681,6 +681,10 @@ export class NativeSrtWorkspaceCommandSandbox implements WorkspaceCommandSandbox 'Workspace command execution aborted', 'EXECUTION_ABORTED', true, + // POSIX commands run in a detached process group, so its + // observed close follows a group-wide SIGKILL. The Windows + // fallback cannot yet prove descendant termination. + this.platform === 'win32', ), ); return; diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 545b96ca..fde84768 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -1562,7 +1562,7 @@ export class BridgeWorker { !( error instanceof WorkspaceToolError && this.options.workspaceTools?.mutationFailuresAreAtomic === true && - !error.mutationMayHaveCommitted + !error.requiresQuarantine )) ) { ambiguousWorkspaceMutationError = error; diff --git a/packages/code/src/workspace-worker.test.ts b/packages/code/src/workspace-worker.test.ts index 43c820fe..14bf2d61 100644 --- a/packages/code/src/workspace-worker.test.ts +++ b/packages/code/src/workspace-worker.test.ts @@ -1410,6 +1410,76 @@ test('worker clears quarantine after a composed command is cleanly rejected', as assert.deepEqual(lifecycle, ['arm', 'execute', 'settle', 'clear']); }); +test('worker clears quarantine after a command cancellation confirms process termination', async () => { + const lifecycle: string[] = []; + const baseCapabilities = { + protocolVersion: 1 as const, + operations: ['read_file' as const], + workspaces: [{ id: 'primary', operations: ['read_file' as const] }], + }; + const workspaceTools = new SandboxWorkspaceTools({ + workspaceTools: { + capabilities: baseCapabilities, + mutationFailuresAreAtomic: true, + async execute() { throw new Error('base executor must not run'); }, + }, + commandWorkspaces: ['primary'], + commandSandbox: { + mutationFailuresAreAtomic: true, + async execute() { + lifecycle.push('execute'); + throw new WorkspaceToolError( + 'Workspace command execution aborted', + 'EXECUTION_ABORTED', + true, + false, + ); + }, + }, + }); + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + workspaceTools: workspaceTools.capabilities, + }, + workspaceTools, + workspaceMutationQuarantine: mutationQuarantine( + () => lifecycle.push('quarantine'), + () => lifecycle.push('arm'), + () => lifecycle.push('clear'), + ), + fetchImpl: async () => { + lifecycle.push('settle'); + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-command-cancelled-cleanly', + workerId: 'vm-1', + incarnationId, + generation: 4, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 5_000).toISOString(), + executionKind: 'workspace_tool', + request: { + protocolVersion: 1, + operation: 'execute_command', + workspaceId: 'primary', + command: 'sleep 30', + }, + }); + assert.deepEqual(lifecycle, ['arm', 'execute', 'settle', 'clear']); +}); + test('worker retains quarantine when an atomic executor cannot confirm durability', async () => { const lifecycle: string[] = []; const workspaceCapabilities = { diff --git a/packages/code/src/workspace.ts b/packages/code/src/workspace.ts index f7770e76..5bfc07fe 100644 --- a/packages/code/src/workspace.ts +++ b/packages/code/src/workspace.ts @@ -176,6 +176,8 @@ export class WorkspaceToolError extends Error { message: string, public readonly code: WorkspaceToolErrorCode, public readonly mutationMayHaveCommitted = false, + /** Retain the durable mutation guard when process or write settlement is uncertain. */ + public readonly requiresQuarantine = mutationMayHaveCommitted, ) { super(message); this.name = 'WorkspaceToolError'; diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 41e11839..32205571 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -23,6 +23,7 @@ import { BridgeWorkspaceSlots } from './slots'; const PREFIX = 'codeapi:bridge:v1'; const POLL_INTERVAL_MS = 100; +const CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS = 5_000; const DEFAULT_WORKER_TTL_SECONDS = 60; const DEFAULT_REDIS_COMMAND_TIMEOUT_MS = 1_000; @@ -1901,6 +1902,55 @@ export class RedisBridgeStore { // settlement before returning an error, even when the caller is gone. pollError = error; } + const workspaceRequest = + assignment.executionKind === 'workspace_tool' && + isWorkspaceToolRequest(assignment.request) + ? assignment.request + : undefined; + const cancelledMutation = + signal.aborted && + workspaceRequest != null && + (workspaceRequest.operation === 'write_file' || + workspaceRequest.operation === 'edit_file' || + workspaceRequest.operation === 'execute_command'); + if (cancelledMutation) { + try { + // Keep the acknowledged assignment available long enough for the + // worker to terminate its process tree and commit a clean rejection. + // Closing it first makes that rejection impossible to acknowledge and + // leaves the worker's durable mutation guard armed. + await this.cancel(assignment.assignmentId, assignment); + // Rejected settlements remain valid after the execution deadline. + // Give Stop its own grace so a near-timeout cancellation is not + // misclassified as an ambiguous timeout. + const cancellationDeadlineAtMs = + Date.now() + CANCELLED_WORKSPACE_SETTLEMENT_GRACE_MS; + let cancellationPollMs = POLL_INTERVAL_MS; + while (Date.now() < cancellationDeadlineAtMs) { + const raw = await boundedCommand( + this.redis.get(settlementKey(assignment.assignmentId)), + Math.max( + 1, + Math.min( + this.redisCommandTimeoutMs, + cancellationDeadlineAtMs - Date.now(), + ), + ), + 'Bridge cancelled workspace settlement poll', + ); + if (raw != null) return JSON.parse(raw) as CodeBridgeSettlement; + await delay( + Math.min( + cancellationPollMs, + Math.max(0, cancellationDeadlineAtMs - Date.now()), + ), + ); + cancellationPollMs = Math.min(cancellationPollMs * 2, 500); + } + } catch (error) { + pollError ??= error; + } + } const closeKeys = [ assignmentKey(assignment.assignmentId), settlementKey(assignment.assignmentId), diff --git a/service/src/bridge/workspace-store.test.ts b/service/src/bridge/workspace-store.test.ts index a23274d8..5b0186d0 100644 --- a/service/src/bridge/workspace-store.test.ts +++ b/service/src/bridge/workspace-store.test.ts @@ -77,6 +77,108 @@ test('dispatches a workspace tool only to a worker advertising its workspace and }); }); +test('drains an acknowledged workspace mutation cancellation before releasing it', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'workspace-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'native-srt', + runtimes: [], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations: ['execute_command'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + const controller = new AbortController(); + const completion = store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'execute_command', + workspaceId: 'primary', + command: 'sleep 30; touch delayed.txt', + }, + deadlineAtMs: Date.now() + 5_000, + executionTimeoutMs: 500, + signal: controller.signal, + }); + const assignment = (await store.lease( + 'workspace-worker', + incarnationId, + 1_000, + ))!; + await store.acknowledgeLease( + 'workspace-worker', + incarnationId, + assignment.assignmentId, + assignment.generation, + assignment.leaseToken, + ); + + await new Promise((resolve) => setTimeout(resolve, 250)); + controller.abort(); + for ( + let attempt = 0; + attempt < 100 && + !(await store.cancelled( + 'workspace-worker', + incarnationId, + assignment.assignmentId, + )); + attempt += 1 + ) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + await new Promise((resolve) => setTimeout(resolve, 350)); + expect(Date.now()).toBeGreaterThan(Date.parse(assignment.expiresAt)); + await store.settle('workspace-worker', assignment.assignmentId, { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + incarnationId, + status: 'rejected', + errorCode: 'EXECUTION_ABORTED', + error: 'Workspace command execution aborted', + }); + + await expect(completion).resolves.toMatchObject({ + status: 'rejected', + errorCode: 'EXECUTION_ABORTED', + }); + + const reuse = store.dispatchWorkspaceTool({ + workerId: 'workspace-worker', + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'execute_command', + workspaceId: 'primary', + command: 'printf reused', + }, + deadlineAtMs: Date.now() + 5_000, + executionTimeoutMs: 5_000, + signal: new AbortController().signal, + }); + const next = (await store.lease( + 'workspace-worker', + incarnationId, + 1_000, + ))!; + expect(next.request).toMatchObject({ command: 'printf reused' }); + await store.settle('workspace-worker', next.assignmentId, { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: next.generation, + leaseToken: next.leaseToken, + incarnationId, + status: 'rejected', + error: 'fixture completion', + }); + await expect(reuse).resolves.toMatchObject({ status: 'rejected' }); +}); + test('rejects a workspace tool that the selected worker did not advertise', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION,