Merge pull request #1073 from ephraimduncan/fix/stream-stall-auto-retry

fix(ai): recover stalled lazy provider streams

packages/ai/CHANGELOG.md

@@ -25,6 +25,11 @@
 
 - Fixed OAuth credentials being silently disabled when two omp processes (or any two `AuthStorage` instances sharing a `agent.db`) race on token refresh. Anthropic rotates refresh tokens on every use, so the loser's `invalid_grant` response previously soft-deleted the row that the winner just rotated, forcing the user to `/login` again. `#tryOAuthCredential` now re-reads the row from disk before declaring a definitive failure: if the persisted `refresh` differs from the snapshot it tried, the peer-rotated credential is reloaded and the request retries against the fresh token instead of disabling the live row.
 - Closed a remaining race window in OAuth refresh-failure handling: between re-reading the credential row to check for peer rotation and the subsequent soft-delete, another process could still complete a refresh and rotate the row, leaving us to disable the freshly-rotated credential by `id`. The disable now runs as a single CAS update conditioned on the row's `data` still matching the snapshot we tried to refresh, and on `disabled_cause IS NULL`. If the CAS reports 0 rows changed (peer rotation, or row already disabled by a concurrent failure on the same snapshot), we reload from disk and retry instead of mutating the wrong row or emitting a spurious `credential_disabled` event.
+### Changed
+- Lowered the default steady-state stream idle timeout from 120s to 30s while preserving the existing environment overrides.
+
+### Fixed
+- Lazy built-in provider streams now enforce the shared idle watchdog and abort stalled provider requests, so session auto-retry can continue after transient network drops instead of remaining stuck. Caller aborts still terminate as aborted.
 
 ## [14.9.3] - 2026-05-10
 

packages/ai/src/providers/register-builtins.ts

@@ -19,7 +19,9 @@ import type {
 	Model,
 	OptionsForApi,
 } from "../types";
+import { type AbortSourceTracker, createAbortSourceTracker } from "../utils/abort";
 import { AssistantMessageEventStream as EventStreamImpl } from "../utils/event-stream";
+import { getStreamFirstEventTimeoutMs, getStreamIdleTimeoutMs, iterateWithIdleTimeout } from "../utils/idle-iterator";
 import type { BedrockOptions } from "./amazon-bedrock";
 import type { AnthropicOptions } from "./anthropic";
 import type { AzureOpenAIResponsesOptions } from "./azure-openai-responses";
