diff options
| -rw-r--r-- | packages/kernel/src/contracts/index.ts | 1 | ||||
| -rw-r--r-- | packages/kernel/src/contracts/provider.ts | 27 | ||||
| -rw-r--r-- | packages/openai-stream/src/getUsage.test.ts | 139 | ||||
| -rw-r--r-- | packages/openai-stream/src/getUsage.ts | 73 | ||||
| -rw-r--r-- | packages/openai-stream/src/index.ts | 1 | ||||
| -rw-r--r-- | packages/openai-stream/src/provider.test.ts | 21 | ||||
| -rw-r--r-- | packages/openai-stream/src/provider.ts | 9 | ||||
| -rw-r--r-- | packages/provider-concurrency/src/concurrency-manager.test.ts | 483 | ||||
| -rw-r--r-- | packages/provider-concurrency/src/concurrency-manager.ts | 475 | ||||
| -rw-r--r-- | packages/provider-concurrency/src/extension.ts | 179 | ||||
| -rw-r--r-- | packages/transport-contract/src/index.ts | 38 | ||||
| -rw-r--r-- | packages/transport-http/src/app.ts | 49 |
12 files changed, 1442 insertions, 53 deletions
diff --git a/packages/kernel/src/contracts/index.ts b/packages/kernel/src/contracts/index.ts index 28e0a0b..fc19267 100644 --- a/packages/kernel/src/contracts/index.ts +++ b/packages/kernel/src/contracts/index.ts @@ -103,6 +103,7 @@ export type { ProviderEvent, ProviderStreamOptions, ProviderToolCallEvent, + ProviderUsage, ReasoningDeltaEvent, ReasoningEffort, TextDeltaEvent, diff --git a/packages/kernel/src/contracts/provider.ts b/packages/kernel/src/contracts/provider.ts index 3137073..dea6c17 100644 --- a/packages/kernel/src/contracts/provider.ts +++ b/packages/kernel/src/contracts/provider.ts @@ -104,6 +104,19 @@ export interface ProviderStreamOptions { } /** + * A snapshot of the provider's current upstream usage. Returned by a + * provider's optional `getUsage` so a concurrency limiter can gate slot grants + * on the REAL upstream in-flight count (not just the limiter's local accounting, + * which lags the upstream `concurrent_sessions` counter by the release + * cooldown). `concurrentSessions` is the number of requests the provider counts + * as currently in flight. + */ +export interface ProviderUsage { + /** Upstream count of currently in-flight (generating) sessions. */ + readonly concurrentSessions: number; +} + +/** * Metadata describing a single model a provider can serve. Returned by * `listModels` so a catalog (e.g. the credential-store) can enumerate the * `<credentialName>/<model>` choices a client may select. Kept minimal — `id` @@ -154,4 +167,18 @@ export interface ProviderContract { * credentials in; today the provider uses the key it resolved at activate. */ readonly listModels?: () => Promise<readonly ModelInfo[]>; + + /** + * Fetch the provider's current upstream usage snapshot. Optional: a provider + * that cannot (or chooses not to) report usage omits it, and a concurrency + * limiter falls back to cooldown-only slot recycling (no usage gate). When + * present, the limiter polls this before admitting a QUEUED agent and grants + * only when `concurrentSessions` is below the configured limit — preventing an + * N+1 overshoot from the upstream accounting lag. + * + * May return `undefined` (e.g. the endpoint returned an unexpected shape or a + * non-200) — the limiter treats `undefined` as "no usage info available" and + * falls back to granting (cooldown-only behavior) for that poll. + */ + readonly getUsage?: () => Promise<ProviderUsage | undefined>; } diff --git a/packages/openai-stream/src/getUsage.test.ts b/packages/openai-stream/src/getUsage.test.ts new file mode 100644 index 0000000..c319b73 --- /dev/null +++ b/packages/openai-stream/src/getUsage.test.ts @@ -0,0 +1,139 @@ +import type { ProviderUsage } from "@dispatch/kernel"; +import type { FetchLike } from "@dispatch/trace-replay"; +import { describe, expect, it, vi } from "vitest"; +import { getUsage } from "./getUsage.js"; + +function jsonResponse(body: unknown, status = 200): Response { + return new Response(JSON.stringify(body), { + status, + headers: { "Content-Type": "application/json" }, + }); +} + +describe("getUsage", () => { + it("extracts concurrent_sessions from the Umans /v1/usage shape", async () => { + const fetchFn = vi.fn( + () => + jsonResponse({ + usage: { concurrent_sessions: 3 }, + limits: { concurrency: { limit: 4, hard_cap: 5 } }, + }) as unknown as ReturnType<FetchLike>, + ); + + const result = await getUsage({ + baseURL: "https://api.umans.ai/v1", + apiKey: "sk-test-1234567890abcdef", + providerId: "umans", + fetchFn, + }); + + expect(result).toEqual<ProviderUsage>({ concurrentSessions: 3 }); + expect(fetchFn).toHaveBeenCalledOnce(); + const call = fetchFn.mock.calls[0]; + expect(call?.[0]).toBe("https://api.umans.ai/v1/usage"); + const init = call?.[1] as RequestInit; + expect(init.method).toBe("GET"); + expect((init.headers as Record<string, string>).Authorization).toBe( + "Bearer sk-test-1234567890abcdef", + ); + }); + + it("returns undefined on non-200 (endpoint unsupported)", async () => { + const fetchFn = vi.fn( + () => jsonResponse({ error: "not found" }, 404) as unknown as ReturnType<FetchLike>, + ); + + const result = await getUsage({ + baseURL: "https://api.example.com/v1", + apiKey: "sk-test", + providerId: "openai", + fetchFn, + }); + + expect(result).toBeUndefined(); + }); + + it("returns undefined on a network error (never throws)", async () => { + const fetchFn = vi.fn( + () => Promise.reject(new Error("ECONNREFUSED")) as unknown as Promise<Response>, + ); + + const result = await getUsage({ + baseURL: "https://api.example.com/v1", + apiKey: "sk-test", + providerId: "openai", + fetchFn, + }); + + expect(result).toBeUndefined(); + }); + + it("returns undefined when the response shape is unexpected", async () => { + const fetchFn = vi.fn( + () => jsonResponse({ unexpected: true }) as unknown as ReturnType<FetchLike>, + ); + + const result = await getUsage({ + baseURL: "https://api.example.com/v1", + apiKey: "sk-test", + providerId: "openai", + fetchFn, + }); + + expect(result).toBeUndefined(); + }); + + it("returns undefined when concurrent_sessions is not a number", async () => { + const fetchFn = vi.fn( + () => + jsonResponse({ + usage: { concurrent_sessions: "three" }, + }) as unknown as ReturnType<FetchLike>, + ); + + const result = await getUsage({ + baseURL: "https://api.example.com/v1", + apiKey: "sk-test", + providerId: "openai", + fetchFn, + }); + + expect(result).toBeUndefined(); + }); + + it("truncates a fractional concurrent_sessions to an integer", async () => { + const fetchFn = vi.fn( + () => + jsonResponse({ usage: { concurrent_sessions: 2.9 } }) as unknown as ReturnType<FetchLike>, + ); + + const result = await getUsage({ + baseURL: "https://api.example.com/v1", + apiKey: "sk-test", + providerId: "openai", + fetchFn, + }); + + expect(result).toEqual<ProviderUsage>({ concurrentSessions: 2 }); + }); + + it("falls back to globalThis.fetch when fetchFn is absent", async () => { + const original = globalThis.fetch; + const stub = vi.fn( + () => jsonResponse({ usage: { concurrent_sessions: 1 } }) as unknown as ReturnType<FetchLike>, + ); + globalThis.fetch = stub as unknown as typeof globalThis.fetch; + + try { + const result = await getUsage({ + baseURL: "https://api.example.com/v1", + apiKey: "sk-test", + providerId: "openai", + }); + expect(result).toEqual<ProviderUsage>({ concurrentSessions: 1 }); + expect(stub).toHaveBeenCalledOnce(); + } finally { + globalThis.fetch = original; + } + }); +}); diff --git a/packages/openai-stream/src/getUsage.ts b/packages/openai-stream/src/getUsage.ts new file mode 100644 index 0000000..5da7fd7 --- /dev/null +++ b/packages/openai-stream/src/getUsage.ts @@ -0,0 +1,73 @@ +import type { ProviderUsage } from "@dispatch/kernel"; +import type { FetchLike } from "@dispatch/trace-replay"; + +/** + * Generic OpenAI-compatible usage fetch. The Umans `/v1/usage` endpoint returns: + * + * { usage: { concurrent_sessions: number }, limits: { concurrency: { limit, hard_cap } } } + * + * We extract only `concurrent_sessions` (the count a concurrency limiter gates + * on). Lives in this library (`@dispatch/openai-stream`) so any OpenAI-compatible + * provider extension reuses it without cross-extension code import + * (isolation-over-DRY: coupling is via this typed library surface). A provider + * supplies its own `id` (used in error labels) via `createOpenAICompatProvider`. + */ + +/** The raw shape of the `/v1/usage` response (only the fields we read). */ +interface UsageResponse { + readonly usage?: { + readonly concurrent_sessions?: number; + }; +} + +export interface GetUsageConfig { + readonly baseURL: string; + readonly apiKey: string; + readonly fetchFn?: FetchLike; + readonly providerId: string; +} + +/** + * Fetch + map the upstream usage snapshot. Returns `undefined` on any error, + * non-200, or unexpected shape so the caller (the concurrency limiter) falls + * back to cooldown-only slot recycling (no usage gate) — never throws. + * + * Pure-ish I/O wrapper: the only effect is the injected fetch. Extracted for + * direct unit testing with a fake fetch. + */ +export async function getUsage(config: GetUsageConfig): Promise<ProviderUsage | undefined> { + const effectiveFetch: FetchLike = config.fetchFn ?? fetch; + const url = `${config.baseURL}/usage`; + + let response: Response; + try { + response = await effectiveFetch(url, { + method: "GET", + headers: { + Authorization: `Bearer ${config.apiKey}`, + }, + }); + } catch { + // Network error — the upstream is unreachable; treat as "no usage info". + return undefined; + } + + if (!response.ok) { + // 404 / 401 / 5xx — the endpoint is unsupported or rejected the request. + return undefined; + } + + let body: UsageResponse; + try { + body = (await response.json()) as UsageResponse; + } catch { + return undefined; + } + + const raw = body.usage?.concurrent_sessions; + if (typeof raw !== "number" || !Number.isFinite(raw) || raw < 0) { + return undefined; + } + + return { concurrentSessions: Math.trunc(raw) }; +} diff --git a/packages/openai-stream/src/index.ts b/packages/openai-stream/src/index.ts index 3f76b99..ff6b4f5 100644 --- a/packages/openai-stream/src/index.ts +++ b/packages/openai-stream/src/index.ts @@ -8,6 +8,7 @@ export type { export { convertMessages } from "./convert-messages.js"; export type { OpenAITool } from "./convert-tools.js"; export { convertTools } from "./convert-tools.js"; +export { getUsage } from "./getUsage.js"; export { isVisionModelId, parseModelList } from "./listModels.js"; export { parseSSELines } from "./parse-sse.js"; export type { CreateOpenAICompatProviderOpts } from "./provider.js"; diff --git a/packages/openai-stream/src/provider.test.ts b/packages/openai-stream/src/provider.test.ts index 8bc6e98..13a303e 100644 --- a/packages/openai-stream/src/provider.test.ts +++ b/packages/openai-stream/src/provider.test.ts @@ -81,6 +81,27 @@ describe("createOpenAICompatProvider stamps the given id on the ProviderContract await expect(listModels()).rejects.toThrow("listModels[my-custom-id]: HTTP 401 — Unauthorized"); }); + + it("exposes getUsage that returns the upstream concurrent_sessions", async () => { + const fetchFn = vi.fn( + () => + new Response(JSON.stringify({ usage: { concurrent_sessions: 2 } }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }) as unknown as ReturnType<FetchLike>, + ); + const provider = createOpenAICompatProvider({ + credentials: makeCreds(), + model: "test-model", + id: "umans", + fetchFn, + }); + const getUsage = provider.getUsage; + if (!getUsage) throw new Error("getUsage not defined"); + + const usage = await getUsage(); + expect(usage).toEqual({ concurrentSessions: 2 }); + }); }); describe("transformBody", () => { diff --git a/packages/openai-stream/src/provider.ts b/packages/openai-stream/src/provider.ts index df5a4ed..9a9369f 100644 --- a/packages/openai-stream/src/provider.ts +++ b/packages/openai-stream/src/provider.ts @@ -4,9 +4,11 @@ import type { ModelInfo, ProviderContract, ProviderStreamOptions, + ProviderUsage, ToolContract, } from "@dispatch/kernel"; import type { FetchLike } from "@dispatch/trace-replay"; +import { getUsage as fetchUsage } from "./getUsage.js"; import { listModels as fetchModels } from "./listModels.js"; import { streamChat } from "./stream.js"; @@ -69,5 +71,12 @@ export function createOpenAICompatProvider(opts: CreateOpenAICompatProviderOpts) providerId: opts.id, ...(fetchFn !== undefined ? { fetchFn } : {}), }), + getUsage: (): Promise<ProviderUsage | undefined> => + fetchUsage({ + baseURL, + apiKey, + providerId: opts.id, + ...(fetchFn !== undefined ? { fetchFn } : {}), + }), }; } diff --git a/packages/provider-concurrency/src/concurrency-manager.test.ts b/packages/provider-concurrency/src/concurrency-manager.test.ts index 1e45eec..357a5d1 100644 --- a/packages/provider-concurrency/src/concurrency-manager.test.ts +++ b/packages/provider-concurrency/src/concurrency-manager.test.ts @@ -1,3 +1,4 @@ +import type { ProviderUsage } from "@dispatch/kernel"; import { describe, expect, it } from "vitest"; import { type ConcurrencyService, createConcurrencyManager } from "./concurrency-manager.js"; @@ -59,7 +60,12 @@ function createFakeTimers() { }; } -function createManager(opts?: { releaseCooldownMs?: number }): { +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>; } { @@ -70,6 +76,9 @@ function createManager(opts?: { releaseCooldownMs?: number }): { 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, @@ -100,6 +109,8 @@ describe("createConcurrencyManager", () => { inFlight: 1, queued: 0, paused: false, + cooldownMs: 0, + autoReduced: false, }); release1(); expect(manager.getStatus("umans")?.inFlight).toBe(0); @@ -342,6 +353,8 @@ describe("createConcurrencyManager", () => { inFlight: 0, queued: 0, paused: false, + cooldownMs: 0, + autoReduced: false, }); }); @@ -484,6 +497,474 @@ describe("createConcurrencyManager", () => { 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", 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", 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", 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", 0), + manager.acquire("umans", "c2", 0), + manager.acquire("umans", "c3", 0), + manager.acquire("umans", "c4", 0), + ]); + expect(manager.getStatus("umans")?.inFlight).toBe(4); + + // 5th agent queues. + let resolved = false; + const promise5 = manager.acquire("umans", "c5", 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", 0), + manager.acquire("umans", "c2", 0), + manager.acquire("umans", "c3", 0), + manager.acquire("umans", "c4", 0), + ]); + + let resolved = false; + const promise5 = manager.acquire("umans", "c5", 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", 0); + let resolved = false; + const promise2 = manager.acquire("umans", "c2", 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", 0), + manager.acquire("umans", "c2", 0), + manager.acquire("umans", "c3", 0), + manager.acquire("umans", "c4", 0), + ]); + + // Queue two waiters. + let r5 = false; + let r6 = false; + const p5 = manager.acquire("umans", "c5", 10).then((r) => { + r5 = true; + return r; + }); + const p6 = manager.acquire("umans", "c6", 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", 0).then(async (release1) => { + // Queue a waiter (arms the 1s fallback timer). + const p2 = manager.acquire("umans", "c2", 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", 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", 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", 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", 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", 0); + // Queue a waiter; release → recycle → poll THROWS. + const p2 = manager.acquire("umans", "c2", 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", 0); + manager.acquire("umans", "c2", 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 ─────────────────────────────────────── diff --git a/packages/provider-concurrency/src/concurrency-manager.ts b/packages/provider-concurrency/src/concurrency-manager.ts index e69edb5..ea66a49 100644 --- a/packages/provider-concurrency/src/concurrency-manager.ts +++ b/packages/provider-concurrency/src/concurrency-manager.ts @@ -7,14 +7,32 @@ * (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. + * 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. * - * This module is the PURE decision logic. It takes an injected clock (`now`) - * and injected timers (`setTimeout`/`clearTimeout`/`setInterval`/`clearInterval`) - * so it is fully testable with deterministic fake time. The extension layer - * wires real timers. + * ── 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. */ @@ -30,6 +48,25 @@ export interface ProviderConcurrencyStatus { 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; } /** @@ -79,8 +116,10 @@ export interface ConcurrencyLimiter { /** * Report a 429 from a provider. Pauses the queue for that provider for - * `retryAfterMs` (or a default duration when omitted). Queued and in-flight - * requests are unaffected; new `acquire` calls block until the pause expires. + * `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; } @@ -89,14 +128,33 @@ export interface ConcurrencyLimiter { * The full service surface (limiter + config + status) for HTTP routes. */ export interface ConcurrencyService extends ConcurrencyLimiter { - /** Set the concurrency limit for a provider. Creates the state if new. */ + /** 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. */ @@ -136,6 +194,19 @@ interface ProviderState { 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 { @@ -148,22 +219,40 @@ export interface ConcurrencyManagerOpts { /** Default pause duration when a 429 arrives without Retry-After (ms). */ readonly defaultPauseMs: number; /** - * 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. + * 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 events. */ + /** 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 @@ -174,6 +263,11 @@ export interface ConcurrencyManagerOpts { 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. } @@ -182,7 +276,8 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre const now = opts.now; const slotTimeoutMs = opts.slotTimeoutMs; const defaultPauseMs = opts.defaultPauseMs; - const releaseCooldownMs = opts.releaseCooldownMs ?? 0; + 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); @@ -199,9 +294,34 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre 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 { @@ -212,20 +332,35 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre released = true; state.slots.delete(id); - // Recycle the slot: decrement inFlight + 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 + // 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 = () => { - state.inFlight--; - tryGrantNext(providerId); + 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 (releaseCooldownMs > 0) { + if (state.cooldownMs > 0) { const timer = setTimeout(() => { cooldownTimers.delete(timer); recycle(); - }, releaseCooldownMs); + }, state.cooldownMs); cooldownTimers.add(timer); } else { recycle(); @@ -241,6 +376,12 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre } /** + * 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. + */ + function grantLoop(state: ProviderState, providerId: string): void { * 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 @@ -267,6 +408,159 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre 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; + 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 ────────────────────────────────────────────────────────────────── @@ -286,6 +580,14 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre 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 = { @@ -297,7 +599,19 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre } if (!state.paused && state.inFlight < state.limit) { - return Promise.resolve(grantSlot(state, providerId, conversationId)); + // 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. @@ -307,6 +621,15 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre // Queue (starred-workspace-first, then oldest-agent-first). return new Promise<() => void>((resolve) => { + state.queue.push({ conversationId, promptStartedAt, resolve }); + // Keep sorted ascending by promptStartedAt (oldest first). + state.queue.sort((a, b) => a.promptStartedAt - b.promptStartedAt); + // 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); + } 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. @@ -326,32 +649,66 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre 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; - tryGrantNext(providerId); + void tryGrantNext(providerId); }, pauseDuration); }, setLimit(providerId, limit) { let state = states.get(providerId); if (state === undefined) { - state = { - limit, - inFlight: 0, - slots: new Map(), - queue: [], - paused: false, - pausedUntil: undefined, - pauseTimer: 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. - tryGrantNext(providerId); + 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) { @@ -369,6 +726,12 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre 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) { @@ -392,6 +755,41 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre })); }, + 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; @@ -401,7 +799,11 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre 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 } : {}), }; }, @@ -437,6 +839,9 @@ export function createConcurrencyManager(opts: ConcurrencyManagerOpts): Concurre if (state.pauseTimer !== undefined) { clearTimeout(state.pauseTimer); } + if (state.gateRepollTimer !== undefined) { + clearTimeout(state.gateRepollTimer); + } } states.clear(); }, diff --git a/packages/provider-concurrency/src/extension.ts b/packages/provider-concurrency/src/extension.ts index 48fbdef..48c019b 100644 --- a/packages/provider-concurrency/src/extension.ts +++ b/packages/provider-concurrency/src/extension.ts @@ -26,24 +26,40 @@ export const manifest: Manifest = { * - `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. - * - `RELEASE_COOLDOWN_MS` (200ms): after a slot is released, hold it for this + * 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. 200ms is the - * default most concurrency proxies use for AI/LLM APIs. + * 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 = 200; +const RELEASE_COOLDOWN_MS = 350; /** - * Wrap a `ConcurrencyService` so `setLimit`/`removeLimit` 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 limit is already set). + * 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, @@ -61,7 +77,13 @@ function createPersistedService( 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) => @@ -70,9 +92,23 @@ function createPersistedService( 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), @@ -81,8 +117,14 @@ function createPersistedService( } /** - * Load saved limits from storage and apply them to the manager. - * Called during activate, before the service is registered. + * 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, @@ -90,21 +132,92 @@ async function loadLimits( logger: Logger, ): Promise<void> { const keys = await storage.keys(); - for (const providerId of 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.setLimit(providerId, limit); + 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 @@ -121,6 +234,7 @@ export async function activate(host: HostAPI): Promise<void> { 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, @@ -134,13 +248,44 @@ export async function activate(host: HostAPI): Promise<void> { 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 before registering the service so the first - // request sees the correct configuration. + // 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 diff --git a/packages/transport-contract/src/index.ts b/packages/transport-contract/src/index.ts index d5f3000..797ad22 100644 --- a/packages/transport-contract/src/index.ts +++ b/packages/transport-contract/src/index.ts @@ -1062,6 +1062,17 @@ export interface ConcurrencyLimitResponse { * - `queued`: how many agents are waiting for a slot. * - `paused`: whether the queue is paused due to a 429 backoff. * - `pausedUntil`: when the pause expires (epoch-ms), present only when paused. + * - `cooldownMs`: the per-slot release cooldown (ms). A recycled slot is held + * this long before the next waiter is admitted — covers the upstream + * provider's accounting lag. Configurable + persisted per provider. + * - `autoReduced`: whether the limit was auto-reduced by 1 after a 429 + * (adaptive headroom, one-way, persisted). The user restores the limit + * manually via `PUT /concurrency/limits/:providerId`, which clears the flag. + * When `true`, the frontend renders a visible notice/banner. + * - `autoReducedFrom`: the original limit before auto-reduction (present only + * when `autoReduced` is true). + * - `notice`: a human-readable notice string for the frontend to render as a + * banner when the limit was auto-reduced (present only when `autoReduced`). */ export interface ConcurrencyStatusEntry { readonly providerId: string; @@ -1070,6 +1081,10 @@ export interface ConcurrencyStatusEntry { readonly queued: number; readonly paused: boolean; readonly pausedUntil?: number; + readonly cooldownMs: number; + readonly autoReduced: boolean; + readonly autoReducedFrom?: number; + readonly notice?: string; } /** @@ -1079,3 +1094,26 @@ export interface ConcurrencyStatusEntry { export interface ConcurrencyStatusResponse { readonly providers: readonly ConcurrencyStatusEntry[]; } + +// ─── Provider concurrency cooldown ──────────────────────────────────────────── + +/** + * Response of `GET /concurrency/cooldown/:providerId` — the per-slot release + * cooldown (ms) for a provider. A recycled slot is held this long before the + * next waiter is admitted, covering the upstream provider's accounting lag. + * When no cooldown was explicitly set, the server default (350ms) is returned. + */ +export interface ConcurrencyCooldownResponse { + readonly providerId: string; + readonly cooldownMs: number; +} + +/** + * Body of `PUT /concurrency/cooldown/:providerId` — set the release cooldown + * (ms) for a provider. `cooldownMs` must be a non-negative integer (0 = no + * cooldown, instant re-admission). The value is persisted and applied to + * subsequently recycled slots. + */ +export interface SetConcurrencyCooldownRequest { + readonly cooldownMs: number; +} diff --git a/packages/transport-http/src/app.ts b/packages/transport-http/src/app.ts index ebbf536..656be9d 100644 --- a/packages/transport-http/src/app.ts +++ b/packages/transport-http/src/app.ts @@ -8,6 +8,7 @@ import type { ComputerListResponse, ComputerResponse, ComputerStatusResponse, + ConcurrencyCooldownResponse, ConcurrencyLimitResponse, ConcurrencyLimitsResponse, ConcurrencyStatusResponse, @@ -31,6 +32,7 @@ import type { QueueResponse, ReasoningEffortResponse, SetCompactPercentRequest, + SetConcurrencyCooldownRequest, SetConcurrencyLimitRequest, SetConversationComputerRequest, SetSystemPromptTemplateRequest, @@ -687,6 +689,53 @@ export function createApp(opts: CreateServerOptions): Hono { return c.json({ ok: true, providerId }, 200); }); + app.get("/concurrency/cooldown/:providerId", (c) => { + const providerId = c.req.param("providerId"); + if (opts.concurrencyService === undefined) { + return c.json({ error: "Concurrency service not available" }, 503); + } + // A cooldown may be the default (when a limit is configured but no explicit + // cooldown was set) or explicitly set. getCooldown returns undefined only + // when the provider has NO state at all (no limit, no cooldown) — treat that + // as "not configured". + const cooldownMs = opts.concurrencyService.getCooldown(providerId); + if (cooldownMs === undefined) { + return c.json({ error: "No concurrency configuration for this provider" }, 404); + } + const body: ConcurrencyCooldownResponse = { providerId, cooldownMs }; + return c.json(body, 200); + }); + + app.put("/concurrency/cooldown/:providerId", async (c) => { + const providerId = c.req.param("providerId"); + if (opts.concurrencyService === undefined) { + return c.json({ error: "Concurrency service not available" }, 503); + } + + let body: unknown; + try { + body = await c.req.json(); + } catch { + log.warn("concurrency: invalid JSON body"); + return c.json({ error: "Invalid JSON body" }, 400); + } + + const parsed = body as SetConcurrencyCooldownRequest; + if ( + parsed === null || + typeof parsed !== "object" || + typeof parsed.cooldownMs !== "number" || + !Number.isInteger(parsed.cooldownMs) || + parsed.cooldownMs < 0 + ) { + return c.json({ error: "Body must be { cooldownMs: <non-negative integer> }" }, 400); + } + + opts.concurrencyService.setCooldown(providerId, parsed.cooldownMs); + const responseBody: ConcurrencyCooldownResponse = { providerId, cooldownMs: parsed.cooldownMs }; + return c.json(responseBody, 200); + }); + app.get("/concurrency/status", (c) => { if (opts.concurrencyService === undefined) { const body: ConcurrencyStatusResponse = { providers: [] }; |