diff options
Diffstat (limited to 'packages/openai-stream/src')
| -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 |
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 } : {}), + }), }; } |