@@ -155,6 +157,9 @@ export function setBedrockProviderModule(module: BedrockProviderModule): void {
 // Stream forwarding / error helpers
 // ---------------------------------------------------------------------------
 
+const LAZY_STREAM_IDLE_TIMEOUT_ERROR = "Provider stream stalled while waiting for the next event";
+const LAZY_STREAM_FIRST_EVENT_TIMEOUT_ERROR = "Provider stream timed out while waiting for the first event";
+
 function hasFinalResult(
 	source: AsyncIterable<AssistantMessageEvent>,
 ): source is AsyncIterable<AssistantMessageEvent> & { result(): Promise<AssistantMessage> } {
@@ -165,10 +170,23 @@ function forwardStream<TApi extends Api>(
 	target: EventStreamImpl,
 	source: AsyncIterable<AssistantMessageEvent>,
 	model: Model<TApi>,
+	options: OptionsForApi<TApi>,
+	abortTracker: AbortSourceTracker,
 ): void {
 	(async () => {
 		try {
-			for await (const event of source) {
+			const idleTimeoutMs = options.streamIdleTimeoutMs ?? getStreamIdleTimeoutMs();
+			const watchedSource = iterateWithIdleTimeout(source, {
+				idleTimeoutMs,
+				firstItemTimeoutMs: options.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(idleTimeoutMs),
+				errorMessage: LAZY_STREAM_IDLE_TIMEOUT_ERROR,
+				firstItemErrorMessage: LAZY_STREAM_FIRST_EVENT_TIMEOUT_ERROR,
+				onIdle: () => abortTracker.abortLocally(new Error(LAZY_STREAM_IDLE_TIMEOUT_ERROR)),
+				onFirstItemTimeout: () => abortTracker.abortLocally(new Error(LAZY_STREAM_FIRST_EVENT_TIMEOUT_ERROR)),
+				abortSignal: options.signal,
+			});
+
+			for await (const event of watchedSource) {
 				target.push(event);
 			}
 			if (hasFinalResult(source)) {
@@ -177,14 +195,19 @@ function forwardStream<TApi extends Api>(
 				target.end();
 			}
 		} catch (error) {
-			const message = createLazyLoadErrorMessage(model, error);
-			target.push({ type: "error", reason: "error", error: message });
+			const stopReason = abortTracker.wasCallerAbort() ? "aborted" : "error";
+			const message = createLazyLoadErrorMessage(model, error, stopReason);
+			target.push({ type: "error", reason: stopReason, error: message });
 			target.end(message);
 		}
 	})();
 }
 
-function createLazyLoadErrorMessage<TApi extends Api>(model: Model<TApi>, error: unknown): AssistantMessage {
+function createLazyLoadErrorMessage<TApi extends Api>(
+	model: Model<TApi>,
+	error: unknown,
+	stopReason: Extract<AssistantMessage["stopReason"], "aborted" | "error"> = "error",
+): AssistantMessage {
 	return {
 		role: "assistant",
 		content: [],
@@ -199,8 +222,9 @@ function createLazyLoadErrorMessage<TApi extends Api>(model: Model<TApi>, error:
 			totalTokens: 0,
 			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
 		},
-		stopReason: "error",
-		errorMessage: error instanceof Error ? error.message : String(error),
+		stopReason,
+		errorMessage:
+			stopReason === "aborted" ? "Request was aborted" : error instanceof Error ? error.message : String(error),
 		timestamp: Date.now(),
 	};
 }
@@ -214,11 +238,14 @@ function createLazyStream<TApi extends Api>(
 ): (model: Model<TApi>, context: Context, options: OptionsForApi<TApi>) => EventStreamImpl {
 	return (model, context, options) => {
 		const outer = new EventStreamImpl();
+		const streamOptions = (options ?? {}) as OptionsForApi<TApi>;
 
 		loadModule()
 			.then(module => {
-				const inner = module.stream(model, context, options);
-				forwardStream(outer, inner, model);
+				const abortTracker = createAbortSourceTracker(streamOptions.signal);
+				const providerOptions = { ...streamOptions, signal: abortTracker.requestSignal } as OptionsForApi<TApi>;
+				const inner = module.stream(model, context, providerOptions);
+				forwardStream(outer, inner, model, streamOptions, abortTracker);
 			})
 			.catch(error => {
 				const message = createLazyLoadErrorMessage(model, error);

packages/ai/src/utils/idle-iterator.ts

@@ -1,6 +1,6 @@
 import { $env } from "@oh-my-pi/pi-utils";
 
-const DEFAULT_STREAM_IDLE_TIMEOUT_MS = 120_000;
+const DEFAULT_STREAM_IDLE_TIMEOUT_MS = 30_000;
 const DEFAULT_STREAM_FIRST_EVENT_TIMEOUT_MS = 100_000;
 
 function normalizeIdleTimeoutMs(value: string | undefined, fallback: number): number | undefined {

packages/ai/test/register-builtins.test.ts

@@ -92,4 +92,84 @@ describe("register-builtins lazy streams", () => {
 		expect(result.stopReason).toBe("error");
 		expect(result.errorMessage).toContain("bedrock exploded");
 	});
+
+	it("turns idle lazy provider streams into retryable terminal errors", async () => {
+		const partialMessage = createAssistantMessage("stop");
+		let providerSignal: AbortSignal | undefined;
+		const source = {
+			async *[Symbol.asyncIterator]() {
+				yield { type: "start", partial: partialMessage } as const;
+				const { promise, reject } = Promise.withResolvers<never>();
+				if (providerSignal?.aborted) {
+					reject(new Error("Request was aborted"));
+				}
+				providerSignal?.addEventListener("abort", () => reject(new Error("Request was aborted")), {
+					once: true,
+				});
+				await promise;
+			},
+		} as unknown as AssistantMessageEventStream;
+
+		setBedrockProviderModule({
+			streamBedrock: (_model, _context, options) => {
+				providerSignal = options.signal;
+				return source;
+			},
+		});
+
+		const stream = streamBedrock(createModel(), baseContext, { streamIdleTimeoutMs: 10 });
+		const result = await Promise.race([stream.result(), Bun.sleep(500).then(() => "timeout" as const)]);
+
+		expect(result).not.toBe("timeout");
+		if (result === "timeout") {
+			throw new Error("Timed out waiting for forwarded stream stall result");
+		}
+		expect(providerSignal?.aborted).toBe(true);
+		expect(result.stopReason).toBe("error");
+		expect(result.errorMessage).toBe("Provider stream stalled while waiting for the next event");
+	});
+
+	it("preserves caller aborts while forwarding lazy provider streams", async () => {
+		const abortController = new AbortController();
+		const partialMessage = createAssistantMessage("stop");
+		let providerSignal: AbortSignal | undefined;
+		const source = {
+			async *[Symbol.asyncIterator]() {
+				yield { type: "start", partial: partialMessage } as const;
+				const { promise, reject } = Promise.withResolvers<never>();
+				if (providerSignal?.aborted) {
+					reject(new Error("Request was aborted"));
+				}
+				providerSignal?.addEventListener("abort", () => reject(new Error("Request was aborted")), {
+					once: true,
+				});
+				await promise;
+			},
+		} as unknown as AssistantMessageEventStream;
+
+		setBedrockProviderModule({
+			streamBedrock: (_model, _context, options) => {
+				providerSignal = options.signal;
+				return source;
+			},
+		});
+
+		const stream = streamBedrock(createModel(), baseContext, {
+			signal: abortController.signal,
+			streamIdleTimeoutMs: 500,
+		});
+		const iterator = stream[Symbol.asyncIterator]();
+		const firstEvent = await iterator.next();
+		expect(firstEvent.value?.type).toBe("start");
+
+		abortController.abort();
+		const result = await Promise.race([stream.result(), Bun.sleep(500).then(() => "timeout" as const)]);
+
+		expect(result).not.toBe("timeout");
+		if (result === "timeout") {
+			throw new Error("Timed out waiting for forwarded caller abort result");
+		}
+		expect(result.stopReason).toBe("aborted");
+		expect(result.errorMessage).toBe("Request was aborted");
+	});
 });

packages/coding-agent/test/agent-session-retry-fallback.test.ts

@@ -337,6 +337,88 @@ describe("AgentSession retry fallback", () => {
 		expect(lastAssistant.content).toContainEqual({ type: "text", text: "Recovered after OpenAI timeout" });
 	});
 
+	it("auto-retries stream stall errors", async () => {
+		const model = getBundledModel("openai", "gpt-4o-mini");
+		if (!model) {
+			throw new Error("Expected bundled OpenAI test model to exist");
+		}
+
+		const stallMessage = "Provider stream stalled while waiting for the next event";
+		const requestedModels: string[] = [];
+		let attemptCount = 0;
+
+		const agent = new Agent({
+			getApiKey: provider => `${provider}-test-key`,
+			initialState: {
+				model,
+				systemPrompt: ["Test"],
+				tools: [],
+				messages: [],
+			},
+			streamFn: requestedModel => {
+				requestedModels.push(`${requestedModel.provider}/${requestedModel.id}`);
+				const stream = new MockAssistantStream();
+				queueMicrotask(() => {
+					attemptCount += 1;
+					if (attemptCount === 1) {
+						const message = createAssistantMessage(requestedModel, {
+							stopReason: "error",
+							errorMessage: stallMessage,
+						});
+						stream.push({ type: "start", partial: message });
+						stream.push({ type: "error", reason: "error", error: message });
+						return;
+					}
+					if (attemptCount === 2) {
+						const message = createAssistantMessage(requestedModel, {
+							text: "Recovered after stream stall",
+							stopReason: "stop",
+						});
+						stream.push({
+							type: "start",
+							partial: createAssistantMessage(requestedModel, { text: "", stopReason: "stop" }),
+						});
+						stream.push({ type: "done", reason: "stop", message });
+						return;
+					}
+					throw new Error(`Unexpected retry attempt in stream stall test: ${attemptCount}`);
+				});
+				return stream;
+			},
+		});
+
+		const settings = Settings.isolated({
+			"compaction.enabled": false,
+			"retry.baseDelayMs": 5,
+			"retry.maxRetries": 1,
+		});
+		settings.setModelRole("default", `${model.provider}/${model.id}`);
+
+		session = new AgentSession({
+			agent,
+			sessionManager: SessionManager.inMemory(),
+			settings,
+			modelRegistry,
+		});
+		const { retryStartEvents, retryEndEvents } = trackRetryEvents(session);
+
+		await session.prompt("Retry stream stall");
+		await session.waitForIdle();
+
+		expect(requestedModels).toEqual([`${model.provider}/${model.id}`, `${model.provider}/${model.id}`]);
+		expect(retryStartEvents).toHaveLength(1);
+		expect(retryStartEvents[0]).toMatchObject({
+			attempt: 1,
+			maxAttempts: 1,
+			errorMessage: stallMessage,
+		});
+		expect(retryEndEvents).toHaveLength(1);
+		expect(retryEndEvents[0]).toMatchObject({ success: true, attempt: 1 });
+		const lastAssistant = getLastAssistantMessage(session);
+		expect(lastAssistant.stopReason).toBe("stop");
+		expect(lastAssistant.content).toContainEqual({ type: "text", text: "Recovered after stream stall" });
+	});
+
 	it("auto-retries OpenAI processing-request transient errors", async () => {
 		const model = getBundledModel("openai", "gpt-4o-mini");
 		if (!model) {