diff --git a/README.md b/README.md index 716941d9..384e013a 100644 --- a/README.md +++ b/README.md @@ -37,7 +37,9 @@ Set `CODEAPI_EXECUTION_PROFILE` consistently on an API deployment and its workers. The default profile keeps the existing `python-queue` and `other-queue`; the stateful profile uses `stateful-python-queue` and `stateful-other-queue`. This allows both deployments to share Redis without -cross-consuming jobs. +cross-consuming jobs. The `remote-bridge` backend additionally uses +`remote-bridge-python-queue` and `remote-bridge-other-queue`, fencing attached +worker jobs from Lambda consumers during rolling deployments. An existing Lambda MicroVM deployment upgraded from a pre-profile release may leave `CODEAPI_EXECUTION_PROFILE` unset for its first binary rollout. An diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index e5bddc91..649a8257 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -32,6 +32,23 @@ session hint. In hardened mode, startup requires the bridge token to be at least 32 bytes. `PTC_MODE=blocking` is rejected; replay mode is required because a remote execution cannot retain an open Code API process across tool callbacks. +To attach multiple principal-owned workers to one Code API deployment, enable +dynamic routing. A compatibility default worker is optional in this mode: + +```dotenv +CODEAPI_BRIDGE_DYNAMIC_WORKERS=true +CODEAPI_BRIDGE_AUTH_MODE=paired +# CODEAPI_BRIDGE_WORKER_ID=my-default-vm +``` + +Dynamic routing is accepted only with paired authentication. LibreChat signs +the selected worker into the short-lived Code API JWT as `code_worker_id`. +`X-LibreChat-Code-Worker-ID` remains the transport header, but Code API accepts +it only when it exactly matches that authenticated claim. The resolved worker +is persisted across the queue and programmatic replay boundaries, and Code API +requires both its stored tenant binding and registered worker credential before +creating a lease. + Create a single-use pairing code with the administrator secret: ```bash @@ -41,6 +58,28 @@ curl -fsS https://code.example.com/v1/bridge/pairings \ --data '{"workerId":"my-vm"}' ``` +With dynamic routing enabled, the trusted control plane must bind each pairing +to one tenant and generic principal. Code API treats the principal as lifecycle +and audit metadata; LibreChat remains responsible for resolving user, role, and +group membership before selecting the worker: + +```bash +curl -fsS https://code.example.com/v1/bridge/pairings \ + -H "Authorization: Bearer $CODEAPI_BRIDGE_TOKEN" \ + -H 'Content-Type: application/json' \ + --data '{ + "workerId":"user-vm", + "binding":{ + "tenantId":"tenant-1", + "principal":{"type":"user","id":"user-1"} + } + }' +``` + +Principal types are `deployment`, `tenant`, `user`, `role`, and `group`. +Pairing and registration bodies from the VM cannot replace the server-issued +binding, and credential rotation preserves it. + Redeem the returned code on the VM using [`@librechat/code`](../../packages/code/README.md). The CLI generates its key locally, proves possession on every request, and rotates its short-lived @@ -84,7 +123,14 @@ execution. digest, timestamp, nonce, and credential. - Accepted proof nonces cannot be replayed, credentials rotate before expiry, and an administrator can revoke the active worker identity immediately. -- Code API permits one active assignment per configured worker. +- Assignment leases bind to a stable paired identity rather than an individual + short-lived credential. Rotation preserves that identity; pairing again + replaces it and fences work queued for the previous owner. +- Remote bridge deployments use backend-specific BullMQ queues and serialize + the expected backend on every new job, preventing Lambda or HTTP consumers + from accepting attached-worker executions. +- Code API permits one active assignment per worker. +- Dynamic workers are fenced to their server-issued tenant before assignment. - Each assignment has an absolute deadline, generation, and random lease token. - Settlements with the wrong worker, generation, token, or expired deadline are rejected. @@ -104,6 +150,6 @@ administrator token immediately. Pairing secures worker transport identity; it cannot attest that a compromised VM truthfully reports or enforces its sandbox capabilities. -The next control-plane layer can add owner-scoped environment records and a -multi-worker directory without changing the execution protocol or moving code -tools into the Agents SDK. +LibreChat's owner-scoped environment registry can issue these principal-bound +pairings without changing the worker execution protocol or moving code tools +into the Agents SDK. diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 75eb78c6..e4753c14 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -413,7 +413,7 @@ test('paired worker refreshes before an assignment that outlives its credential' identity: { privateKey: key.privateKey, credential: 'credential-too-short-for-assignment', - expiresAt: new Date(Date.now() + 90_000).toISOString(), + expiresAt: new Date(Date.now() + 30_000).toISOString(), }, capabilities: { statefulWorkspace: true, @@ -469,6 +469,88 @@ test('worker shutdown interrupts reconnect backoff', async () => { await run; }); +test('sandbox completion does not cancel an in-flight credential rotation', async () => { + const key = createBridgeIdentity(); + const identity = { + privateKey: key.privateKey, + credential: 'credential-before-in-flight-rotation', + expiresAt: new Date(Date.now() + 40).toISOString(), + }; + let refreshStarted!: () => void; + const refreshStartedPromise = new Promise((resolve) => { + refreshStarted = resolve; + }); + let refreshCount = 0; + let settleAuthorization = ''; + const fetchImpl: typeof fetch = async (input, init) => { + const url = String(input); + if (url.endsWith('/credentials/refresh')) { + refreshCount += 1; + if (refreshCount > 1) { + return Response.json({ error: 'stale credential' }, { status: 401 }); + } + refreshStarted(); + await new Promise((resolve, reject) => { + const timer = setTimeout(resolve, 30); + init?.signal?.addEventListener( + 'abort', + () => { + clearTimeout(timer); + reject(new DOMException('Aborted', 'AbortError')); + }, + { once: true }, + ); + }); + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'credential-after-in-flight-rotation', + expiresAt: new Date(Date.now() + 300_000).toISOString(), + }); + } + if (url.endsWith('/execute')) { + await refreshStartedPromise; + return Response.json({ session_id: 'run-rotation-race', files: [] }); + } + settleAuthorization = ( + init?.headers as Record + ).Authorization; + return Response.json({ protocolVersion: 1, accepted: true }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + credentialRefreshWindowMs: 30, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-rotation-race', + workerId: 'vm-1', + incarnationId, + generation: 5, + leaseToken: 'assignment-rotation-race-lease-token', + expiresAt: new Date(Date.now() + 600_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(refreshCount, 1); + assert.equal(identity.credential, 'credential-after-in-flight-rotation'); + assert.equal( + settleAuthorization, + 'Bridge credential-after-in-flight-rotation', + ); +}); + 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); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 5e528d61..0f22d22b 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -32,6 +32,7 @@ export interface BridgeWorkerOptions { onError?: (error: unknown) => void; onIdentityChange?: (identity: BridgeWorkerIdentity) => void | Promise; incarnationId?: string; + credentialRefreshWindowMs?: number; } export interface BridgeWorkerIdentity { @@ -181,13 +182,14 @@ export class BridgeWorker { async refreshCredential( signal?: AbortSignal, - validThroughMs = Date.now() + CREDENTIAL_REFRESH_WINDOW_MS, + validThroughMs = + Date.now() + + (this.options.credentialRefreshWindowMs ?? + CREDENTIAL_REFRESH_WINDOW_MS), ): Promise { const identity = this.options.identity; if (identity == null) return; - if ( - Date.parse(identity.expiresAt) > validThroughMs - ) { + if (Date.parse(identity.expiresAt) > validThroughMs) { return; } const credential = await this.request( @@ -201,7 +203,8 @@ export class BridgeWorker { credential.workerId !== this.options.workerId || typeof credential.credential !== 'string' || credential.credential.length < 32 || - !Number.isFinite(Date.parse(credential.expiresAt)) + !Number.isFinite(Date.parse(credential.expiresAt)) || + Date.parse(credential.expiresAt) <= validThroughMs ) { throw new BridgeProtocolError( 'Code API returned an invalid rotated worker credential', @@ -217,16 +220,55 @@ export class BridgeWorker { identity.expiresAt = rotatedIdentity.expiresAt; } + private async maintainCredential( + assignment: BridgeAssignment, + stopSignal: AbortSignal, + requestSignal?: AbortSignal, + ): Promise { + const identity = this.options.identity; + if (identity == null) return; + const refreshWindowMs = + this.options.credentialRefreshWindowMs ?? + CREDENTIAL_REFRESH_WINDOW_MS; + const assignmentDeadlineMs = Date.parse(assignment.expiresAt); + while (!stopSignal.aborted && Date.now() < assignmentDeadlineMs) { + const refreshAtMs = Date.parse(identity.expiresAt) - refreshWindowMs; + const waitMs = Math.max( + 0, + Math.min(refreshAtMs - Date.now(), assignmentDeadlineMs - Date.now()), + ); + if (waitMs > 0) { + await new Promise((resolve) => { + const timer = setTimeout(resolve, waitMs); + stopSignal.addEventListener( + 'abort', + () => { + clearTimeout(timer); + resolve(); + }, + { once: true }, + ); + }); + } + if (stopSignal.aborted || Date.now() >= assignmentDeadlineMs) return; + /* Sandbox completion stops the maintenance loop, but it must not abort a + * refresh already accepted by the server. Only cancellation of the + * whole worker operation may cancel that request. */ + await this.refreshCredential(requestSignal); + } + } + async executeAndSettle( assignment: BridgeAssignment, signal?: AbortSignal, ): Promise { - await this.refreshCredential( - signal, - Date.parse(assignment.expiresAt) + CREDENTIAL_REFRESH_WINDOW_MS, - ); + await this.refreshCredential(signal); const executionController = new AbortController(); - const abortExecution = (): void => executionController.abort(); + const credentialController = new AbortController(); + const abortExecution = (): void => { + executionController.abort(); + credentialController.abort(); + }; signal?.addEventListener('abort', abortExecution, { once: true }); const deadlineDelay = Math.max( 0, @@ -251,8 +293,18 @@ export class BridgeWorker { executionController, cancellationController.signal, ); + let credentialMaintenanceError: unknown; + let credentialMaintenance: Promise | undefined; let settlement: BridgeSettlement; try { + credentialMaintenance = this.maintainCredential( + assignment, + credentialController.signal, + signal, + ).catch((error) => { + credentialMaintenanceError = error; + executionController.abort(); + }); const headers = { ...assignment.request.headers, ...(assignment.runtimeSessionId @@ -273,6 +325,9 @@ export class BridgeWorker { ); const payload = (await response.json()) as object; if (heartbeatError != null) throw heartbeatError; + if (credentialMaintenanceError != null) { + throw credentialMaintenanceError; + } if (!response.ok) { throw new BridgeProtocolError( errorMessage(payload) ?? @@ -303,9 +358,12 @@ export class BridgeWorker { clearTimeout(deadlineTimer); heartbeatController.abort(); await heartbeat; + credentialController.abort(); + await credentialMaintenance; cancellationController.abort(); await cancellationWatcher; signal?.removeEventListener('abort', abortExecution); + await this.refreshCredential(signal); await this.request( this.assignmentUrl(assignment, 'settle'), settlement, diff --git a/service/src/auth/librechat-jwt.test.ts b/service/src/auth/librechat-jwt.test.ts index 2030b2e7..d20124f5 100644 --- a/service/src/auth/librechat-jwt.test.ts +++ b/service/src/auth/librechat-jwt.test.ts @@ -48,6 +48,7 @@ type JwtClaims = { chc_user_id?: string; auth_context_hash?: string; plan_id?: string; + code_worker_id?: string; }; const originalEnv = new Map(); @@ -75,6 +76,7 @@ function baseClaims(overrides: Partial = {}): JwtClaims { external_user_id: 'chc_123', auth_context_hash: 'hash_123', plan_id: 'prod_plan_123', + code_worker_id: 'code-user_123', ...overrides, }; } @@ -166,6 +168,7 @@ describe('LibreChat JWT auth provider', () => { principalSource: 'openid_reuse', authContextHash: 'hash_123', planId: 'prod_plan_123', + codeWorkerId: 'code-user_123', }); }); diff --git a/service/src/auth/librechat-jwt.ts b/service/src/auth/librechat-jwt.ts index 1e7e8079..e249ce0c 100644 --- a/service/src/auth/librechat-jwt.ts +++ b/service/src/auth/librechat-jwt.ts @@ -39,6 +39,7 @@ interface LibreChatJwtClaims { chc_user_id?: string; // leak-check:allow auth_context_hash?: string; plan_id?: string; + code_worker_id?: string; } interface PublicKeyEntry { @@ -394,6 +395,7 @@ function validateClaims(claims: LibreChatJwtClaims, config: VerificationConfig): const nbf = assertNumericDate(claims.nbf, 'nbf'); const exp = assertNumericDate(claims.exp, 'exp'); const planId = optionalString(claims.plan_id, 'plan_id'); + const codeWorkerId = optionalString(claims.code_worker_id, 'code_worker_id'); const principalSource = assertPrincipalSource(claims.principal_source); const authContextHash = assertString(claims.auth_context_hash, 'auth_context_hash'); @@ -433,6 +435,7 @@ function validateClaims(claims: LibreChatJwtClaims, config: VerificationConfig): principalSource, authContextHash, planId, + codeWorkerId, }; } diff --git a/service/src/auth/principal.ts b/service/src/auth/principal.ts index 94b785f3..615a0ff2 100644 --- a/service/src/auth/principal.ts +++ b/service/src/auth/principal.ts @@ -13,6 +13,7 @@ export type CodeApiPrincipal = { authContextHash?: string; credentialId?: string; planId?: string; + codeWorkerId?: string; }; export function applyPrincipal(req: t.AuthenticatedRequest, principal: CodeApiPrincipal): void { diff --git a/service/src/bridge/index.ts b/service/src/bridge/index.ts index dfba2a8b..08ef1b07 100644 --- a/service/src/bridge/index.ts +++ b/service/src/bridge/index.ts @@ -13,4 +13,5 @@ export default createBridgeRouter({ authMode: env.BRIDGE_AUTH_MODE, adminToken: env.BRIDGE_TOKEN, configuredWorkerId: env.BRIDGE_WORKER_ID, + allowDynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, }); diff --git a/service/src/bridge/pairing.test.ts b/service/src/bridge/pairing.test.ts index 1366d478..310ad7d9 100644 --- a/service/src/bridge/pairing.test.ts +++ b/service/src/bridge/pairing.test.ts @@ -1,4 +1,5 @@ import { afterEach, describe, expect, test } from 'bun:test'; +import { createHash } from 'crypto'; import RedisMock from 'ioredis-mock'; import type Redis from 'ioredis'; @@ -17,6 +18,94 @@ afterEach(async () => { }); describe('RedisBridgePairingStore', () => { + test('preserves a tenant and generic principal binding across credential rotation', async () => { + const identity = createBridgeIdentity(); + const binding = { + tenantId: 'tenant-1', + principal: { type: 'group' as const, id: 'engineering' }, + }; + const pairing = await pairings.issue('vm-bound', binding); + const issued = await pairings.redeem({ + workerId: 'vm-bound', + code: pairing.code, + publicKey: identity.publicKey, + }); + const requestFor = ( + credential: string, + nonce: string, + ): Parameters[0] => { + const proof = { + credential, + method: 'POST', + path: '/v1/bridge/workers/vm-bound/lease', + timestamp: new Date().toISOString(), + nonce, + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + return { + ...proof, + workerId: 'vm-bound', + signature: signBridgeRequest(identity.privateKey, proof), + }; + }; + + const originalAuthorization = await pairings.authorize( + requestFor(issued.credential, 'original-bound-worker-proof'), + ); + const rotated = await pairings.rotate('vm-bound'); + + const rotatedAuthorization = await pairings.authorize( + requestFor(rotated.credential, 'bound-worker-proof'), + ); + expect(rotatedAuthorization).toMatchObject({ workerId: 'vm-bound', binding }); + expect(typeof originalAuthorization.identityId).toBe('string'); + expect(rotatedAuthorization.identityId).toBe( + originalAuthorization.identityId, + ); + await expect( + pairings.authorize(requestFor(issued.credential, 'overlap-bound-proof')), + ).resolves.toMatchObject({ + workerId: 'vm-bound', + identityId: originalAuthorization.identityId, + }); + }); + + test('preserves a legacy unmarked identity across its first rotation', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('legacy-vm'); + const issued = await pairings.redeem({ + workerId: 'legacy-vm', + code: pairing.code, + publicKey: identity.publicKey, + }); + const issuedDigest = createHash('sha256') + .update(issued.credential) + .digest('hex'); + const credentialKey = `codeapi:bridge:v1:credential:${issuedDigest}`; + const stored = JSON.parse((await redis.get(credentialKey)) ?? '{}') as { + identityId?: string; + }; + delete stored.identityId; + await redis.set(credentialKey, JSON.stringify(stored), 'EX', 300); + + const rotated = await pairings.rotate('legacy-vm'); + const proof = { + credential: rotated.credential, + method: 'POST', + path: '/v1/bridge/workers/legacy-vm/lease', + timestamp: new Date().toISOString(), + nonce: 'legacy-rotation-proof', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + const authorization = await pairings.authorize({ + ...proof, + workerId: 'legacy-vm', + signature: signBridgeRequest(identity.privateKey, proof), + }); + + expect(authorization.identityId).toBeUndefined(); + }); + test('redeems a pairing code exactly once for the intended worker identity', async () => { const identity = createBridgeIdentity(); const pairing = await pairings.issue('vm-1'); @@ -38,6 +127,108 @@ describe('RedisBridgePairingStore', () => { ).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); }); + test('only the newest pairing code can rebind a worker identity', async () => { + const identity = createBridgeIdentity(); + const older = await pairings.issue('vm-1', { + tenantId: 'tenant-a', + principal: { type: 'user', id: 'user-a' }, + }); + const newer = await pairings.issue('vm-1', { + tenantId: 'tenant-b', + principal: { type: 'user', id: 'user-b' }, + }); + + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: older.code, + publicKey: identity.publicKey, + }), + ).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: newer.code, + publicKey: identity.publicKey, + }), + ).resolves.toMatchObject({ workerId: 'vm-1' }); + }); + + test('does not let a paused redemption overwrite a newer pairing identity', async () => { + const firstIdentity = createBridgeIdentity(); + const secondIdentity = createBridgeIdentity(); + const firstPairing = await pairings.issue('vm-race', { + tenantId: 'tenant-a', + principal: { type: 'user', id: 'user-a' }, + }); + const originalEval = redis.eval.bind(redis); + let releaseFirst!: () => void; + let firstRedeemed!: () => void; + const firstRedeemedPromise = new Promise((resolve) => { + firstRedeemed = resolve; + }); + const releaseFirstPromise = new Promise((resolve) => { + releaseFirst = resolve; + }); + let paused = false; + redis.eval = (async (script: string, ...args: unknown[]) => { + const result = await (originalEval as (...evalArgs: unknown[]) => Promise)( + script, + ...args, + ); + if (!paused && script.includes('return pairing')) { + paused = true; + firstRedeemed(); + await releaseFirstPromise; + } + return result; + }) as typeof redis.eval; + + try { + const staleRedemption = pairings.redeem({ + workerId: 'vm-race', + code: firstPairing.code, + publicKey: firstIdentity.publicKey, + }); + await firstRedeemedPromise; + const secondPairing = await pairings.issue('vm-race', { + tenantId: 'tenant-b', + principal: { type: 'user', id: 'user-b' }, + }); + const current = await pairings.redeem({ + workerId: 'vm-race', + code: secondPairing.code, + publicKey: secondIdentity.publicKey, + }); + releaseFirst(); + + await expect(staleRedemption).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + const proof = { + credential: current.credential, + method: 'POST', + path: '/v1/bridge/workers/vm-race/lease', + timestamp: new Date().toISOString(), + nonce: 'current-race-proof', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + await expect( + pairings.authorize({ + ...proof, + workerId: 'vm-race', + signature: signBridgeRequest(secondIdentity.privateKey, proof), + }), + ).resolves.toMatchObject({ + binding: { + tenantId: 'tenant-b', + principal: { type: 'user', id: 'user-b' }, + }, + }); + } finally { + redis.eval = originalEval as typeof redis.eval; + releaseFirst(); + } + }); + test('authorizes a credential only with proof from its worker key', async () => { const identity = createBridgeIdentity(); const pairing = await pairings.issue('vm-1'); @@ -226,6 +417,41 @@ describe('RedisBridgePairingStore', () => { ).resolves.toMatchObject({ workerId: 'vm-1' }); }); + test('rejects a stale credential refresh after the worker is paired again', async () => { + const originalIdentity = createBridgeIdentity(); + const originalPairing = await pairings.issue('vm-1'); + const original = await pairings.redeem({ + workerId: 'vm-1', + code: originalPairing.code, + publicKey: originalIdentity.publicKey, + }); + const proof = { + credential: original.credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/credentials/refresh', + timestamp: new Date().toISOString(), + nonce: 'authorized-before-repairing', + body: JSON.stringify({ protocolVersion: 1 }), + }; + const staleAuthorization = await pairings.authorize({ + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(originalIdentity.privateKey, proof), + }); + + const replacementIdentity = createBridgeIdentity(); + const replacementPairing = await pairings.issue('vm-1'); + await pairings.redeem({ + workerId: 'vm-1', + code: replacementPairing.code, + publicKey: replacementIdentity.publicKey, + }); + + await expect( + pairings.rotate('vm-1', staleAuthorization.credentialId), + ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); + }); + test('repairing a worker invalidates its previously paired credential', async () => { const firstIdentity = createBridgeIdentity(); const firstPairing = await pairings.issue('vm-1'); diff --git a/service/src/bridge/pairing.ts b/service/src/bridge/pairing.ts index 54aaf5ab..fd1e8cff 100644 --- a/service/src/bridge/pairing.ts +++ b/service/src/bridge/pairing.ts @@ -19,25 +19,80 @@ local previous = redis.call('GET', KEYS[2]) if not activeDigest or not previous then return 0 end -if activeDigest ~= ARGV[1] and redis.call('GET', KEYS[4]) ~= ARGV[5] then - return 0 +if activeDigest ~= ARGV[1] then + if ARGV[5] == '' or redis.call('GET', KEYS[4]) ~= ARGV[5] then + return 0 + end end redis.call('SET', KEYS[3], ARGV[3], 'EX', ARGV[4]) redis.call('SET', KEYS[1], ARGV[2], 'EX', ARGV[4]) -redis.call('SET', KEYS[4], ARGV[5], 'EX', ARGV[4]) +if ARGV[5] ~= '' then + redis.call('SET', KEYS[4], ARGV[5], 'EX', ARGV[4]) +else + redis.call('DEL', KEYS[4]) +end +return 1 +`; +const ISSUE_PAIRING_SCRIPT = ` +local previous = redis.call('GET', KEYS[1]) +if previous then + redis.call('DEL', previous) +end +redis.call('DEL', KEYS[3]) +redis.call('SET', KEYS[1], KEYS[2], 'EX', ARGV[2]) +redis.call('SET', KEYS[2], ARGV[1], 'EX', ARGV[2]) return 1 `; +const REDEEM_PAIRING_SCRIPT = ` +local pairing = redis.call('GET', KEYS[1]) +if not pairing then + return nil +end +if redis.call('GET', KEYS[2]) ~= KEYS[1] then + redis.call('DEL', KEYS[1]) + return nil +end +redis.call('DEL', KEYS[1], KEYS[2]) +redis.call('SET', KEYS[3], ARGV[1], 'EX', ARGV[2]) +return pairing +`; +const INSTALL_REDEEMED_CREDENTIAL_SCRIPT = ` +if redis.call('GET', KEYS[1]) ~= ARGV[1] then + return 0 +end +redis.call('SET', KEYS[2], ARGV[3], 'EX', ARGV[4]) +redis.call('SET', KEYS[3], ARGV[2], 'EX', ARGV[4]) +if ARGV[5] ~= '' then + redis.call('SET', KEYS[4], ARGV[5], 'EX', ARGV[4]) +else + redis.call('DEL', KEYS[4]) +end +redis.call('DEL', KEYS[1]) +return 1 +`; + +export type BridgePrincipalType = 'deployment' | 'tenant' | 'user' | 'role' | 'group'; + +export interface BridgeWorkerBinding { + tenantId: string; + principal: { + type: BridgePrincipalType; + id: string; + }; +} interface StoredPairing { workerId: string; expiresAt: string; + binding?: BridgeWorkerBinding; } interface StoredCredential { workerId: string; - identityId: string; + identityId?: string; publicKey: string; expiresAt: string; + binding?: BridgeWorkerBinding; } export interface BridgePairing { @@ -75,10 +130,6 @@ function pairingKey(code: string): string { return `${PREFIX}:pairing:${digest(code)}`; } -function credentialKey(credential: string): string { - return credentialDigestKey(digest(credential)); -} - function credentialDigestKey(credentialDigest: string): string { return `${PREFIX}:credential:${credentialDigest}`; } @@ -91,6 +142,14 @@ function workerStableIdentityKey(workerId: string): string { return `${PREFIX}:stable-identity:${workerId}`; } +function workerPairingIndexKey(workerId: string): string { + return `${PREFIX}:pairing-index:${workerId}`; +} + +function workerRedemptionKey(workerId: string): string { + return `${PREFIX}:redemption:${workerId}`; +} + function proofNonceKey(credential: string, nonce: string): string { return `${PREFIX}:proof:${digest(credential)}:${digest(nonce)}`; } @@ -110,17 +169,24 @@ export class RedisBridgePairingStore { private readonly credentialTtlSeconds = DEFAULT_CREDENTIAL_TTL_SECONDS, ) {} - async issue(workerId: string): Promise { + async issue( + workerId: string, + binding?: BridgeWorkerBinding, + ): Promise { const code = randomBytes(24).toString('base64url'); const expiresAt = new Date( Date.now() + this.pairingTtlSeconds * 1000, ).toISOString(); - const pairing: StoredPairing = { workerId, expiresAt }; - await this.redis.set( - pairingKey(code), + const pairing: StoredPairing = { workerId, expiresAt, binding }; + const codeKey = pairingKey(code); + await this.redis.eval( + ISSUE_PAIRING_SCRIPT, + 3, + workerPairingIndexKey(workerId), + codeKey, + workerRedemptionKey(workerId), JSON.stringify(pairing), - 'EX', - this.pairingTtlSeconds, + String(this.pairingTtlSeconds), ); return { workerId, code, expiresAt }; } @@ -130,8 +196,24 @@ export class RedisBridgePairingStore { code: string; publicKey: string; }): Promise { - const raw = await this.redis.getdel(pairingKey(args.code)); - if (raw == null) { + if (!validEd25519PublicKey(args.publicKey)) { + throw new BridgePairingError( + 'PUBLIC_KEY_INVALID', + 'Worker public key must be an Ed25519 key', + ); + } + const codeKey = pairingKey(args.code); + const redemptionId = randomBytes(18).toString('base64url'); + const raw = await this.redis.eval( + REDEEM_PAIRING_SCRIPT, + 3, + codeKey, + workerPairingIndexKey(args.workerId), + workerRedemptionKey(args.workerId), + redemptionId, + String(this.pairingTtlSeconds), + ); + if (typeof raw !== 'string') { throw new BridgePairingError( 'PAIRING_INVALID', 'Pairing code is invalid or expired', @@ -144,14 +226,14 @@ export class RedisBridgePairingStore { 'Pairing code does not authorize this worker', ); } - if (!validEd25519PublicKey(args.publicKey)) { - throw new BridgePairingError( - 'PUBLIC_KEY_INVALID', - 'Worker public key must be an Ed25519 key', - ); - } - - return await this.issueCredential(args.workerId, args.publicKey); + return await this.issueCredential( + args.workerId, + args.publicKey, + undefined, + undefined, + pairing.binding, + redemptionId, + ); } async authorize(args: { @@ -167,7 +249,8 @@ export class RedisBridgePairingStore { workerId: string; credentialId: string; activeCredentialId: string; - identityId: string; + identityId?: string; + binding?: BridgeWorkerBinding; }> { const proofTime = Date.parse(args.timestamp); if ( @@ -198,7 +281,11 @@ export class RedisBridgePairingStore { const active = activeRaw == null ? undefined : JSON.parse(activeRaw) as StoredCredential; - if (active?.identityId !== stored.identityId) { + if ( + stored.identityId == null || + active?.identityId == null || + stored.identityId !== active.identityId + ) { throw new BridgePairingError( 'CREDENTIAL_INVALID', 'Worker credential is invalid or expired', @@ -234,7 +321,8 @@ export class RedisBridgePairingStore { workerId: stored.workerId, credentialId: credentialDigest, activeCredentialId: activeDigest, - identityId: stored.identityId, + ...(stored.identityId != null ? { identityId: stored.identityId } : {}), + ...(stored.binding ? { binding: stored.binding } : {}), }; } @@ -254,7 +342,8 @@ export class RedisBridgePairingStore { expectedCredentialId?: string, ): Promise { const identityKey = workerIdentityKey(workerId); - const previousDigest = expectedCredentialId ?? await this.redis.get(identityKey); + const previousDigest = + expectedCredentialId ?? (await this.redis.get(identityKey)); const previousRaw = previousDigest == null ? null @@ -270,7 +359,8 @@ export class RedisBridgePairingStore { workerId, previous.publicKey, previousDigest, - previous.identityId, + previous.identityId ?? null, + previous.binding, ); } @@ -278,14 +368,23 @@ export class RedisBridgePairingStore { workerId: string, publicKey: string, previousDigest?: string, - identityId = randomBytes(18).toString('base64url'), + identityId: string | null | undefined = randomBytes(18).toString('base64url'), + binding?: BridgeWorkerBinding, + redemptionId?: string, ): Promise { const credential = randomBytes(32).toString('base64url'); const credentialDigest = digest(credential); const expiresAt = new Date( Date.now() + this.credentialTtlSeconds * 1000, ).toISOString(); - const stored: StoredCredential = { workerId, identityId, publicKey, expiresAt }; + const stableIdentityId = identityId ?? undefined; + const stored: StoredCredential = { + workerId, + ...(stableIdentityId != null ? { identityId: stableIdentityId } : {}), + publicKey, + expiresAt, + binding, + }; if (previousDigest !== undefined) { const rotated = await this.redis.eval( ROTATE_CREDENTIAL_SCRIPT, @@ -298,7 +397,7 @@ export class RedisBridgePairingStore { credentialDigest, JSON.stringify(stored), String(this.credentialTtlSeconds), - identityId, + stableIdentityId ?? '', ); if (rotated !== 1) { throw new BridgePairingError( @@ -308,26 +407,31 @@ export class RedisBridgePairingStore { } return { workerId, credential, expiresAt }; } - const transaction = this.redis.multi(); - transaction.set( + if (redemptionId == null) { + throw new BridgePairingError( + 'PAIRING_INVALID', + 'Pairing redemption was not fenced', + ); + } + const installed = await this.redis.eval( + INSTALL_REDEEMED_CREDENTIAL_SCRIPT, + 4, + workerRedemptionKey(workerId), credentialDigestKey(credentialDigest), - JSON.stringify(stored), - 'EX', - this.credentialTtlSeconds, - ); - transaction.set( workerIdentityKey(workerId), - credentialDigest, - 'EX', - this.credentialTtlSeconds, - ); - transaction.set( workerStableIdentityKey(workerId), - identityId, - 'EX', - this.credentialTtlSeconds, + redemptionId, + credentialDigest, + JSON.stringify(stored), + String(this.credentialTtlSeconds), + stableIdentityId ?? '', ); - await transaction.exec(); + if (installed !== 1) { + throw new BridgePairingError( + 'PAIRING_INVALID', + 'Pairing code was superseded before credential installation', + ); + } return { workerId, credential, expiresAt }; } } diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts index 077dcd1f..bb59e991 100644 --- a/service/src/bridge/router.test.ts +++ b/service/src/bridge/router.test.ts @@ -25,6 +25,171 @@ afterEach(async () => { }); describe('paired bridge HTTP API', () => { + test('rejects a malformed optional binding for a configured worker', async () => { + const app = express(); + app.use(json()); + app.use( + '/v1/bridge', + createBridgeRouter({ + store: new RedisBridgeStore(redis), + pairings: new RedisBridgePairingStore(redis), + authMode: 'paired', + adminToken: 'strong-administrator-bootstrap-token', + configuredWorkerId: 'vm-1', + }), + ); + server = createServer(app); + await new Promise((resolve) => server?.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (address == null || typeof address === 'string') { + throw new Error('Expected TCP listener'); + } + + const response = await fetch( + `http://127.0.0.1:${address.port}/v1/bridge/pairings`, + { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + workerId: 'vm-1', + binding: { tenantId: 'tenant-1', principal: { type: 'user' } }, + }), + }, + ); + + expect(response.status).toBe(400); + }); + + test('requires and persists a trusted principal binding for dynamic workers', async () => { + const store = new RedisBridgeStore(redis); + const app = express(); + app.use(json()); + app.use( + '/v1/bridge', + createBridgeRouter({ + store, + pairings: new RedisBridgePairingStore(redis), + authMode: 'paired', + adminToken: 'strong-administrator-bootstrap-token', + allowDynamicWorkers: true, + }), + ); + server = createServer(app); + await new Promise((resolve) => server?.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (address == null || typeof address === 'string') { + throw new Error('Expected TCP listener'); + } + const baseUrl = `http://127.0.0.1:${address.port}/v1/bridge`; + const unboundResponse = await fetch(`${baseUrl}/pairings`, { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ workerId: 'user-vm' }), + }); + expect(unboundResponse.status).toBe(400); + + const binding = { + tenantId: 'tenant-1', + principal: { type: 'user' as const, id: 'user-1' }, + }; + const pairingResponse = await fetch(`${baseUrl}/pairings`, { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ workerId: 'user-vm', binding }), + }); + const pairing = (await pairingResponse.json()) as { code: string }; + expect(pairingResponse.status).toBe(200); + + const identity = createBridgeIdentity(); + const redemptionResponse = await fetch(`${baseUrl}/pairings/redeem`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'user-vm', + code: pairing.code, + publicKey: identity.publicKey, + }), + }); + const issued = (await redemptionResponse.json()) as { credential: string }; + expect(redemptionResponse.status).toBe(200); + + const path = '/v1/bridge/workers/register'; + const body = JSON.stringify({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'user-vm', + incarnationId: 'incarnation-00000001', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + binding: { + tenantId: 'tenant-2', + principal: { type: 'user', id: 'attacker-selected-user' }, + }, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path, + timestamp: new Date().toISOString(), + nonce: 'dynamic-registration-nonce', + body, + }; + const registrationResponse = await fetch( + `http://127.0.0.1:${address.port}${path}`, + { + method: 'POST', + headers: { + Authorization: `Bridge ${issued.credential}`, + 'Content-Type': 'application/json', + 'X-LibreChat-Code-Timestamp': proof.timestamp, + 'X-LibreChat-Code-Nonce': proof.nonce, + 'X-LibreChat-Code-Signature': signBridgeRequest( + identity.privateKey, + proof, + ), + }, + body, + }, + ); + expect(registrationResponse.status).toBe(200); + await expect( + store.dispatch({ + workerId: 'user-vm', + tenantId: binding.tenantId, + requireTenantBinding: true, + body: { language: 'bash' } as never, + headers: {}, + runtimeSessionId: 'stateful-session', + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_MISMATCH' }); + + await expect( + store.dispatch({ + workerId: 'user-vm', + tenantId: 'tenant-2', + requireTenantBinding: true, + body: { language: 'bash' } as never, + headers: {}, + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_UNAUTHORIZED' }); + }); + test('pairs a worker and accepts its proof-of-possession registration', async () => { const app = express(); app.use(json()); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 8ce58bfe..0c2771b1 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -4,15 +4,24 @@ import { Router } from 'express'; import type { NextFunction, Request, Response } from 'express'; import type { BridgeWorkerRegistration } from '../../../packages/code/src/protocol'; +import type { BridgeWorkerBinding, BridgePrincipalType } from './pairing'; import type { CodeBridgeSettlement } from './store'; import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; import { BridgePairingError, RedisBridgePairingStore } from './pairing'; +import { BRIDGE_WORKER_ID_PATTERN } from './selection'; import { BridgeStoreError, RedisBridgeStore } from './store'; -const WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/; const INCARNATION_ID_PATTERN = /^[A-Za-z0-9_-]{16,128}$/; const MAX_LEASE_WAIT_MS = 30_000; +const BRIDGE_BINDING_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/; +const PRINCIPAL_TYPES = new Set([ + 'deployment', + 'tenant', + 'user', + 'role', + 'group', +]); export type BridgeAuthMode = 'static' | 'paired'; @@ -22,6 +31,7 @@ export interface BridgeRouterOptions { authMode: BridgeAuthMode; adminToken: string; configuredWorkerId?: string; + allowDynamicWorkers?: boolean; } function sameToken(left: string, right: string): boolean { @@ -34,7 +44,7 @@ function sameToken(left: string, right: string): boolean { } function validWorkerId(value: string): boolean { - return WORKER_ID_PATTERN.test(value); + return BRIDGE_WORKER_ID_PATTERN.test(value); } function validIncarnationId(value: unknown): value is string { @@ -45,8 +55,32 @@ function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null; } +function parseBinding(value: unknown): BridgeWorkerBinding | undefined { + if (!isRecord(value) || !isRecord(value.principal)) return undefined; + const { tenantId, principal } = value; + if ( + typeof tenantId !== 'string' || + !BRIDGE_BINDING_ID_PATTERN.test(tenantId) || + typeof principal.type !== 'string' || + !PRINCIPAL_TYPES.has(principal.type as BridgePrincipalType) || + typeof principal.id !== 'string' || + !BRIDGE_BINDING_ID_PATTERN.test(principal.id) + ) { + return undefined; + } + return { + tenantId, + principal: { + type: principal.type as BridgePrincipalType, + id: principal.id, + }, + }; +} + function sendStoreError(error: BridgeStoreError, res: Response): void { - const status = error.code === 'ASSIGNMENT_NOT_FOUND' ? 404 : 409; + let status = 409; + if (error.code === 'ASSIGNMENT_NOT_FOUND') status = 404; + if (error.code === 'WORKER_UNAUTHORIZED') status = 403; res.status(status).json({ error: error.message, code: error.code }); } @@ -77,8 +111,7 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { const router = Router(); const configuredWorker = (workerId: string): boolean => - options.configuredWorkerId == null || - options.configuredWorkerId === '' || + options.allowDynamicWorkers === true || workerId === options.configuredWorkerId; const bearerToken = (req: Request): string => @@ -186,7 +219,18 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { res.status(400).json({ error: 'Invalid bridge worker ID' }); return; } - const pairing = await options.pairings.issue(workerId); + const hasBinding = isRecord(req.body) && + Object.prototype.hasOwnProperty.call(req.body, 'binding'); + const binding = isRecord(req.body) ? parseBinding(req.body.binding) : undefined; + if (hasBinding && binding == null) { + res.status(400).json({ error: 'Invalid bridge worker principal binding' }); + return; + } + if (options.allowDynamicWorkers === true && binding == null) { + res.status(400).json({ error: 'Dynamic bridge workers require a valid principal binding' }); + return; + } + const pairing = await options.pairings.issue(workerId, binding); res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...pairing }); }); @@ -304,10 +348,42 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { }); return; } + const authorization = res.locals.bridgeWorkerAuthorization as + | { + workerId: string; + credentialId: string; + activeCredentialId: string; + identityId?: string; + binding?: BridgeWorkerBinding; + } + | undefined; + const capabilities = registration.capabilities; + const trustedRegistration: BridgeWorkerRegistration = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: registration.workerId, + incarnationId: registration.incarnationId, + capabilities: { + statefulWorkspace: capabilities.statefulWorkspace as boolean, + sandboxProfile: capabilities.sandboxProfile as string, + runtimes: capabilities.runtimes as string[], + ...(typeof capabilities.policyDigest === 'string' + ? { policyDigest: capabilities.policyDigest } + : {}), + }, + }; try { - await options.store.register( - registration as unknown as BridgeWorkerRegistration, - ); + await options.store.register({ + ...trustedRegistration, + ...(authorization?.credentialId != null + ? { credentialId: authorization.credentialId } + : {}), + ...(authorization?.identityId != null + ? { identityId: authorization.identityId } + : {}), + ...(authorization?.binding != null + ? { binding: authorization.binding } + : {}), + }, authorization?.activeCredentialId); } catch (error) { if (error instanceof BridgeStoreError) { sendStoreError(error, res); @@ -353,6 +429,12 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { workerId, body.incarnationId, Math.min(requestedWait, MAX_LEASE_WAIT_MS), + undefined, + ( + res.locals.bridgeWorkerAuthorization as + | { identityId: string } + | undefined + )?.identityId, ); res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); } catch (error) { @@ -379,6 +461,11 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { req.params.workerId, req.params.assignmentId, settlement, + ( + res.locals.bridgeWorkerAuthorization as + | { identityId: string } + | undefined + )?.identityId, ); res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/selection.test.ts b/service/src/bridge/selection.test.ts new file mode 100644 index 00000000..b6309702 --- /dev/null +++ b/service/src/bridge/selection.test.ts @@ -0,0 +1,101 @@ +import { describe, expect, test } from 'bun:test'; + +import { + BridgeWorkerSelectionError, + resolveBridgeWorkerSelection, +} from './selection'; + +describe('bridge worker request selection', () => { + test('uses the configured compatibility worker when no dynamic worker is requested', () => { + expect( + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: true, + }), + ).toEqual({ workerId: 'deployment-worker', explicit: false }); + }); + + test('selects only the worker authenticated by the LibreChat JWT', () => { + expect( + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: true, + requestedWorkerId: 'code-user_1', + trustedWorkerId: 'code-user_1', + }), + ).toEqual({ workerId: 'code-user_1', explicit: true }); + + expect( + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: true, + trustedWorkerId: 'code-user_1', + }), + ).toEqual({ workerId: 'code-user_1', explicit: true }); + }); + + test('rejects a caller-controlled worker header without a matching trusted claim', () => { + expect(() => + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: true, + requestedWorkerId: 'victim-worker', + }), + ).toThrow('Code bridge worker selection is not authenticated'); + expect(() => + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: true, + requestedWorkerId: 'victim-worker', + trustedWorkerId: 'caller-worker', + }), + ).toThrow('Code bridge worker selection does not match the authenticated claim'); + }); + + test('rejects dynamic routing on the wrong backend or when it is disabled', () => { + expect(() => + resolveBridgeWorkerSelection({ + backend: 'http', + configuredWorkerId: '', + dynamicWorkers: true, + requestedWorkerId: 'code-user-1', + trustedWorkerId: 'code-user-1', + }), + ).toThrow(BridgeWorkerSelectionError); + expect(() => + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: false, + requestedWorkerId: 'code-user-1', + trustedWorkerId: 'code-user-1', + }), + ).toThrow('Dynamic code bridge workers are disabled'); + }); + + test('rejects malformed worker IDs before they cross the queue boundary', () => { + expect(() => + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: '', + dynamicWorkers: true, + requestedWorkerId: '../worker', + trustedWorkerId: '../worker', + }), + ).toThrow('Invalid code bridge worker ID'); + expect(() => + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: '', + dynamicWorkers: true, + requestedWorkerId: 'victim:assignments', + trustedWorkerId: 'victim:assignments', + }), + ).toThrow('Invalid code bridge worker ID'); + }); +}); diff --git a/service/src/bridge/selection.ts b/service/src/bridge/selection.ts new file mode 100644 index 00000000..0959279f --- /dev/null +++ b/service/src/bridge/selection.ts @@ -0,0 +1,72 @@ +export const CODEAPI_BRIDGE_WORKER_HEADER = 'X-LibreChat-Code-Worker-ID'; +export const BRIDGE_WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$/; + +export class BridgeWorkerSelectionError extends Error { + constructor( + message: string, + public readonly status: 400 | 403 | 503, + ) { + super(message); + this.name = 'BridgeWorkerSelectionError'; + } +} + +export function resolveBridgeWorkerSelection(args: { + backend: SandboxBackendName; + configuredWorkerId: string; + dynamicWorkers: boolean; + requestedWorkerId?: string; + trustedWorkerId?: string; +}): { workerId: string; explicit: boolean } | undefined { + const requestedWorkerId = args.requestedWorkerId?.trim(); + const trustedWorkerId = args.trustedWorkerId?.trim(); + const hasRequestedWorker = requestedWorkerId != null && requestedWorkerId.length > 0; + const hasTrustedWorker = trustedWorkerId != null && trustedWorkerId.length > 0; + if (hasRequestedWorker || hasTrustedWorker) { + if (args.backend !== 'remote-bridge') { + throw new BridgeWorkerSelectionError( + 'Code bridge worker routing requires the remote-bridge backend', + 400, + ); + } + if (hasRequestedWorker && !hasTrustedWorker) { + throw new BridgeWorkerSelectionError( + 'Code bridge worker selection is not authenticated', + 403, + ); + } + if ( + hasRequestedWorker && + hasTrustedWorker && + requestedWorkerId !== trustedWorkerId + ) { + throw new BridgeWorkerSelectionError( + 'Code bridge worker selection does not match the authenticated claim', + 403, + ); + } + const selectedWorkerId = trustedWorkerId as string; + if (!BRIDGE_WORKER_ID_PATTERN.test(selectedWorkerId)) { + throw new BridgeWorkerSelectionError('Invalid code bridge worker ID', 400); + } + if (!args.dynamicWorkers && selectedWorkerId !== args.configuredWorkerId) { + throw new BridgeWorkerSelectionError('Dynamic code bridge workers are disabled', 403); + } + return { + workerId: selectedWorkerId, + explicit: true, + }; + } + + if (args.backend !== 'remote-bridge') return undefined; + const configuredWorkerId = args.configuredWorkerId.trim(); + if (configuredWorkerId.length === 0) { + throw new BridgeWorkerSelectionError('No code bridge worker was selected', 503); + } + if (!BRIDGE_WORKER_ID_PATTERN.test(configuredWorkerId)) { + throw new BridgeWorkerSelectionError('Invalid configured code bridge worker ID', 503); + } + return { workerId: configuredWorkerId, explicit: false }; +} + +type SandboxBackendName = 'http' | 'lambda-microvm' | 'remote-bridge'; diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index c35a8fd8..684381ad 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -14,6 +14,400 @@ afterEach(async () => { }); describe('RedisBridgeStore', () => { + test('rejects a registration whose authenticated identity was replaced', async () => { + await redis.set( + 'codeapi:bridge:v1:identity:fenced-worker', + 'replacement-credential-digest', + ); + + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'fenced-worker', + incarnationId, + credentialId: 'stale-credential-digest', + identityId: 'stale-identity', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }, 'stale-credential-digest'), + ).rejects.toMatchObject({ code: 'WORKER_UNAUTHORIZED' }); + + await expect( + redis.get('codeapi:bridge:v1:worker:fenced-worker'), + ).resolves.toBeNull(); + }); + + test('accepts registration after a same-identity credential rotation', async () => { + await redis.set( + 'codeapi:bridge:v1:identity:rotating-registration-worker', + 'new-active-credential-digest', + ); + await redis.set( + 'codeapi:bridge:v1:stable-identity:rotating-registration-worker', + 'stable-worker-identity', + ); + + await expect( + store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'rotating-registration-worker', + incarnationId, + credentialId: 'old-authenticated-credential-digest', + identityId: 'stable-worker-identity', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }, 'old-authenticated-credential-digest'), + ).resolves.toBeUndefined(); + }); + + test('rejects a dynamic worker lease outside its bound tenant', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'tenant-worker', + incarnationId, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + binding: { + tenantId: 'tenant-1', + principal: { type: 'user', id: 'user-1' }, + }, + }); + + await expect( + store.dispatch({ + workerId: 'tenant-worker', + tenantId: 'tenant-2', + requireTenantBinding: true, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 1_000, + signal: new AbortController().signal, + }), + ).rejects.toMatchObject({ code: 'WORKER_UNAUTHORIZED' }); + }); + + test('does not lease an assignment to a newly rebound worker identity', async () => { + await redis.set( + 'codeapi:bridge:v1:stable-identity:rebound-worker', + 'tenant-a-identity', + ); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'rebound-worker', + incarnationId, + identityId: 'tenant-a-identity', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + binding: { + tenantId: 'tenant-a', + principal: { type: 'user', id: 'user-a' }, + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'rebound-worker', + tenantId: 'tenant-a', + requireTenantBinding: true, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + + await expect( + store.lease( + 'rebound-worker', + incarnationId, + 1_000, + undefined, + 'tenant-b-identity', + ), + ).resolves.toBeUndefined(); + await expect( + store.lease( + 'rebound-worker', + incarnationId, + 1_000, + undefined, + 'tenant-a-identity', + ), + ).resolves.toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + + test('a stale identity poll cannot consume work queued for the replacement identity', async () => { + await redis.set( + 'codeapi:bridge:v1:stable-identity:replacement-worker', + 'replacement-identity', + ); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'replacement-worker', + incarnationId, + identityId: 'replacement-identity', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + binding: { + tenantId: 'tenant-a', + principal: { type: 'user', id: 'user-a' }, + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'replacement-worker', + tenantId: 'tenant-a', + requireTenantBinding: true, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + + await expect( + store.lease( + 'replacement-worker', + incarnationId, + 100, + undefined, + 'stale-identity', + ), + ).resolves.toBeUndefined(); + await expect( + store.lease( + 'replacement-worker', + incarnationId, + 1_000, + undefined, + 'replacement-identity', + ), + ).resolves.toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + + test('a revoked identity poll cannot consume its previously queued work', async () => { + const workerId = 'revoked-lease-worker'; + const stableIdentityKey = `codeapi:bridge:v1:stable-identity:${workerId}`; + await redis.set(stableIdentityKey, 'revoked-identity'); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId, + incarnationId, + identityId: 'revoked-identity', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + binding: { + tenantId: 'tenant-a', + principal: { type: 'user', id: 'user-a' }, + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId, + tenantId: 'tenant-a', + requireTenantBinding: true, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + await redis.set(stableIdentityKey, 'replacement-identity'); + + await expect( + store.lease(workerId, incarnationId, 100, undefined, 'revoked-identity'), + ).resolves.toBeUndefined(); + await expect( + redis.llen(`codeapi:bridge:v1:worker:${workerId}:assignments`), + ).resolves.toBe(1); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + + test('a stale incarnation poll cannot consume replacement incarnation work', async () => { + const replacementIncarnationId = 'incarnation-00000002'; + await redis.set( + 'codeapi:bridge:v1:stable-identity:restarted-worker', + 'stable-restarted-identity', + ); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'restarted-worker', + incarnationId: replacementIncarnationId, + identityId: 'stable-restarted-identity', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'restarted-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + + await expect( + store.lease( + 'restarted-worker', + incarnationId, + 100, + undefined, + 'stable-restarted-identity', + ), + ).resolves.toBeUndefined(); + await expect( + store.lease( + 'restarted-worker', + replacementIncarnationId, + 1_000, + undefined, + 'stable-restarted-identity', + ), + ).resolves.toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + + test('leases an assignment queued by the prior identity-only encoding', async () => { + const identityId = 'rollout-compatible-identity'; + await redis.set( + 'codeapi:bridge:v1:stable-identity:rollout-worker', + identityId, + ); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'rollout-worker', + incarnationId, + identityId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'rollout-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + const [storedKey] = await redis.keys('codeapi:bridge:v1:assignment:*'); + const assignmentId = storedKey.replace('codeapi:bridge:v1:assignment:', ''); + const queue = 'codeapi:bridge:v1:worker:rollout-worker:assignments'; + await redis.del(queue); + await redis.rpush(queue, `${identityId.length}:${identityId}${assignmentId}`); + + await expect( + store.lease('rollout-worker', incarnationId, 1_000, undefined, identityId), + ).resolves.toMatchObject({ assignmentId }); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + + test('leases an assignment queued by the original raw encoding', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'legacy-queue-worker', + incarnationId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'legacy-queue-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + await new Promise((resolve) => setTimeout(resolve, 10)); + const [storedKey] = await redis.keys('codeapi:bridge:v1:assignment:*'); + const assignmentId = storedKey.replace('codeapi:bridge:v1:assignment:', ''); + const queue = 'codeapi:bridge:v1:worker:legacy-queue-worker:assignments'; + await redis.del(queue); + await redis.rpush(queue, assignmentId); + + await expect( + store.lease('legacy-queue-worker', incarnationId, 1_000), + ).resolves.toMatchObject({ assignmentId }); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + + test('leases queued work after credential refresh preserves the paired identity', async () => { + await redis.set( + 'codeapi:bridge:v1:stable-identity:rotating-worker', + 'stable-paired-identity', + ); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'rotating-worker', + incarnationId, + identityId: 'stable-paired-identity', + credentialId: 'credential-before-refresh', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + binding: { + tenantId: 'tenant-a', + principal: { type: 'user', id: 'user-a' }, + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'rotating-worker', + tenantId: 'tenant-a', + requireTenantBinding: true, + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + + const assignment = await store.lease( + 'rotating-worker', + incarnationId, + 1_000, + undefined, + 'stable-paired-identity', + ); + + expect(assignment).toBeDefined(); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + test('delivers and settles one fenced stateful assignment', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -58,6 +452,62 @@ describe('RedisBridgeStore', () => { }); }); + test('rejects settlement after the paired identity is replaced', async () => { + const identityId = 'settlement-owner-identity'; + await redis.set( + 'codeapi:bridge:v1:stable-identity:settlement-worker', + identityId, + ); + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'settlement-worker', + incarnationId, + identityId, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }); + const controller = new AbortController(); + const completion = store.dispatch({ + workerId: 'settlement-worker', + body: { language: 'bash' } as t.PayloadBody, + headers: {}, + deadlineAtMs: Date.now() + 5_000, + signal: controller.signal, + }); + const assignment = await store.lease( + 'settlement-worker', + incarnationId, + 1_000, + undefined, + identityId, + ); + await redis.set( + 'codeapi:bridge:v1:stable-identity:settlement-worker', + 'replacement-owner-identity', + ); + + await expect( + store.settle( + 'settlement-worker', + assignment?.assignmentId ?? '', + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + generation: assignment?.generation ?? 0, + leaseToken: assignment?.leaseToken ?? '', + incarnationId, + status: 'rejected', + error: 'sandbox failed', + }, + identityId, + ), + ).rejects.toMatchObject({ code: 'ASSIGNMENT_FENCED' }); + controller.abort(); + await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); + }); + test('rejects dispatch to an offline worker', async () => { const controller = new AbortController(); await expect( diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 04fb8e29..2d3975ea 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -9,6 +9,7 @@ import type { } from '../../../packages/code/src/protocol'; import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import type { BridgeWorkerBinding } from './pairing'; const PREFIX = 'codeapi:bridge:v1'; const POLL_INTERVAL_MS = 100; @@ -28,6 +29,7 @@ export class BridgeStoreError extends Error { constructor( public readonly code: | 'WORKER_OFFLINE' + | 'WORKER_UNAUTHORIZED' | 'WORKER_BUSY' | 'ASSIGNMENT_EXPIRED' | 'ASSIGNMENT_FENCED' @@ -45,12 +47,23 @@ export class BridgeStoreError extends Error { interface StoredAssignment extends CodeBridgeAssignment { leaseTokenHash: string; + workerIdentityId?: string; +} + +export interface RegisteredBridgeWorker extends BridgeWorkerRegistration { + binding?: BridgeWorkerBinding; + credentialId?: string; + identityId?: string; } function workerKey(workerId: string): string { return `${PREFIX}:worker:${workerId}`; } +function workerStableIdentityKey(workerId: string): string { + return `${PREFIX}:stable-identity:${workerId}`; +} + function workerIncarnationKey(workerId: string): string { return `${PREFIX}:worker:${workerId}:incarnation`; } @@ -89,6 +102,18 @@ function assignmentKey(assignmentId: string): string { return `${PREFIX}:assignment:${assignmentId}`; } +function queuedAssignment( + assignmentId: string, + incarnationId: string, + workerIdentityId?: string, +): string { + const identity = workerIdentityId ?? ''; + return ( + `${identity.length}:${identity}` + + `${incarnationId.length}:${incarnationId}${assignmentId}` + ); +} + function settlementKey(assignmentId: string): string { return `${PREFIX}:assignment:${assignmentId}:settlement`; } @@ -132,8 +157,18 @@ export class RedisBridgeStore { private readonly workerTtlSeconds = DEFAULT_WORKER_TTL_SECONDS, ) {} - async register(registration: BridgeWorkerRegistration): Promise { + async register( + registration: RegisteredBridgeWorker, + expectedActiveCredentialId?: string, + ): Promise { const script = [ + 'if ARGV[5] ~= "" then', + ' if ARGV[6] ~= "" then', + ' if redis.call(\'GET\', KEYS[5]) ~= ARGV[6] then return -3 end', + ' elseif redis.call(\'GET\', KEYS[6]) ~= ARGV[5] then', + ' return -3', + ' end', + 'end', 'if redis.call(\'EXISTS\', KEYS[3]) == 1 then return -2 end', 'if redis.call(\'EXISTS\', KEYS[2]) == 1 then return -1 end', 'local current = redis.call(\'GET\', KEYS[4])', @@ -149,15 +184,19 @@ export class RedisBridgeStore { const result = Number( await this.redis.eval( script, - 4, + 6, workerKey(registration.workerId), incarnationFenceKey(registration.workerId, registration.incarnationId), quarantineKey(registration.workerId, registration.incarnationId), workerIncarnationKey(registration.workerId), + workerStableIdentityKey(registration.workerId), + `${PREFIX}:identity:${registration.workerId}`, registration.incarnationId, JSON.stringify(registration), String(this.workerTtlSeconds), `${PREFIX}:worker:${registration.workerId}:incarnation:`, + expectedActiveCredentialId ?? '', + registration.identityId ?? '', ), ); if (result === -2) { @@ -172,10 +211,18 @@ export class RedisBridgeStore { 'Bridge worker incarnation was replaced', ); } + if (result === -3) { + throw new BridgeStoreError( + 'WORKER_UNAUTHORIZED', + `Bridge worker ${registration.workerId} identity changed during registration`, + ); + } } async dispatch(args: { workerId: string; + tenantId?: string; + requireTenantBinding?: boolean; body: t.PayloadBody; headers: Record; runtimeSessionId?: string; @@ -192,6 +239,18 @@ export class RedisBridgeStore { `Bridge worker ${args.workerId} is offline`, ); } + if ( + (args.requireTenantBinding === true && registration.binding == null) || + (registration.binding != null && + (args.tenantId == null || + args.tenantId.length === 0 || + registration.binding.tenantId !== args.tenantId)) + ) { + throw new BridgeStoreError( + 'WORKER_UNAUTHORIZED', + `Bridge worker ${args.workerId} is not authorized for this tenant`, + ); + } if ( args.runtimeSessionId !== undefined && registration.capabilities.statefulWorkspace !== true @@ -241,6 +300,9 @@ export class RedisBridgeStore { generation, leaseToken, leaseTokenHash: tokenHash(leaseToken), + ...(registration.identityId != null + ? { workerIdentityId: registration.identityId } + : {}), expiresAt: new Date(args.deadlineAtMs).toISOString(), runtimeSessionId: args.runtimeSessionId, request: { @@ -255,7 +317,14 @@ export class RedisBridgeStore { 'EX', ttlSeconds, ); - transaction.rpush(queueKey(args.workerId), assignmentId); + transaction.rpush( + queueKey(args.workerId), + queuedAssignment( + assignmentId, + assignment.incarnationId, + assignment.workerIdentityId, + ), + ); transaction.expire(queueKey(args.workerId), ttlSeconds); await transaction.exec(); const settlement = await this.waitForSettlement( @@ -291,19 +360,73 @@ export class RedisBridgeStore { incarnationId: string, waitMs: number, signal?: AbortSignal, + identityId?: string, ): Promise { const deadline = Date.now() + waitMs; while (signal?.aborted !== true && Date.now() < deadline) { - const assignmentId = await this.redis.lpop(queueKey(workerId)); - if (assignmentId == null) { + const raw = await this.redis.eval( + [ + "if ARGV[3] ~= '' then", + " if redis.call('GET', KEYS[2]) ~= ARGV[3] then return nil end", + "elseif redis.call('EXISTS', KEYS[2]) == 1 then", + ' return nil', + 'end', + "local entries = redis.call('LRANGE', KEYS[1], 0, -1)", + 'for _, entry in ipairs(entries) do', + " local separator = string.find(entry, ':', 1, true)", + ' local id = nil', + " local identity = ''", + " local incarnation = ''", + ' local legacy = false', + ' if separator then', + ' local identityLength = tonumber(string.sub(entry, 1, separator - 1))', + ' if identityLength then', + ' identity = string.sub(entry, separator + 1, separator + identityLength)', + ' local incarnationLengthStart = separator + identityLength + 1', + " local incarnationSeparator = string.find(entry, ':', incarnationLengthStart, true)", + ' if incarnationSeparator then', + ' local incarnationLength = tonumber(string.sub(entry, incarnationLengthStart, incarnationSeparator - 1))', + ' if incarnationLength then', + ' incarnation = string.sub(entry, incarnationSeparator + 1, incarnationSeparator + incarnationLength)', + ' id = string.sub(entry, incarnationSeparator + incarnationLength + 1)', + ' end', + ' else', + ' id = string.sub(entry, incarnationLengthStart)', + ' legacy = true', + ' end', + ' end', + " elseif ARGV[3] == '' then", + ' id = entry', + ' legacy = true', + ' end', + ' if id and identity == ARGV[3] then', + " local raw = redis.call('GET', ARGV[1] .. id)", + ' if not raw then', + " redis.call('LREM', KEYS[1], 1, entry)", + " elseif incarnation == ARGV[4] or (legacy and string.find(raw, '\"incarnationId\":\"' .. ARGV[4] .. '\"', 1, true)) then", + " redis.call('LREM', KEYS[1], 1, entry)", + ' return raw', + ' end', + ' end', + 'end', + 'return nil', + ].join('\n'), + 2, + queueKey(workerId), + workerStableIdentityKey(workerId), + `${PREFIX}:assignment:`, + workerId, + identityId ?? '', + incarnationId, + ); + if (raw == null) { await delay( Math.min(POLL_INTERVAL_MS, Math.max(0, deadline - Date.now())), signal, ); continue; } - const assignment = await this.readAssignment(assignmentId); - if (assignment == null || assignment.workerId !== workerId) continue; + const assignment = JSON.parse(String(raw)) as StoredAssignment; if (assignment.incarnationId !== incarnationId) continue; const registration = await this.registration(workerId); if (registration?.incarnationId !== incarnationId) { @@ -312,8 +435,18 @@ export class RedisBridgeStore { 'Bridge worker incarnation was replaced', ); } + if ( + assignment.workerId !== workerId || + assignment.workerIdentityId !== identityId + ) { + continue; + } if (Date.parse(assignment.expiresAt) <= Date.now()) continue; - const { leaseTokenHash: _leaseTokenHash, ...wireAssignment } = assignment; + const { + leaseTokenHash: _leaseTokenHash, + workerIdentityId: _workerIdentityId, + ...wireAssignment + } = assignment; return wireAssignment; } return undefined; @@ -323,26 +456,32 @@ export class RedisBridgeStore { workerId: string, assignmentId: string, settlement: CodeBridgeSettlement, + identityId?: string, ): Promise { - const assignment = await this.readAssignment(assignmentId); - if (assignment == null) { + const rawAssignment = await this.redis.get(assignmentKey(assignmentId)); + if (rawAssignment == null) { throw new BridgeStoreError( 'ASSIGNMENT_NOT_FOUND', 'Bridge assignment was not found', ); } + const assignment = JSON.parse(rawAssignment) as StoredAssignment; if (assignment.workerId !== workerId) { throw new BridgeStoreError( 'WORKER_MISMATCH', 'Bridge assignment belongs to another worker', ); } - const registration = await this.registration(workerId); + const rawRegistration = await this.redis.get(workerKey(workerId)); + const registration = rawRegistration == null + ? undefined + : JSON.parse(rawRegistration) as RegisteredBridgeWorker; if ( settlement.incarnationId !== assignment.incarnationId || registration?.incarnationId !== settlement.incarnationId || settlement.generation !== assignment.generation || - tokenHash(settlement.leaseToken) !== assignment.leaseTokenHash + tokenHash(settlement.leaseToken) !== assignment.leaseTokenHash || + assignment.workerIdentityId !== identityId ) { throw new BridgeStoreError( 'ASSIGNMENT_FENCED', @@ -356,12 +495,37 @@ export class RedisBridgeStore { ); } const ttlSeconds = assignmentTtlSeconds(Date.parse(assignment.expiresAt)); - await this.redis.set( - settlementKey(assignmentId), - JSON.stringify(settlement), - 'EX', - ttlSeconds, + const accepted = Number( + await this.redis.eval( + [ + "if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end", + "if redis.call('GET', KEYS[2]) ~= ARGV[2] then return 0 end", + 'if ARGV[3] ~= "" then', + " if redis.call('GET', KEYS[3]) ~= ARGV[3] then return 0 end", + "elseif redis.call('EXISTS', KEYS[3]) == 1 then", + ' return 0', + 'end', + "redis.call('SET', KEYS[4], ARGV[4], 'EX', ARGV[5])", + 'return 1', + ].join('\n'), + 4, + assignmentKey(assignmentId), + workerKey(workerId), + workerStableIdentityKey(workerId), + settlementKey(assignmentId), + rawAssignment, + rawRegistration ?? '', + identityId ?? '', + JSON.stringify(settlement), + String(ttlSeconds), + ), ); + if (accepted !== 1) { + throw new BridgeStoreError( + 'ASSIGNMENT_FENCED', + 'Bridge assignment lease is stale', + ); + } } async cancelled( @@ -403,9 +567,9 @@ export class RedisBridgeStore { private async registration( workerId: string, - ): Promise { + ): Promise { const raw = await this.redis.get(workerKey(workerId)); - return raw == null ? undefined : (JSON.parse(raw) as BridgeWorkerRegistration); + return raw == null ? undefined : (JSON.parse(raw) as RegisteredBridgeWorker); } private async readAssignment( diff --git a/service/src/config.ts b/service/src/config.ts index 605c701b..95df4eba 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -365,6 +365,8 @@ export const env = { * - `remote-bridge`: dispatch to an outbound-connected @librechat/code worker. */ SANDBOX_BACKEND: sandboxBackend, + /** Permit trusted callers to route each execution to a paired worker ID. */ + BRIDGE_DYNAMIC_WORKERS: process.env.CODEAPI_BRIDGE_DYNAMIC_WORKERS === 'true', /** Outbound worker selected by the remote-bridge backend. */ BRIDGE_WORKER_ID: process.env.CODEAPI_BRIDGE_WORKER_ID ?? '', /** Static compatibility auth or short-lived proof-of-possession credentials. */ diff --git a/service/src/execution-profile.test.ts b/service/src/execution-profile.test.ts index da20bce2..7ba4640f 100644 --- a/service/src/execution-profile.test.ts +++ b/service/src/execution-profile.test.ts @@ -2,8 +2,11 @@ import { describe, expect, test } from 'bun:test'; import { checkExecutionProfileExpectation, queueNamesForExecutionProfile, + queueNameForExecution, resolveExecutionProfile, resolveExecutionProfileSource, + resolveQueuedSandboxBackend, + validateQueuedSandboxBackend, validateQueuedExecutionProfile, } from './execution-profile'; @@ -55,6 +58,42 @@ describe('execution profile queue isolation', () => { other: 'stateful-other-queue', }); }); + + test('routes a persisted remote bridge replay to the bridge queue on a lambda API', () => { + expect( + queueNameForExecution( + 'python', + 'stateful', + 'explicit', + 'remote-bridge', + ), + ).toBe('remote-bridge-python-queue'); + }); + + test('isolates outbound bridge jobs from Lambda consumers', () => { + expect( + queueNamesForExecutionProfile('stateful', 'explicit', 'remote-bridge'), + ).toEqual({ + python: 'remote-bridge-python-queue', + other: 'remote-bridge-other-queue', + }); + expect( + queueNamesForExecutionProfile('stateful', 'explicit', 'lambda-microvm'), + ).toEqual({ + python: 'stateful-python-queue', + other: 'stateful-other-queue', + }); + }); + + test('labels API-only stateful jobs with their Lambda worker backend', () => { + expect(resolveQueuedSandboxBackend('stateful', 'http')).toBe('lambda-microvm'); + expect(resolveQueuedSandboxBackend('default', 'http', 'explicit')).toBe('http'); + expect(resolveQueuedSandboxBackend('stateful', 'remote-bridge')).toBe('remote-bridge'); + }); + + test('leaves the backend unfenced for inferred stateless legacy queues', () => { + expect(resolveQueuedSandboxBackend('default', 'http', 'inferred')).toBeUndefined(); + }); }); describe('execution profile request assertion', () => { @@ -107,3 +146,29 @@ describe('queued execution profile validation', () => { ); }); }); + +describe('queued sandbox backend validation', () => { + test('accepts matching and legacy jobs', () => { + expect(() => + validateQueuedSandboxBackend('remote-bridge', 'remote-bridge'), + ).not.toThrow(); + expect(() => validateQueuedSandboxBackend(undefined, 'http')).not.toThrow(); + expect(() => + validateQueuedSandboxBackend(undefined, 'remote-bridge', 'legacy-bridge-worker'), + ).not.toThrow(); + }); + + test('rejects invalid and cross-backend jobs', () => { + expect(() => validateQueuedSandboxBackend('invalid', 'http')).toThrow( + 'Queued job has invalid sandbox backend', + ); + expect(() => + validateQueuedSandboxBackend('remote-bridge', 'lambda-microvm'), + ).toThrow( + 'Queued job targets the remote-bridge sandbox backend, but worker serves lambda-microvm', + ); + expect(() => + validateQueuedSandboxBackend(undefined, 'lambda-microvm', 'legacy-bridge-worker'), + ).toThrow('Legacy queued bridge job cannot run on the lambda-microvm sandbox backend'); + }); +}); diff --git a/service/src/execution-profile.ts b/service/src/execution-profile.ts index c4951903..38e9bc8c 100644 --- a/service/src/execution-profile.ts +++ b/service/src/execution-profile.ts @@ -1,4 +1,9 @@ export const EXECUTION_PROFILES = ['default', 'stateful'] as const; +export const SANDBOX_BACKENDS = [ + 'http', + 'lambda-microvm', + 'remote-bridge', +] as const; export type ExecutionProfile = typeof EXECUTION_PROFILES[number]; export type ExecutionProfileSource = 'explicit' | 'inferred'; @@ -11,6 +16,30 @@ export interface ExecutionProfileQueueNames { other: string; } +export type SandboxBackendName = typeof SANDBOX_BACKENDS[number]; + +/** Resolve the backend owned by the queue consumer rather than the API pod. + * Stateful API-only pods intentionally retain the HTTP local default while + * dispatching to Lambda workers. */ +export function resolveQueuedSandboxBackend( + profile: ExecutionProfile, + apiBackend: SandboxBackendName, + source: ExecutionProfileSource = 'explicit', +): SandboxBackendName | undefined { + if (profile === 'stateful' && apiBackend === 'http') { + return 'lambda-microvm'; + } + /* An inferred default profile still uses the pre-fencing legacy queues. + * Its API-only process cannot distinguish the supported HTTP and Lambda + * consumers because Lambda-only configuration belongs to the worker pod. + * Preserve that rollout topology by leaving the backend absent, exactly as + * pre-fencing producers did; explicit profiles regain strict fencing. */ + if (profile === 'default' && source === 'inferred' && apiBackend === 'http') { + return undefined; + } + return apiBackend; +} + export function resolveExecutionProfile( raw: string | undefined, runtimeSessionMode: 'stateless' | 'affinity' | 'strict', @@ -54,10 +83,17 @@ const EXPLICIT_PROFILE_QUEUE_NAMES: Record { logger.info('Starting API service (no workers)...'); validateApiHardenedConfig(); validateExecutionProfilePolicy({ requireBackendMatch: false }); + validateApiSandboxBackendPolicy(); validateApiBridgePolicy(); - /* No validateSandboxBackendPolicy() here: an API-only pod authenticates and + /* No full validateSandboxBackendPolicy() here: an API-only pod authenticates and * enqueues jobs, it never constructs the Lambda backend or checkpoint store. * Validating that policy would force worker-only config (LAMBDA_MICROVM_* and * the MINIO_* checkpoint creds) into API pods just to boot. The worker and @@ -252,12 +259,7 @@ export async function gracefulShutdown(): Promise { } // Close queue connections (both API and Worker need this) - await Promise.all([ - pyQueue.close(), - otherQueue.close(), - pyQueueEvents.close(), - otherQueueEvents.close() - ]); + await closeQueueConnections(); logger.info('Queue connections closed'); // Only disconnect Redis if explicitly requested diff --git a/service/src/local-api.ts b/service/src/local-api.ts index 701270d7..84e24352 100644 --- a/service/src/local-api.ts +++ b/service/src/local-api.ts @@ -13,7 +13,7 @@ import programmaticRouter from './service/programmatic-router'; import { requestErrorLogger, requestNotFoundLogger } from './middleware/request-error-logger'; import { executionProfileMiddleware } from './middleware/execution-profile'; import { localAuth } from './auth/local'; -import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection } from './queue'; +import { pyQueue, otherQueue, connection, closeQueueConnections } from './queue'; import { setStartupComplete } from './lifecycle'; // Workers are imported to ensure they're started with the process import './workers'; @@ -90,12 +90,7 @@ async function localShutdown(): Promise { localShuttingDown = true; logger.info('Shutting down local server...'); try { - await Promise.all([ - pyQueue.close(), - otherQueue.close(), - pyQueueEvents.close(), - otherQueueEvents.close() - ]); + await closeQueueConnections(); try { await shutdownTelemetry(); } catch (telemetryError) { diff --git a/service/src/queue.ts b/service/src/queue.ts index 91f97c36..54fea308 100644 --- a/service/src/queue.ts +++ b/service/src/queue.ts @@ -7,7 +7,15 @@ import type * as tls from 'tls'; import type * as t from './types'; import { Jobs } from './enum'; import { env } from './config'; -import { queueNamesForExecutionProfile } from './execution-profile'; +import { + queueNameForExecution, + queueNamesForExecutionProfile, +} from './execution-profile'; +import type { + ExecutionProfile, + ExecutionProfileSource, + SandboxBackendName, +} from './execution-profile'; import logger from './logger'; import { redisKeepAliveOptions } from './redis-options'; import { bullmqQueueJobs, registerBullmqQueueMetricsCollector } from './metrics'; @@ -60,18 +68,52 @@ const connection = new IORedis({ const queueNames = queueNamesForExecutionProfile( env.EXECUTION_PROFILE, env.EXECUTION_PROFILE_SOURCE, + env.SANDBOX_BACKEND, ); -const pyQueue = new Queue(queueNames.python, { connection }); -const otherQueue = new Queue(queueNames.other, { connection }); +export interface QueueBinding { + queue: Queue; + events: QueueEvents; + language: 'python' | 'bash'; +} + +const queueResources = new Map< + string, + { queue: Queue; events: QueueEvents } +>(); -const pyQueueEvents = new QueueEvents(queueNames.python, { connection }); -const otherQueueEvents = new QueueEvents(queueNames.other, { connection }); +function getQueueResources( + name: string, +): { queue: Queue; events: QueueEvents } { + const existing = queueResources.get(name); + if (existing != null) return existing; + + const queue = new Queue(name, { connection }); + const events = new QueueEvents(name, { connection }); + setMaxListeners(0, queue, events); + const resources = { queue, events }; + queueResources.set(name, resources); + return resources; +} + +export function getExecutionQueueBinding( + language: 'python' | 'bash', + backend: SandboxBackendName | undefined = env.SANDBOX_BACKEND, + profile: ExecutionProfile = env.EXECUTION_PROFILE, + source: ExecutionProfileSource = env.EXECUTION_PROFILE_SOURCE, +): QueueBinding { + const name = queueNameForExecution( + language, + profile, + source, + backend, + ); + return { ...getQueueResources(name), language }; +} + +const { queue: pyQueue, events: pyQueueEvents } = getQueueResources(queueNames.python); +const { queue: otherQueue, events: otherQueueEvents } = getQueueResources(queueNames.other); const queueMetricStates = ['waiting', 'active', 'delayed'] as const; -const queueMetricSources = [ - { name: queueNames.python, queue: pyQueue }, - { name: queueNames.other, queue: otherQueue }, -] as const; const QUEUE_METRICS_TIMEOUT_MS = 1000; async function withTimeout(promise: Promise, timeoutMs: number, message: string): Promise { @@ -90,7 +132,7 @@ async function withTimeout(promise: Promise, timeoutMs: number, message: s } registerBullmqQueueMetricsCollector(async () => { - await Promise.all(queueMetricSources.map(async ({ name, queue }) => { + await Promise.all([...queueResources.entries()].map(async ([name, { queue }]) => { try { const counts = await withTimeout( queue.getJobCounts(...queueMetricStates), @@ -116,4 +158,13 @@ registerBullmqQueueMetricsCollector(async () => { * BullMQ coordination objects. */ setMaxListeners(0, pyQueue, otherQueue, pyQueueEvents, otherQueueEvents); +export async function closeQueueConnections(): Promise { + await Promise.all( + [...queueResources.values()].flatMap(({ queue, events }) => [ + queue.close(), + events.close(), + ]), + ); +} + export { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, queueNames, connection }; diff --git a/service/src/sandbox-backend/remote-bridge.test.ts b/service/src/sandbox-backend/remote-bridge.test.ts new file mode 100644 index 00000000..c1b9e2c9 --- /dev/null +++ b/service/src/sandbox-backend/remote-bridge.test.ts @@ -0,0 +1,153 @@ +import { describe, expect, test } from 'bun:test'; + +import type { SandboxExecuteContext, SandboxTransportRequest } from './types'; +import type { RedisBridgeStore } from '../bridge/store'; + +import { BridgeStoreError } from '../bridge/store'; +import { RemoteBridgeSandboxBackend } from './remote-bridge'; + +function request(): SandboxTransportRequest { + return { + body: { language: 'bash' } as never, + headers: {}, + }; +} + +function context(): SandboxExecuteContext { + return { + executionId: 'execution-1', + language: 'bash', + isSynthetic: false, + signal: new AbortController().signal, + tenantId: 'tenant-1', + bridgeWorkerId: 'user-vm', + runtimeSessionMode: 'strict', + }; +} + +describe('RemoteBridgeSandboxBackend', () => { + test('dispatches a dynamically selected worker with a required tenant binding', async () => { + let dispatched: Parameters[0] | undefined; + const store = { + dispatch: async ( + args: Parameters[0], + ): ReturnType => { + dispatched = args; + return { + protocolVersion: 1 as const, + generation: 1, + leaseToken: 'a'.repeat(32), + incarnationId: 'incarnation-00000001', + status: 'fulfilled' as const, + result: { + session_id: 'session-1', + language: 'bash', + version: '5.2.0', + files: [], + }, + }; + }, + } satisfies Pick; + const backend = new RemoteBridgeSandboxBackend(store, 'default-vm'); + + await expect(backend.execute(request(), context())).resolves.toMatchObject({ + session_id: 'session-1', + }); + expect(dispatched).toMatchObject({ + workerId: 'user-vm', + tenantId: 'tenant-1', + requireTenantBinding: true, + }); + }); + + test('maps tenant authorization rejection to a bridge backend error', async () => { + const store = { + dispatch: async (): ReturnType => { + throw new BridgeStoreError('WORKER_UNAUTHORIZED', 'private tenant detail'); + }, + } satisfies Pick; + const backend = new RemoteBridgeSandboxBackend(store, 'default-vm'); + + await expect(backend.execute(request(), context())).rejects.toMatchObject({ + code: 'BRIDGE_WORKER_UNAUTHORIZED', + }); + }); + + test('keeps an explicitly selected singleton on its unbound compatibility route', async () => { + let dispatched: Parameters[0] | undefined; + const store = { + dispatch: async ( + args: Parameters[0], + ): ReturnType => { + dispatched = args; + return { + protocolVersion: 1 as const, + generation: 1, + leaseToken: 'a'.repeat(32), + incarnationId: 'incarnation-00000001', + status: 'fulfilled' as const, + result: { + session_id: 'session-1', + language: 'bash', + version: '5.2.0', + files: [], + }, + }; + }, + } satisfies Pick; + const backend = new RemoteBridgeSandboxBackend( + store, + 'deployment-worker', + false, + ); + + await backend.execute(request(), { + ...context(), + bridgeWorkerId: 'deployment-worker', + }); + + expect(dispatched).toMatchObject({ + workerId: 'deployment-worker', + requireTenantBinding: false, + }); + }); + + test('requires a binding for the selected default worker in dynamic mode', async () => { + let dispatched: Parameters[0] | undefined; + const store = { + dispatch: async ( + args: Parameters[0], + ): ReturnType => { + dispatched = args; + return { + protocolVersion: 1 as const, + generation: 1, + leaseToken: 'a'.repeat(32), + incarnationId: 'incarnation-00000001', + status: 'fulfilled' as const, + result: { + session_id: 'session-1', + language: 'bash', + version: '5.2.0', + files: [], + }, + }; + }, + } satisfies Pick; + const backend = new RemoteBridgeSandboxBackend( + store, + 'deployment-worker', + true, + ); + + await backend.execute(request(), { + ...context(), + bridgeWorkerId: 'deployment-worker', + }); + + expect(dispatched).toMatchObject({ + workerId: 'deployment-worker', + requireTenantBinding: true, + }); + }); +}); diff --git a/service/src/sandbox-backend/remote-bridge.ts b/service/src/sandbox-backend/remote-bridge.ts index a30ab748..b1a94ae0 100644 --- a/service/src/sandbox-backend/remote-bridge.ts +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -15,15 +15,17 @@ export class RemoteBridgeSandboxBackend implements SandboxBackend { readonly name = 'remote-bridge' as const; constructor( - private readonly store: RedisBridgeStore = bridgeStore, + private readonly store: Pick = bridgeStore, private readonly workerId: string = env.BRIDGE_WORKER_ID, + private readonly dynamicWorkers: boolean = env.BRIDGE_DYNAMIC_WORKERS, ) {} async execute( req: SandboxTransportRequest, ctx: SandboxExecuteContext, ): Promise { - if (!this.workerId) { + const workerId = ctx.bridgeWorkerId ?? this.workerId; + if (workerId.length === 0) { throw new SandboxBackendError( 'BRIDGE_WORKER_OFFLINE', 'No bridge worker is configured', @@ -32,7 +34,11 @@ export class RemoteBridgeSandboxBackend implements SandboxBackend { const sessionResultFinalizer = ctx.sessionResultFinalizer; try { const settlement = await this.store.dispatch({ - workerId: this.workerId, + workerId, + tenantId: ctx.tenantId, + requireTenantBinding: + ctx.bridgeWorkerId != null && + (this.dynamicWorkers || ctx.bridgeWorkerId !== this.workerId), body: req.body, headers: req.headers, runtimeSessionId: ctx.runtimeSessionId, @@ -57,6 +63,13 @@ export class RemoteBridgeSandboxBackend implements SandboxBackend { return settlement.result as SandboxRawResponse; } catch (error) { if (!(error instanceof BridgeStoreError)) throw error; + if (error.code === 'WORKER_UNAUTHORIZED') { + throw new SandboxBackendError( + 'BRIDGE_WORKER_UNAUTHORIZED', + error.message, + error, + ); + } if (error.code === 'WORKER_BUSY') { throw new SandboxBackendError( 'BRIDGE_WORKER_BUSY', diff --git a/service/src/sandbox-backend/types.ts b/service/src/sandbox-backend/types.ts index 15bb6942..64e5b6e0 100644 --- a/service/src/sandbox-backend/types.ts +++ b/service/src/sandbox-backend/types.ts @@ -37,6 +37,8 @@ export interface SandboxExecuteContext { deadlineAtMs?: number; tenantId?: string; canonicalUserId?: string; + /** Trusted API-selected outbound worker. Presence requires a tenant-bound credential. */ + bridgeWorkerId?: string; /** Absent ⇒ stateless execution (no runtime session affinity). */ runtimeSessionId?: string; runtimeSessionMode: t.RuntimeSessionMode; @@ -64,6 +66,7 @@ export interface SandboxBackend { export type SandboxBackendErrorCode = | 'RUNTIME_SESSION_BUSY' | 'BRIDGE_WORKER_OFFLINE' + | 'BRIDGE_WORKER_UNAUTHORIZED' | 'BRIDGE_WORKER_BUSY' | 'BRIDGE_EXECUTION_FAILED' | 'BRIDGE_DEADLINE_EXCEEDED' diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 600f9756..65ac63d2 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -3,6 +3,7 @@ import { env } from './config'; import { validateApiHardenedConfig, validateApiBridgePolicy, + validateApiSandboxBackendPolicy, validateEgressGatewayHardenedConfig, validateExecutionProfilePolicy, validateSandboxBackendPolicy, @@ -15,6 +16,7 @@ const saved = { executionProfile: env.EXECUTION_PROFILE, executionProfileSource: env.EXECUTION_PROFILE_SOURCE, sandboxBackend: env.SANDBOX_BACKEND, + bridgeDynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, bridgeWorkerId: env.BRIDGE_WORKER_ID, bridgeAuthMode: env.BRIDGE_AUTH_MODE, bridgeToken: env.BRIDGE_TOKEN, @@ -56,6 +58,7 @@ function restore(): void { env.EXECUTION_PROFILE = saved.executionProfile; env.EXECUTION_PROFILE_SOURCE = saved.executionProfileSource; env.SANDBOX_BACKEND = saved.sandboxBackend; + env.BRIDGE_DYNAMIC_WORKERS = saved.bridgeDynamicWorkers; env.BRIDGE_WORKER_ID = saved.bridgeWorkerId; env.BRIDGE_AUTH_MODE = saved.bridgeAuthMode; env.BRIDGE_TOKEN = saved.bridgeToken; @@ -313,6 +316,36 @@ describe('sandbox backend policy', () => { expect(() => validateSandboxBackendPolicy()).not.toThrow(); }); + test('allows dynamic-only paired workers without a configured default', () => { + env.SANDBOX_BACKEND = 'remote-bridge'; + env.RUNTIME_SESSION_MODE = 'strict'; + env.PTC_MODE = 'replay'; + env.BRIDGE_DYNAMIC_WORKERS = true; + env.BRIDGE_WORKER_ID = ''; + env.BRIDGE_TOKEN = 'development-bridge-token'; + env.BRIDGE_AUTH_MODE = 'static'; + + expect(() => validateSandboxBackendPolicy()).toThrow( + 'CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + + env.BRIDGE_AUTH_MODE = 'paired'; + expect(() => validateSandboxBackendPolicy()).not.toThrow(); + }); + + test('requires paired dynamic worker auth in an API-only process', () => { + env.SANDBOX_BACKEND = 'http'; + env.BRIDGE_DYNAMIC_WORKERS = true; + env.BRIDGE_AUTH_MODE = 'static'; + + expect(() => validateApiSandboxBackendPolicy()).toThrow( + 'CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + + env.BRIDGE_AUTH_MODE = 'paired'; + expect(() => validateApiSandboxBackendPolicy()).not.toThrow(); + }); + test('hardened remote bridge requires replay PTC, paired auth, and a strong administrator token', () => { env.SANDBOX_BACKEND = 'remote-bridge'; env.RUNTIME_SESSION_MODE = 'affinity'; diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index b213dbf4..d4e08752 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -107,11 +107,20 @@ export function validateExecutionProfilePolicy(options: { } } +export function validateApiSandboxBackendPolicy(): void { + if (env.BRIDGE_DYNAMIC_WORKERS && env.BRIDGE_AUTH_MODE !== 'paired') { + throw new SecureStartupConfigError( + 'Dynamic remote bridge workers require CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + } +} + /** * Backend-selection policy. Unlike the hardened-mode validators, this runs * unconditionally: a misconfigured backend must never half-start. */ export function validateSandboxBackendPolicy(): void { + validateApiSandboxBackendPolicy(); if (env.RUNTIME_SESSION_MODE !== 'stateless' && env.SANDBOX_BACKEND === 'http') { throw new SecureStartupConfigError( `CODEAPI_RUNTIME_SESSION_MODE=${env.RUNTIME_SESSION_MODE} requires ` @@ -119,7 +128,9 @@ export function validateSandboxBackendPolicy(): void { ); } if (env.SANDBOX_BACKEND === 'remote-bridge') { - requireValue('CODEAPI_BRIDGE_WORKER_ID', env.BRIDGE_WORKER_ID); + if (!env.BRIDGE_DYNAMIC_WORKERS) { + requireValue('CODEAPI_BRIDGE_WORKER_ID', env.BRIDGE_WORKER_ID); + } if (env.HARDENED_SANDBOX_MODE) { requireStrongSecret('CODEAPI_BRIDGE_TOKEN', env.BRIDGE_TOKEN); if (env.BRIDGE_AUTH_MODE !== 'paired') { diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index eade27fb..baa350e6 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -2,11 +2,15 @@ import axios from 'axios'; import { nanoid } from 'nanoid'; import { Router } from 'express'; import type { Response } from 'express'; -import type { Queue, QueueEvents } from 'bullmq'; import type * as t from '../types'; import { checkServiceStartUp, checkServiceShutDown } from '../lifecycle'; import { executionLimiter } from '../middleware/limits'; -import { pyQueue, pyQueueEvents, otherQueue, otherQueueEvents, connection } from '../queue'; +import { + pyQueue, + pyQueueEvents, + connection, + getExecutionQueueBinding, +} from '../queue'; import { createProgrammaticPayload, extractPendingFromStdout } from '../preamble'; import { findBashToolNameCollision } from '../preamble-bash'; import type { LCTool } from '../preamble'; @@ -25,6 +29,8 @@ import { } from '../metrics'; import { Jobs } from '../enum'; import { env, jobCompletionWaitTimeoutMs } from '../config'; +import { resolveQueuedSandboxBackend } from '../execution-profile'; +import { publicExecutionFailure } from '../utils'; import { normalizeEgressGatewayUrl, normalizeProgrammaticTimeoutMs, @@ -35,7 +41,15 @@ import { import { findUnregisteredToolCall } from '../tool-scope'; import { summarizeRequestedFiles } from '../execution-log'; import { FileRefAuthorizationError, authorizeRequestedFiles } from './file-authorization'; -import { buildReplayExecutionState } from './programmatic-state'; +import { + buildReplayExecutionState, + resolveReplayStateSandboxBackend, +} from './programmatic-state'; +import { + BridgeWorkerSelectionError, + CODEAPI_BRIDGE_WORKER_HEADER, + resolveBridgeWorkerSelection, +} from '../bridge/selection'; import logger from '../logger'; import { type ExecutionState, @@ -328,19 +342,6 @@ async function waitForExecutionState( // Replay mode helpers // --------------------------------------------------------------------------- -interface QueueBinding { - queue: Queue; - events: QueueEvents; - language: 'python' | 'bash'; -} - -function pickQueue(language: 'python' | 'bash'): QueueBinding { - if (language === 'bash') { - return { queue: otherQueue, events: otherQueueEvents, language: 'bash' }; - } - return { queue: pyQueue, events: pyQueueEvents, language: 'python' }; -} - function buildReplayPayload( req: t.AuthenticatedRequest, state: ExecutionState, @@ -396,7 +397,21 @@ async function runReplayIteration( }); } - const { queue, events, language } = pickQueue(state.language ?? 'python'); + const replayBackend = + state.sandboxBackend ?? + (state.bridgeWorkerId != null + ? 'remote-bridge' + : resolveQueuedSandboxBackend( + env.EXECUTION_PROFILE, + env.SANDBOX_BACKEND, + env.EXECUTION_PROFILE_SOURCE, + )); + const { queue, events, language } = getExecutionQueueBinding( + state.language ?? 'python', + replayBackend, + state.executionProfile ?? env.EXECUTION_PROFILE, + state.executionProfileSource ?? env.EXECUTION_PROFILE_SOURCE, + ); const job = await queue.add(Jobs.execute, { code: state.userCode ?? '', userId, @@ -407,7 +422,9 @@ async function runReplayIteration( executionId: state.execution_id, tenantId: state.tenantId, canonicalUserId: state.canonicalUserId, - executionProfile: env.EXECUTION_PROFILE, + executionProfile: state.executionProfile ?? env.EXECUTION_PROFILE, + sandboxBackend: replayBackend, + ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, executionManifestClaims: sandboxSecurity.executionManifestClaims, @@ -439,9 +456,10 @@ async function handleReplayInitial( params: { apiKeyId: string; userId: string; + bridgeWorkerId?: string; }, ): Promise { - const { apiKeyId, userId } = params; + const { apiKeyId, userId, bridgeWorkerId } = params; const { code, tools, @@ -560,6 +578,15 @@ async function handleReplayInitial( isPyPlot, timeout, language, + bridgeWorkerId, + executionProfile: env.EXECUTION_PROFILE, + executionProfileSource: env.EXECUTION_PROFILE_SOURCE, + sandboxBackend: resolveReplayStateSandboxBackend({ + executionProfile: env.EXECUTION_PROFILE, + executionProfileSource: env.EXECUTION_PROFILE_SOURCE, + apiSandboxBackend: env.SANDBOX_BACKEND, + bridgeWorkerId, + }), }); /** Replay mode persists the full request (`userCode` + `tools` + `files`) * inside `ExecutionState` so continuations can re-enqueue without the @@ -832,7 +859,8 @@ async function runAndRespond( logger.error('Replay iteration failed', { execution_id: state.execution_id, err }); await cleanupExecution(state.execution_id, 'replay'); if (!isDisconnected()) { - const message = (err as Error).message; + const publicFailure = publicExecutionFailure(err); + const message = publicFailure?.body.message ?? (err as Error).message; res.status(200).json({ status: 'error', error: message !== '' ? message : 'Sandbox execution failed', @@ -1023,6 +1051,26 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR } = req.body as t.ProgrammaticRequestBody; const rawBody = req.body as Record; const requestedLanguage: unknown = rawBody.language ?? rawBody.lang; + let bridgeWorkerId: string | undefined; + if (continuation_token == null || continuation_token === '') { + try { + const bridgeSelection = resolveBridgeWorkerSelection({ + backend: env.SANDBOX_BACKEND, + configuredWorkerId: env.BRIDGE_WORKER_ID, + dynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, + requestedWorkerId: req.header(CODEAPI_BRIDGE_WORKER_HEADER), + trustedWorkerId: principal.codeWorkerId, + }); + bridgeWorkerId = bridgeSelection?.explicit === true + ? bridgeSelection.workerId + : undefined; + } catch (error) { + if (error instanceof BridgeWorkerSelectionError) { + return res.status(error.status).json({ error: error.message }); + } + throw error; + } + } if ( requestedLanguage !== undefined && @@ -1079,9 +1127,9 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR }); } if (env.PTC_MODE === 'replay') { - return await handleReplayInitial(req, res, { apiKeyId, userId }); + return await handleReplayInitial(req, res, { apiKeyId, userId, bridgeWorkerId }); } - return await handleBlocking(req, res, { apiKeyId, userId }); + return await handleBlocking(req, res, { apiKeyId, userId, bridgeWorkerId }); } catch (err) { logger.error(`[${INSTANCE_ID}] Programmatic routing error:`, err); if (!res.headersSent) { @@ -1099,9 +1147,9 @@ router.post('/exec/programmatic', executionLimiter, async (req: t.AuthenticatedR async function handleBlocking( req: t.AuthenticatedRequest, res: Response, - params: { apiKeyId: string; userId: string }, + params: { apiKeyId: string; userId: string; bridgeWorkerId?: string }, ): Promise> { - const { apiKeyId, userId } = params; + const { apiKeyId, userId, bridgeWorkerId } = params; const { code, tools, @@ -1282,6 +1330,7 @@ async function handleBlocking( principalSource: identity.principalSource, authContextHash: identity.authContextHash, apiKeyId, + bridgeWorkerId, startTime: Date.now(), lastActivity: Date.now(), mode: 'blocking', @@ -1378,6 +1427,12 @@ async function handleBlocking( tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, + sandboxBackend: resolveQueuedSandboxBackend( + env.EXECUTION_PROFILE, + env.SANDBOX_BACKEND, + env.EXECUTION_PROFILE_SOURCE, + ), + ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, executionManifestClaims: sandboxSecurity.executionManifestClaims, diff --git a/service/src/service/programmatic-state.test.ts b/service/src/service/programmatic-state.test.ts index f710668a..fc84d8f8 100644 --- a/service/src/service/programmatic-state.test.ts +++ b/service/src/service/programmatic-state.test.ts @@ -1,7 +1,10 @@ import { describe, expect, test } from 'bun:test'; import type { CodeApiAuthContext, RequestFile } from '../types'; import type { LCTool } from '../preamble'; -import { buildReplayExecutionState } from './programmatic-state'; +import { + buildReplayExecutionState, + resolveReplayStateSandboxBackend, +} from './programmatic-state'; const TOOLS = [ { @@ -22,7 +25,9 @@ const FILES = [ }, ] as RequestFile[]; -function build(overrides: Partial[0]> = {}) { +function build( + overrides: Partial[0]> = {}, +): ReturnType { return buildReplayExecutionState({ executionId: 'exec_123', sessionId: 'session_123', @@ -35,12 +40,35 @@ function build(overrides: Partial[0 isPyPlot: false, timeout: 300000, language: 'python', + executionProfile: 'default', + executionProfileSource: 'inferred', now: 1778250000000, ...overrides, }); } describe('buildReplayExecutionState', () => { + test('persists the resolved queue consumer backend for split stateful deployments', () => { + expect( + resolveReplayStateSandboxBackend({ + executionProfile: 'stateful', + executionProfileSource: 'explicit', + apiSandboxBackend: 'http', + }), + ).toBe('lambda-microvm'); + }); + + test('pins bridge replay state to the remote bridge backend', () => { + expect( + resolveReplayStateSandboxBackend({ + executionProfile: 'stateful', + executionProfileSource: 'explicit', + apiSandboxBackend: 'http', + bridgeWorkerId: 'worker-1', + }), + ).toBe('remote-bridge'); + }); + test('persists canonical LibreChat auth context for replay continuations', () => { const authContext: CodeApiAuthContext = { userId: 'user_canonical', @@ -52,7 +80,13 @@ describe('buildReplayExecutionState', () => { authContextHash: 'hash_123', }; - const state = build({ authContext }); + const state = build({ + authContext, + bridgeWorkerId: 'code-user_123', + sandboxBackend: 'remote-bridge', + executionProfile: 'stateful', + executionProfileSource: 'explicit', + }); expect(state).toMatchObject({ execution_id: 'exec_123', @@ -67,6 +101,10 @@ describe('buildReplayExecutionState', () => { principalSource: 'openid_reuse', authContextHash: 'hash_123', apiKeyId: 'key_legacy', + bridgeWorkerId: 'code-user_123', + sandboxBackend: 'remote-bridge', + executionProfile: 'stateful', + executionProfileSource: 'explicit', mode: 'replay', userCode: 'print("hello")', tools: TOOLS, diff --git a/service/src/service/programmatic-state.ts b/service/src/service/programmatic-state.ts index f469606f..25571fed 100644 --- a/service/src/service/programmatic-state.ts +++ b/service/src/service/programmatic-state.ts @@ -2,6 +2,26 @@ import type * as t from '../types'; import type { LCTool } from '../preamble'; import type { ExecutionState } from './replay-state'; import { buildExecutionIdentity, type ExecutionIdentity } from '../execution-identity'; +import { resolveQueuedSandboxBackend } from '../execution-profile'; +import type { + ExecutionProfile, + ExecutionProfileSource, + SandboxBackendName, +} from '../execution-profile'; + +export function resolveReplayStateSandboxBackend(params: { + executionProfile: ExecutionProfile; + executionProfileSource: ExecutionProfileSource; + apiSandboxBackend: SandboxBackendName; + bridgeWorkerId?: string; +}): SandboxBackendName | undefined { + if (params.bridgeWorkerId != null) return 'remote-bridge'; + return resolveQueuedSandboxBackend( + params.executionProfile, + params.apiSandboxBackend, + params.executionProfileSource, + ); +} export interface BuildReplayExecutionStateParams { executionId: string; @@ -17,6 +37,10 @@ export interface BuildReplayExecutionStateParams { isPyPlot: boolean; timeout: number; language: 'python' | 'bash'; + bridgeWorkerId?: string; + sandboxBackend?: SandboxBackendName; + executionProfile: ExecutionProfile; + executionProfileSource: ExecutionProfileSource; now?: number; } @@ -41,6 +65,10 @@ export function buildReplayExecutionState( principalSource: identity.principalSource, authContextHash: identity.authContextHash, apiKeyId: params.apiKeyId, + bridgeWorkerId: params.bridgeWorkerId, + sandboxBackend: params.sandboxBackend, + executionProfile: params.executionProfile, + executionProfileSource: params.executionProfileSource, startTime: now, lastActivity: now, mode: 'replay', diff --git a/service/src/service/replay-state.ts b/service/src/service/replay-state.ts index 562e06dd..4d2a0501 100644 --- a/service/src/service/replay-state.ts +++ b/service/src/service/replay-state.ts @@ -24,6 +24,11 @@ import { nanoid } from 'nanoid'; import type { Redis } from 'ioredis'; import type * as t from '../types'; import type { LCTool } from '../preamble'; +import type { + ExecutionProfile, + ExecutionProfileSource, + SandboxBackendName, +} from '../execution-profile'; import { connection } from '../queue'; import { env } from '../config'; import { internalServiceHeaders } from '../internal-service-auth'; @@ -109,6 +114,14 @@ export interface ExecutionState { * after one `EXECUTION_STATE_TTL` window post a trusted-source * apiKeyId invariant. */ apiKeyId?: string; + /** Authenticated worker selection retained across every replay iteration. */ + bridgeWorkerId?: string; + /** Original queue/backend target retained across replay continuations. */ + sandboxBackend?: SandboxBackendName; + /** Original producer profile retained so continuations use the same queue. */ + executionProfile?: ExecutionProfile; + /** Original profile source retained because inferred profiles use legacy queues. */ + executionProfileSource?: ExecutionProfileSource; startTime: number; /** * Wall-clock ms of the last interaction that advanced this execution (initial diff --git a/service/src/service/router.ts b/service/src/service/router.ts index f355c2dd..542fe4c0 100644 --- a/service/src/service/router.ts +++ b/service/src/service/router.ts @@ -25,7 +25,13 @@ import { Jobs, Languages } from '../enum'; import { FileRefAuthorizationError, authorizeRequestedFiles } from './file-authorization'; import { createUploadSessionRegistrar } from './upload-session'; import { prepareSandboxJobSecurity } from '../sandbox-egress'; +import { + BridgeWorkerSelectionError, + CODEAPI_BRIDGE_WORKER_HEADER, + resolveBridgeWorkerSelection, +} from '../bridge/selection'; import logger from '../logger'; +import { resolveQueuedSandboxBackend } from '../execution-profile'; const { INSTANCE_ID } = env; const JOB_COMPLETION_WAIT_TIMEOUT_MS = jobCompletionWaitTimeoutMs( @@ -140,6 +146,25 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) return res.status(400).json({ error: `Unsupported language: ${rawLang}` }); } + let bridgeWorkerId: string | undefined; + try { + const bridgeSelection = resolveBridgeWorkerSelection({ + backend: env.SANDBOX_BACKEND, + configuredWorkerId: env.BRIDGE_WORKER_ID, + dynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, + requestedWorkerId: req.header(CODEAPI_BRIDGE_WORKER_HEADER), + trustedWorkerId: principal.codeWorkerId, + }); + bridgeWorkerId = bridgeSelection?.explicit === true + ? bridgeSelection.workerId + : undefined; + } catch (error) { + if (error instanceof BridgeWorkerSelectionError) { + return res.status(error.status).json({ error: error.message }); + } + throw error; + } + let runtimeSessionId: string | undefined; try { runtimeSessionId = resolveRuntimeSessionIdForExecRequest({ @@ -247,6 +272,12 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, + sandboxBackend: resolveQueuedSandboxBackend( + env.EXECUTION_PROFILE, + env.SANDBOX_BACKEND, + env.EXECUTION_PROFILE_SOURCE, + ), + ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), ...(runtimeSessionId != null ? { runtimeSessionId } : {}), runtimeSessionMode, executionManifestClaims: sandboxSecurity.executionManifestClaims, diff --git a/service/src/types/service.ts b/service/src/types/service.ts index a642dda3..555056d5 100644 --- a/service/src/types/service.ts +++ b/service/src/types/service.ts @@ -3,7 +3,7 @@ import type { Request } from 'express'; import type { ExecutionManifestClaims } from '../execution-manifest'; import type { ExecutionIdentity } from '../execution-identity'; import type { CodeApiPrincipal } from '../auth/principal'; -import type { ExecutionProfile } from '../execution-profile'; +import type { ExecutionProfile, SandboxBackendName } from '../execution-profile'; import { Jobs } from '@/enum/service'; /** @@ -251,8 +251,12 @@ export type JobData = { executionId?: string; tenantId?: string; canonicalUserId?: string; + /** Trusted dynamic outbound worker selection. */ + bridgeWorkerId?: string; /** Producer deployment identity. Optional only for pre-profile queued jobs. */ executionProfile?: ExecutionProfile; + /** Required sandbox transport. Optional only for jobs queued before fencing. */ + sandboxBackend?: SandboxBackendName; /** * Server-derived runtime session identity. Absence is stateless unless * strict mode requires it; explicit exemptions document intentional gaps. diff --git a/service/src/utils.test.ts b/service/src/utils.test.ts index 8aa0dac0..be082040 100644 --- a/service/src/utils.test.ts +++ b/service/src/utils.test.ts @@ -48,7 +48,7 @@ describe('isValidResourceId (heterogeneous resource identifiers)', () => { expect(isValidResourceId('682f49b90f07376815c38ef2')).toBe(true); }); - test("accepts 17-char `agent_` slug", () => { + test('accepts 17-char `agent_` slug', () => { expect(isValidResourceId('agent_abc12345678')).toBe(true); }); @@ -145,15 +145,39 @@ describe('sandbox error formatting', () => { }); }); - test('maps remote bridge failures without exposing worker details', () => { + test('maps bridge authorization and availability failures without leaking worker details', () => { + const unauthorized = publicExecutionFailure( + new Error('BRIDGE_WORKER_UNAUTHORIZED: Worker private-vm belongs to tenant-secret'), + ); + expect(unauthorized).toEqual({ + status: 403, + body: { + error: 'bridge_worker_unauthorized', + message: 'Code environment is not authorized for this tenant', + }, + }); + expect(JSON.stringify(unauthorized)).not.toContain('private-vm'); + expect(JSON.stringify(unauthorized)).not.toContain('tenant-secret'); + + expect( + publicExecutionFailure( + new Error('BRIDGE_WORKER_OFFLINE: Worker private-vm has not checked in'), + ), + ).toEqual({ + status: 503, + body: { + error: 'bridge_worker_offline', + message: 'Code environment is offline', + }, + }); + const cases = [ - ['BRIDGE_WORKER_OFFLINE', 503, 'Remote code worker is unavailable'], - ['BRIDGE_WORKER_BUSY', 409, 'Remote code worker is busy'], - ['BRIDGE_EXECUTION_FAILED', 502, 'Remote code execution failed'], + ['BRIDGE_WORKER_BUSY', 409, 'Code environment is busy'], + ['BRIDGE_EXECUTION_FAILED', 502, 'Code environment execution failed'], [ 'BRIDGE_DEADLINE_EXCEEDED', 504, - 'Remote code execution deadline exceeded', + 'Code environment execution timed out', ], ] as const; for (const [code, status, message] of cases) { diff --git a/service/src/utils.ts b/service/src/utils.ts index b94d8fd2..e34d8411 100644 --- a/service/src/utils.ts +++ b/service/src/utils.ts @@ -131,12 +131,13 @@ export function publicExecutionFailure(error: unknown): { status: number; body: * MicroVM, and bridge codes describe sandbox availability; SESSION_INPUT_* codes * describe the caller's declared input set or its upstream object source. */ const backendMatch = message.match( - /^(RUNTIME_SESSION_BUSY|MICROVM_[A-Z_]+|BRIDGE_[A-Z_]+|SESSION_INPUT_[A-Z_]+):\s*(.+)$/, + /^(RUNTIME_SESSION_BUSY|BRIDGE_[A-Z_]+|MICROVM_[A-Z_]+|SESSION_INPUT_[A-Z_]+):\s*(.+)$/, ); if (backendMatch) { const code = backendMatch[1]; const statuses: Record = { RUNTIME_SESSION_BUSY: 409, + BRIDGE_WORKER_UNAUTHORIZED: 403, BRIDGE_WORKER_OFFLINE: 503, BRIDGE_WORKER_BUSY: 409, BRIDGE_EXECUTION_FAILED: 502, @@ -151,10 +152,11 @@ export function publicExecutionFailure(error: unknown): { status: number; body: const status = statuses[code] ?? (sessionInputFailure ? 500 : 503); const publicMessages: Record = { RUNTIME_SESSION_BUSY: 'Runtime session is busy', - BRIDGE_WORKER_OFFLINE: 'Remote code worker is unavailable', - BRIDGE_WORKER_BUSY: 'Remote code worker is busy', - BRIDGE_EXECUTION_FAILED: 'Remote code execution failed', - BRIDGE_DEADLINE_EXCEEDED: 'Remote code execution deadline exceeded', + BRIDGE_WORKER_UNAUTHORIZED: 'Code environment is not authorized for this tenant', + BRIDGE_WORKER_OFFLINE: 'Code environment is offline', + BRIDGE_WORKER_BUSY: 'Code environment is busy', + BRIDGE_EXECUTION_FAILED: 'Code environment execution failed', + BRIDGE_DEADLINE_EXCEEDED: 'Code environment execution timed out', MICROVM_LAUNCH_FAILED: 'Sandbox launch failed', MICROVM_LAUNCH_THROTTLED: 'Sandbox capacity is temporarily unavailable', MICROVM_UNHEALTHY: 'Sandbox runtime is unavailable', diff --git a/service/src/workers.ts b/service/src/workers.ts index ed2a46dd..d2048dc8 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -17,7 +17,10 @@ import { isSyntheticPrincipalSource } from './auth/synthetic'; import { withSpan, withTraceContext } from './telemetry'; import { workerDeadlineFailure } from './worker-error'; import logger from './logger'; -import { validateQueuedExecutionProfile } from './execution-profile'; +import { + validateQueuedExecutionProfile, + validateQueuedSandboxBackend, +} from './execution-profile'; const { INSTANCE_ID } = env; const WORKER_ID = `${INSTANCE_ID}-${process.pid}`; @@ -38,7 +41,7 @@ async function processJob(job: t.ExecuteJob): Promise { } async function processJobInner(job: t.ExecuteJob): Promise { - const { code, payload, isPyPlot } = job.data; + const { payload, isPyPlot } = job.data; const isSyntheticJob = job.data.isSynthetic === true || isSyntheticPrincipalSource(job.data.principalSource); const language = payload?.language ?? 'unknown'; const endTimer = jobProcessingDuration.startTimer({ language }); @@ -60,6 +63,11 @@ async function processJobInner(job: t.ExecuteJob): Promise { throw new Error(`Job timed out after ${env.JOB_TIMEOUT}ms`); } validateQueuedExecutionProfile(job.data.executionProfile, env.EXECUTION_PROFILE); + validateQueuedSandboxBackend( + job.data.sandboxBackend, + env.SANDBOX_BACKEND, + job.data.bridgeWorkerId, + ); let sandboxPayload = payload; let executionManifestClaims = job.data.executionManifestClaims; let egressGrantToken = job.data.egressGrantToken; @@ -139,6 +147,7 @@ async function processJobInner(job: t.ExecuteJob): Promise { deadlineAtMs, tenantId: job.data.tenantId, canonicalUserId: job.data.canonicalUserId, + bridgeWorkerId: job.data.bridgeWorkerId, runtimeSessionId: runtimeSession.runtimeSessionId, runtimeSessionMode: runtimeSession.runtimeSessionMode, /* Stateful backends run this as a commit barrier after user code but