feat(concurrency-fixes): usage-gate + adaptive headroom + configurable cooldown

5 files changed, 243 insertions, 0 deletions

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