diff --git a/src/main/worker/runtime/credential-concurrency-gate.ts b/src/main/worker/runtime/credential-concurrency-gate.ts new file mode 100644 index 00000000..9740393b --- /dev/null +++ b/src/main/worker/runtime/credential-concurrency-gate.ts @@ -0,0 +1,132 @@ +import { createHash } from "node:crypto"; + +/** + * Per-credential async concurrency gate. + * + * The gate hands out a bounded number of concurrent slots per + * `(provider, apiKey)` pair. When the cap is reached, further + * `acquire` calls park until an in-flight caller releases its slot. + * + * The Boss build seeds a single Fire Pass key shared by every agent. + * Without serialization, the worker fans out three or more concurrent + * `streamSimple` calls against the personal-tier router and trips + * HTTP 429 on every burst. Holding a slot per credential prevents + * the burst at the source instead of recovering from it after the + * fact, which keeps the runtime aligned with the project's + * "let it crash" rule (no retry/backoff fallback layer is added). + * + * The gate is process-local. Only one orchestrator worker runs at a + * time, so a single gate instance correctly serializes every agent + * sharing a credential. + */ +export type ConcurrencyLimitResolver = (provider: string) => number | null; + +export interface CredentialConcurrencyGate { + acquire(provider: string, apiKey: string, signal?: AbortSignal): Promise<() => void>; +} + +interface Waiter { + notify: () => void; + cancel: (err: Error) => void; + signal?: AbortSignal; + abortListener?: () => void; +} + +interface SlotPool { + inFlight: number; + readonly limit: number; + readonly waiters: Waiter[]; +} + +const NOOP_RELEASE = () => undefined; + +function poolKey(provider: string, apiKey: string): string { + // Hash the apiKey so the Map key cannot echo a raw credential into + // heap dumps or future debug output — the gate only needs equality. + const digest = createHash("sha1").update(apiKey).digest("hex").slice(0, 16); + return `${provider}:${digest}`; +} + +function createAbortError(): Error { + const err = new Error("Aborted while waiting for credential concurrency slot."); + err.name = "AbortError"; + return err; +} + +export class InMemoryCredentialConcurrencyGate implements CredentialConcurrencyGate { + private readonly pools = new Map(); + private readonly limitFor: ConcurrencyLimitResolver; + + constructor(limitFor: ConcurrencyLimitResolver) { + this.limitFor = limitFor; + } + + async acquire(provider: string, apiKey: string, signal?: AbortSignal): Promise<() => void> { + const limit = this.limitFor(provider); + if (limit === null) return NOOP_RELEASE; + if (signal?.aborted) throw createAbortError(); + const key = poolKey(provider, apiKey); + const pool = this.getOrCreatePool(key, limit); + if (pool.inFlight < pool.limit) { + pool.inFlight += 1; + return this.makeReleaser(key); + } + await this.waitForSlot(pool, signal); + pool.inFlight += 1; + return this.makeReleaser(key); + } + + private getOrCreatePool(key: string, limit: number): SlotPool { + const existing = this.pools.get(key); + if (existing) return existing; + const pool: SlotPool = { inFlight: 0, limit, waiters: [] }; + this.pools.set(key, pool); + return pool; + } + + private waitForSlot(pool: SlotPool, signal?: AbortSignal): Promise { + return new Promise((resolve, reject) => { + const waiter: Waiter = { + notify: () => { + this.detachAbort(waiter); + resolve(); + }, + cancel: (err) => { + const idx = pool.waiters.indexOf(waiter); + if (idx >= 0) pool.waiters.splice(idx, 1); + this.detachAbort(waiter); + reject(err); + }, + signal, + }; + if (signal) { + const onAbort = () => waiter.cancel(createAbortError()); + waiter.abortListener = onAbort; + signal.addEventListener("abort", onAbort, { once: true }); + } + pool.waiters.push(waiter); + }); + } + + private detachAbort(waiter: Waiter): void { + if (waiter.abortListener && waiter.signal) { + waiter.signal.removeEventListener("abort", waiter.abortListener); + } + } + + private makeReleaser(key: string): () => void { + let released = false; + return () => { + if (released) return; + released = true; + const pool = this.pools.get(key); + if (!pool) return; + pool.inFlight = Math.max(0, pool.inFlight - 1); + const next = pool.waiters.shift(); + if (next) next.notify(); + if (pool.inFlight === 0 && pool.waiters.length === 0) { + this.pools.delete(key); + } + }; + } +} diff --git a/src/main/worker/runtime/pi-safe-stream.ts b/src/main/worker/runtime/pi-safe-stream.ts index a353ad4b..98c1430c 100644 --- a/src/main/worker/runtime/pi-safe-stream.ts +++ b/src/main/worker/runtime/pi-safe-stream.ts @@ -10,6 +10,11 @@ import { type StopReason, type Usage, } from "@mariozechner/pi-ai"; +import { + InMemoryCredentialConcurrencyGate, + type CredentialConcurrencyGate, +} from "./credential-concurrency-gate"; +import { getProviderConcurrencyLimit } from "./provider-rate-limits"; type PiStreamSource = ( model: Model, @@ -67,16 +72,36 @@ function pushStreamError( stream.end(message); } -export function createSafePiStream(source: PiStreamSource): SafePiStreamSource { +/** + * Build a safe stream wrapper. When `gate` is provided, every call + * acquires a slot keyed by `(model.provider, options.apiKey)` before + * the underlying source is invoked, and releases the slot once the + * stream finishes (success, error, or abort). Calls without an + * `apiKey` skip the gate — they target keyless transports such as + * local Ollama / LM Studio. + */ +export function createSafePiStream( + source: PiStreamSource, + gate?: CredentialConcurrencyGate, +): SafePiStreamSource { return (model, context, options) => { const target = createAssistantMessageEventStream(); + const apiKey = typeof options?.apiKey === "string" ? options.apiKey : ""; + const acquireSlot = gate && apiKey + ? () => gate.acquire(model.provider, apiKey, options?.signal) + : () => Promise.resolve<() => void>(() => undefined); void Promise.resolve() - .then(() => source(model, context, options)) - .then(async (sourceStream) => { - for await (const event of sourceStream) { - target.push(event); + .then(acquireSlot) + .then(async (release) => { + try { + const sourceStream = await source(model, context, options); + for await (const event of sourceStream) { + target.push(event); + } + target.end(); + } finally { + release(); } - target.end(); }) .catch((error) => { pushStreamError(target, model, error, options?.signal); @@ -85,4 +110,8 @@ export function createSafePiStream(source: PiStreamSource): SafePiStreamSource { }; } -export const safeStreamSimple = createSafePiStream(streamSimple); +const defaultGate: CredentialConcurrencyGate = new InMemoryCredentialConcurrencyGate( + getProviderConcurrencyLimit, +); + +export const safeStreamSimple = createSafePiStream(streamSimple, defaultGate); diff --git a/src/main/worker/runtime/provider-rate-limits.ts b/src/main/worker/runtime/provider-rate-limits.ts new file mode 100644 index 00000000..4623388c --- /dev/null +++ b/src/main/worker/runtime/provider-rate-limits.ts @@ -0,0 +1,31 @@ +/** + * Per-provider concurrency caps applied to outbound model calls. + * + * Each entry limits the number of in-flight `streamSimple` invocations + * that share a single credential (keyed by `provider` + `apiKey`). + * Providers absent from the map are treated as uncapped — the gate + * short-circuits and the call proceeds immediately. + * + * `fireworks-firepass` is the load-bearing entry. The Fire Pass router + * `accounts/fireworks/routers/kimi-k2p6-turbo` is a personal-tier route: + * the subscription documents zero per-token cost but the upstream + * router rejects bursts on a single key with HTTP 429. In the Boss + * build every agent shares the seeded Fire Pass key, so without a cap + * three concurrent agents trip three 429s — the exact failure mode + * captured in the overview dashboard. A cap of `1` serializes runs on + * the shared key without dropping any work. + * + * Other providers (`anthropic`, `openai`, `google`, …) advertise + * generous tier-based budgets and have no evidence of contention in + * the boss or public builds. Leaving them uncapped honours the + * project's "let it crash" stance — only providers with a known + * upstream constraint are throttled here. + */ +export const PROVIDER_CONCURRENCY_LIMITS: Readonly> = { + "fireworks-firepass": 1, +}; + +export function getProviderConcurrencyLimit(provider: string): number | null { + const value = PROVIDER_CONCURRENCY_LIMITS[provider]; + return typeof value === "number" && value > 0 ? value : null; +} diff --git a/tests/unit/credential-concurrency-gate.test.ts b/tests/unit/credential-concurrency-gate.test.ts new file mode 100644 index 00000000..c1fb1b05 --- /dev/null +++ b/tests/unit/credential-concurrency-gate.test.ts @@ -0,0 +1,106 @@ +import { describe, expect, it } from "vitest"; +import { InMemoryCredentialConcurrencyGate } from "../../src/main/worker/runtime/credential-concurrency-gate"; + +function makeGate(limit: number | null) { + return new InMemoryCredentialConcurrencyGate(() => limit); +} + +describe("InMemoryCredentialConcurrencyGate", () => { + it("returns a no-op release when the provider is uncapped", async () => { + const gate = makeGate(null); + const release = await gate.acquire("anthropic", "key-x"); + expect(typeof release).toBe("function"); + release(); + release(); + }); + + it("hands out slots up to the limit immediately", async () => { + const gate = makeGate(2); + const r1 = await gate.acquire("fireworks-firepass", "shared"); + const r2 = await gate.acquire("fireworks-firepass", "shared"); + expect(r1).not.toBe(r2); + r1(); + r2(); + }); + + it("queues the third caller until a slot is released", async () => { + const gate = makeGate(1); + const release = await gate.acquire("fireworks-firepass", "shared"); + const order: string[] = []; + + const second = (async () => { + const r = await gate.acquire("fireworks-firepass", "shared"); + order.push("second"); + r(); + })(); + const third = (async () => { + const r = await gate.acquire("fireworks-firepass", "shared"); + order.push("third"); + r(); + })(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + expect(order).toEqual([]); + + release(); + await second; + expect(order).toEqual(["second"]); + + await third; + expect(order).toEqual(["second", "third"]); + }); + + it("scopes pools by (provider, apiKey) so different keys do not block each other", async () => { + const gate = makeGate(1); + const r1 = await gate.acquire("fireworks-firepass", "key-a"); + // Different apiKey -> new pool, no waiting. + const r2 = await gate.acquire("fireworks-firepass", "key-b"); + expect(typeof r1).toBe("function"); + expect(typeof r2).toBe("function"); + r1(); + r2(); + }); + + it("releases the pool when no callers remain", async () => { + const gate = makeGate(1); + const release = await gate.acquire("fireworks-firepass", "shared"); + release(); + // Re-acquire should not be queued. + const next = await gate.acquire("fireworks-firepass", "shared"); + next(); + }); + + it("rejects when the abort signal fires while waiting", async () => { + const gate = makeGate(1); + const held = await gate.acquire("fireworks-firepass", "shared"); + + const controller = new AbortController(); + const pending = gate.acquire("fireworks-firepass", "shared", controller.signal); + controller.abort(); + + await expect(pending).rejects.toMatchObject({ name: "AbortError" }); + + // Slot must remain available after the aborted waiter is removed. + held(); + const next = await gate.acquire("fireworks-firepass", "shared"); + next(); + }); + + it("rejects immediately when the abort signal is already aborted", async () => { + const gate = makeGate(1); + const controller = new AbortController(); + controller.abort(); + await expect( + gate.acquire("fireworks-firepass", "shared", controller.signal), + ).rejects.toMatchObject({ name: "AbortError" }); + }); + + it("treats double-release as a no-op", async () => { + const gate = makeGate(1); + const release = await gate.acquire("fireworks-firepass", "shared"); + release(); + release(); + const next = await gate.acquire("fireworks-firepass", "shared"); + next(); + }); +}); diff --git a/tests/unit/pi-safe-stream.test.ts b/tests/unit/pi-safe-stream.test.ts index 5549fe65..a689cc78 100644 --- a/tests/unit/pi-safe-stream.test.ts +++ b/tests/unit/pi-safe-stream.test.ts @@ -7,6 +7,7 @@ import { } from "@mariozechner/pi-ai"; import { describe, expect, it } from "vitest"; import { createSafePiStream } from "../../src/main/worker/runtime/pi-safe-stream"; +import { InMemoryCredentialConcurrencyGate } from "../../src/main/worker/runtime/credential-concurrency-gate"; function createModel(): Model<"openai-responses"> { return { @@ -67,6 +68,102 @@ describe("pi safe stream", () => { await expect(stream.result()).resolves.toMatchObject({ stopReason: "stop" }); }); + it("serializes calls that share a gated credential", async () => { + const gate = new InMemoryCredentialConcurrencyGate(() => 1); + let activeCalls = 0; + let peakActive = 0; + + const safeStream = createSafePiStream((_model, _context) => { + activeCalls += 1; + peakActive = Math.max(peakActive, activeCalls); + const stream = createAssistantMessageEventStream(); + setTimeout(() => { + const message = createAssistantMessage(); + stream.push({ type: "done", reason: "stop", message }); + stream.end(message); + activeCalls -= 1; + }, 5); + return stream; + }, gate); + + const drain = async (stream: ReturnType) => { + for await (const _event of stream) { + // drain + } + await stream.result(); + }; + + const model = createModel(); + const a = safeStream(model, emptyContext, { apiKey: "shared" }); + const b = safeStream(model, emptyContext, { apiKey: "shared" }); + const c = safeStream(model, emptyContext, { apiKey: "shared" }); + + await Promise.all([drain(a), drain(b), drain(c)]); + expect(peakActive).toBe(1); + }); + + it("does not gate calls without an apiKey", async () => { + const gate = new InMemoryCredentialConcurrencyGate(() => 1); + let activeCalls = 0; + let peakActive = 0; + + const safeStream = createSafePiStream((_model, _context) => { + activeCalls += 1; + peakActive = Math.max(peakActive, activeCalls); + const stream = createAssistantMessageEventStream(); + setTimeout(() => { + const message = createAssistantMessage(); + stream.push({ type: "done", reason: "stop", message }); + stream.end(message); + activeCalls -= 1; + }, 5); + return stream; + }, gate); + + const drain = async (stream: ReturnType) => { + for await (const _event of stream) { + // drain + } + await stream.result(); + }; + + const model = createModel(); + await Promise.all([ + drain(safeStream(model, emptyContext)), + drain(safeStream(model, emptyContext)), + ]); + expect(peakActive).toBe(2); + }); + + it("releases the slot when the upstream stream throws", async () => { + const gate = new InMemoryCredentialConcurrencyGate(() => 1); + let attempt = 0; + + const safeStream = createSafePiStream(() => { + attempt += 1; + if (attempt === 1) throw new Error("first call boom"); + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + const message = createAssistantMessage(); + stream.push({ type: "done", reason: "stop", message }); + stream.end(message); + }); + return stream; + }, gate); + + const failing = safeStream(createModel(), emptyContext, { apiKey: "shared" }); + for await (const _event of failing) { + // drain + } + await failing.result(); + + const succeeding = safeStream(createModel(), emptyContext, { apiKey: "shared" }); + for await (const _event of succeeding) { + // drain + } + await expect(succeeding.result()).resolves.toMatchObject({ stopReason: "stop" }); + }); + it("turns rejected stream setup into an error event", async () => { const safeStream = createSafePiStream(async () => { throw new Error("stream setup failed"); diff --git a/tests/unit/provider-rate-limits.test.ts b/tests/unit/provider-rate-limits.test.ts new file mode 100644 index 00000000..36608a57 --- /dev/null +++ b/tests/unit/provider-rate-limits.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, it } from "vitest"; +import { + PROVIDER_CONCURRENCY_LIMITS, + getProviderConcurrencyLimit, +} from "../../src/main/worker/runtime/provider-rate-limits"; + +describe("provider rate limits", () => { + it("caps fireworks-firepass at one concurrent call", () => { + expect(PROVIDER_CONCURRENCY_LIMITS["fireworks-firepass"]).toBe(1); + expect(getProviderConcurrencyLimit("fireworks-firepass")).toBe(1); + }); + + it("returns null for providers without an explicit cap", () => { + expect(getProviderConcurrencyLimit("anthropic")).toBeNull(); + expect(getProviderConcurrencyLimit("openai")).toBeNull(); + expect(getProviderConcurrencyLimit("ollama")).toBeNull(); + expect(getProviderConcurrencyLimit("unknown-provider")).toBeNull(); + }); +});