smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ai/stream-timeouts): added deadline-based first-event timeout handling in stream providers

Browse Source 
- Refactored stream timeout handling to allow providers to manage first-event and idle timeouts directly.
- Extracted firstEventTimeoutMs calculations into reusable variables across OpenAI, Azure OpenAI, and Codex providers.
- Consolidated timeout enforcement logic in idle-iterator to use absolute deadline checking and early callback invocation.
- Added providerHandlesStreamTimeouts flag to LazyStreamLimits interface to skip framework-level timeouts when providers handle them.
- Exported isOpenAIResponsesProgressEvent function and removed deprecated MockModelHandle type alias.
This commit is contained in:
can1357
2026-05-26 15:27:06 +02:00
parent 076ca4e71d
commit 6db59ccbac
8 changed files with 165 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ai/src/providers/azure-openai-responses.ts
+14 -5
View File
@@ -29,7 +29,11 @@ import {
import { sanitizeSchemaForOpenAIResponses, toolWireSchema } from "../utils/schema";
import { wrapFetchForSseDebug } from "../utils/sse-debug";
import { mapToOpenAIResponsesToolChoice } from "../utils/tool-choice";
import { normalizeOpenAIResponsesPromptCacheKey, supportsDeveloperRole } from "./openai-responses";
import {
isOpenAIResponsesProgressEvent,
normalizeOpenAIResponsesPromptCacheKey,
supportsDeveloperRole,
} from "./openai-responses";
import {
appendResponsesToolResultMessages,
applyCommonResponsesSamplingParams,
@@ -120,7 +124,8 @@ export const streamAzureOpenAIResponses: StreamFunction<"azure-openai-responses"
const { baseUrl } = resolveAzureConfig(model, options);
const params = buildParams(model, context, options, deploymentName, baseUrl);
options?.onPayload?.(params);
const idleTimeoutMs = getOpenAIStreamIdleTimeoutMs();
const idleTimeoutMs = options?.streamIdleTimeoutMs ?? getOpenAIStreamIdleTimeoutMs();
const firstEventTimeoutMs = options?.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(idleTimeoutMs);
rawRequestDump = {
provider: model.provider,
api: output.api,
@@ -130,9 +135,8 @@ export const streamAzureOpenAIResponses: StreamFunction<"azure-openai-responses"
body: params,
};
const openaiStream = await client.responses.create(params, { signal: requestSignal });
const firstEventWatchdog = createWatchdog(
options?.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(idleTimeoutMs),
() => abortTracker.abortLocally(firstEventTimeoutAbortError),
const firstEventWatchdog = createWatchdog(firstEventTimeoutMs, () =>
abortTracker.abortLocally(firstEventTimeoutAbortError),
);
stream.push({ type: "start", partial: output });
@@ -140,8 +144,13 @@ export const streamAzureOpenAIResponses: StreamFunction<"azure-openai-responses"
iterateWithIdleTimeout(openaiStream, {
watchdog: firstEventWatchdog,
idleTimeoutMs,
firstItemTimeoutMs: firstEventTimeoutMs,
firstItemErrorMessage: AZURE_OPENAI_RESPONSES_FIRST_EVENT_TIMEOUT_MESSAGE,
errorMessage: "Azure OpenAI responses stream stalled while waiting for the next event",
onIdle: () => requestAbortController.abort(),
onFirstItemTimeout: () => abortTracker.abortLocally(firstEventTimeoutAbortError),
abortSignal: options?.signal,
isProgressItem: isOpenAIResponsesProgressEvent,
}),
output,
stream,
packages/ai/src/providers/mock.ts
-4
View File
@@ -216,10 +216,6 @@ export class MockModel implements Model<MockApi> {
this.toolCallCounter = 0;
}
}
/** @deprecated Use {@link MockModel}; the class IS the handle. */
export type MockModelHandle = MockModel;
/** Check whether `model` was produced by `createMockModel`. */
export function isMockModel(model: Model<Api>): model is MockModel {
return model instanceof MockModel;
packages/ai/src/providers/openai-codex-responses.ts
+11 -2
View File
@@ -48,7 +48,11 @@ import {
} from "../utils";
import { AssistantMessageEventStream } from "../utils/event-stream";
import { finalizeErrorMessage, type RawHttpRequestDump } from "../utils/http-inspector";
import { getOpenAIStreamIdleTimeoutMs, iterateWithIdleTimeout } from "../utils/idle-iterator";
import {
getOpenAIStreamIdleTimeoutMs,
getStreamFirstEventTimeoutMs,
iterateWithIdleTimeout,
} from "../utils/idle-iterator";
import { parseStreamingJson } from "../utils/json-parse";
import { adaptSchemaForStrict, NO_STRICT, sanitizeSchemaForOpenAIResponses, toolWireSchema } from "../utils/schema";
import { notifyRawSseEvent } from "../utils/sse-debug";
@@ -553,13 +557,18 @@ function createRequestSetup(options: OpenAICodexResponsesOptions | undefined): C
const requestSignal = options?.signal
? AbortSignal.any([options.signal, requestAbortController.signal])
: requestAbortController.signal;
const idleTimeoutMs = options?.streamIdleTimeoutMs ?? getOpenAIStreamIdleTimeoutMs();
const firstEventTimeoutMs = options?.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(idleTimeoutMs);
const wrapCodexSseStream = (
source: AsyncGenerator<Record<string, unknown>>,
): AsyncGenerator<Record<string, unknown>> =>
iterateWithIdleTimeout(source, {
idleTimeoutMs: options?.streamIdleTimeoutMs ?? getOpenAIStreamIdleTimeoutMs(),
idleTimeoutMs,
firstItemTimeoutMs: firstEventTimeoutMs,
firstItemErrorMessage: "OpenAI Codex SSE stream timed out while waiting for the first event",
errorMessage: "OpenAI Codex SSE stream stalled while waiting for the next event",
onIdle: () => requestAbortController.abort(),
onFirstItemTimeout: () => requestAbortController.abort(),
abortSignal: options?.signal,
isProgressItem: isCodexStreamProgressEvent,
});
packages/ai/src/providers/openai-completions.ts
+6 -3
View File
@@ -476,9 +476,9 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions"> = (
openaiStream = await createCompletionsStream("none");
}
}
const firstEventWatchdog = createWatchdog(
options?.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(idleTimeoutMs),
() => abortTracker.abortLocally(firstEventTimeoutAbortError),
const firstEventTimeoutMs = options?.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(idleTimeoutMs);
const firstEventWatchdog = createWatchdog(firstEventTimeoutMs, () =>
abortTracker.abortLocally(firstEventTimeoutAbortError),
);
if (premiumRequestsTotal !== undefined) {
output.usage.premiumRequests = premiumRequestsTotal;
@@ -662,8 +662,11 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions"> = (
for await (const chunk of iterateWithIdleTimeout(openaiStream, {
watchdog: firstEventWatchdog,
idleTimeoutMs,
firstItemTimeoutMs: firstEventTimeoutMs,
firstItemErrorMessage: OPENAI_COMPLETIONS_FIRST_EVENT_TIMEOUT_MESSAGE,
errorMessage: "OpenAI completions stream stalled while waiting for the next event",
onIdle: () => requestAbortController.abort(),
onFirstItemTimeout: () => abortTracker.abortLocally(firstEventTimeoutAbortError),
abortSignal: options?.signal,
isProgressItem: isOpenAICompletionsProgressChunk,
})) {
packages/ai/src/providers/openai-responses.ts
+7 -4
View File
@@ -123,7 +123,7 @@ const OPENAI_RESPONSES_PROGRESS_EVENT_TYPES = new Set([
"error",
]);
function isOpenAIResponsesProgressEvent(event: unknown): boolean {
export function isOpenAIResponsesProgressEvent(event: unknown): boolean {
if (!event || typeof event !== "object") return false;
const type = (event as { type?: unknown }).type;
return typeof type === "string" && OPENAI_RESPONSES_PROGRESS_EVENT_TYPES.has(type);
@@ -220,6 +220,7 @@ export const streamOpenAIResponses: StreamFunction<"openai-responses"> = (
const providerSessionState = getOpenAIResponsesProviderSessionState(model, options?.providerSessionState);
const { params } = buildParams(model, context, options, providerSessionState, baseUrl);
const idleTimeoutMs = options?.streamIdleTimeoutMs ?? getOpenAIStreamIdleTimeoutMs();
const firstEventTimeoutMs = options?.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(idleTimeoutMs);
options?.onPayload?.(params);
rawRequestDump = {
provider: model.provider,
@@ -240,9 +241,8 @@ export const streamOpenAIResponses: StreamFunction<"openai-responses"> = (
{ provider: model.provider, signal: requestSignal },
);
if (premiumRequestsTotal !== undefined) output.usage.premiumRequests = premiumRequestsTotal;
const firstEventWatchdog = createWatchdog(
options?.streamFirstEventTimeoutMs ?? getStreamFirstEventTimeoutMs(idleTimeoutMs),
() => abortTracker.abortLocally(firstEventTimeoutAbortError),
const firstEventWatchdog = createWatchdog(firstEventTimeoutMs, () =>
abortTracker.abortLocally(firstEventTimeoutAbortError),
);
stream.push({ type: "start", partial: output });
@@ -250,8 +250,11 @@ export const streamOpenAIResponses: StreamFunction<"openai-responses"> = (
await processResponsesStream(
iterateWithIdleTimeout(openaiStream, {
idleTimeoutMs,
firstItemTimeoutMs: firstEventTimeoutMs,
firstItemErrorMessage: OPENAI_RESPONSES_FIRST_EVENT_TIMEOUT_MESSAGE,
watchdog: firstEventWatchdog,
errorMessage: "OpenAI responses stream stalled while waiting for the next event",
onFirstItemTimeout: () => abortTracker.abortLocally(firstEventTimeoutAbortError),
onIdle: () => requestAbortController.abort(),
abortSignal: options?.signal,
isProgressItem: isOpenAIResponsesProgressEvent,
packages/ai/src/providers/register-builtins.ts
+26 -8
View File
@@ -177,6 +177,11 @@ interface LazyStreamLimits {
defaultFirstEventTimeoutMs?: number;
defaultIdleTimeoutMs?: number;
}
/**
* The provider implementation already wraps its upstream transport with
* stream timeouts. Keep the lazy loader from racing it with generic errors.
*/
providerHandlesStreamTimeouts?: boolean;
/**
* Cloud Code Assist (google-gemini-cli / google-antigravity) routinely takes
@@ -191,6 +196,9 @@ const GOOGLE_GEMINI_CLI_LAZY_STREAM_LIMITS: LazyStreamLimits = {
defaultFirstEventTimeoutMs: 300_000,
};
const PROVIDER_HANDLED_STREAM_TIMEOUTS: LazyStreamLimits = {
providerHandlesStreamTimeouts: true,
};
function forwardStream<TApi extends Api>(
target: EventStreamImpl,
source: AsyncIterable<AssistantMessageEvent>,
@@ -201,12 +209,16 @@ function forwardStream<TApi extends Api>(
): void {
(async () => {
try {
const idleTimeoutMs = options.streamIdleTimeoutMs ?? getStreamIdleTimeoutMs(limits?.defaultIdleTimeoutMs);
const providerHandlesStreamTimeouts = limits?.providerHandlesStreamTimeouts === true;
const idleTimeoutMs = providerHandlesStreamTimeouts
? undefined
: (options.streamIdleTimeoutMs ?? getStreamIdleTimeoutMs(limits?.defaultIdleTimeoutMs));
const watchedSource = iterateWithIdleTimeout(source, {
idleTimeoutMs,
firstItemTimeoutMs:
options.streamFirstEventTimeoutMs ??
getStreamFirstEventTimeoutMs(idleTimeoutMs, limits?.defaultFirstEventTimeoutMs),
firstItemTimeoutMs: providerHandlesStreamTimeouts
? 0
: (options.streamFirstEventTimeoutMs ??
getStreamFirstEventTimeoutMs(idleTimeoutMs, limits?.defaultFirstEventTimeoutMs)),
errorMessage: LAZY_STREAM_IDLE_TIMEOUT_ERROR,
firstItemErrorMessage: LAZY_STREAM_FIRST_EVENT_TIMEOUT_ERROR,
onIdle: () => abortTracker.abortLocally(new Error(LAZY_STREAM_IDLE_TIMEOUT_ERROR)),
@@ -395,8 +407,11 @@ function loadBedrockProviderModule(): Promise<LazyProviderModule<"bedrock-conver
// providers, the lazy loading will take effect on the main code path.
// ---------------------------------------------------------------------------
export const streamAnthropic = createLazyStream(loadAnthropicProviderModule);
export const streamAzureOpenAIResponses = createLazyStream(loadAzureOpenAIResponsesProviderModule);
export const streamAnthropic = createLazyStream(loadAnthropicProviderModule, PROVIDER_HANDLED_STREAM_TIMEOUTS);
export const streamAzureOpenAIResponses = createLazyStream(
loadAzureOpenAIResponsesProviderModule,
PROVIDER_HANDLED_STREAM_TIMEOUTS,
);
export const streamGoogle = createLazyStream(loadGoogleProviderModule);
export const streamGoogleGeminiCli = createLazyStream(
loadGoogleGeminiCliProviderModule,
@@ -404,8 +419,11 @@ export const streamGoogleGeminiCli = createLazyStream(
);
export const streamGoogleVertex = createLazyStream(loadGoogleVertexProviderModule);
export const streamOpenAICodexResponses = createLazyStream(loadOpenAICodexResponsesProviderModule);
export const streamOpenAICompletions = createLazyStream(loadOpenAICompletionsProviderModule);
export const streamOpenAIResponses = createLazyStream(loadOpenAIResponsesProviderModule);
export const streamOpenAICompletions = createLazyStream(
loadOpenAICompletionsProviderModule,
PROVIDER_HANDLED_STREAM_TIMEOUTS,
);
export const streamOpenAIResponses = createLazyStream(loadOpenAIResponsesProviderModule, PROVIDER_HANDLED_STREAM_TIMEOUTS);
export const streamCursor = createLazyStream(loadCursorProviderModule);
export const streamOllama = createLazyStream(loadOllamaProviderModule);
packages/ai/src/stream.ts
+2
View File
@@ -565,6 +565,8 @@ function mapOptionsForApi<TApi extends Api>(
metadata: options?.metadata,
sessionId: options?.sessionId,
promptCacheKey: options?.promptCacheKey,
streamFirstEventTimeoutMs: options?.streamFirstEventTimeoutMs,
streamIdleTimeoutMs: options?.streamIdleTimeoutMs,
providerSessionState: options?.providerSessionState,
onPayload: options?.onPayload,
onResponse: options?.onResponse,
packages/ai/src/utils/idle-iterator.ts
+99 -83
View File
@@ -109,7 +109,9 @@ export async function* iterateWithIdleTimeout<T>(
options: IdleTimeoutIteratorOptions,
): AsyncGenerator<T> {
let watchdog = options.watchdog;
const firstItemTimeoutMs = options.firstItemTimeoutMs ?? options.idleTimeoutMs;
const firstItemTimeoutMs = options.firstItemTimeoutMs ?? (watchdog ? undefined : options.idleTimeoutMs);
const firstItemDeadlineMs =
firstItemTimeoutMs !== undefined && firstItemTimeoutMs > 0 ? Date.now() + firstItemTimeoutMs : undefined;
const abortSignal = options.abortSignal;
const iterator = iterable[Symbol.asyncIterator]();
@@ -151,93 +153,107 @@ export async function* iterateWithIdleTimeout<T>(
(firstItemTimeoutMs === undefined || firstItemTimeoutMs <= 0) &&
(options.idleTimeoutMs === undefined || options.idleTimeoutMs <= 0);
while (true) {
let activeTimeoutMs: number | undefined;
if (awaitingFirstItem) {
activeTimeoutMs = firstItemTimeoutMs;
} else if (options.idleTimeoutMs !== undefined && options.idleTimeoutMs > 0) {
activeTimeoutMs = options.idleTimeoutMs - (Date.now() - lastProgressAt);
if (activeTimeoutMs <= 0) {
options.onIdle?.();
closeIterator();
throw new Error(options.errorMessage);
}
}
const nextResultPromise = withRacy(iterator.next());
const racers: Array<
Promise<
| { kind: "next"; result: IteratorResult<T> }
| { kind: "error"; error: unknown }
| { kind: "timeout" }
| { kind: "abort" }
>
> = [nextResultPromise];
let timer: NodeJS.Timeout | undefined;
let resolveTimeout: ((value: { kind: "timeout" }) => void) | undefined;
const enforceTimeout = !noTimeoutEnforced && activeTimeoutMs !== undefined && activeTimeoutMs > 0;
if (enforceTimeout) {
const { promise, resolve } = Promise.withResolvers<{ kind: "timeout" }>();
resolveTimeout = resolve;
timer = setTimeout(() => resolve({ kind: "timeout" }), activeTimeoutMs);
racers.push(promise);
}
let abortListener: (() => void) | undefined;
let resolveAbort: ((value: { kind: "abort" }) => void) | undefined;
if (abortSignal) {
const { promise, resolve } = Promise.withResolvers<{ kind: "abort" }>();
resolveAbort = resolve;
abortListener = () => resolve({ kind: "abort" });
abortSignal.addEventListener("abort", abortListener, { once: true });
racers.push(promise);
}
try {
const outcome = await Promise.race(racers);
if (outcome.kind === "abort") {
closeIterator();
throw abortReason(abortSignal!);
}
if (outcome.kind === "timeout") {
if (!awaitingFirstItem) {
options.onIdle?.();
} else {
options.onFirstItemTimeout?.();
try {
while (true) {
let activeTimeoutMs: number | undefined;
if (awaitingFirstItem) {
if (firstItemDeadlineMs !== undefined) {
activeTimeoutMs = firstItemDeadlineMs - Date.now();
if (activeTimeoutMs <= 0) {
options.onFirstItemTimeout?.();
closeIterator();
throw new Error(options.firstItemErrorMessage ?? options.errorMessage);
}
}
} else if (options.idleTimeoutMs !== undefined && options.idleTimeoutMs > 0) {
activeTimeoutMs = options.idleTimeoutMs - (Date.now() - lastProgressAt);
if (activeTimeoutMs <= 0) {
options.onIdle?.();
closeIterator();
throw new Error(options.errorMessage);
}
closeIterator();
throw new Error(
!awaitingFirstItem ? options.errorMessage : (options.firstItemErrorMessage ?? options.errorMessage),
);
}
if (outcome.kind === "error") {
throw outcome.error;
const nextResultPromise = withRacy(iterator.next());
const racers: Array<
Promise<
| { kind: "next"; result: IteratorResult<T> }
| { kind: "error"; error: unknown }
| { kind: "timeout" }
| { kind: "abort" }
>
> = [nextResultPromise];
let timer: NodeJS.Timeout | undefined;
let resolveTimeout: ((value: { kind: "timeout" }) => void) | undefined;
const enforceTimeout = !noTimeoutEnforced && activeTimeoutMs !== undefined && activeTimeoutMs > 0;
if (enforceTimeout) {
const { promise, resolve } = Promise.withResolvers<{ kind: "timeout" }>();
resolveTimeout = resolve;
timer = setTimeout(() => resolve({ kind: "timeout" }), activeTimeoutMs);
racers.push(promise);
}
if (outcome.result.done) {
markFirstItemReceived();
return;
let abortListener: (() => void) | undefined;
let resolveAbort: ((value: { kind: "abort" }) => void) | undefined;
if (abortSignal) {
const { promise, resolve } = Promise.withResolvers<{ kind: "abort" }>();
resolveAbort = resolve;
abortListener = () => resolve({ kind: "abort" });
abortSignal.addEventListener("abort", abortListener, { once: true });
racers.push(promise);
}
const item = outcome.result.value;
// Non-progress items (e.g. provider keepalives, synthetic `start` events that
// arrive before the model has produced any tokens) MUST NOT flip us out of
// `awaitingFirstItem`. Otherwise the next iteration switches from the (longer)
// first-item watchdog to the (shorter) idle watchdog while we're still waiting
// on the model's first real output.
if (isProgressItem(item)) {
markFirstItemReceived();
lastProgressAt = Date.now();
try {
const outcome = await Promise.race(racers);
if (outcome.kind === "abort") {
closeIterator();
throw abortReason(abortSignal!);
}
if (outcome.kind === "timeout") {
if (!awaitingFirstItem) {
options.onIdle?.();
} else {
options.onFirstItemTimeout?.();
}
closeIterator();
throw new Error(
!awaitingFirstItem ? options.errorMessage : (options.firstItemErrorMessage ?? options.errorMessage),
);
}
if (outcome.kind === "error") {
throw outcome.error;
}
if (outcome.result.done) {
markFirstItemReceived();
return;
}
const item = outcome.result.value;
// Non-progress items (e.g. provider keepalives, synthetic `start` events that
// arrive before the model has produced any tokens) MUST NOT flip us out of
// `awaitingFirstItem`. Otherwise the next iteration switches from the (longer)
// first-item watchdog to the (shorter) idle watchdog while we're still waiting
// on the model's first real output.
if (isProgressItem(item)) {
markFirstItemReceived();
lastProgressAt = Date.now();
}
yield item;
} finally {
if (timer !== undefined) clearTimeout(timer);
// Resolve dangling promises so the racers don't leak (Promise.race is one-shot).
resolveTimeout?.({ kind: "timeout" });
if (abortListener && abortSignal) {
abortSignal.removeEventListener("abort", abortListener);
}
resolveAbort?.({ kind: "abort" });
}
yield item;
} finally {
if (timer !== undefined) clearTimeout(timer);
// Resolve dangling promises so the racers don't leak (Promise.race is one-shot).
resolveTimeout?.({ kind: "timeout" });
if (abortListener && abortSignal) {
abortSignal.removeEventListener("abort", abortListener);
}
resolveAbort?.({ kind: "abort" });
}
} finally {
if (watchdog) {
clearTimeout(watchdog);
watchdog = undefined;
}
}
}