summaryrefslogtreecommitdiffhomepage
path: root/packages/provider-concurrency
diff options
context:
space:
mode:
Diffstat (limited to 'packages/provider-concurrency')
-rw-r--r--packages/provider-concurrency/package.json12
-rw-r--r--packages/provider-concurrency/src/concurrency-manager.test.ts1260
-rw-r--r--packages/provider-concurrency/src/concurrency-manager.ts851
-rw-r--r--packages/provider-concurrency/src/extension.ts319
-rw-r--r--packages/provider-concurrency/src/index.ts10
-rw-r--r--packages/provider-concurrency/src/provider-wrapper.test.ts262
-rw-r--r--packages/provider-concurrency/src/provider-wrapper.ts77
-rw-r--r--packages/provider-concurrency/src/service.ts11
-rw-r--r--packages/provider-concurrency/tsconfig.json6
9 files changed, 2808 insertions, 0 deletions
diff --git a/packages/provider-concurrency/package.json b/packages/provider-concurrency/package.json
new file mode 100644
index 0000000..b53d599
--- /dev/null
+++ b/packages/provider-concurrency/package.json
@@ -0,0 +1,12 @@
+{
+ "name": "@dispatch/provider-concurrency",
+ "version": "0.0.0",
+ "type": "module",
+ "private": true,
+ "main": "dist/index.js",
+ "types": "dist/index.d.ts",
+ "dependencies": {
+ "@dispatch/conversation-store": "workspace:*",
+ "@dispatch/kernel": "workspace:*"
+ }
+}
diff --git a/packages/provider-concurrency/src/concurrency-manager.test.ts b/packages/provider-concurrency/src/concurrency-manager.test.ts
new file mode 100644
index 0000000..36c0ea3
--- /dev/null
+++ b/packages/provider-concurrency/src/concurrency-manager.test.ts
@@ -0,0 +1,1260 @@
+import type { ProviderUsage } from "@dispatch/kernel";
+import { describe, expect, it } from "vitest";
+import { type ConcurrencyService, createConcurrencyManager } from "./concurrency-manager.js";
+
+// ─── Fake timers ──────────────────────────────────────────────────────────────
+
+interface FakeTimer {
+ fire: () => void;
+ cleared: boolean;
+}
+
+function createFakeTimers() {
+ let currentTime = 0;
+ const intervals: FakeTimer[] = [];
+ const timeouts: { time: number; fire: () => void; cleared: boolean }[] = [];
+
+ const setInterval = ((_fn: () => void, _ms: number) => {
+ const timer: FakeTimer = { fire: () => _fn(), cleared: false };
+ intervals.push(timer);
+ return timer as unknown as ReturnType<typeof setInterval>;
+ }) as typeof setInterval;
+
+ const clearInterval = ((timer: ReturnType<typeof setInterval>) => {
+ const t = timer as unknown as FakeTimer;
+ t.cleared = true;
+ }) as typeof clearInterval;
+
+ const setTimeout = ((_fn: () => void, ms: number) => {
+ const entry = { time: currentTime + ms, fire: () => _fn(), cleared: false };
+ timeouts.push(entry);
+ return entry as unknown as ReturnType<typeof setTimeout>;
+ }) as typeof setTimeout;
+
+ const clearTimeout = ((timer: ReturnType<typeof setTimeout>) => {
+ const t = timer as unknown as { cleared: boolean };
+ t.cleared = true;
+ }) as typeof clearTimeout;
+
+ return {
+ now: () => currentTime,
+ advance(ms: number) {
+ currentTime += ms;
+ // Fire any due timeouts.
+ for (const entry of timeouts) {
+ if (!entry.cleared && entry.time <= currentTime) {
+ entry.cleared = true;
+ entry.fire();
+ }
+ }
+ },
+ fireIntervals() {
+ for (const timer of intervals) {
+ if (!timer.cleared) timer.fire();
+ }
+ },
+ setInterval,
+ clearInterval,
+ setTimeout,
+ clearTimeout,
+ };
+}
+
+function createManager(opts?: {
+ releaseCooldownMs?: number;
+ fetchUsage?: (providerId: string) => Promise<ProviderUsage | undefined>;
+ onLimitReduced?: (providerId: string, newLimit: number, oldLimit: number) => void;
+ onUsagePollError?: (providerId: string, err: unknown) => void;
+}): {
+ manager: ConcurrencyService;
+ timers: ReturnType<typeof createFakeTimers>;
+} {
+ const timers = createFakeTimers();
+ const manager = createConcurrencyManager({
+ now: timers.now,
+ slotTimeoutMs: 5000,
+ watchdogIntervalMs: 1000,
+ defaultPauseMs: 30000,
+ ...(opts?.releaseCooldownMs !== undefined ? { releaseCooldownMs: opts.releaseCooldownMs } : {}),
+ ...(opts?.fetchUsage !== undefined ? { fetchUsage: opts.fetchUsage } : {}),
+ ...(opts?.onLimitReduced !== undefined ? { onLimitReduced: opts.onLimitReduced } : {}),
+ ...(opts?.onUsagePollError !== undefined ? { onUsagePollError: opts.onUsagePollError } : {}),
+ setTimeout: timers.setTimeout,
+ clearTimeout: timers.clearTimeout,
+ setInterval: timers.setInterval,
+ clearInterval: timers.clearInterval,
+ });
+ return { manager, timers };
+}
+
+describe("createConcurrencyManager", () => {
+ it("returns no-op release for providers with no configured limit", async () => {
+ const { manager } = createManager();
+ const release = await manager.acquire("unknown", "conv1", "default", 0);
+ expect(typeof release).toBe("function");
+ // No state → release is a no-op, no error.
+ release();
+ expect(manager.getStatus("unknown")).toBeUndefined();
+ });
+
+ it("grants immediately when under the limit", async () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 4);
+
+ const release1 = await manager.acquire("umans", "conv1", "default", 0);
+ const status = manager.getStatus("umans");
+ expect(status).toEqual({
+ providerId: "umans",
+ limit: 4,
+ inFlight: 1,
+ queued: 0,
+ paused: false,
+ cooldownMs: 0,
+ autoReduced: false,
+ });
+ release1();
+ expect(manager.getStatus("umans")?.inFlight).toBe(0);
+ });
+
+ it("queues when at the limit and grants on release (FIFO when same priority)", async () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 1);
+
+ const release1 = await manager.acquire("umans", "conv1", "default", 100);
+
+ // Second request should block (at limit).
+ let resolved = false;
+ const promise2 = manager.acquire("umans", "conv2", "default", 200).then((r) => {
+ resolved = true;
+ return r;
+ });
+
+ // Let microtasks settle.
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolved).toBe(false);
+ expect(manager.getStatus("umans")?.queued).toBe(1);
+
+ // Release the first slot.
+ release1();
+
+ const release2 = await promise2;
+ expect(resolved).toBe(true);
+ expect(manager.getStatus("umans")?.inFlight).toBe(1);
+ expect(manager.getStatus("umans")?.queued).toBe(0);
+ release2();
+ });
+
+ it("grants to the oldest agent first (priority queue by promptStartedAt)", async () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 1);
+
+ // Hold the single slot.
+ const release0 = await manager.acquire("umans", "holder", "default", 0);
+
+ // Three agents queue with different prompt start times.
+ // Agent C started latest (t=300), Agent A started earliest (t=100).
+ const results: string[] = [];
+ const acquireAndRecord = (conv: string, promptAt: number) =>
+ manager.acquire("umans", conv, "default", promptAt).then((r) => {
+ results.push(conv);
+ return r;
+ });
+
+ // Queue in non-sorted order: B (t=200), A (t=100), C (t=300).
+ const pB = acquireAndRecord("convB", 200);
+ const pA = acquireAndRecord("convA", 100);
+ const pC = acquireAndRecord("convC", 300);
+
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(results).toEqual([]); // none resolved yet.
+
+ // Release the holder. The oldest agent (A, t=100) should get the slot first.
+ release0();
+
+ const rA = await pA;
+ expect(results).toEqual(["convA"]);
+
+ rA.release ? rA.release() : rA();
+
+ // Now B (t=200) should be next.
+ const rB = await pB;
+ expect(results).toEqual(["convA", "convB"]);
+ rB.release ? rB.release() : rB();
+
+ // Then C (t=300).
+ const rC = await pC;
+ expect(results).toEqual(["convA", "convB", "convC"]);
+ rC.release ? rC.release() : rC();
+ });
+
+ it("does not grant slots while paused (429 backoff)", async () => {
+ const { manager, timers } = createManager();
+ manager.setLimit("umans", 1);
+
+ const release1 = await manager.acquire("umans", "conv1", "default", 0);
+ release1();
+
+ // Simulate a 429 → queue pauses.
+ manager.reportRateLimit("umans");
+ const status = manager.getStatus("umans");
+ expect(status?.paused).toBe(true);
+ expect(status?.pausedUntil).toBe(30000);
+
+ // A new acquire should block (paused, even though under limit).
+ let resolved = false;
+ const promise = manager.acquire("umans", "conv2", "default", 0).then((r) => {
+ resolved = true;
+ return r;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolved).toBe(false);
+
+ // Advance past the pause duration.
+ timers.advance(30000);
+
+ const release2 = await promise;
+ expect(resolved).toBe(true);
+ expect(manager.getStatus("umans")?.paused).toBe(false);
+ release2();
+ });
+
+ it("respects retryAfterMs for 429 backoff", () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 2);
+
+ manager.reportRateLimit("umans", 5000);
+ expect(manager.getStatus("umans")?.pausedUntil).toBe(5000);
+ });
+
+ it("watchdog reclaims slots held beyond the timeout", async () => {
+ const { manager, timers } = createManager();
+ manager.setLimit("umans", 1);
+
+ const release = await manager.acquire("umans", "conv1", "default", 0);
+ expect(manager.getStatus("umans")?.inFlight).toBe(1);
+
+ // Advance past the slot timeout (5000ms) and fire the watchdog.
+ timers.advance(5001);
+ timers.fireIntervals();
+
+ // The watchdog should have force-released the slot.
+ expect(manager.getStatus("umans")?.inFlight).toBe(0);
+
+ // Calling release again (from the holder) should be a no-op (idempotent).
+ release();
+ expect(manager.getStatus("umans")?.inFlight).toBe(0);
+ });
+
+ it("watchdog grants the next waiter after reclaiming a stale slot", async () => {
+ const { manager, timers } = createManager();
+ manager.setLimit("umans", 1);
+
+ // Hold the slot.
+ await manager.acquire("umans", "holder", "default", 0);
+
+ // Queue a waiter.
+ let resolved = false;
+ const promise = manager.acquire("umans", "waiter", "default", 10).then((r) => {
+ resolved = true;
+ return r;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolved).toBe(false);
+
+ // Watchdog reclaims the held slot.
+ timers.advance(5001);
+ timers.fireIntervals();
+
+ // The waiter should now be granted.
+ const release = await promise;
+ expect(resolved).toBe(true);
+ expect(manager.getStatus("umans")?.inFlight).toBe(1);
+ release();
+ });
+
+ it("setLimit grants queued requests when the limit increases", async () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 1);
+
+ const release1 = await manager.acquire("umans", "conv1", "default", 0);
+
+ // Queue a waiter.
+ let resolved = false;
+ const promise = manager.acquire("umans", "conv2", "default", 100).then((r) => {
+ resolved = true;
+ return r;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolved).toBe(false);
+
+ // Increase the limit → the queued request should be granted.
+ manager.setLimit("umans", 2);
+
+ const release2 = await promise;
+ expect(resolved).toBe(true);
+ expect(manager.getStatus("umans")?.inFlight).toBe(2);
+
+ release2();
+ release1();
+ });
+
+ it("removeLimit grants all queued requests and removes the state", async () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 1);
+
+ const release1 = await manager.acquire("umans", "conv1", "default", 0);
+
+ // Queue two waiters.
+ const p2 = manager.acquire("umans", "conv2", "default", 100);
+ const p3 = manager.acquire("umans", "conv3", "default", 200);
+ await Promise.resolve();
+ await Promise.resolve();
+
+ // Remove the limit → all queued requests should be granted.
+ manager.removeLimit("umans");
+
+ const r2 = await p2;
+ const r3 = await p3;
+ expect(manager.getStatus("umans")).toBeUndefined();
+
+ // Releases work (no error after state removal).
+ r2();
+ r3();
+ release1();
+ });
+
+ it("getLimits returns all configured limits", () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 4);
+ manager.setLimit("openai-compat", 5);
+
+ const limits = manager.getLimits();
+ expect(limits).toHaveLength(2);
+ expect(limits).toContainEqual({ providerId: "umans", limit: 4 });
+ expect(limits).toContainEqual({ providerId: "openai-compat", limit: 5 });
+ });
+
+ it("getStatusAll returns status for all configured providers", () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 4);
+ manager.setLimit("anthropic", 3);
+
+ const statuses = manager.getStatusAll();
+ expect(statuses).toHaveLength(2);
+ const umans = statuses.find((s) => s.providerId === "umans");
+ expect(umans).toEqual({
+ providerId: "umans",
+ limit: 4,
+ inFlight: 0,
+ queued: 0,
+ paused: false,
+ cooldownMs: 0,
+ autoReduced: false,
+ });
+ });
+
+ it("destroy clears timers without error", () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 4);
+ manager.reportRateLimit("umans", 5000);
+ expect(() => manager.destroy()).not.toThrow();
+ });
+
+ it("release is idempotent (double-release does not overshoot)", async () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 2);
+
+ const release = await manager.acquire("umans", "conv1", "default", 0);
+ expect(manager.getStatus("umans")?.inFlight).toBe(1);
+
+ release();
+ expect(manager.getStatus("umans")?.inFlight).toBe(0);
+
+ // Double-release should not decrement below 0.
+ release();
+ expect(manager.getStatus("umans")?.inFlight).toBe(0);
+ });
+
+ it("multiple concurrent acquires up to the limit all resolve immediately", async () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 3);
+
+ const releases = await Promise.all([
+ manager.acquire("umans", "conv1", "default", 0),
+ manager.acquire("umans", "conv2", "default", 0),
+ manager.acquire("umans", "conv3", "default", 0),
+ ]);
+
+ expect(manager.getStatus("umans")?.inFlight).toBe(3);
+
+ for (const release of releases) {
+ release();
+ }
+ expect(manager.getStatus("umans")?.inFlight).toBe(0);
+ });
+
+ it("release cooldown delays slot recycling (inFlight stays incremented during cooldown)", async () => {
+ const { manager, timers } = createManager({ releaseCooldownMs: 200 });
+ manager.setLimit("umans", 1);
+
+ const release1 = await manager.acquire("umans", "conv1", "default", 0);
+ expect(manager.getStatus("umans")?.inFlight).toBe(1);
+
+ // Queue a waiter.
+ let resolved = false;
+ const promise2 = manager.acquire("umans", "conv2", "default", 100).then((r) => {
+ resolved = true;
+ return r;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolved).toBe(false);
+ expect(manager.getStatus("umans")?.queued).toBe(1);
+
+ // Release the slot — inFlight should stay 1 (cooldown active).
+ release1();
+ expect(manager.getStatus("umans")?.inFlight).toBe(1);
+ expect(resolved).toBe(false); // waiter NOT granted yet
+
+ // Advance past the cooldown.
+ timers.advance(200);
+
+ // Now the slot is recycled and the waiter is granted.
+ const release2 = await promise2;
+ expect(resolved).toBe(true);
+ expect(manager.getStatus("umans")?.inFlight).toBe(1);
+ expect(manager.getStatus("umans")?.queued).toBe(0);
+ release2();
+ });
+
+ it("release cooldown is idempotent (double-release only schedules one cooldown)", async () => {
+ const { manager, timers } = createManager({ releaseCooldownMs: 200 });
+ manager.setLimit("umans", 2);
+
+ const release = await manager.acquire("umans", "conv1", "default", 0);
+ expect(manager.getStatus("umans")?.inFlight).toBe(1);
+
+ release();
+ expect(manager.getStatus("umans")?.inFlight).toBe(1); // still 1 (cooldown)
+
+ // Double-release should not schedule a second cooldown.
+ release();
+
+ // After cooldown, inFlight should drop by exactly 1 (to 0), not 2.
+ timers.advance(200);
+ expect(manager.getStatus("umans")?.inFlight).toBe(0);
+ });
+
+ it("destroy clears cooldown timers without error", () => {
+ const { manager } = createManager({ releaseCooldownMs: 200 });
+ manager.setLimit("umans", 1);
+ // Acquire + release to schedule a cooldown timer.
+ manager.acquire("umans", "conv1", "default", 0).then((release) => {
+ release();
+ // Now there's a pending cooldown timer — destroy should clean it up.
+ expect(() => manager.destroy()).not.toThrow();
+ });
+ });
+
+ it("onQueued is called when the request is enqueued (not granted immediately)", async () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 1);
+
+ // Hold the single slot.
+ const release1 = await manager.acquire("umans", "conv1", "default", 0);
+
+ // Second request should trigger onQueued.
+ let queuedCalled = false;
+ const promise = manager.acquire("umans", "conv2", "default", 100, () => {
+ queuedCalled = true;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+
+ expect(queuedCalled).toBe(true);
+ expect(manager.getStatus("umans")?.queued).toBe(1);
+
+ // Release the slot — the queued request should be granted.
+ release1();
+ const release2 = await promise;
+ release2();
+ });
+
+ it("onQueued is NOT called when the slot is granted immediately", async () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 2);
+
+ let queuedCalled = false;
+ const release = await manager.acquire("umans", "conv1", "default", 0, () => {
+ queuedCalled = true;
+ });
+
+ expect(queuedCalled).toBe(false);
+ release();
+ });
+
+ // ─── Configurable cooldown ──────────────────────────────────────────────
+
+ it("setCooldown changes the cooldown applied to subsequently recycled slots", async () => {
+ const { manager, timers } = createManager({ releaseCooldownMs: 200 });
+ manager.setLimit("umans", 1);
+
+ const release1 = await manager.acquire("umans", "conv1", "default", 0);
+ expect(manager.getStatus("umans")?.cooldownMs).toBe(200);
+
+ // Bump the cooldown to 500ms.
+ manager.setCooldown("umans", 500);
+ expect(manager.getStatus("umans")?.cooldownMs).toBe(500);
+
+ // Queue a waiter.
+ let resolved = false;
+ const promise2 = manager.acquire("umans", "conv2", "default", 100).then((r) => {
+ resolved = true;
+ return r;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolved).toBe(false);
+
+ // Release — the NEW cooldown (500ms) applies.
+ release1();
+ expect(resolved).toBe(false);
+ timers.advance(200); // old cooldown elapsed — still cooling (500ms now).
+ expect(resolved).toBe(false);
+ timers.advance(300); // 500ms total → slot recycled, waiter granted.
+ const release2 = await promise2;
+ expect(resolved).toBe(true);
+ release2();
+ });
+
+ it("getCooldowns returns all configured cooldowns", () => {
+ const { manager } = createManager({ releaseCooldownMs: 350 });
+ manager.setLimit("umans", 4);
+ manager.setCooldown("openai-compat", 100);
+
+ const cooldowns = manager.getCooldowns();
+ expect(cooldowns).toContainEqual({ providerId: "umans", cooldownMs: 350 });
+ expect(cooldowns).toContainEqual({ providerId: "openai-compat", cooldownMs: 100 });
+ });
+
+ it("setCooldown does NOT impose a limit when none is configured (override seeds on setLimit)", async () => {
+ const { manager } = createManager({ releaseCooldownMs: 350 });
+
+ // Set a cooldown with NO limit configured yet.
+ manager.setCooldown("umans", 500);
+ expect(manager.getCooldown("umans")).toBe(500);
+ // No limit → acquire must be unlimited (no state with a limit imposed).
+ const release = await manager.acquire("umans", "conv1", "default", 0);
+ expect(typeof release).toBe("function");
+ release();
+ expect(manager.getStatus("umans")).toBeUndefined(); // no limit state created
+ expect(manager.getLimit("umans")).toBeUndefined();
+
+ // Now set a limit — the pending cooldown override seeds the new state.
+ manager.setLimit("umans", 4);
+ expect(manager.getStatus("umans")?.cooldownMs).toBe(500);
+ });
+
+ // ─── Adaptive headroom (reduce limit by 1 on 429) ────────────────────────
+
+ it("reportRateLimit reduces the limit by 1 (one-way) and sets autoReduced notice", () => {
+ const reduced: { providerId: string; newLimit: number; oldLimit: number }[] = [];
+ const { manager } = createManager({
+ onLimitReduced: (p, n, o) => reduced.push({ providerId: p, newLimit: n, oldLimit: o }),
+ });
+ manager.setLimit("umans", 4);
+
+ manager.reportRateLimit("umans");
+
+ expect(manager.getLimit("umans")).toBe(3);
+ const status = manager.getStatus("umans");
+ expect(status?.autoReduced).toBe(true);
+ expect(status?.autoReducedFrom).toBe(4);
+ expect(status?.notice).toContain("auto-reduced to 3");
+ expect(reduced).toEqual([{ providerId: "umans", newLimit: 3, oldLimit: 4 }]);
+ });
+
+ it("repeated 429s keep reducing (4 -> 3 -> 2 -> 1) and never go below 1", () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 4);
+
+ manager.reportRateLimit("umans");
+ expect(manager.getLimit("umans")).toBe(3);
+ manager.reportRateLimit("umans");
+ expect(manager.getLimit("umans")).toBe(2);
+ manager.reportRateLimit("umans");
+ expect(manager.getLimit("umans")).toBe(1);
+ // Already at the floor — stays 1.
+ manager.reportRateLimit("umans");
+ expect(manager.getLimit("umans")).toBe(1);
+ const status = manager.getStatus("umans");
+ expect(status?.autoReduced).toBe(true);
+ // autoReducedFrom records the FIRST reduction's original limit (4).
+ expect(status?.autoReducedFrom).toBe(4);
+ });
+
+ it("a MANUAL setLimit clears the auto-reduce notice", () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 4);
+ manager.reportRateLimit("umans"); // 4 -> 3, autoReduced
+ expect(manager.getStatus("umans")?.autoReduced).toBe(true);
+
+ // User restores the limit manually.
+ manager.setLimit("umans", 4);
+ const status = manager.getStatus("umans");
+ expect(status?.autoReduced).toBe(false);
+ expect(status?.autoReducedFrom).toBeUndefined();
+ expect(status?.notice).toBeUndefined();
+ });
+
+ it("removeLimit clears the auto-reduce state", () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 4);
+ manager.reportRateLimit("umans"); // auto-reduced
+ expect(manager.getStatus("umans")?.autoReduced).toBe(true);
+
+ manager.removeLimit("umans");
+ expect(manager.getStatus("umans")).toBeUndefined();
+ });
+
+ // ─── Usage gate (poll concurrent_sessions before granting queued agents) ─
+
+ it("usage gate blocks a queued waiter while upstream concurrent_sessions >= limit", async () => {
+ // Upstream always reports AT the limit (4) → the gate never admits.
+ const { manager, timers } = createManager({
+ fetchUsage: async () => ({ concurrentSessions: 4 }),
+ });
+ manager.setLimit("umans", 4);
+ manager.setCooldown("umans", 0); // no cooldown — isolate the gate
+
+ // Fill all 4 slots (fast-path, no gate).
+ const releases = await Promise.all([
+ manager.acquire("umans", "c1", "default", 0),
+ manager.acquire("umans", "c2", "default", 0),
+ manager.acquire("umans", "c3", "default", 0),
+ manager.acquire("umans", "c4", "default", 0),
+ ]);
+ expect(manager.getStatus("umans")?.inFlight).toBe(4);
+
+ // 5th agent queues.
+ let resolved = false;
+ const promise5 = manager.acquire("umans", "c5", "default", 10).then((r) => {
+ resolved = true;
+ return r;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolved).toBe(false);
+ expect(manager.getStatus("umans")?.queued).toBe(1);
+
+ // Release one slot. Cooldown is 0 → recycle polls upstream → 4 >= 4 → NOT granted.
+ releases[0]?.();
+ // Let the async poll settle.
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolved).toBe(false);
+ expect(manager.getStatus("umans")?.queued).toBe(1);
+
+ // Advance past the 1s fallback repoll — still 4 → still blocked.
+ timers.advance(1000);
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolved).toBe(false);
+
+ for (const r of releases.slice(1)) r?.();
+ void promise5;
+ });
+
+ it("usage gate admits a queued waiter once upstream concurrent_sessions < limit", async () => {
+ // Upstream starts at the limit; drops to 3 after the release.
+ let upstream = 4;
+ const { manager } = createManager({
+ fetchUsage: async () => ({ concurrentSessions: upstream }),
+ });
+ manager.setLimit("umans", 4);
+ manager.setCooldown("umans", 0);
+
+ const releases = await Promise.all([
+ manager.acquire("umans", "c1", "default", 0),
+ manager.acquire("umans", "c2", "default", 0),
+ manager.acquire("umans", "c3", "default", 0),
+ manager.acquire("umans", "c4", "default", 0),
+ ]);
+
+ let resolved = false;
+ const promise5 = manager.acquire("umans", "c5", "default", 10).then((r) => {
+ resolved = true;
+ return r;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolved).toBe(false);
+
+ // Upstream now drops to 3 (the released session finally decremented).
+ upstream = 3;
+ // Release a slot → cooldown 0 → poll → 3 < 4 → admit the waiter.
+ releases[0]?.();
+ const release5 = await promise5;
+ expect(resolved).toBe(true);
+ expect(manager.getStatus("umans")?.queued).toBe(0);
+
+ release5();
+ for (const r of releases.slice(1)) r?.();
+ });
+
+ it("usage gate falls back to granting when fetchUsage returns undefined", async () => {
+ const { manager } = createManager({
+ fetchUsage: async () => undefined, // no usage info available
+ });
+ manager.setLimit("umans", 1);
+ manager.setCooldown("umans", 0);
+
+ const release1 = await manager.acquire("umans", "c1", "default", 0);
+ let resolved = false;
+ const promise2 = manager.acquire("umans", "c2", "default", 10).then((r) => {
+ resolved = true;
+ return r;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolved).toBe(false);
+
+ // Release → poll returns undefined → fall back to cooldown-only (grant).
+ release1();
+ const release2 = await promise2;
+ expect(resolved).toBe(true);
+ release2();
+ });
+
+ it("usage gate admits at most ONE queued waiter per successful poll", async () => {
+ let upstream = 4;
+ const { manager } = createManager({
+ fetchUsage: async () => ({ concurrentSessions: upstream }),
+ });
+ manager.setLimit("umans", 4);
+ manager.setCooldown("umans", 0);
+
+ const releases = await Promise.all([
+ manager.acquire("umans", "c1", "default", 0),
+ manager.acquire("umans", "c2", "default", 0),
+ manager.acquire("umans", "c3", "default", 0),
+ manager.acquire("umans", "c4", "default", 0),
+ ]);
+
+ // Queue two waiters.
+ let r5 = false;
+ let r6 = false;
+ const p5 = manager.acquire("umans", "c5", "default", 10).then((r) => {
+ r5 = true;
+ return r;
+ });
+ const p6 = manager.acquire("umans", "c6", "default", 20).then((r) => {
+ r6 = true;
+ return r;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(manager.getStatus("umans")?.queued).toBe(2);
+
+ // Upstream drops to 3. Release one slot → poll 3 < 4 → admit ONE (c5).
+ upstream = 3;
+ releases[0]?.();
+ await Promise.resolve();
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(r5).toBe(true);
+ expect(r6).toBe(false); // c6 still queued — needs another poll.
+ expect(manager.getStatus("umans")?.queued).toBe(1);
+
+ const release5 = await p5;
+ release5();
+ void p6;
+ for (const r of releases.slice(1)) r?.();
+ });
+
+ it("usage gate clears the fallback repoll timer when the queue drains", () => {
+ const { manager, timers } = createManager({
+ fetchUsage: async () => ({ concurrentSessions: 0 }),
+ });
+ manager.setLimit("umans", 1);
+ manager.setCooldown("umans", 0);
+
+ return manager.acquire("umans", "c1", "default", 0).then(async (release1) => {
+ // Queue a waiter (arms the 1s fallback timer).
+ const p2 = manager.acquire("umans", "c2", "default", 10);
+ await Promise.resolve();
+ await Promise.resolve();
+
+ // Release → poll 0 < 1 → grant → queue drains → fallback timer cleared.
+ release1();
+ const release2 = await p2;
+ release2();
+
+ // Advancing past 1s must NOT throw or fire at a drained state.
+ expect(() => timers.advance(1000)).not.toThrow();
+ });
+ });
+
+ // ─── Bug 1: usage-gate fast-path anti-overshoot ─────────────────────────
+
+ it("Bug 1: a concurrent acquire during a recycle-poll queues instead of fast-pathing (no overshoot)", async () => {
+ // The poll resolves only on an explicit microtask flush (deferred), so a
+ // concurrent acquire arriving mid-poll must see gatePolling/inflated inFlight.
+ let resolvePoll: (snap: ProviderUsage) => void = () => {};
+ const pollCalled: number[] = [];
+ const { manager } = createManager({
+ fetchUsage: () =>
+ new Promise<ProviderUsage>((resolve) => {
+ pollCalled.push(1);
+ resolvePoll = resolve;
+ }),
+ });
+ manager.setLimit("umans", 1);
+ manager.setCooldown("umans", 0);
+
+ // Hold the single slot.
+ const release1 = await manager.acquire("umans", "c1", "default", 0);
+ expect(manager.getStatus("umans")?.inFlight).toBe(1);
+
+ // Queue a waiter (c2). Cooldown is 0, but the gate defers admission until a poll.
+ let c2Granted = false;
+ const p2 = manager.acquire("umans", "c2", "default", 10).then((r) => {
+ c2Granted = true;
+ return r;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+
+ // Release c1 → recycle → poll started (inFlight held inflated during poll).
+ release1();
+ await Promise.resolve(); // let recycle schedule the poll
+ await Promise.resolve();
+ expect(pollCalled.length).toBeGreaterThanOrEqual(1);
+ // inFlight is still 1 (the recycle's decrement is deferred until the poll).
+ expect(manager.getStatus("umans")?.inFlight).toBe(1);
+
+ // A NEW acquire arriving mid-poll: inFlight is 1 (== limit) → must QUEUE,
+ // not fast-path. Even if it saw inFlight < limit, gatePolling would route it
+ // through the queue. Either way it must NOT be granted yet.
+ let c3Granted = false;
+ const p3 = manager.acquire("umans", "c3", "default", 20).then((r) => {
+ c3Granted = true;
+ return r;
+ });
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(c3Granted).toBe(false);
+ expect(manager.getStatus("umans")?.queued).toBeGreaterThanOrEqual(1);
+
+ // Resolve the poll with room (0 < 1) → c2 admitted (inFlight: decrement then
+ // re-increment for the grant). c3 stays queued (one admission per poll).
+ resolvePoll({ concurrentSessions: 0 });
+ const release2 = await p2;
+ expect(c2Granted).toBe(true);
+ // c3 NOT admitted by this poll (one per poll).
+ expect(c3Granted).toBe(false);
+
+ release2();
+ void p3;
+ });
+
+ it("Bug 1: when no poll is in flight, the fast-path still grants immediately (common-case throughput preserved)", async () => {
+ const { manager } = createManager({
+ fetchUsage: async () => ({ concurrentSessions: 0 }),
+ });
+ manager.setLimit("umans", 4);
+
+ // Nowhere near the limit, no recycle in progress → fast-path, no poll.
+ const release = await manager.acquire("umans", "c1", "default", 0);
+ expect(manager.getStatus("umans")?.inFlight).toBe(1);
+ release();
+ });
+
+ // ─── Bug 2: fetchUsage exceptions don't become unhandled rejections ──────
+
+ it("Bug 2: a throwing fetchUsage is treated as undefined (cooldown-only fallback) and fires onUsagePollError", async () => {
+ let pollError: { providerId: string; err: unknown } | undefined;
+ const { manager } = createManager({
+ fetchUsage: async () => {
+ throw new Error("usage endpoint exploded");
+ },
+ onUsagePollError: (providerId, err) => {
+ pollError = { providerId, err };
+ },
+ });
+ manager.setLimit("umans", 1);
+ manager.setCooldown("umans", 0);
+
+ const release1 = await manager.acquire("umans", "c1", "default", 0);
+ // Queue a waiter; release → recycle → poll THROWS.
+ const p2 = manager.acquire("umans", "c2", "default", 10);
+ await Promise.resolve();
+ await Promise.resolve();
+
+ // Must NOT reject / throw unhandled — swallow + fall back to granting.
+ release1();
+ const release2 = await p2; // resolves (cooldown-only fallback grants).
+ expect(release2).toBeTypeOf("function");
+ expect(pollError?.providerId).toBe("umans");
+ expect(pollError?.err).toBeInstanceOf(Error);
+ release2();
+ });
+
+ it("Bug 2: no unhandled promise rejection is left when fetchUsage throws (process stays clean)", async () => {
+ const rejections: unknown[] = [];
+ const handler = (reason: unknown) => rejections.push(reason);
+ process.on("unhandledRejection", handler);
+ try {
+ const { manager } = createManager({
+ fetchUsage: async () => {
+ throw new Error("boom");
+ },
+ });
+ manager.setLimit("umans", 1);
+ manager.setCooldown("umans", 0);
+
+ const release1 = await manager.acquire("umans", "c1", "default", 0);
+ manager.acquire("umans", "c2", "default", 10).then((r) => r()); // queue + auto-release
+ await Promise.resolve();
+ await Promise.resolve();
+ release1();
+ // Let the swallowed poll + grant settle fully.
+ await new Promise((r) => setTimeout(r, 5));
+ await new Promise((r) => setTimeout(r, 5));
+ expect(rejections).toEqual([]);
+ } finally {
+ process.off("unhandledRejection", handler);
+ }
+ });
+
+ // ─── Bug 3: persisted auto-reduced limit keeps its notice across restart ─
+
+ it("Bug 3: restoreLimit (startup) preserves the auto-reduce notice that setLimit (manual) clears", () => {
+ const { manager } = createManager();
+ manager.setLimit("umans", 4);
+ manager.reportRateLimit("umans"); // 4 -> 3, autoReduced
+ expect(manager.getStatus("umans")?.autoReduced).toBe(true);
+ expect(manager.getStatus("umans")?.autoReducedFrom).toBe(4);
+
+ // Simulate a restart: a fresh manager restores the persisted limit (3) +
+ // the auto-reduce marker (autoReducedFrom=4) via restoreLimit.
+ const { manager: restarted } = createManager();
+ restarted.restoreLimit("umans", 3, 4);
+ const status = restarted.getStatus("umans");
+ expect(status?.limit).toBe(3);
+ expect(status?.autoReduced).toBe(true);
+ expect(status?.autoReducedFrom).toBe(4);
+ expect(status?.notice).toContain("auto-reduced to 3");
+
+ // Contrast: a MANUAL setLimit clears the notice (user took control).
+ restarted.setLimit("umans", 4);
+ expect(restarted.getStatus("umans")?.autoReduced).toBe(false);
+ expect(restarted.getStatus("umans")?.autoReducedFrom).toBeUndefined();
+ });
+
+ it("Bug 3: restoreLimit without autoReducedFrom does not synthesize a notice", () => {
+ const { manager } = createManager();
+ manager.restoreLimit("umans", 4);
+ const status = manager.getStatus("umans");
+ expect(status?.limit).toBe(4);
+ expect(status?.autoReduced).toBe(false);
+ expect(status?.notice).toBeUndefined();
+ });
+});
+
+// ─── Starred-workspace priority tests ───────────────────────────────────────
+
+describe("starred-workspace priority", () => {
+ it("starred-workspace agents are admitted before non-starred (regardless of promptStartedAt)", async () => {
+ // Use a callback backed by a Set so we can star/unstar at runtime.
+ const starred = new Set<string>();
+ const timers = createFakeTimers();
+ const manager = createConcurrencyManager({
+ now: timers.now,
+ slotTimeoutMs: 5000,
+ watchdogIntervalMs: 1000,
+ defaultPauseMs: 30000,
+ isWorkspaceStarred: (wsId: string) => starred.has(wsId),
+ setTimeout: timers.setTimeout,
+ clearTimeout: timers.clearTimeout,
+ setInterval: timers.setInterval,
+ clearInterval: timers.clearInterval,
+ });
+ manager.setLimit("umans", 1);
+
+ // Hold the single slot.
+ const release0 = await manager.acquire("umans", "holder", "default", 0);
+
+ // Three agents queue:
+ // - convA (workspace "ws-normal", promptAt=100) — non-starred, earliest
+ // - convB (workspace "ws-starred", promptAt=200) — starred, later
+ // - convC (workspace "ws-normal", promptAt=300) — non-starred, latest
+ starred.add("ws-starred");
+
+ const results: string[] = [];
+ const acquireAndRecord = (conv: string, wsId: string, promptAt: number) =>
+ manager.acquire("umans", conv, wsId, promptAt).then((r) => {
+ results.push(conv);
+ return r;
+ });
+
+ const pA = acquireAndRecord("convA", "ws-normal", 100);
+ const pB = acquireAndRecord("convB", "ws-starred", 200);
+ const pC = acquireAndRecord("convC", "ws-normal", 300);
+
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(results).toEqual([]); // none resolved yet.
+
+ // Release the holder. The starred agent (convB, t=200) should get the
+ // slot FIRST, even though convA (t=100) started earlier.
+ release0();
+
+ const rB = await pB;
+ expect(results).toEqual(["convB"]);
+ rB();
+
+ // Now the oldest non-starred (convA, t=100) should be next.
+ const rA = await pA;
+ expect(results).toEqual(["convB", "convA"]);
+ rA();
+
+ // Then convC (t=300).
+ const rC = await pC;
+ expect(results).toEqual(["convB", "convA", "convC"]);
+ rC();
+
+ manager.destroy();
+ });
+
+ it("within the starred group, oldest-agent-first is preserved", async () => {
+ const starred = new Set<string>(["ws-starred"]);
+ const timers = createFakeTimers();
+ const manager = createConcurrencyManager({
+ now: timers.now,
+ slotTimeoutMs: 5000,
+ watchdogIntervalMs: 1000,
+ defaultPauseMs: 30000,
+ isWorkspaceStarred: (wsId: string) => starred.has(wsId),
+ setTimeout: timers.setTimeout,
+ clearTimeout: timers.clearTimeout,
+ setInterval: timers.setInterval,
+ clearInterval: timers.clearInterval,
+ });
+ manager.setLimit("umans", 1);
+
+ const release0 = await manager.acquire("umans", "holder", "default", 0);
+
+ const results: string[] = [];
+ const acquireAndRecord = (conv: string, wsId: string, promptAt: number) =>
+ manager.acquire("umans", conv, wsId, promptAt).then((r) => {
+ results.push(conv);
+ return r;
+ });
+
+ // Two starred agents: convLate (t=300) queues first, convEarly (t=100) second.
+ const pLate = acquireAndRecord("convLate", "ws-starred", 300);
+ const pEarly = acquireAndRecord("convEarly", "ws-starred", 100);
+
+ await Promise.resolve();
+ await Promise.resolve();
+
+ release0();
+
+ // convEarly (t=100) should win within the starred group (oldest-first).
+ const rEarly = await pEarly;
+ expect(results).toEqual(["convEarly"]);
+ rEarly();
+
+ const rLate = await pLate;
+ expect(results).toEqual(["convEarly", "convLate"]);
+ rLate();
+
+ manager.destroy();
+ });
+
+ it("starring a workspace while agents are queued re-prioritizes them immediately", async () => {
+ const timers = createFakeTimers();
+ const manager = createConcurrencyManager({
+ now: timers.now,
+ slotTimeoutMs: 5000,
+ watchdogIntervalMs: 1000,
+ defaultPauseMs: 30000,
+ setTimeout: timers.setTimeout,
+ clearTimeout: timers.clearTimeout,
+ setInterval: timers.setInterval,
+ clearInterval: timers.clearInterval,
+ });
+ manager.setLimit("umans", 1);
+
+ const release0 = await manager.acquire("umans", "holder", "default", 0);
+
+ // convA (non-starred, t=100) queues first.
+ let resolvedA = false;
+ const pA = manager.acquire("umans", "convA", "ws-normal", 100).then((r) => {
+ resolvedA = true;
+ return r;
+ });
+ // convB (non-starred, t=200) queues second.
+ let resolvedB = false;
+ const pB = manager.acquire("umans", "convB", "ws-to-star", 200).then((r) => {
+ resolvedB = true;
+ return r;
+ });
+
+ await Promise.resolve();
+ await Promise.resolve();
+ expect(resolvedA).toBe(false);
+ expect(resolvedB).toBe(false);
+
+ // Now star convB's workspace AFTER it's queued. notifyWorkspaceStarred
+ // updates the internal cache + re-sorts + tries to grant.
+ manager.notifyWorkspaceStarred("ws-to-star", true);
+
+ // Release the holder — convB (now starred) should jump ahead of convA.
+ release0();
+
+ const rB = await pB;
+ expect(resolvedB).toBe(true);
+ expect(resolvedA).toBe(false);
+ rB();
+
+ // Now convA gets the next slot.
+ const rA = await pA;
+ expect(resolvedA).toBe(true);
+ rA();
+
+ manager.destroy();
+ });
+
+ it("unstar a workspace demotes its queued agents", async () => {
+ const timers = createFakeTimers();
+ const manager = createConcurrencyManager({
+ now: timers.now,
+ slotTimeoutMs: 5000,
+ watchdogIntervalMs: 1000,
+ defaultPauseMs: 30000,
+ setTimeout: timers.setTimeout,
+ clearTimeout: timers.clearTimeout,
+ setInterval: timers.setInterval,
+ clearInterval: timers.clearInterval,
+ });
+ manager.setLimit("umans", 1);
+
+ // Initially star "ws-starred" via the internal cache.
+ manager.notifyWorkspaceStarred("ws-starred", true);
+
+ const release0 = await manager.acquire("umans", "holder", "default", 0);
+
+ // convA (starred, t=200) queues first.
+ const pA = manager.acquire("umans", "convA", "ws-starred", 200);
+ // convB (non-starred, t=100) queues second but is older.
+ const pB = manager.acquire("umans", "convB", "ws-normal", 100);
+
+ await Promise.resolve();
+ await Promise.resolve();
+
+ // Unstar convA's workspace — it should now be behind convB (which is older).
+ manager.notifyWorkspaceStarred("ws-starred", false);
+
+ release0();
+
+ // convB (t=100, now non-starred but oldest) should win.
+ const rB = await pB;
+ expect(rB).toBeDefined();
+ rB();
+
+ const rA = await pA;
+ rA();
+
+ manager.destroy();
+ });
+
+ it("notifyWorkspaceStarred re-sorts queues and tries to grant when capacity is free", async () => {
+ const timers = createFakeTimers();
+ const manager = createConcurrencyManager({
+ now: timers.now,
+ slotTimeoutMs: 5000,
+ watchdogIntervalMs: 1000,
+ defaultPauseMs: 30000,
+ setTimeout: timers.setTimeout,
+ clearTimeout: timers.clearTimeout,
+ setInterval: timers.setInterval,
+ clearInterval: timers.clearInterval,
+ });
+ manager.setLimit("umans", 1);
+
+ // Slot is held. Two agents queued (both non-starred).
+ const release0 = await manager.acquire("umans", "holder", "default", 0);
+ const pA = manager.acquire("umans", "convA", "ws-normal", 100);
+ const pB = manager.acquire("umans", "convB", "ws-to-star", 200);
+
+ await Promise.resolve();
+ await Promise.resolve();
+
+ // Star convB's workspace — notifyWorkspaceStarred re-sorts + tries to
+ // grant. But the slot is still held, so no one is granted yet.
+ manager.notifyWorkspaceStarred("ws-to-star", true);
+
+ // Release the slot — convB (now starred) should get it.
+ release0();
+
+ const rB = await pB;
+ expect(rB).toBeDefined();
+ rB();
+
+ const rA = await pA;
+ rA();
+
+ manager.destroy();
+ });
+
+ it("without isWorkspaceStarred callback, all agents are non-starred (backward compatible)", async () => {
+ const timers = createFakeTimers();
+ const manager = createConcurrencyManager({
+ now: timers.now,
+ slotTimeoutMs: 5000,
+ watchdogIntervalMs: 1000,
+ defaultPauseMs: 30000,
+ setTimeout: timers.setTimeout,
+ clearTimeout: timers.clearTimeout,
+ setInterval: timers.setInterval,
+ clearInterval: timers.clearInterval,
+ });
+ manager.setLimit("umans", 1);
+
+ const release0 = await manager.acquire("umans", "holder", "default", 0);
+
+ const results: string[] = [];
+ const acquireAndRecord = (conv: string, wsId: string, promptAt: number) =>
+ manager.acquire("umans", conv, wsId, promptAt).then((r) => {
+ results.push(conv);
+ return r;
+ });
+
+ // Queue in non-sorted order: B (t=200), A (t=100).
+ const pB = acquireAndRecord("convB", "ws-any", 200);
+ const pA = acquireAndRecord("convA", "ws-any", 100);
+
+ await Promise.resolve();
+ await Promise.resolve();
+
+ release0();
+
+ // Without a callback, oldest-first ordering applies (no starred priority).
+ const rA = await pA;
+ expect(results).toEqual(["convA"]);
+ rA();
+
+ const rB = await pB;
+ expect(results).toEqual(["convA", "convB"]);
+ rB();
+
+ manager.destroy();
+ });
+});
diff --git a/packages/provider-concurrency/src/concurrency-manager.ts b/packages/provider-concurrency/src/concurrency-manager.ts
new file mode 100644
index 0000000..986493d
--- /dev/null
+++ b/packages/provider-concurrency/src/concurrency-manager.ts
@@ -0,0 +1,851 @@
+/**
+ * In-memory per-provider concurrency limiter.
+ *
+ * Tracks and limits how many concurrent API requests (token-generating
+ * requests) are in flight per provider. When the limit is reached, additional
+ * requests queue and are granted slots based on oldest-agent-first priority
+ * (the agent whose current prompt started the longest ago wins the next slot).
+ *
+ * A watchdog reclaims slots held beyond a timeout (deadlock / stuck-agent
+ * recovery). 429 backoff pauses a provider's queue for a configurable duration
+ * AND adaptively reduces the effective limit by 1 (one-way, persisted) so the
+ * resumed queue runs with headroom instead of re-overshooting.
+ *
+ * ── Usage gate (anti-overshoot) ──
+ * When a `fetchUsage` callback is injected, before admitting a QUEUED agent the
+ * manager polls the provider's upstream `concurrent_sessions` count and grants
+ * only when it is below the configured limit. This composes with the release
+ * cooldown: release → cooldown delay → usage-gate poll → grant (only if upstream
+ * has room). A waiter is re-checked on two triggers (either one): another agent
+ * releases a slot (immediate re-poll, restarting the 1s countdown) or a 1s
+ * fallback timer elapses (in case the upstream count drops on its own). Each
+ * successful poll admits at most ONE queued waiter (each admission pushes the
+ * upstream count back toward the limit), so additional waiters are admitted on
+ * subsequent repolls. When `fetchUsage` is absent or returns `undefined`, the
+ * gate is skipped and the manager falls back to cooldown-only recycling.
+ *
+ * This module is the PURE decision logic. It takes an injected clock (`now`),
+ * injected timers (`setTimeout`/`clearTimeout`/`setInterval`/`clearInterval`),
+ * and an injected usage-poll effect (`fetchUsage`) so it is fully testable with
+ * deterministic fake time + a fake fetcher. The extension layer wires real
+ * timers + the host's provider registry.
+ */
+
+import type { ProviderUsage } from "@dispatch/kernel";
+
+// ─── Types ───────────────────────────────────────────────────────────────────
+
+/** Status snapshot for a single provider's concurrency state. */
+export interface ProviderConcurrencyStatus {
+ readonly providerId: string;
+ /** Configured concurrency limit. Always present (status is only returned for providers with a limit). */
+ readonly limit: number;
+ /** Currently in-flight (held) slots. */
+ readonly inFlight: number;
+ /** Agents waiting in the queue for a slot. */
+ readonly queued: number;
+ /** Whether the queue is paused (429 backoff). */
+ readonly paused: boolean;
+ /** When the pause expires (epoch-ms). Present only when paused. */
+ readonly pausedUntil?: number;
+ /**
+ * Per-slot release cooldown (ms) — how long a recycled slot is held before the
+ * next waiter is admitted. Covers the upstream provider's accounting lag.
+ * Configurable + persisted per provider.
+ */
+ readonly cooldownMs: number;
+ /**
+ * Whether the limit was auto-reduced by a 429 (adaptive headroom). The user
+ * restores the limit manually (PUT /concurrency/limits/:providerId) which
+ * clears this flag. The frontend renders a visible notice when `true`.
+ */
+ readonly autoReduced: boolean;
+ /** The original limit before auto-reduction (present only when autoReduced). */
+ readonly autoReducedFrom?: number;
+ /**
+ * A human-readable notice string for the frontend to render as a banner when
+ * the limit was auto-reduced. Present only when `autoReduced` is true.
+ */
+ readonly notice?: string;
+}
+
+/**
+ * The limiter surface a consumer (session-orchestrator) needs: acquire a
+ * slot before a provider stream starts, release it when the stream completes,
+ * and report rate-limit (429) events so the manager can back off.
+ */
+export interface ConcurrencyLimiter {
+ /**
+ * Acquire a concurrency slot for `providerId`. Resolves immediately when a
+ * slot is available; otherwise blocks (queued by starred-workspace-first,
+ * then oldest-agent-first) until one frees up. The returned function MUST be
+ * called when the response stream completes (in a `finally` block). For
+ * providers with no configured limit, resolves instantly with a no-op
+ * release.
+ *
+ * **Priority:** agents from **starred workspaces** are always admitted before
+ * agents from non-starred workspaces (regardless of `promptStartedAt`).
+ * Within each group (starred vs non-starred), oldest-agent-first ordering is
+ * preserved. The starred status is looked up via the injected
+ * `isWorkspaceStarred` callback at sort time, so starring a workspace while
+ * agents are queued takes effect on the next sort (new acquire or slot
+ * release).
+ *
+ * If `onQueued` is provided and the request cannot be granted immediately
+ * (at limit or paused), it is called synchronously BEFORE the Promise is
+ * created. This lets the caller emit a "queued" status signal. If the slot
+ * is granted immediately, `onQueued` is NOT called.
+ *
+ * @param providerId The provider to limit (e.g. "umans", "openai-compat").
+ * @param conversationId The agent requesting the slot.
+ * @param workspaceId The workspace the agent belongs to (for starred
+ * priority scheduling). Defaults to `"default"`.
+ * @param promptStartedAt When the agent's current prompt (turn) started
+ * (epoch-ms). Used for oldest-agent-first scheduling
+ * within each starred group.
+ * @param onQueued Called synchronously when the request is enqueued
+ * (not granted immediately). Optional.
+ */
+ acquire(
+ providerId: string,
+ conversationId: string,
+ workspaceId: string,
+ promptStartedAt: number,
+ onQueued?: () => void,
+ ): Promise<() => void>;
+
+ /**
+ * Report a 429 from a provider. Pauses the queue for that provider for
+ * `retryAfterMs` (or a default duration when omitted), AND reduces the
+ * provider's effective limit by 1 (one-way, down to a minimum of 1) so the
+ * resumed queue runs with headroom. Queued and in-flight requests are
+ * otherwise unaffected; new `acquire` calls block until the pause expires.
+ */
+ reportRateLimit(providerId: string, retryAfterMs?: number): void;
+}
+
+/**
+ * The full service surface (limiter + config + status) for HTTP routes.
+ */
+export interface ConcurrencyService extends ConcurrencyLimiter {
+ /** Set the concurrency limit for a provider (MANUAL — clears the auto-reduce notice). Creates the state if new. */
+ setLimit(providerId: string, limit: number): void;
+ /**
+ * Restore a persisted limit on startup WITHOUT clearing the auto-reduce
+ * notice (Bug 3). Unlike {@link setLimit} (a manual user action that signals
+ * "the user took control"), this seeds state from disk: it applies the limit
+ * and, when `autoReducedFrom` is provided, re-marks the state as auto-reduced
+ * so the frontend banner survives a restart. Used by the extension's
+ * `loadLimits`/`loadAutoReduce` on activate.
+ */
+ restoreLimit(providerId: string, limit: number, autoReducedFrom?: number): void;
+ /** Get the configured limit, or `undefined` when none. */
+ getLimit(providerId: string): number | undefined;
+ /** Remove the limit for a provider (makes it unlimited). */
+ removeLimit(providerId: string): void;
+ /** All configured limits as `{ providerId, limit }` entries. */
+ getLimits(): readonly { providerId: string; limit: number }[];
+ /**
+ * Set the release cooldown (ms) for a provider. Applied to subsequently
+ * recycled slots; in-flight cooldown timers keep their original duration.
+ * Creates the state if new (with no limit — unlimited but cooldown-gated).
+ */
+ setCooldown(providerId: string, cooldownMs: number): void;
+ /** Get the configured cooldown (ms), or `undefined` when none was set. */
+ getCooldown(providerId: string): number | undefined;
+ /** All configured cooldowns as `{ providerId, cooldownMs }` entries. */
+ getCooldowns(): readonly { providerId: string; cooldownMs: number }[];
+ /** Status for one provider, or `undefined` when no limit is configured. */
+ getStatus(providerId: string): ProviderConcurrencyStatus | undefined;
+ /** Status for every provider with a configured limit. */
+ getStatusAll(): readonly ProviderConcurrencyStatus[];
+ /**
+ * Notify the limiter that a workspace's starred state changed. Updates the
+ * in-memory starred cache so subsequent queue sorts re-evaluate priority
+ * (a newly-starred workspace's already-queued agents jump ahead). Called by
+ * the transport layer after persisting the starred toggle.
+ */
+ notifyWorkspaceStarred(workspaceId: string, starred: boolean): void;
+ /** Stop the watchdog + clear all timers. */
+ destroy(): void;
+}
+
+// ─── Internal state ───────────────────────────────────────────────────────────
+
+interface Slot {
+ readonly conversationId: string;
+ readonly acquiredAt: number;
+ /** Idempotent release — safe to call from the holder or the watchdog. */
+ readonly releaseFn: () => void;
+}
+
+interface QueuedWaiter {
+ readonly conversationId: string;
+ readonly workspaceId: string;
+ readonly promptStartedAt: number;
+ readonly resolve: (release: () => void) => void;
+}
+
+interface ProviderState {
+ limit: number;
+ inFlight: number;
+ slots: Map<number, Slot>;
+ queue: QueuedWaiter[];
+ paused: boolean;
+ pausedUntil: number | undefined;
+ pauseTimer: ReturnType<typeof setTimeout> | undefined;
+ /** Per-provider release cooldown (ms). Defaults to the manager opt; settable at runtime. */
+ cooldownMs: number;
+ // ── Adaptive headroom ──
+ autoReduced: boolean;
+ autoReducedFrom: number | undefined;
+ notice: string | undefined;
+ // ── Usage-gate state ──
+ /** A usage poll is in flight for this provider (prevents overlapping polls). */
+ gatePolling: boolean;
+ /** Another repoll trigger fired while a poll was in flight → re-poll on completion. */
+ gateRepollRequested: boolean;
+ /** The 1s fallback repoll timer (re-checked periodically even without releases). */
+ gateRepollTimer: ReturnType<typeof setTimeout> | undefined;
+}
+
+export interface ConcurrencyManagerOpts {
+ /** Monotonic-ish clock (epoch-ms). */
+ readonly now: () => number;
+ /** Max time a slot may be held before the watchdog reclaims it (ms). */
+ readonly slotTimeoutMs: number;
+ /** How often the watchdog sweeps (ms). */
+ readonly watchdogIntervalMs: number;
+ /** Default pause duration when a 429 arrives without Retry-After (ms). */
+ readonly defaultPauseMs: number;
+ /**
+ * Default delay after a slot is released before the slot is recycled (ms).
+ * During this window `inFlight` stays incremented — a new `acquire` sees the
+ * slot as still held and queues. This covers the upstream provider's
+ * accounting lag: the provider's `concurrent_sessions` counter may not
+ * decrement the instant our stream completes, so re-admitting immediately
+ * risks an N+1 overshoot. 0 = instant re-admission (no cooldown). Default: 0.
+ * Per-provider overrides via `setCooldown`.
+ */
+ readonly releaseCooldownMs?: number;
+ /**
+ * Injected usage-poll effect. When present, before admitting a QUEUED agent
+ * the manager calls this and grants only when `concurrentSessions` is below
+ * the configured limit (usage gate). When absent, the manager falls back to
+ * cooldown-only slot recycling. Injected (like `now`/`setTimeout`) so the
+ * manager stays unit-testable with a fake fetcher; never hardcodes `fetch`.
+ */
+ readonly fetchUsage?: (providerId: string) => Promise<ProviderUsage | undefined>;
+ /** Injected timers (default: global). Override in tests for deterministic time. */
+ readonly setTimeout?: typeof setTimeout;
+ readonly clearTimeout?: typeof clearTimeout;
+ readonly setInterval?: typeof setInterval;
+ readonly clearInterval?: typeof clearInterval;
+ /** Optional logger for watchdog + pause + auto-reduce events. */
+ readonly onWatchdogReclaim?: (providerId: string, conversationId: string, heldMs: number) => void;
+ readonly onPause?: (providerId: string, durationMs: number) => void;
+ /** Fired when a 429 adaptively reduces a provider's limit (for persistence + logging). */
+ readonly onLimitReduced?: (providerId: string, newLimit: number, oldLimit: number) => void;
+ /**
+ * Fired when the injected `fetchUsage` throws (network/parse failure beyond the
+ * graceful-undefined path). The manager treats a thrown poll as "no usage info"
+ * (cooldown-only fallback) — this callback is for WARN-level logging only. The
+ * poll never becomes an unhandled rejection.
+ */
+ readonly onUsagePollError?: (providerId: string, err: unknown) => void;
+ /**
+ * Injected callback: returns whether a workspace is starred (for priority
+ * scheduling). When provided, agents from starred workspaces jump ahead of
+ * non-starred agents in the queue. When omitted (or returns `false`), all
+ * agents are treated as non-starred (backward-compatible). This is an I/O
+ * effect injected so the manager stays pure + unit-testable with a fake.
+ */
+ readonly isWorkspaceStarred?: (workspaceId: string) => boolean;
+}
+
+/** Min interval between usage-gate fallback repolls (ms). The release trigger is immediate. */
+const USAGE_REPOLL_INTERVAL_MS = 1000;
+/** Minimum the limit may be auto-reduced to (never 0). */
+const MIN_LIMIT = 1;
+
+function noopRelease(): void {
+ // No limit configured → nothing to release.
+}
+
+export function createConcurrencyManager(opts: ConcurrencyManagerOpts): ConcurrencyService {
+ const now = opts.now;
+ const slotTimeoutMs = opts.slotTimeoutMs;
+ const defaultPauseMs = opts.defaultPauseMs;
+ const defaultCooldownMs = opts.releaseCooldownMs ?? 0;
+ const fetchUsage = opts.fetchUsage;
+ const setTimeout = opts.setTimeout ?? globalThis.setTimeout.bind(globalThis);
+ const clearTimeout = opts.clearTimeout ?? globalThis.clearTimeout.bind(globalThis);
+ const setInterval = opts.setInterval ?? globalThis.setInterval.bind(globalThis);
+ const clearInterval = opts.clearInterval ?? globalThis.clearInterval.bind(globalThis);
+
+ // In-memory cache of starred workspace IDs. Populated by the extension on
+ // activation (from the conversation store) + updated via
+ // `notifyWorkspaceStarred`. The `isWorkspaceStarred` callback reads this
+ // synchronously so the queue sort comparator (sync) can re-evaluate priority
+ // on every sort — a newly-starred workspace's already-queued agents jump
+ // ahead on the next sort (new acquire or slot release).
+ const starredWorkspaces = new Set<string>();
+ const isWorkspaceStarred =
+ opts.isWorkspaceStarred ?? ((wsId: string) => starredWorkspaces.has(wsId));
+
+ const states = new Map<string, ProviderState>();
+ const cooldownOverrides = new Map<string, number>();
+ const cooldownTimers = new Set<ReturnType<typeof setTimeout>>();
+ let slotIdCounter = 0;
+
+ function makeState(limit: number, cooldownMs: number): ProviderState {
+ return {
+ limit,
+ inFlight: 0,
+ slots: new Map(),
+ queue: [],
+ paused: false,
+ pausedUntil: undefined,
+ pauseTimer: undefined,
+ cooldownMs,
+ autoReduced: false,
+ autoReducedFrom: undefined,
+ notice: undefined,
+ gatePolling: false,
+ gateRepollRequested: false,
+ gateRepollTimer: undefined,
+ };
+ }
+
+ /** Seed the cooldown for new state from any pending override (else the default). */
+ function seedCooldown(providerId: string): number {
+ return cooldownOverrides.get(providerId) ?? defaultCooldownMs;
+ }
+
+ // ── Slot granting ──────────────────────────────────────────────────────────
+
+ function grantSlot(state: ProviderState, providerId: string, conversationId: string): () => void {
+ const id = slotIdCounter++;
+ let released = false;
+ const releaseFn = () => {
+ if (released) return;
+ released = true;
+ state.slots.delete(id);
+
+ // Recycle the slot: free its inFlight count + attempt to grant the next
+ // waiter. With a release cooldown > 0, defer this by the cooldown duration
+ // so the upstream provider has time to decrement its concurrent_sessions
+ // counter — preventing an N+1 overshoot from accounting lag. During the
+ // cooldown, inFlight stays incremented, so new acquires queue.
+ const recycle = () => {
+ if (fetchUsage === undefined || state.queue.length === 0) {
+ // No usage gate, OR no one waiting (the lag window is irrelevant when
+ // there is no waiter to admit) → free the slot immediately. With no
+ // gate, also drain the queue (grant all that fit).
+ state.inFlight--;
+ if (fetchUsage === undefined) grantLoop(state, providerId);
+ return;
+ }
+ // Usage gate configured + a waiter exists → hold inFlight inflated
+ // DURING the poll window (gatePolling is set synchronously inside
+ // pollAndGrant, the inFlight decrement is deferred until the poll
+ // resolves). This closes the overshoot gap: a concurrent acquire arriving
+ // between the cooldown firing and the poll resolving sees the slot as
+ // still occupied (inFlight >= limit) and queues instead of fast-pathing.
+ // pollAndGrant(decrementOnPoll=true) decrements inFlight after observing
+ // the post-release upstream state, then admits one waiter if there is room.
+ void pollAndGrant(providerId, state, true);
+ };
+ if (state.cooldownMs > 0) {
+ const timer = setTimeout(() => {
+ cooldownTimers.delete(timer);
+ recycle();
+ }, state.cooldownMs);
+ cooldownTimers.add(timer);
+ } else {
+ recycle();
+ }
+ };
+ state.slots.set(id, {
+ conversationId,
+ acquiredAt: now(),
+ releaseFn,
+ });
+ state.inFlight++;
+ return releaseFn;
+ }
+
+ /**
+ * Priority comparator for queued waiters: starred-workspace agents first,
+ * then oldest-agent-first (ascending `promptStartedAt`) within each group.
+ * Called at sort time (both on insert and before granting) so a workspace
+ * starred AFTER an agent queued is re-evaluated on the next sort.
+ */
+ function compareWaiters(a: QueuedWaiter, b: QueuedWaiter): number {
+ const aStarred = isWorkspaceStarred(a.workspaceId);
+ const bStarred = isWorkspaceStarred(b.workspaceId);
+ if (aStarred !== bStarred) return aStarred ? -1 : 1; // starred first
+ return a.promptStartedAt - b.promptStartedAt; // oldest first within group
+ }
+
+ /**
+ * Grant queued waiters WITHOUT the usage gate (the fast path used when no
+ * `fetchUsage` is configured, or as the cooldown-only fallback when a poll
+ * returns no usage info). Grants while there is internal room
+ * (`inFlight < limit`). Synchronous. Re-sorts with {@link compareWaiters}
+ * (starred-workspace-first, then oldest-agent-first) before granting so a
+ * workspace starred AFTER an agent queued is re-evaluated.
+ */
+ function grantLoop(state: ProviderState, providerId: string): void {
+ // Re-sort before granting: a workspace may have been starred/unstarred
+ // since the waiters were enqueued, so priority may have changed.
+ state.queue.sort(compareWaiters);
+ while (state.queue.length > 0 && state.inFlight < state.limit) {
+ const waiter = state.queue[0];
+ if (waiter === undefined) break;
+ state.queue.shift();
+ const releaseFn = grantSlot(state, providerId, waiter.conversationId);
+ waiter.resolve(releaseFn);
+ }
+ // If the queue drained, no need to keep the usage-gate fallback timer armed.
+ if (state.queue.length === 0 && state.gateRepollTimer !== undefined) {
+ clearTimeout(state.gateRepollTimer);
+ state.gateRepollTimer = undefined;
+ }
+ }
+
+ /**
+ * Admit exactly ONE queued waiter (the front of the queue), if there is
+ * internal room. Used by the usage-gated path so each admission is confirmed
+ * by a FRESH upstream poll — admitting multiple from a single (possibly stale)
+ * poll risks an N+1 overshoot when the upstream count lags. Additional waiters
+ * are admitted on subsequent repolls.
+ */
+ function grantOne(state: ProviderState, providerId: string): void {
+ if (state.queue.length === 0) return;
+ if (state.inFlight >= state.limit) return;
+ // Re-sort before picking the front: starred-workspace agents must be
+ // admitted first, even on the usage-gated path (a workspace may have been
+ // starred since the waiters were enqueued).
+ state.queue.sort(compareWaiters);
+ const waiter = state.queue[0];
+ if (waiter === undefined) return;
+ state.queue.shift();
+ const releaseFn = grantSlot(state, providerId, waiter.conversationId);
+ waiter.resolve(releaseFn);
+ // If the queue drained, disarm the fallback timer.
+ if (state.queue.length === 0 && state.gateRepollTimer !== undefined) {
+ clearTimeout(state.gateRepollTimer);
+ state.gateRepollTimer = undefined;
+ }
+ }
+
+ /**
+ * Invoke the injected `fetchUsage`, treating ANY thrown error as "no usage
+ * info available" (cooldown-only fallback) — so a throwing `getUsage()` never
+ * becomes an unhandled rejection. The `onUsagePollError` opt is fired for
+ * WARN-level logging. Returns `undefined` on throw (Bug 2 fix).
+ */
+ async function safeFetchUsage(providerId: string): Promise<ProviderUsage | undefined> {
+ if (fetchUsage === undefined) return undefined;
+ try {
+ return await fetchUsage(providerId);
+ } catch (err) {
+ opts.onUsagePollError?.(providerId, err);
+ return undefined;
+ }
+ }
+
+ /**
+ * Drain the queue, gated on the upstream usage poll when `fetchUsage` is
+ * configured. Called from setLimit, pause-expiry, and the repoll timer (NOT
+ * from release — that goes through {@link recycleGated}, which holds inFlight
+ * inflated during the poll). Async because the usage poll is an injected I/O
+ * effect; callers fire-and-forget the returned promise.
+ *
+ * The fast-path immediate grant in `acquire` (when `inFlight < limit`) is
+ * disabled while `gatePolling` is true — `acquire` queues instead, so a
+ * concurrent caller cannot sneak through the accounting-lag / poll window
+ * (anti-overshoot). When no poll is in flight the fast-path is safe: the
+ * cooldown keeps `inFlight` inflated during the lag window, and a recycle
+ * sets `gatePolling` synchronously before decrementing.
+ *
+ * Each successful poll admits at most ONE queued waiter (each admission pushes
+ * the upstream count back toward the limit); additional waiters are admitted
+ * on subsequent repolls (release triggers an immediate re-poll; the 1s
+ * fallback timer covers an upstream count that drops on its own).
+ */
+ async function tryGrantNext(providerId: string): Promise<void> {
+ const state = states.get(providerId);
+ if (state === undefined) return;
+ if (state.paused) return;
+ if (state.queue.length === 0) return;
+ if (state.inFlight >= state.limit) return; // no internal room
+
+ // No usage gate → immediate grant loop (original behavior).
+ if (fetchUsage === undefined) {
+ grantLoop(state, providerId);
+ return;
+ }
+
+ // Avoid overlapping polls for this provider. A poll is already in flight;
+ // mark that another trigger fired so it re-polls on completion.
+ if (state.gatePolling) {
+ state.gateRepollRequested = true;
+ return;
+ }
+
+ await pollAndGrant(providerId, state);
+ }
+
+ /**
+ * Shared poll-then-admit. `decrementOnPoll` is true for the recycle path
+ * (the released slot's inFlight decrement is deferred until the poll resolves,
+ * holding inFlight inflated so concurrent acquires queue — anti-overshoot) and
+ * false for the drain path (setLimit/pause-expiry/repoll — no slot to account).
+ * Admits at most ONE waiter on a successful poll.
+ */
+ async function pollAndGrant(
+ providerId: string,
+ state: ProviderState,
+ decrementOnPoll = false,
+ ): Promise<void> {
+ state.gatePolling = true;
+ try {
+ const snapshot = await safeFetchUsage(providerId);
+
+ // For the recycle path, the released slot is now truly freed (the poll
+ // has observed the post-release upstream state).
+ if (decrementOnPoll) {
+ state.inFlight--;
+ }
+
+ // Conditions may have changed during the async poll — re-check.
+ if (state.paused) return;
+ if (state.queue.length === 0) return;
+
+ if (snapshot === undefined) {
+ // No usage info available → fall back to cooldown-only (grant one).
+ grantOne(state, providerId);
+ return;
+ }
+
+ if (snapshot.concurrentSessions < state.limit) {
+ // Upstream has room — admit exactly ONE queued waiter.
+ grantOne(state, providerId);
+ }
+ // else: upstream at/over limit → keep queued; repoll timer handles retry.
+ } finally {
+ state.gatePolling = false;
+ // (Re)arm the 1s fallback timer while waiters remain queued, so an
+ // upstream count that drops on its own is still detected.
+ armGateRepoll(providerId, state);
+ if (state.gateRepollRequested) {
+ state.gateRepollRequested = false;
+ // A release (or other trigger) fired during the poll → re-poll now.
+ void tryGrantNext(providerId);
+ }
+ }
+ }
+
+ function armGateRepoll(providerId: string, state: ProviderState): void {
+ // Only arm while there are queued waiters (otherwise no work to re-check).
+ if (state.queue.length === 0) {
+ if (state.gateRepollTimer !== undefined) {
+ clearTimeout(state.gateRepollTimer);
+ state.gateRepollTimer = undefined;
+ }
+ return;
+ }
+ if (state.gateRepollTimer !== undefined) {
+ clearTimeout(state.gateRepollTimer);
+ }
+ state.gateRepollTimer = setTimeout(() => {
+ state.gateRepollTimer = undefined;
+ void tryGrantNext(providerId);
+ }, USAGE_REPOLL_INTERVAL_MS);
+ }
+
+ // ── Watchdog ──────────────────────────────────────────────────────────────────
+
+ function sweep(): void {
+ const currentNow = now();
+ for (const [providerId, state] of states) {
+ for (const [, slot] of state.slots) {
+ const heldMs = currentNow - slot.acquiredAt;
+ if (heldMs > slotTimeoutMs) {
+ opts.onWatchdogReclaim?.(providerId, slot.conversationId, heldMs);
+ slot.releaseFn();
+ }
+ }
+ }
+ }
+
+ const watchdogTimer = setInterval(sweep, opts.watchdogIntervalMs);
+
+ // ── Adaptive headroom ──────────────────────────────────────────────────────
+
+ function clearAutoReduce(state: ProviderState): void {
+ state.autoReduced = false;
+ state.autoReducedFrom = undefined;
+ state.notice = undefined;
+ }
+
+ // ── Public API ─────────────────────────────────────────────────────────────
+
+ const manager: ConcurrencyService = {
+ acquire(providerId, conversationId, workspaceId, promptStartedAt, onQueued) {
+ const state = states.get(providerId);
+ if (state === undefined) {
+ // No limit configured → unlimited.
+ return Promise.resolve(noopRelease);
+ }
+
+ if (!state.paused && state.inFlight < state.limit) {
+ // Usage-gate anti-overshoot: while a recycle-poll is in flight, the
+ // inFlight count is momentarily unreliable (a released slot's decrement
+ // is deferred until the poll resolves — see pollAndGrant). A concurrent
+ // caller that fast-pathed now could overshoot the upstream limit before
+ // the poll confirms room. So route it through the queue instead; the
+ // in-flight poll will re-check (gateRepollRequested) and admit it once
+ // upstream confirms room. When no poll is in flight the fast-path is
+ // safe (the cooldown keeps inFlight inflated through the lag window).
+ if (fetchUsage !== undefined && state.gatePolling) {
+ // falls through to the queue path below
+ } else {
+ return Promise.resolve(grantSlot(state, providerId, conversationId));
+ }
+ }
+
+ // Cannot grant immediately — the request will be queued.
+ // Notify the caller BEFORE creating the Promise so they can emit a
+ // "queued" status signal while we're still synchronous.
+ onQueued?.();
+
+ // Queue (starred-workspace-first, then oldest-agent-first).
+ return new Promise<() => void>((resolve) => {
+ state.queue.push({ conversationId, workspaceId, promptStartedAt, resolve });
+ // Keep sorted by priority (starred first, then oldest-agent-first).
+ // The queue is typically tiny (<20), so a simple sort is fine.
+ state.queue.sort(compareWaiters);
+ // If the usage gate is active, ensure the fallback repoll timer is
+ // armed (a release may not come for a while; the 1s timer covers an
+ // upstream count that drops on its own).
+ if (fetchUsage !== undefined) {
+ armGateRepoll(providerId, state);
+ }
+ });
+ },
+
+ reportRateLimit(providerId, retryAfterMs) {
+ const state = states.get(providerId);
+ if (state === undefined) return;
+
+ const pauseDuration = retryAfterMs ?? defaultPauseMs;
+ state.paused = true;
+ state.pausedUntil = now() + pauseDuration;
+
+ if (state.pauseTimer !== undefined) {
+ clearTimeout(state.pauseTimer);
+ }
+ opts.onPause?.(providerId, pauseDuration);
+
+ // Adaptive headroom: reduce the effective limit by 1 (one-way, min 1) so
+ // the resumed queue runs with headroom instead of re-overshooting. The
+ // reduction is persisted + surfaced (via onLimitReduced + status).
+ if (state.limit > MIN_LIMIT) {
+ const oldLimit = state.limit;
+ state.limit = Math.max(MIN_LIMIT, state.limit - 1);
+ if (!state.autoReduced) {
+ state.autoReduced = true;
+ state.autoReducedFrom = oldLimit;
+ }
+ state.notice =
+ `Concurrency limit auto-reduced to ${state.limit} after a 429 — ` +
+ "restore manually when ready.";
+ opts.onLimitReduced?.(providerId, state.limit, oldLimit);
+ }
+
+ state.pauseTimer = setTimeout(() => {
+ state.paused = false;
+ state.pausedUntil = undefined;
+ state.pauseTimer = undefined;
+ void tryGrantNext(providerId);
+ }, pauseDuration);
+ },
+
+ setLimit(providerId, limit) {
+ let state = states.get(providerId);
+ if (state === undefined) {
+ state = makeState(limit, seedCooldown(providerId));
+ states.set(providerId, state);
+ } else {
+ state.limit = limit;
+ // A MANUAL limit set clears the auto-reduce notice (the user took control).
+ clearAutoReduce(state);
+ }
+ // A higher limit may let queued requests through.
+ void tryGrantNext(providerId);
+ },
+
+ restoreLimit(providerId, limit, autoReducedFrom) {
+ // Startup restoration (Bug 3): seed state from disk WITHOUT the manual
+ // "user took control" semantics, so a persisted auto-reduced limit keeps
+ // its notice/banner across a restart. When autoReducedFrom is provided,
+ // re-mark the state as auto-reduced (rebuild the notice).
+ let state = states.get(providerId);
+ if (state === undefined) {
+ state = makeState(limit, seedCooldown(providerId));
+ states.set(providerId, state);
+ } else {
+ state.limit = limit;
+ }
+ if (autoReducedFrom !== undefined && autoReducedFrom > limit) {
+ state.autoReduced = true;
+ state.autoReducedFrom = autoReducedFrom;
+ state.notice =
+ `Concurrency limit auto-reduced to ${limit} after a 429 — ` +
+ "restore manually when ready.";
+ }
+ // A higher limit may let queued requests through.
+ void tryGrantNext(providerId);
+ },
+
+ getLimit(providerId) {
+ return states.get(providerId)?.limit;
+ },
+
+ removeLimit(providerId) {
+ const state = states.get(providerId);
+ if (state === undefined) return;
+
+ // Clear pause.
+ state.paused = false;
+ state.pausedUntil = undefined;
+ if (state.pauseTimer !== undefined) {
+ clearTimeout(state.pauseTimer);
+ state.pauseTimer = undefined;
+ }
+ // Clear usage-gate fallback timer.
+ if (state.gateRepollTimer !== undefined) {
+ clearTimeout(state.gateRepollTimer);
+ state.gateRepollTimer = undefined;
+ }
+ clearAutoReduce(state);
+
+ // Grant all queued requests (they become unlimited now).
+ while (state.queue.length > 0) {
+ const waiter = state.queue[0];
+ if (waiter === undefined) break;
+ state.queue.shift();
+ const releaseFn = grantSlot(state, providerId, waiter.conversationId);
+ waiter.resolve(releaseFn);
+ }
+
+ // Remove the state. In-flight slots' release functions still work —
+ // they close over `state` and call `tryGrantNext` which finds no state
+ // and returns early. The watchdog won't sweep removed states.
+ states.delete(providerId);
+ },
+
+ getLimits() {
+ return [...states.entries()].map(([providerId, s]) => ({
+ providerId,
+ limit: s.limit,
+ }));
+ },
+
+ setCooldown(providerId, cooldownMs) {
+ // A cooldown is only meaningful WITH a limit (it gates slot recycling,
+ // which only happens under a limit). But we store the override regardless
+ // so it applies the moment a limit IS set — and so a persisted cooldown
+ // restored before a limit does NOT impose a limit (setCooldown never
+ // creates a state). If a state already exists, update it live.
+ cooldownOverrides.set(providerId, cooldownMs);
+ const state = states.get(providerId);
+ if (state !== undefined) {
+ state.cooldownMs = cooldownMs;
+ }
+ },
+
+ getCooldown(providerId) {
+ const state = states.get(providerId);
+ if (state !== undefined) return state.cooldownMs;
+ return cooldownOverrides.get(providerId);
+ },
+
+ getCooldowns() {
+ // Merge: states (cooldown from state.cooldownMs) + pending overrides with no state.
+ const seen = new Set<string>();
+ const out: { providerId: string; cooldownMs: number }[] = [];
+ for (const [providerId, s] of states) {
+ seen.add(providerId);
+ out.push({ providerId, cooldownMs: s.cooldownMs });
+ }
+ for (const [providerId, cooldownMs] of cooldownOverrides) {
+ if (!seen.has(providerId)) {
+ out.push({ providerId, cooldownMs });
+ }
+ }
+ return out;
+ },
+
+ getStatus(providerId) {
+ const state = states.get(providerId);
+ if (state === undefined) return undefined;
+ return {
+ providerId,
+ limit: state.limit,
+ inFlight: state.inFlight,
+ queued: state.queue.length,
+ paused: state.paused,
+ cooldownMs: state.cooldownMs,
+ autoReduced: state.autoReduced,
+ ...(state.pausedUntil !== undefined ? { pausedUntil: state.pausedUntil } : {}),
+ ...(state.autoReducedFrom !== undefined ? { autoReducedFrom: state.autoReducedFrom } : {}),
+ ...(state.notice !== undefined ? { notice: state.notice } : {}),
+ };
+ },
+
+ getStatusAll() {
+ return [...states.keys()]
+ .map((providerId) => manager.getStatus(providerId))
+ .filter((s): s is ProviderConcurrencyStatus => s !== undefined);
+ },
+
+ notifyWorkspaceStarred(workspaceId, starred) {
+ if (starred) {
+ starredWorkspaces.add(workspaceId);
+ } else {
+ starredWorkspaces.delete(workspaceId);
+ }
+ // Re-sort all queues so a newly-starred workspace's already-queued
+ // agents jump ahead immediately (no need to wait for the next acquire).
+ for (const [providerId, state] of states) {
+ if (state.queue.length > 0) {
+ state.queue.sort(compareWaiters);
+ tryGrantNext(providerId);
+ }
+ }
+ },
+
+ destroy() {
+ clearInterval(watchdogTimer);
+ for (const timer of cooldownTimers) {
+ clearTimeout(timer);
+ }
+ cooldownTimers.clear();
+ for (const state of states.values()) {
+ if (state.pauseTimer !== undefined) {
+ clearTimeout(state.pauseTimer);
+ }
+ if (state.gateRepollTimer !== undefined) {
+ clearTimeout(state.gateRepollTimer);
+ }
+ }
+ states.clear();
+ },
+ };
+
+ return manager;
+}
diff --git a/packages/provider-concurrency/src/extension.ts b/packages/provider-concurrency/src/extension.ts
new file mode 100644
index 0000000..48c019b
--- /dev/null
+++ b/packages/provider-concurrency/src/extension.ts
@@ -0,0 +1,319 @@
+import { conversationStoreHandle } from "@dispatch/conversation-store";
+import type { Extension, HostAPI, Logger, Manifest, StorageNamespace } from "@dispatch/kernel";
+import type { ConcurrencyManagerOpts, ConcurrencyService } from "./concurrency-manager.js";
+import { createConcurrencyManager } from "./concurrency-manager.js";
+import { concurrencyServiceHandle } from "./service.js";
+
+export const manifest: Manifest = {
+ id: "provider-concurrency",
+ name: "Provider Concurrency Limits",
+ version: "0.0.0",
+ apiVersion: "^0.1.0",
+ trust: "bundled",
+ activation: "eager",
+ capabilities: { db: true },
+ dependsOn: ["conversation-store"],
+ contributes: { services: ["provider-concurrency/service"] },
+};
+
+/**
+ * Default tuning constants.
+ *
+ * - `SLOT_TIMEOUT_MS` (5 min): a slot held longer than this is force-reclaimed
+ * by the watchdog (deadlock / stuck-agent recovery). Generation streams
+ * rarely exceed 2–3 minutes; 5 min is a generous safety margin.
+ * - `WATCHDOG_INTERVAL_MS` (30s): how often the watchdog sweeps for stale slots.
+ * - `DEFAULT_PAUSE_MS` (30s): default 429 backoff when no Retry-After is given.
+ * Umans docs note each concurrency 429 deprioritizes the account for ~30 min,
+ * but a 30s queue pause prevents immediate re-overshoot while still allowing
+ * recovery. Combined with adaptive headroom (limit reduced by 1) + the usage
+ * gate, the resumed queue no longer re-overshoots — so the pause is kept
+ * (gives upstream a breather) rather than dropped.
+ * - `RELEASE_COOLDOWN_MS` (350ms): after a slot is released, hold it for this
+ * duration before recycling it to the next waiter. Covers the upstream
+ * provider's accounting lag — the provider's concurrent_sessions counter may
+ * not decrement the instant our stream completes, so re-admitting immediately
+ * risks an N+1 overshoot that triggers a 429. Raised from 200ms to 350ms
+ * (Umans's accounting lag exceeded the 200ms cooldown, causing overshoot at 4
+ * connections). Configurable + persisted per provider (PUT
+ * /concurrency/cooldown/:providerId).
+ */
+const SLOT_TIMEOUT_MS = 5 * 60 * 1000;
+const WATCHDOG_INTERVAL_MS = 30 * 1000;
+const DEFAULT_PAUSE_MS = 30 * 1000;
+const RELEASE_COOLDOWN_MS = 350;
+
+/**
+ * Storage key prefixes. Limits are stored under the bare `<providerId>` key
+ * (unchanged for backward compatibility). Cooldowns + the adaptive-headroom
+ * auto-reduce marker are stored under their own prefixed keys so they persist
+ * independently without loadLimits misreading them as limits.
+ */
+const COOLDOWN_KEY_PREFIX = "cooldown:";
+const AUTOREDUCE_KEY_PREFIX = "auto-reduce:";
+
+/**
+ * Wrap a `ConcurrencyService` so `setLimit`/`removeLimit`/`setCooldown` persist
+ * to the given `StorageNamespace`. All other methods delegate directly to the
+ * inner service. Persistence is fire-and-forget — a storage write failure logs
+ * a warning but does NOT fail the API call (the in-memory value is already set).
+ *
+ * `restoreLimit` is NOT persisted here — it is a startup restore FROM disk, so
+ * it delegates straight through (the value is already on disk).
+ */
+function createPersistedService(
+ inner: ConcurrencyService,
+ storage: StorageNamespace,
+ logger: Logger,
+): ConcurrencyService {
+ return {
+ acquire: inner.acquire.bind(inner),
+ reportRateLimit: inner.reportRateLimit.bind(inner),
+ setLimit(providerId, limit) {
+ inner.setLimit(providerId, limit);
+ storage.set(providerId, String(limit)).catch((err) =>
+ logger.warn("provider-concurrency: failed to persist limit", {
+ providerId,
+ err: err instanceof Error ? err.message : String(err),
+ }),
+ );
+ // A MANUAL limit set clears the auto-reduce notice (the user took
+ // control) → drop the persisted auto-reduce marker too.
+ storage.delete(`${AUTOREDUCE_KEY_PREFIX}${providerId}`).catch(() => {
+ /* absent marker is fine */
+ });
+ },
+ restoreLimit: inner.restoreLimit.bind(inner),
+ removeLimit(providerId) {
+ inner.removeLimit(providerId);
+ storage.delete(providerId).catch((err) =>
+ logger.warn("provider-concurrency: failed to delete persisted limit", {
+ providerId,
+ err: err instanceof Error ? err.message : String(err),
+ }),
+ );
+ storage.delete(`${AUTOREDUCE_KEY_PREFIX}${providerId}`).catch(() => {
+ /* absent marker is fine */
+ });
+ },
+ setCooldown(providerId, cooldownMs) {
+ inner.setCooldown(providerId, cooldownMs);
+ storage.set(`${COOLDOWN_KEY_PREFIX}${providerId}`, String(cooldownMs)).catch((err) =>
+ logger.warn("provider-concurrency: failed to persist cooldown", {
+ providerId,
+ err: err instanceof Error ? err.message : String(err),
+ }),
+ );
+ },
+ getLimit: inner.getLimit.bind(inner),
+ getLimits: inner.getLimits.bind(inner),
+ getCooldown: inner.getCooldown.bind(inner),
+ getCooldowns: inner.getCooldowns.bind(inner),
+ getStatus: inner.getStatus.bind(inner),
+ getStatusAll: inner.getStatusAll.bind(inner),
+ notifyWorkspaceStarred: inner.notifyWorkspaceStarred.bind(inner),
+ destroy: inner.destroy.bind(inner),
+ };
+}
+
+/**
+ * Load saved limits from storage and apply them to the manager via
+ * `restoreLimit` (NOT `setLimit` — Bug 3). `setLimit` is a MANUAL user action
+ * that clears the auto-reduce notice; using it at startup would wipe the
+ * persisted auto-reduce banner. `restoreLimit` seeds the limit WITHOUT clearing
+ * the notice, and `loadAutoReduce` re-applies the notice afterward.
+ *
+ * Skips prefixed keys (cooldown:/auto-reduce:) — those are loaded by their
+ * own loaders.
+ */
+async function loadLimits(
+ storage: StorageNamespace,
+ manager: ConcurrencyService,
+ logger: Logger,
+): Promise<void> {
+ const keys = await storage.keys();
+ for (const key of keys) {
+ if (key.startsWith(COOLDOWN_KEY_PREFIX)) continue; // cooldown settings
+ if (key.startsWith(AUTOREDUCE_KEY_PREFIX)) continue; // auto-reduce markers
+ const providerId = key;
+ const raw = await storage.get(providerId);
+ if (raw === null) continue;
+ const limit = Number.parseInt(raw, 10);
+ if (!Number.isNaN(limit) && limit > 0) {
+ manager.restoreLimit(providerId, limit);
+ logger.info(`provider-concurrency: restored limit ${limit} for "${providerId}"`);
+ }
+ }
+}
+
+/**
+ * Load saved auto-reduce markers and re-apply them via `restoreLimit` so the
+ * frontend banner survives a restart (Bug 3). A marker is stored under
+ * `auto-reduce:<providerId>` with the value = the ORIGINAL limit before
+ * reduction (autoReducedFrom). The current (reduced) limit was already restored
+ * by {@link loadLimits}; this call re-marks it as auto-reduced.
+ */
+async function loadAutoReduce(
+ storage: StorageNamespace,
+ manager: ConcurrencyService,
+ logger: Logger,
+): Promise<void> {
+ const keys = await storage.keys(AUTOREDUCE_KEY_PREFIX);
+ for (const key of keys) {
+ const providerId = key.slice(AUTOREDUCE_KEY_PREFIX.length);
+ if (providerId.length === 0) continue;
+ const raw = await storage.get(key);
+ if (raw === null) continue;
+ const autoReducedFrom = Number.parseInt(raw, 10);
+ if (!Number.isNaN(autoReducedFrom) && autoReducedFrom > 0) {
+ const currentLimit = manager.getLimit(providerId);
+ if (currentLimit !== undefined && currentLimit < autoReducedFrom) {
+ manager.restoreLimit(providerId, currentLimit, autoReducedFrom);
+ logger.info(
+ `provider-concurrency: restored auto-reduce notice for "${providerId}" ` +
+ `(${autoReducedFrom} -> ${currentLimit})`,
+ );
+ }
+ }
+ }
+}
+
+/**
+ * Load saved cooldowns from storage and apply them to the manager.
+ * Cooldowns are stored under `cooldown:<providerId>` keys (distinct from the
+ * bare-`<providerId>` limit keys) so the two settings persist independently.
+ */
+async function loadCooldowns(
+ storage: StorageNamespace,
+ manager: ConcurrencyService,
+ logger: Logger,
+): Promise<void> {
+ const keys = await storage.keys(COOLDOWN_KEY_PREFIX);
+ for (const key of keys) {
+ const providerId = key.slice(COOLDOWN_KEY_PREFIX.length);
+ if (providerId.length === 0) continue;
+ const raw = await storage.get(key);
+ if (raw === null) continue;
+ const cooldownMs = Number.parseInt(raw, 10);
+ if (!Number.isNaN(cooldownMs) && cooldownMs >= 0) {
+ manager.setCooldown(providerId, cooldownMs);
+ logger.info(`provider-concurrency: restored cooldown ${cooldownMs}ms for "${providerId}"`);
+ }
+ }
+}
+
+export async function activate(host: HostAPI): Promise<void> {
+ const logger = host.logger;
+ const storage = host.storage("provider-concurrency");
+
+ // Build the injected usage-poll effect from the host's provider registry.
+ // Lazy (called at poll time, not activate time) so activation order with the
+ // provider extensions doesn't matter. A provider that doesn't expose
+ // `getUsage` (or isn't registered) → returns undefined → the manager's usage
+ // gate falls back to cooldown-only recycling for that provider. This keeps
+ // the manager pure (the HTTP poll is an injected effect, not hardcoded fetch).
+ const fetchUsage = async (providerId: string) => {
+ const provider = host.getProviders().get(providerId);
+ if (provider === undefined || provider.getUsage === undefined) return undefined;
+ return provider.getUsage();
+ };
+
+ // Resolve the conversation store to seed the in-memory starred-workspace
+ // cache. The `isWorkspaceStarred` callback reads this cache synchronously
+ // (the queue sort comparator is sync), so we must populate it before the
+ // manager handles its first acquire. `dependsOn: ["conversation-store"]`
+ // in the manifest guarantees the store is registered before we activate.
+ const conversationStore = host.getService(conversationStoreHandle);
+
+ // The manager owns the in-memory `starredWorkspaces` set internally (the
+ // default `isWorkspaceStarred` callback checks it). We seed it by calling
+ // `notifyWorkspaceStarred` for each starred workspace found in the store.
+ const managerOpts: ConcurrencyManagerOpts = {
+ now: () => Date.now(),
+ slotTimeoutMs: SLOT_TIMEOUT_MS,
+ watchdogIntervalMs: WATCHDOG_INTERVAL_MS,
+ defaultPauseMs: DEFAULT_PAUSE_MS,
+ releaseCooldownMs: RELEASE_COOLDOWN_MS,
+ fetchUsage,
+ onWatchdogReclaim: (providerId, conversationId, heldMs) => {
+ logger.warn("provider-concurrency: watchdog reclaimed stale slot", {
+ providerId,
+ conversationId,
+ heldMs,
+ });
+ },
+ onPause: (providerId, durationMs) => {
+ logger.warn("provider-concurrency: 429 backoff — pausing queue", {
+ providerId,
+ durationMs,
+ });
+ },
+ onLimitReduced: (providerId, newLimit, oldLimit) => {
+ logger.warn("provider-concurrency: 429 adaptive headroom — limit reduced", {
+ providerId,
+ oldLimit,
+ newLimit,
+ });
+ // Persist the reduced (one-way) limit so it survives a restart, AND the
+ // auto-reduce marker (autoReducedFrom) so the banner survives too (Bug 3).
+ storage.set(providerId, String(newLimit)).catch((err) =>
+ logger.warn("provider-concurrency: failed to persist auto-reduced limit", {
+ providerId,
+ err: err instanceof Error ? err.message : String(err),
+ }),
+ );
+ storage.set(`${AUTOREDUCE_KEY_PREFIX}${providerId}`, String(oldLimit)).catch((err) =>
+ logger.warn("provider-concurrency: failed to persist auto-reduce marker", {
+ providerId,
+ err: err instanceof Error ? err.message : String(err),
+ }),
+ );
+ },
+ onUsagePollError: (providerId, err) => {
+ // A throwing getUsage() is treated as "no usage info" (cooldown-only
+ // fallback) by the manager — this is WARN-level observability only (Bug 2).
+ logger.warn("provider-concurrency: usage poll failed — falling back to cooldown-only", {
+ providerId,
+ err: err instanceof Error ? err.message : String(err),
+ });
+ },
+ };
+
+ const inner = createConcurrencyManager(managerOpts);
+
+ // Restore persisted limits + auto-reduce notices + cooldowns before registering
+ // the service so the first request sees the correct configuration.
+ await loadLimits(storage, inner, logger);
+ await loadAutoReduce(storage, inner, logger);
+ await loadCooldowns(storage, inner, logger);
+
+ // Seed the in-memory starred cache from the conversation store so the
+ // priority scheduling is correct on a fresh server start (previously-starred
+ // workspaces are respected without requiring the user to re-star them).
+ try {
+ const workspaces = await conversationStore.listWorkspaces();
+ for (const ws of workspaces) {
+ if (ws.starred) {
+ inner.notifyWorkspaceStarred(ws.id, true);
+ }
+ }
+ if (workspaces.some((w) => w.starred)) {
+ logger.info("provider-concurrency: restored starred workspaces", {
+ count: workspaces.filter((w) => w.starred).length,
+ });
+ }
+ } catch (err) {
+ logger.warn("provider-concurrency: failed to load starred workspaces", {
+ err: err instanceof Error ? err.message : String(err),
+ });
+ }
+
+ const service = createPersistedService(inner, storage, logger);
+ host.provideService(concurrencyServiceHandle, service);
+ logger.info("provider-concurrency: registered");
+}
+
+export const extension: Extension = {
+ manifest,
+ activate,
+};
diff --git a/packages/provider-concurrency/src/index.ts b/packages/provider-concurrency/src/index.ts
new file mode 100644
index 0000000..f35c070
--- /dev/null
+++ b/packages/provider-concurrency/src/index.ts
@@ -0,0 +1,10 @@
+export {
+ type ConcurrencyLimiter,
+ type ConcurrencyManagerOpts,
+ type ConcurrencyService,
+ createConcurrencyManager,
+ type ProviderConcurrencyStatus,
+} from "./concurrency-manager.js";
+export { extension, manifest } from "./extension.js";
+export { wrapProviderWithConcurrency } from "./provider-wrapper.js";
+export { concurrencyServiceHandle } from "./service.js";
diff --git a/packages/provider-concurrency/src/provider-wrapper.test.ts b/packages/provider-concurrency/src/provider-wrapper.test.ts
new file mode 100644
index 0000000..7554e64
--- /dev/null
+++ b/packages/provider-concurrency/src/provider-wrapper.test.ts
@@ -0,0 +1,262 @@
+import type { ProviderContract, ProviderEvent } from "@dispatch/kernel";
+import { describe, expect, it } from "vitest";
+import type { ConcurrencyLimiter } from "./concurrency-manager.js";
+import { wrapProviderWithConcurrency } from "./provider-wrapper.js";
+
+/** Build a fake provider that yields a sequence of events. */
+function fakeProvider(events: ProviderEvent[]): ProviderContract {
+ return {
+ id: "test-provider",
+ stream: async function* (): AsyncIterable<ProviderEvent> {
+ for (const e of events) {
+ yield e;
+ }
+ },
+ };
+}
+
+/** A fake limiter that records acquire/release calls. */
+function recordingLimiter(): ConcurrencyLimiter & {
+ acquireCalls: {
+ providerId: string;
+ conversationId: string;
+ workspaceId: string;
+ promptStartedAt: number;
+ }[];
+ releaseCalls: number;
+ rateLimitReports: string[];
+} {
+ const acquireCalls: {
+ providerId: string;
+ conversationId: string;
+ workspaceId: string;
+ promptStartedAt: number;
+ }[] = [];
+ const releaseCalls: { count: number } = { count: 0 };
+ const rateLimitReports: string[] = [];
+
+ return {
+ acquireCalls,
+ get releaseCalls() {
+ return releaseCalls.count;
+ },
+ rateLimitReports,
+ acquire(providerId, conversationId, workspaceId, promptStartedAt) {
+ acquireCalls.push({ providerId, conversationId, workspaceId, promptStartedAt });
+ return Promise.resolve(() => {
+ releaseCalls.count++;
+ });
+ },
+ reportRateLimit(providerId) {
+ rateLimitReports.push(providerId);
+ },
+ };
+}
+
+describe("wrapProviderWithConcurrency", () => {
+ it("acquires a slot before streaming and releases after the stream completes", async () => {
+ const provider = fakeProvider([
+ { type: "text-delta", delta: "hello" },
+ { type: "finish", reason: "stop" },
+ ]);
+ const limiter = recordingLimiter();
+
+ const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 12345);
+
+ const events: ProviderEvent[] = [];
+ for await (const e of wrapped.stream([], [])) {
+ events.push(e);
+ }
+
+ // Slot acquired before stream, released after.
+ expect(limiter.acquireCalls).toEqual([
+ {
+ providerId: "test-provider",
+ conversationId: "conv1",
+ workspaceId: "default",
+ promptStartedAt: 12345,
+ },
+ ]);
+ expect(limiter.releaseCalls).toBe(1);
+ expect(events).toEqual([
+ { type: "text-delta", delta: "hello" },
+ { type: "finish", reason: "stop" },
+ ]);
+ });
+
+ it("releases the slot even when the stream throws", async () => {
+ const provider: ProviderContract = {
+ id: "err-provider",
+ stream: async function* (): AsyncIterable<ProviderEvent> {
+ yield { type: "text-delta", delta: "partial" };
+ throw new Error("stream exploded");
+ },
+ };
+ const limiter = recordingLimiter();
+ const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 0);
+
+ await expect(async () => {
+ for await (const _e of wrapped.stream([], [])) {
+ // consume
+ }
+ }).rejects.toThrow("stream exploded");
+
+ expect(limiter.releaseCalls).toBe(1);
+ });
+
+ it("reports 429 errors to the limiter", async () => {
+ const provider = fakeProvider([
+ { type: "error", message: "Too many requests", code: "429", retryable: true },
+ ]);
+ const limiter = recordingLimiter();
+ const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 0);
+
+ const events: ProviderEvent[] = [];
+ for await (const e of wrapped.stream([], [])) {
+ events.push(e);
+ }
+
+ expect(limiter.rateLimitReports).toEqual(["test-provider"]);
+ // The 429 error event is still yielded to the consumer (kernel handles retry).
+ expect(events).toHaveLength(1);
+ expect(events[0]?.type).toBe("error");
+ });
+
+ it("does not report non-429 errors", async () => {
+ const provider = fakeProvider([
+ { type: "error", message: "Internal error", code: "500", retryable: true },
+ ]);
+ const limiter = recordingLimiter();
+ const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 0);
+
+ for await (const _e of wrapped.stream([], [])) {
+ // consume
+ }
+
+ expect(limiter.rateLimitReports).toEqual([]);
+ });
+
+ it("preserves the provider id and listModels", async () => {
+ const provider: ProviderContract = {
+ id: "my-provider",
+ stream: async function* (): AsyncIterable<ProviderEvent> {
+ yield { type: "finish", reason: "stop" };
+ },
+ listModels: async () => [{ id: "model-1" }],
+ };
+ const limiter = recordingLimiter();
+ const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 0);
+
+ expect(wrapped.id).toBe("my-provider");
+ expect(wrapped.listModels).toBeDefined();
+ const models = await wrapped.listModels?.();
+ expect(models).toEqual([{ id: "model-1" }]);
+ });
+
+ it("calls onQueued when the request blocks and onAcquired when the slot is granted", async () => {
+ let queuedCalled = false;
+ let acquiredCalled = false;
+
+ const blockingLimiter: ConcurrencyLimiter = {
+ acquire(_providerId, _convId, _wsId, _promptAt, onQueued) {
+ // Simulate a queued request: call onQueued, then resolve on next tick.
+ onQueued?.();
+ return new Promise((resolve) => {
+ setTimeout(() => {
+ resolve(() => {});
+ }, 0);
+ });
+ },
+ reportRateLimit() {},
+ };
+
+ const provider = fakeProvider([{ type: "finish", reason: "stop" }]);
+ const wrapped = wrapProviderWithConcurrency(
+ provider,
+ blockingLimiter,
+ "conv1",
+ "default",
+ 0,
+ () => {
+ queuedCalled = true;
+ },
+ () => {
+ acquiredCalled = true;
+ },
+ );
+
+ for await (const _e of wrapped.stream([], [])) {
+ // consume
+ }
+
+ expect(queuedCalled).toBe(true);
+ expect(acquiredCalled).toBe(true);
+ });
+
+ it("does NOT call onQueued when the slot is granted immediately", async () => {
+ let queuedCalled = false;
+ let acquiredCalled = false;
+
+ const immediateLimiter: ConcurrencyLimiter = {
+ acquire(_providerId, _convId, _wsId, _promptAt, _onQueued) {
+ // Grant immediately — do NOT call onQueued.
+ return Promise.resolve(() => {});
+ },
+ reportRateLimit() {},
+ };
+
+ const provider = fakeProvider([{ type: "finish", reason: "stop" }]);
+ const wrapped = wrapProviderWithConcurrency(
+ provider,
+ immediateLimiter,
+ "conv1",
+ "default",
+ 0,
+ () => {
+ queuedCalled = true;
+ },
+ () => {
+ acquiredCalled = true;
+ },
+ );
+
+ for await (const _e of wrapped.stream([], [])) {
+ // consume
+ }
+
+ expect(queuedCalled).toBe(false);
+ expect(acquiredCalled).toBe(true);
+ });
+
+ it("passes through messages, tools, and opts to the inner stream", async () => {
+ let receivedArgs:
+ | {
+ messages: unknown;
+ tools: unknown;
+ opts: unknown;
+ }
+ | undefined;
+
+ const provider: ProviderContract = {
+ id: "passthrough",
+ stream: async function* (messages, tools, opts): AsyncIterable<ProviderEvent> {
+ receivedArgs = { messages, tools, opts };
+ yield { type: "finish", reason: "stop" };
+ },
+ };
+ const limiter = recordingLimiter();
+ const wrapped = wrapProviderWithConcurrency(provider, limiter, "conv1", "default", 0);
+
+ const messages = [{ role: "user" as const, chunks: [{ type: "text" as const, text: "hi" }] }];
+ const tools = [{ name: "test_tool", description: "test", parameters: {} }];
+ const opts = { model: "gpt-4" };
+
+ for await (const _e of wrapped.stream(messages, tools, opts)) {
+ // consume
+ }
+
+ expect(receivedArgs?.messages).toBe(messages);
+ expect(receivedArgs?.tools).toBe(tools);
+ expect(receivedArgs?.opts).toBe(opts);
+ });
+});
diff --git a/packages/provider-concurrency/src/provider-wrapper.ts b/packages/provider-concurrency/src/provider-wrapper.ts
new file mode 100644
index 0000000..1e3f2c0
--- /dev/null
+++ b/packages/provider-concurrency/src/provider-wrapper.ts
@@ -0,0 +1,77 @@
+import type {
+ ChatMessage,
+ ProviderContract,
+ ProviderEvent,
+ ProviderStreamOptions,
+ ToolContract,
+} from "@dispatch/kernel";
+import type { ConcurrencyLimiter } from "./concurrency-manager.js";
+
+/**
+ * Wrap a provider's `stream` method with concurrency limiting.
+ *
+ * A slot is acquired BEFORE the first event is yielded (before the HTTP
+ * request is sent — the `await limiter.acquire()` runs before the generator
+ * body starts iterating the inner stream). The slot is released in a `finally`
+ * block AFTER the inner stream completes (the full response stream, not just
+ * HTTP headers — matching the Umans concurrency model where a slot is held
+ * only while tokens are actually generating).
+ *
+ * 429 detection: if the provider yields an `error` event with `code: "429"`,
+ * the limiter is notified so it can pause the queue for that provider.
+ *
+ * @param provider The underlying provider to wrap.
+ * @param limiter The concurrency limiter (acquire/release/reportRateLimit).
+ * @param conversationId The agent requesting the stream (for slot attribution).
+ * @param workspaceId The workspace the agent belongs to (for starred
+ * priority scheduling in the limiter queue).
+ * @param promptStartedAt When the agent's current prompt (turn) started
+ * (epoch-ms, for oldest-agent-first scheduling).
+ * @param onQueued Called synchronously when `acquire()` decides to
+ * queue the request (cannot grant immediately).
+ * Lets the caller emit a "queued" status signal.
+ * @param onAcquired Called when `acquire()` resolves (slot granted,
+ * whether immediately or after queueing). Lets the
+ * caller emit an "active" status signal.
+ */
+export function wrapProviderWithConcurrency(
+ provider: ProviderContract,
+ limiter: ConcurrencyLimiter,
+ conversationId: string,
+ workspaceId: string,
+ promptStartedAt: number,
+ onQueued?: () => void,
+ onAcquired?: () => void,
+): ProviderContract {
+ const innerStream = provider.stream;
+ const providerId = provider.id;
+
+ return {
+ id: provider.id,
+ stream: async function* (
+ messages: readonly ChatMessage[],
+ tools: readonly ToolContract[],
+ opts?: ProviderStreamOptions,
+ ): AsyncIterable<ProviderEvent> {
+ const release = await limiter.acquire(
+ providerId,
+ conversationId,
+ workspaceId,
+ promptStartedAt,
+ onQueued,
+ );
+ onAcquired?.();
+ try {
+ for await (const event of innerStream(messages, tools, opts)) {
+ if (event.type === "error" && event.code === "429") {
+ limiter.reportRateLimit(providerId);
+ }
+ yield event;
+ }
+ } finally {
+ release();
+ }
+ },
+ ...(provider.listModels !== undefined ? { listModels: provider.listModels } : {}),
+ };
+}
diff --git a/packages/provider-concurrency/src/service.ts b/packages/provider-concurrency/src/service.ts
new file mode 100644
index 0000000..aa578e8
--- /dev/null
+++ b/packages/provider-concurrency/src/service.ts
@@ -0,0 +1,11 @@
+import { defineService } from "@dispatch/kernel";
+import type { ConcurrencyService } from "./concurrency-manager.js";
+
+/**
+ * Typed service handle for the provider-concurrency service. The
+ * `provider-concurrency` extension provides the implementation; the
+ * session-orchestrator + transport-http consume it.
+ */
+export const concurrencyServiceHandle = defineService<ConcurrencyService>(
+ "provider-concurrency/service",
+);
diff --git a/packages/provider-concurrency/tsconfig.json b/packages/provider-concurrency/tsconfig.json
new file mode 100644
index 0000000..a4b26ef
--- /dev/null
+++ b/packages/provider-concurrency/tsconfig.json
@@ -0,0 +1,6 @@
+{
+ "extends": "../../tsconfig.base.json",
+ "compilerOptions": { "rootDir": "src", "outDir": "dist", "composite": true },
+ "include": ["src/**/*.ts"],
+ "references": [{ "path": "../conversation-store" }, { "path": "../kernel" }]
+}
generated by cgit v1.2.3 (git 2.39.1) at 2026-09-12 05:12:21 +0000