From b2ca6fef6d08c33ab060d74d3171bc0b2ae92b76 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 12:24:38 -0400 Subject: [PATCH 01/13] feat: add principal-bound bridge workers --- docs/remote-bridge/README.md | 46 ++++++- service/src/bridge/index.ts | 1 + service/src/bridge/pairing.test.ts | 40 ++++++ service/src/bridge/pairing.ts | 43 +++++-- service/src/bridge/router.test.ts | 114 ++++++++++++++++++ service/src/bridge/router.ts | 74 ++++++++++-- service/src/bridge/selection.test.ts | 59 +++++++++ service/src/bridge/selection.ts | 51 ++++++++ service/src/bridge/store.test.ts | 28 +++++ service/src/bridge/store.ts | 26 +++- service/src/config.ts | 2 + .../src/sandbox-backend/remote-bridge.test.ts | 74 ++++++++++++ service/src/sandbox-backend/remote-bridge.ts | 16 ++- service/src/sandbox-backend/types.ts | 3 + service/src/secure-startup.test.ts | 19 +++ service/src/secure-startup.ts | 10 +- service/src/service/router.ts | 24 ++++ service/src/types/service.ts | 2 + service/src/utils.test.ts | 36 +++++- service/src/utils.ts | 12 +- service/src/workers.ts | 3 +- 21 files changed, 644 insertions(+), 39 deletions(-) create mode 100644 service/src/bridge/selection.test.ts create mode 100644 service/src/bridge/selection.ts create mode 100644 service/src/sandbox-backend/remote-bridge.test.ts diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index e5bddc91..86107e2e 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -32,6 +32,21 @@ 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. The trusted +LibreChat-to-Code-API request selects a worker with +`X-LibreChat-Code-Worker-ID`; Code API validates the identifier before it +crosses the queue boundary and requires that worker's stored tenant binding +before creating a lease. + Create a single-use pairing code with the administrator secret: ```bash @@ -41,6 +56,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 +121,8 @@ 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. +- 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 +142,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/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..4f974bab 100644 --- a/service/src/bridge/pairing.test.ts +++ b/service/src/bridge/pairing.test.ts @@ -17,6 +17,46 @@ 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 rotated = await pairings.rotate('vm-bound'); + 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), + }; + }; + + await expect( + pairings.authorize(requestFor(rotated.credential, 'bound-worker-proof')), + ).resolves.toEqual({ workerId: 'vm-bound', binding }); + await expect( + pairings.authorize(requestFor(issued.credential, 'superseded-bound-proof')), + ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); + }); + test('redeems a pairing code exactly once for the intended worker identity', async () => { const identity = createBridgeIdentity(); const pairing = await pairings.issue('vm-1'); diff --git a/service/src/bridge/pairing.ts b/service/src/bridge/pairing.ts index 54aaf5ab..b1d7cd4b 100644 --- a/service/src/bridge/pairing.ts +++ b/service/src/bridge/pairing.ts @@ -28,9 +28,20 @@ redis.call('SET', KEYS[4], ARGV[5], 'EX', ARGV[4]) 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 { @@ -38,6 +49,7 @@ interface StoredCredential { identityId: string; publicKey: string; expiresAt: string; + binding?: BridgeWorkerBinding; } export interface BridgePairing { @@ -75,10 +87,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}`; } @@ -110,12 +118,15 @@ 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 }; + const pairing: StoredPairing = { workerId, expiresAt, binding }; await this.redis.set( pairingKey(code), JSON.stringify(pairing), @@ -151,7 +162,13 @@ export class RedisBridgePairingStore { ); } - return await this.issueCredential(args.workerId, args.publicKey); + return await this.issueCredential( + args.workerId, + args.publicKey, + undefined, + undefined, + pairing.binding, + ); } async authorize(args: { @@ -168,6 +185,7 @@ export class RedisBridgePairingStore { credentialId: string; activeCredentialId: string; identityId: string; + binding?: BridgeWorkerBinding; }> { const proofTime = Date.parse(args.timestamp); if ( @@ -235,6 +253,7 @@ export class RedisBridgePairingStore { credentialId: credentialDigest, activeCredentialId: activeDigest, identityId: stored.identityId, + ...(stored.binding ? { binding: stored.binding } : {}), }; } @@ -271,6 +290,7 @@ export class RedisBridgePairingStore { previous.publicKey, previousDigest, previous.identityId, + previous.binding, ); } @@ -279,13 +299,20 @@ export class RedisBridgePairingStore { publicKey: string, previousDigest?: string, identityId = randomBytes(18).toString('base64url'), + binding?: BridgeWorkerBinding, ): 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 stored: StoredCredential = { + workerId, + identityId, + publicKey, + expiresAt, + binding, + }; if (previousDigest !== undefined) { const rotated = await this.redis.eval( ROTATE_CREDENTIAL_SCRIPT, diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts index 077dcd1f..f9e77a20 100644 --- a/service/src/bridge/router.test.ts +++ b/service/src/bridge/router.test.ts @@ -25,6 +25,120 @@ afterEach(async () => { }); describe('paired bridge HTTP API', () => { + 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', + capabilities: { + statefulWorkspace: true, + 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: '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..730097b2 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -4,15 +4,23 @@ 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 PRINCIPAL_TYPES = new Set([ + 'deployment', + 'tenant', + 'user', + 'role', + 'group', +]); export type BridgeAuthMode = 'static' | 'paired'; @@ -22,6 +30,7 @@ export interface BridgeRouterOptions { authMode: BridgeAuthMode; adminToken: string; configuredWorkerId?: string; + allowDynamicWorkers?: boolean; } function sameToken(left: string, right: string): boolean { @@ -34,7 +43,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,6 +54,28 @@ 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_WORKER_ID_PATTERN.test(tenantId) || + typeof principal.type !== 'string' || + !PRINCIPAL_TYPES.has(principal.type as BridgePrincipalType) || + typeof principal.id !== 'string' || + !BRIDGE_WORKER_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; res.status(status).json({ error: error.message, code: error.code }); @@ -77,8 +108,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 +216,12 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { res.status(400).json({ error: 'Invalid bridge worker ID' }); return; } - const pairing = await options.pairings.issue(workerId); + const binding = isRecord(req.body) ? parseBinding(req.body.binding) : undefined; + 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 +339,33 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { }); return; } + const authorization = res.locals.bridgeWorkerAuthorization as + | { + workerId: string; + binding?: BridgeWorkerBinding; + } + | undefined; + const capabilities = registration.capabilities; + const trustedRegistration: BridgeWorkerRegistration = { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: registration.workerId, + incarnationId: registration.incarnationId, + capabilities: { + statefulWorkspace: true, + 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?.binding != null + ? { binding: authorization.binding } + : {}), + }); } catch (error) { if (error instanceof BridgeStoreError) { sendStoreError(error, res); diff --git a/service/src/bridge/selection.test.ts b/service/src/bridge/selection.test.ts new file mode 100644 index 00000000..29376d9d --- /dev/null +++ b/service/src/bridge/selection.test.ts @@ -0,0 +1,59 @@ +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', dynamic: false }); + }); + + test('selects a valid dynamic worker only when dynamic routing is enabled', () => { + expect( + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: true, + requestedWorkerId: 'code-user_1', + }), + ).toEqual({ workerId: 'code-user_1', dynamic: true }); + }); + + test('rejects dynamic routing on the wrong backend or when it is disabled', () => { + expect(() => + resolveBridgeWorkerSelection({ + backend: 'http', + configuredWorkerId: '', + dynamicWorkers: true, + requestedWorkerId: 'code-user-1', + }), + ).toThrow(BridgeWorkerSelectionError); + expect(() => + resolveBridgeWorkerSelection({ + backend: 'remote-bridge', + configuredWorkerId: 'deployment-worker', + dynamicWorkers: false, + requestedWorkerId: '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', + }), + ).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..4ade6a53 --- /dev/null +++ b/service/src/bridge/selection.ts @@ -0,0 +1,51 @@ +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; +}): { workerId: string; dynamic: boolean } | undefined { + const requestedWorkerId = args.requestedWorkerId?.trim(); + if (requestedWorkerId != null && requestedWorkerId.length > 0) { + if (args.backend !== 'remote-bridge') { + throw new BridgeWorkerSelectionError( + 'Code bridge worker routing requires the remote-bridge backend', + 400, + ); + } + if (!BRIDGE_WORKER_ID_PATTERN.test(requestedWorkerId)) { + throw new BridgeWorkerSelectionError('Invalid code bridge worker ID', 400); + } + if (!args.dynamicWorkers && requestedWorkerId !== args.configuredWorkerId) { + throw new BridgeWorkerSelectionError('Dynamic code bridge workers are disabled', 403); + } + return { + workerId: requestedWorkerId, + dynamic: requestedWorkerId !== args.configuredWorkerId, + }; + } + + 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, dynamic: 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..4956b5b8 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -14,6 +14,34 @@ afterEach(async () => { }); describe('RedisBridgeStore', () => { + test('rejects a dynamic worker lease outside its bound tenant', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'tenant-worker', + 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('delivers and settles one fenced stateful assignment', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 04fb8e29..9d16dd07 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' @@ -47,6 +49,10 @@ interface StoredAssignment extends CodeBridgeAssignment { leaseTokenHash: string; } +export interface RegisteredBridgeWorker extends BridgeWorkerRegistration { + binding?: BridgeWorkerBinding; +} + function workerKey(workerId: string): string { return `${PREFIX}:worker:${workerId}`; } @@ -132,7 +138,7 @@ export class RedisBridgeStore { private readonly workerTtlSeconds = DEFAULT_WORKER_TTL_SECONDS, ) {} - async register(registration: BridgeWorkerRegistration): Promise { + async register(registration: RegisteredBridgeWorker): Promise { const script = [ 'if redis.call(\'EXISTS\', KEYS[3]) == 1 then return -2 end', 'if redis.call(\'EXISTS\', KEYS[2]) == 1 then return -1 end', @@ -176,6 +182,8 @@ export class RedisBridgeStore { async dispatch(args: { workerId: string; + tenantId?: string; + requireTenantBinding?: boolean; body: t.PayloadBody; headers: Record; runtimeSessionId?: string; @@ -192,6 +200,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 @@ -403,9 +423,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/sandbox-backend/remote-bridge.test.ts b/service/src/sandbox-backend/remote-bridge.test.ts new file mode 100644 index 00000000..4ff52fb9 --- /dev/null +++ b/service/src/sandbox-backend/remote-bridge.test.ts @@ -0,0 +1,74 @@ +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), + 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', + }); + }); +}); diff --git a/service/src/sandbox-backend/remote-bridge.ts b/service/src/sandbox-backend/remote-bridge.ts index a30ab748..1134e4e2 100644 --- a/service/src/sandbox-backend/remote-bridge.ts +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -15,7 +15,7 @@ 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, ) {} @@ -23,7 +23,8 @@ export class RemoteBridgeSandboxBackend implements SandboxBackend { 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 +33,9 @@ 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, body: req.body, headers: req.headers, runtimeSessionId: ctx.runtimeSessionId, @@ -57,6 +60,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..c970a0fa 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -15,6 +15,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 +57,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 +315,23 @@ 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('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..7e4e1ee4 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -119,7 +119,15 @@ export function validateSandboxBackendPolicy(): void { ); } if (env.SANDBOX_BACKEND === 'remote-bridge') { - requireValue('CODEAPI_BRIDGE_WORKER_ID', env.BRIDGE_WORKER_ID); + if (env.BRIDGE_DYNAMIC_WORKERS) { + if (env.BRIDGE_AUTH_MODE !== 'paired') { + throw new SecureStartupConfigError( + 'Dynamic remote bridge workers require CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + } + } else { + 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/router.ts b/service/src/service/router.ts index f355c2dd..65a3dc38 100644 --- a/service/src/service/router.ts +++ b/service/src/service/router.ts @@ -25,6 +25,11 @@ 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'; const { INSTANCE_ID } = env; @@ -140,6 +145,24 @@ 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), + }); + bridgeWorkerId = bridgeSelection?.dynamic === 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 +270,7 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, + ...(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..68127b8e 100644 --- a/service/src/types/service.ts +++ b/service/src/types/service.ts @@ -251,6 +251,8 @@ 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; /** 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..d0c9debc 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -38,7 +38,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 }); @@ -139,6 +139,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 From 7c6d6a7a56f42983ba57c912b99147ee4d40fb0d Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 15:51:23 -0400 Subject: [PATCH 02/13] fix: authenticate principal worker routing --- docs/remote-bridge/README.md | 12 +++-- service/src/auth/librechat-jwt.test.ts | 3 ++ service/src/auth/librechat-jwt.ts | 3 ++ service/src/auth/principal.ts | 1 + service/src/bridge/pairing.test.ts | 29 ++++++++++- service/src/bridge/pairing.ts | 46 +++++++++++++++--- service/src/bridge/router.ts | 15 +++++- service/src/bridge/selection.test.ts | 48 +++++++++++++++++-- service/src/bridge/selection.ts | 37 ++++++++++---- service/src/bridge/store.test.ts | 33 +++++++++++++ service/src/bridge/store.ts | 13 ++++- service/src/service/programmatic-router.ts | 40 ++++++++++++++-- .../src/service/programmatic-state.test.ts | 7 ++- service/src/service/programmatic-state.ts | 2 + service/src/service/replay-state.ts | 2 + service/src/service/router.ts | 3 +- 16 files changed, 260 insertions(+), 34 deletions(-) diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index 86107e2e..2f4e71a6 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -41,11 +41,13 @@ CODEAPI_BRIDGE_AUTH_MODE=paired # CODEAPI_BRIDGE_WORKER_ID=my-default-vm ``` -Dynamic routing is accepted only with paired authentication. The trusted -LibreChat-to-Code-API request selects a worker with -`X-LibreChat-Code-Worker-ID`; Code API validates the identifier before it -crosses the queue boundary and requires that worker's stored tenant binding -before creating a lease. +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: 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/pairing.test.ts b/service/src/bridge/pairing.test.ts index 4f974bab..950ba849 100644 --- a/service/src/bridge/pairing.test.ts +++ b/service/src/bridge/pairing.test.ts @@ -51,7 +51,7 @@ describe('RedisBridgePairingStore', () => { await expect( pairings.authorize(requestFor(rotated.credential, 'bound-worker-proof')), - ).resolves.toEqual({ workerId: 'vm-bound', binding }); + ).resolves.toMatchObject({ workerId: 'vm-bound', binding }); await expect( pairings.authorize(requestFor(issued.credential, 'superseded-bound-proof')), ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); @@ -78,6 +78,33 @@ 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('authorizes a credential only with proof from its worker key', async () => { const identity = createBridgeIdentity(); const pairing = await pairings.issue('vm-1'); diff --git a/service/src/bridge/pairing.ts b/service/src/bridge/pairing.ts index b1d7cd4b..892ff5c2 100644 --- a/service/src/bridge/pairing.ts +++ b/service/src/bridge/pairing.ts @@ -27,6 +27,27 @@ redis.call('SET', KEYS[1], ARGV[2], 'EX', ARGV[4]) redis.call('SET', KEYS[4], ARGV[5], 'EX', ARGV[4]) return 1 `; +const ISSUE_PAIRING_SCRIPT = ` +local previous = redis.call('GET', KEYS[1]) +if previous then + redis.call('DEL', previous) +end +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]) +return pairing +`; export type BridgePrincipalType = 'deployment' | 'tenant' | 'user' | 'role' | 'group'; @@ -99,6 +120,10 @@ function workerStableIdentityKey(workerId: string): string { return `${PREFIX}:stable-identity:${workerId}`; } +function workerPairingIndexKey(workerId: string): string { + return `${PREFIX}:pairing-index:${workerId}`; +} + function proofNonceKey(credential: string, nonce: string): string { return `${PREFIX}:proof:${digest(credential)}:${digest(nonce)}`; } @@ -127,11 +152,14 @@ export class RedisBridgePairingStore { Date.now() + this.pairingTtlSeconds * 1000, ).toISOString(); const pairing: StoredPairing = { workerId, expiresAt, binding }; - await this.redis.set( - pairingKey(code), + const codeKey = pairingKey(code); + await this.redis.eval( + ISSUE_PAIRING_SCRIPT, + 2, + workerPairingIndexKey(workerId), + codeKey, JSON.stringify(pairing), - 'EX', - this.pairingTtlSeconds, + String(this.pairingTtlSeconds), ); return { workerId, code, expiresAt }; } @@ -141,8 +169,14 @@ export class RedisBridgePairingStore { code: string; publicKey: string; }): Promise { - const raw = await this.redis.getdel(pairingKey(args.code)); - if (raw == null) { + const codeKey = pairingKey(args.code); + const raw = await this.redis.eval( + REDEEM_PAIRING_SCRIPT, + 2, + codeKey, + workerPairingIndexKey(args.workerId), + ); + if (typeof raw !== 'string') { throw new BridgePairingError( 'PAIRING_INVALID', 'Pairing code is invalid or expired', diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 730097b2..b426cc05 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -14,6 +14,7 @@ import { BridgeStoreError, RedisBridgeStore } from './store'; 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', @@ -59,11 +60,11 @@ function parseBinding(value: unknown): BridgeWorkerBinding | undefined { const { tenantId, principal } = value; if ( typeof tenantId !== 'string' || - !BRIDGE_WORKER_ID_PATTERN.test(tenantId) || + !BRIDGE_BINDING_ID_PATTERN.test(tenantId) || typeof principal.type !== 'string' || !PRINCIPAL_TYPES.has(principal.type as BridgePrincipalType) || typeof principal.id !== 'string' || - !BRIDGE_WORKER_ID_PATTERN.test(principal.id) + !BRIDGE_BINDING_ID_PATTERN.test(principal.id) ) { return undefined; } @@ -342,6 +343,7 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { const authorization = res.locals.bridgeWorkerAuthorization as | { workerId: string; + credentialId: string; binding?: BridgeWorkerBinding; } | undefined; @@ -362,6 +364,9 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { try { await options.store.register({ ...trustedRegistration, + ...(authorization?.credentialId != null + ? { credentialId: authorization.credentialId } + : {}), ...(authorization?.binding != null ? { binding: authorization.binding } : {}), @@ -411,6 +416,12 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { workerId, body.incarnationId, Math.min(requestedWait, MAX_LEASE_WAIT_MS), + undefined, + ( + res.locals.bridgeWorkerAuthorization as + | { credentialId: string } + | undefined + )?.credentialId, ); res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); } catch (error) { diff --git a/service/src/bridge/selection.test.ts b/service/src/bridge/selection.test.ts index 29376d9d..b6309702 100644 --- a/service/src/bridge/selection.test.ts +++ b/service/src/bridge/selection.test.ts @@ -13,18 +13,48 @@ describe('bridge worker request selection', () => { configuredWorkerId: 'deployment-worker', dynamicWorkers: true, }), - ).toEqual({ workerId: 'deployment-worker', dynamic: false }); + ).toEqual({ workerId: 'deployment-worker', explicit: false }); }); - test('selects a valid dynamic worker only when dynamic routing is enabled', () => { + 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', dynamic: true }); + ).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', () => { @@ -34,6 +64,7 @@ describe('bridge worker request selection', () => { configuredWorkerId: '', dynamicWorkers: true, requestedWorkerId: 'code-user-1', + trustedWorkerId: 'code-user-1', }), ).toThrow(BridgeWorkerSelectionError); expect(() => @@ -42,6 +73,7 @@ describe('bridge worker request selection', () => { configuredWorkerId: 'deployment-worker', dynamicWorkers: false, requestedWorkerId: 'code-user-1', + trustedWorkerId: 'code-user-1', }), ).toThrow('Dynamic code bridge workers are disabled'); }); @@ -53,6 +85,16 @@ describe('bridge worker request selection', () => { 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 index 4ade6a53..0959279f 100644 --- a/service/src/bridge/selection.ts +++ b/service/src/bridge/selection.ts @@ -1,5 +1,5 @@ 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 const BRIDGE_WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$/; export class BridgeWorkerSelectionError extends Error { constructor( @@ -16,24 +16,45 @@ export function resolveBridgeWorkerSelection(args: { configuredWorkerId: string; dynamicWorkers: boolean; requestedWorkerId?: string; -}): { workerId: string; dynamic: boolean } | undefined { + trustedWorkerId?: string; +}): { workerId: string; explicit: boolean } | undefined { const requestedWorkerId = args.requestedWorkerId?.trim(); - if (requestedWorkerId != null && requestedWorkerId.length > 0) { + 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 (!BRIDGE_WORKER_ID_PATTERN.test(requestedWorkerId)) { + 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 && requestedWorkerId !== args.configuredWorkerId) { + if (!args.dynamicWorkers && selectedWorkerId !== args.configuredWorkerId) { throw new BridgeWorkerSelectionError('Dynamic code bridge workers are disabled', 403); } return { - workerId: requestedWorkerId, - dynamic: requestedWorkerId !== args.configuredWorkerId, + workerId: selectedWorkerId, + explicit: true, }; } @@ -45,7 +66,7 @@ export function resolveBridgeWorkerSelection(args: { if (!BRIDGE_WORKER_ID_PATTERN.test(configuredWorkerId)) { throw new BridgeWorkerSelectionError('Invalid configured code bridge worker ID', 503); } - return { workerId: configuredWorkerId, dynamic: false }; + 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 4956b5b8..0abb42a7 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -42,6 +42,39 @@ describe('RedisBridgeStore', () => { ).rejects.toMatchObject({ code: 'WORKER_UNAUTHORIZED' }); }); + test('does not lease an assignment to a newly rebound worker credential', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'rebound-worker', + credentialId: 'tenant-a-credential', + 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', 1_000, undefined, 'tenant-b-credential'), + ).resolves.toBeUndefined(); + 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, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 9d16dd07..356492e2 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -47,10 +47,12 @@ export class BridgeStoreError extends Error { interface StoredAssignment extends CodeBridgeAssignment { leaseTokenHash: string; + workerCredentialId?: string; } export interface RegisteredBridgeWorker extends BridgeWorkerRegistration { binding?: BridgeWorkerBinding; + credentialId?: string; } function workerKey(workerId: string): string { @@ -261,6 +263,9 @@ export class RedisBridgeStore { generation, leaseToken, leaseTokenHash: tokenHash(leaseToken), + ...(registration.credentialId != null + ? { workerCredentialId: registration.credentialId } + : {}), expiresAt: new Date(args.deadlineAtMs).toISOString(), runtimeSessionId: args.runtimeSessionId, request: { @@ -311,6 +316,7 @@ export class RedisBridgeStore { incarnationId: string, waitMs: number, signal?: AbortSignal, + credentialId?: string, ): Promise { const deadline = Date.now() + waitMs; while (signal?.aborted !== true && Date.now() < deadline) { @@ -332,8 +338,13 @@ export class RedisBridgeStore { 'Bridge worker incarnation was replaced', ); } + if (assignment.workerCredentialId !== credentialId) continue; if (Date.parse(assignment.expiresAt) <= Date.now()) continue; - const { leaseTokenHash: _leaseTokenHash, ...wireAssignment } = assignment; + const { + leaseTokenHash: _leaseTokenHash, + workerCredentialId: _workerCredentialId, + ...wireAssignment + } = assignment; return wireAssignment; } return undefined; diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index eade27fb..e6fb51fd 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -36,6 +36,11 @@ import { findUnregisteredToolCall } from '../tool-scope'; import { summarizeRequestedFiles } from '../execution-log'; import { FileRefAuthorizationError, authorizeRequestedFiles } from './file-authorization'; import { buildReplayExecutionState } from './programmatic-state'; +import { + BridgeWorkerSelectionError, + CODEAPI_BRIDGE_WORKER_HEADER, + resolveBridgeWorkerSelection, +} from '../bridge/selection'; import logger from '../logger'; import { type ExecutionState, @@ -408,6 +413,7 @@ async function runReplayIteration( tenantId: state.tenantId, canonicalUserId: state.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, + ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, executionManifestClaims: sandboxSecurity.executionManifestClaims, @@ -439,9 +445,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 +567,7 @@ async function handleReplayInitial( isPyPlot, timeout, language, + bridgeWorkerId, }); /** Replay mode persists the full request (`userCode` + `tools` + `files`) * inside `ExecutionState` so continuations can re-enqueue without the @@ -1023,6 +1031,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 +1107,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 +1127,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 +1310,7 @@ async function handleBlocking( principalSource: identity.principalSource, authContextHash: identity.authContextHash, apiKeyId, + bridgeWorkerId, startTime: Date.now(), lastActivity: Date.now(), mode: 'blocking', @@ -1378,6 +1407,7 @@ async function handleBlocking( tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, + ...(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..ec548476 100644 --- a/service/src/service/programmatic-state.test.ts +++ b/service/src/service/programmatic-state.test.ts @@ -22,7 +22,9 @@ const FILES = [ }, ] as RequestFile[]; -function build(overrides: Partial[0]> = {}) { +function build( + overrides: Partial[0]> = {}, +): ReturnType { return buildReplayExecutionState({ executionId: 'exec_123', sessionId: 'session_123', @@ -52,7 +54,7 @@ describe('buildReplayExecutionState', () => { authContextHash: 'hash_123', }; - const state = build({ authContext }); + const state = build({ authContext, bridgeWorkerId: 'code-user_123' }); expect(state).toMatchObject({ execution_id: 'exec_123', @@ -67,6 +69,7 @@ describe('buildReplayExecutionState', () => { principalSource: 'openid_reuse', authContextHash: 'hash_123', apiKeyId: 'key_legacy', + bridgeWorkerId: 'code-user_123', 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..3a5f2f5b 100644 --- a/service/src/service/programmatic-state.ts +++ b/service/src/service/programmatic-state.ts @@ -17,6 +17,7 @@ export interface BuildReplayExecutionStateParams { isPyPlot: boolean; timeout: number; language: 'python' | 'bash'; + bridgeWorkerId?: string; now?: number; } @@ -41,6 +42,7 @@ export function buildReplayExecutionState( principalSource: identity.principalSource, authContextHash: identity.authContextHash, apiKeyId: params.apiKeyId, + bridgeWorkerId: params.bridgeWorkerId, startTime: now, lastActivity: now, mode: 'replay', diff --git a/service/src/service/replay-state.ts b/service/src/service/replay-state.ts index 562e06dd..0d6b6709 100644 --- a/service/src/service/replay-state.ts +++ b/service/src/service/replay-state.ts @@ -109,6 +109,8 @@ 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; 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 65a3dc38..68489d6f 100644 --- a/service/src/service/router.ts +++ b/service/src/service/router.ts @@ -152,8 +152,9 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) configuredWorkerId: env.BRIDGE_WORKER_ID, dynamicWorkers: env.BRIDGE_DYNAMIC_WORKERS, requestedWorkerId: req.header(CODEAPI_BRIDGE_WORKER_HEADER), + trustedWorkerId: principal.codeWorkerId, }); - bridgeWorkerId = bridgeSelection?.dynamic === true + bridgeWorkerId = bridgeSelection?.explicit === true ? bridgeSelection.workerId : undefined; } catch (error) { From 488981ae23cb3c259c68ed78fdbde26daf9cdaae Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 16:10:50 -0400 Subject: [PATCH 03/13] fix: fence bridge identity and backend routing --- README.md | 4 +- docs/remote-bridge/README.md | 6 ++ service/src/bridge/pairing.test.ts | 89 +++++++++++++++++++++- service/src/bridge/pairing.ts | 17 ++++- service/src/bridge/router.ts | 8 +- service/src/bridge/store.test.ts | 45 ++++++++++- service/src/bridge/store.ts | 13 ++-- service/src/execution-profile.test.ts | 36 +++++++++ service/src/execution-profile.ts | 31 ++++++++ service/src/queue.ts | 1 + service/src/service/programmatic-router.ts | 2 + service/src/service/router.ts | 1 + service/src/types/service.ts | 4 +- service/src/workers.ts | 6 +- 14 files changed, 244 insertions(+), 19 deletions(-) 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 2f4e71a6..649a8257 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -123,6 +123,12 @@ 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. +- 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. diff --git a/service/src/bridge/pairing.test.ts b/service/src/bridge/pairing.test.ts index 950ba849..847c2563 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'; @@ -29,7 +30,6 @@ describe('RedisBridgePairingStore', () => { code: pairing.code, publicKey: identity.publicKey, }); - const rotated = await pairings.rotate('vm-bound'); const requestFor = ( credential: string, nonce: string, @@ -49,14 +49,60 @@ describe('RedisBridgePairingStore', () => { }; }; - await expect( - pairings.authorize(requestFor(rotated.credential, 'bound-worker-proof')), - ).resolves.toMatchObject({ workerId: 'vm-bound', binding }); + 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, 'superseded-bound-proof')), ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); }); + 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'); @@ -293,6 +339,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 892ff5c2..e7d3cd03 100644 --- a/service/src/bridge/pairing.ts +++ b/service/src/bridge/pairing.ts @@ -48,6 +48,18 @@ end redis.call('DEL', KEYS[1], KEYS[2]) return pairing `; +const ROTATE_CREDENTIAL_SCRIPT = ` +if redis.call('GET', KEYS[1]) ~= ARGV[1] then + return 0 +end +if redis.call('EXISTS', KEYS[2]) ~= 1 then + return 0 +end +redis.call('SET', KEYS[3], ARGV[3], 'EX', ARGV[4]) +redis.call('SET', KEYS[1], ARGV[2], 'EX', ARGV[4]) +redis.call('DEL', KEYS[2]) +return 1 +`; export type BridgePrincipalType = 'deployment' | 'tenant' | 'user' | 'role' | 'group'; @@ -307,7 +319,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 @@ -325,6 +338,7 @@ export class RedisBridgePairingStore { previousDigest, previous.identityId, previous.binding, + previous.identityId ?? null, ); } @@ -334,6 +348,7 @@ export class RedisBridgePairingStore { previousDigest?: string, identityId = randomBytes(18).toString('base64url'), binding?: BridgeWorkerBinding, + identityId?: string | null, ): Promise { const credential = randomBytes(32).toString('base64url'); const credentialDigest = digest(credential); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index b426cc05..a28caeb2 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -344,6 +344,7 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { | { workerId: string; credentialId: string; + identityId?: string; binding?: BridgeWorkerBinding; } | undefined; @@ -367,6 +368,9 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { ...(authorization?.credentialId != null ? { credentialId: authorization.credentialId } : {}), + ...(authorization?.identityId != null + ? { identityId: authorization.identityId } + : {}), ...(authorization?.binding != null ? { binding: authorization.binding } : {}), @@ -419,9 +423,9 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { undefined, ( res.locals.bridgeWorkerAuthorization as - | { credentialId: string } + | { identityId: string } | undefined - )?.credentialId, + )?.identityId, ); res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); } catch (error) { diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 0abb42a7..f968d795 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -42,11 +42,11 @@ describe('RedisBridgeStore', () => { ).rejects.toMatchObject({ code: 'WORKER_UNAUTHORIZED' }); }); - test('does not lease an assignment to a newly rebound worker credential', async () => { + test('does not lease an assignment to a newly rebound worker identity', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: 'rebound-worker', - credentialId: 'tenant-a-credential', + identityId: 'tenant-a-identity', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -69,12 +69,51 @@ describe('RedisBridgeStore', () => { }); await expect( - store.lease('rebound-worker', 1_000, undefined, 'tenant-b-credential'), + store.lease('rebound-worker', 1_000, undefined, 'tenant-b-identity'), ).resolves.toBeUndefined(); controller.abort(); await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); }); + test('leases queued work after credential refresh preserves the paired identity', async () => { + await store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'rotating-worker', + 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', + 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, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 356492e2..197cf293 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -47,12 +47,13 @@ export class BridgeStoreError extends Error { interface StoredAssignment extends CodeBridgeAssignment { leaseTokenHash: string; - workerCredentialId?: string; + workerIdentityId?: string; } export interface RegisteredBridgeWorker extends BridgeWorkerRegistration { binding?: BridgeWorkerBinding; credentialId?: string; + identityId?: string; } function workerKey(workerId: string): string { @@ -263,8 +264,8 @@ export class RedisBridgeStore { generation, leaseToken, leaseTokenHash: tokenHash(leaseToken), - ...(registration.credentialId != null - ? { workerCredentialId: registration.credentialId } + ...(registration.identityId != null + ? { workerIdentityId: registration.identityId } : {}), expiresAt: new Date(args.deadlineAtMs).toISOString(), runtimeSessionId: args.runtimeSessionId, @@ -316,7 +317,7 @@ export class RedisBridgeStore { incarnationId: string, waitMs: number, signal?: AbortSignal, - credentialId?: string, + identityId?: string, ): Promise { const deadline = Date.now() + waitMs; while (signal?.aborted !== true && Date.now() < deadline) { @@ -338,11 +339,11 @@ export class RedisBridgeStore { 'Bridge worker incarnation was replaced', ); } - if (assignment.workerCredentialId !== credentialId) continue; + if (assignment.workerIdentityId !== identityId) continue; if (Date.parse(assignment.expiresAt) <= Date.now()) continue; const { leaseTokenHash: _leaseTokenHash, - workerCredentialId: _workerCredentialId, + workerIdentityId: _workerIdentityId, ...wireAssignment } = assignment; return wireAssignment; diff --git a/service/src/execution-profile.test.ts b/service/src/execution-profile.test.ts index da20bce2..0e572a55 100644 --- a/service/src/execution-profile.test.ts +++ b/service/src/execution-profile.test.ts @@ -4,6 +4,7 @@ import { queueNamesForExecutionProfile, resolveExecutionProfile, resolveExecutionProfileSource, + validateQueuedSandboxBackend, validateQueuedExecutionProfile, } from './execution-profile'; @@ -55,6 +56,21 @@ describe('execution profile queue isolation', () => { other: 'stateful-other-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', + }); + }); }); describe('execution profile request assertion', () => { @@ -107,3 +123,23 @@ 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(); + }); + + 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', + ); + }); +}); diff --git a/service/src/execution-profile.ts b/service/src/execution-profile.ts index c4951903..3d39dd80 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,8 @@ export interface ExecutionProfileQueueNames { other: string; } +export type SandboxBackendName = typeof SANDBOX_BACKENDS[number]; + export function resolveExecutionProfile( raw: string | undefined, runtimeSessionMode: 'stateless' | 'affinity' | 'strict', @@ -54,10 +61,17 @@ const EXPLICIT_PROFILE_QUEUE_NAMES: Record(queueNames.python, { connection }); const otherQueue = new Queue(queueNames.other, { connection }); diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index e6fb51fd..daf042f8 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -413,6 +413,7 @@ async function runReplayIteration( tenantId: state.tenantId, canonicalUserId: state.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, + sandboxBackend: env.SANDBOX_BACKEND, ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, @@ -1407,6 +1408,7 @@ async function handleBlocking( tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, + sandboxBackend: env.SANDBOX_BACKEND, ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, diff --git a/service/src/service/router.ts b/service/src/service/router.ts index 68489d6f..01894793 100644 --- a/service/src/service/router.ts +++ b/service/src/service/router.ts @@ -271,6 +271,7 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, + sandboxBackend: env.SANDBOX_BACKEND, ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), ...(runtimeSessionId != null ? { runtimeSessionId } : {}), runtimeSessionMode, diff --git a/service/src/types/service.ts b/service/src/types/service.ts index 68127b8e..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'; /** @@ -255,6 +255,8 @@ export type JobData = { 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/workers.ts b/service/src/workers.ts index d0c9debc..e053c3bc 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}`; @@ -60,6 +63,7 @@ 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); let sandboxPayload = payload; let executionManifestClaims = job.data.executionManifestClaims; let egressGrantToken = job.data.egressGrantToken; From ad0acd52c80b41303f7b419b74e6e888aca53fe8 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 16:29:51 -0400 Subject: [PATCH 04/13] fix: fence bridge redemption and queue routing --- service/src/bridge/pairing.test.ts | 75 ++++++++++++++++++ service/src/bridge/pairing.ts | 79 +++++++++++-------- service/src/execution-profile.test.ts | 13 +++ service/src/execution-profile.ts | 22 +++++- .../src/sandbox-backend/remote-bridge.test.ts | 34 ++++++++ service/src/sandbox-backend/remote-bridge.ts | 3 +- service/src/service/programmatic-router.ts | 11 ++- service/src/service/router.ts | 6 +- service/src/workers.ts | 6 +- 9 files changed, 210 insertions(+), 39 deletions(-) diff --git a/service/src/bridge/pairing.test.ts b/service/src/bridge/pairing.test.ts index 847c2563..5fe0d0db 100644 --- a/service/src/bridge/pairing.test.ts +++ b/service/src/bridge/pairing.test.ts @@ -151,6 +151,81 @@ describe('RedisBridgePairingStore', () => { ).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'); diff --git a/service/src/bridge/pairing.ts b/service/src/bridge/pairing.ts index e7d3cd03..3ff3d30d 100644 --- a/service/src/bridge/pairing.ts +++ b/service/src/bridge/pairing.ts @@ -32,6 +32,7 @@ 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 @@ -46,18 +47,17 @@ if redis.call('GET', KEYS[2]) ~= KEYS[1] then return nil end redis.call('DEL', KEYS[1], KEYS[2]) +redis.call('SET', KEYS[3], ARGV[1], 'EX', ARGV[2]) return pairing `; -const ROTATE_CREDENTIAL_SCRIPT = ` +const INSTALL_REDEEMED_CREDENTIAL_SCRIPT = ` if redis.call('GET', KEYS[1]) ~= ARGV[1] then return 0 end -if redis.call('EXISTS', KEYS[2]) ~= 1 then - return 0 -end -redis.call('SET', KEYS[3], ARGV[3], 'EX', ARGV[4]) -redis.call('SET', KEYS[1], ARGV[2], 'EX', ARGV[4]) -redis.call('DEL', KEYS[2]) +redis.call('SET', KEYS[2], ARGV[3], 'EX', ARGV[4]) +redis.call('SET', KEYS[3], ARGV[2], 'EX', ARGV[4]) +redis.call('SET', KEYS[4], ARGV[5], 'EX', ARGV[4]) +redis.call('DEL', KEYS[1]) return 1 `; @@ -136,6 +136,10 @@ 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)}`; } @@ -167,9 +171,10 @@ export class RedisBridgePairingStore { const codeKey = pairingKey(code); await this.redis.eval( ISSUE_PAIRING_SCRIPT, - 2, + 3, workerPairingIndexKey(workerId), codeKey, + workerRedemptionKey(workerId), JSON.stringify(pairing), String(this.pairingTtlSeconds), ); @@ -181,12 +186,22 @@ export class RedisBridgePairingStore { code: string; publicKey: string; }): Promise { + 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, - 2, + 3, codeKey, workerPairingIndexKey(args.workerId), + workerRedemptionKey(args.workerId), + redemptionId, + String(this.pairingTtlSeconds), ); if (typeof raw !== 'string') { throw new BridgePairingError( @@ -201,19 +216,13 @@ 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, undefined, undefined, pairing.binding, + redemptionId, ); } @@ -338,7 +347,6 @@ export class RedisBridgePairingStore { previousDigest, previous.identityId, previous.binding, - previous.identityId ?? null, ); } @@ -348,7 +356,7 @@ export class RedisBridgePairingStore { previousDigest?: string, identityId = randomBytes(18).toString('base64url'), binding?: BridgeWorkerBinding, - identityId?: string | null, + redemptionId?: string, ): Promise { const credential = randomBytes(32).toString('base64url'); const credentialDigest = digest(credential); @@ -384,26 +392,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), + redemptionId, + credentialDigest, + JSON.stringify(stored), + String(this.credentialTtlSeconds), identityId, - 'EX', - this.credentialTtlSeconds, ); - 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/execution-profile.test.ts b/service/src/execution-profile.test.ts index 0e572a55..2a7cc073 100644 --- a/service/src/execution-profile.test.ts +++ b/service/src/execution-profile.test.ts @@ -4,6 +4,7 @@ import { queueNamesForExecutionProfile, resolveExecutionProfile, resolveExecutionProfileSource, + resolveQueuedSandboxBackend, validateQueuedSandboxBackend, validateQueuedExecutionProfile, } from './execution-profile'; @@ -71,6 +72,12 @@ describe('execution profile queue isolation', () => { 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')).toBe('http'); + expect(resolveQueuedSandboxBackend('stateful', 'remote-bridge')).toBe('remote-bridge'); + }); }); describe('execution profile request assertion', () => { @@ -130,6 +137,9 @@ describe('queued sandbox backend validation', () => { 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', () => { @@ -141,5 +151,8 @@ describe('queued sandbox backend validation', () => { ).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 3d39dd80..0c0c730b 100644 --- a/service/src/execution-profile.ts +++ b/service/src/execution-profile.ts @@ -18,6 +18,18 @@ export interface ExecutionProfileQueueNames { 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, +): SandboxBackendName { + return profile === 'stateful' && apiBackend === 'http' + ? 'lambda-microvm' + : apiBackend; +} + export function resolveExecutionProfile( raw: string | undefined, runtimeSessionMode: 'stateless' | 'affinity' | 'strict', @@ -153,8 +165,16 @@ export function validateQueuedExecutionProfile( export function validateQueuedSandboxBackend( jobBackend: unknown, workerBackend: SandboxBackendName, + bridgeWorkerId?: string, ): void { - if (jobBackend == null) return; + if (jobBackend == null) { + if (bridgeWorkerId != null && workerBackend !== 'remote-bridge') { + throw new Error( + `Legacy queued bridge job cannot run on the ${workerBackend} sandbox backend`, + ); + } + return; + } if (!SANDBOX_BACKENDS.includes(jobBackend as SandboxBackendName)) { throw new Error(`Queued job has invalid sandbox backend: ${String(jobBackend)}`); } diff --git a/service/src/sandbox-backend/remote-bridge.test.ts b/service/src/sandbox-backend/remote-bridge.test.ts index 4ff52fb9..0fb85856 100644 --- a/service/src/sandbox-backend/remote-bridge.test.ts +++ b/service/src/sandbox-backend/remote-bridge.test.ts @@ -71,4 +71,38 @@ describe('RemoteBridgeSandboxBackend', () => { 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), + 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'); + + await backend.execute(request(), { + ...context(), + bridgeWorkerId: 'deployment-worker', + }); + + expect(dispatched).toMatchObject({ + workerId: 'deployment-worker', + requireTenantBinding: false, + }); + }); }); diff --git a/service/src/sandbox-backend/remote-bridge.ts b/service/src/sandbox-backend/remote-bridge.ts index 1134e4e2..6a2dd26d 100644 --- a/service/src/sandbox-backend/remote-bridge.ts +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -35,7 +35,8 @@ export class RemoteBridgeSandboxBackend implements SandboxBackend { const settlement = await this.store.dispatch({ workerId, tenantId: ctx.tenantId, - requireTenantBinding: ctx.bridgeWorkerId != null, + requireTenantBinding: + ctx.bridgeWorkerId != null && ctx.bridgeWorkerId !== this.workerId, body: req.body, headers: req.headers, runtimeSessionId: ctx.runtimeSessionId, diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index daf042f8..a2832a37 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -25,6 +25,7 @@ import { } from '../metrics'; import { Jobs } from '../enum'; import { env, jobCompletionWaitTimeoutMs } from '../config'; +import { resolveQueuedSandboxBackend } from '../execution-profile'; import { normalizeEgressGatewayUrl, normalizeProgrammaticTimeoutMs, @@ -413,7 +414,10 @@ async function runReplayIteration( tenantId: state.tenantId, canonicalUserId: state.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, - sandboxBackend: env.SANDBOX_BACKEND, + sandboxBackend: resolveQueuedSandboxBackend( + env.EXECUTION_PROFILE, + env.SANDBOX_BACKEND, + ), ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, @@ -1408,7 +1412,10 @@ async function handleBlocking( tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, - sandboxBackend: env.SANDBOX_BACKEND, + sandboxBackend: resolveQueuedSandboxBackend( + env.EXECUTION_PROFILE, + env.SANDBOX_BACKEND, + ), ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, diff --git a/service/src/service/router.ts b/service/src/service/router.ts index 01894793..c4304405 100644 --- a/service/src/service/router.ts +++ b/service/src/service/router.ts @@ -31,6 +31,7 @@ import { resolveBridgeWorkerSelection, } from '../bridge/selection'; import logger from '../logger'; +import { resolveQueuedSandboxBackend } from '../execution-profile'; const { INSTANCE_ID } = env; const JOB_COMPLETION_WAIT_TIMEOUT_MS = jobCompletionWaitTimeoutMs( @@ -271,7 +272,10 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) tenantId: identity.storageNamespace, canonicalUserId: identity.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, - sandboxBackend: env.SANDBOX_BACKEND, + sandboxBackend: resolveQueuedSandboxBackend( + env.EXECUTION_PROFILE, + env.SANDBOX_BACKEND, + ), ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), ...(runtimeSessionId != null ? { runtimeSessionId } : {}), runtimeSessionMode, diff --git a/service/src/workers.ts b/service/src/workers.ts index e053c3bc..d2048dc8 100644 --- a/service/src/workers.ts +++ b/service/src/workers.ts @@ -63,7 +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); + validateQueuedSandboxBackend( + job.data.sandboxBackend, + env.SANDBOX_BACKEND, + job.data.bridgeWorkerId, + ); let sandboxPayload = payload; let executionManifestClaims = job.data.executionManifestClaims; let egressGrantToken = job.data.egressGrantToken; From c9eb86bef54abd10ff4b4b9fcfb6a52f7fa8b119 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 16:47:33 -0400 Subject: [PATCH 05/13] fix: preserve legacy routing and assignment auth --- packages/code/src/worker.ts | 11 +++++++---- service/src/execution-profile.test.ts | 6 +++++- service/src/execution-profile.ts | 18 ++++++++++++++---- service/src/service/programmatic-router.ts | 2 ++ service/src/service/router.ts | 1 + 5 files changed, 29 insertions(+), 9 deletions(-) diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 5e528d61..4f88b982 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -185,9 +185,7 @@ export class BridgeWorker { ): 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 +199,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', @@ -253,6 +252,10 @@ export class BridgeWorker { ); let settlement: BridgeSettlement; try { + await this.refreshCredential( + signal, + Date.parse(assignment.expiresAt), + ); const headers = { ...assignment.request.headers, ...(assignment.runtimeSessionId diff --git a/service/src/execution-profile.test.ts b/service/src/execution-profile.test.ts index 2a7cc073..bebdd97b 100644 --- a/service/src/execution-profile.test.ts +++ b/service/src/execution-profile.test.ts @@ -75,9 +75,13 @@ describe('execution profile queue isolation', () => { test('labels API-only stateful jobs with their Lambda worker backend', () => { expect(resolveQueuedSandboxBackend('stateful', 'http')).toBe('lambda-microvm'); - expect(resolveQueuedSandboxBackend('default', 'http')).toBe('http'); + 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', () => { diff --git a/service/src/execution-profile.ts b/service/src/execution-profile.ts index 0c0c730b..eac5628f 100644 --- a/service/src/execution-profile.ts +++ b/service/src/execution-profile.ts @@ -24,10 +24,20 @@ export type SandboxBackendName = typeof SANDBOX_BACKENDS[number]; export function resolveQueuedSandboxBackend( profile: ExecutionProfile, apiBackend: SandboxBackendName, -): SandboxBackendName { - return profile === 'stateful' && apiBackend === 'http' - ? 'lambda-microvm' - : apiBackend; + 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( diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index a2832a37..11a0ea33 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -417,6 +417,7 @@ async function runReplayIteration( sandboxBackend: resolveQueuedSandboxBackend( env.EXECUTION_PROFILE, env.SANDBOX_BACKEND, + env.EXECUTION_PROFILE_SOURCE, ), ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', @@ -1415,6 +1416,7 @@ async function handleBlocking( sandboxBackend: resolveQueuedSandboxBackend( env.EXECUTION_PROFILE, env.SANDBOX_BACKEND, + env.EXECUTION_PROFILE_SOURCE, ), ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', diff --git a/service/src/service/router.ts b/service/src/service/router.ts index c4304405..542fe4c0 100644 --- a/service/src/service/router.ts +++ b/service/src/service/router.ts @@ -275,6 +275,7 @@ router.post('/exec', executionLimiter, async (req: t.AuthenticatedRequest, res) sandboxBackend: resolveQueuedSandboxBackend( env.EXECUTION_PROFILE, env.SANDBOX_BACKEND, + env.EXECUTION_PROFILE_SOURCE, ), ...(bridgeWorkerId != null ? { bridgeWorkerId } : {}), ...(runtimeSessionId != null ? { runtimeSessionId } : {}), From a8e08cb93a8aee0d0fbcb885b2bdc63af341e7e9 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 17:18:19 -0400 Subject: [PATCH 06/13] fix: preserve long-lived bridge assignments --- packages/code/src/worker.ts | 66 +++++++++++++++++-- service/src/bridge/store.test.ts | 39 +++++++++++ service/src/bridge/store.ts | 60 +++++++++++++++-- service/src/service/programmatic-router.ts | 15 +++-- .../src/service/programmatic-state.test.ts | 7 +- service/src/service/programmatic-state.ts | 3 + service/src/service/replay-state.ts | 3 + 7 files changed, 175 insertions(+), 18 deletions(-) diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 4f88b982..2e71b9dd 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,7 +182,10 @@ 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; @@ -216,6 +220,40 @@ export class BridgeWorker { identity.expiresAt = rotatedIdentity.expiresAt; } + private async maintainCredential( + assignment: BridgeAssignment, + signal: 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 (!signal.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); + signal.addEventListener( + 'abort', + () => { + clearTimeout(timer); + resolve(); + }, + { once: true }, + ); + }); + } + if (signal.aborted || Date.now() >= assignmentDeadlineMs) return; + await this.refreshCredential(signal); + } + } + async executeAndSettle( assignment: BridgeAssignment, signal?: AbortSignal, @@ -225,7 +263,11 @@ export class BridgeWorker { Date.parse(assignment.expiresAt) + CREDENTIAL_REFRESH_WINDOW_MS, ); 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, @@ -250,12 +292,18 @@ export class BridgeWorker { executionController, cancellationController.signal, ); + let credentialMaintenanceError: unknown; + let credentialMaintenance: Promise | undefined; let settlement: BridgeSettlement; try { - await this.refreshCredential( - signal, - Date.parse(assignment.expiresAt), - ); + await this.refreshCredential(signal); + credentialMaintenance = this.maintainCredential( + assignment, + credentialController.signal, + ).catch((error) => { + credentialMaintenanceError = error; + executionController.abort(); + }); const headers = { ...assignment.request.headers, ...(assignment.runtimeSessionId @@ -276,6 +324,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) ?? @@ -306,9 +357,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/bridge/store.test.ts b/service/src/bridge/store.test.ts index f968d795..3f843182 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -71,6 +71,45 @@ describe('RedisBridgeStore', () => { await expect( store.lease('rebound-worker', 1_000, undefined, 'tenant-b-identity'), ).resolves.toBeUndefined(); + await expect( + store.lease('rebound-worker', 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 store.register({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'replacement-worker', + 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', 100, undefined, 'stale-identity'), + ).resolves.toBeUndefined(); + await expect( + store.lease('replacement-worker', 1_000, undefined, 'replacement-identity'), + ).resolves.toBeDefined(); controller.abort(); await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); }); diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index 197cf293..f4f7ebbc 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -98,6 +98,14 @@ function assignmentKey(assignmentId: string): string { return `${PREFIX}:assignment:${assignmentId}`; } +function queuedAssignment( + assignmentId: string, + workerIdentityId?: string, +): string { + const identity = workerIdentityId ?? ''; + return `${identity.length}:${identity}${assignmentId}`; +} + function settlementKey(assignmentId: string): string { return `${PREFIX}:assignment:${assignmentId}:settlement`; } @@ -281,7 +289,10 @@ export class RedisBridgeStore { 'EX', ttlSeconds, ); - transaction.rpush(queueKey(args.workerId), assignmentId); + transaction.rpush( + queueKey(args.workerId), + queuedAssignment(assignmentId, assignment.workerIdentityId), + ); transaction.expire(queueKey(args.workerId), ttlSeconds); await transaction.exec(); const settlement = await this.waitForSettlement( @@ -321,16 +332,48 @@ export class RedisBridgeStore { ): 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( + [ + "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 = ''", + ' if separator then', + ' local identityLength = tonumber(string.sub(entry, 1, separator - 1))', + ' if identityLength then', + ' identity = string.sub(entry, separator + 1, separator + identityLength)', + ' id = string.sub(entry, separator + identityLength + 1)', + ' end', + " elseif ARGV[3] == '' then", + ' id = entry', + ' 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)", + ' else', + " redis.call('LREM', KEYS[1], 1, entry)", + ' return raw', + ' end', + ' end', + 'end', + 'return nil', + ].join('\n'), + 1, + queueKey(workerId), + `${PREFIX}:assignment:`, + workerId, + identityId ?? '', + ); + 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) { @@ -339,7 +382,12 @@ export class RedisBridgeStore { 'Bridge worker incarnation was replaced', ); } - if (assignment.workerIdentityId !== identityId) continue; + if ( + assignment.workerId !== workerId || + assignment.workerIdentityId !== identityId + ) { + continue; + } if (Date.parse(assignment.expiresAt) <= Date.now()) continue; const { leaseTokenHash: _leaseTokenHash, diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index 11a0ea33..bd3b71be 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -414,11 +414,15 @@ async function runReplayIteration( tenantId: state.tenantId, canonicalUserId: state.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, - sandboxBackend: resolveQueuedSandboxBackend( - env.EXECUTION_PROFILE, - env.SANDBOX_BACKEND, - env.EXECUTION_PROFILE_SOURCE, - ), + sandboxBackend: + state.sandboxBackend ?? + (state.bridgeWorkerId != null + ? 'remote-bridge' + : resolveQueuedSandboxBackend( + env.EXECUTION_PROFILE, + env.SANDBOX_BACKEND, + env.EXECUTION_PROFILE_SOURCE, + )), ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, @@ -574,6 +578,7 @@ async function handleReplayInitial( timeout, language, bridgeWorkerId, + sandboxBackend: env.SANDBOX_BACKEND, }); /** Replay mode persists the full request (`userCode` + `tools` + `files`) * inside `ExecutionState` so continuations can re-enqueue without the diff --git a/service/src/service/programmatic-state.test.ts b/service/src/service/programmatic-state.test.ts index ec548476..315cf009 100644 --- a/service/src/service/programmatic-state.test.ts +++ b/service/src/service/programmatic-state.test.ts @@ -54,7 +54,11 @@ describe('buildReplayExecutionState', () => { authContextHash: 'hash_123', }; - const state = build({ authContext, bridgeWorkerId: 'code-user_123' }); + const state = build({ + authContext, + bridgeWorkerId: 'code-user_123', + sandboxBackend: 'remote-bridge', + }); expect(state).toMatchObject({ execution_id: 'exec_123', @@ -70,6 +74,7 @@ describe('buildReplayExecutionState', () => { authContextHash: 'hash_123', apiKeyId: 'key_legacy', bridgeWorkerId: 'code-user_123', + sandboxBackend: 'remote-bridge', mode: 'replay', userCode: 'print("hello")', tools: TOOLS, diff --git a/service/src/service/programmatic-state.ts b/service/src/service/programmatic-state.ts index 3a5f2f5b..acb18101 100644 --- a/service/src/service/programmatic-state.ts +++ b/service/src/service/programmatic-state.ts @@ -2,6 +2,7 @@ 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 type { SandboxBackendName } from '../execution-profile'; export interface BuildReplayExecutionStateParams { executionId: string; @@ -18,6 +19,7 @@ export interface BuildReplayExecutionStateParams { timeout: number; language: 'python' | 'bash'; bridgeWorkerId?: string; + sandboxBackend?: SandboxBackendName; now?: number; } @@ -43,6 +45,7 @@ export function buildReplayExecutionState( authContextHash: identity.authContextHash, apiKeyId: params.apiKeyId, bridgeWorkerId: params.bridgeWorkerId, + sandboxBackend: params.sandboxBackend, startTime: now, lastActivity: now, mode: 'replay', diff --git a/service/src/service/replay-state.ts b/service/src/service/replay-state.ts index 0d6b6709..0e35e401 100644 --- a/service/src/service/replay-state.ts +++ b/service/src/service/replay-state.ts @@ -24,6 +24,7 @@ import { nanoid } from 'nanoid'; import type { Redis } from 'ioredis'; import type * as t from '../types'; import type { LCTool } from '../preamble'; +import type { SandboxBackendName } from '../execution-profile'; import { connection } from '../queue'; import { env } from '../config'; import { internalServiceHeaders } from '../internal-service-auth'; @@ -111,6 +112,8 @@ export interface ExecutionState { apiKeyId?: string; /** Authenticated worker selection retained across every replay iteration. */ bridgeWorkerId?: string; + /** Original queue/backend target retained across replay continuations. */ + sandboxBackend?: SandboxBackendName; startTime: number; /** * Wall-clock ms of the last interaction that advanced this execution (initial From 509bc47bb1644c419048045b4b954b7c7d7ae523 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 17:43:08 -0400 Subject: [PATCH 07/13] fix: persist replay queue backend --- service/src/service/programmatic-router.ts | 12 +++++++-- .../src/service/programmatic-state.test.ts | 26 ++++++++++++++++++- service/src/service/programmatic-state.ts | 21 ++++++++++++++- 3 files changed, 55 insertions(+), 4 deletions(-) diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index bd3b71be..eb4cf75d 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -36,7 +36,10 @@ 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, @@ -578,7 +581,12 @@ async function handleReplayInitial( timeout, language, bridgeWorkerId, - sandboxBackend: env.SANDBOX_BACKEND, + 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 diff --git a/service/src/service/programmatic-state.test.ts b/service/src/service/programmatic-state.test.ts index 315cf009..a814e9e6 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 = [ { @@ -43,6 +46,27 @@ function build( } 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', diff --git a/service/src/service/programmatic-state.ts b/service/src/service/programmatic-state.ts index acb18101..551094f3 100644 --- a/service/src/service/programmatic-state.ts +++ b/service/src/service/programmatic-state.ts @@ -2,7 +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 type { SandboxBackendName } from '../execution-profile'; +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; From 4fe9147b0ba41925ea56b547aa8e7f0bd753047c Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 18:06:40 -0400 Subject: [PATCH 08/13] fix: fence bridge replay and credential rotation --- packages/code/src/worker.test.ts | 80 ++++++++++++++++++++++ packages/code/src/worker.ts | 15 ++-- service/src/execution-profile.test.ts | 12 ++++ service/src/execution-profile.ts | 10 +++ service/src/lifecycle.ts | 18 ++--- service/src/local-api.ts | 9 +-- service/src/queue.ts | 64 ++++++++++++++--- service/src/secure-startup.test.ts | 14 ++++ service/src/secure-startup.ts | 17 +++-- service/src/service/programmatic-router.ts | 45 ++++++------ 10 files changed, 222 insertions(+), 62 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 75eb78c6..128dba08 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -469,6 +469,86 @@ 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', + 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', + 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 2e71b9dd..5530b1fb 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -222,7 +222,8 @@ export class BridgeWorker { private async maintainCredential( assignment: BridgeAssignment, - signal: AbortSignal, + stopSignal: AbortSignal, + requestSignal?: AbortSignal, ): Promise { const identity = this.options.identity; if (identity == null) return; @@ -230,7 +231,7 @@ export class BridgeWorker { this.options.credentialRefreshWindowMs ?? CREDENTIAL_REFRESH_WINDOW_MS; const assignmentDeadlineMs = Date.parse(assignment.expiresAt); - while (!signal.aborted && Date.now() < assignmentDeadlineMs) { + while (!stopSignal.aborted && Date.now() < assignmentDeadlineMs) { const refreshAtMs = Date.parse(identity.expiresAt) - refreshWindowMs; const waitMs = Math.max( 0, @@ -239,7 +240,7 @@ export class BridgeWorker { if (waitMs > 0) { await new Promise((resolve) => { const timer = setTimeout(resolve, waitMs); - signal.addEventListener( + stopSignal.addEventListener( 'abort', () => { clearTimeout(timer); @@ -249,8 +250,11 @@ export class BridgeWorker { ); }); } - if (signal.aborted || Date.now() >= assignmentDeadlineMs) return; - await this.refreshCredential(signal); + 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); } } @@ -300,6 +304,7 @@ export class BridgeWorker { credentialMaintenance = this.maintainCredential( assignment, credentialController.signal, + signal, ).catch((error) => { credentialMaintenanceError = error; executionController.abort(); diff --git a/service/src/execution-profile.test.ts b/service/src/execution-profile.test.ts index bebdd97b..7ba4640f 100644 --- a/service/src/execution-profile.test.ts +++ b/service/src/execution-profile.test.ts @@ -2,6 +2,7 @@ import { describe, expect, test } from 'bun:test'; import { checkExecutionProfileExpectation, queueNamesForExecutionProfile, + queueNameForExecution, resolveExecutionProfile, resolveExecutionProfileSource, resolveQueuedSandboxBackend, @@ -58,6 +59,17 @@ describe('execution profile queue isolation', () => { }); }); + 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'), diff --git a/service/src/execution-profile.ts b/service/src/execution-profile.ts index eac5628f..38e9bc8c 100644 --- a/service/src/execution-profile.ts +++ b/service/src/execution-profile.ts @@ -105,6 +105,16 @@ export function queueNamesForExecutionProfile( : LEGACY_QUEUE_NAMES; } +export function queueNameForExecution( + language: 'python' | 'bash', + profile: ExecutionProfile, + source: ExecutionProfileSource, + backend?: SandboxBackendName, +): string { + const names = queueNamesForExecutionProfile(profile, source, backend); + return language === 'bash' ? names.other : names.python; +} + export type ExecutionProfileExpectation = | { ok: true } | { diff --git a/service/src/lifecycle.ts b/service/src/lifecycle.ts index 596193fd..456265b6 100644 --- a/service/src/lifecycle.ts +++ b/service/src/lifecycle.ts @@ -1,10 +1,16 @@ import type { Queue } from 'bullmq'; import type { Express } from 'express'; -import { pyQueue, otherQueue, pyQueueEvents, otherQueueEvents, connection } from './queue'; +import { + pyQueue, + otherQueue, + connection, + closeQueueConnections, +} from './queue'; import { validateStartupAuthConfig } from './auth/startup'; import { env } from './config'; import { validateApiHardenedConfig, + validateApiSandboxBackendPolicy, validateExecutionProfilePolicy, validateSandboxBackendPolicy, validateApiBridgePolicy, @@ -91,8 +97,9 @@ export async function startupApiOnly(): Promise { 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 ce7e5b77..2fc14fab 100644 --- a/service/src/queue.ts +++ b/service/src/queue.ts @@ -7,7 +7,11 @@ 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 { SandboxBackendName } from './execution-profile'; import logger from './logger'; import { redisKeepAliveOptions } from './redis-options'; import { bullmqQueueJobs, registerBullmqQueueMetricsCollector } from './metrics'; @@ -62,17 +66,48 @@ const queueNames = queueNamesForExecutionProfile( 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, +): QueueBinding { + const name = queueNameForExecution( + language, + env.EXECUTION_PROFILE, + env.EXECUTION_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 { @@ -91,7 +126,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), @@ -117,4 +152,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/secure-startup.test.ts b/service/src/secure-startup.test.ts index c970a0fa..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, @@ -332,6 +333,19 @@ describe('sandbox backend policy', () => { 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 7e4e1ee4..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,13 +128,7 @@ export function validateSandboxBackendPolicy(): void { ); } if (env.SANDBOX_BACKEND === 'remote-bridge') { - if (env.BRIDGE_DYNAMIC_WORKERS) { - if (env.BRIDGE_AUTH_MODE !== 'paired') { - throw new SecureStartupConfigError( - 'Dynamic remote bridge workers require CODEAPI_BRIDGE_AUTH_MODE=paired', - ); - } - } else { + if (!env.BRIDGE_DYNAMIC_WORKERS) { requireValue('CODEAPI_BRIDGE_WORKER_ID', env.BRIDGE_WORKER_ID); } if (env.HARDENED_SANDBOX_MODE) { diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index eb4cf75d..da7f4fc8 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'; @@ -337,19 +341,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, @@ -405,7 +396,19 @@ 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, + ); const job = await queue.add(Jobs.execute, { code: state.userCode ?? '', userId, @@ -417,15 +420,7 @@ async function runReplayIteration( tenantId: state.tenantId, canonicalUserId: state.canonicalUserId, executionProfile: env.EXECUTION_PROFILE, - sandboxBackend: - state.sandboxBackend ?? - (state.bridgeWorkerId != null - ? 'remote-bridge' - : resolveQueuedSandboxBackend( - env.EXECUTION_PROFILE, - env.SANDBOX_BACKEND, - env.EXECUTION_PROFILE_SOURCE, - )), + sandboxBackend: replayBackend, ...(state.bridgeWorkerId != null ? { bridgeWorkerId: state.bridgeWorkerId } : {}), runtimeSessionMode: 'stateless', runtimeSessionExemption: PROGRAMMATIC_RUNTIME_SESSION_EXEMPTION, From 172d0b1a744bbf4f2f610a570c7e036a72b25080 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 20:35:34 -0400 Subject: [PATCH 09/13] fix: address principal worker review findings --- service/src/bridge/pairing.test.ts | 7 +++- service/src/bridge/router.test.ts | 38 +++++++++++++++++++ service/src/bridge/router.ts | 13 ++++++- service/src/bridge/store.test.ts | 25 ++++++++++++ service/src/bridge/store.ts | 16 +++++++- service/src/queue.ts | 12 ++++-- service/src/service/programmatic-router.ts | 10 ++++- .../src/service/programmatic-state.test.ts | 6 +++ service/src/service/programmatic-state.ts | 4 ++ service/src/service/replay-state.ts | 10 ++++- 10 files changed, 129 insertions(+), 12 deletions(-) diff --git a/service/src/bridge/pairing.test.ts b/service/src/bridge/pairing.test.ts index 5fe0d0db..310ad7d9 100644 --- a/service/src/bridge/pairing.test.ts +++ b/service/src/bridge/pairing.test.ts @@ -63,8 +63,11 @@ describe('RedisBridgePairingStore', () => { originalAuthorization.identityId, ); await expect( - pairings.authorize(requestFor(issued.credential, 'superseded-bound-proof')), - ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); + 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 () => { diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts index f9e77a20..4d9b7693 100644 --- a/service/src/bridge/router.test.ts +++ b/service/src/bridge/router.test.ts @@ -25,6 +25,44 @@ 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(); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index a28caeb2..167db3ed 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -78,7 +78,9 @@ function parseBinding(value: unknown): BridgeWorkerBinding | undefined { } 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 }); } @@ -217,7 +219,13 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { res.status(400).json({ error: 'Invalid bridge worker ID' }); return; } + 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; @@ -344,6 +352,7 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { | { workerId: string; credentialId: string; + activeCredentialId: string; identityId?: string; binding?: BridgeWorkerBinding; } @@ -374,7 +383,7 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { ...(authorization?.binding != null ? { binding: authorization.binding } : {}), - }); + }, authorization?.activeCredentialId); } catch (error) { if (error instanceof BridgeStoreError) { sendStoreError(error, res); diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 3f843182..a1960704 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -14,6 +14,31 @@ 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', + 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('rejects a dynamic worker lease outside its bound tenant', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index f4f7ebbc..d4aa2ac3 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -149,8 +149,12 @@ export class RedisBridgeStore { private readonly workerTtlSeconds = DEFAULT_WORKER_TTL_SECONDS, ) {} - async register(registration: RegisteredBridgeWorker): Promise { + async register( + registration: RegisteredBridgeWorker, + expectedActiveCredentialId?: string, + ): Promise { const script = [ + 'if ARGV[5] ~= "" and redis.call(\'GET\', KEYS[5]) ~= ARGV[5] then return -3 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])', @@ -166,15 +170,17 @@ export class RedisBridgeStore { const result = Number( await this.redis.eval( script, - 4, + 5, workerKey(registration.workerId), incarnationFenceKey(registration.workerId, registration.incarnationId), quarantineKey(registration.workerId, registration.incarnationId), workerIncarnationKey(registration.workerId), + `${PREFIX}:identity:${registration.workerId}`, registration.incarnationId, JSON.stringify(registration), String(this.workerTtlSeconds), `${PREFIX}:worker:${registration.workerId}:incarnation:`, + expectedActiveCredentialId ?? '', ), ); if (result === -2) { @@ -189,6 +195,12 @@ 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: { diff --git a/service/src/queue.ts b/service/src/queue.ts index 2fc14fab..54fea308 100644 --- a/service/src/queue.ts +++ b/service/src/queue.ts @@ -11,7 +11,11 @@ import { queueNameForExecution, queueNamesForExecutionProfile, } from './execution-profile'; -import type { SandboxBackendName } 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'; @@ -94,11 +98,13 @@ function getQueueResources( 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, - env.EXECUTION_PROFILE, - env.EXECUTION_PROFILE_SOURCE, + profile, + source, backend, ); return { ...getQueueResources(name), language }; diff --git a/service/src/service/programmatic-router.ts b/service/src/service/programmatic-router.ts index da7f4fc8..baa350e6 100644 --- a/service/src/service/programmatic-router.ts +++ b/service/src/service/programmatic-router.ts @@ -30,6 +30,7 @@ import { import { Jobs } from '../enum'; import { env, jobCompletionWaitTimeoutMs } from '../config'; import { resolveQueuedSandboxBackend } from '../execution-profile'; +import { publicExecutionFailure } from '../utils'; import { normalizeEgressGatewayUrl, normalizeProgrammaticTimeoutMs, @@ -408,6 +409,8 @@ async function runReplayIteration( 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 ?? '', @@ -419,7 +422,7 @@ 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', @@ -576,6 +579,8 @@ async function handleReplayInitial( timeout, language, bridgeWorkerId, + executionProfile: env.EXECUTION_PROFILE, + executionProfileSource: env.EXECUTION_PROFILE_SOURCE, sandboxBackend: resolveReplayStateSandboxBackend({ executionProfile: env.EXECUTION_PROFILE, executionProfileSource: env.EXECUTION_PROFILE_SOURCE, @@ -854,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', diff --git a/service/src/service/programmatic-state.test.ts b/service/src/service/programmatic-state.test.ts index a814e9e6..fc84d8f8 100644 --- a/service/src/service/programmatic-state.test.ts +++ b/service/src/service/programmatic-state.test.ts @@ -40,6 +40,8 @@ function build( isPyPlot: false, timeout: 300000, language: 'python', + executionProfile: 'default', + executionProfileSource: 'inferred', now: 1778250000000, ...overrides, }); @@ -82,6 +84,8 @@ describe('buildReplayExecutionState', () => { authContext, bridgeWorkerId: 'code-user_123', sandboxBackend: 'remote-bridge', + executionProfile: 'stateful', + executionProfileSource: 'explicit', }); expect(state).toMatchObject({ @@ -99,6 +103,8 @@ describe('buildReplayExecutionState', () => { 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 551094f3..25571fed 100644 --- a/service/src/service/programmatic-state.ts +++ b/service/src/service/programmatic-state.ts @@ -39,6 +39,8 @@ export interface BuildReplayExecutionStateParams { language: 'python' | 'bash'; bridgeWorkerId?: string; sandboxBackend?: SandboxBackendName; + executionProfile: ExecutionProfile; + executionProfileSource: ExecutionProfileSource; now?: number; } @@ -65,6 +67,8 @@ export function buildReplayExecutionState( 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 0e35e401..4d2a0501 100644 --- a/service/src/service/replay-state.ts +++ b/service/src/service/replay-state.ts @@ -24,7 +24,11 @@ import { nanoid } from 'nanoid'; import type { Redis } from 'ioredis'; import type * as t from '../types'; import type { LCTool } from '../preamble'; -import type { SandboxBackendName } from '../execution-profile'; +import type { + ExecutionProfile, + ExecutionProfileSource, + SandboxBackendName, +} from '../execution-profile'; import { connection } from '../queue'; import { env } from '../config'; import { internalServiceHeaders } from '../internal-service-auth'; @@ -114,6 +118,10 @@ export interface ExecutionState { 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 From 80ec7bd4728e88debdfc6b481ab1125908b3a807 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 21:24:31 -0400 Subject: [PATCH 10/13] fix: reconcile principal workers with bridge fencing --- packages/code/src/worker.test.ts | 4 +- packages/code/src/worker.ts | 6 +-- service/src/bridge/pairing.ts | 41 +++++++++++++------ service/src/bridge/router.test.ts | 1 + service/src/bridge/store.test.ts | 38 +++++++++++++++-- .../src/sandbox-backend/remote-bridge.test.ts | 2 + 6 files changed, 69 insertions(+), 23 deletions(-) diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 128dba08..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, @@ -520,6 +520,7 @@ test('sandbox completion does not cancel an in-flight credential rotation', asyn const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', workerId: 'vm-1', + incarnationId, sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', identity, capabilities: { @@ -535,6 +536,7 @@ test('sandbox completion does not cancel an in-flight credential rotation', asyn 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(), diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 5530b1fb..0f22d22b 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -262,10 +262,7 @@ export class BridgeWorker { 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 credentialController = new AbortController(); const abortExecution = (): void => { @@ -300,7 +297,6 @@ export class BridgeWorker { let credentialMaintenance: Promise | undefined; let settlement: BridgeSettlement; try { - await this.refreshCredential(signal); credentialMaintenance = this.maintainCredential( assignment, credentialController.signal, diff --git a/service/src/bridge/pairing.ts b/service/src/bridge/pairing.ts index 3ff3d30d..fd1e8cff 100644 --- a/service/src/bridge/pairing.ts +++ b/service/src/bridge/pairing.ts @@ -19,12 +19,18 @@ 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 = ` @@ -56,7 +62,11 @@ if redis.call('GET', KEYS[1]) ~= ARGV[1] then end redis.call('SET', KEYS[2], ARGV[3], 'EX', ARGV[4]) redis.call('SET', KEYS[3], 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 redis.call('DEL', KEYS[1]) return 1 `; @@ -79,7 +89,7 @@ interface StoredPairing { interface StoredCredential { workerId: string; - identityId: string; + identityId?: string; publicKey: string; expiresAt: string; binding?: BridgeWorkerBinding; @@ -239,7 +249,7 @@ export class RedisBridgePairingStore { workerId: string; credentialId: string; activeCredentialId: string; - identityId: string; + identityId?: string; binding?: BridgeWorkerBinding; }> { const proofTime = Date.parse(args.timestamp); @@ -271,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', @@ -307,7 +321,7 @@ export class RedisBridgePairingStore { workerId: stored.workerId, credentialId: credentialDigest, activeCredentialId: activeDigest, - identityId: stored.identityId, + ...(stored.identityId != null ? { identityId: stored.identityId } : {}), ...(stored.binding ? { binding: stored.binding } : {}), }; } @@ -345,7 +359,7 @@ export class RedisBridgePairingStore { workerId, previous.publicKey, previousDigest, - previous.identityId, + previous.identityId ?? null, previous.binding, ); } @@ -354,7 +368,7 @@ 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 { @@ -363,9 +377,10 @@ export class RedisBridgePairingStore { const expiresAt = new Date( Date.now() + this.credentialTtlSeconds * 1000, ).toISOString(); + const stableIdentityId = identityId ?? undefined; const stored: StoredCredential = { workerId, - identityId, + ...(stableIdentityId != null ? { identityId: stableIdentityId } : {}), publicKey, expiresAt, binding, @@ -382,7 +397,7 @@ export class RedisBridgePairingStore { credentialDigest, JSON.stringify(stored), String(this.credentialTtlSeconds), - identityId, + stableIdentityId ?? '', ); if (rotated !== 1) { throw new BridgePairingError( @@ -409,7 +424,7 @@ export class RedisBridgePairingStore { credentialDigest, JSON.stringify(stored), String(this.credentialTtlSeconds), - identityId, + stableIdentityId ?? '', ); if (installed !== 1) { throw new BridgePairingError( diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts index 4d9b7693..72123fa5 100644 --- a/service/src/bridge/router.test.ts +++ b/service/src/bridge/router.test.ts @@ -127,6 +127,7 @@ describe('paired bridge HTTP API', () => { const body = JSON.stringify({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: 'user-vm', + incarnationId: 'incarnation-00000001', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index a1960704..496f5a38 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -24,6 +24,7 @@ describe('RedisBridgeStore', () => { store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: 'fenced-worker', + incarnationId, credentialId: 'stale-credential-digest', identityId: 'stale-identity', capabilities: { @@ -43,6 +44,7 @@ describe('RedisBridgeStore', () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: 'tenant-worker', + incarnationId, capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -71,6 +73,7 @@ describe('RedisBridgeStore', () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: 'rebound-worker', + incarnationId, identityId: 'tenant-a-identity', capabilities: { statefulWorkspace: true, @@ -94,10 +97,22 @@ describe('RedisBridgeStore', () => { }); await expect( - store.lease('rebound-worker', 1_000, undefined, 'tenant-b-identity'), + store.lease( + 'rebound-worker', + incarnationId, + 1_000, + undefined, + 'tenant-b-identity', + ), ).resolves.toBeUndefined(); await expect( - store.lease('rebound-worker', 1_000, undefined, 'tenant-a-identity'), + store.lease( + 'rebound-worker', + incarnationId, + 1_000, + undefined, + 'tenant-a-identity', + ), ).resolves.toBeDefined(); controller.abort(); await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); @@ -107,6 +122,7 @@ describe('RedisBridgeStore', () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: 'replacement-worker', + incarnationId, identityId: 'replacement-identity', capabilities: { statefulWorkspace: true, @@ -130,10 +146,22 @@ describe('RedisBridgeStore', () => { }); await expect( - store.lease('replacement-worker', 100, undefined, 'stale-identity'), + store.lease( + 'replacement-worker', + incarnationId, + 100, + undefined, + 'stale-identity', + ), ).resolves.toBeUndefined(); await expect( - store.lease('replacement-worker', 1_000, undefined, 'replacement-identity'), + store.lease( + 'replacement-worker', + incarnationId, + 1_000, + undefined, + 'replacement-identity', + ), ).resolves.toBeDefined(); controller.abort(); await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); @@ -143,6 +171,7 @@ describe('RedisBridgeStore', () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: 'rotating-worker', + incarnationId, identityId: 'stable-paired-identity', credentialId: 'credential-before-refresh', capabilities: { @@ -168,6 +197,7 @@ describe('RedisBridgeStore', () => { const assignment = await store.lease( 'rotating-worker', + incarnationId, 1_000, undefined, 'stable-paired-identity', diff --git a/service/src/sandbox-backend/remote-bridge.test.ts b/service/src/sandbox-backend/remote-bridge.test.ts index 0fb85856..50cb8744 100644 --- a/service/src/sandbox-backend/remote-bridge.test.ts +++ b/service/src/sandbox-backend/remote-bridge.test.ts @@ -37,6 +37,7 @@ describe('RemoteBridgeSandboxBackend', () => { protocolVersion: 1 as const, generation: 1, leaseToken: 'a'.repeat(32), + incarnationId: 'incarnation-00000001', status: 'fulfilled' as const, result: { session_id: 'session-1', @@ -83,6 +84,7 @@ describe('RemoteBridgeSandboxBackend', () => { protocolVersion: 1 as const, generation: 1, leaseToken: 'a'.repeat(32), + incarnationId: 'incarnation-00000001', status: 'fulfilled' as const, result: { session_id: 'session-1', From 45a9b6d6f1a942a588251af259d264ac897de30b Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 21:36:06 -0400 Subject: [PATCH 11/13] fix: fence principal worker lifecycle transitions --- service/src/bridge/router.test.ts | 14 ++++++- service/src/bridge/router.ts | 2 +- service/src/bridge/store.test.ts | 70 +++++++++++++++++++++++++++++++ service/src/bridge/store.ts | 40 ++++++++++++++---- 4 files changed, 116 insertions(+), 10 deletions(-) diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts index 72123fa5..bb59e991 100644 --- a/service/src/bridge/router.test.ts +++ b/service/src/bridge/router.test.ts @@ -129,7 +129,7 @@ describe('paired bridge HTTP API', () => { workerId: 'user-vm', incarnationId: 'incarnation-00000001', capabilities: { - statefulWorkspace: true, + statefulWorkspace: false, sandboxProfile: 'nsjail', runtimes: ['bash'], }, @@ -164,6 +164,18 @@ describe('paired bridge HTTP API', () => { }, ); 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({ diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 167db3ed..de4cd5e0 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -363,7 +363,7 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { workerId: registration.workerId, incarnationId: registration.incarnationId, capabilities: { - statefulWorkspace: true, + statefulWorkspace: capabilities.statefulWorkspace as boolean, sandboxProfile: capabilities.sandboxProfile as string, runtimes: capabilities.runtimes as string[], ...(typeof capabilities.policyDigest === 'string' diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 496f5a38..9c126822 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -40,6 +40,32 @@ describe('RedisBridgeStore', () => { ).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, @@ -167,6 +193,50 @@ describe('RedisBridgeStore', () => { await expect(completion).rejects.toMatchObject({ code: 'ASSIGNMENT_EXPIRED' }); }); + test('a stale incarnation poll cannot consume replacement incarnation work', async () => { + const replacementIncarnationId = 'incarnation-00000002'; + 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 queued work after credential refresh preserves the paired identity', async () => { await store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index d4aa2ac3..5103dcf7 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -100,10 +100,14 @@ function assignmentKey(assignmentId: string): string { function queuedAssignment( assignmentId: string, + incarnationId: string, workerIdentityId?: string, ): string { const identity = workerIdentityId ?? ''; - return `${identity.length}:${identity}${assignmentId}`; + return ( + `${identity.length}:${identity}` + + `${incarnationId.length}:${incarnationId}${assignmentId}` + ); } function settlementKey(assignmentId: string): string { @@ -154,7 +158,13 @@ export class RedisBridgeStore { expectedActiveCredentialId?: string, ): Promise { const script = [ - 'if ARGV[5] ~= "" and redis.call(\'GET\', KEYS[5]) ~= ARGV[5] then return -3 end', + '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])', @@ -170,17 +180,19 @@ export class RedisBridgeStore { const result = Number( await this.redis.eval( script, - 5, + 6, workerKey(registration.workerId), incarnationFenceKey(registration.workerId, registration.incarnationId), quarantineKey(registration.workerId, registration.incarnationId), workerIncarnationKey(registration.workerId), + `${PREFIX}:stable-identity:${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) { @@ -303,7 +315,11 @@ export class RedisBridgeStore { ); transaction.rpush( queueKey(args.workerId), - queuedAssignment(assignmentId, assignment.workerIdentityId), + queuedAssignment( + assignmentId, + assignment.incarnationId, + assignment.workerIdentityId, + ), ); transaction.expire(queueKey(args.workerId), ttlSeconds); await transaction.exec(); @@ -351,16 +367,23 @@ export class RedisBridgeStore { " local separator = string.find(entry, ':', 1, true)", ' local id = nil', " local identity = ''", + " local incarnation = ''", ' if separator then', ' local identityLength = tonumber(string.sub(entry, 1, separator - 1))', ' if identityLength then', ' identity = string.sub(entry, separator + 1, separator + identityLength)', - ' id = string.sub(entry, separator + identityLength + 1)', + ' 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', + ' end', ' end', - " elseif ARGV[3] == '' then", - ' id = entry', ' end', - ' if id and identity == ARGV[3] then', + ' if id and identity == ARGV[3] and incarnation == ARGV[4] then', " local raw = redis.call('GET', ARGV[1] .. id)", ' if not raw then', " redis.call('LREM', KEYS[1], 1, entry)", @@ -377,6 +400,7 @@ export class RedisBridgeStore { `${PREFIX}:assignment:`, workerId, identityId ?? '', + incarnationId, ); if (raw == null) { await delay( From d654c486e1f97d24f295698c2ecf6da8a9b85e61 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 21:46:13 -0400 Subject: [PATCH 12/13] fix: fence bridge settlement ownership --- service/src/bridge/router.ts | 5 ++ service/src/bridge/store.test.ts | 124 +++++++++++++++++++++++++++++++ service/src/bridge/store.ts | 60 ++++++++++++--- 3 files changed, 178 insertions(+), 11 deletions(-) diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index de4cd5e0..0c2771b1 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -461,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/store.test.ts b/service/src/bridge/store.test.ts index 9c126822..4d6b96a2 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -237,6 +237,74 @@ describe('RedisBridgeStore', () => { 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 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 store.register({ protocolVersion: BRIDGE_PROTOCOL_VERSION, @@ -322,6 +390,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 5103dcf7..ce6e26ea 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -368,6 +368,7 @@ export class RedisBridgeStore { ' 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', @@ -380,14 +381,20 @@ export class RedisBridgeStore { ' 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] and incarnation == ARGV[4] then', + ' 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)", - ' else', + " 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', @@ -439,26 +446,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', @@ -472,12 +485,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), + `${PREFIX}:stable-identity:${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( From a8b3c177017c4c10056c3cd758ebf62157f7d1af Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 22:01:26 -0400 Subject: [PATCH 13/13] fix: fence bridge leases to active principals --- service/src/bridge/store.test.ts | 62 +++++++++++++++++++ service/src/bridge/store.ts | 16 ++++- .../src/sandbox-backend/remote-bridge.test.ts | 45 +++++++++++++- service/src/sandbox-backend/remote-bridge.ts | 4 +- 4 files changed, 122 insertions(+), 5 deletions(-) diff --git a/service/src/bridge/store.test.ts b/service/src/bridge/store.test.ts index 4d6b96a2..684381ad 100644 --- a/service/src/bridge/store.test.ts +++ b/service/src/bridge/store.test.ts @@ -96,6 +96,10 @@ describe('RedisBridgeStore', () => { }); 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', @@ -145,6 +149,10 @@ describe('RedisBridgeStore', () => { }); 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', @@ -193,8 +201,54 @@ describe('RedisBridgeStore', () => { 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', @@ -239,6 +293,10 @@ describe('RedisBridgeStore', () => { 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', @@ -306,6 +364,10 @@ describe('RedisBridgeStore', () => { }); 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', diff --git a/service/src/bridge/store.ts b/service/src/bridge/store.ts index ce6e26ea..2d3975ea 100644 --- a/service/src/bridge/store.ts +++ b/service/src/bridge/store.ts @@ -60,6 +60,10 @@ 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`; } @@ -185,7 +189,7 @@ export class RedisBridgeStore { incarnationFenceKey(registration.workerId, registration.incarnationId), quarantineKey(registration.workerId, registration.incarnationId), workerIncarnationKey(registration.workerId), - `${PREFIX}:stable-identity:${registration.workerId}`, + workerStableIdentityKey(registration.workerId), `${PREFIX}:identity:${registration.workerId}`, registration.incarnationId, JSON.stringify(registration), @@ -362,6 +366,11 @@ export class RedisBridgeStore { while (signal?.aborted !== true && Date.now() < deadline) { 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)", @@ -402,8 +411,9 @@ export class RedisBridgeStore { 'end', 'return nil', ].join('\n'), - 1, + 2, queueKey(workerId), + workerStableIdentityKey(workerId), `${PREFIX}:assignment:`, workerId, identityId ?? '', @@ -501,7 +511,7 @@ export class RedisBridgeStore { 4, assignmentKey(assignmentId), workerKey(workerId), - `${PREFIX}:stable-identity:${workerId}`, + workerStableIdentityKey(workerId), settlementKey(assignmentId), rawAssignment, rawRegistration ?? '', diff --git a/service/src/sandbox-backend/remote-bridge.test.ts b/service/src/sandbox-backend/remote-bridge.test.ts index 50cb8744..c1b9e2c9 100644 --- a/service/src/sandbox-backend/remote-bridge.test.ts +++ b/service/src/sandbox-backend/remote-bridge.test.ts @@ -95,7 +95,11 @@ describe('RemoteBridgeSandboxBackend', () => { }; }, } satisfies Pick; - const backend = new RemoteBridgeSandboxBackend(store, 'deployment-worker'); + const backend = new RemoteBridgeSandboxBackend( + store, + 'deployment-worker', + false, + ); await backend.execute(request(), { ...context(), @@ -107,4 +111,43 @@ describe('RemoteBridgeSandboxBackend', () => { 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 6a2dd26d..b1a94ae0 100644 --- a/service/src/sandbox-backend/remote-bridge.ts +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -17,6 +17,7 @@ export class RemoteBridgeSandboxBackend implements SandboxBackend { constructor( private readonly store: Pick = bridgeStore, private readonly workerId: string = env.BRIDGE_WORKER_ID, + private readonly dynamicWorkers: boolean = env.BRIDGE_DYNAMIC_WORKERS, ) {} async execute( @@ -36,7 +37,8 @@ export class RemoteBridgeSandboxBackend implements SandboxBackend { workerId, tenantId: ctx.tenantId, requireTenantBinding: - ctx.bridgeWorkerId != null && ctx.bridgeWorkerId !== this.workerId, + ctx.bridgeWorkerId != null && + (this.dynamicWorkers || ctx.bridgeWorkerId !== this.workerId), body: req.body, headers: req.headers, runtimeSessionId: ctx.runtimeSessionId,