fix(ai-gateway-provider): isolate per-request fetch via Proxy

.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`.

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<Response>,
-		};
+		const responseModel = modelWithFetch(
+			this.models[step],
+			(_url, _req) => resp as unknown as Promise<Response>,
+		);
 
-		return this.models[step][modelMethod](options) as Promise<Awaited<ReturnType<T>>>;
+		return responseModel[modelMethod](options) as Promise<Awaited<ReturnType<T>>>;
 	}
 
 	async doStream(

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"]]));
+	});
+});