🪪 feat: Add Principal-Bound Bridge Workers

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

docs/remote-bridge/README.md

@@ -32,6 +32,23 @@ session hint. In hardened mode, startup requires the bridge token to be at least
 32 bytes. `PTC_MODE=blocking` is rejected; replay mode is required because a
 remote execution cannot retain an open Code API process across tool callbacks.
 
+To attach multiple principal-owned workers to one Code API deployment, enable
+dynamic routing. A compatibility default worker is optional in this mode:
+
+```dotenv
+CODEAPI_BRIDGE_DYNAMIC_WORKERS=true
+CODEAPI_BRIDGE_AUTH_MODE=paired
+# CODEAPI_BRIDGE_WORKER_ID=my-default-vm
+```
+
+Dynamic routing is accepted only with paired authentication. LibreChat signs
+the selected worker into the short-lived Code API JWT as `code_worker_id`.
+`X-LibreChat-Code-Worker-ID` remains the transport header, but Code API accepts
+it only when it exactly matches that authenticated claim. The resolved worker
+is persisted across the queue and programmatic replay boundaries, and Code API
+requires both its stored tenant binding and registered worker credential before
+creating a lease.
+
 Create a single-use pairing code with the administrator secret:
 
 ```bash
@@ -41,6 +58,28 @@ curl -fsS https://code.example.com/v1/bridge/pairings \
   --data '{"workerId":"my-vm"}'
 ```
 
+With dynamic routing enabled, the trusted control plane must bind each pairing
+to one tenant and generic principal. Code API treats the principal as lifecycle
+and audit metadata; LibreChat remains responsible for resolving user, role, and
+group membership before selecting the worker:
+
+```bash
+curl -fsS https://code.example.com/v1/bridge/pairings \
+  -H "Authorization: Bearer $CODEAPI_BRIDGE_TOKEN" \
+  -H 'Content-Type: application/json' \
+  --data '{
+    "workerId":"user-vm",
+    "binding":{
+      "tenantId":"tenant-1",
+      "principal":{"type":"user","id":"user-1"}
+    }
+  }'
+```
+
+Principal types are `deployment`, `tenant`, `user`, `role`, and `group`.
+Pairing and registration bodies from the VM cannot replace the server-issued
+binding, and credential rotation preserves it.
+
 Redeem the returned code on the VM using
 [`@librechat/code`](../../packages/code/README.md). The CLI generates its key
 locally, proves possession on every request, and rotates its short-lived
@@ -84,7 +123,14 @@ execution.
   digest, timestamp, nonce, and credential.
 - Accepted proof nonces cannot be replayed, credentials rotate before expiry,
   and an administrator can revoke the active worker identity immediately.
-- Code API permits one active assignment per configured worker.
+- Assignment leases bind to a stable paired identity rather than an individual
+  short-lived credential. Rotation preserves that identity; pairing again
+  replaces it and fences work queued for the previous owner.
+- Remote bridge deployments use backend-specific BullMQ queues and serialize
+  the expected backend on every new job, preventing Lambda or HTTP consumers
+  from accepting attached-worker executions.
+- Code API permits one active assignment per worker.
+- Dynamic workers are fenced to their server-issued tenant before assignment.
 - Each assignment has an absolute deadline, generation, and random lease token.
 - Settlements with the wrong worker, generation, token, or expired deadline are
   rejected.
@@ -104,6 +150,6 @@ administrator token immediately. Pairing secures worker transport identity; it
 cannot attest that a compromised VM truthfully reports or enforces its sandbox
 capabilities.
 
-The next control-plane layer can add owner-scoped environment records and a
-multi-worker directory without changing the execution protocol or moving code
-tools into the Agents SDK.
+LibreChat's owner-scoped environment registry can issue these principal-bound
+pairings without changing the worker execution protocol or moving code tools
+into the Agents SDK.

packages/code/src/worker.test.ts

