new_pull

packages/ai/test/openai-completions-retry.test.ts

@@ -1,9 +1,10 @@
-import { beforeEach, describe, expect, it, vi } from "vitest";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { stream as streamOpenAICompletions } from "../src/api/openai-completions.ts";
 import type { Context, Model } from "../src/types.ts";
 
 const mockState = vi.hoisted(() => ({
 	requestOptions: [] as unknown[],
+	requestErrors: [] as Error[],
 }));
 
 vi.mock("openai", () => {
@@ -30,10 +31,14 @@ vi.mock("openai", () => {
 							response: { status: number; headers: Headers };
 						}>;
 					};
-					promise.withResponse = async () => ({
-						data: stream,
-						response: { status: 200, headers: new Headers() },
-					});
+					promise.withResponse = async () => {
+						const error = mockState.requestErrors.shift();
+						if (error) throw error;
+						return {
+							data: stream,
+							response: { status: 200, headers: new Headers() },
+						};
+					};
 					return promise;
 				},
 			},
@@ -61,7 +66,7 @@ const context: Context = {
 	tools: [],
 };
 
-async function consume(options?: { maxRetries?: number }) {
+async function consume(options?: { maxRetries?: number; maxRetryDelayMs?: number }) {
 	const stream = streamOpenAICompletions(model, context, { apiKey: "test", ...options });
 	for await (const _event of stream) {
 		void _event;
@@ -72,6 +77,11 @@ async function consume(options?: { maxRetries?: number }) {
 describe("openai-completions provider retries", () => {
 	beforeEach(() => {
 		mockState.requestOptions = [];
+		mockState.requestErrors = [];
+	});
+
+	afterEach(() => {
+		vi.useRealTimers();
 	});
 
 	it("disables SDK retries by default", async () => {
@@ -79,8 +89,51 @@ describe("openai-completions provider retries", () => {
 		expect(mockState.requestOptions).toEqual([expect.objectContaining({ maxRetries: 0 })]);
 	});
 
-	it("honors explicit provider retry settings", async () => {
-		await consume({ maxRetries: 2 });
-		expect(mockState.requestOptions).toEqual([expect.objectContaining({ maxRetries: 2 })]);
+	it("honors provider retries while keeping SDK retries disabled", async () => {
+		vi.useFakeTimers();
+		mockState.requestErrors = [
+			Object.assign(new Error("rate limited"), {
+				status: 429,
+				headers: new Headers({ "retry-after-ms": "100" }),
+			}),
+			Object.assign(new Error("server error"), {
+				status: 500,
+				headers: new Headers({ "retry-after-ms": "100" }),
+			}),
+		];
+
+		const result = consume({ maxRetries: 2, maxRetryDelayMs: 100 });
+		await vi.advanceTimersByTimeAsync(0);
+		expect(mockState.requestOptions).toHaveLength(1);
+		await vi.advanceTimersByTimeAsync(99);
+		expect(mockState.requestOptions).toHaveLength(1);
+		await vi.advanceTimersByTimeAsync(1);
+		expect(mockState.requestOptions).toHaveLength(2);
+		await vi.advanceTimersByTimeAsync(99);
+		expect(mockState.requestOptions).toHaveLength(2);
+		await vi.advanceTimersByTimeAsync(1);
+		await result;
+
+		expect(mockState.requestOptions).toEqual([
+			expect.objectContaining({ maxRetries: 0 }),
+			expect.objectContaining({ maxRetries: 0 }),
+			expect.objectContaining({ maxRetries: 0 }),
+		]);
+	});
+
+	it("fails immediately when a provider-requested retry delay exceeds the limit", async () => {
+		mockState.requestErrors = [
+			Object.assign(new Error("rate limited"), {
+				status: 429,
+				headers: new Headers({ "retry-after": "277403" }),
+			}),
+		];
+
+		const result = await consume({ maxRetries: 2, maxRetryDelayMs: 1000 });
+
+		expect(result.stopReason).toBe("error");
+		expect(result.errorMessage).toContain("Server requested 277403s retry delay (max: 1s)");
+		expect(result.errorMessage).toContain("rate limited");
+		expect(mockState.requestOptions).toEqual([expect.objectContaining({ maxRetries: 0 })]);
 	});
 });