fix(runtime): serialize Fire Pass calls per credential to stop boss-build 429s
Boss build seeds one Fire Pass key for every agent. Three concurrent agents -> three concurrent streamSimple calls against the Kimi K2.6 Turbo router -> three HTTP 429s on the overview dashboard. Add an in-memory async semaphore keyed by (provider, sha1(apiKey)) and acquire a slot inside `createSafePiStream` before invoking the source. The slot is released when the stream finishes, errors, or aborts. `fireworks-firepass` is capped at 1 concurrent call; every other provider stays uncapped (no fallback layer, no 429 retry — flow control prevents the burst at the source instead of recovering after). Same code path applies to public and boss builds: only the seeded key differs, the gate behaviour is identical.
This commit is contained in:
1 parent
a532120649
commit
fccb0d8586
6 files changed
+421
-7
No files matched your search
@@ -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<typeof safeStream>) => {
|
||||
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<typeof safeStream>) => {
|
||||
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");
|
||||
|
||||
Reference in new issue
Block a user