fix(ai): retry incomplete Codex tool calls once

packages/ai/changelog.d/codex-timeout-incomplete-toolcall-retry.md

@@ -0,0 +1,3 @@
+### Fixed
+
+- Retry Codex Responses streams once when a request timeout leaves only incomplete tool-call output.

packages/ai/src/providers/openai-codex-responses.ts

@@ -489,6 +489,7 @@ interface CodexStreamRuntime {
 	partialToolCallReplayAttempted: boolean;
 	bufferedEvents: AssistantMessageEvent[];
 	eventsReleased: boolean;
+	incompleteToolCallRetryAttempted: boolean;
 	toolChoiceFallbackAttempted: boolean;
 	/**
 	 * Stale-anchor recovery is one-shot. Once the anchor is cleared the replay no
@@ -1258,6 +1259,7 @@ function createCodexStreamRuntime(initial: {
 		partialToolCallReplayAttempted: false,
 		bufferedEvents: [],
 		eventsReleased: false,
+		incompleteToolCallRetryAttempted: false,
 		toolChoiceFallbackAttempted: initial.toolChoiceFallbackApplied === true,
 		previousResponseRecoveryAttempted: false,
 		sentPreviousResponseId: initial.sentPreviousResponseId === true,
@@ -2349,12 +2351,51 @@ async function recoverCodexStreamError(
 	if (await tryReplayWebsocketFailureOverSse(context, runtime, error)) {
 		return true;
 	}
+	if (isCodexIncompleteToolCallTimeoutRetryable(context, runtime, error)) {
+		runtime.incompleteToolCallRetryAttempted = true;
+		const retried = await tryRetryCodexProviderError(context, runtime, error, true);
+		if (retried) return true;
+	}
 	if (await tryRetryCodexProviderError(context, runtime, error)) {
 		return true;
 	}
 	return false;
 }
 
+function isCodexIncompleteToolCallTimeoutRetryable(
+	context: CodexStreamProcessingContext,
+	runtime: CodexStreamRuntime,
+	error: unknown,
+): boolean {
+	if (
+		runtime.incompleteToolCallRetryAttempted ||
+		runtime.partialToolCallReplayAttempted ||
+		runtime.eventsReleased ||
+		!isCodexTransientStreamClose(error)
+	) {
+		return false;
+	}
+	const providerCode =
+		error instanceof CodexProviderStreamError
+			? error.code?.toLowerCase()
+			: (error as { providerCode?: string }).providerCode?.toLowerCase();
+	if (providerCode !== "request_timeout") return false;
+	const toolCalls = context.output.content.filter((block): block is ToolCall => block.type === "toolCall");
+	return (
+		toolCalls.length > 0 &&
+		runtime.finalizedToolCallIds.size === 0 &&
+		!runtime.toolArgumentCorrelationFailed &&
+		context.output.content.every(
+			block =>
+				block.type === "thinking" ||
+				(block.type === "toolCall" &&
+					!runtime.finalizedToolCallIds.has(block.id) &&
+					(block as ToolCall & { argumentsComplete?: boolean }).argumentsComplete !== true) ||
+				(block.type === "text" && block.text === ""),
+		)
+	);
+}
+
 async function tryRetryWithoutForcedToolChoice(
 	context: CodexStreamProcessingContext,
 	runtime: CodexStreamRuntime,
@@ -2634,6 +2675,7 @@ async function tryRetryCodexProviderError(
 	context: CodexStreamProcessingContext,
 	runtime: CodexStreamRuntime,
 	error: unknown,
+	allowIncompleteToolCallOutput = false,
 ): Promise<boolean> {
 	const hasUnfinalizedToolCall = context.output.content.some(
 		block => block.type === "toolCall" && !runtime.finalizedToolCallIds.has(block.id),
@@ -2644,9 +2686,12 @@ async function tryRetryCodexProviderError(
 		hasUnfinalizedToolCall &&
 		!runtime.partialToolCallReplayAttempted &&
 		isCodexPartialOutputReplayEligible(context.output, runtime);
+	const canRetryWithOutput =
+		canReplayPartialToolCall ||
+		(allowIncompleteToolCallOutput && !runtime.partialToolCallReplayAttempted && !runtime.eventsReleased);
 	if (
-		(!isRetryableCodexProviderError(error) && !canReplayPartialToolCall) ||
-		(context.output.content.length > 0 && !canReplayPartialToolCall) ||
+		(!isRetryableCodexProviderError(error) && !canRetryWithOutput) ||
+		(context.output.content.length > 0 && !canRetryWithOutput) ||
 		runtime.providerRetryAttempt >= resolveRetryBudget(context.options?.streamMaxRetries, CODEX_MAX_RETRIES) ||
 		context.options?.signal?.aborted ||
 		context.options?.fallbackManaged ||
@@ -2656,7 +2701,7 @@ async function tryRetryCodexProviderError(
 	}
 
 	runtime.providerRetryAttempt += 1;
-	if (context.output.content.length > 0 && canReplayPartialToolCall) runtime.partialToolCallReplayAttempted = true;
+	if (context.output.content.length > 0 && canRetryWithOutput) runtime.partialToolCallReplayAttempted = true;
 	const websocketState = context.requestContext.websocketState;
 	if (runtime.transport === "websocket" && websocketState) {
 		resetCodexWebSocketAppendState(websocketState);
@@ -2677,7 +2722,7 @@ async function tryRetryCodexProviderError(
 	runtime.nativeOutputItemOutputIndexes.length = 0;
 	runtime.finalizedToolCallIds.clear();
 	runtime.toolArgumentCorrelationFailed = false;
-	if (canReplayPartialToolCall) discardCodexEvents(runtime, true);
+	if (canRetryWithOutput) discardCodexEvents(runtime, true);
 	resetOutputState(context.output);
 	context.firstTokenTime = undefined;
 	await scheduler.wait(CODEX_RETRY_DELAY_MS * runtime.providerRetryAttempt, {

packages/ai/test/openai-codex-stream.test.ts

@@ -789,6 +789,84 @@ describe("openai-codex streaming", () => {
 		]);
 	});
 
+	it("does not replay a released partial call after a timeout retry", async () => {
+		const firstPartial = [
+			{
+				type: "response.output_item.added",
+				item: { type: "function_call", id: "fc_1", call_id: "call_1", name: "old_tool", arguments: "" },
+			},
+			{ type: "response.function_call_arguments.delta", item_id: "fc_1", delta: '{"old":' },
+		];
+		const secondTimeout = createCodexErrorSse([
+			{
+				type: "response.output_item.added",
+				item: { type: "function_call", id: "fc_2", call_id: "call_2", name: "old_tool", arguments: "" },
+			},
+			{ type: "response.function_call_arguments.delta", item_id: "fc_2", delta: '{"old":' },
+			{ type: "error", code: "request_timeout", message: "stream closed before response.completed" },
+		]);
+		let requestCount = 0;
+		global.fetch = vi.fn(async (_input, _init) => {
+			requestCount += 1;
+			if (requestCount === 1) return createSocketCloseCodexSse(undefined, firstPartial);
+			return new Response(secondTimeout, { status: 200, headers: { "content-type": "text/event-stream" } });
+		}) as unknown as typeof fetch;
+		vi.spyOn(scheduler, "wait").mockResolvedValue(undefined);
+
+		const stream = streamOpenAICodexResponses(
+			{ ...createCodexTestModel("https://chatgpt.com/backend-api"), preferWebsockets: false },
+			createCodexTestContext(),
+			{ apiKey: createCodexTestToken() },
+		);
+		const events: AssistantMessageEvent[] = [];
+		for await (const event of stream) events.push(event);
+		const result = await stream.result();
+
+		expect(requestCount).toBe(2);
+		expect(result.stopReason).toBe("error");
+		expect(events.filter(event => event.type === "toolcall_start")).toHaveLength(1);
+	});
+
+	it("does not replay a partial call after timeout retry consumes the replay", async () => {
+		const firstTimeout = createCodexErrorSse([
+			{
+				type: "response.output_item.added",
+				item: { type: "function_call", id: "fc_1", call_id: "call_1", name: "old_tool", arguments: "" },
+			},
+			{ type: "response.function_call_arguments.delta", item_id: "fc_1", delta: '{"old":' },
+			{ type: "error", code: "request_timeout", message: "stream closed before response.completed" },
+		]);
+		const secondPartial = [
+			{
+				type: "response.output_item.added",
+				item: { type: "function_call", id: "fc_2", call_id: "call_2", name: "old_tool", arguments: "" },
+			},
+			{ type: "response.function_call_arguments.delta", item_id: "fc_2", delta: '{"old":' },
+		];
+		let requestCount = 0;
+		global.fetch = vi.fn(async (_input, _init) => {
+			requestCount += 1;
+			if (requestCount === 1) {
+				return new Response(firstTimeout, { status: 200, headers: { "content-type": "text/event-stream" } });
+			}
+			return createSocketCloseCodexSse(undefined, secondPartial);
+		}) as unknown as typeof fetch;
+		vi.spyOn(scheduler, "wait").mockResolvedValue(undefined);
+
+		const stream = streamOpenAICodexResponses(
+			{ ...createCodexTestModel("https://chatgpt.com/backend-api"), preferWebsockets: false },
+			createCodexTestContext(),
+			{ apiKey: createCodexTestToken() },
+		);
+		const events: AssistantMessageEvent[] = [];
+		for await (const event of stream) events.push(event);
+		const result = await stream.result();
+
+		expect(requestCount).toBe(2);
+		expect(result.stopReason).toBe("error");
+		expect(events.filter(event => event.type === "toolcall_start")).toHaveLength(1);
+	});
+
 	it("does not replay an ECONNRESET socket close after visible text", async () => {
 		const partial = [
 			{