fix: Release Unassigned Workspace Slots When Dispatch Cleanup Fails (LibreChat-AI#189)

Closes LibreChat-AI#170

Author: danny-avila

service/src/bridge/concurrent-store.test.ts

@@ -384,6 +384,46 @@ test('queued cancellation never leases and does not block another root', async (
   ).toBeUndefined();
 });
 
+test('unassigned slot releases even when dispatch cleanup fails', async () => {
+  await register();
+  const originalIncr = redis.incr.bind(redis);
+  const originalSet = redis.set.bind(redis);
+  const set = originalSet as (...args: unknown[]) => unknown;
+  redis.incr = ((key: string) =>
+    key.endsWith(':generation')
+      ? Promise.reject(new Error('injected generation outage'))
+      : originalIncr(key)) as typeof redis.incr;
+  redis.set = ((key: string, ...args: unknown[]) =>
+    key.endsWith(':cancelled')
+      ? Promise.reject(new Error('injected cancellation outage'))
+      : set(key, ...args)) as typeof redis.set;
+  try {
+    // The reservation succeeds, then dispatch fails before storing an assignment.
+    await expect(dispatch('a')).rejects.toThrow('injected cancellation outage');
+  } finally {
+    redis.incr = originalIncr;
+    redis.set = originalSet;
+  }
+  expect(
+    await redis.hlen(`codeapi:bridge:v1:worker:${workerId}:workspace-slots`),
+  ).toBe(0);
+  expect(
+    await redis.get(`codeapi:bridge:v1:worker:${workerId}:lock`),
+  ).toBeNull();
+  const next = dispatch('a');
+  const assignment = (await store.lease(
+    workerId,
+    incarnationId,
+    1000,
+    undefined,
+    undefined,
+    0,
+  ))!;
+  expect(assignment.request).toMatchObject({ workspaceId: 'a' });
+  await settle(assignment);
+  await expect(next).resolves.toMatchObject({ status: 'rejected' });
+});
+
 test('late quarantine releases its slot after caller cancellation and retains only its root fence', async () => {
   await register();
   const controller = new AbortController();

service/src/bridge/store.ts

@@ -1060,20 +1060,15 @@ export class RedisBridgeStore {
           // already committed result rather than turning cleanup availability
           // into a client-visible failure that could prompt duplicate work.
         }
+      } else if (workspaceSlots != null && assignment == null) {
+        await this.cleanupUnassignedSlot(
+          args.workerId,
+          lockIncarnationId,
+          assignmentId,
+        );
       } else {
         await this.cleanupDispatch(args.workerId, assignmentId, assignment);
       }
-      if (workspaceSlots != null && assignment == null) {
-        await boundedCommand(
-          workspaceSlots.release(
-            args.workerId,
-            lockIncarnationId,
-            assignmentId,
-          ),
-          this.redisCommandTimeoutMs,
-          'Bridge unassigned slot cleanup',
-        );
-      }
     }
   }
 
@@ -2138,6 +2133,29 @@ export class RedisBridgeStore {
     ]);
   }
 
+  private async cleanupUnassignedSlot(
+    workerId: string,
+    incarnationId: string,
+    assignmentId: string,
+  ): Promise<void> {
+    // No stored assignment owns this reservation, so a cancellation outage
+    // must not leave the slot and its root busy until TTL expiry.
+    const [cleanup, release] = await Promise.allSettled([
+      this.cleanupDispatch(workerId, assignmentId, undefined),
+      boundedCommand(
+        new BridgeWorkspaceSlots(this.redis).release(
+          workerId,
+          incarnationId,
+          assignmentId,
+        ),
+        this.redisCommandTimeoutMs,
+        'Bridge unassigned slot cleanup',
+      ),
+    ]);
+    if (cleanup.status === 'rejected') throw cleanup.reason;
+    if (release.status === 'rejected') throw release.reason;
+  }
+
   private async commitPendingWorkspace(
     assignment: StoredAssignment,
     settlement: AnyCodeBridgeSettlement,