From 4f91a9b00b4fd06b87978986dbaa1748186968ff Mon Sep 17 00:00:00 2001 From: Patrick Wozniak Date: Thu, 20 Aug 2026 00:13:56 +0200 Subject: [PATCH] fix(stream): isolate transport state by credential --- src/transport.ts | 5 ++-- tests/test-transport.ts | 64 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/src/transport.ts b/src/transport.ts index 9d83d5d..c30171c 100644 --- a/src/transport.ts +++ b/src/transport.ts @@ -71,6 +71,7 @@ export function createCommandCodeTransportRouter(deps: TransportDependencies) { apiKey = options?.apiKey transport = "unknown" } + const requestApiKey = options?.apiKey if (transport === "generate") return deps.streamGenerate(model, context, options) const output = deps.createStream() @@ -94,13 +95,13 @@ export function createCommandCodeTransportRouter(deps: TransportDependencies) { for await (const event of providerStream) { if (!upgradeRequired) { - transport = "provider" + if (apiKey === requestApiKey) transport = "provider" output.push(event) } } if (upgradeRequired) { - transport = "generate" + if (apiKey === requestApiKey) transport = "generate" await pipe(deps.streamGenerate(model, context, options), output) } output.end() diff --git a/tests/test-transport.ts b/tests/test-transport.ts index 59ea51b..82e9771 100644 --- a/tests/test-transport.ts +++ b/tests/test-transport.ts @@ -159,6 +159,70 @@ describe("Command Code transport router", () => { assert.equal(generateCalls, 1) }) + it("does not let a stale request overwrite the transport for a new API key", async () => { + let releaseGoRequest: (() => void) | undefined + const goRequestGate = new Promise((resolve) => { + releaseGoRequest = resolve + }) + let providerCalls = 0 + let generateCalls = 0 + const upgradeBody = JSON.stringify({ error: { code: "upgrade_required" } }) + const router = createCommandCodeTransportRouter({ + createStream: createTestEventStream, + streamProvider: (_model, _context, options) => { + providerCalls += 1 + const response = + options?.apiKey === "go-key" + ? new Response(upgradeBody, { status: 403 }) + : new Response("ok", { status: 200 }) + const stream = createTestEventStream() + const run = async () => { + if (options?.apiKey === "go-key") await goRequestGate + const received = await (options?.fetch ?? fetch)("https://provider.test", {}) + await options?.onResponse?.( + { status: received.status, headers: {} }, + makeModel({ api: "openai-completions" }), + ) + if (response.ok) { + for await (const event of completedStream("provider")) stream.push(event) + } + stream.end() + } + run().catch(() => stream.end()) + return stream + }, + streamGenerate: () => { + generateCalls += 1 + return completedStream("generate") + }, + }) + + const staleGoRequest = collectEvents( + router.stream(makeModel(), makeContext(), { + apiKey: "go-key", + fetch: () => Promise.resolve(new Response(upgradeBody, { status: 403 })), + }), + ) + await collectEvents( + router.stream(makeModel(), makeContext(), { + apiKey: "provider-key", + fetch: () => Promise.resolve(new Response("ok", { status: 200 })), + }), + ) + releaseGoRequest?.() + await staleGoRequest + await collectEvents( + router.stream(makeModel(), makeContext(), { + apiKey: "provider-key", + fetch: () => Promise.resolve(new Response("ok", { status: 200 })), + }), + ) + + assert.equal(router.getTransport(), "provider") + assert.equal(providerCalls, 3) + assert.equal(generateCalls, 1) + }) + it("does not fall back for other 403 errors", async () => { let generateCalls = 0 const responseBody = JSON.stringify({ error: { code: "permission_denied" } })