diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index 99a0dece..b9de3ac2 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -256,6 +256,23 @@ execution. the currently registered incarnation. - Request cancellation is polled by the worker and aborts the local sandbox request. +- Replay PTC clients may attach a fresh `X-LibreChat-Code-Request-ID` to each + `/exec/programmatic` request and send that same opaque ID to + `POST /v1/exec/programmatic/cancel`. Code API binds the short-lived request + record to the authenticated principal, durably marks cancellation in Redis, + and publishes it to the worker process holding the BullMQ job. This explicit + path avoids relying on HTTP connection teardown, frees waiting jobs + immediately, and interrupts active remote-bridge assignments without polling + once per active job. + Cancellation and completed-result publication use an atomic Redis decision: + a late cancel returns `already_completed` instead of acknowledging Stop after + completion won. Ambiguous enqueue/cancellation errors retain replay ownership + until a durable fence or the original job deadline. Completed results are + retained temporarily (bounded to 16 MiB) so a lost BullMQ completion reply + does not cause sandbox effects to be repeated. Reconnect reconciliation reads + only small status markers, using one subscriber per process. + Roll out the matching Code API queue-worker processes before enabling this + endpoint on API replicas; pre-cancellation workers do not observe its markers. - A leased assignment remains in a Redis-backed delivery claim until the worker explicitly acknowledges it; reconnecting before acknowledgement redelivers the same fenced assignment instead of losing it after an HTTP disconnect. diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 68d95b28..e26f212a 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -252,7 +252,10 @@ test('worker asks its supervisor to quarantine an ambiguous stateful runtime', a const quarantined: Array<{ sessionId: string; reason: string }> = []; const supervisor: RuntimeSupervisor = { async acquire() { - return { endpoint: 'http://127.0.0.1:3000/runtime', sessionId: 'rt-user-1' }; + return { + endpoint: 'http://127.0.0.1:3000/runtime', + sessionId: 'rt-user-1', + }; }, async reset() {}, async quarantine(sessionId, reason) { @@ -388,7 +391,10 @@ test('worker continues after an assignment-scoped settlement conflict', async () registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (init?.signal?.aborted === true) { @@ -397,15 +403,25 @@ test('worker continues after an assignment-scoped settlement conflict', async () if (url.endsWith('/lease')) { leases += 1; return new Response( - JSON.stringify({ protocolVersion: 1, serverElapsedMs: 0, assignment }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + JSON.stringify({ + protocolVersion: 1, + serverElapsedMs: 0, + assignment, + }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/ack')) { leaseAcknowledged = true; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/execute')) { @@ -511,7 +527,10 @@ test('worker refreshes its registration during a long assignment', async () => { registeredAt: new Date().toISOString(), leaseTtlMs: 100, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/execute')) { @@ -576,7 +595,10 @@ test('worker schedules registration freshness from request start', async () => { registeredAt: new Date().toISOString(), leaseTtlMs: 50, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/execute')) { @@ -589,7 +611,10 @@ test('worker schedules registration freshness from request start', async () => { if (url.endsWith('/cancelled')) { return new Response( JSON.stringify({ protocolVersion: 1, cancelled: false }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } return new Response( @@ -658,10 +683,13 @@ test('worker continues cancellation polling after a stalled response', async () }); } settlementAttempted = true; - return new Response(JSON.stringify({ protocolVersion: 1, accepted: true }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); }; const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', @@ -694,6 +722,175 @@ test('worker continues cancellation polling after a stalled response', async () assert.equal(settlementAttempted, true); }); +test('worker stops an outstanding cancellation request before settling completed work', async () => { + let startCancellation!: () => void; + const cancellationStarted = new Promise((resolve) => { + startCancellation = resolve; + }); + let cancellationAborted = false; + let settlementAttempted = false; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/execute')) { + await cancellationStarted; + return Response.json({ session_id: 'run-1', files: [] }); + } + if (url.endsWith('/cancellation')) { + startCancellation(); + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => { + cancellationAborted = true; + reject(new DOMException('aborted', 'AbortError')); + }, + { once: true }, + ); + }); + } + settlementAttempted = true; + return Response.json({ protocolVersion: 1, accepted: true }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + cancellationPollIntervalMs: 1, + cancellationTransportTimeoutMs: 10_000, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'complete-while-cancellation-polling', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + remainingMs: 1_000, + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(cancellationAborted, true); + assert.equal(settlementAttempted, true); +}); + +test('worker aborts a retryable cancellation error body before settling completed work', async () => { + let cancellationBodyStarted!: () => void; + const cancellationStarted = new Promise((resolve) => { + cancellationBodyStarted = resolve; + }); + let cancellationBodyAborted = false; + let settlementAttempted = false; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/execute')) { + await cancellationStarted; + return Response.json({ session_id: 'run-1', files: [] }); + } + if (url.endsWith('/cancellation')) { + return new Response( + new ReadableStream({ + start(controller) { + cancellationBodyStarted(); + init?.signal?.addEventListener( + 'abort', + () => { + cancellationBodyAborted = true; + controller.error(new DOMException('aborted', 'AbortError')); + }, + { once: true }, + ); + }, + }), + { status: 500 }, + ); + } + settlementAttempted = true; + return Response.json({ protocolVersion: 1, accepted: true }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + cancellationPollIntervalMs: 1, + cancellationTransportTimeoutMs: 10_000, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'complete-during-retryable-cancellation-response', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + remainingMs: 1_000, + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(cancellationBodyAborted, true); + assert.equal(settlementAttempted, true); +}); + +test('worker stops its cancellation delay before settling immediately completed work', async () => { + let settlementAttempted = false; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + cancellationPollIntervalMs: 10_000, + fetchImpl: async (input) => { + const url = String(input); + if (url.endsWith('/execute')) { + return Response.json({ session_id: 'run-1', files: [] }); + } + if (url.endsWith('/cancellation')) { + throw new Error('cancellation transport should not start'); + } + settlementAttempted = true; + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'complete-before-cancellation-polling', + workerId: 'vm-1', + incarnationId: 'incarnation-00000001', + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + remainingMs: 1_000, + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(settlementAttempted, true); +}); + test('worker routes a hintless assignment to an ephemeral template session', async () => { let executeUrl = ''; let runtimeSessionHeader = ''; @@ -852,10 +1049,13 @@ test('worker preserves status for a non-JSON settlement rejection', async () => }, fetchImpl: async (input) => { if (String(input).endsWith('/execute')) { - return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ session_id: 'run-1', files: [] }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); } settlementAttempts += 1; return new Response('assignment fenced', { @@ -1032,7 +1232,10 @@ test('worker retries a known-clean rejection after shutdown until acknowledged', registeredAt: new Date().toISOString(), leaseTtlMs: 50, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (String(input).endsWith('/execute')) { @@ -1046,18 +1249,20 @@ test('worker retries a known-clean rejection after shutdown until acknowledged', controller.abort(); throw new TypeError('connection reset'); } - return new Response(JSON.stringify({ protocolVersion: 1, accepted: true }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); }; const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1097,8 +1302,7 @@ test('worker preserves a definite rejection when its heartbeat fails', async () token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1118,7 +1322,10 @@ test('worker preserves a definite rejection when its heartbeat fails', async () registeredAt: new Date().toISOString(), leaseTtlMs: 50, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (String(input).endsWith('/execute')) { @@ -1139,7 +1346,10 @@ test('worker preserves a definite rejection when its heartbeat fails', async () } return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1170,8 +1380,7 @@ test('worker quarantines a stateful workspace after a sandbox 5xx response', asy token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1187,7 +1396,10 @@ test('worker quarantines a stateful workspace after a sandbox 5xx response', asy settlementAttempted = true; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1217,8 +1429,7 @@ test('worker treats a non-JSON sandbox 4xx as a definite rejection', async () => token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1235,7 +1446,10 @@ test('worker treats a non-JSON sandbox 4xx as a definite rejection', async () => JSON.parse(String(init?.body) || '{}').status === 'rejected'; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1268,18 +1482,20 @@ test('worker quarantines a stateful workspace after the sandbox request aborts', }); } settlementAttempted = true; - return new Response(JSON.stringify({ protocolVersion: 1, accepted: true }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); }; const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1331,13 +1547,23 @@ test('worker surfaces quarantine when shutdown aborts stateful execution', async registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/lease')) { return new Response( - JSON.stringify({ protocolVersion: 1, serverElapsedMs: 0, assignment }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + JSON.stringify({ + protocolVersion: 1, + serverElapsedMs: 0, + assignment, + }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (url.endsWith('/execute')) { @@ -1351,18 +1577,20 @@ test('worker surfaces quarantine when shutdown aborts stateful execution', async ); }); } - return new Response(JSON.stringify({ protocolVersion: 1, accepted: true }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ protocolVersion: 1, accepted: true }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); }; const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1510,7 +1738,10 @@ test('worker subtracts lease response transit from the server budget', async () if (String(input).endsWith('/ack')) { return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } now += 50; @@ -1527,10 +1758,16 @@ test('worker subtracts lease response transit from the server budget', async () leaseToken: 'lease-token-that-is-long-enough-for-testing', expiresAt: new Date(0).toISOString(), remainingMs: 1_000, - request: { body: { language: 'bash' }, headers: {} }, + request: { + body: { language: 'bash' }, + headers: {}, + }, }, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1572,21 +1809,28 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( registeredAt: new Date().toISOString(), leaseTtlMs: 50, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (String(input).endsWith('/ack')) { now += 10; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (String(input).endsWith('/settle')) { settlementAttempts += 1; - abandonedSettlement = JSON.parse( - String(init?.body), - ) as Record; + abandonedSettlement = JSON.parse(String(init?.body)) as Record< + string, + unknown + >; if (settlementAttempts === 1) { return new Response(JSON.stringify({ error: 'unavailable' }), { status: 503, @@ -1595,7 +1839,10 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( } return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } return new Response( @@ -1611,15 +1858,24 @@ test('worker rejects a lease whose acknowledgement exhausts its budget', async ( leaseToken: 'lease-token-that-is-long-enough-for-testing', expiresAt: new Date(0).toISOString(), remainingMs: 10, - request: { body: { language: 'bash' }, headers: {} }, + request: { + body: { language: 'bash' }, + headers: {}, + }, }, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); - await assert.rejects(worker.lease(), /expired during lease acknowledgement/); + await assert.rejects( + worker.lease(), + /expired during lease acknowledgement/, + ); assert.equal(abandonedSettlement?.status, 'rejected'); assert.ok(registrations > 0); assert.equal(settlementAttempts, 2); @@ -1653,7 +1909,10 @@ test('worker rejects an assignment after ambiguous acknowledgement delivery', as JSON.parse(String(init?.body) || '{}').status === 'rejected'; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } if (String(input).endsWith('/workers/register')) { @@ -1665,7 +1924,10 @@ test('worker rejects an assignment after ambiguous acknowledgement delivery', as registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); } return new Response( @@ -1685,7 +1947,10 @@ test('worker rejects an assignment after ambiguous acknowledgement delivery', as request: { body: { language: 'bash' }, headers: {} }, }, }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1702,8 +1967,7 @@ test('worker clamps rejected settlement errors to the protocol limit', async () token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1716,11 +1980,16 @@ test('worker clamps rejected settlement errors to the protocol limit', async () headers: { 'Content-Type': 'application/json' }, }); } - const settlement = JSON.parse(String(init?.body)) as { error: string }; + const settlement = JSON.parse(String(init?.body)) as { + error: string; + }; rejection = settlement.error; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1747,8 +2016,7 @@ test('worker quarantines an explicitly dirty stateful sandbox response', async ( token: 'worker-secret', workerId: 'vm-1', incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -1761,13 +2029,19 @@ test('worker quarantines an explicitly dirty stateful sandbox response', async ( error: 'session_workspace_dirty', message: 'restore required', }), - { status: 409, headers: { 'Content-Type': 'application/json' } }, + { + status: 409, + headers: { 'Content-Type': 'application/json' }, + }, ); } settlementAttempted = true; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1859,15 +2133,21 @@ test('worker uses the server-relative lease budget despite VM clock skew', async }, fetchImpl: async (input) => { if (String(input).endsWith('/execute')) { - return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, - }); + return new Response( + JSON.stringify({ session_id: 'run-1', files: [] }), + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, + ); } settlementAttempted = true; return new Response( JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, + { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }, ); }, }); @@ -1886,7 +2166,6 @@ test('worker uses the server-relative lease budget despite VM clock skew', async assert.equal(settlementAttempted, true); }); - test('worker continues after an expired assignment settlement conflict', async () => { const controller = new AbortController(); let registrations = 0; @@ -1921,18 +2200,22 @@ test('worker continues after an expired assignment settlement conflict', async ( leases += 1; return Response.json({ protocolVersion: 1, - assignment: leases === 1 - ? { - protocolVersion: 1, - assignmentId: 'assignment-expired', - workerId: 'vm-1', - incarnationId, - generation: 1, - leaseToken: 'lease-token-that-is-long-enough-for-testing', - expiresAt: new Date(Date.now() + 10_000).toISOString(), - request: { body: { language: 'bash' }, headers: {} }, - } - : undefined, + assignment: + leases === 1 + ? { + protocolVersion: 1, + assignmentId: 'assignment-expired', + workerId: 'vm-1', + incarnationId, + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 10_000).toISOString(), + request: { + body: { language: 'bash' }, + headers: {}, + }, + } + : undefined, }); } if (url.endsWith('/execute')) { @@ -1940,7 +2223,10 @@ test('worker continues after an expired assignment settlement conflict', async ( } if (url.endsWith('/settle')) { return Response.json( - { error: 'Bridge assignment has expired', code: 'ASSIGNMENT_EXPIRED' }, + { + error: 'Bridge assignment has expired', + code: 'ASSIGNMENT_EXPIRED', + }, { status: 409 }, ); } @@ -2172,7 +2458,10 @@ test('paired worker rotates credentials throughout a long assignment', async () } if (url.endsWith('/execute')) { await new Promise((resolve) => setTimeout(resolve, 55)); - return Response.json({ session_id: 'run-long-rotation', files: [] }); + return Response.json({ + session_id: 'run-long-rotation', + files: [], + }); } return Response.json({ protocolVersion: 1, accepted: true }); }; @@ -2276,6 +2565,117 @@ test('paired worker cancels a stalled credential refresh after execution', async assert.equal(refreshAborted, true); }); +test('one concurrent caller cannot abort a credential refresh another caller still needs', async () => { + const key = createBridgeIdentity(); + const first = new AbortController(); + const second = new AbortController(); + let releaseRefresh!: () => void; + let refreshStarted!: () => void; + const started = new Promise((resolve) => { + refreshStarted = resolve; + }); + const released = new Promise((resolve) => { + releaseRefresh = resolve; + }); + let transportAborted = false; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'credential-before-shared-refresh', + expiresAt: new Date(Date.now() + 5).toISOString(), + }, + credentialRefreshWindowMs: 10, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (input, init) => { + assert.match(String(input), /credentials\/refresh$/); + refreshStarted(); + init?.signal?.addEventListener('abort', () => { + transportAborted = true; + }); + await released; + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'credential-after-shared-refresh-value', + expiresAt: new Date(Date.now() + 120_000).toISOString(), + }); + }, + }); + + const firstRefresh = worker.refreshCredential(first.signal); + await started; + const secondRefresh = worker.refreshCredential(second.signal); + first.abort(); + + await assert.rejects(firstRefresh, { name: 'AbortError' }); + assert.equal(transportAborted, false); + releaseRefresh(); + await secondRefresh; + assert.equal(transportAborted, false); +}); + +test('a new caller starts a fresh credential refresh after the last waiter aborts', async () => { + const key = createBridgeIdentity(); + const first = new AbortController(); + let refreshCount = 0; + let firstRefreshStarted!: () => void; + const started = new Promise((resolve) => { + firstRefreshStarted = resolve; + }); + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'credential-before-replacement-refresh', + expiresAt: new Date(Date.now() + 5).toISOString(), + }, + credentialRefreshWindowMs: 10, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async (_input, init) => { + refreshCount += 1; + if (refreshCount === 1) { + firstRefreshStarted(); + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }); + } + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'credential-after-replacement-refresh', + expiresAt: new Date(Date.now() + 120_000).toISOString(), + }); + }, + }); + + const abandoned = worker.refreshCredential(first.signal, Date.now() + 1_000); + await started; + first.abort(); + await assert.rejects(abandoned, { name: 'AbortError' }); + await worker.refreshCredential(undefined, Date.now() + 1_000); + + assert.equal(refreshCount, 2); +}); + test('paired worker refreshes conservatively before server clock calibration', async () => { const key = createBridgeIdentity(); let refreshCount = 0; @@ -2347,8 +2747,7 @@ test('paired worker charges initial credential refresh against the assignment de sandboxStarted = true; } if (url.endsWith('/settle')) { - rejected = - JSON.parse(String(init?.body)).status === 'rejected'; + rejected = JSON.parse(String(init?.body)).status === 'rejected'; } return Response.json({ protocolVersion: 1, @@ -2389,8 +2788,7 @@ test('paired worker rechecks the deadline after request serialization', async () codeApiUrl: 'https://code.example/v1', workerId: 'vm-1', incarnationId, - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', + sandboxEndpoint: 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2', identity: { privateKey: key.privateKey, credential: 'credential-valid-during-serialization', @@ -2408,8 +2806,7 @@ test('paired worker rechecks the deadline after request serialization', async () sandboxStarted = true; } if (url.endsWith('/settle')) { - rejected = - JSON.parse(String(init?.body)).status === 'rejected'; + rejected = JSON.parse(String(init?.body)).status === 'rejected'; } return Response.json({ protocolVersion: 1, @@ -2471,8 +2868,7 @@ test('paired worker keeps endpoint validation failures known-clean', async () => const url = String(input); if (url.endsWith('/execute')) sandboxStarted = true; if (url.endsWith('/settle')) { - rejected = - JSON.parse(String(init?.body)).status === 'rejected'; + rejected = JSON.parse(String(init?.body)).status === 'rejected'; } return Response.json({ protocolVersion: 1, @@ -2758,11 +3154,13 @@ test('sandbox completion does not cancel an in-flight credential rotation', asyn } if (url.endsWith('/execute')) { await refreshStartedPromise; - return Response.json({ session_id: 'run-rotation-race', files: [] }); + return Response.json({ + session_id: 'run-rotation-race', + files: [], + }); } - settleAuthorization = ( - init?.headers as Record - ).Authorization; + settleAuthorization = (init?.headers as Record) + .Authorization; return Response.json({ protocolVersion: 1, accepted: true }); }; const worker = new BridgeWorker({ @@ -2799,8 +3197,41 @@ test('sandbox completion does not cancel an in-flight credential rotation', asyn ); }); +test('settlement does not drain another lane credential renewal', async () => { + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', workerId: 'vm-1', token: 'fixture', + incarnationId, sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', runtimes: ['bash'] }, + fetchImpl: async (input) => String(input).endsWith('/execute') + ? Response.json({ session_id: 'independent-lane', files: [] }) + : Response.json({ protocolVersion: 1, accepted: true }), + }); + // A different lane owns this pending renewal. The settling lane has no + // maintenance waiter and must not consume its own lease on that promise. + Object.assign(worker, { credentialInFlight: { + promise: new Promise(() => {}), controller: new AbortController(), waiters: 1, + }, refreshCredential: async () => {} }); + const startedAt = Date.now(); + await worker.executeAndSettle({ + protocolVersion: 1, assignmentId: 'independent-lane', workerId: 'vm-1', + incarnationId, generation: 5, leaseToken: 'independent-lane-lease-token', + expiresAt: new Date(Date.now() + 600_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + assert.ok(Date.now() - startedAt < 500, 'unrelated renewal must not add a one-second drain'); +}); + test('reconnect delay uses bounded exponential jitter', () => { - assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 0), 500); - assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 1), 1_000); - assert.equal(reconnectDelayMs(10, 1_000, 30_000, () => 1), 30_000); + assert.equal( + reconnectDelayMs(0, 1_000, 30_000, () => 0), + 500, + ); + assert.equal( + reconnectDelayMs(0, 1_000, 30_000, () => 1), + 1_000, + ); + assert.equal( + reconnectDelayMs(10, 1_000, 30_000, () => 1), + 30_000, + ); }); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index a0517bd9..ffe1b350 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -38,6 +38,11 @@ export interface BridgeWorkerOptions { capabilities: BridgeWorkerCapabilities; workspaceTools?: WorkspaceToolExecutor; workspaceProgrammatic?: { + /** + * True when a WorkspaceToolError without mutation uncertainty proves the + * selected workspace was not changed. + */ + mutationFailuresAreAtomic?: true; executeProgrammatic( workspaceId: string, request: BridgeWorkspaceProgrammaticRequest, @@ -99,6 +104,7 @@ const DEFAULT_REGISTRATION_TRANSPORT_TIMEOUT_MS = 10_000; const DEFAULT_CONTROL_TRANSPORT_TIMEOUT_MS = 10_000; const DEFAULT_CANCELLATION_POLL_INTERVAL_MS = 500; const DEFAULT_CANCELLATION_TRANSPORT_TIMEOUT_MS = 2_000; +const CREDENTIAL_REFRESH_SETTLEMENT_GRACE_MS = 1_000; const MIN_REGISTRATION_HEARTBEAT_MS = 25; const REGISTRATION_RETRY_DELAY_MS = 100; const CREDENTIAL_REFRESH_RETRY_DELAY_MS = 100; @@ -158,26 +164,29 @@ function workspaceCapabilitiesMatch( advertised.writeFileModes?.length === executor.writeFileModes?.length && (advertised.writeFileModes?.every( (mode, index) => mode === executor.writeFileModes?.[index], - ) ?? executor.writeFileModes == null) && + ) ?? + executor.writeFileModes == null) && advertised.editFileModes?.length === executor.editFileModes?.length && (advertised.editFileModes?.every( (mode, index) => mode === executor.editFileModes?.[index], - ) ?? executor.editFileModes == null) && - advertised.editFileFeatures?.length === - executor.editFileFeatures?.length && + ) ?? + executor.editFileModes == null) && + advertised.editFileFeatures?.length === executor.editFileFeatures?.length && (advertised.editFileFeatures?.every( (feature, index) => feature === executor.editFileFeatures?.[index], - ) ?? executor.editFileFeatures == null) && - advertised.listFileFeatures?.length === - executor.listFileFeatures?.length && + ) ?? + executor.editFileFeatures == null) && + advertised.listFileFeatures?.length === executor.listFileFeatures?.length && (advertised.listFileFeatures?.every( (feature, index) => feature === executor.listFileFeatures?.[index], - ) ?? executor.listFileFeatures == null) && + ) ?? + executor.listFileFeatures == null) && advertised.programmaticLanguages?.length === executor.programmaticLanguages?.length && (advertised.programmaticLanguages?.every( (language, index) => language === executor.programmaticLanguages?.[index], - ) ?? executor.programmaticLanguages == null) && + ) ?? + executor.programmaticLanguages == null) && advertised.workspaces.length === executor.workspaces.length && advertised.workspaces.every( (workspace, index) => @@ -189,7 +198,8 @@ function workspaceCapabilitiesMatch( (operation, operationIndex) => operation === executor.workspaces[index]?.operations?.[operationIndex], - ) ?? executor.workspaces[index]?.operations == null), + ) ?? + executor.workspaces[index]?.operations == null), ) ); } @@ -201,8 +211,7 @@ function registrationCompatibleCapabilities( if ( workspaceTools == null || (workspaceTools.operations.every( - (operation) => - operation === 'read_file' || operation === 'search_text', + (operation) => operation === 'read_file' || operation === 'search_text', ) && workspaceTools.workspaces.every( (workspace) => workspace.operations == null, @@ -211,8 +220,7 @@ function registrationCompatibleCapabilities( return capabilities; } const operations = workspaceTools.operations.filter( - (operation) => - operation === 'read_file' || operation === 'search_text', + (operation) => operation === 'read_file' || operation === 'search_text', ); if (operations.length === 0) { const { workspaceTools: _workspaceTools, ...compatible } = capabilities; @@ -221,7 +229,9 @@ function registrationCompatibleCapabilities( const workspaces = workspaceTools.workspaces.flatMap((workspace) => { if ( workspace.operations != null && - !operations.every((operation) => workspace.operations?.includes(operation)) + !operations.every((operation) => + workspace.operations?.includes(operation), + ) ) { return []; } @@ -386,7 +396,11 @@ export class BridgeWorker { private negotiatedWorkspaceSlots = 1; private concurrentRunning = false; private registrationInFlight?: Promise; - private credentialInFlight?: Promise; + private credentialInFlight?: { + promise: Promise; + controller: AbortController; + waiters: number; + }; private serverClockOffsetMs = MAX_PROOF_CLOCK_SKEW_MS; constructor(private readonly options: BridgeWorkerOptions) { @@ -434,7 +448,8 @@ export class BridgeWorker { (options.workspaceProgrammatic != null) !== (options.capabilities.workspaceTools?.programmaticLanguages?.includes( 'bash', - ) === true) + ) === + true) ) { throw new BridgeProtocolError( 'Workspace programmatic capability requires a matching executor', @@ -553,7 +568,9 @@ export class BridgeWorker { if (signal?.aborted) { abortRegistration(); } else { - signal?.addEventListener('abort', abortRegistration, { once: true }); + signal?.addEventListener('abort', abortRegistration, { + once: true, + }); } const timeoutMs = Math.min( Math.max(1, this.registrationTtlMs - 1), @@ -575,7 +592,10 @@ export class BridgeWorker { workerId: this.options.workerId, incarnationId: this.incarnationId, capabilities: this.maintenanceOnly - ? { ...capabilities, requiresReadyConfirmation: true } + ? { + ...capabilities, + requiresReadyConfirmation: true, + } : capabilities, }, registrationController.signal, @@ -709,7 +729,9 @@ export class BridgeWorker { } await this.runtimeSupervisor.reset(runtimeSessionId, signal); await this.timedRequest( - `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}/workspaces/reset`, + `${this.codeApiUrl}${bridgeWorkerPath( + this.options.workerId, + )}/workspaces/reset`, { protocolVersion: BRIDGE_PROTOCOL_VERSION, incarnationId: this.incarnationId, @@ -745,7 +767,9 @@ export class BridgeWorker { // machine-local guard before the remote fence can be removed. await guard.assertAvailable(); await this.timedRequest( - `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}/workspaces/reset`, + `${this.codeApiUrl}${bridgeWorkerPath( + this.options.workerId, + )}/workspaces/reset`, { protocolVersion: BRIDGE_PROTOCOL_VERSION, incarnationId: this.incarnationId, @@ -1051,20 +1075,58 @@ export class BridgeWorker { transportTimeoutMs = Number.POSITIVE_INFINITY, ): Promise { while (this.credentialInFlight) { - await this.credentialInFlight; + await this.waitForCredentialRefresh(this.credentialInFlight, signal); // A longer-lived caller may still need another refresh after this one. } + const controller = new AbortController(); const pending = this.refreshCredentialOwned( - signal, + controller.signal, validThroughMs, transportTimeoutMs, ); - this.credentialInFlight = pending; + const entry = { promise: pending, controller, waiters: 0 }; + this.credentialInFlight = entry; + void pending.then( + () => { + if (this.credentialInFlight === entry) + this.credentialInFlight = undefined; + }, + () => { + if (this.credentialInFlight === entry) + this.credentialInFlight = undefined; + }, + ); + await this.waitForCredentialRefresh(entry, signal); + } + + private async waitForCredentialRefresh( + entry: NonNullable, + signal?: AbortSignal, + ): Promise { + entry.waiters += 1; + let removeAbortListener = (): void => {}; + const aborted = new Promise((_, reject) => { + if (signal == null) return; + const abort = (): void => + reject( + signal.reason instanceof Error + ? signal.reason + : new DOMException('aborted', 'AbortError'), + ); + removeAbortListener = (): void => + signal.removeEventListener('abort', abort); + signal.addEventListener('abort', abort, { once: true }); + if (signal.aborted) abort(); + }); try { - await pending; + await Promise.race([entry.promise, aborted]); } finally { - if (this.credentialInFlight === pending) + removeAbortListener(); + entry.waiters -= 1; + if (entry.waiters === 0 && this.credentialInFlight === entry) { this.credentialInFlight = undefined; + entry.controller.abort(); + } } } @@ -1118,7 +1180,7 @@ export class BridgeWorker { assignment: BridgeAssignment, stopSignal: AbortSignal, serverClockOffsetMs: number, - requestSignal?: AbortSignal, + maintenance: { refresh?: Promise }, ): Promise { const identity = this.options.identity; if (identity == null) return; @@ -1136,18 +1198,18 @@ export class BridgeWorker { await abortableDelay(waitMs, stopSignal); if (stopSignal.aborted || Date.now() >= assignmentDeadlineMs) return; try { - await this.refreshCredential( - requestSignal, + maintenance.refresh = this.refreshCredential( + stopSignal, Date.now() + serverClockOffsetMs + refreshWindowMs, ); + await maintenance.refresh; } catch (error) { if (stopSignal.aborted) return; const terminal = error instanceof BridgeProtocolError && (error.status === 401 || error.status === 403); const credentialRemainingMs = - Date.parse(identity.expiresAt) - - (Date.now() + serverClockOffsetMs); + Date.parse(identity.expiresAt) - (Date.now() + serverClockOffsetMs); if (terminal || credentialRemainingMs <= 0) throw error; await abortableDelay( Math.min( @@ -1156,6 +1218,8 @@ export class BridgeWorker { ), stopSignal, ); + } finally { + maintenance.refresh = undefined; } } } @@ -1352,6 +1416,7 @@ export class BridgeWorker { ); let credentialMaintenanceError: unknown; let credentialMaintenance: Promise | undefined; + const ownCredentialMaintenance: { refresh?: Promise } = {}; let settlement: BridgeSettlement; let ambiguousSandboxError: unknown; let ambiguousWorkspaceMutationError: unknown; @@ -1368,7 +1433,7 @@ export class BridgeWorker { assignment, credentialController.signal, serverClockOffsetMs, - signal, + ownCredentialMaintenance, ).catch((error) => { credentialMaintenanceError = error; executionController.abort(); @@ -1674,14 +1739,22 @@ export class BridgeWorker { ) { workspaceMutationGuardError = error; } + const knownAtomicWorkspaceToolFailure = + assignment.executionKind === 'workspace_tool' && + error instanceof WorkspaceToolError && + this.options.workspaceTools?.mutationFailuresAreAtomic === true && + !error.requiresQuarantine; + const knownAtomicProgrammaticFailure = + assignment.executionKind === 'workspace_programmatic' && + error instanceof WorkspaceToolError && + this.options.workspaceProgrammatic?.mutationFailuresAreAtomic === + true && + !error.requiresQuarantine; if ( workspaceMutationApplied || (workspaceMutationArmed && - !( - error instanceof WorkspaceToolError && - this.options.workspaceTools?.mutationFailuresAreAtomic === true && - !error.requiresQuarantine - )) + !knownAtomicWorkspaceToolFailure && + !knownAtomicProgrammaticFailure) ) { ambiguousWorkspaceMutationError = error; } @@ -1713,12 +1786,30 @@ export class BridgeWorker { clearTimeout(deadlineTimer); cancellationController.abort(); await cancellationWatcher; + // Only drain renewal joined by this assignment, never an unrelated lane's + // refresh. Leave settlement time inside the original assignment budget. + const credentialInFlight = ownCredentialMaintenance.refresh; + if (credentialInFlight != null && !credentialController.signal.aborted) { + let drainTimer: ReturnType | undefined; + await Promise.race([ + credentialInFlight.catch(() => undefined), + new Promise((resolve) => { + drainTimer = setTimeout( + resolve, + Math.min(CREDENTIAL_REFRESH_SETTLEMENT_GRACE_MS, + Math.max(0, Date.parse(assignment.expiresAt) - serverClockOffsetMs - Date.now() - 5_000)), + ); + }), + ]); + if (drainTimer != null) clearTimeout(drainTimer); + } credentialController.abort(); await credentialMaintenance; try { if (workspaceMutationGuardError != null) throw workspaceMutationGuardError; if (ambiguousWorkspaceMutationError != null) { + this.options.onError?.(ambiguousWorkspaceMutationError); throw await this.quarantineWorkspace( undefined, 'Worker stopped after a workspace mutation completed without a fulfilled settlement', @@ -1913,7 +2004,9 @@ export class BridgeWorker { return await lease.execute({ body, headers, signal }); } if (lease.endpoint == null) { - throw new BridgeProtocolError('Runtime lease does not provide an execution transport'); + throw new BridgeProtocolError( + 'Runtime lease does not provide an execution transport', + ); } const endpoint = lease.endpoint.replace(/\/+$/, ''); const response = await this.fetchImpl(`${endpoint}/execute`, { @@ -2170,17 +2263,23 @@ export class BridgeWorker { signal: AbortSignal, ): Promise { while (!signal.aborted && !executionController.signal.aborted) { - await this.delay( - Math.max( - 1, - this.options.cancellationPollIntervalMs ?? - DEFAULT_CANCELLATION_POLL_INTERVAL_MS, - ), - signal, - ); + try { + await this.delay( + Math.max( + 1, + this.options.cancellationPollIntervalMs ?? + DEFAULT_CANCELLATION_POLL_INTERVAL_MS, + ), + signal, + ); + } catch (error) { + if (signal.aborted || executionController.signal.aborted) return; + throw error; + } if (signal.aborted || executionController.signal.aborted) return; const pollController = new AbortController(); const abortPoll = (): void => pollController.abort(); + signal.addEventListener('abort', abortPoll, { once: true }); executionController.signal.addEventListener('abort', abortPoll, { once: true, }); @@ -2200,6 +2299,15 @@ export class BridgeWorker { incarnationId: this.incarnationId, }, pollController.signal, + (response) => { + // Once response headers arrive, drain the bounded body before a + // successful execution can settle. Otherwise a cancellation=true + // response racing command completion can be discarded. The + // transport timer and execution signal still cap the drain. + if (response.ok || response.status === 404) { + signal.removeEventListener('abort', abortPoll); + } + }, ); if (response.cancelled) { executionController.abort(); @@ -2213,6 +2321,7 @@ export class BridgeWorker { if (signal.aborted) return; } finally { clearTimeout(timeout); + signal.removeEventListener('abort', abortPoll); executionController.signal.removeEventListener('abort', abortPoll); } } @@ -2222,6 +2331,7 @@ export class BridgeWorker { url: string, body: object, signal?: AbortSignal, + onResponseHeaders?: (response: Response) => void, ): Promise { const requestBody = JSON.stringify(body); const response = await this.fetchImpl(url, { @@ -2233,6 +2343,7 @@ export class BridgeWorker { body: requestBody, signal, }); + onResponseHeaders?.(response); let payload: unknown; try { payload = await response.json(); diff --git a/packages/code/src/workspace-worker.test.ts b/packages/code/src/workspace-worker.test.ts index fa08a84c..f6205cd4 100644 --- a/packages/code/src/workspace-worker.test.ts +++ b/packages/code/src/workspace-worker.test.ts @@ -1032,6 +1032,250 @@ test('worker executes programmatic Bash in the selected workspace and preserves ]); }); +test('worker keeps a selected workspace usable after an atomic programmatic setup failure', async () => { + const lifecycle: string[] = []; + let settlement: Record | undefined; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['execute_command' as const], + programmaticLanguages: ['bash' as const], + workspaces: [{ id: 'primary' }], + }; + 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: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + mutationFailuresAreAtomic: true, + async execute() { + throw new Error('workspace tool executor must not run'); + }, + }, + workspaceProgrammatic: { + mutationFailuresAreAtomic: true, + async executeProgrammatic() { + throw new WorkspaceToolError( + 'Programmatic input download failed', + 'COMMAND_UNAVAILABLE', + ); + }, + }, + workspaceQuarantines: new Map([ + [ + 'primary', + mutationQuarantine( + () => lifecycle.push('quarantine'), + () => lifecycle.push('arm'), + () => lifecycle.push('clear'), + ), + ], + ]), + fetchImpl: async (_input, init) => { + settlement = JSON.parse(String(init?.body)) as Record; + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-programmatic-setup-failure', + 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_programmatic', + workspaceId: 'primary', + request: { + body: { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [{ name: 'main.sh', content: 'echo ready' }], + }, + headers: {}, + }, + }); + + assert.deepEqual(lifecycle, ['arm', 'clear']); + assert.equal(settlement?.status, 'rejected'); + assert.equal(settlement?.errorCode, 'COMMAND_UNAVAILABLE'); +}); + +test('worker keeps a selected workspace usable after confirmed programmatic cancellation cleanup', async () => { + const lifecycle: string[] = []; + let settlement: Record | undefined; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['execute_command' as const], + programmaticLanguages: ['bash' as const], + workspaces: [{ id: 'primary' }], + }; + 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: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + mutationFailuresAreAtomic: true, + async execute() { + throw new Error('workspace tool executor must not run'); + }, + }, + workspaceProgrammatic: { + mutationFailuresAreAtomic: true, + async executeProgrammatic() { + throw new WorkspaceToolError( + 'Workspace command execution aborted', + 'EXECUTION_ABORTED', + true, + false, + ); + }, + }, + workspaceQuarantines: new Map([ + [ + 'primary', + mutationQuarantine( + () => lifecycle.push('quarantine'), + () => lifecycle.push('arm'), + () => lifecycle.push('clear'), + ), + ], + ]), + fetchImpl: async (_input, init) => { + settlement = JSON.parse(String(init?.body)) as Record; + return Response.json({ protocolVersion: 1, accepted: true }); + }, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-programmatic-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_programmatic', + workspaceId: 'primary', + request: { + body: { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [{ name: 'main.sh', content: 'sleep 30' }], + }, + headers: {}, + }, + }); + + assert.deepEqual(lifecycle, ['arm', 'clear']); + assert.equal(settlement?.status, 'rejected'); + assert.equal(settlement?.errorCode, 'EXECUTION_ABORTED'); +}); + +test('worker reports the underlying cause before quarantining an uncertain programmatic mutation', async () => { + const rootCause = new WorkspaceToolError( + 'Programmatic output upload failed', + 'COMMAND_UNAVAILABLE', + true, + true, + ); + let reported: unknown; + let quarantined = false; + const workspaceCapabilities = { + protocolVersion: 1 as const, + operations: ['execute_command' as const], + programmaticLanguages: ['bash' as const], + workspaces: [{ id: 'primary' }], + }; + 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: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: workspaceCapabilities, + }, + workspaceTools: { + capabilities: workspaceCapabilities, + mutationFailuresAreAtomic: true, + async execute() { + throw new Error('workspace tool executor must not run'); + }, + }, + workspaceProgrammatic: { + mutationFailuresAreAtomic: true, + async executeProgrammatic() { + throw rootCause; + }, + }, + workspaceQuarantines: new Map([ + [ + 'primary', + mutationQuarantine(() => { + quarantined = true; + }), + ], + ]), + onError(error) { + reported = error; + }, + fetchImpl: async () => { + throw new Error('settlement must not run'); + }, + }); + + await assert.rejects( + worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-programmatic-uncertain-failure', + 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_programmatic', + workspaceId: 'primary', + request: { + body: { + language: 'bash', + version: '5.2', + session_id: 'session-1', + files: [{ name: 'main.sh', content: 'echo ready' }], + }, + headers: {}, + }, + }), + BridgeWorkspaceQuarantinedError, + ); + + assert.equal(reported, rootCause); + assert.equal(quarantined, true); +}); + test('worker stops after Code API rejects a fulfilled workspace mutation', async () => { let quarantinedReason: string | undefined; let armed = 0; diff --git a/service/src/bridge/concurrent-worker.test.ts b/service/src/bridge/concurrent-worker.test.ts index 5aec49cf..ea5ac601 100644 --- a/service/src/bridge/concurrent-worker.test.ts +++ b/service/src/bridge/concurrent-worker.test.ts @@ -243,11 +243,25 @@ for (const failure of [ status: 'fulfilled', value: { status: 'fulfilled' }, }); - for (let i = 0; i < 300 && errors.length === 0; i++) + const diagnosticCount = cleanupFailure ? 1 : 2; + const quarantineAttemptCount = + failure === 'lost-response' + ? 2 + : failure === 'all-responses-lost' || failure === 'delivery-outage' + ? 3 + : 1; + for ( + let i = 0; + i < 300 && + (errors.length < diagnosticCount || + (!cleanupFailure && quarantineAttempts < quarantineAttemptCount)); + i++ + ) { await new Promise((resolve) => setTimeout(resolve, 5)); + } if (failure === 'delivery-outage') - expect(errors.length).toBeGreaterThanOrEqual(1); - else expect(errors.length).toBe(1); + expect(errors.length).toBeGreaterThanOrEqual(2); + else expect(errors.length).toBe(diagnosticCount); if (failure === 'lost-response') expect(quarantineAttempts).toBe(2); if (failure === 'delivery-outage') expect(quarantineAttempts).toBeGreaterThanOrEqual(3); diff --git a/service/src/config.spec.ts b/service/src/config.spec.ts index 44b64b8b..69b416d8 100644 --- a/service/src/config.spec.ts +++ b/service/src/config.spec.ts @@ -99,6 +99,11 @@ describe('egress grant TTL configuration', () => { }); describe('job deadline accounting', () => { + it('never extends the producer deadline when worker configuration differs', () => { + expect(jobDeadlineAtMs(1_000, 300_000, 50_000, 91_000)).toBe(91_000); + expect(jobDeadlineAtMs(1_000, 30_000, 50_000, 91_000)).toBe(31_000); + expect(jobDeadlineAtMs(1_000, 300_000, 50_000, Number.NaN)).toBe(0); + }); it('counts time spent waiting in BullMQ against JOB_TIMEOUT', () => { expect(jobDeadlineAtMs(1_000, 300_000, 50_000)).toBe(301_000); }); diff --git a/service/src/config.ts b/service/src/config.ts index d025831e..55570dfe 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -95,10 +95,17 @@ export function jobDeadlineAtMs( enqueuedAtMs: number | undefined, timeoutMs: number, nowMs: number = Date.now(), + producerDeadlineAtMs?: number, ): number { - return Number.isFinite(enqueuedAtMs) && (enqueuedAtMs as number) > 0 + const localDeadline = Number.isFinite(enqueuedAtMs) && (enqueuedAtMs as number) > 0 ? (enqueuedAtMs as number) + timeoutMs : nowMs + timeoutMs; + if (producerDeadlineAtMs === undefined) return localDeadline; + // A worker with a larger JOB_TIMEOUT must not outlive the admission fence + // retained by its API producer. Malformed explicit deadlines fail closed. + return Number.isFinite(producerDeadlineAtMs) + ? Math.min(localDeadline, producerDeadlineAtMs) + : 0; } /** The worker stops user work at JOB_TIMEOUT, then may still need to terminate diff --git a/service/src/job-cancellation-commit.test.ts b/service/src/job-cancellation-commit.test.ts new file mode 100644 index 00000000..8ded29f2 --- /dev/null +++ b/service/src/job-cancellation-commit.test.ts @@ -0,0 +1,431 @@ +import { afterEach, beforeEach, expect, test } from 'bun:test'; +import { startTestRedis } from './test/redis'; +import { RedisBridgeStore } from './bridge/store'; +import { + commitJobResult, + readCommittedJobResult, + requestJobCancellation, + JobCancellationRegistry, + jobCancellationInternals, + fenceJobCancellation, + waitForJobWithCancellation, + jobCancellationRetentionSeconds, + claimJobExecution, +} from './job-cancellation'; + +let redis: Awaited>; +beforeEach(async () => { + redis = await startTestRedis(); +}); +afterEach(async () => { + await redis.closeTestServer(); +}); +const target = { queueName: 'other', jobId: 'commit-race' }; + +for (const outcome of ['commit', 'stop', 'duplicate']) + test(`native mutation handoff commits or quarantines before root release (${outcome})`, async () => { + const store = new RedisBridgeStore(redis); + const workerId = 'handoff-worker'; + const incarnationId = 'incarnation-handoff-01'; + await store.register({ + protocolVersion: 1, + workerId, + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'anthropic-srt', + runtimes: [], + workspaceTools: { + protocolVersion: 1, + operations: ['execute_command'], + programmaticLanguages: ['bash'], + workspaces: [{ id: 'primary' }], + }, + }, + }); + const controller = new AbortController(); + const dispatchArgs = { + workerId, + workspaceId: 'primary', + headers: {}, + body: { + language: 'bash', + version: '5.2', + session_id: 'handoff-session', + files: [{ name: 'main.sh', content: 'echo mutation' }], + }, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }; + const completion = store.dispatch({ + ...dispatchArgs, + finalize: async settlement => { + if (outcome === 'stop') await requestJobCancellation(redis, target, 60); + if (outcome === 'duplicate') + await commitJobResult( + redis, + target, + { stdout: 'first mutation' }, + 60, + ); + if ( + (await commitJobResult( + redis, + target, + { stdout: 'mutation settled' }, + 60, + )) !== 'committed' + ) + throw new Error('handoff did not win'); + // This represents Stop during post-handoff egress cleanup. It must no + // longer turn the applied mutation into an acknowledged cancellation. + expect(await requestJobCancellation(redis, target, 60)).toBe(false); + return settlement; + }, + }); + void completion.catch(() => undefined); + const assignment = await store.lease(workerId, incarnationId, 1_000); + if (assignment == null) throw new Error('Missing assignment'); + await store.settle(workerId, assignment.assignmentId, { + protocolVersion: 1, + incarnationId, + generation: assignment.generation, + leaseToken: assignment.leaseToken, + status: 'fulfilled', + result: { + session_id: 'handoff-session', + language: 'bash', + version: '5.2', + files: [], + }, + }); + if (outcome !== 'commit') { + await expect(completion).rejects.toThrow('handoff did not win'); + await expect(store.dispatch(dispatchArgs)).rejects.toMatchObject({ + code: 'WORKSPACE_QUARANTINED', + }); + } else { + await expect(completion).resolves.toMatchObject({ + status: 'fulfilled', + }); + expect(await readCommittedJobResult(redis, target)).toEqual({ + result: { stdout: 'mutation settled' }, + }); + } + }); + +test('concurrent stalled-job redelivery claims at most one sandbox execution', async () => { + let executions = 0; + const attempt = async () => { + const claim = await claimJobExecution(redis, target, 60); + if (claim.status === 'claimed') executions += 1; + return claim; + }; + const results = await Promise.allSettled([attempt(), attempt()]); + expect(executions).toBe(1); + expect(results.filter(result => result.status === 'fulfilled')).toHaveLength( + 1, + ); + expect(results.filter(result => result.status === 'rejected')).toHaveLength( + 1, + ); + await expect(attempt()).rejects.toThrow('already claimed'); + expect(executions).toBe(1); + await commitJobResult(redis, target, { stdout: 'first result' }, 60); + expect(await attempt()).toEqual({ + status: 'completed', + result: { stdout: 'first result' }, + }); + expect(executions).toBe(1); + expect( + await commitJobResult(redis, target, { stdout: 'different result' }, 60), + ).toBe('already_completed'); + expect(await readCommittedJobResult(redis, target)).toEqual({ + result: { stdout: 'first result' }, + }); +}); + +test('a lost execution-claim reply never authorizes a second attempt', async () => { + const lostReply = { + eval: async (...args: Parameters) => { + await redis.eval(...args); + throw new Error('claim reply lost'); + }, + } as unknown as typeof redis; + await expect(claimJobExecution(lostReply, target, 60)).rejects.toThrow( + 'claim reply lost', + ); + await expect(claimJobExecution(redis, target, 60)).rejects.toThrow( + 'already claimed', + ); +}); + +test('cancel-before-claim and missing completion payload fail closed', async () => { + await requestJobCancellation(redis, target, 60); + await expect(claimJobExecution(redis, target, 60)).rejects.toThrow( + 'cancelled', + ); + const completedTarget = { ...target, jobId: 'missing-payload-claim' }; + await commitJobResult(redis, completedTarget, { stdout: 'done' }, 60); + await redis.del( + `${jobCancellationInternals.cancellationKey(completedTarget)}:result`, + ); + await expect(claimJobExecution(redis, completedTarget, 60)).rejects.toThrow( + 'refusing re-execution', + ); +}); + +for (const corrupt of [false, true]) + test(`invalid committed result fails immediately without Redis retries (corrupt=${corrupt})`, async () => { + await commitJobResult(redis, target, { stdout: 'done' }, 60); + const key = jobCancellationInternals.cancellationKey(target); + if (corrupt) await redis.set(`${key}:result`, '{invalid'); + else await redis.del(`${key}:result`); + let calls = 0; + const commands = { + eval: (...args: Parameters) => { + calls += 1; + return redis.eval(...args); + }, + } as unknown as typeof redis; + await expect( + fenceJobCancellation({ + commands, + target, + ttlSeconds: 60, + deadlineAtMs: Date.now() + 30_000, + }), + ).rejects.toThrow(); + expect(calls).toBe(1); + }); + +test('completion retention includes the API producer across timeout configuration drift', async () => { + const ttl = jobCancellationRetentionSeconds(30_000, 430); + expect(ttl).toBe(430); + expect(jobCancellationRetentionSeconds(300_000, 430)).toBe(780); + await commitJobResult(redis, target, { stdout: 'done' }, ttl); + const key = jobCancellationInternals.cancellationKey(target); + expect(await redis.ttl(key)).toBeGreaterThanOrEqual(429); + expect(await redis.ttl(`${key}:result`)).toBeGreaterThanOrEqual(429); +}); + +test('a late Stop renews completion evidence along with its request tombstone', async () => { + await commitJobResult(redis, target, { stdout: 'done' }, 1); + expect(await requestJobCancellation(redis, target, 60)).toBe(false); + const key = jobCancellationInternals.cancellationKey(target); + expect(await redis.ttl(key)).toBeGreaterThanOrEqual(59); + expect(await redis.ttl(`${key}:result`)).toBeGreaterThanOrEqual(59); +}); + +test('retention renewal does not lose subsecond time to rounded TTL readings', async () => { + await commitJobResult(redis, target, { stdout: 'done' }, 60); + const key = jobCancellationInternals.cancellationKey(target); + await redis.pexpire(key, 59_900); + const expiration = async () => + Number( + await redis.eval( + ` + local now = redis.call('TIME') + return tonumber(now[1]) * 1000 + math.floor(tonumber(now[2]) / 1000) + redis.call('PTTL', KEYS[1]) + `, + 1, + key, + ), + ); + const before = await expiration(); + await requestJobCancellation(redis, target, 60); + expect(await expiration()).toBeGreaterThan(before); +}); + +test('fencing returns the committed result without a vulnerable second Redis read', async () => { + const result = { stdout: 'one committed effect' }; + await commitJobResult(redis, target, result, 60); + let calls = 0; + const connectionDropsAfterDecision = { + eval: async (...args: Parameters) => { + calls += 1; + return redis.eval(...args); + }, + get: async () => { + throw new Error('connection lost after decision'); + }, + mget: async () => { + throw new Error('connection lost after decision'); + }, + } as unknown as typeof redis; + expect( + await fenceJobCancellation({ + commands: connectionDropsAfterDecision, + target, + ttlSeconds: 60, + deadlineAtMs: Date.now() + 1_000, + }), + ).toEqual({ status: 'completed', result }); + expect(calls).toBe(1); +}); + +test('disconnect returns a known completed result without waiting for a lost queue event', async () => { + const result = { stdout: 'done' }; + await commitJobResult(redis, target, result, 60); + const registry = new JobCancellationRegistry(redis); + const controller = new AbortController(); + controller.abort(); + const job = { + id: target.jobId, + queueName: target.queueName, + waitUntilFinished: () => new Promise(() => {}), + } as unknown as Parameters[0]['job']; + try { + expect( + await waitForJobWithCancellation({ + commands: redis, + registry, + job, + events: {} as Parameters< + typeof waitForJobWithCancellation + >[0]['events'], + timeoutMs: 1_000, + cancellationTtlSeconds: 60, + signal: controller.signal, + }), + ).toEqual(result); + } finally { + await registry.close(); + } +}); + +test('durable cancellation wins even before its subscriber notification arrives', async () => { + expect(await requestJobCancellation(redis, target, 60)).toBe(true); + expect(await commitJobResult(redis, target, { stdout: 'late' }, 60)).toBe( + 'cancelled', + ); + expect(await readCommittedJobResult(redis, target)).toBeUndefined(); +}); + +test('committed results reject late Stop and survive a lost BullMQ completion reply', async () => { + const result = { stdout: 'one mutation', files: [] }; + expect(await commitJobResult(redis, target, result, 60)).toBe('committed'); + expect(await requestJobCancellation(redis, target, 60)).toBe(false); + expect(await readCommittedJobResult(redis, target)).toEqual({ result }); + expect( + await redis.get(jobCancellationInternals.cancellationKey(target)), + ).toBe('completed'); + const registry = new JobCancellationRegistry(redis); + const controller = new AbortController(); + try { + await registry.register(target, controller); + expect(controller.signal.aborted).toBe(false); + } finally { + await registry.close(); + } +}); + +test('concurrent cancellation and completion have exactly one winner', async () => { + const [cancelled, committed] = await Promise.all([ + requestJobCancellation(redis, target, 60), + commitJobResult(redis, target, { stdout: 'result' }, 60), + ]); + expect(Number(cancelled) + Number(committed === 'committed')).toBe(1); +}); + +test('a missing committed result fails closed instead of re-executing', async () => { + await commitJobResult(redis, target, { stdout: 'already applied' }, 60); + await redis.del(`${jobCancellationInternals.cancellationKey(target)}:result`); + await expect(readCommittedJobResult(redis, target)).rejects.toThrow( + 'refusing re-execution', + ); +}); + +test('an enqueue failure can recover a result that won cancellation fencing', async () => { + const result = { stdout: 'effect already applied' }; + await commitJobResult(redis, target, result, 60); + expect( + await fenceJobCancellation({ + commands: redis, + target, + ttlSeconds: 60, + deadlineAtMs: Date.now() + 5_000, + }), + ).toEqual({ status: 'completed', result }); + expect(await readCommittedJobResult(redis, target)).toEqual({ result }); +}); + +test('enqueue fencing still recovers completion after the original deadline', async () => { + await commitJobResult(redis, target, { stdout: 'done' }, 60); + expect( + await fenceJobCancellation({ + commands: redis, + target, + ttlSeconds: 60, + deadlineAtMs: Date.now() - 1_000, + }), + ).toEqual({ status: 'completed', result: { stdout: 'done' } }); +}); + +test('Redis rejects commitment when recovery happens after the producer deadline', async () => { + const delayed = { + eval: async (...args: Parameters) => { + await new Promise(resolve => setTimeout(resolve, 150)); + return redis.eval(...args); + }, + } as unknown as typeof redis; + await expect( + commitJobResult(delayed, target, { stdout: 'late' }, 60, Date.now() + 100), + ).rejects.toThrow('exceeded its deadline'); + expect(await readCommittedJobResult(redis, target)).toBeUndefined(); +}); + +test('a timely durable commit remains successful when only its acknowledgement is late', async () => { + const delayedReply = { + eval: async (...args: Parameters) => { + const value = await redis.eval(...args); + await new Promise(resolve => setTimeout(resolve, 150)); + return value; + }, + } as unknown as typeof redis; + expect( + await commitJobResult( + delayedReply, + target, + { stdout: 'committed' }, + 60, + Date.now() + 100, + ), + ).toBe('committed'); + expect(await readCommittedJobResult(redis, target)).toEqual({ + result: { stdout: 'committed' }, + }); +}); + +for (const failedStage of ['subscription', 'completion'] as const) { + test(`a lost ${failedStage} reply recovers the committed result instead of reporting failure`, async () => { + const result = { stdout: 'already applied once' }; + await commitJobResult(redis, target, result, 60); + const registry = new JobCancellationRegistry(redis); + if (failedStage === 'subscription') + registry.register = async () => { + throw new Error('lost reply'); + }; + const job = { + id: target.jobId, + queueName: target.queueName, + waitUntilFinished: () => Promise.reject(new Error('lost result event')), + } as unknown as Parameters[0]['job']; + try { + expect( + await waitForJobWithCancellation({ + commands: redis, + registry, + job, + events: {} as Parameters< + typeof waitForJobWithCancellation + >[0]['events'], + timeoutMs: 1_000, + cancellationTtlSeconds: 60, + }), + ).toEqual(result); + } finally { + await registry.close(); + } + }); +} diff --git a/service/src/job-cancellation.test.ts b/service/src/job-cancellation.test.ts new file mode 100644 index 00000000..e23186e4 --- /dev/null +++ b/service/src/job-cancellation.test.ts @@ -0,0 +1,599 @@ +import { expect, test } from 'bun:test'; +import { EventEmitter } from 'node:events'; +import type IORedis from 'ioredis'; +import type { Job, QueueEvents } from 'bullmq'; +import { + CLIENT_DISCONNECT_REASON, + JobCancellationRegistry, + jobResultCommitFailure, + jobCancellationInternals, + removeJobIfWaiting, + requestJobCancellation, + throwIfJobAborted, + waitForJobWithCancellation, + fenceJobCancellation, +} from './job-cancellation'; + +class FakeSubscriber extends EventEmitter { + subscribed?: string; + closed = false; + subscribeFailures = 0; + + async subscribe(channel: string): Promise { + if (this.subscribeFailures > 0) { + this.subscribeFailures -= 1; + throw new Error('subscriber unavailable'); + } + this.subscribed = channel; + return 1; + } + + async quit(): Promise<'OK'> { + this.closed = true; + return 'OK'; + } + + disconnect(): void { + this.closed = true; + } +} + +class FakeTransaction { + readonly operations: unknown[][] = []; + + set(...args: unknown[]): this { + this.operations.push(['set', ...args]); + return this; + } + + publish(...args: unknown[]): this { + this.operations.push(['publish', ...args]); + return this; + } + + async exec(): Promise> { + return this.operations.map(() => [null, 'OK']); + } +} + +class FakeRedis { + readonly subscriber = new FakeSubscriber(); + duplicateCalls = 0; + readonly existing = new Set(); + readonly deleted: string[] = []; + readonly transactions: FakeTransaction[] = []; + mgetFailures = 0; + cancellationFailures = 0; + cancellationAttempts = 0; + + duplicate(): FakeSubscriber { + this.duplicateCalls += 1; + return this.subscriber; + } + + async get(key: string): Promise { + return this.existing.has(key) ? '1' : null; + } + + async eval( + _script: string, + _keys: number, + key: string, + _resultKey: string, + ttl: number, + channel: string, + payload: string, + ): Promise { + this.cancellationAttempts += 1; + if (this.cancellationFailures-- > 0) throw new Error('Redis unavailable'); + const transaction = this.multi(); + transaction.set(key, '1', 'EX', ttl); + transaction.publish(channel, payload); + await transaction.exec(); + return [1]; + } + + async mget(...keys: string[]): Promise> { + if (this.mgetFailures > 0) { + this.mgetFailures -= 1; + throw new Error('command connection unavailable'); + } + return keys.map(key => (this.existing.has(key) ? '1' : null)); + } + + async del(key: string): Promise { + this.deleted.push(key); + this.existing.delete(key); + return 1; + } + + multi(): FakeTransaction { + const transaction = new FakeTransaction(); + this.transactions.push(transaction); + return transaction; + } +} + +function redis(fake: FakeRedis): IORedis { + return fake as unknown as IORedis; +} + +test('idle registries allocate no subscriber connection', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + + await registry.close(); + + expect(fake.duplicateCalls).toBe(0); +}); + +test('shutdown disconnects a subscriber whose startup is still waiting for Redis', async () => { + const fake = new FakeRedis(); + fake.subscriber.subscribe = async () => new Promise(() => {}); + const registry = new JobCancellationRegistry(redis(fake)); + void registry + .register( + { queueName: 'other', jobId: 'shutdown-startup' }, + new AbortController(), + ) + .catch(() => undefined); + await registry.close(); + expect(fake.subscriber.closed).toBe(true); + expect(fake.subscriber.listenerCount('message')).toBe(0); +}, 1_000); + +test('failed subscription startup removes handlers before a bounded retry', async () => { + const fake = new FakeRedis(); + fake.subscriber.subscribeFailures = 1; + const registry = new JobCancellationRegistry(redis(fake)); + const first = new AbortController(); + + await expect( + registry.register({ queueName: 'other', jobId: 'job-failed-start' }, first), + ).rejects.toThrow('subscriber unavailable'); + expect(fake.subscriber.listenerCount('message')).toBe(0); + expect(fake.subscriber.listenerCount('ready')).toBe(0); + expect(fake.subscriber.listenerCount('error')).toBe(0); + + const second = new AbortController(); + await registry.register({ queueName: 'other', jobId: 'job-retry' }, second); + expect(fake.duplicateCalls).toBe(2); + expect(fake.subscriber.listenerCount('message')).toBe(1); + await registry.close(); +}); + +test('registry catches durable cancellation before subscriber registration', async () => { + const fake = new FakeRedis(); + const target = { queueName: 'other', jobId: 'job-1' }; + fake.existing.add(jobCancellationInternals.cancellationKey(target)); + const registry = new JobCancellationRegistry(redis(fake)); + const controller = new AbortController(); + + await registry.register(target, controller); + + expect(controller.signal.aborted).toBe(true); + expect(controller.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + expect(fake.subscriber.subscribed).toBe(jobCancellationInternals.channel); + await registry.unregister(target); + await registry.close(); + expect(fake.subscriber.closed).toBe(true); +}); + +test('one pubsub listener cancels only the matching active job', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const first = new AbortController(); + const second = new AbortController(); + await registry.register({ queueName: 'other', jobId: 'job-1' }, first); + await registry.register({ queueName: 'other', jobId: 'job-2' }, second); + + fake.subscriber.emit( + 'message', + jobCancellationInternals.channel, + JSON.stringify({ queueName: 'other', jobId: 'job-2' }), + ); + + expect(first.signal.aborted).toBe(false); + expect(second.signal.aborted).toBe(true); + await registry.close(); +}); + +test('one pubsub listener wakes every local waiter for the same job', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const target = { queueName: 'other', jobId: 'job-shared' }; + const first = new AbortController(); + const second = new AbortController(); + await registry.register(target, first); + await registry.register(target, second); + + fake.subscriber.emit( + 'message', + jobCancellationInternals.channel, + JSON.stringify(target), + ); + + expect(first.signal.aborted).toBe(true); + expect(second.signal.aborted).toBe(true); + expect(fake.duplicateCalls).toBe(1); + await registry.close(); +}); + +test('unregistering one local waiter preserves other waiters for the job', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const target = { queueName: 'other', jobId: 'job-shared-unregister' }; + const first = new AbortController(); + const second = new AbortController(); + await registry.register(target, first); + await registry.register(target, second); + await registry.unregister(target, first); + + fake.subscriber.emit( + 'message', + jobCancellationInternals.channel, + JSON.stringify(target), + ); + + expect(first.signal.aborted).toBe(false); + expect(second.signal.aborted).toBe(true); + await registry.close(); +}); + +test('subscriber reconnect reconciles active jobs against durable markers', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const target = { queueName: 'other', jobId: 'job-reconnect' }; + const controller = new AbortController(); + await registry.register(target, controller); + fake.existing.add(jobCancellationInternals.cancellationKey(target)); + + fake.subscriber.emit('ready'); + await new Promise(resolve => setTimeout(resolve, 10)); + + expect(controller.signal.aborted).toBe(true); + expect(controller.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + await registry.close(); +}); + +test('subscriber reconnect retries durable-marker reconciliation', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const target = { queueName: 'other', jobId: 'job-retry-reconcile' }; + const controller = new AbortController(); + await registry.register(target, controller); + fake.existing.add(jobCancellationInternals.cancellationKey(target)); + fake.mgetFailures = 1; + + fake.subscriber.emit('ready'); + await new Promise(resolve => setTimeout(resolve, 150)); + + expect(controller.signal.aborted).toBe(true); + expect(controller.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + await registry.close(); +}); + +test('terminal subscriber disconnect rebuilds the subscription and reconciles markers', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const target = { queueName: 'other', jobId: 'job-terminal-reconnect' }; + const controller = new AbortController(); + await registry.register(target, controller); + fake.existing.add(jobCancellationInternals.cancellationKey(target)); + + fake.subscriber.emit('end'); + await new Promise(resolve => setTimeout(resolve, 10)); + + expect(fake.duplicateCalls).toBe(2); + expect(controller.signal.aborted).toBe(true); + expect(controller.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + await registry.close(); +}); + +test('cancellation writes a durable marker before publishing', async () => { + const fake = new FakeRedis(); + const target = { queueName: 'other', jobId: 'job-3' }; + + await requestJobCancellation(redis(fake), target, 42); + + expect(fake.transactions).toHaveLength(1); + expect(fake.transactions[0]?.operations).toEqual([ + ['set', jobCancellationInternals.cancellationKey(target), '1', 'EX', 42], + ['publish', jobCancellationInternals.channel, JSON.stringify(target)], + ]); +}); + +test('result commit barrier rejects cancellation observed after execution', () => { + const controller = new AbortController(); + expect(() => throwIfJobAborted(controller.signal)).not.toThrow(); + controller.abort(CLIENT_DISCONNECT_REASON); + expect(() => throwIfJobAborted(controller.signal)).toThrow( + CLIENT_DISCONNECT_REASON, + ); +}); + +test('result cleanup maps late cancellation to stable worker failures', () => { + const disconnected = new AbortController(); + disconnected.abort(CLIENT_DISCONNECT_REASON); + expect(jobResultCommitFailure(disconnected.signal, 30_000)?.message).toBe( + 'Job cancelled after client disconnected', + ); + + const deadline = new AbortController(); + deadline.abort('deadline'); + expect(jobResultCommitFailure(deadline.signal, 30_000)?.message).toBe( + 'Job timed out after 30000ms', + ); + expect( + jobResultCommitFailure(new AbortController().signal, 30_000), + ).toBeUndefined(); +}); + +test('disconnect frees a waiting job and rejects promptly', async () => { + const fake = new FakeRedis(); + const controller = new AbortController(); + let removed = false; + const never = new Promise(() => {}); + const job = { + id: 'job-4', + queueName: 'other', + waitUntilFinished: () => never, + getState: async () => 'waiting', + remove: async () => { + removed = true; + }, + } as unknown as Job; + + const registry = new JobCancellationRegistry(redis(fake)); + const waiting = waitForJobWithCancellation({ + commands: redis(fake), + registry, + job, + events: {} as QueueEvents, + timeoutMs: 60_000, + cancellationTtlSeconds: 120, + signal: controller.signal, + }); + controller.abort(CLIENT_DISCONNECT_REASON); + + await expect(waiting).rejects.toMatchObject({ name: 'AbortError' }); + expect(removed).toBe(true); + expect(fake.cancellationAttempts).toBe(1); + expect(fake.transactions[0]?.operations[0]).toEqual([ + 'set', + jobCancellationInternals.cancellationKey({ + queueName: 'other', + jobId: 'job-4', + }), + '1', + 'EX', + 120, + ]); + await registry.close(); +}); + +test('registration failure fences and removes the already-enqueued job', async () => { + const fake = new FakeRedis(); + fake.subscriber.subscribeFailures = 1; + let removed = false; + const job = { + id: 'job-register-failure', + queueName: 'other', + waitUntilFinished: () => new Promise(() => {}), + getState: async () => 'waiting', + remove: async () => { + removed = true; + }, + } as unknown as Job; + const registry = new JobCancellationRegistry(redis(fake)); + + await expect( + waitForJobWithCancellation({ + commands: redis(fake), + registry, + job, + events: {} as QueueEvents, + timeoutMs: 60_000, + cancellationTtlSeconds: 120, + }), + ).rejects.toThrow('subscriber unavailable'); + + expect(removed).toBe(true); + expect(fake.transactions[0]?.operations[0]).toEqual([ + 'set', + jobCancellationInternals.cancellationKey({ + queueName: 'other', + jobId: 'job-register-failure', + }), + '1', + 'EX', + 120, + ]); + await registry.close(); +}); + +test('a result rejection is owned while subscription registration is pending', async () => { + const fake = new FakeRedis(); + let release!: () => void; + fake.subscriber.subscribe = async () => { + await new Promise(resolve => { + release = resolve; + }); + return 1; + }; + const registry = new JobCancellationRegistry(redis(fake)); + const job = { + id: 'pending-registration', + queueName: 'other', + waitUntilFinished: () => Promise.reject(new Error('completion timeout')), + getState: async () => 'active', + remove: async () => {}, + } as unknown as Job; + const waiting = waitForJobWithCancellation({ + commands: redis(fake), + registry, + job, + events: {} as QueueEvents, + timeoutMs: 1_000, + cancellationTtlSeconds: 60, + }); + const rejection = waiting.catch((error: Error) => error); + // An unowned rejection fails the test runner on this event-loop turn. + await new Promise(resolve => setImmediate(resolve)); + release(); + expect(await rejection).toMatchObject({ message: 'completion timeout' }); + await registry.close(); +}); + +test('external cancellation frees a waiting job before rejecting its waiter', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + let removed = false; + const job = { + id: 'job-external-waiting', + queueName: 'other', + waitUntilFinished: () => new Promise(() => {}), + getState: async () => 'waiting', + remove: async () => { + removed = true; + }, + } as unknown as Job; + const waiting = waitForJobWithCancellation({ + commands: redis(fake), + registry, + job, + events: {} as QueueEvents, + timeoutMs: 60_000, + cancellationTtlSeconds: 120, + }); + await new Promise(resolve => setImmediate(resolve)); + + fake.subscriber.emit( + 'message', + jobCancellationInternals.channel, + JSON.stringify({ queueName: 'other', jobId: 'job-external-waiting' }), + ); + + await expect(waiting).rejects.toMatchObject({ name: 'AbortError' }); + expect(removed).toBe(true); + await registry.close(); +}); + +test('a separate cancellation request wakes the original job waiter', async () => { + const fake = new FakeRedis(); + const registry = new JobCancellationRegistry(redis(fake)); + const never = new Promise(() => {}); + const job = { + id: 'job-external-cancel', + queueName: 'other', + waitUntilFinished: () => never, + getState: async () => 'active', + remove: async () => undefined, + } as unknown as Job; + + const waiting = waitForJobWithCancellation({ + commands: redis(fake), + registry, + job, + events: {} as QueueEvents, + timeoutMs: 60_000, + cancellationTtlSeconds: 120, + }); + await new Promise(resolve => setImmediate(resolve)); + fake.subscriber.emit( + 'message', + jobCancellationInternals.channel, + JSON.stringify({ queueName: 'other', jobId: 'job-external-cancel' }), + ); + + await expect(waiting).rejects.toMatchObject({ name: 'AbortError' }); + expect(fake.deleted).toEqual([]); + await registry.close(); +}); + +test('queued removal never removes an active job', async () => { + let removed = false; + const job = { + getState: async () => 'active' as const, + remove: async () => { + removed = true; + }, + }; + + expect(await removeJobIfWaiting(job)).toBe(false); + expect(removed).toBe(false); +}); + +test('a failed cancellation write retains ownership until a durable retry succeeds', async () => { + const fake = new FakeRedis(); + fake.cancellationFailures = 2; + const target = { queueName: 'other', jobId: 'ambiguous-enqueue' }; + let released = false; + const fencing = fenceJobCancellation({ + commands: redis(fake), + target, + ttlSeconds: 60, + deadlineAtMs: Date.now() + 500, + }).then(() => { + released = true; + }); + await new Promise(resolve => setTimeout(resolve, 10)); + expect(released).toBe(false); + await fencing; + expect(fake.cancellationAttempts).toBe(3); + expect(released).toBe(true); +}); + +test('an unavailable Redis cannot release ownership before the fixed job deadline', async () => { + const fake = new FakeRedis(); + fake.cancellationFailures = 1_000; + const startedAt = Date.now(); + await fenceJobCancellation({ + commands: redis(fake), + target: { queueName: 'other', jobId: 'offline' }, + ttlSeconds: 60, + deadlineAtMs: startedAt + 80, + }); + expect(Date.now() - startedAt).toBeGreaterThanOrEqual(80); + expect(fake.cancellationAttempts).toBeLessThanOrEqual(4); +}); + +test('a pending Redis write allocates no retry backlog and waits until the deadline', async () => { + const fake = new FakeRedis(); + let calls = 0; + fake.eval = async () => { + calls += 1; + return new Promise(() => {}); + }; + const startedAt = Date.now(); + await fenceJobCancellation({ + commands: redis(fake), + target: { queueName: 'other', jobId: 'pending-write' }, + ttlSeconds: 60, + deadlineAtMs: startedAt + 40, + }); + expect(Date.now() - startedAt).toBeGreaterThanOrEqual(39); + expect(calls).toBe(1); +}); + +test('queued removal frees waiting capacity and tolerates an activation race', async () => { + let removals = 0; + expect( + await removeJobIfWaiting({ + getState: async () => 'waiting', + remove: async () => { + removals += 1; + }, + }), + ).toBe(true); + expect( + await removeJobIfWaiting({ + getState: async () => 'waiting', + remove: async () => { + removals += 1; + throw new Error('job is active'); + }, + }), + ).toBe(false); + expect(removals).toBe(2); +}); diff --git a/service/src/job-cancellation.ts b/service/src/job-cancellation.ts new file mode 100644 index 00000000..4da3253d --- /dev/null +++ b/service/src/job-cancellation.ts @@ -0,0 +1,664 @@ +import type IORedis from 'ioredis'; +import type { Job, QueueEvents } from 'bullmq'; + +const JOB_CANCELLATION_PREFIX = 'codeapi:job-cancellation:v1'; +const JOB_CANCELLATION_CHANNEL = `${JOB_CANCELLATION_PREFIX}:events`; +export const CLIENT_DISCONNECT_REASON = 'client_disconnected'; +export const JOB_CANCELLED_MESSAGE = 'Job cancelled after client disconnected'; + +interface JobTarget { + queueName: string; + jobId: string; +} + +function targetKey(target: JobTarget): string { + return `${target.queueName}:${target.jobId}`; +} + +function cancellationKey(target: JobTarget): string { + return `${JOB_CANCELLATION_PREFIX}:${encodeURIComponent( + target.queueName, + )}:${encodeURIComponent(target.jobId)}`; +} + +function parseTarget(raw: string): JobTarget | undefined { + try { + const parsed = JSON.parse(raw) as Partial; + if ( + typeof parsed.queueName !== 'string' || + parsed.queueName.length === 0 || + parsed.queueName.length > 256 || + typeof parsed.jobId !== 'string' || + parsed.jobId.length === 0 || + parsed.jobId.length > 256 + ) { + return undefined; + } + return { queueName: parsed.queueName, jobId: parsed.jobId }; + } catch { + return undefined; + } +} + +/** + * Cross-process cancellation for BullMQ work. + * + * The durable marker closes the publish-before-subscribe race while one + * process-wide pub/sub connection makes active cancellation O(events), not + * O(active jobs) Redis polling. Only explicitly cancellable replay jobs use + * this path, so ordinary queue traffic pays no extra Redis round trips. + */ +export class JobCancellationRegistry { + private subscriber?: IORedis; + private readonly controllers = new Map< + string, + { target: JobTarget; controllers: Set } + >(); + private startPromise?: Promise; + private readonly subscriberEndHandlers = new WeakMap void>(); + private reconcileTimer?: ReturnType; + private subscriberRestartTimer?: ReturnType; + private reconcileRetryMs = 100; + private closed = false; + + constructor(private readonly commands: IORedis) {} + + private readonly onSubscriberError = (): void => { + // ioredis reconnects using the shared policy. The listener prevents a + // transient subscriber outage from becoming an uncaught process error. + }; + + private readonly onSubscriberReady = (): void => { + this.scheduleReconcile(0); + }; + + private readonly onSubscriberMessage = ( + channel: string, + raw: string, + ): void => { + if (channel !== JOB_CANCELLATION_CHANNEL) return; + const target = parseTarget(raw); + if (target == null) return; + for (const controller of this.controllers.get(targetKey(target)) + ?.controllers ?? []) { + controller.abort(CLIENT_DISCONNECT_REASON); + } + }; + + private detachSubscriber(subscriber: IORedis): void { + subscriber.removeListener('error', this.onSubscriberError); + subscriber.removeListener('ready', this.onSubscriberReady); + subscriber.removeListener('message', this.onSubscriberMessage); + const onEnd = this.subscriberEndHandlers.get(subscriber); + if (onEnd != null) subscriber.removeListener('end', onEnd); + this.subscriberEndHandlers.delete(subscriber); + } + + private restartAfterTerminalDisconnect(subscriber: IORedis): void { + if (this.closed || this.subscriber !== subscriber) return; + this.detachSubscriber(subscriber); + this.subscriber = undefined; + this.startPromise = undefined; + if (this.controllers.size === 0) return; + void this.start().then( + () => this.scheduleReconcile(0), + () => this.scheduleSubscriberRestart(), + ); + } + + private scheduleSubscriberRestart(): void { + if ( + this.closed || + this.controllers.size === 0 || + this.startPromise != null || + this.subscriberRestartTimer != null + ) + return; + const retryMs = this.reconcileRetryMs; + this.reconcileRetryMs = Math.min(2_000, retryMs * 2); + this.subscriberRestartTimer = setTimeout(() => { + this.subscriberRestartTimer = undefined; + if ( + this.closed || + this.controllers.size === 0 || + this.startPromise != null + ) + return; + void this.start().then( + () => { + this.reconcileRetryMs = 100; + this.scheduleReconcile(0); + }, + () => this.scheduleSubscriberRestart(), + ); + }, retryMs); + } + + private async reconcile(): Promise { + const entries = [...this.controllers.values()]; + if (entries.length === 0) return; + const cancelled = await this.commands.mget( + ...entries.map(({ target }) => cancellationKey(target)), + ); + cancelled.forEach((value, index) => { + if (value === '1') { + for (const controller of entries[index]?.controllers ?? []) { + controller.abort(CLIENT_DISCONNECT_REASON); + } + } + }); + } + + private scheduleReconcile(delayMs: number): void { + if ( + this.closed || + this.controllers.size === 0 || + this.reconcileTimer != null + ) { + return; + } + this.reconcileTimer = setTimeout(() => { + this.reconcileTimer = undefined; + void this.reconcile().then( + () => { + this.reconcileRetryMs = 100; + }, + () => { + const retryMs = this.reconcileRetryMs; + this.reconcileRetryMs = Math.min(2_000, retryMs * 2); + this.scheduleReconcile(retryMs); + }, + ); + }, delayMs); + } + + private start(): Promise { + if (this.closed) { + return Promise.reject(new Error('Job cancellation registry is closed')); + } + if (this.startPromise != null) return this.startPromise; + const starting = (async (): Promise => { + const subscriber = this.commands.duplicate(); + this.subscriber = subscriber; + subscriber.on('error', this.onSubscriberError); + subscriber.on('ready', this.onSubscriberReady); + subscriber.on('message', this.onSubscriberMessage); + const onEnd = (): void => this.restartAfterTerminalDisconnect(subscriber); + this.subscriberEndHandlers.set(subscriber, onEnd); + subscriber.on('end', onEnd); + try { + await subscriber.subscribe(JOB_CANCELLATION_CHANNEL); + } catch (error) { + this.detachSubscriber(subscriber); + if (this.subscriber === subscriber) this.subscriber = undefined; + subscriber.disconnect(false); + throw error; + } + })(); + this.startPromise = starting; + void starting.catch(() => { + if (this.startPromise === starting) this.startPromise = undefined; + }); + return starting; + } + + async register( + target: JobTarget, + controller: AbortController, + ): Promise { + const key = targetKey(target); + const entry = this.controllers.get(key) ?? { + target, + controllers: new Set(), + }; + entry.controllers.add(controller); + this.controllers.set(key, entry); + try { + await this.start(); + if ((await this.commands.get(cancellationKey(target))) === '1') { + controller.abort(CLIENT_DISCONNECT_REASON); + } + } catch (error) { + entry.controllers.delete(controller); + if (entry.controllers.size === 0) this.controllers.delete(key); + throw error; + } + } + + async unregister( + target: JobTarget, + controller?: AbortController, + ): Promise { + const key = targetKey(target); + const entry = this.controllers.get(key); + if (controller == null) { + this.controllers.delete(key); + } else if (entry != null) { + entry.controllers.delete(controller); + if (entry.controllers.size === 0) this.controllers.delete(key); + } + // Markers expire by TTL. Deleting one here can erase the only evidence + // needed by another replica whose subscriber was reconnecting. + } + + async close(): Promise { + this.closed = true; + this.controllers.clear(); + if (this.reconcileTimer != null) clearTimeout(this.reconcileTimer); + this.reconcileTimer = undefined; + if (this.subscriberRestartTimer != null) + clearTimeout(this.subscriberRestartTimer); + this.subscriberRestartTimer = undefined; + const subscriber = this.subscriber; + this.subscriber = undefined; + this.startPromise = undefined; + if (subscriber == null) return; + this.detachSubscriber(subscriber); + // This socket only carries notifications. Disconnect it before awaiting + // anything: subscribe() may be queued through an indefinite Redis outage. + subscriber.disconnect(false); + } +} + +async function cancelJobInRedis( + commands: IORedis, + target: JobTarget, + ttlSeconds: number, + includeResult: boolean, +): Promise { + // Cancellation and result publication have ONE durable winner. Pub/sub is + // only a notification; it must not decide whether Stop was accepted. + return commands.eval( + ` + local state = redis.call('GET', KEYS[1]) + if state == 'completed' then + -- Keep completion evidence at least as long as the requesting process + -- requires, even across API/worker config differences. Attached request + -- tombstones never renew independently of this decision. + local requestedTtlMs = tonumber(ARGV[1]) * 1000 + for i = 1, 2 do + if redis.call('PTTL', KEYS[i]) < requestedTtlMs then + redis.call('PEXPIRE', KEYS[i], requestedTtlMs) + end + end + if ARGV[4] == '1' then return {0, redis.call('GET', KEYS[2])} end + return {0} + end + if state and state ~= '1' then return {-1} end + redis.call('SET', KEYS[1], '1', 'EX', ARGV[1]) + redis.call('PUBLISH', ARGV[2], ARGV[3]) + return {1} + `, + 2, + cancellationKey(target), + `${cancellationKey(target)}:result`, + Math.max(1, ttlSeconds), + JOB_CANCELLATION_CHANNEL, + JSON.stringify(target), + includeResult ? '1' : '0', + ); +} + +export async function requestJobCancellation( + commands: IORedis, + target: JobTarget, + ttlSeconds: number, +): Promise { + const decision = await cancelJobInRedis(commands, target, ttlSeconds, false); + if (!Array.isArray(decision) || ![0, 1].includes(decision[0])) { + throw new Error('Invalid durable cancellation decision'); + } + return decision[0] === 1; +} + +export type JobFenceOutcome = + | { status: 'cancelled' | 'expired' } + | { status: 'completed'; result: T }; + +function decodeCommittedResult(value: unknown): { result: T } { + if (typeof value !== 'string') { + throw new Error( + 'Committed programmatic result expired; refusing re-execution', + ); + } + const decoded: unknown = JSON.parse(value); + if ( + decoded == null || + typeof decoded !== 'object' || + !Object.prototype.hasOwnProperty.call(decoded, 'result') + ) { + throw new Error( + 'Invalid committed programmatic result; refusing re-execution', + ); + } + return decoded as { result: T }; +} + +export function jobCancellationRetentionSeconds( + localTimeoutMs: number, + producerTtlSeconds = 0, +): number { + return Math.max( + Math.ceil(localTimeoutMs / 1_000) * 2 + 180, + Number.isFinite(producerTtlSeconds) ? producerTtlSeconds : 0, + ); +} + +/** Retain the actual result so a BullMQ retry after a lost completion reply + * cannot repeat sandbox mutations. Keep the status small: reconnect MGETs must + * never load every active job's output into each API/worker replica. */ +export async function commitJobResult( + commands: IORedis, + target: JobTarget, + result: T, + ttlSeconds: number, + deadlineAtMs = Number.MAX_SAFE_INTEGER, +): Promise<'committed' | 'cancelled' | 'already_completed'> { + const serialized = JSON.stringify({ result }); + if (Buffer.byteLength(serialized) > 16 * 1024 * 1024) { + throw new Error('Programmatic completion exceeds the 16 MiB result limit'); + } + const decision = await commands.eval( + ` + local state = redis.call('GET', KEYS[1]) + if state == '1' then return 0 end + if state == 'completed' then return 2 end + if state then return -2 end + if not state then + local now = redis.call('TIME') + if tonumber(now[1]) * 1000 + math.floor(tonumber(now[2]) / 1000) >= tonumber(ARGV[3]) then + return -1 + end + -- One write command, so an OOM cannot publish just half the decision. + redis.call('MSET', KEYS[1], 'completed', KEYS[2], ARGV[1]) + redis.call('EXPIRE', KEYS[1], ARGV[2]) + redis.call('EXPIRE', KEYS[2], ARGV[2]) + end + return 1 + `, + 2, + cancellationKey(target), + `${cancellationKey(target)}:result`, + serialized, + Math.max(1, ttlSeconds), + deadlineAtMs, + ); + if (decision === -1) + throw new Error('Job result commitment exceeded its deadline'); + if (decision === 1) return 'committed'; + if (decision === 0) return 'cancelled'; + if (decision === 2) return 'already_completed'; + throw new Error('Invalid durable result commitment'); +} + +/** BullMQ lock loss can redeliver a job while its first processor still runs. + * Claim once before any sandbox work, retaining the claim through the job's + * recovery horizon. An ambiguous/stalled attempt is never permission to rerun. + * Completion lookup and claim are atomic, so there is no read-then-start gap. */ +export async function claimJobExecution( + commands: IORedis, + target: JobTarget, + ttlSeconds: number, +): Promise<{ status: 'claimed' } | { status: 'completed'; result: T }> { + const key = cancellationKey(target); + const decision = await commands.eval( + ` + local state = redis.call('GET', KEYS[1]) + if state == 'completed' then return {0, redis.call('GET', KEYS[2])} end + if state == '1' then return {-1} end + if state then return {-3} end + if redis.call('SET', KEYS[3], '1', 'NX', 'EX', ARGV[1]) then return {1} end + return {-2} + `, + 3, + key, + `${key}:result`, + `${key}:execution`, + Math.max(1, ttlSeconds), + ); + if (!Array.isArray(decision)) throw new Error('Invalid execution claim'); + if (decision[0] === 1) return { status: 'claimed' }; + if (decision[0] === 0) + return { + status: 'completed', + result: decodeCommittedResult(decision[1]).result, + }; + if (decision[0] === -1) throw new Error(JOB_CANCELLED_MESSAGE); + if (decision[0] === -2) + throw new Error( + 'Programmatic job already claimed; refusing duplicate execution', + ); + throw new Error('Invalid durable execution claim'); +} + +export async function readCommittedJobResult( + commands: IORedis, + target: JobTarget, +): Promise<{ result: T } | undefined> { + const [state, value] = await commands.mget( + cancellationKey(target), + `${cancellationKey(target)}:result`, + ); + if (state !== 'completed') return undefined; + return decodeCommittedResult(value); +} + +/** Do not release replay ownership on an ambiguous Redis failure. Keep one + * outstanding marker write, retry rejected writes with bounded backoff, and + * retain ownership until it succeeds or the job's ORIGINAL deadline expires. + * A delayed queue.add must carry that same timestamp into the worker. */ +export async function fenceJobCancellation(args: { + commands: IORedis; + target: JobTarget; + ttlSeconds: number; + deadlineAtMs: number; +}): Promise> { + let retryMs = 25; + let firstAttempt = true; + while (firstAttempt || Date.now() < args.deadlineAtMs) { + firstAttempt = false; + // If a lost enqueue reply arrives after the execution deadline, still + // give a healthy Redis one bounded opportunity to return completion's + // winning decision. Never translate a known committed effect to failure. + const remainingMs = args.deadlineAtMs - Date.now(); + let timer: ReturnType | undefined; + let decision: unknown; + try { + decision = await Promise.race([ + cancelJobInRedis(args.commands, args.target, args.ttlSeconds, true), + new Promise(resolve => { + timer = setTimeout( + () => resolve(undefined), + remainingMs > 0 ? remainingMs : 1_000, + ); + }), + ]); + } catch { + await new Promise(resolve => + setTimeout( + resolve, + Math.min(retryMs, Math.max(0, args.deadlineAtMs - Date.now())), + ), + ); + retryMs = Math.min(1_000, retryMs * 2); + continue; + } finally { + if (timer != null) clearTimeout(timer); + } + // Only transport failures retry. Corrupt/missing durable results are + // deterministic invariant failures, not an invitation to extend their TTL. + if (decision === undefined) return { status: 'expired' }; + if (!Array.isArray(decision)) + throw new Error('Invalid durable cancellation decision'); + if (decision[0] === 1) return { status: 'cancelled' }; + if (decision[0] === 0) + return { + status: 'completed', + result: decodeCommittedResult(decision[1]).result, + }; + throw new Error('Invalid durable cancellation decision'); + } + return { status: 'expired' }; +} + +const REMOVABLE_JOB_STATES = new Set([ + 'waiting', + 'delayed', + 'prioritized', + 'waiting-children', +]); + +/** Frees queued capacity without ever removing an active or settled job. */ +export async function removeJobIfWaiting( + job: Pick, +): Promise { + if (!REMOVABLE_JOB_STATES.has(await job.getState())) return false; + try { + await job.remove(); + return true; + } catch { + // A worker may have activated the job between getState() and remove(). + // The durable marker remains authoritative for that race. + return false; + } +} + +export function programmaticCancellationError(): Error { + return new DOMException( + 'Programmatic execution request disconnected', + 'AbortError', + ); +} + +/** Commit barrier for result-processing stages that may yield after execution. */ +export function throwIfJobAborted(signal: AbortSignal): void { + if (!signal.aborted) return; + if (signal.reason instanceof Error) throw signal.reason; + throw new DOMException( + typeof signal.reason === 'string' ? signal.reason : 'Job aborted', + 'AbortError', + ); +} + +/** Maps cancellation observed during asynchronous result cleanup to the same + * stable worker failure used by the main execution catch path. */ +export function jobResultCommitFailure( + signal: AbortSignal, + jobTimeoutMs: number, +): Error | undefined { + if (!signal.aborted) return undefined; + return new Error( + signal.reason === CLIENT_DISCONNECT_REASON + ? JOB_CANCELLED_MESSAGE + : `Job timed out after ${jobTimeoutMs}ms`, + ); +} + +export async function waitForJobWithCancellation(args: { + commands: IORedis; + registry: JobCancellationRegistry; + job: Job; + events: QueueEvents; + timeoutMs: number; + cancellationTtlSeconds: number; + deadlineAtMs?: number; + signal?: AbortSignal; +}): Promise { + const { + commands, + registry, + job, + events, + timeoutMs, + cancellationTtlSeconds, + signal, + } = args; + const completion = job.waitUntilFinished(events, timeoutMs); + // Subscription startup can itself wait for Redis recovery. Own the losing + // promise immediately, before any await, rather than after registration. + void completion.catch(() => undefined); + const target = { queueName: job.queueName, jobId: String(job.id) }; + const deadlineAtMs = args.deadlineAtMs ?? Date.now() + timeoutMs; + let fencing: Promise> | undefined; + const fence = (): Promise> => + (fencing ??= fenceJobCancellation({ + commands, + target, + ttlSeconds: cancellationTtlSeconds, + deadlineAtMs, + })); + const externalController = new AbortController(); + try { + await registry.register(target, externalController); + } catch (error) { + void completion.catch(() => undefined); + const outcome = await fence(); + if (outcome.status === 'completed') return outcome.result; + await removeJobIfWaiting(job).catch(() => false); + throw error; + } + + let removeAbortListener = (): void => {}; + const disconnected = new Promise((resolve, reject) => { + let cancelling = false; + const cancel = (): void => { + if (cancelling) return; + cancelling = true; + void fence() + .then(async outcome => { + if (outcome.status === 'completed') { + resolve(outcome.result); + return; + } + // Removing a waiting job immediately frees queue capacity. An active + // job cannot be removed; its worker observes the durable marker or + // pub/sub event and aborts the sandbox transport instead. + await removeJobIfWaiting(job).catch(() => false); + reject(programmaticCancellationError()); + }) + .catch(reject); + }; + if (signal != null) { + removeAbortListener = (): void => + signal.removeEventListener('abort', cancel); + signal.addEventListener('abort', cancel, { once: true }); + if (signal.aborted) cancel(); + } + }); + const cancelled = new Promise((_, reject) => { + const cancel = (): void => { + void removeJobIfWaiting(job).then( + () => reject(programmaticCancellationError()), + () => reject(programmaticCancellationError()), + ); + }; + externalController.signal.addEventListener('abort', cancel, { + once: true, + }); + if (externalController.signal.aborted) cancel(); + }); + + // A cancelled request stops awaiting the BullMQ result, so attach a sink to + // the losing promise before racing it to avoid an unhandled late rejection. + void completion.catch(() => undefined); + try { + return await Promise.race([completion, disconnected, cancelled]); + } catch (error) { + // Includes waitUntilFinished timeouts and registration/transport errors, + // not only explicit Stop. Replay cleanup is unsafe until this barrier. + const outcome = await fence(); + if (outcome.status === 'completed') return outcome.result; + throw error; + } finally { + removeAbortListener(); + await registry + .unregister(target, externalController) + .catch(() => undefined); + } +} + +export const jobCancellationInternals = { + channel: JOB_CANCELLATION_CHANNEL, + cancellationKey, + parseTarget, +}; diff --git a/service/src/metrics.ts b/service/src/metrics.ts index adfd9872..42c5536f 100644 --- a/service/src/metrics.ts +++ b/service/src/metrics.ts @@ -117,6 +117,12 @@ export const jobsFailed = new Counter({ labelNames: ['language'] as const, }); +export const jobsCancelled = new Counter({ + name: 'codeapi_jobs_cancelled_total', + help: 'Total number of jobs cancelled after the calling client disconnected', + labelNames: ['language'] as const, +}); + export const activeJobs = new Gauge({ name: 'codeapi_active_jobs', help: 'Number of jobs currently being processed', diff --git a/service/src/middleware/limits.ts b/service/src/middleware/limits.ts index 099261a1..9ffaeb9d 100644 --- a/service/src/middleware/limits.ts +++ b/service/src/middleware/limits.ts @@ -182,6 +182,19 @@ export const executionLimiter = createRateLimiter( } ); +/** Keep Stop available when execution admission is full, while independently + * bounding request-id churn in the cancellation registry. */ +export const cancellationLimiter = createRateLimiter( + 'exec-cancel', + env.EXEC_LIMIT_WINDOW, + Math.max(80, env.EXEC_MAX_REQUESTS * 4), + { + message: 'Too many CodeAPI cancellation requests.', + structuredBody: true, + logRejections: true, + } +); + export const uploadLimiter = createRateLimiter( 'upload', env.UPLOAD_LIMIT_WINDOW, diff --git a/service/src/programmatic-cancellation.test.ts b/service/src/programmatic-cancellation.test.ts new file mode 100644 index 00000000..53a4340e --- /dev/null +++ b/service/src/programmatic-cancellation.test.ts @@ -0,0 +1,234 @@ +import { afterEach, beforeEach, expect, test } from 'bun:test'; +import type IORedis from 'ioredis'; +import { startTestRedis } from './test/redis'; +import { commitJobResult, requestJobCancellation } from './job-cancellation'; +import { + attachProgrammaticCancellationTarget, + cancelProgrammaticRequest, + normalizeProgrammaticRequestId, + programmaticCancellationInternals, + releaseProgrammaticCancellation, + reserveProgrammaticCancellation, +} from './programmatic-cancellation'; + +let redis: IORedis & { closeTestServer(): Promise }; + +beforeEach(async () => { + redis = await startTestRedis(); +}); + +afterEach(async () => { + await redis.closeTestServer(); +}); + +test('normalizes only bounded opaque request IDs', () => { + expect(normalizeProgrammaticRequestId('request_123456789')).toBe( + 'request_123456789', + ); + expect(normalizeProgrammaticRequestId(' short ')).toBeUndefined(); + expect( + normalizeProgrammaticRequestId('../request_123456789'), + ).toBeUndefined(); + expect(normalizeProgrammaticRequestId('a'.repeat(129))).toBeUndefined(); +}); + +test('Stop cannot extend an attached tombstone before the outcome command succeeds', async () => { + const requestId = 'request_no_split_renewal'; + const owner = 'owner-a'; + const target = { queueName: 'other', jobId: 'split-renewal' }; + await reserveProgrammaticCancellation({ + redis, + requestId, + owner, + ttlSeconds: 60, + }); + await attachProgrammaticCancellationTarget({ + redis, + requestId, + owner, + target, + ttlSeconds: 60, + }); + await commitJobResult(redis, target, { stdout: 'done' }, 60); + const key = programmaticCancellationInternals.requestKey(requestId); + await redis.pexpire(key, 5_000); + const before = await redis.pttl(key); + expect( + await cancelProgrammaticRequest({ + redis, + requestId, + owner, + ttlSeconds: 600, + }), + ).toEqual({ status: 'accepted', target }); + // Simulate losing Redis before requestJobCancellation: no second command. + expect(await redis.pttl(key)).toBeLessThanOrEqual(before); + expect(await requestJobCancellation(redis, target, 600)).toBe(false); +}); + +test('cancellation before queue attachment is retained atomically', async () => { + const requestId = 'request_early_cancel_123'; + const owner = 'owner-a'; + expect( + await cancelProgrammaticRequest({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toEqual({ status: 'accepted' }); + + expect( + await reserveProgrammaticCancellation({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toBe('cancelled'); + expect( + await attachProgrammaticCancellationTarget({ + redis, + requestId, + owner, + target: { queueName: 'other', jobId: '42' }, + ttlSeconds: 60, + }), + ).toBe('cancelled'); +}); + +test('cancellation after attachment returns the exact queue target', async () => { + const requestId = 'request_attached_cancel_1'; + const owner = 'owner-a'; + expect( + await reserveProgrammaticCancellation({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toBe('active'); + expect( + await attachProgrammaticCancellationTarget({ + redis, + requestId, + owner, + target: { queueName: 'other', jobId: '43' }, + ttlSeconds: 60, + }), + ).toBe('active'); + + expect( + await cancelProgrammaticRequest({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toEqual({ + status: 'accepted', + target: { queueName: 'other', jobId: '43' }, + }); +}); + +test('overlapping requests from the same owner cannot share cancellation state', async () => { + const requestId = 'request_duplicate_owner_1'; + const owner = 'owner-a'; + expect( + await reserveProgrammaticCancellation({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toBe('active'); + + expect( + await reserveProgrammaticCancellation({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toBe('duplicate'); +}); + +test('a different principal cannot reserve, attach, cancel, or release a request', async () => { + const requestId = 'request_owned_cancel_123'; + await reserveProgrammaticCancellation({ + redis, + requestId, + owner: 'owner-a', + ttlSeconds: 60, + }); + + expect( + await reserveProgrammaticCancellation({ + redis, + requestId, + owner: 'owner-b', + ttlSeconds: 60, + }), + ).toBe('forbidden'); + expect( + await attachProgrammaticCancellationTarget({ + redis, + requestId, + owner: 'owner-b', + target: { queueName: 'other', jobId: '44' }, + ttlSeconds: 60, + }), + ).toBe('forbidden'); + expect( + await cancelProgrammaticRequest({ + redis, + requestId, + owner: 'owner-b', + ttlSeconds: 60, + }), + ).toEqual({ status: 'forbidden' }); + await releaseProgrammaticCancellation({ + redis, + requestId, + owner: 'owner-b', + }); + expect( + await redis.exists(programmaticCancellationInternals.requestKey(requestId)), + ).toBe(1); +}); + +test('settlement retains a bounded target tombstone for late Stop classification', async () => { + const requestId = 'request_release_cancel_1'; + await reserveProgrammaticCancellation({ + redis, + requestId, + owner: 'owner-a', + ttlSeconds: 60, + }); + const target = { queueName: 'other', jobId: 'settled-job' }; + await attachProgrammaticCancellationTarget({ + redis, + requestId, + owner: 'owner-a', + target, + ttlSeconds: 60, + }); + await commitJobResult(redis, target, { stdout: 'done' }, 60); + await releaseProgrammaticCancellation({ + redis, + requestId, + owner: 'owner-a', + }); + const key = programmaticCancellationInternals.requestKey(requestId); + expect(await redis.exists(key)).toBe(1); + expect(await redis.ttl(key)).toBeGreaterThan(0); + expect(await redis.ttl(key)).toBeLessThanOrEqual(60); + const cancelled = await cancelProgrammaticRequest({ + redis, + requestId, + owner: 'owner-a', + ttlSeconds: 60, + }); + expect(cancelled).toEqual({ status: 'accepted', target }); + expect(await requestJobCancellation(redis, target, 60)).toBe(false); +}); diff --git a/service/src/programmatic-cancellation.ts b/service/src/programmatic-cancellation.ts new file mode 100644 index 00000000..12a6694d --- /dev/null +++ b/service/src/programmatic-cancellation.ts @@ -0,0 +1,182 @@ +import { createHash } from 'node:crypto'; +import type IORedis from 'ioredis'; +import type { AuthenticatedRequest } from './types'; +import { getCredentialId } from './auth/principal'; +import { getExecutionIdentity } from './execution-identity'; + +export const CODEAPI_PROGRAMMATIC_REQUEST_HEADER = + 'X-LibreChat-Code-Request-ID'; +const REQUEST_PREFIX = 'codeapi:programmatic-cancellation:v1'; +const REQUEST_ID_PATTERN = /^[A-Za-z0-9_-]{16,128}$/; + +interface CancellationTarget { + queueName: string; + jobId: string; +} + +export type CancellationRequestResult = + | { status: 'accepted'; target?: CancellationTarget } + | { status: 'forbidden' }; + +function requestKey(requestId: string): string { + return `${REQUEST_PREFIX}:${requestId}`; +} + +export function normalizeProgrammaticRequestId( + value: unknown, +): string | undefined { + if (typeof value !== 'string') return undefined; + const trimmed = value.trim(); + return REQUEST_ID_PATTERN.test(trimmed) ? trimmed : undefined; +} + +export function programmaticCancellationOwner( + req: AuthenticatedRequest, + userId: string, +): string { + const identity = getExecutionIdentity(req, userId); + return createHash('sha256') + .update( + JSON.stringify([ + identity.storageNamespace, + identity.canonicalUserId, + getCredentialId(req), + identity.authContextHash ?? '', + ]), + ) + .digest('hex'); +} + +const RESERVE_SCRIPT = ` +local key = KEYS[1] +local owner = ARGV[1] +local ttl = tonumber(ARGV[2]) +local existing = redis.call('HGET', key, 'owner') +if existing and existing ~= owner then return -1 end +if redis.call('HGET', key, 'reserved') == '1' then return -2 end +if not existing then + redis.call('HSET', key, 'owner', owner, 'cancelled', '0') +end +redis.call('HSET', key, 'reserved', '1') +redis.call('EXPIRE', key, ttl) +return tonumber(redis.call('HGET', key, 'cancelled') or '0') +`; + +const ATTACH_SCRIPT = ` +local key = KEYS[1] +local owner = ARGV[1] +local queueName = ARGV[2] +local jobId = ARGV[3] +local ttl = tonumber(ARGV[4]) +if redis.call('HGET', key, 'owner') ~= owner then return -1 end +redis.call('HSET', key, 'queueName', queueName, 'jobId', jobId) +redis.call('EXPIRE', key, ttl) +return tonumber(redis.call('HGET', key, 'cancelled') or '0') +`; + +const CANCEL_SCRIPT = ` +local key = KEYS[1] +local owner = ARGV[1] +local ttl = tonumber(ARGV[2]) +local existing = redis.call('HGET', key, 'owner') +if existing and existing ~= owner then return {-1} end +if not existing then redis.call('HSET', key, 'owner', owner) end +redis.call('HSET', key, 'cancelled', '1') +local queueName = redis.call('HGET', key, 'queueName') +local jobId = redis.call('HGET', key, 'jobId') +-- Once attached, never extend this mapping independently of the job decision. +-- Its original admission TTL already covers execution and late cancellation. +if queueName and jobId then return {1, queueName, jobId} end +redis.call('EXPIRE', key, ttl) +return {1} +`; + +const RELEASE_SCRIPT = ` +if redis.call('HGET', KEYS[1], 'owner') == ARGV[1] then + -- Keep the owner/target tombstone through its existing bounded TTL. A Stop + -- racing response delivery must still reach the job's completion decision. + return redis.call('HSET', KEYS[1], 'finished', '1') +end +return 0 +`; + +export async function reserveProgrammaticCancellation(args: { + redis: IORedis; + requestId: string; + owner: string; + ttlSeconds: number; +}): Promise<'active' | 'cancelled' | 'duplicate' | 'forbidden'> { + const result = Number( + await args.redis.eval( + RESERVE_SCRIPT, + 1, + requestKey(args.requestId), + args.owner, + Math.max(1, args.ttlSeconds), + ), + ); + if (result === -1) return 'forbidden'; + if (result === -2) return 'duplicate'; + return result === 1 ? 'cancelled' : 'active'; +} + +export async function attachProgrammaticCancellationTarget(args: { + redis: IORedis; + requestId: string; + owner: string; + target: CancellationTarget; + ttlSeconds: number; +}): Promise<'active' | 'cancelled' | 'forbidden'> { + const result = Number( + await args.redis.eval( + ATTACH_SCRIPT, + 1, + requestKey(args.requestId), + args.owner, + args.target.queueName, + args.target.jobId, + Math.max(1, args.ttlSeconds), + ), + ); + if (result < 0) return 'forbidden'; + return result === 1 ? 'cancelled' : 'active'; +} + +export async function cancelProgrammaticRequest(args: { + redis: IORedis; + requestId: string; + owner: string; + ttlSeconds: number; +}): Promise { + const raw = await args.redis.eval( + CANCEL_SCRIPT, + 1, + requestKey(args.requestId), + args.owner, + Math.max(1, args.ttlSeconds), + ); + const result = Array.isArray(raw) ? raw.map(String) : []; + if (result[0] === '-1') return { status: 'forbidden' }; + if (result.length >= 3) { + return { + status: 'accepted', + target: { queueName: result[1]!, jobId: result[2]! }, + }; + } + return { status: 'accepted' }; +} + +export async function releaseProgrammaticCancellation(args: { + redis: IORedis; + requestId: string; + owner: string; +}): Promise { + await args.redis.eval( + RELEASE_SCRIPT, + 1, + requestKey(args.requestId), + args.owner, + ); +} + +export const programmaticCancellationInternals = { requestKey }; diff --git a/service/src/queue.ts b/service/src/queue.ts index 54fea308..72c2fbb1 100644 --- a/service/src/queue.ts +++ b/service/src/queue.ts @@ -1,6 +1,7 @@ // src/queue.ts import IORedis from 'ioredis'; import { Queue, QueueEvents } from 'bullmq'; +import type { Job } from 'bullmq'; import { setMaxListeners } from 'events'; import type { CommonRedisOptions } from 'ioredis'; import type * as tls from 'tls'; @@ -17,19 +18,13 @@ import type { SandboxBackendName, } from './execution-profile'; import logger from './logger'; -import { redisKeepAliveOptions } from './redis-options'; +import { redisKeepAliveOptions, redisReconnectDelay } from './redis-options'; import { bullmqQueueJobs, registerBullmqQueueMetricsCollector } from './metrics'; - -const MAX_RECONNECT_ATTEMPTS = 5; -const RECONNECT_DELAY = 2000; +import { JobCancellationRegistry } from './job-cancellation'; const retryStrategy: CommonRedisOptions['retryStrategy'] = (times) => { - if (times > MAX_RECONNECT_ATTEMPTS) { - logger.error(`Failed to connect to Redis after ${times} attempts`); - return null; - } logger.warn(`Retrying Redis connection attempt ${times}`); - return RECONNECT_DELAY; + return redisReconnectDelay(times); }; const reconnectOnError: CommonRedisOptions['reconnectOnError'] = (err) => { @@ -60,6 +55,7 @@ const connection = new IORedis({ ? { dnsLookup: (address: string, callback: (err: Error | null, addr: string) => void): void => callback(null, address) } : {}) }); +const jobCancellationRegistry = new JobCancellationRegistry(connection); // Global queues - no INSTANCE_ID prefix // This enables horizontal scaling where any worker can process any job @@ -110,6 +106,19 @@ export function getExecutionQueueBinding( return { ...getQueueResources(name), language }; } +/** + * Resolve a job only from this deployment's already-open queue set. Every + * homogeneous API replica opens both execution queues at startup, so this + * supports cross-replica cancellation without allocating attacker-shaped + * QueueEvents connections for arbitrary names recovered from Redis. + */ +export async function getExistingExecutionJob( + queueName: string, + jobId: string, +): Promise | undefined> { + return queueResources.get(queueName)?.queue.getJob(jobId); +} + const { queue: pyQueue, events: pyQueueEvents } = getQueueResources(queueNames.python); const { queue: otherQueue, events: otherQueueEvents } = getQueueResources(queueNames.other); @@ -163,8 +172,16 @@ export async function closeQueueConnections(): Promise { [...queueResources.values()].flatMap(({ queue, events }) => [ queue.close(), events.close(), - ]), + ]).concat(jobCancellationRegistry.close()), ); } -export { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, queueNames, connection }; +export { + pyQueue, + otherQueue, + pyQueueEvents, + otherQueueEvents, + queueNames, + connection, + jobCancellationRegistry, +}; diff --git a/service/src/redis-options.test.ts b/service/src/redis-options.test.ts index 1bb942e6..9795e279 100644 --- a/service/src/redis-options.test.ts +++ b/service/src/redis-options.test.ts @@ -1,5 +1,15 @@ import { afterEach, describe, expect, test } from 'bun:test'; -import { redisKeepAliveMs, redisKeepAliveOptions } from './redis-options'; +import { + redisKeepAliveMs, + redisKeepAliveOptions, + redisReconnectDelay, +} from './redis-options'; + +test('long-lived command connections keep recovering with a bounded retry delay', () => { + expect(redisReconnectDelay(1)).toBe(100); + expect(redisReconnectDelay(6)).toBe(600); + expect(redisReconnectDelay(1_000)).toBe(2_000); +}); describe('Redis keepalive options', () => { afterEach(() => { diff --git a/service/src/redis-options.ts b/service/src/redis-options.ts index 98d3f5b1..1f2c8e42 100644 --- a/service/src/redis-options.ts +++ b/service/src/redis-options.ts @@ -1,5 +1,11 @@ import type { CommonRedisOptions } from 'ioredis'; +/** Long-lived queue/cancellation command and subscriber connections must both + * recover after an outage. Never leave a live process with a terminal client. */ +export function redisReconnectDelay(attempt: number): number { + return Math.min(2_000, 100 * Math.max(1, attempt)); +} + export function redisKeepAliveMs(): number { const raw = process.env.REDIS_KEEP_ALIVE_MS; const trimmed = raw?.trim(); diff --git a/service/src/request-disconnect.test.ts b/service/src/request-disconnect.test.ts new file mode 100644 index 00000000..15660aca --- /dev/null +++ b/service/src/request-disconnect.test.ts @@ -0,0 +1,60 @@ +import { expect, test } from 'bun:test'; +import { EventEmitter } from 'node:events'; +import type { Response } from 'express'; +import type { AuthenticatedRequest } from './types'; +import { CLIENT_DISCONNECT_REASON } from './job-cancellation'; +import { observeRequestDisconnect } from './request-disconnect'; + +function requestAndResponse(options: { + requestAborted?: boolean; + requestDestroyed?: boolean; + responseDestroyed?: boolean; +} = {}): { + req: AuthenticatedRequest & EventEmitter; + res: Response & EventEmitter; +} { + const req = Object.assign(new EventEmitter(), { + aborted: options.requestAborted ?? false, + destroyed: options.requestDestroyed ?? false, + }) as AuthenticatedRequest & EventEmitter; + const res = Object.assign(new EventEmitter(), { + destroyed: options.responseDestroyed ?? false, + writableFinished: false, + }) as Response & EventEmitter; + return { req, res }; +} + +test('a consumed Bun request stream is not mistaken for a disconnect', () => { + const { req, res } = requestAndResponse({ requestDestroyed: true }); + const observer = observeRequestDisconnect(req, res); + + expect(observer.isDisconnected()).toBe(false); + expect(observer.signal.aborted).toBe(false); + observer.dispose(); +}); + +test('current and future transport abandonment abort exactly once', () => { + const current = requestAndResponse({ requestAborted: true }); + const currentObserver = observeRequestDisconnect(current.req, current.res); + expect(currentObserver.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + + const future = requestAndResponse(); + const futureObserver = observeRequestDisconnect(future.req, future.res); + future.res.emit('close'); + future.req.emit('aborted'); + expect(futureObserver.signal.reason).toBe(CLIENT_DISCONNECT_REASON); + expect(future.req.listenerCount('aborted')).toBe(0); + expect(future.res.listenerCount('close')).toBe(0); +}); + +test('a completed response disposes listeners without aborting', () => { + const { req, res } = requestAndResponse(); + const observer = observeRequestDisconnect(req, res); + (res as unknown as { writableFinished: boolean }).writableFinished = true; + res.emit('finish'); + res.emit('close'); + + expect(observer.isDisconnected()).toBe(false); + expect(req.listenerCount('aborted')).toBe(0); + expect(res.listenerCount('close')).toBe(0); +}); diff --git a/service/src/request-disconnect.ts b/service/src/request-disconnect.ts new file mode 100644 index 00000000..2a4d7a9d --- /dev/null +++ b/service/src/request-disconnect.ts @@ -0,0 +1,49 @@ +import type { Response } from 'express'; +import type { AuthenticatedRequest } from './types'; +import { CLIENT_DISCONNECT_REASON } from './job-cancellation'; + +export interface RequestDisconnectObserver { + signal: AbortSignal; + isDisconnected(): boolean; + dispose(): void; +} + +/** + * Observe a genuinely abandoned HTTP response across Node and Bun. + * + * Bun may mark the consumed IncomingMessage stream as `destroyed` while the + * response remains healthy, so request stream destruction is deliberately not + * treated as a disconnect. Express' `aborted` event and ServerResponse's + * pre-finish `close` event are the portable abandonment signals. + */ +export function observeRequestDisconnect( + req: AuthenticatedRequest, + res: Response, +): RequestDisconnectObserver { + const controller = new AbortController(); + let disposed = false; + const dispose = (): void => { + if (disposed) return; + disposed = true; + req.removeListener('aborted', disconnect); + res.removeListener('close', disconnect); + res.removeListener('finish', dispose); + }; + const disconnect = (): void => { + if (!res.writableFinished && !controller.signal.aborted) { + controller.abort(CLIENT_DISCONNECT_REASON); + } + dispose(); + }; + + req.once('aborted', disconnect); + res.once('close', disconnect); + res.once('finish', dispose); + if (req.aborted || res.destroyed) disconnect(); + + return { + signal: controller.signal, + isDisconnected: () => controller.signal.aborted, + dispose, + }; +} diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index 93bbc0c5..bb140013 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -4,13 +4,32 @@ import { Router } from 'express'; import type { Response } from 'express'; import type * as t from '../types'; import { checkServiceStartUp, checkServiceShutDown } from '../lifecycle'; -import { executionLimiter } from '../middleware/limits'; +import { cancellationLimiter, executionLimiter } from '../middleware/limits'; import { pyQueue, pyQueueEvents, connection, + jobCancellationRegistry, getExecutionQueueBinding, + getExistingExecutionJob, } from '../queue'; +import { + JOB_CANCELLED_MESSAGE, + programmaticCancellationError, + removeJobIfWaiting, + requestJobCancellation, + fenceJobCancellation, + waitForJobWithCancellation, +} from '../job-cancellation'; +import { + CODEAPI_PROGRAMMATIC_REQUEST_HEADER, + attachProgrammaticCancellationTarget, + cancelProgrammaticRequest, + normalizeProgrammaticRequestId, + programmaticCancellationOwner, + releaseProgrammaticCancellation, + reserveProgrammaticCancellation, +} from '../programmatic-cancellation'; import { createProgrammaticPayload, extractPendingFromControlPayload, @@ -38,6 +57,7 @@ import { Jobs } from '../enum'; import { env, jobCompletionWaitTimeoutMs } from '../config'; import { resolveQueuedSandboxBackend } from '../execution-profile'; import { publicExecutionFailure } from '../utils'; +import { observeRequestDisconnect } from '../request-disconnect'; import { normalizeEgressGatewayUrl, normalizeProgrammaticTimeoutMs, @@ -107,6 +127,18 @@ const JOB_COMPLETION_WAIT_TIMEOUT_MS = jobCompletionWaitTimeoutMs( env.LAMBDA_MICROVM_LAUNCH_TIMEOUT_MS, env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS, ); +const PROGRAMMATIC_CANCELLATION_TTL_SECONDS = + Math.ceil(JOB_COMPLETION_WAIT_TIMEOUT_MS / 1000) + 60; + +interface ReplayRequestCancellation { + signal: AbortSignal; + isDisconnected(): boolean; + request?: { + requestId: string; + owner: string; + cancelledBeforeStart: boolean; + }; +} const router = Router(); @@ -322,7 +354,10 @@ async function runReplayIteration( state: ExecutionState, apiKeyId: string, userId: string, + signal?: AbortSignal, + cancellation?: { requestId: string; owner: string }, ): Promise { + if (signal?.aborted) throw programmaticCancellationError(); const history = await loadToolHistory(state.execution_id); const rawPayload = buildReplayPayload(req, state, history); const sessionKey = state.sessionKey ?? state.userId; @@ -369,41 +404,83 @@ async function runReplayIteration( state.executionProfile ?? env.EXECUTION_PROFILE, state.executionProfileSource ?? env.EXECUTION_PROFILE_SOURCE, ); - const job = await queue.add( - Jobs.execute, - { - code: state.userCode ?? '', - userId, - payload: sandboxSecurity.payload, - apiKeyId, - isPyPlot: state.isPyPlot ?? false, - principalSource: state.principalSource, - executionId: state.execution_id, - tenantId: state.tenantId, - canonicalUserId: state.canonicalUserId, - executionProfile: state.executionProfile ?? env.EXECUTION_PROFILE, - sandboxBackend: replayBackend, - ...(state.bridgeWorkerId != null - ? { bridgeWorkerId: state.bridgeWorkerId } - : {}), - ...(state.workspaceId != null - ? { workspaceId: state.workspaceId } - : {}), - runtimeSessionMode: 'stateless', - runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, - executionManifestClaims: sandboxSecurity.executionManifestClaims, - egressGrantClaims: sandboxSecurity.egressGrantClaims, - egressGrantToken: sandboxSecurity.egressGrantToken, - }, - { - removeOnComplete: { age: 60, count: 1 }, - removeOnFail: { age: 180, count: 1 }, - attempts: 1, - }, + if (signal?.aborted) throw programmaticCancellationError(); + const cancellationTarget = { queueName: queue.name, jobId: nanoid() }; + if (cancellation != null) { + const attachment = await attachProgrammaticCancellationTarget({ + redis: connection, + requestId: cancellation.requestId, + owner: cancellation.owner, + target: cancellationTarget, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + }); + if (attachment === 'forbidden') { + throw new Error('Programmatic cancellation request ownership changed'); + } + if (attachment === 'cancelled') { + throw programmaticCancellationError(); + } + } + const submittedAtMs = Date.now(); + const deadlineAtMs = submittedAtMs + env.JOB_TIMEOUT; + let job: Awaited>; + try { + job = await queue.add( + Jobs.execute, + { + code: state.userCode ?? '', + userId, + payload: sandboxSecurity.payload, + apiKeyId, + isPyPlot: state.isPyPlot ?? false, + principalSource: state.principalSource, + executionId: state.execution_id, + tenantId: state.tenantId, + canonicalUserId: state.canonicalUserId, + executionProfile: state.executionProfile ?? env.EXECUTION_PROFILE, + sandboxBackend: replayBackend, + ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), + ...(state.workspaceId != null ? { workspaceId: state.workspaceId } : {}), + cancellable: true, + deadlineAtMs, + cancellationTtlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + runtimeSessionMode: 'stateless', + runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, + executionManifestClaims: sandboxSecurity.executionManifestClaims, + egressGrantClaims: sandboxSecurity.egressGrantClaims, + egressGrantToken: sandboxSecurity.egressGrantToken, + }, + { + removeOnComplete: { age: 60, count: 1 }, + removeOnFail: { age: 180, count: 1 }, + attempts: 1, + jobId: cancellationTarget.jobId, + timestamp: submittedAtMs, + }, ); + } catch (error) { + // Redis may have enqueued the job even though its reply was lost. + // Preserve replay ownership until cancellation is durable or the job's + // fixed worker deadline prevents a late admission from executing. + const outcome = await fenceJobCancellation({ + commands: connection, target: cancellationTarget, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, deadlineAtMs, + }); + if (outcome.status === 'completed') return outcome.result; + throw error; + } jobsSubmitted.inc({ language }); - return job.waitUntilFinished(events, JOB_COMPLETION_WAIT_TIMEOUT_MS); + return waitForJobWithCancellation({ + commands: connection, + registry: jobCancellationRegistry, + job, + events, + timeoutMs: JOB_COMPLETION_WAIT_TIMEOUT_MS, + cancellationTtlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + deadlineAtMs, + signal, + }); } function isSandboxRunSuccess(result: t.ExecuteResult): boolean { @@ -425,6 +502,7 @@ async function handleReplayInitial( bridgeWorkerId?: string; workspaceId?: string; }, + cancellation: ReplayRequestCancellation, ): Promise { const { apiKeyId, userId, bridgeWorkerId, workspaceId } = params; const { code, tools, user_id, files } = @@ -543,6 +621,19 @@ async function handleReplayInitial( throw error; } + if ( + cancellation.signal.aborted || + cancellation.request?.cancelledBeforeStart === true + ) { + if (!cancellation.isDisconnected()) { + res.status(200).json({ + status: 'error', + error: 'Programmatic execution request cancelled', + }); + } + return; + } + const session_id = nanoid(); const execution_id = nanoid(); const authContext = req.codeApiAuthContext; @@ -622,7 +713,7 @@ async function handleReplayInitial( timeout, }); - await runAndRespond(req, res, state, apiKeyId, userId); + await runAndRespond(req, res, state, apiKeyId, userId, cancellation); } async function handleReplayContinuation( @@ -634,6 +725,7 @@ async function handleReplayContinuation( decoded: { execution_id: string }; tool_results: NonNullable; }, + cancellation: ReplayRequestCancellation, ): Promise { const { apiKeyId, userId, decoded, tool_results } = params; @@ -692,6 +784,20 @@ async function handleReplayContinuation( res.status(404).json({ error: 'Execution not found or expired' }); return; } + if ( + cancellation.signal.aborted || + cancellation.request?.cancelledBeforeStart === true + ) { + await cleanupExecution(state.execution_id, 'replay'); + if (!cancellation.isDisconnected()) { + res.status(200).json({ + status: 'error', + error: 'Programmatic execution request cancelled', + session_id: state.session_id, + }); + } + return; + } /** Compute the delta against already-persisted history first so the * cap checks see the real impact of this batch (new call_ids only * advance `callCount`; overwrites may shrink or grow `historyBytes` @@ -845,7 +951,14 @@ async function handleReplayContinuation( }); } - await runAndRespond(req, res, state, apiKeyId, userId); + await runAndRespond( + req, + res, + state, + apiKeyId, + userId, + cancellation, + ); } finally { await releaseExecutionLock(decoded.execution_id, lockToken); } @@ -857,29 +970,29 @@ async function runAndRespond( state: ExecutionState, apiKeyId: string, userId: string, + cancellation: ReplayRequestCancellation, ): Promise { - /** Read disconnect state through `isDisconnected()` rather than a - * direct boolean. The `req.on('close', ...)` handler flips the flag - * during awaits, but `@typescript-eslint/no-unnecessary-condition` - * (correctly per TS semantics) narrows a directly-mutated `let`/object - * member to its literal value after an early-return `if (...) return`, - * even across awaits. A function call is opaque to that narrowing. */ - let disconnected = false; - const isDisconnected = (): boolean => disconnected; - req.on('close', () => { - if (!res.writableEnded) disconnected = true; - }); - let result: t.ExecuteResult; try { - result = await runReplayIteration(req, state, apiKeyId, userId); + result = await runReplayIteration( + req, + state, + apiKeyId, + userId, + cancellation.signal, + cancellation.request, + ); } catch (err) { - logger.error('Replay iteration failed', { - execution_id: state.execution_id, - err, - }); + const cancelled = + (err as Error).name === 'AbortError' || + (err as Error).message === JOB_CANCELLED_MESSAGE; + logger.log(cancelled ? 'info' : 'error', 'Replay iteration failed', { + execution_id: state.execution_id, + cancelled, + err, + }); await cleanupExecution(state.execution_id, 'replay'); - if (!isDisconnected()) { + if (!cancellation.isDisconnected()) { const publicFailure = publicExecutionFailure(err); const message = publicFailure?.body.message ?? (err as Error).message; @@ -892,7 +1005,7 @@ async function runAndRespond( return; } - if (isDisconnected()) { + if (cancellation.isDisconnected()) { logger.info('Client disconnected during replay; cleaning up', { execution_id: state.execution_id, }); @@ -1002,7 +1115,7 @@ async function runAndRespond( await cleanupExecution(state.execution_id, 'replay').catch( () => {}, ); - if (!isDisconnected()) { + if (!cancellation.isDisconnected()) { if (err instanceof ExecutionStateTooLargeError) { /** A continuation that pushes `emittedCallIds` past the * `MAX_EXECUTION_STATE_BYTES` cap is a client-input sizing @@ -1076,6 +1189,67 @@ async function runAndRespond( // Request entrypoint // --------------------------------------------------------------------------- +router.post( + '/exec/programmatic/cancel', + cancellationLimiter, + async (req: t.AuthenticatedRequest, res) => { + const principal = getPrincipalOrReject(req, res); + if (!principal) return; + const requestId = normalizeProgrammaticRequestId( + (req.body as Record)?.request_id, + ); + if (requestId == null) { + res.status(400).json({ error: 'Invalid or missing request_id' }); + return; + } + try { + const owner = programmaticCancellationOwner(req, principal.userId); + const cancellation = await cancelProgrammaticRequest({ + redis: connection, + requestId, + owner, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + }); + if (cancellation.status === 'forbidden') { + res.status(403).json({ error: 'Programmatic request belongs to another principal' }); + return; + } + if (cancellation.target != null) { + const accepted = await requestJobCancellation( + connection, + cancellation.target, + PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + ); + if (!accepted) { + res.status(200).json({ status: 'already_completed' }); + return; + } + try { + const queuedJob = await getExistingExecutionJob( + cancellation.target.queueName, + cancellation.target.jobId, + ); + if (queuedJob != null) await removeJobIfWaiting(queuedJob); + } catch (error) { + logger.warn('Failed to remove cancelled waiting execution', { + requestId, + queueName: cancellation.target?.queueName, + jobId: cancellation.target?.jobId, + error: (error as Error).message, + }); + } + } + res.status(202).json({ status: 'cancellation_requested' }); + } catch (error) { + logger.error('Failed to request programmatic execution cancellation', { + requestId, + error: (error as Error).message, + }); + res.status(503).json({ error: 'Cancellation service unavailable' }); + } + }, +); + router.post( '/exec/programmatic', executionLimiter, @@ -1095,6 +1269,11 @@ router.post( const { continuation_token, tool_results } = req.body as t.ProgrammaticRequestBody; const rawBody = req.body as Record; + const rawRequestId = req.header(CODEAPI_PROGRAMMATIC_REQUEST_HEADER); + const requestId = normalizeProgrammaticRequestId(rawRequestId); + if (rawRequestId != null && requestId == null) { + return res.status(400).json({ error: 'Invalid programmatic request ID' }); + } const requestedLanguage: unknown = rawBody.language ?? rawBody.lang; let bridgeWorkerId: string | undefined; let workspaceId: string | undefined; @@ -1151,7 +1330,51 @@ router.post( }); } + const disconnectObserver = observeRequestDisconnect(req, res); + + const cancellation: ReplayRequestCancellation = { + signal: disconnectObserver.signal, + isDisconnected: disconnectObserver.isDisconnected, + }; + let reservedCancellation: { requestId: string; owner: string } | undefined; + try { + if (requestId != null) { + const owner = programmaticCancellationOwner(req, userId); + let reservation: Awaited>; + try { + reservation = await reserveProgrammaticCancellation({ + redis: connection, + requestId, + owner, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + }); + } catch (error) { + logger.error('Failed to reserve programmatic cancellation request', { + requestId, + error: (error as Error).message, + }); + if (!cancellation.isDisconnected()) { + return res.status(503).json({ error: 'Cancellation service unavailable' }); + } + return; + } + if (reservation === 'forbidden' || reservation === 'duplicate') { + if (!cancellation.isDisconnected()) { + return res.status(409).json({ + error: 'Programmatic request ID is already in use', + }); + } + return; + } + reservedCancellation = { requestId, owner }; + cancellation.request = { + requestId, + owner, + cancelledBeforeStart: reservation === 'cancelled', + }; + } + /** For continuations, peek at the stored execution to route by the * mode it was started in rather than the current process default. * Without this, a replay-mode execution resumed via an instance @@ -1184,7 +1407,7 @@ router.post( userId, decoded, tool_results, - }); + }, cancellation); } return await handleBlocking(req, res, { apiKeyId, userId }); } @@ -1203,7 +1426,7 @@ router.post( userId, bridgeWorkerId, workspaceId, - }); + }, cancellation); } if (workspaceId != null) { return res.status(400).json({ @@ -1221,6 +1444,20 @@ router.post( return res.status(500).json({ error: 'Internal server error' }); } return; + } finally { + disconnectObserver.dispose(); + if (reservedCancellation != null) { + await releaseProgrammaticCancellation({ + redis: connection, + requestId: reservedCancellation.requestId, + owner: reservedCancellation.owner, + }).catch(error => { + logger.warn('Failed to release programmatic cancellation request', { + requestId: reservedCancellation?.requestId, + error: (error as Error).message, + }); + }); + } } }, ); diff --git a/service/src/types/service.ts b/service/src/types/service.ts index d17b5a99..a0c78486 100644 --- a/service/src/types/service.ts +++ b/service/src/types/service.ts @@ -297,6 +297,12 @@ export type JobData = { bridgeWorkerId?: string; /** Trusted selected workspace for native replay-mode PTC. */ workspaceId?: string; + /** Opts replay jobs into durable client-disconnect cancellation. */ + cancellable?: boolean; + /** Absolute producer budget; queue-worker configuration may only tighten it. */ + deadlineAtMs?: number; + /** Producer request tombstones must never outlive the completion decision. */ + cancellationTtlSeconds?: number; /** Producer deployment identity. Optional only for pre-profile queued jobs. */ executionProfile?: ExecutionProfile; /** Required sandbox transport. Optional only for jobs queued before fencing. */ diff --git a/service/src/workers.ts b/service/src/workers.ts index f11a5420..a7153b53 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -1,69 +1,148 @@ import axios from 'axios'; import { Worker } from 'bullmq'; import type * as t from './types'; -import { filterSystemLogs, applySystemReplacements, getAxiosErrorDetails, sandboxErrorMessageFromAxios } from './utils'; -import { jobProcessingDuration, jobsCompleted, jobsFailed, activeJobs, workerRunning } from './metrics'; -import { connection, queueNames } from './queue'; +import { + filterSystemLogs, + applySystemReplacements, + getAxiosErrorDetails, + sandboxErrorMessageFromAxios, +} from './utils'; +import { + jobProcessingDuration, + jobsCancelled, + jobsCompleted, + jobsFailed, + activeJobs, + workerRunning, +} from './metrics'; +import { connection, jobCancellationRegistry, queueNames } from './queue'; import { env, jobDeadlineAtMs } from './config'; import { summarizeSandboxResponse, summarizeText } from './execution-log'; -import { createGatewayEgressGrant, restoreGatewaySandboxResult, revokeGatewayEgressGrant } from './egress-gateway-client'; +import { + createGatewayEgressGrant, + restoreGatewaySandboxResult, + revokeGatewayEgressGrant, +} from './egress-gateway-client'; import { refreshEgressGrantClaims } from './sandbox-egress'; import { buildSandboxExecuteRequest } from './sandbox-dispatch'; import { prepareInputDelivery } from './runtime-session/input-delivery'; import { SessionFilesError } from './runtime-session/files'; import { resolveRuntimeSessionForJob } from './runtime-session/job-policy'; -import { getSandboxBackend, SandboxBackendError, type SandboxRawResponse } from './sandbox-backend'; +import { + getSandboxBackend, + SandboxBackendError, + type SandboxRawResponse, +} from './sandbox-backend'; import { isSyntheticPrincipalSource } from './auth/synthetic'; import { withSpan, withTraceContext } from './telemetry'; import { workerDeadlineFailure } from './worker-error'; +import { + CLIENT_DISCONNECT_REASON, + JOB_CANCELLED_MESSAGE, + jobResultCommitFailure, + commitJobResult, + claimJobExecution, + jobCancellationRetentionSeconds, + throwIfJobAborted, +} from './job-cancellation'; import logger from './logger'; import { validateQueuedExecutionProfile, validateQueuedSandboxBackend, } from './execution-profile'; -import { BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, programmaticTransferReserveMs } from '../../packages/code/src/protocol'; +import { + BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, + programmaticTransferReserveMs, +} from '../../packages/code/src/protocol'; const { INSTANCE_ID } = env; const WORKER_ID = `${INSTANCE_ID}-${process.pid}`; function isAbortError(error: unknown): boolean { - return axios.isAxiosError(error) && (error.name === 'AbortError' || error.code === 'ERR_CANCELED'); + return ( + axios.isAxiosError(error) && + (error.name === 'AbortError' || error.code === 'ERR_CANCELED') + ); } async function processJob(job: t.ExecuteJob): Promise { - return withTraceContext(job.data._otel, () => withSpan('codeapi.job.process', { - 'messaging.system': 'bullmq', - 'messaging.operation.name': 'process', - 'messaging.message.id': typeof job.id === 'string' ? job.id : String(job.id ?? ''), - 'codeapi.language': job.data.payload?.language ?? 'unknown', - 'codeapi.execution_profile': job.data.executionProfile ?? 'legacy', - 'codeapi.worker_execution_profile': env.EXECUTION_PROFILE, - }, () => processJobInner(job), 'CONSUMER')); + return withTraceContext(job.data._otel, () => + withSpan( + 'codeapi.job.process', + { + 'messaging.system': 'bullmq', + 'messaging.operation.name': 'process', + 'messaging.message.id': + typeof job.id === 'string' ? job.id : String(job.id ?? ''), + 'codeapi.language': job.data.payload?.language ?? 'unknown', + 'codeapi.execution_profile': job.data.executionProfile ?? 'legacy', + 'codeapi.worker_execution_profile': env.EXECUTION_PROFILE, + }, + () => processJobInner(job), + 'CONSUMER', + ), + ); } async function processJobInner(job: t.ExecuteJob): Promise { const { payload, isPyPlot } = job.data; - const isSyntheticJob = job.data.isSynthetic === true || isSyntheticPrincipalSource(job.data.principalSource); + const isSyntheticJob = + job.data.isSynthetic === true || + isSyntheticPrincipalSource(job.data.principalSource); const language = payload?.language ?? 'unknown'; const endTimer = jobProcessingDuration.startTimer({ language }); activeJobs.inc({ language }); const controller = new AbortController(); - const deadlineAtMs = jobDeadlineAtMs(job.timestamp, env.JOB_TIMEOUT); + const cancellationTarget = + job.data.cancellable === true && job.id != null + ? { queueName: job.queueName, jobId: String(job.id) } + : undefined; + let cancellationRegistered = false; + const deadlineAtMs = jobDeadlineAtMs( + job.timestamp, + env.JOB_TIMEOUT, + Date.now(), + job.data.deadlineAtMs, + ); const remainingBudgetMs = Math.max(0, deadlineAtMs - Date.now()); - const timer = remainingBudgetMs > 0 - ? setTimeout(() => controller.abort(), remainingBudgetMs) - : undefined; - if (remainingBudgetMs === 0) controller.abort(); + const timer = + remainingBudgetMs > 0 + ? setTimeout(() => controller.abort('deadline'), remainingBudgetMs) + : undefined; + if (remainingBudgetMs === 0) controller.abort('deadline'); let egressGrantId: string | undefined; let egressGrantTokenForRestore: string | undefined; let revokeReason = 'completed'; + let completedResult = false; + let resultToCommit: t.ExecuteResult | undefined; + let resultCommittedAtHandoff = false; + const commitAtHandoff = + cancellationTarget != null && + job.data.workspaceId != null && + env.SANDBOX_BACKEND === 'remote-bridge'; try { + if (cancellationTarget != null) { + await jobCancellationRegistry.register(cancellationTarget, controller); + cancellationRegistered = true; + const claim = await claimJobExecution( + connection, + cancellationTarget, + jobCancellationRetentionSeconds( + env.JOB_TIMEOUT, + job.data.cancellationTtlSeconds, + ), + ); + if (claim.status === 'completed') return claim.result; + } if (controller.signal.aborted) { throw new Error(`Job timed out after ${env.JOB_TIMEOUT}ms`); } - validateQueuedExecutionProfile(job.data.executionProfile, env.EXECUTION_PROFILE); + validateQueuedExecutionProfile( + job.data.executionProfile, + env.EXECUTION_PROFILE, + ); validateQueuedSandboxBackend( job.data.sandboxBackend, env.SANDBOX_BACKEND, @@ -77,7 +156,10 @@ async function processJobInner(job: t.ExecuteJob): Promise { const nowSeconds = Math.floor(Date.now() / 1000); const prepared = await createGatewayEgressGrant({ payload, - claims: refreshEgressGrantClaims(job.data.egressGrantClaims, nowSeconds), + claims: refreshEgressGrantClaims( + job.data.egressGrantClaims, + nowSeconds, + ), isSynthetic: isSyntheticJob, signal: controller.signal, }); @@ -85,19 +167,27 @@ async function processJobInner(job: t.ExecuteJob): Promise { sandboxPayload = prepared.payload; egressGrantToken = prepared.egressGrantToken; egressGrantTokenForRestore = prepared.egressGrantToken; - executionManifestClaims = (env.EXECUTION_MANIFEST_PRIVATE_KEY || env.EXECUTION_MANIFEST_SECRET) - ? prepared.executionManifestClaims - : undefined; + executionManifestClaims = + env.EXECUTION_MANIFEST_PRIVATE_KEY || env.EXECUTION_MANIFEST_SECRET + ? prepared.executionManifestClaims + : undefined; } const delivery = prepareInputDelivery(payload, sandboxPayload); const sandboxRequest = buildSandboxExecuteRequest({ - ...(job.data.workspaceId == null ? {} : { programmaticTransferReserveMs: programmaticTransferReserveMs(env.JOB_TIMEOUT) }), + ...(job.data.workspaceId == null + ? {} + : { + programmaticTransferReserveMs: programmaticTransferReserveMs( + env.JOB_TIMEOUT, + ), + }), payload: delivery.payload, egressGrantToken, executionManifestClaims, maxOutputFileBytes: Math.min( - executionManifestClaims?.max_upload_bytes ?? env.EGRESS_GATEWAY_MAX_FILE_BYTES, + executionManifestClaims?.max_upload_bytes ?? + env.EGRESS_GATEWAY_MAX_FILE_BYTES, env.EGRESS_GATEWAY_MAX_FILE_BYTES, BRIDGE_WORKSPACE_PROGRAMMATIC_MAX_FILE_BYTES, ), @@ -121,21 +211,41 @@ async function processJobInner(job: t.ExecuteJob): Promise { * the transformed object makes that second call an idempotent no-op. */ const resultRestoreToken = egressGrantTokenForRestore; const finalizedSandboxResults = new WeakSet(); - const finalizeSandboxResult = async (result: SandboxRawResponse): Promise => { - if ( - resultRestoreToken === undefined || - resultRestoreToken.length === 0 || - finalizedSandboxResults.has(result) - ) { - return result; + const finalizeSandboxResult = async ( + result: SandboxRawResponse, + ): Promise => { + if (finalizedSandboxResults.has(result)) return result; + const restored = + resultRestoreToken == null || resultRestoreToken.length === 0 + ? result + : await restoreGatewaySandboxResult({ + grantId: egressGrantId, + egressGrantToken: resultRestoreToken, + result, + isSynthetic: isSyntheticJob, + signal: controller.signal, + }); + if (commitAtHandoff && cancellationTarget != null) { + // The bridge still owns its mutation fence here. A failed/ambiguous + // commit quarantines that root before it can serve a caller retry. + throwIfJobAborted(controller.signal); + const mapped = mapSandboxResult(restored); + const committed = await commitJobResult( + connection, + cancellationTarget, + mapped, + jobCancellationRetentionSeconds( + env.JOB_TIMEOUT, + job.data.cancellationTtlSeconds, + ), + deadlineAtMs, + ); + if (committed === 'cancelled') throw new Error(JOB_CANCELLED_MESSAGE); + if (committed === 'already_completed') + throw new Error('Duplicate mutation handoff; quarantining workspace'); + resultToCommit = mapped; + resultCommittedAtHandoff = true; } - const restored = await restoreGatewaySandboxResult({ - grantId: egressGrantId, - egressGrantToken: resultRestoreToken, - result, - isSynthetic: isSyntheticJob, - signal: controller.signal, - }); finalizedSandboxResults.add(restored); return restored; }; @@ -162,76 +272,115 @@ async function processJobInner(job: t.ExecuteJob): Promise { /* Stateful backends run this as a commit barrier after user code but * before checkpointing/reusing the mutated workspace. Stateless/HTTP * paths retain the worker-owned fallback immediately below. */ - sessionResultFinalizer: resultRestoreToken !== undefined && resultRestoreToken.length > 0 - ? finalizeSandboxResult - : undefined, + sessionResultFinalizer: + commitAtHandoff || + (resultRestoreToken !== undefined && resultRestoreToken.length > 0) + ? finalizeSandboxResult + : undefined, }, ); const responseData = await finalizeSandboxResult(responseRaw); + // Cancellation can arrive after sandbox exit while artifact restoration + // yields. Do not let BullMQ commit a success after Stop was acknowledged. + if (!resultCommittedAtHandoff) throwIfJobAborted(controller.signal); - if (!isSyntheticJob) { - logger.info('Sandbox response', summarizeSandboxResponse(responseData)); - } + function mapSandboxResult( + responseData: SandboxRawResponse, + ): t.ExecuteResult { + if (!isSyntheticJob) { + logger.info('Sandbox response', summarizeSandboxResponse(responseData)); + } - const { files } = responseData; - const run = responseData.run; - const stdout = applySystemReplacements(run?.stdout ?? ''); - const stderr = filterSystemLogs(run?.stderr ?? '', isPyPlot); + const { files } = responseData; + const run = responseData.run; + const stdout = applySystemReplacements(run?.stdout ?? ''); + const stderr = filterSystemLogs(run?.stderr ?? '', isPyPlot); - const result: t.ExecuteResult = { - session_id: responseData.session_id, - /* `files` is optional on the sandbox response (e.g. dry-run - * execute with no outputs); the public `ExecuteResult.files` is - * required and downstream callers always iterate it. Default to - * `[]` so the strictened response type from Phase B doesn't - * surface a regression that wasn't there before. */ - files: files ?? [], - ...(responseData.artifact_delivery != null - ? { artifact_delivery: responseData.artifact_delivery } - : {}), - ...(responseData.artifact_truncation != null - ? { artifact_truncation: responseData.artifact_truncation } - : {}), - stdout, - stderr, - ...(responseData.pending_tool_calls_payload != null - ? { pending_tool_calls_payload: responseData.pending_tool_calls_payload } - : {}), - }; + const result: t.ExecuteResult = { + session_id: responseData.session_id, + /* `files` is optional on the sandbox response (e.g. dry-run + * execute with no outputs); the public `ExecuteResult.files` is + * required and downstream callers always iterate it. Default to + * `[]` so the strictened response type from Phase B doesn't + * surface a regression that wasn't there before. */ + files: files ?? [], + ...(responseData.artifact_delivery != null + ? { artifact_delivery: responseData.artifact_delivery } + : {}), + ...(responseData.artifact_truncation != null + ? { artifact_truncation: responseData.artifact_truncation } + : {}), + stdout, + stderr, + ...(responseData.pending_tool_calls_payload != null + ? { + pending_tool_calls_payload: + responseData.pending_tool_calls_payload, + } + : {}), + }; - if (run) { - result.code = run.code ?? null; - result.signal = run.signal != null ? String(run.signal) : null; - result.message = run.message ?? null; - result.status = run.status ?? null; - result.wall_time = (run as Record).wall_time as number | null ?? null; - } + if (run) { + result.code = run.code ?? null; + result.signal = run.signal != null ? String(run.signal) : null; + result.message = run.message ?? null; + result.status = run.status ?? null; + result.wall_time = + ((run as Record).wall_time as number | null) ?? null; + } - if (result.message || result.signal) { - logger.warn('Sandbox execution error metadata', { - session_id: responseData.session_id, - code: result.code, - signal: result.signal, - message: summarizeText(result.message), - status: result.status, - wall_time: result.wall_time, - }); + if (result.message || result.signal) { + logger.warn('Sandbox execution error metadata', { + session_id: responseData.session_id, + code: result.code, + signal: result.signal, + message: summarizeText(result.message), + status: result.status, + wall_time: result.wall_time, + }); + } + + return result; } + const result = resultToCommit ?? mapSandboxResult(responseData); + completedResult = true; + resultToCommit = result; return result; } catch (error) { - revokeReason = controller.signal.aborted || isAbortError(error) ? 'timeout' : 'failed'; + // Bridge fence cleanup can fail after the outcome was durably committed. + // Preserve the winning result; the bridge retains/quarantines its fence. + if (resultCommittedAtHandoff && resultToCommit != null) + return resultToCommit; + const clientDisconnected = + controller.signal.aborted && + controller.signal.reason === CLIENT_DISCONNECT_REASON; + revokeReason = clientDisconnected + ? 'cancelled' + : controller.signal.aborted || isAbortError(error) + ? 'timeout' + : 'failed'; const errorDetails = getAxiosErrorDetails(error); - logger.error('Error processing job', errorDetails); + if (clientDisconnected) { + logger.info('Job cancelled after client disconnected', { + queueName: job.queueName, + jobId: job.id, + executionId: job.data.executionId, + }); + } else { + logger.error('Error processing job', errorDetails); + } const deadlineFailure = workerDeadlineFailure( error, - controller.signal.aborted, + controller.signal.aborted && !clientDisconnected, env.JOB_TIMEOUT, ); if (deadlineFailure) { throw deadlineFailure; + } else if (clientDisconnected) { + throw new Error(JOB_CANCELLED_MESSAGE); } else if (error instanceof SandboxBackendError) { throw new Error(`${error.code}: ${error.message}`); } else if (error instanceof SessionFilesError) { @@ -251,17 +400,67 @@ async function processJobInner(job: t.ExecuteJob): Promise { if (egressGrantId || egressGrantTokenForRestore) { await revokeGatewayEgressGrant({ grantId: egressGrantId, - egressGrantToken: egressGrantId ? undefined : egressGrantTokenForRestore, + egressGrantToken: egressGrantId + ? undefined + : egressGrantTokenForRestore, isSynthetic: isSyntheticJob, reason: revokeReason, timeoutMs: env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS, }).catch(error => { - logger.error('Failed to revoke egress grant', { grantId: egressGrantId, error: getAxiosErrorDetails(error) }); + logger.error('Failed to revoke egress grant', { + grantId: egressGrantId, + error: getAxiosErrorDetails(error), + }); }); } + let lateCommitFailure = + completedResult && !resultCommittedAtHandoff + ? jobResultCommitFailure(controller.signal, env.JOB_TIMEOUT) + : undefined; + if ( + completedResult && + !resultCommittedAtHandoff && + cancellationTarget != null && + lateCommitFailure == null + ) { + try { + const committed = await commitJobResult( + connection, + cancellationTarget, + resultToCommit, + jobCancellationRetentionSeconds( + env.JOB_TIMEOUT, + job.data.cancellationTtlSeconds, + ), + deadlineAtMs, + ); + if (committed === 'cancelled') { + lateCommitFailure = new Error(JOB_CANCELLED_MESSAGE); + } else if (committed === 'already_completed') { + lateCommitFailure = new Error( + 'Duplicate result handoff; refusing replacement', + ); + } + } catch (error) { + lateCommitFailure = + error instanceof Error ? error : new Error('Result commit failed'); + } + } if (timer) clearTimeout(timer); + if (cancellationTarget != null && cancellationRegistered) { + await jobCancellationRegistry + .unregister(cancellationTarget, controller) + .catch(error => { + logger.warn('Failed to clear queued execution cancellation state', { + queueName: cancellationTarget.queueName, + jobId: cancellationTarget.jobId, + error: getAxiosErrorDetails(error), + }); + }); + } endTimer(); activeJobs.dec({ language }); + if (lateCommitFailure != null) throw lateCommitFailure; } } @@ -304,21 +503,31 @@ otherWorker.on('completed', job => { }); pyWorker.on('failed', (job, err) => { + if (err.message === JOB_CANCELLED_MESSAGE) { + logger.info(`[${WORKER_ID}] Python job ${job?.id} cancelled`); + jobsCancelled.inc({ language: 'python' }); + return; + } logger.error(`[${WORKER_ID}] Python job ${job?.id} failed`, err); jobsFailed.inc({ language: 'python' }); }); otherWorker.on('failed', (job, err) => { + if (err.message === JOB_CANCELLED_MESSAGE) { + logger.info(`[${WORKER_ID}] Other job ${job?.id} cancelled`); + jobsCancelled.inc({ language: 'other' }); + return; + } logger.error(`[${WORKER_ID}] Other job ${job?.id} failed`, err); jobsFailed.inc({ language: 'other' }); }); -pyWorker.on('error', (err) => { +pyWorker.on('error', err => { logger.error(`[${WORKER_ID}] Python worker error`, err); workerRunning.set({ worker_type: 'python' }, 0); }); -otherWorker.on('error', (err) => { +otherWorker.on('error', err => { logger.error(`[${WORKER_ID}] Other worker error`, err); workerRunning.set({ worker_type: 'other' }, 0); });