From bea95b056d9a00067ace5987bdc2eb8c807e0d77 Mon Sep 17 00:00:00 2001 From: ephraimduncan Date: Thu, 14 May 2026 15:19:18 +0000 Subject: [PATCH] fix(ai): recover stalled lazy provider streams --- packages/ai/CHANGELOG.md | 5 ++ .../ai/src/providers/register-builtins.ts | 43 ++++++++-- packages/ai/src/utils/idle-iterator.ts | 2 +- packages/ai/test/register-builtins.test.ts | 80 ++++++++++++++++++ .../test/agent-session-retry-fallback.test.ts | 82 +++++++++++++++++++ 5 files changed, 203 insertions(+), 9 deletions(-) diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index e8fbc0668..6c53e986e 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -22,6 +22,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 diff --git a/packages/ai/src/providers/register-builtins.ts b/packages/ai/src/providers/register-builtins.ts index bf9d85103..9b2895cdb 100644 --- a/packages/ai/src/providers/register-builtins.ts +++ b/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, ): source is AsyncIterable & { result(): Promise } { @@ -165,10 +170,23 @@ function forwardStream( target: EventStreamImpl, source: AsyncIterable, model: Model, + options: OptionsForApi, + 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( 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(model: Model, error: unknown): AssistantMessage { +function createLazyLoadErrorMessage( + model: Model, + error: unknown, + stopReason: Extract = "error", +): AssistantMessage { return { role: "assistant", content: [], @@ -199,8 +222,9 @@ function createLazyLoadErrorMessage(model: Model, 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( ): (model: Model, context: Context, options: OptionsForApi) => EventStreamImpl { return (model, context, options) => { const outer = new EventStreamImpl(); + const streamOptions = (options ?? {}) as OptionsForApi; 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; + const inner = module.stream(model, context, providerOptions); + forwardStream(outer, inner, model, streamOptions, abortTracker); }) .catch(error => { const message = createLazyLoadErrorMessage(model, error); diff --git a/packages/ai/src/utils/idle-iterator.ts b/packages/ai/src/utils/idle-iterator.ts index 44bd532da..61d06063e 100644 --- a/packages/ai/src/utils/idle-iterator.ts +++ b/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 { diff --git a/packages/ai/test/register-builtins.test.ts b/packages/ai/test/register-builtins.test.ts index f3911edcb..15bdf6a21 100644 --- a/packages/ai/test/register-builtins.test.ts +++ b/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(); + 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(); + 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"); + }); }); diff --git a/packages/coding-agent/test/agent-session-retry-fallback.test.ts b/packages/coding-agent/test/agent-session-retry-fallback.test.ts index ab29a1187..a94e7886a 100644 --- a/packages/coding-agent/test/agent-session-retry-fallback.test.ts +++ b/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) {