[pull] main from LibreChat-AI:main

docs/byom-worker-admission.md

@@ -0,0 +1,26 @@
+# BYOM worker admission
+
+Workspace tool calls to a busy worker wait in a bounded FIFO shared through Redis.
+The limit is 32 admitted requests per worker, including the active request. When
+the limit is reached, the workspace endpoint returns HTTP 429 with
+`WORKER_QUEUE_FULL`. A different worker has an independent admission queue.
+
+Waiting uses the caller's existing absolute dispatch deadline. It does not reset
+or extend execution timeouts. Disconnecting or cancelling removes the waiting
+request without cancelling the active assignment. Expired entries are pruned;
+Redis key expiry also bounds state left by a crashed API process.
+
+After admission, the API revalidates the worker incarnation, identity, tenant
+binding and workspace operation. A waiting request cannot migrate to a replacement
+worker. Existing execution acknowledgement, fencing, settlement and quarantine
+rules remain responsible for the active assignment.
+
+This is compatible with existing workers and requires only a Code API update.
+Existing workers still execute one assignment at a time. Parallel execution across
+workspaces requires separate lease claims and isolated native sandbox contexts;
+this admission change does not advertise that capability. Queue time and execution
+time currently share the HTTP request deadline. Separate budgets require a matching
+LibreChat client change so that the client does not disconnect while waiting.
+
+Focused regression coverage lives in `service/src/bridge/admission.test.ts` and
+`service/src/bridge/worker-admission.test.ts`.

service/src/bridge/admission.test.ts

@@ -0,0 +1,37 @@
+import { afterEach, expect, test } from 'bun:test';
+import RedisMock from 'ioredis-mock';
+import type Redis from 'ioredis';
+import { BridgeAdmissionQueue } from './admission';
+
+const redis = new RedisMock() as unknown as Redis;
+afterEach(async () => {
+  await redis.flushall();
+});
+
+test('bounds FIFO admission across API replicas and releases cancelled waiters', async () => {
+  const firstReplica = new BridgeAdmissionQueue(redis, 2);
+  const secondReplica = new BridgeAdmissionQueue(redis, 2);
+  const deadline = Date.now() + 5000;
+  expect(await firstReplica.enter('worker', 'first', deadline)).toBe(true);
+  expect(await secondReplica.enter('worker', 'second', deadline)).toBe(true);
+  expect(await firstReplica.enter('worker', 'third', deadline)).toBe(false);
+  expect(await secondReplica.isHead('worker', 'second')).toBe(false);
+  await firstReplica.leave('worker', 'first');
+  expect(await secondReplica.isHead('worker', 'second')).toBe(true);
+  expect(await firstReplica.enter('worker', 'third', deadline)).toBe(true);
+  expect(await firstReplica.isHead('worker', 'third')).toBe(false);
+});
+
+test('expired crashed callers cannot strand the next request or consume capacity', async () => {
+  const queue = new BridgeAdmissionQueue(redis, 1);
+  expect(await queue.enter('worker', 'expired', Date.now() - 1)).toBe(true);
+  expect(await queue.enter('worker', 'live', Date.now() + 5000)).toBe(true);
+  expect(await queue.isHead('worker', 'live')).toBe(true);
+});
+
+test('different machines do not share admission capacity', async () => {
+  const queue = new BridgeAdmissionQueue(redis, 1);
+  expect(await queue.enter('worker-a', 'a', Date.now() + 5000)).toBe(true);
+  expect(await queue.enter('worker-b', 'b', Date.now() + 5000)).toBe(true);
+  expect(await queue.isHead('worker-b', 'b')).toBe(true);
+});

service/src/bridge/admission.ts

