From 169ad29fe438824f2e8307fd2644908ed106021d Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Sep 2026 04:28:33 -0400 Subject: [PATCH 01/11] fix: cancel replay jobs across API and worker processes --- docs/remote-bridge/README.md | 8 + service/src/job-cancellation.test.ts | 179 ++++++++++++++++ service/src/job-cancellation.ts | 194 ++++++++++++++++++ service/src/metrics.ts | 6 + service/src/middleware/limits.ts | 13 ++ service/src/programmatic-cancellation.test.ts | 125 +++++++++++ service/src/programmatic-cancellation.ts | 166 +++++++++++++++ service/src/queue.ts | 14 +- service/src/service/programmatic-router.ts | 189 +++++++++++++++-- service/src/types/service.ts | 2 + service/src/workers.ts | 61 +++++- 11 files changed, 930 insertions(+), 27 deletions(-) create mode 100644 service/src/job-cancellation.test.ts create mode 100644 service/src/job-cancellation.ts create mode 100644 service/src/programmatic-cancellation.test.ts create mode 100644 service/src/programmatic-cancellation.ts diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index 99a0dece..42e3abaa 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -256,6 +256,14 @@ 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. - 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/service/src/job-cancellation.test.ts b/service/src/job-cancellation.test.ts new file mode 100644 index 00000000..70c9e87e --- /dev/null +++ b/service/src/job-cancellation.test.ts @@ -0,0 +1,179 @@ +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, + jobCancellationInternals, + requestJobCancellation, + waitForJobWithCancellation, +} from './job-cancellation'; + +class FakeSubscriber extends EventEmitter { + subscribed?: string; + closed = false; + + async subscribe(channel: string): Promise { + this.subscribed = channel; + return 1; + } + + async quit(): Promise<'OK'> { + this.closed = true; + return 'OK'; + } +} + +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(); + readonly existing = new Set(); + readonly deleted: string[] = []; + readonly transactions: FakeTransaction[] = []; + + duplicate(): FakeSubscriber { + return this.subscriber; + } + + async exists(key: string): Promise { + return this.existing.has(key) ? 1 : 0; + } + + async mget(...keys: string[]): Promise> { + 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('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('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) => setImmediate(resolve)); + + 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('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, + remove: async () => { + removed = true; + }, + } as unknown as Job; + + const waiting = waitForJobWithCancellation({ + commands: redis(fake), + 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.transactions[0]?.operations[0]).toEqual([ + 'set', + jobCancellationInternals.cancellationKey({ queueName: 'other', jobId: 'job-4' }), + '1', + 'EX', + 120, + ]); +}); diff --git a/service/src/job-cancellation.ts b/service/src/job-cancellation.ts new file mode 100644 index 00000000..a1eb9db5 --- /dev/null +++ b/service/src/job-cancellation.ts @@ -0,0 +1,194 @@ +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 readonly subscriber: IORedis; + private readonly controllers = new Map< + string, + { target: JobTarget; controller: AbortController } + >(); + private startPromise?: Promise; + + constructor(private readonly commands: IORedis) { + this.subscriber = commands.duplicate(); + this.subscriber.on('error', () => { + // ioredis reconnects using the shared policy. The listener prevents a + // transient subscriber outage from becoming an uncaught process error. + }); + this.subscriber.on('ready', () => { + void this.reconcile().catch(() => undefined); + }); + } + + 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 != null) { + entries[index]?.controller.abort(CLIENT_DISCONNECT_REASON); + } + }); + } + + private start(): Promise { + if (this.startPromise != null) return this.startPromise; + const starting = (async () => { + this.subscriber.on('message', (channel, raw) => { + if (channel !== JOB_CANCELLATION_CHANNEL) return; + const target = parseTarget(raw); + if (target == null) return; + this.controllers + .get(targetKey(target)) + ?.controller.abort(CLIENT_DISCONNECT_REASON); + }); + await this.subscriber.subscribe(JOB_CANCELLATION_CHANNEL); + })(); + this.startPromise = starting.catch(error => { + this.startPromise = undefined; + throw error; + }); + return this.startPromise; + } + + async register(target: JobTarget, controller: AbortController): Promise { + this.controllers.set(targetKey(target), { target, controller }); + try { + await this.start(); + if (await this.commands.exists(cancellationKey(target))) { + controller.abort(CLIENT_DISCONNECT_REASON); + } + } catch (error) { + this.controllers.delete(targetKey(target)); + throw error; + } + } + + async unregister(target: JobTarget): Promise { + this.controllers.delete(targetKey(target)); + await this.commands.del(cancellationKey(target)); + } + + async close(): Promise { + if (this.startPromise == null) return; + this.controllers.clear(); + await this.subscriber.quit(); + } +} + +export async function requestJobCancellation( + commands: IORedis, + target: JobTarget, + ttlSeconds: number, +): Promise { + const payload = JSON.stringify(target); + const transaction = commands.multi(); + transaction.set(cancellationKey(target), '1', 'EX', Math.max(1, ttlSeconds)); + transaction.publish(JOB_CANCELLATION_CHANNEL, payload); + const result = await transaction.exec(); + if (result == null) { + throw new Error('Redis transaction aborted while cancelling queued execution'); + } + const failure = result.find(([error]) => error != null)?.[0]; + if (failure != null) throw failure; +} + +function abortError(): Error { + return new DOMException('Programmatic execution request disconnected', 'AbortError'); +} + +export async function waitForJobWithCancellation(args: { + commands: IORedis; + job: Job; + events: QueueEvents; + timeoutMs: number; + cancellationTtlSeconds: number; + signal?: AbortSignal; +}): Promise { + const { commands, job, events, timeoutMs, cancellationTtlSeconds, signal } = args; + const completion = job.waitUntilFinished(events, timeoutMs); + if (signal == null) return completion; + + const target = { queueName: job.queueName, jobId: String(job.id) }; + let removeAbortListener = (): void => {}; + const cancelled = new Promise((_, reject) => { + let cancelling = false; + const cancel = (): void => { + if (cancelling) return; + cancelling = true; + void requestJobCancellation(commands, target, cancellationTtlSeconds) + .then(async () => { + // 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 job.remove().catch(() => undefined); + }) + .then(() => reject(abortError()), reject); + }; + removeAbortListener = (): void => signal.removeEventListener('abort', cancel); + signal.addEventListener('abort', cancel, { once: true }); + if (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, cancelled]); + } finally { + removeAbortListener(); + } +} + +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..ae955ee3 --- /dev/null +++ b/service/src/programmatic-cancellation.test.ts @@ -0,0 +1,125 @@ +import { afterEach, beforeEach, expect, test } from 'bun:test'; +import type IORedis from 'ioredis'; +import { startTestRedis } from './test/redis'; +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('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('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('the owner releases cancellation state after settlement', async () => { + const requestId = 'request_release_cancel_1'; + await reserveProgrammaticCancellation({ + redis, + requestId, + owner: 'owner-a', + ttlSeconds: 60, + }); + await releaseProgrammaticCancellation({ redis, requestId, owner: 'owner-a' }); + expect(await redis.exists(programmaticCancellationInternals.requestKey(requestId))).toBe(0); +}); diff --git a/service/src/programmatic-cancellation.ts b/service/src/programmatic-cancellation.ts new file mode 100644 index 00000000..880ba2f6 --- /dev/null +++ b/service/src/programmatic-cancellation.ts @@ -0,0 +1,166 @@ +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 not existing then + redis.call('HSET', key, 'owner', owner, 'cancelled', '0') +end +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') +redis.call('EXPIRE', key, ttl) +local queueName = redis.call('HGET', key, 'queueName') +local jobId = redis.call('HGET', key, 'jobId') +if queueName and jobId then return {1, queueName, jobId} end +return {1} +`; + +const RELEASE_SCRIPT = ` +if redis.call('HGET', KEYS[1], 'owner') == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0 +`; + +export async function reserveProgrammaticCancellation(args: { + redis: IORedis; + requestId: string; + owner: string; + ttlSeconds: number; +}): Promise<'active' | 'cancelled' | 'forbidden'> { + const result = Number(await args.redis.eval( + RESERVE_SCRIPT, + 1, + requestKey(args.requestId), + args.owner, + Math.max(1, args.ttlSeconds), + )); + if (result < 0) return 'forbidden'; + 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..064489af 100644 --- a/service/src/queue.ts +++ b/service/src/queue.ts @@ -19,6 +19,7 @@ import type { import logger from './logger'; import { redisKeepAliveOptions } from './redis-options'; import { bullmqQueueJobs, registerBullmqQueueMetricsCollector } from './metrics'; +import { JobCancellationRegistry } from './job-cancellation'; const MAX_RECONNECT_ATTEMPTS = 5; const RECONNECT_DELAY = 2000; @@ -60,6 +61,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 @@ -163,8 +165,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/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index 93bbc0c5..109aa2ac 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -4,13 +4,28 @@ 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, getExecutionQueueBinding, } from '../queue'; +import { + CLIENT_DISCONNECT_REASON, + JOB_CANCELLED_MESSAGE, + requestJobCancellation, + waitForJobWithCancellation, +} from '../job-cancellation'; +import { + CODEAPI_PROGRAMMATIC_REQUEST_HEADER, + attachProgrammaticCancellationTarget, + cancelProgrammaticRequest, + normalizeProgrammaticRequestId, + programmaticCancellationOwner, + releaseProgrammaticCancellation, + reserveProgrammaticCancellation, +} from '../programmatic-cancellation'; import { createProgrammaticPayload, extractPendingFromControlPayload, @@ -107,6 +122,8 @@ 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; const router = Router(); @@ -322,6 +339,8 @@ async function runReplayIteration( state: ExecutionState, apiKeyId: string, userId: string, + signal?: AbortSignal, + cancellation?: { requestId: string; owner: string }, ): Promise { const history = await loadToolHistory(state.execution_id); const rawPayload = buildReplayPayload(req, state, history); @@ -389,6 +408,7 @@ async function runReplayIteration( ...(state.workspaceId != null ? { workspaceId: state.workspaceId } : {}), + cancellable: true, runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, executionManifestClaims: sandboxSecurity.executionManifestClaims, @@ -403,7 +423,44 @@ async function runReplayIteration( ); jobsSubmitted.inc({ language }); - return job.waitUntilFinished(events, JOB_COMPLETION_WAIT_TIMEOUT_MS); + if (cancellation != null) { + const target = { queueName: job.queueName, jobId: String(job.id) }; + try { + const attachment = await attachProgrammaticCancellationTarget({ + redis: connection, + requestId: cancellation.requestId, + owner: cancellation.owner, + target, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + }); + if (attachment === 'forbidden') { + throw new Error('Programmatic cancellation request ownership changed'); + } + if (attachment === 'cancelled') { + await requestJobCancellation( + connection, + target, + PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + ); + } + } catch (error) { + await requestJobCancellation( + connection, + target, + PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + ).catch(() => undefined); + throw error; + } + } + + return waitForJobWithCancellation({ + commands: connection, + job, + events, + timeoutMs: JOB_COMPLETION_WAIT_TIMEOUT_MS, + cancellationTtlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + signal, + }); } function isSandboxRunSuccess(result: t.ExecuteResult): boolean { @@ -858,26 +915,62 @@ async function runAndRespond( apiKeyId: string, userId: string, ): 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; - }); + const rawRequestId = req.header(CODEAPI_PROGRAMMATIC_REQUEST_HEADER); + const requestId = normalizeProgrammaticRequestId(rawRequestId); + const cancellationOwner = + requestId != null ? programmaticCancellationOwner(req, userId) : undefined; + if (requestId != null && cancellationOwner != null) { + const reservation = await reserveProgrammaticCancellation({ + redis: connection, + requestId, + owner: cancellationOwner, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + }); + if (reservation === 'forbidden') { + res.status(409).json({ error: 'Programmatic request ID is already in use' }); + return; + } + } + const disconnectController = new AbortController(); + const disconnect = (): void => { + if (!res.writableFinished) { + disconnectController.abort(CLIENT_DISCONNECT_REASON); + disposeDisconnectListeners(); + } + }; + const disposeDisconnectListeners = (): void => { + req.removeListener('aborted', disconnect); + req.removeListener('close', disconnect); + res.removeListener('close', disconnect); + res.removeListener('finish', disposeDisconnectListeners); + }; + req.once('aborted', disconnect); + req.once('close', disconnect); + res.once('close', disconnect); + res.once('finish', disposeDisconnectListeners); + const isDisconnected = (): boolean => disconnectController.signal.aborted; let result: t.ExecuteResult; try { - result = await runReplayIteration(req, state, apiKeyId, userId); + result = await runReplayIteration( + req, + state, + apiKeyId, + userId, + disconnectController.signal, + requestId != null && cancellationOwner != null + ? { requestId, owner: cancellationOwner } + : undefined, + ); } 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()) { const publicFailure = publicExecutionFailure(err); @@ -890,6 +983,19 @@ async function runAndRespond( }); } return; + } finally { + if (requestId != null && cancellationOwner != null) { + await releaseProgrammaticCancellation({ + redis: connection, + requestId, + owner: cancellationOwner, + }).catch(error => { + logger.warn('Failed to release programmatic cancellation request', { + requestId, + error: (error as Error).message, + }); + }); + } } if (isDisconnected()) { @@ -1076,6 +1182,49 @@ 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) { + await requestJobCancellation( + connection, + cancellation.target, + PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + ); + } + 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 +1244,10 @@ 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); + if (rawRequestId != null && normalizeProgrammaticRequestId(rawRequestId) == 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; diff --git a/service/src/types/service.ts b/service/src/types/service.ts index d17b5a99..f10405df 100644 --- a/service/src/types/service.ts +++ b/service/src/types/service.ts @@ -297,6 +297,8 @@ 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; /** 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..dd59a9c8 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -2,8 +2,8 @@ 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 { 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'; @@ -16,6 +16,10 @@ import { getSandboxBackend, SandboxBackendError, type SandboxRawResponse } from import { isSyntheticPrincipalSource } from './auth/synthetic'; import { withSpan, withTraceContext } from './telemetry'; import { workerDeadlineFailure } from './worker-error'; +import { + CLIENT_DISCONNECT_REASON, + JOB_CANCELLED_MESSAGE, +} from './job-cancellation'; import logger from './logger'; import { validateQueuedExecutionProfile, @@ -49,17 +53,25 @@ async function processJobInner(job: t.ExecuteJob): Promise { activeJobs.inc({ language }); const controller = new AbortController(); + 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); const remainingBudgetMs = Math.max(0, deadlineAtMs - Date.now()); const timer = remainingBudgetMs > 0 - ? setTimeout(() => controller.abort(), remainingBudgetMs) + ? setTimeout(() => controller.abort('deadline'), remainingBudgetMs) : undefined; - if (remainingBudgetMs === 0) controller.abort(); + if (remainingBudgetMs === 0) controller.abort('deadline'); let egressGrantId: string | undefined; let egressGrantTokenForRestore: string | undefined; let revokeReason = 'completed'; try { + if (cancellationTarget != null) { + await jobCancellationRegistry.register(cancellationTarget, controller); + cancellationRegistered = true; + } if (controller.signal.aborted) { throw new Error(`Job timed out after ${env.JOB_TIMEOUT}ms`); } @@ -221,17 +233,33 @@ async function processJobInner(job: t.ExecuteJob): Promise { return result; } catch (error) { - revokeReason = controller.signal.aborted || isAbortError(error) ? 'timeout' : 'failed'; + 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) { @@ -260,6 +288,15 @@ async function processJobInner(job: t.ExecuteJob): Promise { }); } if (timer) clearTimeout(timer); + if (cancellationTarget != null && cancellationRegistered) { + await jobCancellationRegistry.unregister(cancellationTarget).catch(error => { + logger.warn('Failed to clear queued execution cancellation state', { + queueName: cancellationTarget.queueName, + jobId: cancellationTarget.jobId, + error: getAxiosErrorDetails(error), + }); + }); + } endTimer(); activeJobs.dec({ language }); } @@ -304,11 +341,21 @@ 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' }); }); From e8fb38d8c2ce0248847be40ce546b1d6310ae0ce Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Sep 2026 06:12:15 -0400 Subject: [PATCH 02/11] fix: drain worker cancellation watches promptly --- packages/code/src/worker.test.ts | 103 +++++++++ packages/code/src/worker.ts | 46 ++-- packages/code/src/workspace-worker.test.ts | 244 +++++++++++++++++++++ 3 files changed, 380 insertions(+), 13 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 68d95b28..09f059af 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -694,6 +694,109 @@ 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 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 = ''; diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index a0517bd9..0eefbd0d 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, @@ -1674,14 +1679,21 @@ 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; } @@ -1719,6 +1731,7 @@ export class BridgeWorker { 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', @@ -2170,17 +2183,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, }); @@ -2213,6 +2232,7 @@ export class BridgeWorker { if (signal.aborted) return; } finally { clearTimeout(timeout); + signal.removeEventListener('abort', abortPoll); executionController.signal.removeEventListener('abort', abortPoll); } } 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; From 599cd1102e101e1256e2e30f1eb006c129d473a7 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Sep 2026 06:56:51 -0400 Subject: [PATCH 03/11] fix: close programmatic cancellation races --- service/src/job-cancellation.test.ts | 73 ++++++ service/src/job-cancellation.ts | 105 ++++++--- service/src/programmatic-cancellation.test.ts | 18 ++ service/src/programmatic-cancellation.ts | 7 +- service/src/queue.ts | 14 ++ service/src/request-disconnect.test.ts | 60 +++++ service/src/request-disconnect.ts | 49 ++++ service/src/service/programmatic-router.ts | 213 ++++++++++++------ 8 files changed, 444 insertions(+), 95 deletions(-) create mode 100644 service/src/request-disconnect.test.ts create mode 100644 service/src/request-disconnect.ts diff --git a/service/src/job-cancellation.test.ts b/service/src/job-cancellation.test.ts index 70c9e87e..0582e9c3 100644 --- a/service/src/job-cancellation.test.ts +++ b/service/src/job-cancellation.test.ts @@ -6,6 +6,7 @@ import { CLIENT_DISCONNECT_REASON, JobCancellationRegistry, jobCancellationInternals, + removeJobIfWaiting, requestJobCancellation, waitForJobWithCancellation, } from './job-cancellation'; @@ -13,8 +14,13 @@ import { 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; } @@ -23,6 +29,10 @@ class FakeSubscriber extends EventEmitter { this.closed = true; return 'OK'; } + + disconnect(): void { + this.closed = true; + } } class FakeTransaction { @@ -45,11 +55,13 @@ class FakeTransaction { class FakeRedis { readonly subscriber = new FakeSubscriber(); + duplicateCalls = 0; readonly existing = new Set(); readonly deleted: string[] = []; readonly transactions: FakeTransaction[] = []; duplicate(): FakeSubscriber { + this.duplicateCalls += 1; return this.subscriber; } @@ -78,6 +90,35 @@ 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('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' }; @@ -152,6 +193,7 @@ test('disconnect frees a waiting job and rejects promptly', async () => { id: 'job-4', queueName: 'other', waitUntilFinished: () => never, + getState: async () => 'waiting', remove: async () => { removed = true; }, @@ -177,3 +219,34 @@ test('disconnect frees a waiting job and rejects promptly', async () => { 120, ]); }); + +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('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 index a1eb9db5..c6c2778c 100644 --- a/service/src/job-cancellation.ts +++ b/service/src/job-cancellation.ts @@ -47,22 +47,38 @@ function parseTarget(raw: string): JobTarget | undefined { * this path, so ordinary queue traffic pays no extra Redis round trips. */ export class JobCancellationRegistry { - private readonly subscriber: IORedis; + private subscriber?: IORedis; private readonly controllers = new Map< string, { target: JobTarget; controller: AbortController } >(); private startPromise?: Promise; + private closed = false; - constructor(private readonly commands: IORedis) { - this.subscriber = commands.duplicate(); - this.subscriber.on('error', () => { + 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. - }); - this.subscriber.on('ready', () => { - void this.reconcile().catch(() => undefined); - }); + }; + + private readonly onSubscriberReady = (): void => { + void this.reconcile().catch(() => undefined); + }; + + private readonly onSubscriberMessage = (channel: string, raw: string): void => { + if (channel !== JOB_CANCELLATION_CHANNEL) return; + const target = parseTarget(raw); + if (target == null) return; + this.controllers + .get(targetKey(target)) + ?.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); } private async reconcile(): Promise { @@ -79,23 +95,30 @@ export class JobCancellationRegistry { } 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 () => { - this.subscriber.on('message', (channel, raw) => { - if (channel !== JOB_CANCELLATION_CHANNEL) return; - const target = parseTarget(raw); - if (target == null) return; - this.controllers - .get(targetKey(target)) - ?.controller.abort(CLIENT_DISCONNECT_REASON); - }); - await this.subscriber.subscribe(JOB_CANCELLATION_CHANNEL); + 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); + 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.catch(error => { - this.startPromise = undefined; - throw error; + this.startPromise = starting; + void starting.catch(() => { + if (this.startPromise === starting) this.startPromise = undefined; }); - return this.startPromise; + return starting; } async register(target: JobTarget, controller: AbortController): Promise { @@ -117,9 +140,15 @@ export class JobCancellationRegistry { } async close(): Promise { - if (this.startPromise == null) return; + this.closed = true; this.controllers.clear(); - await this.subscriber.quit(); + await this.startPromise?.catch(() => undefined); + const subscriber = this.subscriber; + this.subscriber = undefined; + this.startPromise = undefined; + if (subscriber == null) return; + this.detachSubscriber(subscriber); + await subscriber.quit(); } } @@ -140,7 +169,29 @@ export async function requestJobCancellation( if (failure != null) throw failure; } -function abortError(): Error { +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'); } @@ -168,9 +219,9 @@ export async function waitForJobWithCancellation(args: { // 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 job.remove().catch(() => undefined); + await removeJobIfWaiting(job).catch(() => false); }) - .then(() => reject(abortError()), reject); + .then(() => reject(programmaticCancellationError()), reject); }; removeAbortListener = (): void => signal.removeEventListener('abort', cancel); signal.addEventListener('abort', cancel, { once: true }); diff --git a/service/src/programmatic-cancellation.test.ts b/service/src/programmatic-cancellation.test.ts index ae955ee3..ac60f7ce 100644 --- a/service/src/programmatic-cancellation.test.ts +++ b/service/src/programmatic-cancellation.test.ts @@ -80,6 +80,24 @@ test('cancellation after attachment returns the exact queue target', async () => }); }); +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({ diff --git a/service/src/programmatic-cancellation.ts b/service/src/programmatic-cancellation.ts index 880ba2f6..7b5af79d 100644 --- a/service/src/programmatic-cancellation.ts +++ b/service/src/programmatic-cancellation.ts @@ -48,9 +48,11 @@ 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') `; @@ -94,7 +96,7 @@ export async function reserveProgrammaticCancellation(args: { requestId: string; owner: string; ttlSeconds: number; -}): Promise<'active' | 'cancelled' | 'forbidden'> { +}): Promise<'active' | 'cancelled' | 'duplicate' | 'forbidden'> { const result = Number(await args.redis.eval( RESERVE_SCRIPT, 1, @@ -102,7 +104,8 @@ export async function reserveProgrammaticCancellation(args: { args.owner, Math.max(1, args.ttlSeconds), )); - if (result < 0) return 'forbidden'; + if (result === -1) return 'forbidden'; + if (result === -2) return 'duplicate'; return result === 1 ? 'cancelled' : 'active'; } diff --git a/service/src/queue.ts b/service/src/queue.ts index 064489af..4bf1c8e1 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'; @@ -112,6 +113,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); 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 109aa2ac..8170ff41 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -10,10 +10,12 @@ import { pyQueueEvents, connection, getExecutionQueueBinding, + getExistingExecutionJob, } from '../queue'; import { - CLIENT_DISCONNECT_REASON, JOB_CANCELLED_MESSAGE, + programmaticCancellationError, + removeJobIfWaiting, requestJobCancellation, waitForJobWithCancellation, } from '../job-cancellation'; @@ -53,6 +55,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, @@ -125,6 +128,16 @@ const JOB_COMPLETION_WAIT_TIMEOUT_MS = jobCompletionWaitTimeoutMs( 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(); function sendFileRefAuthorizationError( @@ -342,6 +355,7 @@ async function runReplayIteration( 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; @@ -388,6 +402,7 @@ async function runReplayIteration( state.executionProfile ?? env.EXECUTION_PROFILE, state.executionProfileSource ?? env.EXECUTION_PROFILE_SOURCE, ); + if (signal?.aborted) throw programmaticCancellationError(); const job = await queue.add( Jobs.execute, { @@ -425,6 +440,7 @@ async function runReplayIteration( if (cancellation != null) { const target = { queueName: job.queueName, jobId: String(job.id) }; + let cancellationPublished = false; try { const attachment = await attachProgrammaticCancellationTarget({ redis: connection, @@ -442,13 +458,18 @@ async function runReplayIteration( target, PROGRAMMATIC_CANCELLATION_TTL_SECONDS, ); + cancellationPublished = true; + await removeJobIfWaiting(job).catch(() => false); + throw programmaticCancellationError(); } } catch (error) { - await requestJobCancellation( - connection, - target, - PROGRAMMATIC_CANCELLATION_TTL_SECONDS, - ).catch(() => undefined); + if (!cancellationPublished) { + await requestJobCancellation( + connection, + target, + PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + ).catch(() => undefined); + } throw error; } } @@ -482,6 +503,7 @@ async function handleReplayInitial( bridgeWorkerId?: string; workspaceId?: string; }, + cancellation: ReplayRequestCancellation, ): Promise { const { apiKeyId, userId, bridgeWorkerId, workspaceId } = params; const { code, tools, user_id, files } = @@ -600,6 +622,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; @@ -679,7 +714,7 @@ async function handleReplayInitial( timeout, }); - await runAndRespond(req, res, state, apiKeyId, userId); + await runAndRespond(req, res, state, apiKeyId, userId, cancellation); } async function handleReplayContinuation( @@ -691,6 +726,7 @@ async function handleReplayContinuation( decoded: { execution_id: string }; tool_results: NonNullable; }, + cancellation: ReplayRequestCancellation, ): Promise { const { apiKeyId, userId, decoded, tool_results } = params; @@ -749,6 +785,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` @@ -902,7 +952,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); } @@ -914,42 +971,8 @@ async function runAndRespond( state: ExecutionState, apiKeyId: string, userId: string, + cancellation: ReplayRequestCancellation, ): Promise { - const rawRequestId = req.header(CODEAPI_PROGRAMMATIC_REQUEST_HEADER); - const requestId = normalizeProgrammaticRequestId(rawRequestId); - const cancellationOwner = - requestId != null ? programmaticCancellationOwner(req, userId) : undefined; - if (requestId != null && cancellationOwner != null) { - const reservation = await reserveProgrammaticCancellation({ - redis: connection, - requestId, - owner: cancellationOwner, - ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, - }); - if (reservation === 'forbidden') { - res.status(409).json({ error: 'Programmatic request ID is already in use' }); - return; - } - } - const disconnectController = new AbortController(); - const disconnect = (): void => { - if (!res.writableFinished) { - disconnectController.abort(CLIENT_DISCONNECT_REASON); - disposeDisconnectListeners(); - } - }; - const disposeDisconnectListeners = (): void => { - req.removeListener('aborted', disconnect); - req.removeListener('close', disconnect); - res.removeListener('close', disconnect); - res.removeListener('finish', disposeDisconnectListeners); - }; - req.once('aborted', disconnect); - req.once('close', disconnect); - res.once('close', disconnect); - res.once('finish', disposeDisconnectListeners); - const isDisconnected = (): boolean => disconnectController.signal.aborted; - let result: t.ExecuteResult; try { result = await runReplayIteration( @@ -957,10 +980,8 @@ async function runAndRespond( state, apiKeyId, userId, - disconnectController.signal, - requestId != null && cancellationOwner != null - ? { requestId, owner: cancellationOwner } - : undefined, + cancellation.signal, + cancellation.request, ); } catch (err) { const cancelled = @@ -972,7 +993,7 @@ async function runAndRespond( 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; @@ -983,22 +1004,9 @@ async function runAndRespond( }); } return; - } finally { - if (requestId != null && cancellationOwner != null) { - await releaseProgrammaticCancellation({ - redis: connection, - requestId, - owner: cancellationOwner, - }).catch(error => { - logger.warn('Failed to release programmatic cancellation request', { - requestId, - error: (error as Error).message, - }); - }); - } } - if (isDisconnected()) { + if (cancellation.isDisconnected()) { logger.info('Client disconnected during replay; cleaning up', { execution_id: state.execution_id, }); @@ -1108,7 +1116,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 @@ -1213,6 +1221,20 @@ router.post( cancellation.target, PROGRAMMATIC_CANCELLATION_TTL_SECONDS, ); + 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) { @@ -1245,7 +1267,8 @@ router.post( req.body as t.ProgrammaticRequestBody; const rawBody = req.body as Record; const rawRequestId = req.header(CODEAPI_PROGRAMMATIC_REQUEST_HEADER); - if (rawRequestId != null && normalizeProgrammaticRequestId(rawRequestId) == null) { + 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; @@ -1304,7 +1327,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 @@ -1337,7 +1404,7 @@ router.post( userId, decoded, tool_results, - }); + }, cancellation); } return await handleBlocking(req, res, { apiKeyId, userId }); } @@ -1356,7 +1423,7 @@ router.post( userId, bridgeWorkerId, workspaceId, - }); + }, cancellation); } if (workspaceId != null) { return res.status(400).json({ @@ -1374,6 +1441,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, + }); + }); + } } }, ); From 3ebec6be02141d94513cce2db2dbd128aa515d0c Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Sep 2026 07:02:11 -0400 Subject: [PATCH 04/11] fix: preserve cancellation response ordering --- packages/code/src/worker.ts | 9 +++++++++ service/src/bridge/concurrent-worker.test.ts | 20 +++++++++++++++++--- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 0eefbd0d..970081fe 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -2219,6 +2219,13 @@ export class BridgeWorker { incarnationId: this.incarnationId, }, pollController.signal, + () => { + // 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. + signal.removeEventListener('abort', abortPoll); + }, ); if (response.cancelled) { executionController.abort(); @@ -2242,6 +2249,7 @@ export class BridgeWorker { url: string, body: object, signal?: AbortSignal, + onResponseHeaders?: () => void, ): Promise { const requestBody = JSON.stringify(body); const response = await this.fetchImpl(url, { @@ -2253,6 +2261,7 @@ export class BridgeWorker { body: requestBody, signal, }); + onResponseHeaders?.(); let payload: unknown; try { payload = await response.json(); 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); From 2ef0779575a27d2900ddc78d156e79ae0c7e68a9 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Sep 2026 07:55:53 -0400 Subject: [PATCH 05/11] fix: close distributed cancellation races --- packages/code/src/worker.test.ts | 66 ++++++++++++ packages/code/src/worker.ts | 26 +++-- service/src/job-cancellation.test.ts | 111 ++++++++++++++++++- service/src/job-cancellation.ts | 119 +++++++++++++++++---- service/src/service/programmatic-router.ts | 58 ++++------ service/src/workers.ts | 6 +- 6 files changed, 322 insertions(+), 64 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 09f059af..ea785dde 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -755,6 +755,72 @@ test('worker stops an outstanding cancellation request before settling completed 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({ diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 970081fe..45c096d2 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -104,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; @@ -1123,7 +1124,6 @@ export class BridgeWorker { assignment: BridgeAssignment, stopSignal: AbortSignal, serverClockOffsetMs: number, - requestSignal?: AbortSignal, ): Promise { const identity = this.options.identity; if (identity == null) return; @@ -1142,7 +1142,7 @@ export class BridgeWorker { if (stopSignal.aborted || Date.now() >= assignmentDeadlineMs) return; try { await this.refreshCredential( - requestSignal, + stopSignal, Date.now() + serverClockOffsetMs + refreshWindowMs, ); } catch (error) { @@ -1373,7 +1373,6 @@ export class BridgeWorker { assignment, credentialController.signal, serverClockOffsetMs, - signal, ).catch((error) => { credentialMaintenanceError = error; executionController.abort(); @@ -1725,6 +1724,17 @@ export class BridgeWorker { clearTimeout(deadlineTimer); cancellationController.abort(); await cancellationWatcher; + const credentialInFlight = this.credentialInFlight; + if (credentialInFlight != null && !credentialController.signal.aborted) { + let drainTimer: ReturnType | undefined; + await Promise.race([ + credentialInFlight.catch(() => undefined), + new Promise((resolve) => { + drainTimer = setTimeout(resolve, CREDENTIAL_REFRESH_SETTLEMENT_GRACE_MS); + }), + ]); + if (drainTimer != null) clearTimeout(drainTimer); + } credentialController.abort(); await credentialMaintenance; try { @@ -2219,12 +2229,14 @@ 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. - signal.removeEventListener('abort', abortPoll); + if (response.ok || response.status === 404) { + signal.removeEventListener('abort', abortPoll); + } }, ); if (response.cancelled) { @@ -2249,7 +2261,7 @@ export class BridgeWorker { url: string, body: object, signal?: AbortSignal, - onResponseHeaders?: () => void, + onResponseHeaders?: (response: Response) => void, ): Promise { const requestBody = JSON.stringify(body); const response = await this.fetchImpl(url, { @@ -2261,7 +2273,7 @@ export class BridgeWorker { body: requestBody, signal, }); - onResponseHeaders?.(); + onResponseHeaders?.(response); let payload: unknown; try { payload = await response.json(); diff --git a/service/src/job-cancellation.test.ts b/service/src/job-cancellation.test.ts index 0582e9c3..2c7329a5 100644 --- a/service/src/job-cancellation.test.ts +++ b/service/src/job-cancellation.test.ts @@ -8,6 +8,7 @@ import { jobCancellationInternals, removeJobIfWaiting, requestJobCancellation, + throwIfJobAborted, waitForJobWithCancellation, } from './job-cancellation'; @@ -59,6 +60,7 @@ class FakeRedis { readonly existing = new Set(); readonly deleted: string[] = []; readonly transactions: FakeTransaction[] = []; + mgetFailures = 0; duplicate(): FakeSubscriber { this.duplicateCalls += 1; @@ -70,6 +72,10 @@ class FakeRedis { } 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)); } @@ -155,6 +161,48 @@ test('one pubsub listener cancels only the matching active job', async () => { 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)); @@ -164,7 +212,24 @@ test('subscriber reconnect reconciles active jobs against durable markers', asyn fake.existing.add(jobCancellationInternals.cancellationKey(target)); fake.subscriber.emit('ready'); - await new Promise((resolve) => setImmediate(resolve)); + 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); @@ -184,6 +249,15 @@ test('cancellation writes a durable marker before publishing', async () => { ]); }); +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('disconnect frees a waiting job and rejects promptly', async () => { const fake = new FakeRedis(); const controller = new AbortController(); @@ -199,8 +273,10 @@ test('disconnect frees a waiting job and rejects promptly', async () => { }, } as unknown as Job; + const registry = new JobCancellationRegistry(redis(fake)); const waiting = waitForJobWithCancellation({ commands: redis(fake), + registry, job, events: {} as QueueEvents, timeoutMs: 60_000, @@ -218,6 +294,39 @@ test('disconnect frees a waiting job and rejects promptly', async () => { 'EX', 120, ]); + 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 () => { diff --git a/service/src/job-cancellation.ts b/service/src/job-cancellation.ts index c6c2778c..bf3eae72 100644 --- a/service/src/job-cancellation.ts +++ b/service/src/job-cancellation.ts @@ -50,9 +50,11 @@ export class JobCancellationRegistry { private subscriber?: IORedis; private readonly controllers = new Map< string, - { target: JobTarget; controller: AbortController } + { target: JobTarget; controllers: Set } >(); private startPromise?: Promise; + private reconcileTimer?: ReturnType; + private reconcileRetryMs = 100; private closed = false; constructor(private readonly commands: IORedis) {} @@ -63,16 +65,17 @@ export class JobCancellationRegistry { }; private readonly onSubscriberReady = (): void => { - void this.reconcile().catch(() => undefined); + 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; - this.controllers - .get(targetKey(target)) - ?.controller.abort(CLIENT_DISCONNECT_REASON); + for (const controller of + this.controllers.get(targetKey(target))?.controllers ?? []) { + controller.abort(CLIENT_DISCONNECT_REASON); + } }; private detachSubscriber(subscriber: IORedis): void { @@ -89,11 +92,36 @@ export class JobCancellationRegistry { ); cancelled.forEach((value, index) => { if (value != null) { - entries[index]?.controller.abort(CLIENT_DISCONNECT_REASON); + 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')); @@ -122,26 +150,46 @@ export class JobCancellationRegistry { } async register(target: JobTarget, controller: AbortController): Promise { - this.controllers.set(targetKey(target), { target, controller }); + 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.exists(cancellationKey(target))) { controller.abort(CLIENT_DISCONNECT_REASON); } } catch (error) { - this.controllers.delete(targetKey(target)); + entry.controllers.delete(controller); + if (entry.controllers.size === 0) this.controllers.delete(key); throw error; } } - async unregister(target: JobTarget): Promise { - this.controllers.delete(targetKey(target)); - await this.commands.del(cancellationKey(target)); + 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; await this.startPromise?.catch(() => undefined); const subscriber = this.subscriber; this.subscriber = undefined; @@ -195,21 +243,46 @@ 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', + ); +} + export async function waitForJobWithCancellation(args: { commands: IORedis; + registry: JobCancellationRegistry; job: Job; events: QueueEvents; timeoutMs: number; cancellationTtlSeconds: number; signal?: AbortSignal; }): Promise { - const { commands, job, events, timeoutMs, cancellationTtlSeconds, signal } = args; + const { + commands, + registry, + job, + events, + timeoutMs, + cancellationTtlSeconds, + signal, + } = args; const completion = job.waitUntilFinished(events, timeoutMs); - if (signal == null) return completion; - const target = { queueName: job.queueName, jobId: String(job.id) }; + const externalController = new AbortController(); + try { + await registry.register(target, externalController); + } catch (error) { + void completion.catch(() => undefined); + throw error; + } + let removeAbortListener = (): void => {}; - const cancelled = new Promise((_, reject) => { + const disconnected = new Promise((_, reject) => { let cancelling = false; const cancel = (): void => { if (cancelling) return; @@ -223,18 +296,26 @@ export async function waitForJobWithCancellation(args: { }) .then(() => reject(programmaticCancellationError()), reject); }; - removeAbortListener = (): void => signal.removeEventListener('abort', cancel); - signal.addEventListener('abort', cancel, { once: true }); - if (signal.aborted) cancel(); + 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 => 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, cancelled]); + return await Promise.race([completion, disconnected, cancelled]); } finally { removeAbortListener(); + await registry.unregister(target, externalController).catch(() => undefined); } } diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index 8170ff41..9e6eb698 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -9,6 +9,7 @@ import { pyQueue, pyQueueEvents, connection, + jobCancellationRegistry, getExecutionQueueBinding, getExistingExecutionJob, } from '../queue'; @@ -403,6 +404,25 @@ async function runReplayIteration( state.executionProfileSource ?? env.EXECUTION_PROFILE_SOURCE, ); if (signal?.aborted) throw programmaticCancellationError(); + let cancellationTarget: + | { queueName: string; jobId: string } + | undefined; + if (cancellation != null) { + cancellationTarget = { queueName: queue.name, jobId: nanoid() }; + 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 job = await queue.add( Jobs.execute, { @@ -434,48 +454,14 @@ async function runReplayIteration( removeOnComplete: { age: 60, count: 1 }, removeOnFail: { age: 180, count: 1 }, attempts: 1, + ...(cancellationTarget != null ? { jobId: cancellationTarget.jobId } : {}), }, ); jobsSubmitted.inc({ language }); - if (cancellation != null) { - const target = { queueName: job.queueName, jobId: String(job.id) }; - let cancellationPublished = false; - try { - const attachment = await attachProgrammaticCancellationTarget({ - redis: connection, - requestId: cancellation.requestId, - owner: cancellation.owner, - target, - ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, - }); - if (attachment === 'forbidden') { - throw new Error('Programmatic cancellation request ownership changed'); - } - if (attachment === 'cancelled') { - await requestJobCancellation( - connection, - target, - PROGRAMMATIC_CANCELLATION_TTL_SECONDS, - ); - cancellationPublished = true; - await removeJobIfWaiting(job).catch(() => false); - throw programmaticCancellationError(); - } - } catch (error) { - if (!cancellationPublished) { - await requestJobCancellation( - connection, - target, - PROGRAMMATIC_CANCELLATION_TTL_SECONDS, - ).catch(() => undefined); - } - throw error; - } - } - return waitForJobWithCancellation({ commands: connection, + registry: jobCancellationRegistry, job, events, timeoutMs: JOB_COMPLETION_WAIT_TIMEOUT_MS, diff --git a/service/src/workers.ts b/service/src/workers.ts index dd59a9c8..b1e35dd4 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -19,6 +19,7 @@ import { workerDeadlineFailure } from './worker-error'; import { CLIENT_DISCONNECT_REASON, JOB_CANCELLED_MESSAGE, + throwIfJobAborted, } from './job-cancellation'; import logger from './logger'; import { @@ -181,6 +182,9 @@ async function processJobInner(job: t.ExecuteJob): Promise { ); 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. + throwIfJobAborted(controller.signal); if (!isSyntheticJob) { logger.info('Sandbox response', summarizeSandboxResponse(responseData)); @@ -289,7 +293,7 @@ async function processJobInner(job: t.ExecuteJob): Promise { } if (timer) clearTimeout(timer); if (cancellationTarget != null && cancellationRegistered) { - await jobCancellationRegistry.unregister(cancellationTarget).catch(error => { + await jobCancellationRegistry.unregister(cancellationTarget, controller).catch(error => { logger.warn('Failed to clear queued execution cancellation state', { queueName: cancellationTarget.queueName, jobId: cancellationTarget.jobId, From eaf6646671eec0eca1f3d4a3d7b4767a711347c0 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Sep 2026 08:55:43 -0400 Subject: [PATCH 06/11] fix: harden cancellation under concurrent load --- packages/code/src/worker.test.ts | 450 ++++++++++++++++++++------- packages/code/src/worker.ts | 125 ++++++-- service/src/job-cancellation.test.ts | 144 ++++++++- service/src/job-cancellation.ts | 118 ++++++- service/src/workers.ts | 152 ++++++--- 5 files changed, 780 insertions(+), 209 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index ea785dde..41c4ab4a 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', @@ -1021,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', { @@ -1201,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')) { @@ -1215,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', @@ -1266,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', @@ -1287,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')) { @@ -1308,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' }, + }, ); }, }); @@ -1339,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', @@ -1356,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' }, + }, ); }, }); @@ -1386,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', @@ -1404,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' }, + }, ); }, }); @@ -1437,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', @@ -1500,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')) { @@ -1520,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', @@ -1679,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; @@ -1696,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' }, + }, ); }, }); @@ -1741,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, @@ -1764,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( @@ -1780,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); @@ -1822,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')) { @@ -1834,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( @@ -1854,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' }, + }, ); }, }); @@ -1871,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', @@ -1885,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' }, + }, ); }, }); @@ -1916,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', @@ -1930,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' }, + }, ); }, }); @@ -2028,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' }, + }, ); }, }); @@ -2055,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; @@ -2090,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')) { @@ -2109,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 }, ); } @@ -2341,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 }); }; @@ -2445,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; @@ -2516,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, @@ -2558,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', @@ -2577,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, @@ -2640,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, @@ -2927,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({ @@ -2969,7 +3198,16 @@ test('sandbox completion does not cancel an in-flight credential rotation', asyn }); 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 45c096d2..ac38cadd 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -164,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) => @@ -195,7 +198,8 @@ function workspaceCapabilitiesMatch( (operation, operationIndex) => operation === executor.workspaces[index]?.operations?.[operationIndex], - ) ?? executor.workspaces[index]?.operations == null), + ) ?? + executor.workspaces[index]?.operations == null), ) ); } @@ -207,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, @@ -217,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; @@ -227,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 []; } @@ -392,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) { @@ -440,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', @@ -559,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), @@ -581,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, @@ -715,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, @@ -751,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, @@ -1057,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(); + } } } @@ -1151,8 +1207,7 @@ export class BridgeWorker { 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( @@ -1686,7 +1741,8 @@ export class BridgeWorker { const knownAtomicProgrammaticFailure = assignment.executionKind === 'workspace_programmatic' && error instanceof WorkspaceToolError && - this.options.workspaceProgrammatic?.mutationFailuresAreAtomic === true && + this.options.workspaceProgrammatic?.mutationFailuresAreAtomic === + true && !error.requiresQuarantine; if ( workspaceMutationApplied || @@ -1724,13 +1780,16 @@ export class BridgeWorker { clearTimeout(deadlineTimer); cancellationController.abort(); await cancellationWatcher; - const credentialInFlight = this.credentialInFlight; + const credentialInFlight = this.credentialInFlight?.promise; if (credentialInFlight != null && !credentialController.signal.aborted) { let drainTimer: ReturnType | undefined; await Promise.race([ credentialInFlight.catch(() => undefined), new Promise((resolve) => { - drainTimer = setTimeout(resolve, CREDENTIAL_REFRESH_SETTLEMENT_GRACE_MS); + drainTimer = setTimeout( + resolve, + CREDENTIAL_REFRESH_SETTLEMENT_GRACE_MS, + ); }), ]); if (drainTimer != null) clearTimeout(drainTimer); @@ -1936,7 +1995,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`, { diff --git a/service/src/job-cancellation.test.ts b/service/src/job-cancellation.test.ts index 2c7329a5..31c556f2 100644 --- a/service/src/job-cancellation.test.ts +++ b/service/src/job-cancellation.test.ts @@ -5,6 +5,7 @@ import type { Job, QueueEvents } from 'bullmq'; import { CLIENT_DISCONNECT_REASON, JobCancellationRegistry, + jobResultCommitFailure, jobCancellationInternals, removeJobIfWaiting, requestJobCancellation, @@ -236,6 +237,23 @@ test('subscriber reconnect retries durable-marker reconciliation', async () => { 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' }; @@ -258,6 +276,23 @@ test('result commit barrier rejects cancellation observed after execution', () = ); }); +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(); @@ -289,7 +324,50 @@ test('disconnect frees a waiting job and rejects promptly', async () => { expect(removed).toBe(true); expect(fake.transactions[0]?.operations[0]).toEqual([ 'set', - jobCancellationInternals.cancellationKey({ queueName: 'other', jobId: 'job-4' }), + 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, @@ -297,6 +375,40 @@ test('disconnect frees a waiting job and rejects promptly', async () => { 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)); @@ -344,18 +456,22 @@ test('queued removal never removes an active job', async () => { 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( + 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 index bf3eae72..c2e695b8 100644 --- a/service/src/job-cancellation.ts +++ b/service/src/job-cancellation.ts @@ -16,7 +16,9 @@ function targetKey(target: JobTarget): string { } function cancellationKey(target: JobTarget): string { - return `${JOB_CANCELLATION_PREFIX}:${encodeURIComponent(target.queueName)}:${encodeURIComponent(target.jobId)}`; + return `${JOB_CANCELLATION_PREFIX}:${encodeURIComponent( + target.queueName, + )}:${encodeURIComponent(target.jobId)}`; } function parseTarget(raw: string): JobTarget | undefined { @@ -53,27 +55,32 @@ export class JobCancellationRegistry { { 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. + // 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 => { + 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 ?? []) { + for (const controller of this.controllers.get(targetKey(target)) + ?.controllers ?? []) { controller.abort(CLIENT_DISCONNECT_REASON); } }; @@ -82,6 +89,49 @@ export class JobCancellationRegistry { 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 { @@ -133,6 +183,9 @@ export class JobCancellationRegistry { 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) { @@ -149,7 +202,10 @@ export class JobCancellationRegistry { return starting; } - async register(target: JobTarget, controller: AbortController): Promise { + async register( + target: JobTarget, + controller: AbortController, + ): Promise { const key = targetKey(target); const entry = this.controllers.get(key) ?? { target, @@ -190,6 +246,9 @@ export class JobCancellationRegistry { this.controllers.clear(); if (this.reconcileTimer != null) clearTimeout(this.reconcileTimer); this.reconcileTimer = undefined; + if (this.subscriberRestartTimer != null) + clearTimeout(this.subscriberRestartTimer); + this.subscriberRestartTimer = undefined; await this.startPromise?.catch(() => undefined); const subscriber = this.subscriber; this.subscriber = undefined; @@ -211,7 +270,9 @@ export async function requestJobCancellation( transaction.publish(JOB_CANCELLATION_CHANNEL, payload); const result = await transaction.exec(); if (result == null) { - throw new Error('Redis transaction aborted while cancelling queued execution'); + throw new Error( + 'Redis transaction aborted while cancelling queued execution', + ); } const failure = result.find(([error]) => error != null)?.[0]; if (failure != null) throw failure; @@ -240,7 +301,10 @@ export async function removeJobIfWaiting( } export function programmaticCancellationError(): Error { - return new DOMException('Programmatic execution request disconnected', 'AbortError'); + return new DOMException( + 'Programmatic execution request disconnected', + 'AbortError', + ); } /** Commit barrier for result-processing stages that may yield after execution. */ @@ -253,6 +317,20 @@ export function throwIfJobAborted(signal: AbortSignal): void { ); } +/** 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; @@ -278,6 +356,10 @@ export async function waitForJobWithCancellation(args: { await registry.register(target, externalController); } catch (error) { void completion.catch(() => undefined); + await Promise.allSettled([ + requestJobCancellation(commands, target, cancellationTtlSeconds), + removeJobIfWaiting(job), + ]); throw error; } @@ -297,14 +379,22 @@ export async function waitForJobWithCancellation(args: { .then(() => reject(programmaticCancellationError()), reject); }; if (signal != null) { - removeAbortListener = (): void => signal.removeEventListener('abort', cancel); + removeAbortListener = (): void => + signal.removeEventListener('abort', cancel); signal.addEventListener('abort', cancel, { once: true }); if (signal.aborted) cancel(); } }); const cancelled = new Promise((_, reject) => { - const cancel = (): void => reject(programmaticCancellationError()); - externalController.signal.addEventListener('abort', cancel, { once: true }); + const cancel = (): void => { + void removeJobIfWaiting(job).then( + () => reject(programmaticCancellationError()), + () => reject(programmaticCancellationError()), + ); + }; + externalController.signal.addEventListener('abort', cancel, { + once: true, + }); if (externalController.signal.aborted) cancel(); }); @@ -315,7 +405,9 @@ export async function waitForJobWithCancellation(args: { return await Promise.race([completion, disconnected, cancelled]); } finally { removeAbortListener(); - await registry.unregister(target, externalController).catch(() => undefined); + await registry + .unregister(target, externalController) + .catch(() => undefined); } } diff --git a/service/src/workers.ts b/service/src/workers.ts index b1e35dd4..9a97b9c1 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -1,24 +1,45 @@ import axios from 'axios'; import { Worker } from 'bullmq'; import type * as t from './types'; -import { filterSystemLogs, applySystemReplacements, getAxiosErrorDetails, sandboxErrorMessageFromAxios } from './utils'; -import { jobProcessingDuration, jobsCancelled, jobsCompleted, jobsFailed, activeJobs, workerRunning } from './metrics'; +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, throwIfJobAborted, } from './job-cancellation'; import logger from './logger'; @@ -32,41 +53,57 @@ 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 cancellationTarget = job.data.cancellable === true && job.id != null - ? { queueName: job.queueName, jobId: String(job.id) } - : undefined; + 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); const remainingBudgetMs = Math.max(0, deadlineAtMs - Date.now()); - const timer = remainingBudgetMs > 0 - ? setTimeout(() => controller.abort('deadline'), remainingBudgetMs) - : undefined; + 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; try { if (cancellationTarget != null) { @@ -76,7 +113,10 @@ async function processJobInner(job: t.ExecuteJob): Promise { 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, @@ -90,7 +130,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, }); @@ -98,9 +141,10 @@ 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); @@ -110,7 +154,8 @@ async function processJobInner(job: t.ExecuteJob): Promise { 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, ), @@ -134,7 +179,9 @@ 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 => { + const finalizeSandboxResult = async ( + result: SandboxRawResponse, + ): Promise => { if ( resultRestoreToken === undefined || resultRestoreToken.length === 0 || @@ -175,9 +222,10 @@ 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: + resultRestoreToken !== undefined && resultRestoreToken.length > 0 + ? finalizeSandboxResult + : undefined, }, ); @@ -212,7 +260,9 @@ async function processJobInner(job: t.ExecuteJob): Promise { stdout, stderr, ...(responseData.pending_tool_calls_payload != null - ? { pending_tool_calls_payload: responseData.pending_tool_calls_payload } + ? { + pending_tool_calls_payload: responseData.pending_tool_calls_payload, + } : {}), }; @@ -221,7 +271,8 @@ async function processJobInner(job: t.ExecuteJob): Promise { 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; + result.wall_time = + ((run as Record).wall_time as number | null) ?? null; } if (result.message || result.signal) { @@ -235,10 +286,12 @@ async function processJobInner(job: t.ExecuteJob): Promise { }); } + completedResult = true; return result; } catch (error) { const clientDisconnected = - controller.signal.aborted && controller.signal.reason === CLIENT_DISCONNECT_REASON; + controller.signal.aborted && + controller.signal.reason === CLIENT_DISCONNECT_REASON; revokeReason = clientDisconnected ? 'cancelled' : controller.signal.aborted || isAbortError(error) @@ -283,26 +336,37 @@ 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) }); + }).catch((error) => { + logger.error('Failed to revoke egress grant', { + grantId: egressGrantId, + error: getAxiosErrorDetails(error), + }); }); } + const lateCommitFailure = completedResult + ? jobResultCommitFailure(controller.signal, env.JOB_TIMEOUT) + : undefined; 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), + 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; } } @@ -330,14 +394,14 @@ export const otherWorker = new Worker(queueNames.other, processJob, { workerRunning.set({ worker_type: 'python' }, 1); workerRunning.set({ worker_type: 'other' }, 1); -pyWorker.on('completed', job => { +pyWorker.on('completed', (job) => { if (job.data.isSynthetic !== true) { logger.info(`[${WORKER_ID}] Python job completed ${job.id}`); } jobsCompleted.inc({ language: 'python' }); }); -otherWorker.on('completed', job => { +otherWorker.on('completed', (job) => { if (job.data.isSynthetic !== true) { logger.info(`[${WORKER_ID}] Other job completed ${job.id}`); } From 34611c2f9fd66ed0fa0b12c161029b778ce334dc Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Sep 2026 10:02:47 -0400 Subject: [PATCH 07/11] fix: make cancellation ownership durable through completion --- docs/remote-bridge/README.md | 9 ++ packages/code/src/worker.test.ts | 24 ++++ packages/code/src/worker.ts | 15 ++- service/src/job-cancellation-commit.test.ts | 56 ++++++++ service/src/job-cancellation.test.ts | 88 ++++++++++-- service/src/job-cancellation.ts | 142 +++++++++++++++++--- service/src/queue.ts | 11 +- service/src/redis-options.test.ts | 12 +- service/src/redis-options.ts | 6 + service/src/service/programmatic-router.ts | 27 +++- service/src/workers.ts | 18 ++- 11 files changed, 360 insertions(+), 48 deletions(-) create mode 100644 service/src/job-cancellation-commit.test.ts diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index 42e3abaa..b9de3ac2 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -264,6 +264,15 @@ execution. 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 41c4ab4a..e26f212a 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -3197,6 +3197,30 @@ 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), diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index ac38cadd..ffe1b350 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -1180,6 +1180,7 @@ export class BridgeWorker { assignment: BridgeAssignment, stopSignal: AbortSignal, serverClockOffsetMs: number, + maintenance: { refresh?: Promise }, ): Promise { const identity = this.options.identity; if (identity == null) return; @@ -1197,10 +1198,11 @@ export class BridgeWorker { await abortableDelay(waitMs, stopSignal); if (stopSignal.aborted || Date.now() >= assignmentDeadlineMs) return; try { - await this.refreshCredential( + maintenance.refresh = this.refreshCredential( stopSignal, Date.now() + serverClockOffsetMs + refreshWindowMs, ); + await maintenance.refresh; } catch (error) { if (stopSignal.aborted) return; const terminal = @@ -1216,6 +1218,8 @@ export class BridgeWorker { ), stopSignal, ); + } finally { + maintenance.refresh = undefined; } } } @@ -1412,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; @@ -1428,6 +1433,7 @@ export class BridgeWorker { assignment, credentialController.signal, serverClockOffsetMs, + ownCredentialMaintenance, ).catch((error) => { credentialMaintenanceError = error; executionController.abort(); @@ -1780,7 +1786,9 @@ export class BridgeWorker { clearTimeout(deadlineTimer); cancellationController.abort(); await cancellationWatcher; - const credentialInFlight = this.credentialInFlight?.promise; + // 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([ @@ -1788,7 +1796,8 @@ export class BridgeWorker { new Promise((resolve) => { drainTimer = setTimeout( resolve, - CREDENTIAL_REFRESH_SETTLEMENT_GRACE_MS, + Math.min(CREDENTIAL_REFRESH_SETTLEMENT_GRACE_MS, + Math.max(0, Date.parse(assignment.expiresAt) - serverClockOffsetMs - Date.now() - 5_000)), ); }), ]); diff --git a/service/src/job-cancellation-commit.test.ts b/service/src/job-cancellation-commit.test.ts new file mode 100644 index 00000000..9a801686 --- /dev/null +++ b/service/src/job-cancellation-commit.test.ts @@ -0,0 +1,56 @@ +import { afterEach, beforeEach, expect, test } from 'bun:test'; +import { startTestRedis } from './test/redis'; +import { + commitJobResult, + readCommittedJobResult, + requestJobCancellation, + JobCancellationRegistry, + jobCancellationInternals, +} from './job-cancellation'; + +let redis: Awaited>; +beforeEach(async () => { + redis = await startTestRedis(); +}); +afterEach(async () => { + await redis.closeTestServer(); +}); +const target = { queueName: 'other', jobId: 'commit-race' }; + +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( + false, + ); + 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(true); + 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)).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'); +}); diff --git a/service/src/job-cancellation.test.ts b/service/src/job-cancellation.test.ts index 31c556f2..b7f1524c 100644 --- a/service/src/job-cancellation.test.ts +++ b/service/src/job-cancellation.test.ts @@ -11,6 +11,7 @@ import { requestJobCancellation, throwIfJobAborted, waitForJobWithCancellation, + fenceJobCancellation, } from './job-cancellation'; class FakeSubscriber extends EventEmitter { @@ -62,14 +63,33 @@ class FakeRedis { readonly deleted: string[] = []; readonly transactions: FakeTransaction[] = []; mgetFailures = 0; + cancellationFailures = 0; + cancellationAttempts = 0; duplicate(): FakeSubscriber { this.duplicateCalls += 1; return this.subscriber; } - async exists(key: string): Promise { - return this.existing.has(key) ? 1 : 0; + async get(key: string): Promise { + return this.existing.has(key) ? '1' : null; + } + + async eval( + _script: string, + _keys: number, + key: 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> { @@ -77,7 +97,7 @@ class FakeRedis { this.mgetFailures -= 1; throw new Error('command connection unavailable'); } - return keys.map((key) => (this.existing.has(key) ? '1' : null)); + return keys.map(key => (this.existing.has(key) ? '1' : null)); } async del(key: string): Promise { @@ -213,7 +233,7 @@ test('subscriber reconnect reconciles active jobs against durable markers', asyn fake.existing.add(jobCancellationInternals.cancellationKey(target)); fake.subscriber.emit('ready'); - await new Promise((resolve) => setTimeout(resolve, 10)); + await new Promise(resolve => setTimeout(resolve, 10)); expect(controller.signal.aborted).toBe(true); expect(controller.signal.reason).toBe(CLIENT_DISCONNECT_REASON); @@ -230,7 +250,7 @@ test('subscriber reconnect retries durable-marker reconciliation', async () => { fake.mgetFailures = 1; fake.subscriber.emit('ready'); - await new Promise((resolve) => setTimeout(resolve, 150)); + await new Promise(resolve => setTimeout(resolve, 150)); expect(controller.signal.aborted).toBe(true); expect(controller.signal.reason).toBe(CLIENT_DISCONNECT_REASON); @@ -246,7 +266,7 @@ test('terminal subscriber disconnect rebuilds the subscription and reconciles ma fake.existing.add(jobCancellationInternals.cancellationKey(target)); fake.subscriber.emit('end'); - await new Promise((resolve) => setTimeout(resolve, 10)); + await new Promise(resolve => setTimeout(resolve, 10)); expect(fake.duplicateCalls).toBe(2); expect(controller.signal.aborted).toBe(true); @@ -396,7 +416,7 @@ test('external cancellation frees a waiting job before rejecting its waiter', as timeoutMs: 60_000, cancellationTtlSeconds: 120, }); - await new Promise((resolve) => setImmediate(resolve)); + await new Promise(resolve => setImmediate(resolve)); fake.subscriber.emit( 'message', @@ -429,7 +449,7 @@ test('a separate cancellation request wakes the original job waiter', async () = timeoutMs: 60_000, cancellationTtlSeconds: 120, }); - await new Promise((resolve) => setImmediate(resolve)); + await new Promise(resolve => setImmediate(resolve)); fake.subscriber.emit( 'message', jobCancellationInternals.channel, @@ -454,6 +474,58 @@ test('queued removal never removes an active job', async () => { 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( diff --git a/service/src/job-cancellation.ts b/service/src/job-cancellation.ts index c2e695b8..8d88fd5d 100644 --- a/service/src/job-cancellation.ts +++ b/service/src/job-cancellation.ts @@ -141,7 +141,7 @@ export class JobCancellationRegistry { ...entries.map(({ target }) => cancellationKey(target)), ); cancelled.forEach((value, index) => { - if (value != null) { + if (value === '1') { for (const controller of entries[index]?.controllers ?? []) { controller.abort(CLIENT_DISCONNECT_REASON); } @@ -215,7 +215,7 @@ export class JobCancellationRegistry { this.controllers.set(key, entry); try { await this.start(); - if (await this.commands.exists(cancellationKey(target))) { + if ((await this.commands.get(cancellationKey(target))) === '1') { controller.abort(CLIENT_DISCONNECT_REASON); } } catch (error) { @@ -263,19 +263,110 @@ export async function requestJobCancellation( commands: IORedis, target: JobTarget, ttlSeconds: number, -): Promise { - const payload = JSON.stringify(target); - const transaction = commands.multi(); - transaction.set(cancellationKey(target), '1', 'EX', Math.max(1, ttlSeconds)); - transaction.publish(JOB_CANCELLATION_CHANNEL, payload); - const result = await transaction.exec(); - if (result == null) { +): Promise { + // Cancellation and result publication have ONE durable winner. Pub/sub is + // only a notification; it must not decide whether Stop was accepted. + return ( + (await commands.eval( + ` + local state = redis.call('GET', KEYS[1]) + if state and state ~= '1' then return 0 end + redis.call('SET', KEYS[1], '1', 'EX', ARGV[1]) + redis.call('PUBLISH', ARGV[2], ARGV[3]) + return 1 + `, + 1, + cancellationKey(target), + Math.max(1, ttlSeconds), + JOB_CANCELLATION_CHANNEL, + JSON.stringify(target), + )) === 1 + ); +} + +/** 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, +): Promise { + const serialized = JSON.stringify({ result }); + if (Buffer.byteLength(serialized) > 16 * 1024 * 1024) { + throw new Error('Programmatic completion exceeds the 16 MiB result limit'); + } + return ( + (await commands.eval( + ` + local state = redis.call('GET', KEYS[1]) + if state == '1' then return 0 end + if not state then + redis.call('SET', KEYS[2], ARGV[1], 'EX', ARGV[2]) + redis.call('SET', KEYS[1], 'completed', 'EX', ARGV[2]) + end + return 1 + `, + 2, + cancellationKey(target), + `${cancellationKey(target)}:result`, + serialized, + Math.max(1, ttlSeconds), + )) === 1 + ); +} + +export async function readCommittedJobResult( + commands: IORedis, + target: JobTarget, +): Promise<{ result: T } | undefined> { + const state = await commands.get(cancellationKey(target)); + if (state !== 'completed') return undefined; + const value = await commands.get(`${cancellationKey(target)}:result`); + if (value == null) throw new Error( - 'Redis transaction aborted while cancelling queued execution', + 'Committed programmatic result expired; refusing re-execution', ); + return JSON.parse(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; + while (Date.now() < args.deadlineAtMs) { + let timer: ReturnType | undefined; + try { + return await Promise.race([ + requestJobCancellation(args.commands, args.target, args.ttlSeconds), + new Promise(resolve => { + timer = setTimeout( + () => resolve(true), + Math.max(1, args.deadlineAtMs - Date.now()), + ); + }), + ]); + } catch { + await new Promise(resolve => + setTimeout( + resolve, + Math.min(retryMs, Math.max(0, args.deadlineAtMs - Date.now())), + ), + ); + retryMs = Math.min(1_000, retryMs * 2); + } finally { + if (timer != null) clearTimeout(timer); + } } - const failure = result.find(([error]) => error != null)?.[0]; - if (failure != null) throw failure; + return true; } const REMOVABLE_JOB_STATES = new Set([ @@ -338,6 +429,7 @@ export async function waitForJobWithCancellation(args: { events: QueueEvents; timeoutMs: number; cancellationTtlSeconds: number; + deadlineAtMs?: number; signal?: AbortSignal; }): Promise { const { @@ -351,15 +443,20 @@ export async function waitForJobWithCancellation(args: { } = args; const completion = job.waitUntilFinished(events, timeoutMs); const target = { queueName: job.queueName, jobId: String(job.id) }; + const fence = (): Promise => + fenceJobCancellation({ + commands, + target, + ttlSeconds: cancellationTtlSeconds, + deadlineAtMs: args.deadlineAtMs ?? Date.now() + timeoutMs, + }); const externalController = new AbortController(); try { await registry.register(target, externalController); } catch (error) { void completion.catch(() => undefined); - await Promise.allSettled([ - requestJobCancellation(commands, target, cancellationTtlSeconds), - removeJobIfWaiting(job), - ]); + await fence(); + await removeJobIfWaiting(job).catch(() => false); throw error; } @@ -369,14 +466,16 @@ export async function waitForJobWithCancellation(args: { const cancel = (): void => { if (cancelling) return; cancelling = true; - void requestJobCancellation(commands, target, cancellationTtlSeconds) - .then(async () => { + void fence() + .then(async accepted => { + if (!accepted) 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()); }) - .then(() => reject(programmaticCancellationError()), reject); + .catch(reject); }; if (signal != null) { removeAbortListener = (): void => @@ -403,6 +502,11 @@ export async function waitForJobWithCancellation(args: { 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. + await fence(); + throw error; } finally { removeAbortListener(); await registry diff --git a/service/src/queue.ts b/service/src/queue.ts index 4bf1c8e1..72c2fbb1 100644 --- a/service/src/queue.ts +++ b/service/src/queue.ts @@ -18,20 +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'; import { JobCancellationRegistry } from './job-cancellation'; -const MAX_RECONNECT_ATTEMPTS = 5; -const RECONNECT_DELAY = 2000; - 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) => { 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/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index 9e6eb698..a4ead752 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -18,6 +18,7 @@ import { programmaticCancellationError, removeJobIfWaiting, requestJobCancellation, + fenceJobCancellation, waitForJobWithCancellation, } from '../job-cancellation'; import { @@ -404,11 +405,8 @@ async function runReplayIteration( state.executionProfileSource ?? env.EXECUTION_PROFILE_SOURCE, ); if (signal?.aborted) throw programmaticCancellationError(); - let cancellationTarget: - | { queueName: string; jobId: string } - | undefined; + const cancellationTarget = { queueName: queue.name, jobId: nanoid() }; if (cancellation != null) { - cancellationTarget = { queueName: queue.name, jobId: nanoid() }; const attachment = await attachProgrammaticCancellationTarget({ redis: connection, requestId: cancellation.requestId, @@ -423,6 +421,8 @@ async function runReplayIteration( throw programmaticCancellationError(); } } + const submittedAtMs = Date.now(); + const deadlineAtMs = submittedAtMs + env.JOB_TIMEOUT; const job = await queue.add( Jobs.execute, { @@ -454,9 +454,17 @@ async function runReplayIteration( removeOnComplete: { age: 60, count: 1 }, removeOnFail: { age: 180, count: 1 }, attempts: 1, - ...(cancellationTarget != null ? { jobId: cancellationTarget.jobId } : {}), + jobId: cancellationTarget.jobId, + timestamp: submittedAtMs, }, - ); + ).catch(async (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. + await fenceJobCancellation({ commands: connection, target: cancellationTarget, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, deadlineAtMs }); + throw error; + }); jobsSubmitted.inc({ language }); return waitForJobWithCancellation({ @@ -466,6 +474,7 @@ async function runReplayIteration( events, timeoutMs: JOB_COMPLETION_WAIT_TIMEOUT_MS, cancellationTtlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, + deadlineAtMs, signal, }); } @@ -1202,11 +1211,15 @@ router.post( return; } if (cancellation.target != null) { - await requestJobCancellation( + 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, diff --git a/service/src/workers.ts b/service/src/workers.ts index 9a97b9c1..9c695979 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -40,6 +40,8 @@ import { CLIENT_DISCONNECT_REASON, JOB_CANCELLED_MESSAGE, jobResultCommitFailure, + commitJobResult, + readCommittedJobResult, throwIfJobAborted, } from './job-cancellation'; import logger from './logger'; @@ -104,11 +106,14 @@ async function processJobInner(job: t.ExecuteJob): Promise { let egressGrantTokenForRestore: string | undefined; let revokeReason = 'completed'; let completedResult = false; + let resultToCommit: t.ExecuteResult | undefined; try { if (cancellationTarget != null) { await jobCancellationRegistry.register(cancellationTarget, controller); cancellationRegistered = true; + const committed = await readCommittedJobResult(connection, cancellationTarget); + if (committed != null) return committed.result; } if (controller.signal.aborted) { throw new Error(`Job timed out after ${env.JOB_TIMEOUT}ms`); @@ -287,6 +292,7 @@ async function processJobInner(job: t.ExecuteJob): Promise { } completedResult = true; + resultToCommit = result; return result; } catch (error) { const clientDisconnected = @@ -349,9 +355,19 @@ async function processJobInner(job: t.ExecuteJob): Promise { }); }); } - const lateCommitFailure = completedResult + let lateCommitFailure = completedResult ? jobResultCommitFailure(controller.signal, env.JOB_TIMEOUT) : undefined; + if (completedResult && cancellationTarget != null && lateCommitFailure == null) { + try { + if (!(await commitJobResult(connection, cancellationTarget, resultToCommit, + Math.ceil(env.JOB_TIMEOUT / 1_000) * 2 + 180))) { + lateCommitFailure = new Error(JOB_CANCELLED_MESSAGE); + } + } catch (error) { + lateCommitFailure = error instanceof Error ? error : new Error('Result commit failed'); + } + } if (timer) clearTimeout(timer); if (cancellationTarget != null && cancellationRegistered) { await jobCancellationRegistry From 0d34a9e2666dae72a8dac35498bf8f3b27bcf763 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Sep 2026 10:25:23 -0400 Subject: [PATCH 08/11] fix: recover durable replay outcomes across lost replies --- service/src/config.spec.ts | 5 + service/src/config.ts | 9 +- service/src/job-cancellation-commit.test.ts | 104 +++++++++++++++++- service/src/job-cancellation.test.ts | 49 +++++++++ service/src/job-cancellation.ts | 67 +++++++---- service/src/programmatic-cancellation.test.ts | 16 ++- service/src/programmatic-cancellation.ts | 4 +- service/src/service/programmatic-router.ts | 91 ++++++++------- service/src/types/service.ts | 2 + service/src/workers.ts | 15 ++- 10 files changed, 291 insertions(+), 71 deletions(-) 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 index 9a801686..eed5db52 100644 --- a/service/src/job-cancellation-commit.test.ts +++ b/service/src/job-cancellation-commit.test.ts @@ -6,6 +6,8 @@ import { requestJobCancellation, JobCancellationRegistry, jobCancellationInternals, + fenceJobCancellation, + waitForJobWithCancellation, } from './job-cancellation'; let redis: Awaited>; @@ -30,7 +32,9 @@ test('committed results reject late Stop and survive a lost BullMQ completion re expect(await commitJobResult(redis, target, result, 60)).toBe(true); expect(await requestJobCancellation(redis, target, 60)).toBe(false); expect(await readCommittedJobResult(redis, target)).toEqual({ result }); - expect(await redis.get(jobCancellationInternals.cancellationKey(target))).toBe('completed'); + expect( + await redis.get(jobCancellationInternals.cancellationKey(target)), + ).toBe('completed'); const registry = new JobCancellationRegistry(redis); const controller = new AbortController(); try { @@ -52,5 +56,101 @@ test('concurrent cancellation and completion have exactly one winner', async () 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'); + 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, + }), + ).toBe(false); + 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, + }), + ).toBe(false); }); + +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(true); + 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 index b7f1524c..e0945f04 100644 --- a/service/src/job-cancellation.test.ts +++ b/service/src/job-cancellation.test.ts @@ -126,6 +126,21 @@ test('idle registries allocate no subscriber connection', async () => { 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; @@ -342,6 +357,7 @@ test('disconnect frees a waiting job and rejects promptly', async () => { 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({ @@ -395,6 +411,39 @@ test('registration failure fences and removes the already-enqueued job', async ( 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)); diff --git a/service/src/job-cancellation.ts b/service/src/job-cancellation.ts index 8d88fd5d..a899ea4d 100644 --- a/service/src/job-cancellation.ts +++ b/service/src/job-cancellation.ts @@ -249,13 +249,14 @@ export class JobCancellationRegistry { if (this.subscriberRestartTimer != null) clearTimeout(this.subscriberRestartTimer); this.subscriberRestartTimer = undefined; - await this.startPromise?.catch(() => undefined); const subscriber = this.subscriber; this.subscriber = undefined; this.startPromise = undefined; if (subscriber == null) return; this.detachSubscriber(subscriber); - await subscriber.quit(); + // This socket only carries notifications. Disconnect it before awaiting + // anything: subscribe() may be queued through an indefinite Redis outage. + subscriber.disconnect(false); } } @@ -292,29 +293,38 @@ export async function commitJobResult( target: JobTarget, result: T, ttlSeconds: number, + deadlineAtMs = Number.MAX_SAFE_INTEGER, ): Promise { const serialized = JSON.stringify({ result }); if (Buffer.byteLength(serialized) > 16 * 1024 * 1024) { throw new Error('Programmatic completion exceeds the 16 MiB result limit'); } - return ( - (await commands.eval( - ` + const decision = await commands.eval( + ` local state = redis.call('GET', KEYS[1]) if state == '1' then return 0 end if not state then - redis.call('SET', KEYS[2], ARGV[1], 'EX', ARGV[2]) - redis.call('SET', KEYS[1], 'completed', 'EX', ARGV[2]) + 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), - )) === 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'); + return decision === 1; } export async function readCommittedJobResult( @@ -342,7 +352,13 @@ export async function fenceJobCancellation(args: { deadlineAtMs: number; }): Promise { let retryMs = 25; - while (Date.now() < args.deadlineAtMs) { + 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; try { return await Promise.race([ @@ -350,7 +366,7 @@ export async function fenceJobCancellation(args: { new Promise(resolve => { timer = setTimeout( () => resolve(true), - Math.max(1, args.deadlineAtMs - Date.now()), + remainingMs > 0 ? remainingMs : 1_000, ); }), ]); @@ -442,20 +458,28 @@ export async function waitForJobWithCancellation(args: { 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 => - fenceJobCancellation({ + (fencing ??= fenceJobCancellation({ commands, target, ttlSeconds: cancellationTtlSeconds, - deadlineAtMs: args.deadlineAtMs ?? Date.now() + timeoutMs, - }); + deadlineAtMs, + })); const externalController = new AbortController(); try { await registry.register(target, externalController); } catch (error) { void completion.catch(() => undefined); - await fence(); + if (!(await fence())) { + const committed = await readCommittedJobResult(commands, target); + if (committed != null) return committed.result; + } await removeJobIfWaiting(job).catch(() => false); throw error; } @@ -505,7 +529,10 @@ export async function waitForJobWithCancellation(args: { } catch (error) { // Includes waitUntilFinished timeouts and registration/transport errors, // not only explicit Stop. Replay cleanup is unsafe until this barrier. - await fence(); + if (!(await fence())) { + const committed = await readCommittedJobResult(commands, target); + if (committed != null) return committed.result; + } throw error; } finally { removeAbortListener(); diff --git a/service/src/programmatic-cancellation.test.ts b/service/src/programmatic-cancellation.test.ts index ac60f7ce..7d444029 100644 --- a/service/src/programmatic-cancellation.test.ts +++ b/service/src/programmatic-cancellation.test.ts @@ -1,6 +1,7 @@ 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, @@ -130,7 +131,7 @@ test('a different principal cannot reserve, attach, cancel, or release a request expect(await redis.exists(programmaticCancellationInternals.requestKey(requestId))).toBe(1); }); -test('the owner releases cancellation state after settlement', async () => { +test('settlement retains a bounded target tombstone for late Stop classification', async () => { const requestId = 'request_release_cancel_1'; await reserveProgrammaticCancellation({ redis, @@ -138,6 +139,17 @@ test('the owner releases cancellation state after settlement', async () => { 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' }); - expect(await redis.exists(programmaticCancellationInternals.requestKey(requestId))).toBe(0); + 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 index 7b5af79d..fd92618a 100644 --- a/service/src/programmatic-cancellation.ts +++ b/service/src/programmatic-cancellation.ts @@ -86,7 +86,9 @@ return {1} const RELEASE_SCRIPT = ` if redis.call('HGET', KEYS[1], 'owner') == ARGV[1] then - return redis.call('DEL', KEYS[1]) + -- 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 `; diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index a4ead752..a8249123 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -19,6 +19,7 @@ import { removeJobIfWaiting, requestJobCancellation, fenceJobCancellation, + readCommittedJobResult, waitForJobWithCancellation, } from '../job-cancellation'; import { @@ -423,48 +424,56 @@ async function runReplayIteration( } const submittedAtMs = Date.now(); const deadlineAtMs = submittedAtMs + env.JOB_TIMEOUT; - 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 } - : {}), - cancellable: true, - 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(async (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. - await fenceJobCancellation({ commands: connection, target: cancellationTarget, - ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, deadlineAtMs }); - throw error; + 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, + 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 cancelled = await fenceJobCancellation({ + commands: connection, target: cancellationTarget, + ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, deadlineAtMs, }); + if (!cancelled) { + const committed = await readCommittedJobResult( + connection, cancellationTarget, + ); + if (committed != null) return committed.result; + } + throw error; + } jobsSubmitted.inc({ language }); return waitForJobWithCancellation({ diff --git a/service/src/types/service.ts b/service/src/types/service.ts index f10405df..e4902640 100644 --- a/service/src/types/service.ts +++ b/service/src/types/service.ts @@ -299,6 +299,8 @@ export type JobData = { 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 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 9c695979..802c35ab 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -95,7 +95,9 @@ async function processJobInner(job: t.ExecuteJob): Promise { ? { queueName: job.queueName, jobId: String(job.id) } : undefined; let cancellationRegistered = false; - const deadlineAtMs = jobDeadlineAtMs(job.timestamp, env.JOB_TIMEOUT); + 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 @@ -358,10 +360,15 @@ async function processJobInner(job: t.ExecuteJob): Promise { let lateCommitFailure = completedResult ? jobResultCommitFailure(controller.signal, env.JOB_TIMEOUT) : undefined; - if (completedResult && cancellationTarget != null && lateCommitFailure == null) { + if ( + completedResult && cancellationTarget != null && lateCommitFailure == null + ) { try { - if (!(await commitJobResult(connection, cancellationTarget, resultToCommit, - Math.ceil(env.JOB_TIMEOUT / 1_000) * 2 + 180))) { + const committed = await commitJobResult( + connection, cancellationTarget, resultToCommit, + Math.ceil(env.JOB_TIMEOUT / 1_000) * 2 + 180, deadlineAtMs, + ); + if (!committed) { lateCommitFailure = new Error(JOB_CANCELLED_MESSAGE); } } catch (error) { From 8f2b15ae415d9285cdb352107b7a8ab833683c89 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Sep 2026 10:42:19 -0400 Subject: [PATCH 09/11] fix: return atomic cancellation outcomes with aligned retention --- service/src/job-cancellation-commit.test.ts | 93 +++++++++++- service/src/job-cancellation.test.ts | 5 +- service/src/job-cancellation.ts | 150 ++++++++++++++------ service/src/service/programmatic-router.ts | 11 +- service/src/types/service.ts | 2 + service/src/workers.ts | 3 +- 6 files changed, 211 insertions(+), 53 deletions(-) diff --git a/service/src/job-cancellation-commit.test.ts b/service/src/job-cancellation-commit.test.ts index eed5db52..ae55dd36 100644 --- a/service/src/job-cancellation-commit.test.ts +++ b/service/src/job-cancellation-commit.test.ts @@ -8,6 +8,7 @@ import { jobCancellationInternals, fenceJobCancellation, waitForJobWithCancellation, + jobCancellationRetentionSeconds, } from './job-cancellation'; let redis: Awaited>; @@ -19,6 +20,94 @@ afterEach(async () => { }); const target = { queueName: 'other', jobId: 'commit-race' }; +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( @@ -71,7 +160,7 @@ test('an enqueue failure can recover a result that won cancellation fencing', as ttlSeconds: 60, deadlineAtMs: Date.now() + 5_000, }), - ).toBe(false); + ).toEqual({ status: 'completed', result }); expect(await readCommittedJobResult(redis, target)).toEqual({ result }); }); @@ -84,7 +173,7 @@ test('enqueue fencing still recovers completion after the original deadline', as ttlSeconds: 60, deadlineAtMs: Date.now() - 1_000, }), - ).toBe(false); + ).toEqual({ status: 'completed', result: { stdout: 'done' } }); }); test('Redis rejects commitment when recovery happens after the producer deadline', async () => { diff --git a/service/src/job-cancellation.test.ts b/service/src/job-cancellation.test.ts index e0945f04..e23186e4 100644 --- a/service/src/job-cancellation.test.ts +++ b/service/src/job-cancellation.test.ts @@ -79,17 +79,18 @@ class FakeRedis { _script: string, _keys: number, key: string, + _resultKey: string, ttl: number, channel: string, payload: string, - ): Promise { + ): 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; + return [1]; } async mget(...keys: string[]): Promise> { diff --git a/service/src/job-cancellation.ts b/service/src/job-cancellation.ts index a899ea4d..75e0de2f 100644 --- a/service/src/job-cancellation.ts +++ b/service/src/job-cancellation.ts @@ -260,28 +260,86 @@ export class JobCancellationRegistry { } } -export async function requestJobCancellation( +async function cancelJobInRedis( commands: IORedis, target: JobTarget, ttlSeconds: number, -): Promise { + 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 ( - (await commands.eval( - ` + return commands.eval( + ` local state = redis.call('GET', KEYS[1]) - if state and state ~= '1' then return 0 end + if state == 'completed' then + -- The request tombstone is renewed by Stop. Keep its decision and + -- result at least as long, even across API/worker config differences. + 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 + return {1} `, - 1, - cancellationKey(target), - Math.max(1, ttlSeconds), - JOB_CANCELLATION_CHANNEL, - JSON.stringify(target), - )) === 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, ); } @@ -331,26 +389,24 @@ export async function readCommittedJobResult( commands: IORedis, target: JobTarget, ): Promise<{ result: T } | undefined> { - const state = await commands.get(cancellationKey(target)); + const [state, value] = await commands.mget( + cancellationKey(target), + `${cancellationKey(target)}:result`, + ); if (state !== 'completed') return undefined; - const value = await commands.get(`${cancellationKey(target)}:result`); - if (value == null) - throw new Error( - 'Committed programmatic result expired; refusing re-execution', - ); - return JSON.parse(value); + 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: { +export async function fenceJobCancellation(args: { commands: IORedis; target: JobTarget; ttlSeconds: number; deadlineAtMs: number; -}): Promise { +}): Promise> { let retryMs = 25; let firstAttempt = true; while (firstAttempt || Date.now() < args.deadlineAtMs) { @@ -362,10 +418,25 @@ export async function fenceJobCancellation(args: { let timer: ReturnType | undefined; try { return await Promise.race([ - requestJobCancellation(args.commands, args.target, args.ttlSeconds), - new Promise(resolve => { + cancelJobInRedis( + args.commands, + args.target, + args.ttlSeconds, + true, + ).then((decision): JobFenceOutcome => { + 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'); + }), + new Promise>(resolve => { timer = setTimeout( - () => resolve(true), + () => resolve({ status: 'expired' }), remainingMs > 0 ? remainingMs : 1_000, ); }), @@ -382,7 +453,7 @@ export async function fenceJobCancellation(args: { if (timer != null) clearTimeout(timer); } } - return true; + return { status: 'expired' }; } const REMOVABLE_JOB_STATES = new Set([ @@ -463,9 +534,9 @@ export async function waitForJobWithCancellation(args: { 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({ + let fencing: Promise> | undefined; + const fence = (): Promise> => + (fencing ??= fenceJobCancellation({ commands, target, ttlSeconds: cancellationTtlSeconds, @@ -476,23 +547,24 @@ export async function waitForJobWithCancellation(args: { await registry.register(target, externalController); } catch (error) { void completion.catch(() => undefined); - if (!(await fence())) { - const committed = await readCommittedJobResult(commands, target); - if (committed != null) return committed.result; - } + 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((_, reject) => { + const disconnected = new Promise((resolve, reject) => { let cancelling = false; const cancel = (): void => { if (cancelling) return; cancelling = true; void fence() - .then(async accepted => { - if (!accepted) return; + .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. @@ -529,10 +601,8 @@ export async function waitForJobWithCancellation(args: { } catch (error) { // Includes waitUntilFinished timeouts and registration/transport errors, // not only explicit Stop. Replay cleanup is unsafe until this barrier. - if (!(await fence())) { - const committed = await readCommittedJobResult(commands, target); - if (committed != null) return committed.result; - } + const outcome = await fence(); + if (outcome.status === 'completed') return outcome.result; throw error; } finally { removeAbortListener(); diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index a8249123..bb140013 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -19,7 +19,6 @@ import { removeJobIfWaiting, requestJobCancellation, fenceJobCancellation, - readCommittedJobResult, waitForJobWithCancellation, } from '../job-cancellation'; import { @@ -444,6 +443,7 @@ async function runReplayIteration( ...(state.workspaceId != null ? { workspaceId: state.workspaceId } : {}), cancellable: true, deadlineAtMs, + cancellationTtlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, executionManifestClaims: sandboxSecurity.executionManifestClaims, @@ -462,16 +462,11 @@ async function runReplayIteration( // 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 cancelled = await fenceJobCancellation({ + const outcome = await fenceJobCancellation({ commands: connection, target: cancellationTarget, ttlSeconds: PROGRAMMATIC_CANCELLATION_TTL_SECONDS, deadlineAtMs, }); - if (!cancelled) { - const committed = await readCommittedJobResult( - connection, cancellationTarget, - ); - if (committed != null) return committed.result; - } + if (outcome.status === 'completed') return outcome.result; throw error; } jobsSubmitted.inc({ language }); diff --git a/service/src/types/service.ts b/service/src/types/service.ts index e4902640..a0c78486 100644 --- a/service/src/types/service.ts +++ b/service/src/types/service.ts @@ -301,6 +301,8 @@ export type JobData = { 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 802c35ab..c71276c2 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -42,6 +42,7 @@ import { jobResultCommitFailure, commitJobResult, readCommittedJobResult, + jobCancellationRetentionSeconds, throwIfJobAborted, } from './job-cancellation'; import logger from './logger'; @@ -366,7 +367,7 @@ async function processJobInner(job: t.ExecuteJob): Promise { try { const committed = await commitJobResult( connection, cancellationTarget, resultToCommit, - Math.ceil(env.JOB_TIMEOUT / 1_000) * 2 + 180, deadlineAtMs, + jobCancellationRetentionSeconds(env.JOB_TIMEOUT, job.data.cancellationTtlSeconds), deadlineAtMs, ); if (!committed) { lateCommitFailure = new Error(JOB_CANCELLED_MESSAGE); From 7059bd4427aee536da635f8ebff73a3a4c77f495 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Sep 2026 10:57:54 -0400 Subject: [PATCH 10/11] fix: commit native results inside the workspace mutation fence --- service/src/job-cancellation-commit.test.ts | 121 +++++++++- service/src/job-cancellation.ts | 42 ++-- service/src/programmatic-cancellation.test.ts | 219 ++++++++++++------ service/src/programmatic-cancellation.ts | 61 +++-- service/src/workers.ts | 216 ++++++++++------- 5 files changed, 463 insertions(+), 196 deletions(-) diff --git a/service/src/job-cancellation-commit.test.ts b/service/src/job-cancellation-commit.test.ts index ae55dd36..1ba2ef93 100644 --- a/service/src/job-cancellation-commit.test.ts +++ b/service/src/job-cancellation-commit.test.ts @@ -1,5 +1,6 @@ import { afterEach, beforeEach, expect, test } from 'bun:test'; import { startTestRedis } from './test/redis'; +import { RedisBridgeStore } from './bridge/store'; import { commitJobResult, readCommittedJobResult, @@ -20,6 +21,115 @@ afterEach(async () => { }); const target = { queueName: 'other', jobId: 'commit-race' }; +for (const stopWins of [false, true]) + test(`native mutation handoff commits or quarantines before root release (stopWins=${stopWins})`, 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 (stopWins) await requestJobCancellation(redis, target, 60); + if ( + !(await commitJobResult( + redis, + target, + { stdout: 'mutation settled' }, + 60, + )) + ) + throw new Error('cancelled before handoff'); + // 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 (stopWins) { + await expect(completion).rejects.toThrow('cancelled before handoff'); + 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' }, + }); + } + }); + +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); @@ -42,10 +152,17 @@ test('retention renewal does not lose subsecond time to rounded TTL readings', a 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(` + 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)); + `, + 1, + key, + ), + ); const before = await expiration(); await requestJobCancellation(redis, target, 60); expect(await expiration()).toBeGreaterThan(before); diff --git a/service/src/job-cancellation.ts b/service/src/job-cancellation.ts index 75e0de2f..19715c9e 100644 --- a/service/src/job-cancellation.ts +++ b/service/src/job-cancellation.ts @@ -272,8 +272,9 @@ async function cancelJobInRedis( ` local state = redis.call('GET', KEYS[1]) if state == 'completed' then - -- The request tombstone is renewed by Stop. Keep its decision and - -- result at least as long, even across API/worker config differences. + -- 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 @@ -416,27 +417,13 @@ export async function fenceJobCancellation(args: { // winning decision. Never translate a known committed effect to failure. const remainingMs = args.deadlineAtMs - Date.now(); let timer: ReturnType | undefined; + let decision: unknown; try { - return await Promise.race([ - cancelJobInRedis( - args.commands, - args.target, - args.ttlSeconds, - true, - ).then((decision): JobFenceOutcome => { - 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'); - }), - new Promise>(resolve => { + decision = await Promise.race([ + cancelJobInRedis(args.commands, args.target, args.ttlSeconds, true), + new Promise(resolve => { timer = setTimeout( - () => resolve({ status: 'expired' }), + () => resolve(undefined), remainingMs > 0 ? remainingMs : 1_000, ); }), @@ -449,9 +436,22 @@ export async function fenceJobCancellation(args: { ), ); 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' }; } diff --git a/service/src/programmatic-cancellation.test.ts b/service/src/programmatic-cancellation.test.ts index 7d444029..53a4340e 100644 --- a/service/src/programmatic-cancellation.test.ts +++ b/service/src/programmatic-cancellation.test.ts @@ -22,60 +22,110 @@ afterEach(async () => { }); test('normalizes only bounded opaque request IDs', () => { - expect(normalizeProgrammaticRequestId('request_123456789')).toBe('request_123456789'); + expect(normalizeProgrammaticRequestId('request_123456789')).toBe( + 'request_123456789', + ); expect(normalizeProgrammaticRequestId(' short ')).toBeUndefined(); - expect(normalizeProgrammaticRequestId('../request_123456789')).toBeUndefined(); + expect( + normalizeProgrammaticRequestId('../request_123456789'), + ).toBeUndefined(); expect(normalizeProgrammaticRequestId('a'.repeat(129))).toBeUndefined(); }); -test('cancellation before queue attachment is retained atomically', async () => { - const requestId = 'request_early_cancel_123'; +test('Stop cannot extend an attached tombstone before the outcome command succeeds', async () => { + const requestId = 'request_no_split_renewal'; const owner = 'owner-a'; - expect(await cancelProgrammaticRequest({ - redis, - requestId, - owner, - ttlSeconds: 60, - })).toEqual({ status: 'accepted' }); - - expect(await reserveProgrammaticCancellation({ + const target = { queueName: 'other', jobId: 'split-renewal' }; + await reserveProgrammaticCancellation({ redis, requestId, owner, ttlSeconds: 60, - })).toBe('cancelled'); - expect(await attachProgrammaticCancellationTarget({ + }); + await attachProgrammaticCancellationTarget({ redis, requestId, owner, - target: { queueName: 'other', jobId: '42' }, + target, ttlSeconds: 60, - })).toBe('cancelled'); + }); + 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 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({ + expect( + await cancelProgrammaticRequest({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toEqual({ status: 'accepted', target: { queueName: 'other', jobId: '43' }, }); @@ -84,19 +134,23 @@ test('cancellation after attachment returns the exact queue target', async () => 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('active'); - expect(await reserveProgrammaticCancellation({ - redis, - requestId, - owner, - ttlSeconds: 60, - })).toBe('duplicate'); + expect( + await reserveProgrammaticCancellation({ + redis, + requestId, + owner, + ttlSeconds: 60, + }), + ).toBe('duplicate'); }); test('a different principal cannot reserve, attach, cancel, or release a request', async () => { @@ -108,27 +162,39 @@ test('a different principal cannot reserve, attach, cancel, or release a request 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({ + 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', - ttlSeconds: 60, - })).toEqual({ status: 'forbidden' }); - await releaseProgrammaticCancellation({ redis, requestId, owner: 'owner-b' }); - expect(await redis.exists(programmaticCancellationInternals.requestKey(requestId))).toBe(1); + }); + expect( + await redis.exists(programmaticCancellationInternals.requestKey(requestId)), + ).toBe(1); }); test('settlement retains a bounded target tombstone for late Stop classification', async () => { @@ -140,16 +206,29 @@ test('settlement retains a bounded target tombstone for late Stop classification ttlSeconds: 60, }); const target = { queueName: 'other', jobId: 'settled-job' }; - await attachProgrammaticCancellationTarget({ redis, requestId, owner: 'owner-a', - target, ttlSeconds: 60 }); + await attachProgrammaticCancellationTarget({ + redis, + requestId, + owner: 'owner-a', + target, + ttlSeconds: 60, + }); await commitJobResult(redis, target, { stdout: 'done' }, 60); - await releaseProgrammaticCancellation({ redis, requestId, owner: 'owner-a' }); + 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 }); + 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 index fd92618a..12a6694d 100644 --- a/service/src/programmatic-cancellation.ts +++ b/service/src/programmatic-cancellation.ts @@ -4,7 +4,8 @@ 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'; +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}$/; @@ -21,7 +22,9 @@ function requestKey(requestId: string): string { return `${REQUEST_PREFIX}:${requestId}`; } -export function normalizeProgrammaticRequestId(value: unknown): string | undefined { +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; @@ -33,12 +36,14 @@ export function programmaticCancellationOwner( ): string { const identity = getExecutionIdentity(req, userId); return createHash('sha256') - .update(JSON.stringify([ - identity.storageNamespace, - identity.canonicalUserId, - getCredentialId(req), - identity.authContextHash ?? '', - ])) + .update( + JSON.stringify([ + identity.storageNamespace, + identity.canonicalUserId, + getCredentialId(req), + identity.authContextHash ?? '', + ]), + ) .digest('hex'); } @@ -77,10 +82,12 @@ 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') -redis.call('EXPIRE', key, ttl) 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} `; @@ -99,13 +106,15 @@ export async function reserveProgrammaticCancellation(args: { 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), - )); + 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'; @@ -118,15 +127,17 @@ export async function attachProgrammaticCancellationTarget(args: { 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), - )); + 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'; } diff --git a/service/src/workers.ts b/service/src/workers.ts index c71276c2..e3e9ba14 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -50,7 +50,10 @@ 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}`; @@ -97,7 +100,10 @@ async function processJobInner(job: t.ExecuteJob): Promise { : undefined; let cancellationRegistered = false; const deadlineAtMs = jobDeadlineAtMs( - job.timestamp, env.JOB_TIMEOUT, Date.now(), job.data.deadlineAtMs, + job.timestamp, + env.JOB_TIMEOUT, + Date.now(), + job.data.deadlineAtMs, ); const remainingBudgetMs = Math.max(0, deadlineAtMs - Date.now()); const timer = @@ -110,12 +116,20 @@ async function processJobInner(job: t.ExecuteJob): Promise { 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 committed = await readCommittedJobResult(connection, cancellationTarget); + const committed = await readCommittedJobResult( + connection, + cancellationTarget, + ); if (committed != null) return committed.result; } if (controller.signal.aborted) { @@ -157,7 +171,13 @@ async function processJobInner(job: t.ExecuteJob): Promise { 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, @@ -190,20 +210,36 @@ async function processJobInner(job: t.ExecuteJob): Promise { const finalizeSandboxResult = async ( result: SandboxRawResponse, ): Promise => { - if ( - resultRestoreToken === undefined || - resultRestoreToken.length === 0 || - finalizedSandboxResults.has(result) - ) { - return result; + 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) throw new Error(JOB_CANCELLED_MESSAGE); + resultToCommit = mapped; + resultCommittedAtHandoff = true; } - const restored = await restoreGatewaySandboxResult({ - grantId: egressGrantId, - egressGrantToken: resultRestoreToken, - result, - isSynthetic: isSyntheticJob, - signal: controller.signal, - }); finalizedSandboxResults.add(restored); return restored; }; @@ -231,7 +267,8 @@ async function processJobInner(job: t.ExecuteJob): Promise { * before checkpointing/reusing the mutated workspace. Stateless/HTTP * paths retain the worker-owned fallback immediately below. */ sessionResultFinalizer: - resultRestoreToken !== undefined && resultRestoreToken.length > 0 + commitAtHandoff || + (resultRestoreToken !== undefined && resultRestoreToken.length > 0) ? finalizeSandboxResult : undefined, }, @@ -240,64 +277,76 @@ async function processJobInner(job: t.ExecuteJob): Promise { 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. - throwIfJobAborted(controller.signal); + 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) { + // 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; @@ -351,36 +400,47 @@ async function processJobInner(job: t.ExecuteJob): Promise { isSynthetic: isSyntheticJob, reason: revokeReason, timeoutMs: env.EGRESS_GATEWAY_REVOKE_TIMEOUT_MS, - }).catch((error) => { + }).catch(error => { logger.error('Failed to revoke egress grant', { grantId: egressGrantId, error: getAxiosErrorDetails(error), }); }); } - let lateCommitFailure = completedResult - ? jobResultCommitFailure(controller.signal, env.JOB_TIMEOUT) - : undefined; + let lateCommitFailure = + completedResult && !resultCommittedAtHandoff + ? jobResultCommitFailure(controller.signal, env.JOB_TIMEOUT) + : undefined; if ( - completedResult && cancellationTarget != null && lateCommitFailure == null + completedResult && + !resultCommittedAtHandoff && + cancellationTarget != null && + lateCommitFailure == null ) { try { const committed = await commitJobResult( - connection, cancellationTarget, resultToCommit, - jobCancellationRetentionSeconds(env.JOB_TIMEOUT, job.data.cancellationTtlSeconds), deadlineAtMs, + connection, + cancellationTarget, + resultToCommit, + jobCancellationRetentionSeconds( + env.JOB_TIMEOUT, + job.data.cancellationTtlSeconds, + ), + deadlineAtMs, ); if (!committed) { lateCommitFailure = new Error(JOB_CANCELLED_MESSAGE); } } catch (error) { - lateCommitFailure = error instanceof Error ? error : new Error('Result commit failed'); + 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) => { + .catch(error => { logger.warn('Failed to clear queued execution cancellation state', { queueName: cancellationTarget.queueName, jobId: cancellationTarget.jobId, @@ -418,14 +478,14 @@ export const otherWorker = new Worker(queueNames.other, processJob, { workerRunning.set({ worker_type: 'python' }, 1); workerRunning.set({ worker_type: 'other' }, 1); -pyWorker.on('completed', (job) => { +pyWorker.on('completed', job => { if (job.data.isSynthetic !== true) { logger.info(`[${WORKER_ID}] Python job completed ${job.id}`); } jobsCompleted.inc({ language: 'python' }); }); -otherWorker.on('completed', (job) => { +otherWorker.on('completed', job => { if (job.data.isSynthetic !== true) { logger.info(`[${WORKER_ID}] Other job completed ${job.id}`); } @@ -452,12 +512,12 @@ otherWorker.on('failed', (job, 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); }); From e8b7b340a8488c31820225c33adf6232e3bb2fca Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 14 Sep 2026 11:08:21 -0400 Subject: [PATCH 11/11] fix: claim programmatic execution before stalled-job redelivery --- service/src/job-cancellation-commit.test.ts | 93 ++++++++++++++++++--- service/src/job-cancellation.ts | 49 ++++++++++- service/src/workers.ts | 20 +++-- 3 files changed, 143 insertions(+), 19 deletions(-) diff --git a/service/src/job-cancellation-commit.test.ts b/service/src/job-cancellation-commit.test.ts index 1ba2ef93..8ded29f2 100644 --- a/service/src/job-cancellation-commit.test.ts +++ b/service/src/job-cancellation-commit.test.ts @@ -10,6 +10,7 @@ import { fenceJobCancellation, waitForJobWithCancellation, jobCancellationRetentionSeconds, + claimJobExecution, } from './job-cancellation'; let redis: Awaited>; @@ -21,8 +22,8 @@ afterEach(async () => { }); const target = { queueName: 'other', jobId: 'commit-race' }; -for (const stopWins of [false, true]) - test(`native mutation handoff commits or quarantines before root release (stopWins=${stopWins})`, async () => { +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'; @@ -59,16 +60,23 @@ for (const stopWins of [false, true]) const completion = store.dispatch({ ...dispatchArgs, finalize: async settlement => { - if (stopWins) await requestJobCancellation(redis, target, 60); + if (outcome === 'stop') await requestJobCancellation(redis, target, 60); + if (outcome === 'duplicate') + await commitJobResult( + redis, + target, + { stdout: 'first mutation' }, + 60, + ); if ( - !(await commitJobResult( + (await commitJobResult( redis, target, { stdout: 'mutation settled' }, 60, - )) + )) !== 'committed' ) - throw new Error('cancelled before handoff'); + 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); @@ -91,8 +99,8 @@ for (const stopWins of [false, true]) files: [], }, }); - if (stopWins) { - await expect(completion).rejects.toThrow('cancelled before handoff'); + if (outcome !== 'commit') { + await expect(completion).rejects.toThrow('handoff did not win'); await expect(store.dispatch(dispatchArgs)).rejects.toMatchObject({ code: 'WORKSPACE_QUARANTINED', }); @@ -106,6 +114,67 @@ for (const stopWins of [false, true]) } }); +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); @@ -228,14 +297,14 @@ test('disconnect returns a known completed result without waiting for a lost que 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( - false, + '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(true); + 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( @@ -256,7 +325,7 @@ test('concurrent cancellation and completion have exactly one winner', async () requestJobCancellation(redis, target, 60), commitJobResult(redis, target, { stdout: 'result' }, 60), ]); - expect(Number(cancelled) + Number(committed)).toBe(1); + expect(Number(cancelled) + Number(committed === 'committed')).toBe(1); }); test('a missing committed result fails closed instead of re-executing', async () => { @@ -322,7 +391,7 @@ test('a timely durable commit remains successful when only its acknowledgement i 60, Date.now() + 100, ), - ).toBe(true); + ).toBe('committed'); expect(await readCommittedJobResult(redis, target)).toEqual({ result: { stdout: 'committed' }, }); diff --git a/service/src/job-cancellation.ts b/service/src/job-cancellation.ts index 19715c9e..4da3253d 100644 --- a/service/src/job-cancellation.ts +++ b/service/src/job-cancellation.ts @@ -353,7 +353,7 @@ export async function commitJobResult( result: T, ttlSeconds: number, deadlineAtMs = Number.MAX_SAFE_INTEGER, -): Promise { +): 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'); @@ -362,6 +362,8 @@ export async function commitJobResult( ` 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 @@ -383,7 +385,50 @@ export async function commitJobResult( ); if (decision === -1) throw new Error('Job result commitment exceeded its deadline'); - return decision === 1; + 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( diff --git a/service/src/workers.ts b/service/src/workers.ts index e3e9ba14..a7153b53 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -41,7 +41,7 @@ import { JOB_CANCELLED_MESSAGE, jobResultCommitFailure, commitJobResult, - readCommittedJobResult, + claimJobExecution, jobCancellationRetentionSeconds, throwIfJobAborted, } from './job-cancellation'; @@ -126,11 +126,15 @@ async function processJobInner(job: t.ExecuteJob): Promise { if (cancellationTarget != null) { await jobCancellationRegistry.register(cancellationTarget, controller); cancellationRegistered = true; - const committed = await readCommittedJobResult( + const claim = await claimJobExecution( connection, cancellationTarget, + jobCancellationRetentionSeconds( + env.JOB_TIMEOUT, + job.data.cancellationTtlSeconds, + ), ); - if (committed != null) return committed.result; + if (claim.status === 'completed') return claim.result; } if (controller.signal.aborted) { throw new Error(`Job timed out after ${env.JOB_TIMEOUT}ms`); @@ -236,7 +240,9 @@ async function processJobInner(job: t.ExecuteJob): Promise { ), deadlineAtMs, ); - if (!committed) throw new Error(JOB_CANCELLED_MESSAGE); + 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; } @@ -428,8 +434,12 @@ async function processJobInner(job: t.ExecuteJob): Promise { ), deadlineAtMs, ); - if (!committed) { + 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 =