diff --git a/.changeset/stream-stall-watchdog.md b/.changeset/stream-stall-watchdog.md new file mode 100644 index 0000000000..cfbfb01ffc --- /dev/null +++ b/.changeset/stream-stall-watchdog.md @@ -0,0 +1,6 @@ +--- +"@moonshot-ai/kimi-code": patch +"@moonshot-ai/kimi-code-sdk": patch +--- + +Fail stalled provider streams instead of hanging the turn forever. `generate()` now watches for stream inactivity (default budget 300s per gap, `KIMI_STREAM_STALL_TIMEOUT_MS` to override, `0` to disable): when no part arrives within the budget the connection is torn down and an `APITimeoutError` is thrown, which the existing retry classification already treats as retryable — transient stalls recover via `chatWithRetry`, persistent ones end the turn with a real error. Cancelling a stalled stream now also aborts promptly instead of hanging until the next part. diff --git a/packages/agent-core-v2/src/kosong/contract/generate.ts b/packages/agent-core-v2/src/kosong/contract/generate.ts index 3fbecb6e15..d0c7b1b05d 100644 --- a/packages/agent-core-v2/src/kosong/contract/generate.ts +++ b/packages/agent-core-v2/src/kosong/contract/generate.ts @@ -6,10 +6,14 @@ * 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`. + * `APIEmptyResponseError`. A stall watchdog bounds the wait for response + * headers and for each streamed part (`GenerateOptions.streamStallTimeoutMs`, + * default `DEFAULT_STREAM_STALL_TIMEOUT_MS`, `0` disables): a silent exchange + * is torn down via a request-linked `AbortController` and rejected with + * `APITimeoutError` (retryable) instead of hanging forever. */ -import { APIEmptyResponseError, createAbortError } from './errors'; +import { APIEmptyResponseError, APITimeoutError, createAbortError } from './errors'; import { isContentPart, isToolCall, @@ -39,6 +43,8 @@ export interface GenerateCallbacks { onToolCall?: (toolCall: ToolCall) => void | Promise; } +export const DEFAULT_STREAM_STALL_TIMEOUT_MS = 300_000; + export async function generate( provider: ChatProvider, systemPrompt: string, @@ -61,7 +67,33 @@ export async function generate( : tools; options?.onRequestStart?.(); - const stream = await provider.generate(systemPrompt, wireTools, history, options); + const stallAbort = new AbortController(); + const requestSignal = + options?.signal === undefined + ? stallAbort.signal + : AbortSignal.any([options.signal, stallAbort.signal]); + const stallTimeoutMs = options?.streamStallTimeoutMs ?? DEFAULT_STREAM_STALL_TIMEOUT_MS; + const generatePromise = provider.generate(systemPrompt, wireTools, history, { + ...options, + signal: requestSignal, + }); + const generateOutcome = await raceStallOrAbort(generatePromise, stallTimeoutMs, requestSignal); + if (generateOutcome === 'aborted' || generateOutcome === 'stalled') { + if (generateOutcome === 'stalled') { + stallAbort.abort(); + } + void generatePromise + .then((lateStream) => cancelStream(lateStream)) + .catch(() => undefined); + if (generateOutcome === 'aborted') { + throw createAbortError(); + } + throw new APITimeoutError( + `The API did not respond within ${stallTimeoutMs}ms (no response headers).` + + ` Provider: ${provider.name}, model: ${provider.modelName}`, + ); + } + const stream = generateOutcome; if (stream.traceId !== undefined) { options?.onTraceId?.(stream.traceId); } @@ -73,50 +105,69 @@ export async function generate( let firstPartAt: number | undefined; let lastResumeAt = 0; - for await (const part of stream) { - const arrivedAt = Date.now(); - if (firstPartAt === undefined) { - firstPartAt = arrivedAt; - } else { - serverDecodeMs += arrivedAt - lastResumeAt; - } - - try { - await throwIfAborted(options?.signal, stream); + const iterator = stream[Symbol.asyncIterator](); + try { + for (;;) { + const next = await nextStreamPart(iterator, stream, stallTimeoutMs, requestSignal, stallAbort); + if (next === 'stalled') { + throw new APITimeoutError( + `The API stream stalled: no data received for ${stallTimeoutMs}ms.` + + formatFinishReasonHint(stream) + + ` Provider: ${provider.name}, model: ${provider.modelName}`, + ); + } + if (next.done === true) { + break; + } + const part = next.value; + const arrivedAt = Date.now(); + if (firstPartAt === undefined) { + firstPartAt = arrivedAt; + } else { + serverDecodeMs += arrivedAt - lastResumeAt; + } - if (callbacks?.onMessagePart !== undefined) { - await callbacks.onMessagePart(deepCopyPart(part)); + try { await throwIfAborted(options?.signal, stream); - } - if ( - isToolCallPart(part) && - part.index !== undefined && - !isPendingToolCallAtIndex(pendingPart, part.index) - ) { - const arrayIdx = toolCallIndexMap.get(part.index); - if (arrayIdx !== undefined) { - const target = message.toolCalls[arrayIdx]; - if (target !== undefined && part.argumentsPart !== null) { - target.arguments = - target.arguments === null - ? part.argumentsPart - : target.arguments + part.argumentsPart; + if (callbacks?.onMessagePart !== undefined) { + await callbacks.onMessagePart(deepCopyPart(part)); + await throwIfAborted(options?.signal, stream); + } + + if ( + isToolCallPart(part) && + part.index !== undefined && + !isPendingToolCallAtIndex(pendingPart, part.index) + ) { + const arrayIdx = toolCallIndexMap.get(part.index); + if (arrayIdx !== undefined) { + const target = message.toolCalls[arrayIdx]; + if (target !== undefined && part.argumentsPart !== null) { + target.arguments = + target.arguments === null + ? part.argumentsPart + : target.arguments + part.argumentsPart; + } + continue; } - continue; } - } - if (pendingPart === null) { - pendingPart = part; - } else if (!mergeInPlace(pendingPart, part)) { - flushPart(message, pendingPart, toolCallIndexMap); - pendingPart = part; + if (pendingPart === null) { + pendingPart = part; + } else if (!mergeInPlace(pendingPart, part)) { + flushPart(message, pendingPart, toolCallIndexMap); + pendingPart = part; + } + } finally { + lastResumeAt = Date.now(); + clientConsumeMs += lastResumeAt - arrivedAt; } - } finally { - lastResumeAt = Date.now(); - clientConsumeMs += lastResumeAt - arrivedAt; } + } catch (error) { + void cancelStream(stream); + teardownIterator(iterator); + throw error; } await throwIfAborted(options?.signal, stream); @@ -198,6 +249,68 @@ async function cancelStream(stream: StreamedMessage): Promise { } catch {} } +async function raceStallOrAbort( + pending: Promise, + stallTimeoutMs: number, + signal: AbortSignal, +): Promise { + if (signal.aborted) { + return 'aborted'; + } + + let timer: ReturnType | undefined; + let onAbort: (() => void) | undefined; + const watchdog = new Promise<'stalled' | 'aborted'>((resolve) => { + if (Number.isFinite(stallTimeoutMs) && stallTimeoutMs > 0) { + timer = setTimeout(() => { + resolve('stalled'); + }, stallTimeoutMs); + (timer as { unref?: () => void }).unref?.(); + } + onAbort = () => { + resolve('aborted'); + }; + signal.addEventListener('abort', onAbort, { once: true }); + }); + + try { + return await Promise.race([pending, watchdog]); + } finally { + if (timer !== undefined) { + clearTimeout(timer); + } + if (onAbort !== undefined) { + signal.removeEventListener('abort', onAbort); + } + } +} + +async function nextStreamPart( + iterator: AsyncIterator, + stream: StreamedMessage, + stallTimeoutMs: number, + signal: AbortSignal, + stallAbort: AbortController, +): Promise | 'stalled'> { + const outcome = await raceStallOrAbort(iterator.next(), stallTimeoutMs, signal); + if (outcome === 'aborted') { + void cancelStream(stream); + teardownIterator(iterator); + throw createAbortError(); + } + if (outcome === 'stalled') { + stallAbort.abort(); + void cancelStream(stream); + teardownIterator(iterator); + return 'stalled'; + } + return outcome; +} + +function teardownIterator(iterator: AsyncIterator): void { + void Promise.resolve(iterator.return?.()).catch(() => undefined); +} + async function throwIfAborted(signal?: AbortSignal, stream?: StreamedMessage): Promise { if (!signal?.aborted) { return; diff --git a/packages/agent-core-v2/src/kosong/contract/provider.ts b/packages/agent-core-v2/src/kosong/contract/provider.ts index 3d79b10039..f5a7407431 100644 --- a/packages/agent-core-v2/src/kosong/contract/provider.ts +++ b/packages/agent-core-v2/src/kosong/contract/provider.ts @@ -102,6 +102,7 @@ export interface GenerateOptions { onRequestSent?: () => void; onStreamEnd?: (stats?: StreamDecodeStats) => void; onTraceId?: (traceId: string | null) => void; + streamStallTimeoutMs?: number; } export interface ChatProvider { 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 68fcfd5a1a..d5d18222ff 100644 --- a/packages/agent-core-v2/test/kosong/contract/generate.test.ts +++ b/packages/agent-core-v2/test/kosong/contract/generate.test.ts @@ -9,7 +9,7 @@ import { describe, expect, it, vi } from 'vitest'; -import { APIEmptyResponseError } from '#/kosong/contract/errors'; +import { APIEmptyResponseError, APITimeoutError } from '#/kosong/contract/errors'; import { generate, type GenerateResult } from '#/kosong/contract/generate'; import type { Message, StreamedMessagePart, ToolCall } from '#/kosong/contract/message'; import type { @@ -314,6 +314,154 @@ 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); + const received = generateSpy.mock.calls[0]?.[3]; + // The intent fields pass through unchanged; `signal` is replaced by the + // generate loop's linked stall-abort signal (never the caller's object, + // and present even when the caller passed none). + expect(received).toMatchObject({ + cacheKey: 'session-42', + sampling: { temperature: 0.7, topP: 0.9 }, + thinking: { effort: 'high', keep: 'all' }, + maxCompletionTokens: 4096, + usedContextTokens: 1000, + maxContextTokens: 128000, + }); + expect(received?.signal).toBeInstanceOf(AbortSignal); + }); +}); + +describe('generate() stream-stall watchdog', () => { + class HangingStreamedMessage implements StreamedMessage { + readonly id: string | null = 'hang-1'; + readonly usage: TokenUsage | null = null; + readonly finishReason: FinishReason | null = null; + readonly rawFinishReason: string | null = null; + cancelCalls = 0; + + constructor(private readonly firstPart?: StreamedMessagePart) {} + + [Symbol.asyncIterator](): AsyncIterator { + const firstPart = this.firstPart; + let yielded = false; + return { + next: (): Promise> => { + if (firstPart !== undefined && !yielded) { + yielded = true; + return Promise.resolve({ done: false, value: firstPart }); + } + return new Promise>(() => {}); + }, + }; + } + + cancel(): void { + this.cancelCalls++; + } + } + + it('fails with APITimeoutError and cancels the stream when it stalls mid-generation', async () => { + const stream = new HangingStreamedMessage({ type: 'text', text: 'partial' }); + const { provider } = createFakeProvider(stream); + const startedAt = Date.now(); + + await expect( + generate(provider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, { + streamStallTimeoutMs: 50, + }), + ).rejects.toBeInstanceOf(APITimeoutError); + expect(Date.now() - startedAt).toBeLessThan(5_000); + expect(stream.cancelCalls).toBeGreaterThan(0); + }); + + it('applies the watchdog to the wait for the first part', async () => { + const stream = new HangingStreamedMessage(); + const { provider } = createFakeProvider(stream); + + await expect( + generate(provider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, { + streamStallTimeoutMs: 50, + }), + ).rejects.toBeInstanceOf(APITimeoutError); + }); + + it('covers the response-headers wait inside provider.generate()', async () => { + const hangingProvider: ChatProvider = { + name: 'hanging-generate', + modelName: 'hanging-model', + thinkingEffort: null, + generate: () => new Promise(() => {}), + }; + const startedAt = Date.now(); + + await expect( + generate(hangingProvider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, { + streamStallTimeoutMs: 50, + }), + ).rejects.toBeInstanceOf(APITimeoutError); + expect(Date.now() - startedAt).toBeLessThan(5_000); + }); + + it('does not hang when stream teardown never settles', async () => { + // A faulty provider whose cancel() pends forever must not block the stall + // path — teardown is best-effort, the timeout still wins. + const stream = new HangingStreamedMessage(); + stream.cancel = () => new Promise(() => {}); + const { provider } = createFakeProvider(stream); + const startedAt = Date.now(); + + await expect( + generate(provider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, { + streamStallTimeoutMs: 50, + }), + ).rejects.toBeInstanceOf(APITimeoutError); + expect(Date.now() - startedAt).toBeLessThan(5_000); + }); + + it('cancels the stream when a part callback throws', async () => { + // `for await` closed the iterator automatically on a throwing body + // (AsyncIteratorClose); the manual watchdog loop must preserve that, or a + // throwing callback leaks the provider connection. + const stream = new FakeStreamedMessage([{ type: 'text', text: 'x' }]); + const { provider } = createFakeProvider(stream); + + await expect( + generate( + provider, + SYSTEM_PROMPT, + NO_TOOLS, + HISTORY, + { + onMessagePart: () => { + throw new Error('boom'); + }, + }, + { streamStallTimeoutMs: 0 }, + ), + ).rejects.toThrow('boom'); + expect(stream.cancelCalls).toBeGreaterThan(0); + }); + + it('aborting mid-stall rejects promptly with the standard abort DOMException', async () => { + const stream = new HangingStreamedMessage(); + const { provider } = createFakeProvider(stream); + const controller = new AbortController(); + const startedAt = Date.now(); + + const pending = generate(provider, SYSTEM_PROMPT, NO_TOOLS, HISTORY, undefined, { + streamStallTimeoutMs: 60_000, + signal: controller.signal, + }); + setTimeout(() => controller.abort(), 30); + + let caught: unknown; + try { + await pending; + } catch (error) { + caught = error; + } + expect(caught).toBeInstanceOf(DOMException); + expect((caught as DOMException).name).toBe('AbortError'); + expect(Date.now() - startedAt).toBeLessThan(5_000); + expect(stream.cancelCalls).toBeGreaterThan(0); }); }); diff --git a/packages/agent-core-v2/test/kosong/model/modelRequester.test.ts b/packages/agent-core-v2/test/kosong/model/modelRequester.test.ts index 7daae99708..b083ac5d5a 100644 --- a/packages/agent-core-v2/test/kosong/model/modelRequester.test.ts +++ b/packages/agent-core-v2/test/kosong/model/modelRequester.test.ts @@ -147,12 +147,12 @@ describe('ModelRequesterImpl request execution', () => { it('maps ModelRequestParams onto GenerateOptions 1:1', async () => { const provider = new FakeChatProvider(); const requester = new ModelRequesterImpl(modelWith(staticAuth('sk-1')), registryReturning(provider)); - const signal = AbortSignal.timeout(1000); + const controller = new AbortController(); await collect( requester.request( { ...INPUT, responseFormat: { type: 'json_object' } }, - signal, + controller.signal, { cacheKey: 'session-1', sampling: { temperature: 0.5, topP: 0.9 }, @@ -167,7 +167,13 @@ describe('ModelRequesterImpl request execution', () => { expect(provider.calls).toHaveLength(1); const options = provider.calls[0]!.options; - expect(options?.signal).toBe(signal); + // The provider receives a composite signal: the caller's signal linked + // with the generate loop's stall-abort controller. Aborting the caller's + // signal must propagate through the link. + expect(options?.signal).toBeInstanceOf(AbortSignal); + expect(options?.signal?.aborted).toBe(false); + controller.abort(); + expect(options?.signal?.aborted).toBe(true); expect(options?.auth).toEqual({ apiKey: 'sk-1' }); expect(options?.cacheKey).toBe('session-1'); expect(options?.sampling).toEqual({ temperature: 0.5, topP: 0.9 }); diff --git a/packages/agent-core/src/agent/turn/kosong-llm.ts b/packages/agent-core/src/agent/turn/kosong-llm.ts index 893ea13fd9..1d5ff69ec0 100644 --- a/packages/agent-core/src/agent/turn/kosong-llm.ts +++ b/packages/agent-core/src/agent/turn/kosong-llm.ts @@ -35,6 +35,7 @@ import type { LLMChatResponse, LLMStreamTiming, } from '../../loop'; +import { parseFloatEnv } from '#/config/resolve'; import { applyCompletionBudget, type CompletionBudgetConfig, @@ -124,6 +125,10 @@ export class KosongLLM implements LLM { onStreamEnd: markStreamEnd, onTraceId: (traceId) => params.trace?.capture(traceId), requestLogFields: params.requestLogFields, + streamStallTimeoutMs: parseFloatEnv( + process.env['KIMI_STREAM_STALL_TIMEOUT_MS'], + 'KIMI_STREAM_STALL_TIMEOUT_MS', + ), }; const result = await this.generate( diff --git a/packages/kosong/src/generate.ts b/packages/kosong/src/generate.ts index 626dc30736..6c71f64b92 100644 --- a/packages/kosong/src/generate.ts +++ b/packages/kosong/src/generate.ts @@ -1,4 +1,4 @@ -import { APIEmptyResponseError } from './errors'; +import { APIEmptyResponseError, APITimeoutError } from './errors'; import { isContentPart, isToolCall, @@ -60,6 +60,15 @@ export interface GenerateCallbacks { onToolCall?: (toolCall: ToolCall) => void | Promise; } +/** + * Default inactivity budget for the stream-stall watchdog: if no stream part + * arrives within this window, the stream is cancelled and the generate call + * fails with `APITimeoutError`. Generous on purpose — a legitimately slow + * first token on a huge context must not trip it; it exists to catch dead + * connections that would otherwise hang the turn forever. + */ +export const DEFAULT_STREAM_STALL_TIMEOUT_MS = 300_000; + /** * Generate one assistant message by streaming from the given provider. * @@ -81,6 +90,9 @@ export interface GenerateCallbacks { * * @throws {DOMException} with name `"AbortError"` when `options.signal` is * aborted before or during streaming. + * @throws {APITimeoutError} when no stream part arrives within the stall + * watchdog budget (`options.streamStallTimeoutMs`, default + * {@link DEFAULT_STREAM_STALL_TIMEOUT_MS}). * @throws {APIEmptyResponseError} when the response contains no content and * no tool calls, or only thinking content without any text or tool calls. */ @@ -117,7 +129,42 @@ export async function generate( : tools; options?.onRequestStart?.(); - const stream = await provider.generate(systemPrompt, wireTools, history, options); + // Link the caller's signal with an internal controller: when the stall + // watchdog fires, aborting this controller is what actually tears down the + // provider's HTTP connection (providers forward the signal to their HTTP + // clients); the in-flight iteration then settles instead of leaking. + const stallAbort = new AbortController(); + const requestSignal = + options?.signal === undefined + ? stallAbort.signal + : AbortSignal.any([options.signal, stallAbort.signal]); + const stallTimeoutMs = options?.streamStallTimeoutMs ?? DEFAULT_STREAM_STALL_TIMEOUT_MS; + // The watchdog covers the whole exchange, starting with the response-headers + // wait inside provider.generate(): an endpoint that accepts the request but + // never answers must not park the turn before the stream even exists. + const generatePromise = provider.generate(systemPrompt, wireTools, history, { + ...options, + signal: requestSignal, + }); + const generateOutcome = await raceStallOrAbort(generatePromise, stallTimeoutMs, requestSignal); + if (generateOutcome === 'aborted' || generateOutcome === 'stalled') { + if (generateOutcome === 'stalled') { + stallAbort.abort(); + } + // The provider call may still be in flight (or have resolved just as the + // watchdog fired): cancel whatever it produced, without blocking on it. + void generatePromise + .then((lateStream) => cancelStream(lateStream)) + .catch(() => undefined); + if (generateOutcome === 'aborted') { + throwAbortError(); + } + throw new APITimeoutError( + `The API did not respond within ${stallTimeoutMs}ms (no response headers).` + + ` Provider: ${provider.name}, model: ${provider.modelName}`, + ); + } + const stream = generateOutcome; // 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,61 +189,84 @@ export async function generate( let firstPartAt: number | undefined; let lastResumeAt = 0; - for await (const part of stream) { - const arrivedAt = Date.now(); - if (firstPartAt === undefined) { - firstPartAt = arrivedAt; - } else { - serverDecodeMs += arrivedAt - lastResumeAt; - } - - try { - await throwIfAborted(options?.signal, stream); + const iterator = stream[Symbol.asyncIterator](); + try { + for (;;) { + const next = await nextStreamPart(iterator, stream, stallTimeoutMs, requestSignal, stallAbort); + if (next === 'stalled') { + throw new APITimeoutError( + `The API stream stalled: no data received for ${stallTimeoutMs}ms.` + + formatFinishReasonHint(stream) + + ` Provider: ${provider.name}, model: ${provider.modelName}`, + ); + } + if (next.done === true) { + break; + } + const part = next.value; + const arrivedAt = Date.now(); + if (firstPartAt === undefined) { + firstPartAt = arrivedAt; + } else { + serverDecodeMs += arrivedAt - lastResumeAt; + } - // Notify raw part callback (deep copy to avoid aliasing mutations). - if (callbacks?.onMessagePart !== undefined) { - await callbacks.onMessagePart(deepCopyPart(part)); + try { await throwIfAborted(options?.signal, stream); - } - // Index-based routing for parallel tool call argument deltas. - // When a ToolCallPart arrives with an index referring to a tool call - // that is NOT the currently-pending one, append it directly to the - // correct ToolCall in message.toolCalls instead of relying on sequential - // merging. This prevents argument cross-contamination across parallel calls. - if ( - isToolCallPart(part) && - part.index !== undefined && - !isPendingToolCallAtIndex(pendingPart, part.index) - ) { - const arrayIdx = toolCallIndexMap.get(part.index); - if (arrayIdx !== undefined) { - const target = message.toolCalls[arrayIdx]; - if (target !== undefined && part.argumentsPart !== null) { - target.arguments = - target.arguments === null - ? part.argumentsPart - : target.arguments + part.argumentsPart; + // Notify raw part callback (deep copy to avoid aliasing mutations). + if (callbacks?.onMessagePart !== undefined) { + await callbacks.onMessagePart(deepCopyPart(part)); + await throwIfAborted(options?.signal, stream); + } + + // Index-based routing for parallel tool call argument deltas. + // When a ToolCallPart arrives with an index referring to a tool call + // that is NOT the currently-pending one, append it directly to the + // correct ToolCall in message.toolCalls instead of relying on sequential + // merging. This prevents argument cross-contamination across parallel calls. + if ( + isToolCallPart(part) && + part.index !== undefined && + !isPendingToolCallAtIndex(pendingPart, part.index) + ) { + const arrayIdx = toolCallIndexMap.get(part.index); + if (arrayIdx !== undefined) { + const target = message.toolCalls[arrayIdx]; + if (target !== undefined && part.argumentsPart !== null) { + target.arguments = + target.arguments === null + ? part.argumentsPart + : target.arguments + part.argumentsPart; + } + continue; } - continue; + // Unknown index — fall through to the sequential logic as a safety net. } - // Unknown index — fall through to the sequential logic as a safety net. - } - if (pendingPart === null) { - pendingPart = part; - } else if (!mergeInPlace(pendingPart, part)) { - // Could not merge — flush the pending part and start a new one. - // For parallel tool calls this happens when a new ToolCall header arrives - // while a previous ToolCall is still pending; the flush finalizes the - // previous tool call into `message.toolCalls`. - flushPart(message, pendingPart, toolCallIndexMap); - pendingPart = part; + if (pendingPart === null) { + pendingPart = part; + } else if (!mergeInPlace(pendingPart, part)) { + // Could not merge — flush the pending part and start a new one. + // For parallel tool calls this happens when a new ToolCall header arrives + // while a previous ToolCall is still pending; the flush finalizes the + // previous tool call into `message.toolCalls`. + flushPart(message, pendingPart, toolCallIndexMap); + pendingPart = part; + } + } finally { + lastResumeAt = Date.now(); + clientConsumeMs += lastResumeAt - arrivedAt; } - } finally { - lastResumeAt = Date.now(); - clientConsumeMs += lastResumeAt - arrivedAt; } + } catch (error) { + // `for await` closed the iterator automatically when the body threw; the + // manual loop must do it itself, or a throwing callback/merge leaks the + // provider connection. Best-effort, never awaited — the original error + // must not be masked by a hanging teardown. + void cancelStream(stream); + teardownIterator(iterator); + throw error; } await throwIfAborted(options?.signal, stream); @@ -286,6 +356,98 @@ async function cancelStream(stream: StreamedMessage): Promise { } catch {} } +/** + * Race a pending promise against the stall watchdog and the abort signal. + * + * Returns the settled value, or the `'stalled'` / `'aborted'` sentinels. The + * watchdog only arms for a finite positive `stallTimeoutMs` (`0` disables it). + * `Promise.race` subscribes to `pending` even when the watchdog wins, so a + * late settlement of the abandoned promise is always observed and cannot + * surface as an unhandled rejection. + */ +async function raceStallOrAbort( + pending: Promise, + stallTimeoutMs: number, + signal: AbortSignal, +): Promise { + if (signal.aborted) { + return 'aborted'; + } + + let timer: ReturnType | undefined; + let onAbort: (() => void) | undefined; + const watchdog = new Promise<'stalled' | 'aborted'>((resolve) => { + if (Number.isFinite(stallTimeoutMs) && stallTimeoutMs > 0) { + timer = setTimeout(() => { + resolve('stalled'); + }, stallTimeoutMs); + // The watchdog must never keep the process (or a test runner) alive. + (timer as { unref?: () => void }).unref?.(); + } + onAbort = () => { + resolve('aborted'); + }; + signal.addEventListener('abort', onAbort, { once: true }); + }); + + try { + return await Promise.race([pending, watchdog]); + } finally { + if (timer !== undefined) { + clearTimeout(timer); + } + if (onAbort !== undefined) { + signal.removeEventListener('abort', onAbort); + } + } +} + +/** + * Await the next stream part, bounded by the stall watchdog and the abort + * signal. + * + * Returns the iterator result, or the `'stalled'` sentinel when no part + * arrived within `stallTimeoutMs`. On stall the linked abort controller fires + * first — providers forward the signal to their HTTP client, so this kills + * the dead connection and settles the in-flight iteration. An abort observed + * while waiting throws the standard abort error — without this race, + * cancelling a stalled stream would hang until the next part arrived. + * + * All teardown here is fire-and-forget: a faulty provider whose `cancel()` or + * `return()` never settles must not hang the stall path — the very failure + * this watchdog exists to escape. + */ +async function nextStreamPart( + iterator: AsyncIterator, + stream: StreamedMessage, + stallTimeoutMs: number, + signal: AbortSignal, + stallAbort: AbortController, +): Promise | 'stalled'> { + const outcome = await raceStallOrAbort(iterator.next(), stallTimeoutMs, signal); + if (outcome === 'aborted') { + void cancelStream(stream); + teardownIterator(iterator); + throwAbortError(); + } + if (outcome === 'stalled') { + stallAbort.abort(); + void cancelStream(stream); + teardownIterator(iterator); + return 'stalled'; + } + return outcome; +} + +/** + * Best-effort generator teardown. Never awaited: a provider iterator that + * ignores the abort signal would otherwise hang the stall path — the very + * failure this watchdog exists to escape. + */ +function teardownIterator(iterator: AsyncIterator): void { + void Promise.resolve(iterator.return?.()).catch(() => undefined); +} + async function throwIfAborted(signal?: AbortSignal, stream?: StreamedMessage): Promise { if (!signal?.aborted) { return; diff --git a/packages/kosong/src/index.ts b/packages/kosong/src/index.ts index cd0440637b..52f91c9081 100644 --- a/packages/kosong/src/index.ts +++ b/packages/kosong/src/index.ts @@ -56,7 +56,7 @@ export type { } from './catalog'; // Core functions -export { generate } from './generate'; +export { generate, DEFAULT_STREAM_STALL_TIMEOUT_MS } from './generate'; export type { GenerateCallbacks, GenerateResult } from './generate'; // Tool wire schema diff --git a/packages/kosong/src/provider.ts b/packages/kosong/src/provider.ts index 905b12c153..734c074bc8 100644 --- a/packages/kosong/src/provider.ts +++ b/packages/kosong/src/provider.ts @@ -179,6 +179,16 @@ export interface GenerateOptions { * least one part was streamed, or `undefined` for an empty stream. */ onStreamEnd?: (stats?: StreamDecodeStats) => void; + /** + * Maximum time in milliseconds to wait for the response — headers first, + * then each streamed part — before declaring the request stalled. The + * budget resets on every part, so slow generations are unaffected — only a + * complete silence longer than this trips the watchdog. On a stall the + * connection is torn down and {@link generate | generate()} throws an + * `APITimeoutError` (retryable). Defaults to + * `DEFAULT_STREAM_STALL_TIMEOUT_MS`; set to `0` to disable. + */ + streamStallTimeoutMs?: number; } /** diff --git a/packages/kosong/test/e2e/stream-stall.test.ts b/packages/kosong/test/e2e/stream-stall.test.ts new file mode 100644 index 0000000000..da82f60382 --- /dev/null +++ b/packages/kosong/test/e2e/stream-stall.test.ts @@ -0,0 +1,82 @@ +import * as node_http from 'node:http'; +import type { AddressInfo } from 'node:net'; + +import { APITimeoutError } from '#/errors'; +import { generate } from '#/generate'; +import type { Message } from '#/message'; +import { KimiChatProvider } from '#/providers/kimi'; +import { describe, expect, it } from 'vitest'; + +/** + * End-to-end stall coverage: a real `KimiChatProvider` (OpenAI SDK over HTTP) + * against a local server that sends response headers and one SSE chunk, then + * holds the socket open forever — the exact mid-stream stall that used to + * wedge sessions. `generate()` must fail with `APITimeoutError` once the + * inactivity budget elapses instead of hanging. + */ + +const USER_MSG: Message = { + role: 'user', + content: [{ type: 'text', text: 'hi' }], + toolCalls: [], +}; + +const FIRST_CHUNK = { + id: 'chatcmpl-stall', + object: 'chat.completion.chunk', + created: 0, + model: 'test', + choices: [{ index: 0, delta: { role: 'assistant', content: 'STALLED: ' } }], +}; + +interface HangingServer { + readonly baseUrl: string; + close(): Promise; +} + +async function createHangingServer(): Promise { + const server = node_http.createServer((req, res) => { + res.writeHead(200, { + 'cache-control': 'no-cache', + connection: 'keep-alive', + 'content-type': 'text/event-stream; charset=utf-8', + }); + res.write(`data: ${JSON.stringify(FIRST_CHUNK)}\n\n`); + // Never write again, never end — a dead mid-stream connection. + }); + await new Promise((resolve, reject) => { + server.once('error', reject); + server.listen(0, '127.0.0.1', () => resolve()); + }); + const { port } = server.address() as AddressInfo; + return { + baseUrl: `http://127.0.0.1:${port}/v1`, + close: () => + new Promise((resolve) => { + server.closeAllConnections(); + server.close(() => resolve()); + }), + }; +} + +describe('generate() stream stall (live HTTP)', () => { + it('times out a stalled chat-completions stream instead of hanging', async () => { + const server = await createHangingServer(); + try { + const provider = new KimiChatProvider({ + model: 'test', + apiKey: 'test-key', + baseUrl: server.baseUrl, + stream: true, + }); + const startedAt = Date.now(); + + await expect( + generate(provider, '', [], [USER_MSG], undefined, { streamStallTimeoutMs: 100 }), + ).rejects.toBeInstanceOf(APITimeoutError); + expect(Date.now() - startedAt).toBeLessThan(10_000); + } finally { + await server.close(); + } + }, 15_000); +}); diff --git a/packages/kosong/test/generate-stall.test.ts b/packages/kosong/test/generate-stall.test.ts new file mode 100644 index 0000000000..7743acf638 --- /dev/null +++ b/packages/kosong/test/generate-stall.test.ts @@ -0,0 +1,231 @@ +import { generate, DEFAULT_STREAM_STALL_TIMEOUT_MS } from '#/generate'; +import { APITimeoutError } from '#/errors'; +import type { Message, StreamedMessagePart } from '#/message'; +import type { + ChatProvider, + FinishReason, + GenerateOptions, + StreamedMessage, + ThinkingEffort, +} from '#/provider'; +import type { Tool } from '#/tool'; +import type { TokenUsage } from '#/usage'; +import { describe, expect, it } from 'vitest'; + +const USER_MSG: Message = { + role: 'user', + content: [{ type: 'text', text: 'hi' }], + toolCalls: [], +}; + +function textPart(text: string): StreamedMessagePart { + return { type: 'text', text }; +} + +/** + * A StreamedMessage whose iteration is fully scripted: each step either + * yields a part after `delayMs` or pends forever (simulating a dead + * connection that never delivers and never closes). Records `return()` + * calls so tests can assert the stream was torn down. + */ +class ScriptedStream implements StreamedMessage { + readonly id = 'scripted'; + readonly usage: TokenUsage | null = null; + readonly finishReason: FinishReason | null = null; + readonly rawFinishReason: string | null = null; + + returned = false; + + constructor(private readonly steps: Array<{ part: StreamedMessagePart; delayMs: number } | 'hang'>) {} + + [Symbol.asyncIterator](): AsyncIterator { + const steps = this.steps; + const onReturn = (): void => { + this.returned = true; + }; + let index = 0; + return { + next(): Promise> { + const step = steps[index]; + index += 1; + if (step === undefined) { + return Promise.resolve({ done: true, value: undefined }); + } + if (step === 'hang') { + return new Promise>(() => {}); + } + return new Promise((resolve) => { + setTimeout(() => resolve({ done: false, value: step.part }), step.delayMs); + }); + }, + return(): Promise> { + onReturn(); + return Promise.resolve({ done: true, value: undefined }); + }, + }; + } +} + +class ScriptedProvider implements ChatProvider { + readonly name = 'scripted'; + readonly modelName = 'scripted-model'; + readonly thinkingEffort: ThinkingEffort | null = null; + + constructor(readonly stream: ScriptedStream) {} + + generate( + _systemPrompt: string, + _tools: Tool[], + _history: Message[], + _options?: GenerateOptions, + ): Promise { + return Promise.resolve(this.stream); + } + + withThinking(): ChatProvider { + return this; + } +} + +/** A provider whose generate() never resolves — the response-headers stall. */ +class HangingGenerateProvider implements ChatProvider { + readonly name = 'hanging-generate'; + readonly modelName = 'hanging-model'; + readonly thinkingEffort: ThinkingEffort | null = null; + + generate(): Promise { + return new Promise(() => {}); + } + + withThinking(): ChatProvider { + return this; + } +} + +describe('generate() stream-stall watchdog', () => { + it('fails with APITimeoutError when the stream stalls mid-generation', async () => { + const stream = new ScriptedStream([{ part: textPart('partial'), delayMs: 0 }, 'hang']); + const startedAt = Date.now(); + + await expect( + generate(new ScriptedProvider(stream), '', [], [USER_MSG], undefined, { + streamStallTimeoutMs: 50, + }), + ).rejects.toBeInstanceOf(APITimeoutError); + + // The watchdog fired near the configured budget, not the 5min default. + expect(Date.now() - startedAt).toBeLessThan(5_000); + // The stalled stream was torn down, not left dangling. + expect(stream.returned).toBe(true); + }); + + it('resets the inactivity budget on every part', async () => { + const stream = new ScriptedStream([ + { part: textPart('a'), delayMs: 30 }, + { part: textPart('b'), delayMs: 30 }, + { part: textPart('c'), delayMs: 30 }, + ]); + + const result = await generate(new ScriptedProvider(stream), '', [], [USER_MSG], undefined, { + streamStallTimeoutMs: 80, + }); + + // Total runtime (~90ms) exceeds the 80ms budget, but no single gap does. + expect(result.message.content).toEqual([{ type: 'text', text: 'abc' }]); + }); + + it('applies the watchdog to the wait for the first part', async () => { + const stream = new ScriptedStream(['hang']); + + await expect( + generate(new ScriptedProvider(stream), '', [], [USER_MSG], undefined, { + streamStallTimeoutMs: 50, + }), + ).rejects.toBeInstanceOf(APITimeoutError); + }); + + it('covers the response-headers wait inside provider.generate()', async () => { + const startedAt = Date.now(); + + await expect( + generate(new HangingGenerateProvider(), '', [], [USER_MSG], undefined, { + streamStallTimeoutMs: 50, + }), + ).rejects.toBeInstanceOf(APITimeoutError); + expect(Date.now() - startedAt).toBeLessThan(5_000); + }); + + it('does not hang when stream teardown never settles', async () => { + // A faulty provider whose cancel()/return() pend forever must not block + // the stall path — teardown is best-effort, the timeout still wins. + const stream = new ScriptedStream(['hang']); + const iterator = stream[Symbol.asyncIterator](); + const neverSettling: AsyncIterator = { + next: () => iterator.next(), + return: () => new Promise>(() => {}), + }; + stream[Symbol.asyncIterator] = () => neverSettling; + + const startedAt = Date.now(); + await expect( + generate(new ScriptedProvider(stream), '', [], [USER_MSG], undefined, { + streamStallTimeoutMs: 50, + }), + ).rejects.toBeInstanceOf(APITimeoutError); + expect(Date.now() - startedAt).toBeLessThan(5_000); + }); + + it('aborting mid-stall rejects promptly with AbortError', async () => { + const stream = new ScriptedStream(['hang']); + const controller = new AbortController(); + const startedAt = Date.now(); + + const pending = generate(new ScriptedProvider(stream), '', [], [USER_MSG], undefined, { + // Far larger than the abort delay: only the abort race can settle this. + streamStallTimeoutMs: 60_000, + signal: controller.signal, + }); + setTimeout(() => controller.abort(), 30); + + await expect(pending).rejects.toMatchObject({ name: 'AbortError' }); + expect(Date.now() - startedAt).toBeLessThan(5_000); + expect(stream.returned).toBe(true); + }); + + it('streamStallTimeoutMs: 0 disables the watchdog', async () => { + const stream = new ScriptedStream([{ part: textPart('ok'), delayMs: 30 }]); + + const result = await generate(new ScriptedProvider(stream), '', [], [USER_MSG], undefined, { + streamStallTimeoutMs: 0, + }); + + expect(result.message.content).toEqual([{ type: 'text', text: 'ok' }]); + }); + + it('exposes a generous default budget', () => { + expect(DEFAULT_STREAM_STALL_TIMEOUT_MS).toBeGreaterThanOrEqual(60_000); + }); + + it('closes the stream iterator when part processing throws', async () => { + // `for await` closed the iterator automatically on a throwing body + // (AsyncIteratorClose); the manual watchdog loop must preserve that, or a + // throwing callback leaks the provider connection. + const stream = new ScriptedStream([{ part: textPart('x'), delayMs: 0 }, 'hang']); + + await expect( + generate( + new ScriptedProvider(stream), + '', + [], + [USER_MSG], + { + onMessagePart: () => { + throw new Error('boom'); + }, + }, + { streamStallTimeoutMs: 0 }, + ), + ).rejects.toThrow('boom'); + expect(stream.returned).toBe(true); + }); +});