@@ -0,0 +1,90 @@
+import type Redis from 'ioredis';
+
+/** Bounded FIFO admission shared by API replicas. Entries expire after caller deadlines. */
+export class BridgeAdmissionQueue {
+  constructor(
+    private readonly redis: Redis,
+    private readonly capacity = 32,
+  ) {}
+
+  private keys(workerId: string): [string, string, string] {
+    const prefix = `codeapi:bridge:v1:worker:${encodeURIComponent(workerId)}:admission`;
+    return [prefix, `${prefix}:deadlines`, `${prefix}:sequence`];
+  }
+
+  async enter(
+    workerId: string,
+    id: string,
+    deadlineAtMs: number,
+  ): Promise<boolean> {
+    return (
+      Number(
+        await this.redis.eval(
+          [
+            "local expired = redis.call('ZRANGEBYSCORE', KEYS[2], '-inf', ARGV[2])",
+            'for _, id in ipairs(expired) do',
+            "  redis.call('ZREM', KEYS[1], id)",
+            "  redis.call('ZREM', KEYS[2], id)",
+            'end',
+            "if redis.call('ZSCORE', KEYS[1], ARGV[1]) then return 1 end",
+            "if redis.call('ZCARD', KEYS[1]) >= tonumber(ARGV[4]) then return 0 end",
+            "local sequence = redis.call('INCR', KEYS[3])",
+            "redis.call('ZADD', KEYS[1], sequence, ARGV[1])",
+            "redis.call('ZADD', KEYS[2], ARGV[3], ARGV[1])",
+            "local latest = redis.call('ZREVRANGE', KEYS[2], 0, 0, 'WITHSCORES')",
+            'for _, key in ipairs(KEYS) do',
+            "  redis.call('PEXPIREAT', key, tonumber(latest[2]) + 30000)",
+            'end',
+            'return 1',
+          ].join('\n'),
+          3,
+          ...this.keys(workerId),
+          id,
+          Date.now(),
+          deadlineAtMs,
+          this.capacity,
+        ),
+      ) === 1
+    );
+  }
+
+  async isHead(workerId: string, id: string): Promise<boolean> {
+    const [order, deadlines] = this.keys(workerId);
+    return (
+      Number(
+        await this.redis.eval(
+          [
+            "local expired = redis.call('ZRANGEBYSCORE', KEYS[2], '-inf', ARGV[2])",
+            'for _, id in ipairs(expired) do',
+            "  redis.call('ZREM', KEYS[1], id)",
+            "  redis.call('ZREM', KEYS[2], id)",
+            'end',
+            "local head = redis.call('ZRANGE', KEYS[1], 0, 0)",
+            'if head[1] == ARGV[1] then return 1 end',
+            'return 0',
+          ].join('\n'),
+          2,
+          order,
+          deadlines,
+          id,
+          Date.now(),
+        ),
+      ) === 1
+    );
+  }
+
+  async leave(workerId: string, id: string): Promise<void> {
+    const [order, deadlines] = this.keys(workerId);
+    await this.redis.eval(
+      [
+        "redis.call('ZREM', KEYS[1], ARGV[1])",
+        "redis.call('ZREM', KEYS[2], ARGV[1])",
+        'return 1',
+      ].join('\n'),
+      2,
+      order,
+      deadlines,
+      id,
+    );
+  }
+}

service/src/bridge/store.ts

@@ -18,6 +18,7 @@ import {
   isWorkspaceToolResult,
 } from '../../../packages/code/src/protocol';
 import type { BridgeWorkerBinding } from './pairing';
+import { BridgeAdmissionQueue } from './admission';
 
 const PREFIX = 'codeapi:bridge:v1';
 const POLL_INTERVAL_MS = 100;
