diff --git a/docs/byom-worker-admission.md b/docs/byom-worker-admission.md new file mode 100644 index 00000000..c88d5669 --- /dev/null +++ b/docs/byom-worker-admission.md @@ -0,0 +1,26 @@ +# BYOM worker admission + +Workspace tool calls to a busy worker wait in a bounded FIFO shared through Redis. +The limit is 32 admitted requests per worker, including the active request. When +the limit is reached, the workspace endpoint returns HTTP 429 with +`WORKER_QUEUE_FULL`. A different worker has an independent admission queue. + +Waiting uses the caller's existing absolute dispatch deadline. It does not reset +or extend execution timeouts. Disconnecting or cancelling removes the waiting +request without cancelling the active assignment. Expired entries are pruned; +Redis key expiry also bounds state left by a crashed API process. + +After admission, the API revalidates the worker incarnation, identity, tenant +binding and workspace operation. A waiting request cannot migrate to a replacement +worker. Existing execution acknowledgement, fencing, settlement and quarantine +rules remain responsible for the active assignment. + +This is compatible with existing workers and requires only a Code API update. +Existing workers still execute one assignment at a time. Parallel execution across +workspaces requires separate lease claims and isolated native sandbox contexts; +this admission change does not advertise that capability. Queue time and execution +time currently share the HTTP request deadline. Separate budgets require a matching +LibreChat client change so that the client does not disconnect while waiting. + +Focused regression coverage lives in `service/src/bridge/admission.test.ts` and +`service/src/bridge/worker-admission.test.ts`. diff --git a/service/src/bridge/admission.test.ts b/service/src/bridge/admission.test.ts new file mode 100644 index 00000000..b2268012 --- /dev/null +++ b/service/src/bridge/admission.test.ts @@ -0,0 +1,37 @@ +import { afterEach, expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; +import type Redis from 'ioredis'; +import { BridgeAdmissionQueue } from './admission'; + +const redis = new RedisMock() as unknown as Redis; +afterEach(async () => { + await redis.flushall(); +}); + +test('bounds FIFO admission across API replicas and releases cancelled waiters', async () => { + const firstReplica = new BridgeAdmissionQueue(redis, 2); + const secondReplica = new BridgeAdmissionQueue(redis, 2); + const deadline = Date.now() + 5000; + expect(await firstReplica.enter('worker', 'first', deadline)).toBe(true); + expect(await secondReplica.enter('worker', 'second', deadline)).toBe(true); + expect(await firstReplica.enter('worker', 'third', deadline)).toBe(false); + expect(await secondReplica.isHead('worker', 'second')).toBe(false); + await firstReplica.leave('worker', 'first'); + expect(await secondReplica.isHead('worker', 'second')).toBe(true); + expect(await firstReplica.enter('worker', 'third', deadline)).toBe(true); + expect(await firstReplica.isHead('worker', 'third')).toBe(false); +}); + +test('expired crashed callers cannot strand the next request or consume capacity', async () => { + const queue = new BridgeAdmissionQueue(redis, 1); + expect(await queue.enter('worker', 'expired', Date.now() - 1)).toBe(true); + expect(await queue.enter('worker', 'live', Date.now() + 5000)).toBe(true); + expect(await queue.isHead('worker', 'live')).toBe(true); +}); + +test('different machines do not share admission capacity', async () => { + const queue = new BridgeAdmissionQueue(redis, 1); + expect(await queue.enter('worker-a', 'a', Date.now() + 5000)).toBe(true); + expect(await queue.enter('worker-b', 'b', Date.now() + 5000)).toBe(true); + expect(await queue.isHead('worker-b', 'b')).toBe(true); +}); diff --git a/service/src/bridge/admission.ts b/service/src/bridge/admission.ts new file mode 100644 index 00000000..1e19286e --- /dev/null +++ b/service/src/bridge/admission.ts @@ -0,0 +1,90 @@ +import type Redis from 'ioredis'; + +/** Bounded FIFO admission shared by API replicas. Entries expire after caller deadlines. */ +export class BridgeAdmissionQueue { + constructor( + private readonly redis: Redis, + private readonly capacity = 32, + ) {} + + private keys(workerId: string): [string, string, string] { + const prefix = `codeapi:bridge:v1:worker:${encodeURIComponent(workerId)}:admission`; + return [prefix, `${prefix}:deadlines`, `${prefix}:sequence`]; + } + + async enter( + workerId: string, + id: string, + deadlineAtMs: number, + ): Promise { + return ( + Number( + await this.redis.eval( + [ + "local expired = redis.call('ZRANGEBYSCORE', KEYS[2], '-inf', ARGV[2])", + 'for _, id in ipairs(expired) do', + " redis.call('ZREM', KEYS[1], id)", + " redis.call('ZREM', KEYS[2], id)", + 'end', + "if redis.call('ZSCORE', KEYS[1], ARGV[1]) then return 1 end", + "if redis.call('ZCARD', KEYS[1]) >= tonumber(ARGV[4]) then return 0 end", + "local sequence = redis.call('INCR', KEYS[3])", + "redis.call('ZADD', KEYS[1], sequence, ARGV[1])", + "redis.call('ZADD', KEYS[2], ARGV[3], ARGV[1])", + "local latest = redis.call('ZREVRANGE', KEYS[2], 0, 0, 'WITHSCORES')", + 'for _, key in ipairs(KEYS) do', + " redis.call('PEXPIREAT', key, tonumber(latest[2]) + 30000)", + 'end', + 'return 1', + ].join('\n'), + 3, + ...this.keys(workerId), + id, + Date.now(), + deadlineAtMs, + this.capacity, + ), + ) === 1 + ); + } + + async isHead(workerId: string, id: string): Promise { + const [order, deadlines] = this.keys(workerId); + return ( + Number( + await this.redis.eval( + [ + "local expired = redis.call('ZRANGEBYSCORE', KEYS[2], '-inf', ARGV[2])", + 'for _, id in ipairs(expired) do', + " redis.call('ZREM', KEYS[1], id)", + " redis.call('ZREM', KEYS[2], id)", + 'end', + "local head = redis.call('ZRANGE', KEYS[1], 0, 0)", + 'if head[1] == ARGV[1] then return 1 end', + 'return 0', + ].join('\n'), + 2, + order, + deadlines, + id, + Date.now(), + ), + ) === 1 + ); + } + + async leave(workerId: string, id: string): Promise { + const [order, deadlines] = this.keys(workerId); + await this.redis.eval( + [ + "redis.call('ZREM', KEYS[1], ARGV[1])", + "redis.call('ZREM', KEYS[2], ARGV[1])", + 'return 1', + ].join('\n'), + 2, + order, + deadlines, + id, + ); + } +} diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index de302654..43634c7c 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -18,6 +18,7 @@ import { isWorkspaceToolResult, } from '../../../packages/code/src/protocol'; import type { BridgeWorkerBinding } from './pairing'; +import { BridgeAdmissionQueue } from './admission'; const PREFIX = 'codeapi:bridge:v1'; const POLL_INTERVAL_MS = 100; @@ -43,6 +44,7 @@ export class BridgeStoreError extends Error { | 'WORKER_OFFLINE' | 'WORKER_UNAUTHORIZED' | 'WORKER_BUSY' + | 'WORKER_QUEUE_FULL' | 'ASSIGNMENT_EXPIRED' | 'ASSIGNMENT_FENCED' | 'ASSIGNMENT_NOT_FOUND' @@ -684,18 +686,42 @@ export class RedisBridgeStore { const lockIncarnationId = registration.incarnationId; let assignment: StoredAssignment | undefined; let resultCommitted = false; + const admission = args.workspaceRequest == null + ? undefined + : new BridgeAdmissionQueue(this.redis); try { - const locked = await this.dispatchCommand( - () => - this.acquireLock( - args.workerId, - assignmentId, - lockIncarnationId, - ttlSeconds, - ), + if (admission != null && !(await this.dispatchCommand( + () => admission.enter(args.workerId, assignmentId, args.deadlineAtMs), args, - 'Bridge assignment lock acquisition', - ); + 'Bridge admission enqueue', + ))) { + throw new BridgeStoreError('WORKER_QUEUE_FULL', 'Bridge worker pending request limit reached'); + } + let locked = false; + do { + if (admission != null && !(await this.dispatchCommand( + () => admission.isHead(args.workerId, assignmentId), + args, + 'Bridge admission position', + ))) { + await delay(Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), args.signal); + continue; + } + locked = await this.dispatchCommand( + () => + this.acquireLock( + args.workerId, + assignmentId, + lockIncarnationId, + ttlSeconds, + ), + args, + 'Bridge assignment lock acquisition', + ); + if (!locked && admission != null) { + await delay(Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), args.signal); + } + } while (!locked && admission != null); if (!locked) { throw new BridgeStoreError( 'WORKER_BUSY', @@ -703,6 +729,25 @@ export class RedisBridgeStore { ); } this.assertDispatchActive(args.signal, args.deadlineAtMs); + if (admission != null) { + // Waiting must not transfer accepted work to a replacement machine or identity. + const current = await this.dispatchCommand( + () => this.dispatchableRegistration(args.workerId), + args, + 'Bridge admitted worker validation', + ); + if ( + current == null || + current.registration.incarnationId !== registration.incarnationId || + current.registration.identityId !== registration.identityId || + current.registration.binding?.tenantId !== registration.binding?.tenantId + ) { + throw new BridgeStoreError('WORKER_OFFLINE', 'Bridge worker changed while the request was waiting'); + } + if (!supportsWorkspaceTool(current.registration, args.workspaceRequest!)) { + throw new BridgeStoreError('WORKER_MISMATCH', 'Bridge worker capabilities changed while the request was waiting'); + } + } const generation = await this.dispatchCommand( () => this.redis.incr(generationKey(args.workerId)), args, @@ -748,6 +793,12 @@ export class RedisBridgeStore { 'Bridge assignment enqueue', ); if (queued) break; + if (admission != null) { + throw new BridgeStoreError( + 'WORKER_FENCED', + 'Bridge worker changed before the waiting request could be dispatched', + ); + } const replacement = await this.dispatchCommand( () => this.dispatchableRegistration(args.workerId), args, @@ -810,6 +861,14 @@ export class RedisBridgeStore { throw error; } } finally { + if (admission != null) { + // Expiry remains the fallback if Redis is unavailable during cancellation. + await boundedCommand( + admission.leave(args.workerId, assignmentId), + this.redisCommandTimeoutMs, + 'Bridge admission cleanup', + ).catch(() => undefined); + } if (resultCommitted) { try { await this.cleanupWithRetry(args.workerId, assignmentId, assignment); diff --git a/service/src/bridge/worker-admission.test.ts b/service/src/bridge/worker-admission.test.ts new file mode 100644 index 00000000..57ad3d84 --- /dev/null +++ b/service/src/bridge/worker-admission.test.ts @@ -0,0 +1,144 @@ +import { afterEach, expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; +import type Redis from 'ioredis'; +import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import { RedisBridgeStore } from './store'; +import type { CodeBridgeAssignment } from './store'; + +const redis = new RedisMock() as unknown as Redis; +const store = new RedisBridgeStore(redis); +const workerId = 'admission-worker'; +const incarnationId = 'incarnation-00000001'; +afterEach(async () => { + await redis.flushall(); +}); + +async function register( + operations: Array<'read_file' | 'list_files'> = ['read_file'], +): Promise { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'native-srt', + runtimes: ['bash'], + workspaceTools: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operations, + workspaces: [{ id: 'primary', name: 'Workspace' }], + }, + }, + }); +} + +function dispatch( + path: string, + controller = new AbortController(), + budgetMs = 5000, +): ReturnType { + return store.dispatchWorkspaceTool({ + workerId, + signal: controller.signal, + deadlineAtMs: Date.now() + budgetMs, + request: { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + operation: 'read_file', + workspaceId: 'primary', + path, + }, + }); +} + +async function settle( + assignment: CodeBridgeAssignment | undefined, +): Promise { + expect(assignment).toBeDefined(); + await store.settle(workerId, assignment!.assignmentId, { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId, + generation: assignment!.generation, + leaseToken: assignment!.leaseToken, + status: 'rejected', + error: 'File not found', + }); +} + +test('a second workspace call waits until the first settles instead of returning WORKER_BUSY', async () => { + await register(); + const first = dispatch('first'); + const assignment = await store.lease(workerId, incarnationId, 1000); + const second = dispatch('second'); + await settle(assignment); + await first; + const next = await store.lease(workerId, incarnationId, 1000); + expect(next?.request).toMatchObject({ path: 'second' }); + await settle(next); + await expect(second).resolves.toMatchObject({ + status: 'rejected', + error: 'File not found', + }); +}); + +test('cancelling a waiting caller does not release or cancel the active assignment', async () => { + await register(); + const first = dispatch('first'); + const assignment = await store.lease(workerId, incarnationId, 1000); + const controller = new AbortController(); + const second = dispatch('second', controller); + void second.catch(() => undefined); + const waitDeadline = Date.now() + 1000; + while ( + (await redis.zcard(`codeapi:bridge:v1:worker:${workerId}:admission`)) < 2 + ) { + if (Date.now() >= waitDeadline) + throw new Error('Second caller never entered admission'); + await new Promise((resolve) => setTimeout(resolve, 5)); + } + controller.abort(); + await expect(second).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + expect(await redis.get(`codeapi:bridge:v1:worker:${workerId}:lock`)).toBe( + assignment!.assignmentId, + ); + await settle(assignment); + await first; + expect(await store.lease(workerId, incarnationId, 50)).toBeUndefined(); +}); + +test('an expired queued call never reaches the worker and does not strand later calls', async () => { + await register(); + const first = dispatch('first'); + const assignment = await store.lease(workerId, incarnationId, 1000); + await expect( + dispatch('expired', new AbortController(), 25), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + const third = dispatch('third'); + await settle(assignment); + await first; + const next = await store.lease(workerId, incarnationId, 1000); + expect(next?.request).toMatchObject({ path: 'third' }); + await settle(next); + await third; +}); + +test('a queued request is rejected if the worker withdraws its capability', async () => { + await register(); + const first = dispatch('first'); + const assignment = await store.lease(workerId, incarnationId, 1000); + const second = dispatch('second'); + void second.catch(() => undefined); + const deadline = Date.now() + 1000; + while ( + (await redis.zcard(`codeapi:bridge:v1:worker:${workerId}:admission`)) < 2 + ) { + if (Date.now() > deadline) + throw new Error('Second caller did not enter admission'); + await new Promise((resolve) => setTimeout(resolve, 5)); + } + await register(['list_files']); + await settle(assignment); + await first; + await expect(second).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + expect(await store.lease(workerId, incarnationId, 50)).toBeUndefined(); +}); diff --git a/service/src/sandbox-backend/remote-bridge.test.ts b/service/src/sandbox-backend/remote-bridge.test.ts index b0121d2e..a65bc89e 100644 --- a/service/src/sandbox-backend/remote-bridge.test.ts +++ b/service/src/sandbox-backend/remote-bridge.test.ts @@ -79,6 +79,7 @@ describe('RemoteBridgeSandboxBackend', () => { WORKER_OFFLINE: ['BRIDGE_WORKER_OFFLINE', true, 503, 'Code environment is offline'], WORKER_UNAUTHORIZED: ['BRIDGE_WORKER_UNAUTHORIZED', false, 403, 'Code environment is not authorized for this tenant'], WORKER_BUSY: ['BRIDGE_WORKER_BUSY', false, 409, 'Code environment is busy'], + WORKER_QUEUE_FULL: ['BRIDGE_WORKER_BUSY', false, 409, 'Code environment is busy'], ASSIGNMENT_EXPIRED: ['BRIDGE_DEADLINE_EXCEEDED', false, 504, 'Code environment execution timed out'], ASSIGNMENT_FENCED: ['BRIDGE_ASSIGNMENT_FENCED', false, 409, 'Code environment assignment is fenced; inspect the execution before retrying'], ASSIGNMENT_NOT_FOUND: ['BRIDGE_ASSIGNMENT_NOT_FOUND', false, 409, 'Code environment assignment is no longer available; inspect the execution before retrying'], diff --git a/service/src/sandbox-backend/remote-bridge.ts b/service/src/sandbox-backend/remote-bridge.ts index 719e7004..0bee0038 100644 --- a/service/src/sandbox-backend/remote-bridge.ts +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -18,6 +18,7 @@ const bridgeErrorCodes = { WORKER_OFFLINE: 'BRIDGE_WORKER_OFFLINE', WORKER_UNAUTHORIZED: 'BRIDGE_WORKER_UNAUTHORIZED', WORKER_BUSY: 'BRIDGE_WORKER_BUSY', + WORKER_QUEUE_FULL: 'BRIDGE_WORKER_BUSY', ASSIGNMENT_EXPIRED: 'BRIDGE_DEADLINE_EXCEEDED', ASSIGNMENT_FENCED: 'BRIDGE_ASSIGNMENT_FENCED', ASSIGNMENT_NOT_FOUND: 'BRIDGE_ASSIGNMENT_NOT_FOUND', diff --git a/service/src/workspace-tools/router.test.ts b/service/src/workspace-tools/router.test.ts index 21aa86c7..b5c38ef8 100644 --- a/service/src/workspace-tools/router.test.ts +++ b/service/src/workspace-tools/router.test.ts @@ -32,6 +32,7 @@ afterEach(() => { test('maps invalid worker results to an upstream failure', () => { expect(bridgeStoreStatus(new BridgeStoreError('RESULT_INVALID', 'invalid worker result'))).toBe(502); + expect(bridgeStoreStatus(new BridgeStoreError('WORKER_QUEUE_FULL', 'queue full'))).toBe(429); }); test('rejects new workspace dispatches while the service is shutting down', async () => { diff --git a/service/src/workspace-tools/router.ts b/service/src/workspace-tools/router.ts index b2b56e06..93d15890 100644 --- a/service/src/workspace-tools/router.ts +++ b/service/src/workspace-tools/router.ts @@ -31,6 +31,7 @@ function asyncRoute(handler: (req: AuthenticatedRequest, res: Response) => Promi } export function bridgeStoreStatus(error: BridgeStoreError): number { + if (error.code === 'WORKER_QUEUE_FULL') return 429; if (error.code === 'WORKER_UNAUTHORIZED') return 403; if (error.code === 'ASSIGNMENT_INVALID') return 400; if (error.code === 'RESULT_INVALID') return 502;