@@ -413,7 +413,7 @@ test('paired worker refreshes before an assignment that outlives its credential'
     identity: {
       privateKey: key.privateKey,
       credential: 'credential-too-short-for-assignment',
-      expiresAt: new Date(Date.now() + 90_000).toISOString(),
+      expiresAt: new Date(Date.now() + 30_000).toISOString(),
     },
     capabilities: {
       statefulWorkspace: true,
@@ -469,6 +469,88 @@ test('worker shutdown interrupts reconnect backoff', async () => {
   await run;
 });
 
+test('sandbox completion does not cancel an in-flight credential rotation', async () => {
+  const key = createBridgeIdentity();
+  const identity = {
+    privateKey: key.privateKey,
+    credential: 'credential-before-in-flight-rotation',
+    expiresAt: new Date(Date.now() + 40).toISOString(),
+  };
+  let refreshStarted!: () => void;
+  const refreshStartedPromise = new Promise<void>((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<void>((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<string, string>
+    ).Authorization;
+    return Response.json({ protocolVersion: 1, accepted: true });
+  };
+  const worker = new BridgeWorker({
+    codeApiUrl: 'https://code.example/v1',
+    workerId: 'vm-1',
+    incarnationId,
+    sandboxEndpoint: 'http://127.0.0.1:2000/api/v2',
+    identity,
+    capabilities: {
+      statefulWorkspace: true,
+      sandboxProfile: 'nsjail',
+      runtimes: ['bash'],
+    },
+    fetchImpl,
+    credentialRefreshWindowMs: 30,
+  });
+
+  await worker.executeAndSettle({
+    protocolVersion: 1,
+    assignmentId: 'assignment-rotation-race',
+    workerId: 'vm-1',
+    incarnationId,
+    generation: 5,
+    leaseToken: 'assignment-rotation-race-lease-token',
+    expiresAt: new Date(Date.now() + 600_000).toISOString(),
+    request: { body: { language: 'bash' }, headers: {} },
+  });
+
+  assert.equal(refreshCount, 1);
+  assert.equal(identity.credential, 'credential-after-in-flight-rotation');
+  assert.equal(
+    settleAuthorization,
+    'Bridge credential-after-in-flight-rotation',
+  );
+});
+
 test('reconnect delay uses bounded exponential jitter', () => {
   assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 0), 500);
   assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 1), 1_000);

packages/code/src/worker.ts

@@ -32,6 +32,7 @@ export interface BridgeWorkerOptions {
   onError?: (error: unknown) => void;
   onIdentityChange?: (identity: BridgeWorkerIdentity) => void | Promise<void>;
   incarnationId?: string;
+  credentialRefreshWindowMs?: number;
 }
 
 export interface BridgeWorkerIdentity {
@@ -181,13 +182,14 @@ export class BridgeWorker {
 
   async refreshCredential(
     signal?: AbortSignal,
-    validThroughMs = Date.now() + CREDENTIAL_REFRESH_WINDOW_MS,
+    validThroughMs =
+      Date.now() +
+      (this.options.credentialRefreshWindowMs ??
+        CREDENTIAL_REFRESH_WINDOW_MS),
   ): Promise<void> {
     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<BridgeWorkerCredentialResponse>(
@@ -201,7 +203,8 @@ export class BridgeWorker {
       credential.workerId !== this.options.workerId ||
       typeof credential.credential !== 'string' ||
       credential.credential.length < 32 ||
-      !Number.isFinite(Date.parse(credential.expiresAt))
+      !Number.isFinite(Date.parse(credential.expiresAt)) ||
+      Date.parse(credential.expiresAt) <= validThroughMs
     ) {
       throw new BridgeProtocolError(
         'Code API returned an invalid rotated worker credential',
@@ -217,16 +220,55 @@ export class BridgeWorker {
     identity.expiresAt = rotatedIdentity.expiresAt;
   }
 
+  private async maintainCredential(
+    assignment: BridgeAssignment,
+    stopSignal: AbortSignal,
+    requestSignal?: AbortSignal,
+  ): Promise<void> {
+    const identity = this.options.identity;
+    if (identity == null) return;
+    const refreshWindowMs =
+      this.options.credentialRefreshWindowMs ??
+      CREDENTIAL_REFRESH_WINDOW_MS;
+    const assignmentDeadlineMs = Date.parse(assignment.expiresAt);
+    while (!stopSignal.aborted && Date.now() < assignmentDeadlineMs) {
+      const refreshAtMs = Date.parse(identity.expiresAt) - refreshWindowMs;
+      const waitMs = Math.max(
+        0,
+        Math.min(refreshAtMs - Date.now(), assignmentDeadlineMs - Date.now()),
+      );
+      if (waitMs > 0) {
+        await new Promise<void>((resolve) => {
+          const timer = setTimeout(resolve, waitMs);
+          stopSignal.addEventListener(
+            'abort',
+            () => {
+              clearTimeout(timer);
+              resolve();
+            },
+            { once: true },
+          );
+        });
+      }
+      if (stopSignal.aborted || Date.now() >= assignmentDeadlineMs) return;
+      /* Sandbox completion stops the maintenance loop, but it must not abort a
+       * refresh already accepted by the server. Only cancellation of the
+       * whole worker operation may cancel that request. */
+      await this.refreshCredential(requestSignal);
+    }
+  }
+
   async executeAndSettle(
     assignment: BridgeAssignment,
     signal?: AbortSignal,
   ): Promise<void> {
-    await this.refreshCredential(
-      signal,
-      Date.parse(assignment.expiresAt) + CREDENTIAL_REFRESH_WINDOW_MS,
-    );
+    await this.refreshCredential(signal);
     const executionController = new AbortController();
-    const abortExecution = (): void => executionController.abort();
+    const credentialController = new AbortController();
+    const abortExecution = (): void => {
+      executionController.abort();
+      credentialController.abort();
+    };
     signal?.addEventListener('abort', abortExecution, { once: true });
     const deadlineDelay = Math.max(
       0,
@@ -251,8 +293,18 @@ export class BridgeWorker {
       executionController,
       cancellationController.signal,
     );
+    let credentialMaintenanceError: unknown;
+    let credentialMaintenance: Promise<void> | undefined;
     let settlement: BridgeSettlement;
     try {
+      credentialMaintenance = this.maintainCredential(
+        assignment,
+        credentialController.signal,
+        signal,
+      ).catch((error) => {
+        credentialMaintenanceError = error;
+        executionController.abort();
+      });
       const headers = {
         ...assignment.request.headers,
         ...(assignment.runtimeSessionId
@@ -273,6 +325,9 @@ export class BridgeWorker {
       );
       const payload = (await response.json()) as object;
       if (heartbeatError != null) throw heartbeatError;
+      if (credentialMaintenanceError != null) {
+        throw credentialMaintenanceError;
+      }
       if (!response.ok) {
         throw new BridgeProtocolError(
           errorMessage(payload) ??
@@ -303,9 +358,12 @@ export class BridgeWorker {
     clearTimeout(deadlineTimer);
     heartbeatController.abort();
     await heartbeat;
+    credentialController.abort();
+    await credentialMaintenance;
     cancellationController.abort();
     await cancellationWatcher;
     signal?.removeEventListener('abort', abortExecution);
+    await this.refreshCredential(signal);
     await this.request<BridgeSettlementResponse>(
       this.assignmentUrl(assignment, 'settle'),
       settlement,

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<string, string | undefined>();
@@ -75,6 +76,7 @@ function baseClaims(overrides: Partial<JwtClaims> = {}): 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',
     });
   });
 

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,
   };
 }
 

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 {

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,
 });