@@ -43,6 +44,7 @@ export class BridgeStoreError extends Error {
       | 'WORKER_OFFLINE'
       | 'WORKER_UNAUTHORIZED'
       | 'WORKER_BUSY'
+      | 'WORKER_QUEUE_FULL'
       | 'ASSIGNMENT_EXPIRED'
       | 'ASSIGNMENT_FENCED'
       | 'ASSIGNMENT_NOT_FOUND'
@@ -684,25 +686,68 @@ export class RedisBridgeStore {
     const lockIncarnationId = registration.incarnationId;
     let assignment: StoredAssignment | undefined;
     let resultCommitted = false;
+    const admission = args.workspaceRequest == null
+      ? undefined
+      : new BridgeAdmissionQueue(this.redis);
     try {
-      const locked = await this.dispatchCommand(
-        () =>
-          this.acquireLock(
-            args.workerId,
-            assignmentId,
-            lockIncarnationId,
-            ttlSeconds,
-          ),
+      if (admission != null && !(await this.dispatchCommand(
+        () => admission.enter(args.workerId, assignmentId, args.deadlineAtMs),
         args,
-        'Bridge assignment lock acquisition',
-      );
+        'Bridge admission enqueue',
+      ))) {
+        throw new BridgeStoreError('WORKER_QUEUE_FULL', 'Bridge worker pending request limit reached');
+      }
+      let locked = false;
+      do {
+        if (admission != null && !(await this.dispatchCommand(
+          () => admission.isHead(args.workerId, assignmentId),
+          args,
+          'Bridge admission position',
+        ))) {
+          await delay(Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), args.signal);
+          continue;
+        }
+        locked = await this.dispatchCommand(
+          () =>
+            this.acquireLock(
+              args.workerId,
+              assignmentId,
+              lockIncarnationId,
+              ttlSeconds,
+            ),
+          args,
+          'Bridge assignment lock acquisition',
+        );
+        if (!locked && admission != null) {
+          await delay(Math.min(POLL_INTERVAL_MS, args.deadlineAtMs - Date.now()), args.signal);
+        }
+      } while (!locked && admission != null);
       if (!locked) {
         throw new BridgeStoreError(
           'WORKER_BUSY',
           `Bridge worker ${args.workerId} is busy`,
         );
       }
       this.assertDispatchActive(args.signal, args.deadlineAtMs);
+      if (admission != null) {
+        // Waiting must not transfer accepted work to a replacement machine or identity.
+        const current = await this.dispatchCommand(
+          () => this.dispatchableRegistration(args.workerId),
+          args,
+          'Bridge admitted worker validation',
+        );
+        if (
+          current == null ||
+          current.registration.incarnationId !== registration.incarnationId ||
+          current.registration.identityId !== registration.identityId ||
+          current.registration.binding?.tenantId !== registration.binding?.tenantId
+        ) {
+          throw new BridgeStoreError('WORKER_OFFLINE', 'Bridge worker changed while the request was waiting');
+        }
+        if (!supportsWorkspaceTool(current.registration, args.workspaceRequest!)) {
+          throw new BridgeStoreError('WORKER_MISMATCH', 'Bridge worker capabilities changed while the request was waiting');
+        }
+      }
       const generation = await this.dispatchCommand(
         () => this.redis.incr(generationKey(args.workerId)),
         args,
@@ -748,6 +793,12 @@ export class RedisBridgeStore {
           'Bridge assignment enqueue',
         );
         if (queued) break;
+        if (admission != null) {
+          throw new BridgeStoreError(
+            'WORKER_FENCED',
+            'Bridge worker changed before the waiting request could be dispatched',
+          );
+        }
         const replacement = await this.dispatchCommand(
           () => this.dispatchableRegistration(args.workerId),
           args,
@@ -810,6 +861,14 @@ export class RedisBridgeStore {
         throw error;
       }
     } finally {
+      if (admission != null) {
+        // Expiry remains the fallback if Redis is unavailable during cancellation.
+        await boundedCommand(
+          admission.leave(args.workerId, assignmentId),
+          this.redisCommandTimeoutMs,
+          'Bridge admission cleanup',
+        ).catch(() => undefined);
+      }
       if (resultCommitted) {
         try {
           await this.cleanupWithRetry(args.workerId, assignmentId, assignment);