diff --git a/.changeset/gateway-concurrent-fetch-isolation.md b/.changeset/gateway-concurrent-fetch-isolation.md new file mode 100644 index 000000000..551d316fd --- /dev/null +++ b/.changeset/gateway-concurrent-fetch-isolation.md @@ -0,0 +1,5 @@ +--- +"ai-gateway-provider": patch +--- + +Isolate per-request `fetch` configuration with a Proxy so concurrent AI Gateway calls cannot cross-wire shared `model.config.fetch`. diff --git a/packages/ai-gateway-provider/src/index.ts b/packages/ai-gateway-provider/src/index.ts index 22cd35363..64ab433f3 100644 --- a/packages/ai-gateway-provider/src/index.ts +++ b/packages/ai-gateway-provider/src/index.ts @@ -46,6 +46,29 @@ type InternalLanguageModelV4 = LanguageModelV4 & { config?: { fetch?: FetchFunction | undefined }; }; +/** + * Return a per-request view of `model` whose `config.fetch` is isolated. + * Mutating the shared `model.config.fetch` races under concurrent calls and + * can cross-wire gateway responses (#620). + */ +function modelWithFetch( + model: InternalLanguageModelV4, + fetchFn: FetchFunction, +): InternalLanguageModelV4 { + return new Proxy(model, { + get(target, prop, receiver) { + if (prop === "config") { + return { ...target.config, fetch: fetchFn }; + } + const value = Reflect.get(target, prop, receiver); + if (typeof value === "function") { + return value.bind(receiver); + } + return value; + }, + }); +} + export class AiGatewayChatLanguageModel implements LanguageModelV4 { readonly specificationVersion = "v4"; readonly defaultObjectGenerationMode = "json"; @@ -94,17 +117,17 @@ export class AiGatewayChatLanguageModel implements LanguageModelV4 { ); } - model.config.fetch = (url, request) => { + const captureModel = modelWithFetch(model, (url, request) => { requests.push({ modelProvider: model.provider, request: request as Request, url: url as string, }); throw new AiGatewayInternalFetchError("Stopping provider execution..."); - }; + }); try { - await model[modelMethod](options); + await captureModel[modelMethod](options); } catch (e) { if (!(e instanceof AiGatewayInternalFetchError)) { throw e; @@ -236,12 +259,12 @@ export class AiGatewayChatLanguageModel implements LanguageModelV4 { } } - this.models[step].config = { - ...this.models[step].config, - fetch: (_url, _req) => resp as unknown as Promise, - }; + const responseModel = modelWithFetch( + this.models[step], + (_url, _req) => resp as unknown as Promise, + ); - return this.models[step][modelMethod](options) as Promise>>; + return responseModel[modelMethod](options) as Promise>>; } async doStream( diff --git a/packages/ai-gateway-provider/test/concurrency.test.ts b/packages/ai-gateway-provider/test/concurrency.test.ts new file mode 100644 index 000000000..94d2bbd49 --- /dev/null +++ b/packages/ai-gateway-provider/test/concurrency.test.ts @@ -0,0 +1,80 @@ +import { createDeepSeek } from "../src/providers/deepseek"; +import { generateText } from "ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { createAiGateway } from "../src"; + +/** + * Regression for https://github.com/cloudflare/ai/issues/620 + * + * Concurrent calls on a shared gateway-wrapped model must not cross-wire + * request bodies / responses via a mutated shared `model.config.fetch`. + */ +describe("Concurrent request isolation", () => { + const originalFetch = globalThis.fetch; + + afterEach(() => { + globalThis.fetch = originalFetch; + vi.restoreAllMocks(); + }); + + it("keeps concurrent generateText calls isolated on a shared model", async () => { + const gatewayBodies: unknown[] = []; + + globalThis.fetch = (async (_url, init) => { + const body = JSON.parse(String(init?.body ?? "null")); + const prompts = Array.isArray(body) + ? body.map( + (request: { query?: { messages?: { content?: string }[] } }) => + request.query?.messages?.at(-1)?.content, + ) + : []; + gatewayBodies.push(prompts); + + const prompt = prompts[0]; + return new Response( + JSON.stringify({ + id: "id", + object: "chat.completion", + created: 0, + model: "deepseek-chat", + choices: [ + { + index: 0, + message: { role: "assistant", content: `ANSWER:${prompt}` }, + finish_reason: "stop", + }, + ], + usage: { + prompt_tokens: 1, + completion_tokens: 1, + total_tokens: 2, + }, + }), + { + headers: { + "content-type": "application/json", + "cf-aig-step": "0", + }, + }, + ); + }) as typeof fetch; + + const gateway = createAiGateway({ + accountId: "test", + gateway: "test", + apiKey: "test", + }); + const deepseek = createDeepSeek(); + const model = gateway(deepseek("deepseek-chat")); + + const ask = (prompt: string) => + generateText({ model, prompt, maxRetries: 0 }).then((result) => result.text); + + const [a, b] = await Promise.all([ask("A"), ask("B")]); + + expect(a).toBe("ANSWER:A"); + expect(b).toBe("ANSWER:B"); + expect(gatewayBodies).toHaveLength(2); + expect(gatewayBodies).toEqual(expect.arrayContaining([["A"], ["B"]])); + }); +});