From ea96d22866480af69a16f66274695cd13dba3f1c Mon Sep 17 00:00:00 2001 From: gcf Date: Fri, 31 Jul 2026 12:22:33 +0800 Subject: [PATCH] fix: time out stalled LLM streams --- .changeset/fix-stalled-llm-streams.md | 5 + .../src/kosong/contract/errors.ts | 24 +++ .../src/kosong/contract/generate.ts | 124 +++++++++++-- .../src/kosong/contract/provider.ts | 7 +- .../test/kosong/contract/generate.test.ts | 153 +++++++++++++++- packages/kosong/src/errors.ts | 27 +++ packages/kosong/src/generate.ts | 125 +++++++++++-- packages/kosong/src/index.ts | 1 + packages/kosong/src/provider.ts | 7 +- packages/kosong/test/generate.test.ts | 167 +++++++++++++++++- 10 files changed, 604 insertions(+), 36 deletions(-) create mode 100644 .changeset/fix-stalled-llm-streams.md diff --git a/.changeset/fix-stalled-llm-streams.md b/.changeset/fix-stalled-llm-streams.md new file mode 100644 index 0000000000..323cc7b98c --- /dev/null +++ b/.changeset/fix-stalled-llm-streams.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fail and retry LLM streams that stop producing data instead of waiting indefinitely. diff --git a/packages/agent-core-v2/src/kosong/contract/errors.ts b/packages/agent-core-v2/src/kosong/contract/errors.ts index 7347afc47e..2b16d57787 100644 --- a/packages/agent-core-v2/src/kosong/contract/errors.ts +++ b/packages/agent-core-v2/src/kosong/contract/errors.ts @@ -56,6 +56,30 @@ export class APITimeoutError extends ChatProviderError { } } +export class StreamIdleTimeoutError extends APITimeoutError { + readonly idleTimeoutMs: number; + readonly elapsedMs: number; + readonly traceId: string | null; + + constructor( + providerName: string, + modelName: string, + idleTimeoutMs: number, + elapsedMs: number, + traceId: string | null, + ) { + const traceHint = traceId === null ? '' : `, traceId: ${traceId}`; + super( + `LLM stream stalled for ${idleTimeoutMs}ms ` + + `(provider: ${providerName}, model: ${modelName}, elapsedMs: ${elapsedMs}${traceHint}).`, + ); + this.name = 'StreamIdleTimeoutError'; + this.idleTimeoutMs = idleTimeoutMs; + this.elapsedMs = elapsedMs; + this.traceId = traceId; + } +} + export class APIStatusError extends ChatProviderError { readonly statusCode: number; readonly requestId: string | null; diff --git a/packages/agent-core-v2/src/kosong/contract/generate.ts b/packages/agent-core-v2/src/kosong/contract/generate.ts index 64be58b60c..453ec4c14b 100644 --- a/packages/agent-core-v2/src/kosong/contract/generate.ts +++ b/packages/agent-core-v2/src/kosong/contract/generate.ts @@ -4,12 +4,11 @@ * `generate()` is the single place that orchestrates "call * `ChatProvider.generate` and normalize the event stream": it merges streamed * deltas into a complete assistant `Message`, fires the caller's callbacks, - * enforces the abort contract (standard abort DOMException, stream cancelled - * on abort), and rejects empty or thinking-only responses with - * `APIEmptyResponseError`. + * enforces the abort and stream-idle contracts, and rejects empty or + * thinking-only responses with `APIEmptyResponseError`. */ -import { APIEmptyResponseError, createAbortError } from './errors'; +import { APIEmptyResponseError, createAbortError, StreamIdleTimeoutError } from './errors'; import { isContentPart, isToolCall, @@ -24,6 +23,7 @@ import type { Tool } from './tool'; import type { TokenUsage } from './usage'; type StoredToolCall = Omit; +const DEFAULT_STREAM_IDLE_TIMEOUT_MS = 180_000; export interface GenerateResult { readonly id: string | null; @@ -60,8 +60,17 @@ export async function generate( ? tools.filter((tool) => tool.deferred !== true) : tools; + const idleTimeoutMs = resolveStreamIdleTimeoutMs(options?.streamIdleTimeoutMs); + const watchdog = new AbortController(); + const providerSignal = + options?.signal === undefined + ? watchdog.signal + : AbortSignal.any([options.signal, watchdog.signal]); + const providerOptions: GenerateOptions = { ...options, signal: providerSignal }; + delete providerOptions.streamIdleTimeoutMs; + options?.onRequestStart?.(); - const stream = await provider.generate(systemPrompt, wireTools, history, options); + const stream = await provider.generate(systemPrompt, wireTools, history, providerOptions); if (stream.traceId !== undefined) { options?.onTraceId?.(stream.traceId); } @@ -73,7 +82,13 @@ export async function generate( let firstPartAt: number | undefined; let lastResumeAt = 0; - for await (const part of stream) { + for await (const part of withStreamIdleTimeout( + stream, + provider, + watchdog, + options?.signal, + idleTimeoutMs, + )) { const arrivedAt = Date.now(); if (firstPartAt === undefined) { firstPartAt = arrivedAt; @@ -186,16 +201,99 @@ type CancelableStream = StreamedMessage & { return?: () => unknown; }; -async function cancelStream(stream: StreamedMessage): Promise { - const cancelable = stream as CancelableStream; - +function settleWithoutWaiting(action: () => unknown): void { try { - await cancelable.cancel?.(); + void Promise.resolve(action()).catch(() => {}); } catch {} +} + +function abandonStream(stream: StreamedMessage): void { + const cancelable = stream as CancelableStream; + settleWithoutWaiting(() => cancelable.cancel?.()); + settleWithoutWaiting(() => cancelable.return?.()); +} + +function resolveStreamIdleTimeoutMs(value?: number): number { + return value !== undefined && Number.isFinite(value) && value > 0 + ? value + : DEFAULT_STREAM_IDLE_TIMEOUT_MS; +} + +async function* withStreamIdleTimeout( + stream: StreamedMessage, + provider: ChatProvider, + watchdog: AbortController, + callerSignal: AbortSignal | undefined, + idleTimeoutMs: number, +): AsyncGenerator { + const iterator = stream[Symbol.asyncIterator](); + const startedAt = Date.now(); + let completed = false; try { - await cancelable.return?.(); - } catch {} + while (true) { + const next = iterator.next(); + let timer: ReturnType | undefined; + let onCallerAbort: (() => void) | undefined; + const timeout = new Promise((_, reject) => { + timer = setTimeout(() => { + reject( + new StreamIdleTimeoutError( + provider.name, + provider.modelName, + idleTimeoutMs, + Date.now() - startedAt, + stream.traceId ?? null, + ), + ); + }, idleTimeoutMs); + }); + const callerAbort = + callerSignal === undefined + ? undefined + : new Promise((_, reject) => { + onCallerAbort = () => reject(createAbortError()); + if (callerSignal.aborted) { + onCallerAbort(); + } else { + callerSignal.addEventListener('abort', onCallerAbort, { once: true }); + } + }); + + try { + const result = + callerAbort === undefined + ? await Promise.race([next, timeout]) + : await Promise.race([next, timeout, callerAbort]); + if (result.done === true) { + completed = true; + return; + } + yield result.value; + } catch (error) { + if (callerSignal?.aborted) { + next.catch(() => {}); + abandonStream(stream); + throw createAbortError(); + } + if (error instanceof StreamIdleTimeoutError) { + next.catch(() => {}); + watchdog.abort(error); + abandonStream(stream); + } + throw error; + } finally { + clearTimeout(timer); + if (onCallerAbort !== undefined) { + callerSignal?.removeEventListener('abort', onCallerAbort); + } + } + } + } finally { + if (!completed) { + settleWithoutWaiting(() => iterator.return?.()); + } + } } async function throwIfAborted(signal?: AbortSignal, stream?: StreamedMessage): Promise { @@ -204,7 +302,7 @@ async function throwIfAborted(signal?: AbortSignal, stream?: StreamedMessage): P } if (stream !== undefined) { - await cancelStream(stream); + abandonStream(stream); } throw createAbortError(); diff --git a/packages/agent-core-v2/src/kosong/contract/provider.ts b/packages/agent-core-v2/src/kosong/contract/provider.ts index 59f695a20b..e23380ab41 100644 --- a/packages/agent-core-v2/src/kosong/contract/provider.ts +++ b/packages/agent-core-v2/src/kosong/contract/provider.ts @@ -9,9 +9,9 @@ * `with*` methods; every per-turn intent (prompt-cache key, sampling * overrides, thinking effort/keep, completion-token budget) flows through * `GenerateOptions` on each `generate` call instead of through morphs. - * - `GenerateOptions` is the per-turn intent carrier. Each wire dialect - * decides how — or whether — to encode an intent (e.g. a cache key may - * become `prompt_cache_key`, `metadata.user_id`, or be silently dropped). + * - `GenerateOptions` is the per-turn intent carrier and stream-control + * contract. Each wire dialect decides how — or whether — to encode an + * intent; the generation driver owns the stream idle deadline. * * Pure types only — no other domain, no I/O, no SDKs. */ @@ -126,6 +126,7 @@ export interface VideoUploadInput { */ export interface GenerateOptions { signal?: AbortSignal; + streamIdleTimeoutMs?: number; auth?: ProviderRequestAuth; responseFormat?: ResponseFormat; /** diff --git a/packages/agent-core-v2/test/kosong/contract/generate.test.ts b/packages/agent-core-v2/test/kosong/contract/generate.test.ts index 20aecc8d42..9c4ddf34ae 100644 --- a/packages/agent-core-v2/test/kosong/contract/generate.test.ts +++ b/packages/agent-core-v2/test/kosong/contract/generate.test.ts @@ -3,13 +3,17 @@ * * Covers event normalization (text/think deltas merged, tool-call argument * deltas routed by stream index), the empty/thinking-only response - * rejections, the abort contract (standard DOMException, stream cancelled), - * callback plumbing, and per-turn intent passthrough via GenerateOptions. + * rejections, abort and stream-idle contracts, callback plumbing, and + * per-turn intent passthrough via GenerateOptions. */ import { describe, expect, it, vi } from 'vitest'; -import { APIEmptyResponseError } from '#/kosong/contract/errors'; +import { + APIEmptyResponseError, + isRetryableGenerateError, + StreamIdleTimeoutError, +} from '#/kosong/contract/errors'; import { generate, type GenerateResult } from '#/kosong/contract/generate'; import type { Message, StreamedMessagePart, ToolCall } from '#/kosong/contract/message'; import type { @@ -56,6 +60,49 @@ class FakeStreamedMessage implements StreamedMessage { } } +function createStalledStream( + options: { + traceId?: string | null; + stallBeforeFirst?: boolean; + onStall?: () => void; + cancel?: () => unknown; + iteratorReturn?: () => Promise>; + } = {}, +): StreamedMessage { + let first = true; + const stream: StreamedMessage & { cancel(): unknown } = { + id: 'stalled-1', + usage: null, + finishReason: null, + rawFinishReason: null, + traceId: options.traceId, + cancel: () => options.cancel?.(), + [Symbol.asyncIterator]() { + return { + next(): Promise> { + if (first) { + first = false; + if (options.stallBeforeFirst !== true) { + return Promise.resolve({ + done: false, + value: { type: 'think', think: 'partial' }, + }); + } + } + options.onStall?.(); + return new Promise(() => {}); + }, + return(): Promise> { + return ( + options.iteratorReturn?.() ?? Promise.resolve({ done: true, value: undefined }) + ); + }, + }; + }, + }; + return stream; +} + interface FakeProvider { readonly provider: ChatProvider; readonly generateSpy: ReturnType; @@ -299,6 +346,103 @@ describe('generate() abort contract', () => { }); }); +describe('generate() stream idle deadline', () => { + it('throws a retryable stream timeout when no part arrives before the idle deadline', async () => { + const stream = createStalledStream({ traceId: 'trace-stalled' }); + const { provider } = createFakeProvider(stream); + let caught: unknown; + + try { + await generate(provider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, { + streamIdleTimeoutMs: 10, + }); + } catch (error) { + caught = error; + } + + expect(caught).toBeInstanceOf(StreamIdleTimeoutError); + expect(caught).toMatchObject({ + idleTimeoutMs: 10, + traceId: 'trace-stalled', + }); + expect(isRetryableGenerateError(caught)).toBe(true); + }); + + it('throws a stream timeout when the first streamed part never arrives', async () => { + const stream = createStalledStream({ stallBeforeFirst: true }); + const { provider } = createFakeProvider(stream); + + await expect( + generate(provider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, { + streamIdleTimeoutMs: 10, + }), + ).rejects.toBeInstanceOf(StreamIdleTimeoutError); + }); + + it('does not forward the stream idle deadline to the provider', async () => { + const stream = new FakeStreamedMessage([{ type: 'text', text: 'ok' }]); + const { provider, generateSpy } = createFakeProvider(stream); + + await generate(provider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, { + streamIdleTimeoutMs: 10, + }); + + expect(generateSpy.mock.calls[0]?.[3]).not.toHaveProperty('streamIdleTimeoutMs'); + }); + + it('aborts the provider signal when the stream idle deadline expires', async () => { + const { provider, generateSpy } = createFakeProvider(createStalledStream()); + + await expect( + generate(provider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, { + streamIdleTimeoutMs: 10, + }), + ).rejects.toBeInstanceOf(StreamIdleTimeoutError); + + const providerOptions = generateSpy.mock.calls[0]?.[3] as GenerateOptions | undefined; + expect(providerOptions?.signal?.aborted).toBe(true); + expect(providerOptions?.signal?.reason).toBeInstanceOf(StreamIdleTimeoutError); + }); + + it('reports a stream timeout without waiting for stalled cleanup', async () => { + const cancel = vi.fn(() => new Promise(() => {})); + const iteratorReturn = vi.fn( + () => new Promise>(() => {}), + ); + const stream = createStalledStream({ cancel, iteratorReturn }); + const { provider } = createFakeProvider(stream); + + await expect( + generate(provider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, { + streamIdleTimeoutMs: 10, + }), + ).rejects.toBeInstanceOf(StreamIdleTimeoutError); + + expect(cancel).toHaveBeenCalledOnce(); + expect(iteratorReturn).toHaveBeenCalledOnce(); + }); + + it('throws AbortError when the caller aborts while the stream is stalled', async () => { + const controller = new AbortController(); + const stream = createStalledStream({ onStall: () => controller.abort() }); + const { provider } = createFakeProvider(stream); + let caught: unknown; + + try { + await generate(provider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, { + signal: controller.signal, + streamIdleTimeoutMs: 10, + }); + } catch (error) { + caught = error; + } + + expect(caught).toBeInstanceOf(DOMException); + expect((caught as DOMException).name).toBe('AbortError'); + expect(isRetryableGenerateError(caught)).toBe(false); + }); +}); + describe('generate() per-turn intent passthrough', () => { it('passes the GenerateOptions intent fields through to the provider', async () => { const stream = new FakeStreamedMessage([{ type: 'text', text: 'ok' }]); @@ -315,6 +459,7 @@ describe('generate() per-turn intent passthrough', () => { await generate(provider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, options); expect(generateSpy).toHaveBeenCalledTimes(1); - expect(generateSpy.mock.calls[0]?.[3]).toBe(options); + expect(generateSpy.mock.calls[0]?.[3]).toMatchObject(options); + expect(generateSpy.mock.calls[0]?.[3]?.signal).toBeInstanceOf(AbortSignal); }); }); diff --git a/packages/kosong/src/errors.ts b/packages/kosong/src/errors.ts index 1ad2c00742..ae7113b15f 100644 --- a/packages/kosong/src/errors.ts +++ b/packages/kosong/src/errors.ts @@ -30,6 +30,33 @@ export class APITimeoutError extends ChatProviderError { } } +/** + * A provider stream stopped producing parts before its idle deadline. + */ +export class StreamIdleTimeoutError extends APITimeoutError { + readonly idleTimeoutMs: number; + readonly elapsedMs: number; + readonly traceId: string | null; + + constructor( + providerName: string, + modelName: string, + idleTimeoutMs: number, + elapsedMs: number, + traceId: string | null, + ) { + const traceHint = traceId === null ? '' : `, traceId: ${traceId}`; + super( + `LLM stream stalled for ${idleTimeoutMs}ms ` + + `(provider: ${providerName}, model: ${modelName}, elapsedMs: ${elapsedMs}${traceHint}).`, + ); + this.name = 'StreamIdleTimeoutError'; + this.idleTimeoutMs = idleTimeoutMs; + this.elapsedMs = elapsedMs; + this.traceId = traceId; + } +} + /** * HTTP status error from the API. */ diff --git a/packages/kosong/src/generate.ts b/packages/kosong/src/generate.ts index 626dc30736..3eb5271e5f 100644 --- a/packages/kosong/src/generate.ts +++ b/packages/kosong/src/generate.ts @@ -1,4 +1,4 @@ -import { APIEmptyResponseError } from './errors'; +import { APIEmptyResponseError, StreamIdleTimeoutError } from './errors'; import { isContentPart, isToolCall, @@ -14,6 +14,7 @@ import type { TokenUsage } from './usage'; /** Snapshot of a ToolCall excluding the internal `_streamIndex` routing field. */ type StoredToolCall = Omit; +const DEFAULT_STREAM_IDLE_TIMEOUT_MS = 180_000; /** * The result of a single {@link generate} call. @@ -116,8 +117,17 @@ export async function generate( ? tools.filter((tool) => tool.deferred !== true) : tools; + const idleTimeoutMs = resolveStreamIdleTimeoutMs(options?.streamIdleTimeoutMs); + const watchdog = new AbortController(); + const providerSignal = + options?.signal === undefined + ? watchdog.signal + : AbortSignal.any([options.signal, watchdog.signal]); + const providerOptions: GenerateOptions = { ...options, signal: providerSignal }; + delete providerOptions.streamIdleTimeoutMs; + options?.onRequestStart?.(); - const stream = await provider.generate(systemPrompt, wireTools, history, options); + const stream = await provider.generate(systemPrompt, wireTools, history, providerOptions); // Early capture: the trace id arrives with the response headers, before the // stream body — and before any mid-stream abort — so hosts can attribute // even a cancelled stream to its server-side request. @@ -142,7 +152,13 @@ export async function generate( let firstPartAt: number | undefined; let lastResumeAt = 0; - for await (const part of stream) { + for await (const part of withStreamIdleTimeout( + stream, + provider, + watchdog, + options?.signal, + idleTimeoutMs, + )) { const arrivedAt = Date.now(); if (firstPartAt === undefined) { firstPartAt = arrivedAt; @@ -270,20 +286,107 @@ type CancelableStream = StreamedMessage & { return?: () => unknown; }; -function throwAbortError(): never { - throw new DOMException('The operation was aborted.', 'AbortError'); +function createAbortError(): DOMException { + return new DOMException('The operation was aborted.', 'AbortError'); } -async function cancelStream(stream: StreamedMessage): Promise { - const cancelable = stream as CancelableStream; +function throwAbortError(): never { + throw createAbortError(); +} +function settleWithoutWaiting(action: () => unknown): void { try { - await cancelable.cancel?.(); + void Promise.resolve(action()).catch(() => {}); } catch {} +} + +function abandonStream(stream: StreamedMessage): void { + const cancelable = stream as CancelableStream; + settleWithoutWaiting(() => cancelable.cancel?.()); + settleWithoutWaiting(() => cancelable.return?.()); +} + +function resolveStreamIdleTimeoutMs(value?: number): number { + return value !== undefined && Number.isFinite(value) && value > 0 + ? value + : DEFAULT_STREAM_IDLE_TIMEOUT_MS; +} + +async function* withStreamIdleTimeout( + stream: StreamedMessage, + provider: ChatProvider, + watchdog: AbortController, + callerSignal: AbortSignal | undefined, + idleTimeoutMs: number, +): AsyncGenerator { + const iterator = stream[Symbol.asyncIterator](); + const startedAt = Date.now(); + let completed = false; try { - await cancelable.return?.(); - } catch {} + while (true) { + const next = iterator.next(); + let timer: ReturnType | undefined; + let onCallerAbort: (() => void) | undefined; + const timeout = new Promise((_, reject) => { + timer = setTimeout(() => { + reject( + new StreamIdleTimeoutError( + provider.name, + provider.modelName, + idleTimeoutMs, + Date.now() - startedAt, + stream.traceId ?? null, + ), + ); + }, idleTimeoutMs); + }); + const callerAbort = + callerSignal === undefined + ? undefined + : new Promise((_, reject) => { + onCallerAbort = () => reject(createAbortError()); + if (callerSignal.aborted) { + onCallerAbort(); + } else { + callerSignal.addEventListener('abort', onCallerAbort, { once: true }); + } + }); + + try { + const result = + callerAbort === undefined + ? await Promise.race([next, timeout]) + : await Promise.race([next, timeout, callerAbort]); + if (result.done === true) { + completed = true; + return; + } + yield result.value; + } catch (error) { + if (callerSignal?.aborted) { + next.catch(() => {}); + abandonStream(stream); + throw createAbortError(); + } + if (error instanceof StreamIdleTimeoutError) { + next.catch(() => {}); + watchdog.abort(error); + abandonStream(stream); + } + throw error; + } finally { + clearTimeout(timer); + if (onCallerAbort !== undefined) { + callerSignal?.removeEventListener('abort', onCallerAbort); + } + } + } + } finally { + if (!completed) { + settleWithoutWaiting(() => iterator.return?.()); + } + } } async function throwIfAborted(signal?: AbortSignal, stream?: StreamedMessage): Promise { @@ -292,7 +395,7 @@ async function throwIfAborted(signal?: AbortSignal, stream?: StreamedMessage): P } if (stream !== undefined) { - await cancelStream(stream); + abandonStream(stream); } throwAbortError(); diff --git a/packages/kosong/src/index.ts b/packages/kosong/src/index.ts index cd0440637b..8be5d267b3 100644 --- a/packages/kosong/src/index.ts +++ b/packages/kosong/src/index.ts @@ -86,6 +86,7 @@ export { isRequestTooLargeStatusError, isRetryableGenerateError, isToolExchangeAdjacencyError, + StreamIdleTimeoutError, throwIfAbortError, } from './errors'; diff --git a/packages/kosong/src/provider.ts b/packages/kosong/src/provider.ts index 905b12c153..1d6ab06949 100644 --- a/packages/kosong/src/provider.ts +++ b/packages/kosong/src/provider.ts @@ -137,8 +137,13 @@ export interface GenerateOptions { * to their underlying HTTP client; the generate loop in * {@link generate | generate()} also checks the signal between streamed * parts. - */ + */ signal?: AbortSignal; + /** + * Maximum time to wait for the next streamed part after response headers + * arrive. Invalid values fall back to the default three-minute deadline. + */ + streamIdleTimeoutMs?: number; /** * Request-scoped provider auth. Hosts should resolve this immediately before * each request/retry so providers never retain mutable credential state. diff --git a/packages/kosong/test/generate.test.ts b/packages/kosong/test/generate.test.ts index 4f84b95724..b39f636d67 100644 --- a/packages/kosong/test/generate.test.ts +++ b/packages/kosong/test/generate.test.ts @@ -1,7 +1,21 @@ -import { APIEmptyResponseError } from '#/errors'; +/** + * kosong generate() — stream assembly, callbacks, aborts, idle deadlines, and + * response metadata through a fake provider boundary. + */ + +import { + APIEmptyResponseError, + isRetryableGenerateError, + StreamIdleTimeoutError, +} from '#/errors'; import { generate } from '#/generate'; import type { Message, StreamedMessagePart, ToolCall } from '#/message'; -import type { ChatProvider, StreamedMessage, ThinkingEffort } from '#/provider'; +import type { + ChatProvider, + GenerateOptions, + StreamedMessage, + ThinkingEffort, +} from '#/provider'; import type { Tool } from '#/tool'; import type { TokenUsage } from '#/usage'; import { describe, expect, it, vi } from 'vitest'; @@ -31,7 +45,10 @@ function createMockStream( }; } -function createMockProvider(stream: StreamedMessage): ChatProvider { +function createMockProvider( + stream: StreamedMessage, + onGenerate?: (options?: GenerateOptions) => void, +): ChatProvider { return { name: 'mock', modelName: 'mock-model', @@ -40,12 +57,60 @@ function createMockProvider(stream: StreamedMessage): ChatProvider { _systemPrompt: string, _tools: Tool[], _history: Message[], - ): Promise => stream, + options?: GenerateOptions, + ): Promise => { + onGenerate?.(options); + return stream; + }, withThinking(_effort: ThinkingEffort): ChatProvider { return this; }, }; } + +function createStalledStream( + options: { + traceId?: string | null; + stallBeforeFirst?: boolean; + onStall?: () => void; + cancel?: () => unknown; + iteratorReturn?: () => Promise>; + } = {}, +): StreamedMessage { + let first = true; + const stream: StreamedMessage & { cancel(): unknown } = { + id: 'stalled-1', + usage: null, + finishReason: null, + rawFinishReason: null, + traceId: options.traceId, + cancel: () => options.cancel?.(), + [Symbol.asyncIterator]() { + return { + next(): Promise> { + if (first) { + first = false; + if (options.stallBeforeFirst !== true) { + return Promise.resolve({ + done: false, + value: { type: 'think', think: 'partial' }, + }); + } + } + options.onStall?.(); + return new Promise(() => {}); + }, + return(): Promise> { + return ( + options.iteratorReturn?.() ?? Promise.resolve({ done: true, value: undefined }) + ); + }, + }; + }, + }; + return stream; +} + describe('generate()', () => { it('omits trace metadata when the provider does not expose it', async () => { const onTraceId = vi.fn(); @@ -881,6 +946,100 @@ describe('generate()', () => { expect(cancel).toHaveBeenCalledTimes(1); }); + it('throws a retryable stream timeout when no part arrives before the idle deadline', async () => { + const stream = createStalledStream({ traceId: 'trace-stalled' }); + let caught: unknown; + + try { + await generate(createMockProvider(stream), '', [], [], undefined, { + streamIdleTimeoutMs: 10, + }); + } catch (error) { + caught = error; + } + + expect(caught).toBeInstanceOf(StreamIdleTimeoutError); + expect(caught).toMatchObject({ + idleTimeoutMs: 10, + traceId: 'trace-stalled', + }); + expect(isRetryableGenerateError(caught)).toBe(true); + }); + + it('throws a stream timeout when the first streamed part never arrives', async () => { + const stream = createStalledStream({ stallBeforeFirst: true }); + + await expect( + generate(createMockProvider(stream), '', [], [], undefined, { + streamIdleTimeoutMs: 10, + }), + ).rejects.toBeInstanceOf(StreamIdleTimeoutError); + }); + + it('does not forward the stream idle deadline to the provider', async () => { + let providerOptions: GenerateOptions | undefined; + const provider = createMockProvider( + createMockStream([{ type: 'text', text: 'ok' }]), + (options) => { + providerOptions = options; + }, + ); + + await generate(provider, '', [], [], undefined, { streamIdleTimeoutMs: 10 }); + + expect(providerOptions).not.toHaveProperty('streamIdleTimeoutMs'); + }); + + it('aborts the provider signal when the stream idle deadline expires', async () => { + let providerSignal: AbortSignal | undefined; + const provider = createMockProvider(createStalledStream(), (options) => { + providerSignal = options?.signal; + }); + + await expect( + generate(provider, '', [], [], undefined, { streamIdleTimeoutMs: 10 }), + ).rejects.toBeInstanceOf(StreamIdleTimeoutError); + + expect(providerSignal?.aborted).toBe(true); + expect(providerSignal?.reason).toBeInstanceOf(StreamIdleTimeoutError); + }); + + it('reports a stream timeout without waiting for stalled cleanup', async () => { + const cancel = vi.fn(() => new Promise(() => {})); + const iteratorReturn = vi.fn( + () => new Promise>(() => {}), + ); + const stream = createStalledStream({ cancel, iteratorReturn }); + + await expect( + generate(createMockProvider(stream), '', [], [], undefined, { + streamIdleTimeoutMs: 10, + }), + ).rejects.toBeInstanceOf(StreamIdleTimeoutError); + + expect(cancel).toHaveBeenCalledOnce(); + expect(iteratorReturn).toHaveBeenCalledOnce(); + }); + + it('throws AbortError when the caller aborts while the stream is stalled', async () => { + const controller = new AbortController(); + const stream = createStalledStream({ onStall: () => controller.abort() }); + let caught: unknown; + + try { + await generate(createMockProvider(stream), '', [], [], undefined, { + signal: controller.signal, + streamIdleTimeoutMs: 10, + }); + } catch (error) { + caught = error; + } + + expect(caught).toBeInstanceOf(DOMException); + expect((caught as DOMException).name).toBe('AbortError'); + expect(isRetryableGenerateError(caught)).toBe(false); + }); + it('onToolCall receives tool calls in message order', async () => { // Stream yields tc0-header, then its full args, then tc1-header, // then its full args. The callback must see them in message order.