diff --git a/.changeset/stream-idle-watchdog.md b/.changeset/stream-idle-watchdog.md new file mode 100644 index 0000000000..25c9e37774 --- /dev/null +++ b/.changeset/stream-idle-watchdog.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Fail loudly and automatically retry when a provider stream stalls mid-flight. The underlying OpenAI-compatible client clears its request timeout as soon as response headers arrive, so a silent SSE stall used to hang `generate()` — and the whole turn — forever. The new inter-chunk idle watchdog (default 180 s, tunable via `KIMI_STREAM_IDLE_TIMEOUT_MS`) is applied to both stream loops in the CLI bundle — `packages/kosong/src/generate.ts` (used by the ACP path: kimi-code → kimi-code-sdk → agent-core → kosong) and `packages/agent-core-v2/src/app/llmProtocol/generate.ts` (used by the kap-server path) — cancelling the stream and raising an `APITimeoutError` subclass (`StreamIdleTimeoutError`) that carries elapsed time and the Kimi `x-trace-id`. The loop's step-retry plugin classifies it as retryable and re-drives the failed step, so the user sees a `step.retrying` recovery instead of a permanent hang. diff --git a/packages/agent-core-v2/src/app/llmProtocol/generate.ts b/packages/agent-core-v2/src/app/llmProtocol/generate.ts index cbe70260eb..1341c8f531 100644 --- a/packages/agent-core-v2/src/app/llmProtocol/generate.ts +++ b/packages/agent-core-v2/src/app/llmProtocol/generate.ts @@ -1,4 +1,17 @@ -import { APIEmptyResponseError } from './errors'; +/** + * Streams one assistant message from a chat provider and assembles the parts + * into a complete `GenerateResult`, merging interleaved content and routing + * parallel tool-call argument deltas to the correct call. + * + * Also owns the stream idle watchdog: the underlying HTTP client's request + * timeout is cleared once response headers arrive, so a stream that goes + * silent mid-flight would otherwise block `for await` forever. When no chunk + * arrives within `KIMI_STREAM_IDLE_TIMEOUT_MS` (default 180s) the watchdog + * aborts the provider request through the merged AbortSignal — actually + * closing the stalled connection — and throws `StreamIdleTimeoutError`, an + * `APITimeoutError` subclass the loop's retry plugin classifies as retryable. + */ +import { APIEmptyResponseError, APITimeoutError } from './errors'; import { isContentPart, isToolCall, @@ -50,7 +63,15 @@ export async function generate( : tools; options?.onRequestStart?.(); - const stream = await provider.generate(systemPrompt, wireTools, history, options); + const watchdog = new AbortController(); + const providerOptions: GenerateOptions = { + ...options, + signal: + options?.signal === undefined + ? watchdog.signal + : AbortSignal.any([options.signal, watchdog.signal]), + }; + const stream = await provider.generate(systemPrompt, wireTools, history, providerOptions); if (stream.traceId !== undefined) { options?.onTraceId?.(stream.traceId); } @@ -62,7 +83,7 @@ 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)) { const arrivedAt = Date.now(); if (firstPartAt === undefined) { firstPartAt = arrivedAt; @@ -175,6 +196,82 @@ type CancelableStream = StreamedMessage & { return?: () => unknown; }; +const STREAM_IDLE_TIMEOUT_MS = (() => { + const raw = Number(process.env['KIMI_STREAM_IDLE_TIMEOUT_MS']); + return Number.isFinite(raw) && raw > 0 ? raw : 180_000; +})(); + +export class StreamIdleTimeoutError extends APITimeoutError { + readonly idleMs: number; + readonly elapsedMs: number; + readonly traceId: string | null; + + constructor( + providerName: string, + modelName: string, + idleMs: number, + elapsedMs: number, + traceId: string | null, + ) { + const traceHint = traceId === null ? '' : ` traceId=${traceId}`; + super( + `LLM stream stalled: no data received for ${Math.round(idleMs / 1000)}s ` + + `(provider: ${providerName}, model: ${modelName}, ` + + `elapsedMs: ${elapsedMs}${traceHint}). ` + + 'The connection was abandoned to avoid hanging forever.', + ); + this.name = 'StreamIdleTimeoutError'; + this.idleMs = idleMs; + this.elapsedMs = elapsedMs; + this.traceId = traceId; + } +} + +async function* withStreamIdleTimeout( + stream: StreamedMessage, + provider: ChatProvider, + watchdog: AbortController, +): AsyncGenerator { + const iterator = stream[Symbol.asyncIterator](); + const startedAt = Date.now(); + try { + while (true) { + let timer: ReturnType | undefined; + const timeout = new Promise((_, reject) => { + timer = setTimeout(() => { + reject( + new StreamIdleTimeoutError( + provider.name, + provider.modelName, + STREAM_IDLE_TIMEOUT_MS, + Date.now() - startedAt, + stream.traceId ?? null, + ), + ); + }, STREAM_IDLE_TIMEOUT_MS); + }); + const next = iterator.next(); + let result: IteratorResult; + try { + result = await Promise.race([next, timeout]); + } catch (error) { + if (error instanceof StreamIdleTimeoutError) { + next.catch(() => {}); + watchdog.abort(error); + await cancelStream(stream); + } + throw error; + } finally { + clearTimeout(timer); + } + if (result.done === true) return; + yield result.value; + } + } finally { + void Promise.resolve(iterator.return?.()).catch(() => {}); + } +} + function throwAbortError(): never { throw new DOMException('The operation was aborted.', 'AbortError'); } diff --git a/packages/agent-core-v2/test/app/llmProtocol/stream-idle-timeout.test.ts b/packages/agent-core-v2/test/app/llmProtocol/stream-idle-timeout.test.ts new file mode 100644 index 0000000000..74a5cbe2cd --- /dev/null +++ b/packages/agent-core-v2/test/app/llmProtocol/stream-idle-timeout.test.ts @@ -0,0 +1,190 @@ +/** + * Stream idle watchdog: a provider stream that goes silent mid-flight must + * fail with StreamIdleTimeoutError instead of hanging `generate()` forever, + * and must be classified as retryable so the loop's step-retry plugin + * re-drives the failed step. + */ +import type { ChatProvider, StreamedMessage } from '#/app/llmProtocol/provider'; +import type { StreamedMessagePart } from '#/app/llmProtocol/message'; +import { + APITimeoutError, + isRetryableGenerateError, +} from '#/app/llmProtocol/errors'; +import { describe, expect, it } from 'vitest'; + +function makeProvider( + stream: StreamedMessage, + onGenerate?: (options?: { signal?: AbortSignal }) => void, +): ChatProvider { + return { + name: 'fake', + modelName: 'fake-model', + generate: async ( + _systemPrompt: unknown, + _tools: unknown, + _history: unknown, + options?: { signal?: AbortSignal }, + ) => { + onGenerate?.(options); + return stream; + }, + } as unknown as ChatProvider; +} + +function stalledStream( + parts: StreamedMessagePart[], + onCancel?: () => void, + traceId: string | null = null, + onIteratorReturn?: () => void, +): StreamedMessage { + let i = 0; + return { + id: 'msg-1', + usage: null, + finishReason: null, + rawFinishReason: null, + traceId, + cancel: () => onCancel?.(), + [Symbol.asyncIterator]() { + return { + next(): Promise> { + if (i < parts.length) { + return Promise.resolve({ done: false, value: parts[i++]! }); + } + return new Promise(() => {}); // silent forever + }, + return(): Promise> { + onIteratorReturn?.(); + return Promise.resolve({ done: true, value: undefined }); + }, + }; + }, + } as StreamedMessage; +} + +function healthyStream(parts: StreamedMessagePart[]): StreamedMessage { + let i = 0; + return { + id: 'msg-1', + usage: null, + finishReason: 'completed', + rawFinishReason: 'stop', + [Symbol.asyncIterator]() { + return { + next(): Promise> { + if (i < parts.length) { + return Promise.resolve({ done: false, value: parts[i++]! }); + } + return Promise.resolve({ done: true, value: undefined }); + }, + }; + }, + } as StreamedMessage; +} + +async function importGenerateWithTimeout(ms: number) { + process.env['KIMI_STREAM_IDLE_TIMEOUT_MS'] = String(ms); + const mod = await import('#/app/llmProtocol/generate'); + return mod; +} + +describe('generate() stream idle watchdog', () => { + it('throws StreamIdleTimeoutError when the stream goes silent after first part', async () => { + const { generate, StreamIdleTimeoutError } = await importGenerateWithTimeout(200); + let cancelled = false; + const stream = stalledStream( + [{ type: 'think', think: 'partial thinking' }], + () => { cancelled = true; }, + ); + await expect( + generate(makeProvider(stream), 'system', [], [ + { role: 'user', content: [{ type: 'text', text: 'hi' }], toolCalls: [] }, + ]), + ).rejects.toBeInstanceOf(StreamIdleTimeoutError); + expect(cancelled).toBe(true); + }); + + it('classifies the timeout error as a retryable APITimeoutError so step-retry drives recovery', async () => { + const { generate, StreamIdleTimeoutError } = await importGenerateWithTimeout(200); + const stream = stalledStream( + [{ type: 'think', think: 'partial thinking' }], + undefined, + 'trace-abc123', + ); + let captured: unknown; + try { + await generate(makeProvider(stream), 'system', [], [ + { role: 'user', content: [{ type: 'text', text: 'hi' }], toolCalls: [] }, + ]); + } catch (error) { + captured = error; + } + expect(captured).toBeInstanceOf(StreamIdleTimeoutError); + expect(captured).toBeInstanceOf(APITimeoutError); + expect(isRetryableGenerateError(captured)).toBe(true); + const err = captured as InstanceType; + expect(err.traceId).toBe('trace-abc123'); + expect(err.idleMs).toBe(200); + expect(err.elapsedMs).toBeGreaterThanOrEqual(0); + expect(err.message).toContain('traceId=trace-abc123'); + }); + + it('does not interfere with a healthy stream', async () => { + const { generate } = await importGenerateWithTimeout(200); + const stream = healthyStream([ + { type: 'text', text: 'hello ' }, + { type: 'text', text: 'world' }, + ]); + const result = await generate(makeProvider(stream), 'system', [], [ + { role: 'user', content: [{ type: 'text', text: 'hi' }], toolCalls: [] }, + ]); + expect(result.message.content.some((p) => p.type === 'text')).toBe(true); + }); + + it('aborts the provider request and closes the iterator when the watchdog fires', async () => { + const { generate, StreamIdleTimeoutError } = await importGenerateWithTimeout(200); + let providerSignal: AbortSignal | undefined; + let iteratorReturned = false; + const stream = stalledStream( + [{ type: 'think', think: 'partial thinking' }], + undefined, + null, + () => { + iteratorReturned = true; + }, + ); + const provider = makeProvider(stream, (options) => { + providerSignal = options?.signal; + }); + await expect( + generate(provider, 'system', [], [ + { role: 'user', content: [{ type: 'text', text: 'hi' }], toolCalls: [] }, + ]), + ).rejects.toBeInstanceOf(StreamIdleTimeoutError); + expect(providerSignal?.aborted).toBe(true); + expect(providerSignal?.reason).toBeInstanceOf(StreamIdleTimeoutError); + expect(iteratorReturned).toBe(true); + }); + + it('still delivers caller aborts to the provider through the merged signal', async () => { + const { generate } = await importGenerateWithTimeout(200); + let providerSignal: AbortSignal | undefined; + const stream = healthyStream([{ type: 'text', text: 'hello' }]); + const provider = makeProvider(stream, (options) => { + providerSignal = options?.signal; + }); + const caller = new AbortController(); + await generate( + provider, + 'system', + [], + [{ role: 'user', content: [{ type: 'text', text: 'hi' }], toolCalls: [] }], + undefined, + { signal: caller.signal }, + ); + expect(providerSignal).toBeDefined(); + expect(providerSignal?.aborted).toBe(false); + caller.abort(); + expect(providerSignal?.aborted).toBe(true); + }); +}); diff --git a/packages/kosong/src/generate.ts b/packages/kosong/src/generate.ts index 626dc30736..5925842400 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, @@ -117,7 +117,19 @@ export async function generate( : tools; options?.onRequestStart?.(); - const stream = await provider.generate(systemPrompt, wireTools, history, options); + // Provider streams expose no cancellation handle of their own, so the idle + // watchdog must tear down a stalled transport through the request's + // AbortSignal: merge the caller's signal with a watchdog-owned controller + // and hand the merged signal to the provider. + const watchdog = new AbortController(); + const providerOptions: GenerateOptions = { + ...options, + signal: + options?.signal === undefined + ? watchdog.signal + : AbortSignal.any([options.signal, watchdog.signal]), + }; + 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 +154,7 @@ 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)) { const arrivedAt = Date.now(); if (firstPartAt === undefined) { firstPartAt = arrivedAt; @@ -270,6 +282,108 @@ type CancelableStream = StreamedMessage & { return?: () => unknown; }; +/** + * Guard the gap between streamed chunks with an inactivity deadline. + * + * The provider SSE stream has no read timeout of its own — the openai-node + * HTTP client's request timeout is cleared as soon as response headers arrive + * (see `openai/src/client.ts` `fetchWithTimeout`, wrapping the fetch call in + * `try { … } finally { clearTimeout(timeout) }` and resolving on headers) — + * so a connection that goes silent mid-stream would otherwise block + * `for await` forever and leave the whole turn hanging. Node's built-in + * `undici` has a `bodyTimeout`, but it is reset by every SSE heartbeat frame + * (empty `choices:[]`), so a heartbeat-only channel that stops emitting real + * deltas never trips it either. + * + * If no chunk arrives within the deadline the stream is cancelled and a + * descriptive `APITimeoutError` subclass is thrown so the step fails loudly + * and the loop's retry plugin can re-drive it (see + * `isRetryableGenerateError`). + */ +const STREAM_IDLE_TIMEOUT_MS = (() => { + const raw = Number(process.env['KIMI_STREAM_IDLE_TIMEOUT_MS']); + return Number.isFinite(raw) && raw > 0 ? raw : 180_000; +})(); + +export class StreamIdleTimeoutError extends APITimeoutError { + readonly idleMs: number; + readonly elapsedMs: number; + readonly traceId: string | null; + + constructor( + providerName: string, + modelName: string, + idleMs: number, + elapsedMs: number, + traceId: string | null, + ) { + const traceHint = traceId === null ? '' : ` traceId=${traceId}`; + super( + `LLM stream stalled: no data received for ${Math.round(idleMs / 1000)}s ` + + `(provider: ${providerName}, model: ${modelName}, ` + + `elapsedMs: ${elapsedMs}${traceHint}). ` + + 'The connection was abandoned to avoid hanging forever.', + ); + this.name = 'StreamIdleTimeoutError'; + this.idleMs = idleMs; + this.elapsedMs = elapsedMs; + this.traceId = traceId; + } +} + +async function* withStreamIdleTimeout( + stream: StreamedMessage, + provider: ChatProvider, + watchdog: AbortController, +): AsyncGenerator { + const iterator = stream[Symbol.asyncIterator](); + const startedAt = Date.now(); + try { + while (true) { + let timer: ReturnType | undefined; + const timeout = new Promise((_, reject) => { + timer = setTimeout(() => { + reject( + new StreamIdleTimeoutError( + provider.name, + provider.modelName, + STREAM_IDLE_TIMEOUT_MS, + Date.now() - startedAt, + stream.traceId ?? null, + ), + ); + }, STREAM_IDLE_TIMEOUT_MS); + }); + const next = iterator.next(); + let result: IteratorResult; + try { + result = await Promise.race([next, timeout]); + } catch (error) { + if (error instanceof StreamIdleTimeoutError) { + // The dangling next() will reject once the transport is aborted + // below; swallow it so it never surfaces as an unhandled rejection. + next.catch(() => {}); + // Abort the underlying request so the stalled connection is truly + // closed instead of leaking until the server gives up on it. + watchdog.abort(error); + await cancelStream(stream); + } + throw error; + } finally { + clearTimeout(timer); + } + if (result.done === true) return; + yield result.value; + } + } finally { + // Close the provider iterator on every exit path (timeout, consumer + // throw, early break) so generator cleanup runs. Never await it: on a + // stalled stream `return()` stays pending until the transport abort + // settles the dangling read, and awaiting would re-introduce the hang. + void Promise.resolve(iterator.return?.()).catch(() => {}); + } +} + function throwAbortError(): never { throw new DOMException('The operation was aborted.', 'AbortError'); } diff --git a/packages/kosong/src/index.ts b/packages/kosong/src/index.ts index 10a7af7304..b726f428df 100644 --- a/packages/kosong/src/index.ts +++ b/packages/kosong/src/index.ts @@ -47,7 +47,7 @@ export { export type { Catalog, CatalogModel, CatalogModelEntry, CatalogProviderEntry } from './catalog'; // Core functions -export { generate } from './generate'; +export { generate, StreamIdleTimeoutError } from './generate'; export type { GenerateCallbacks, GenerateResult } from './generate'; // Tool wire schema diff --git a/packages/kosong/test/stream-idle-timeout.test.ts b/packages/kosong/test/stream-idle-timeout.test.ts new file mode 100644 index 0000000000..e2e2310483 --- /dev/null +++ b/packages/kosong/test/stream-idle-timeout.test.ts @@ -0,0 +1,207 @@ +/** + * Stream idle watchdog on kosong's `generate()`: a provider stream that goes + * silent mid-flight must fail with `StreamIdleTimeoutError` instead of + * hanging `generate()` forever, and must be classified as retryable so the + * agent-core loop's step-retry re-drives the failed step. + * + * kosong is the shared LLM-streaming layer used by the ACP path (kimi-code + * TUI + `kimi acp` → kimi-code-sdk → agent-core → kosong.generate). The + * agent-core-v2 stack ships its own equivalent watchdog in + * `packages/agent-core-v2/src/app/llmProtocol/generate.ts` for the + * kap-server path — the same failure signature must fail loudly on both. + */ +import { APITimeoutError, isRetryableGenerateError } from '#/errors'; +import type { Message, StreamedMessagePart } from '#/message'; +import type { + ChatProvider, + GenerateOptions, + StreamedMessage, + ThinkingEffort, +} from '#/provider'; +import type { Tool } from '#/tool'; +import { describe, expect, it } from 'vitest'; + +function makeProvider( + stream: StreamedMessage, + onGenerate?: (options?: GenerateOptions) => void, +): ChatProvider { + return { + name: 'fake', + modelName: 'fake-model', + thinkingEffort: null, + generate: async ( + _systemPrompt: string, + _tools: Tool[], + _history: Message[], + options?: GenerateOptions, + ): Promise => { + onGenerate?.(options); + return stream; + }, + withThinking(_effort: ThinkingEffort): ChatProvider { + return this; + }, + }; +} + +function stalledStream( + parts: StreamedMessagePart[], + onCancel?: () => void, + traceId: string | null = null, + onIteratorReturn?: () => void, +): StreamedMessage { + let i = 0; + return { + id: 'msg-1', + usage: null, + finishReason: null, + rawFinishReason: null, + traceId, + cancel: () => onCancel?.(), + [Symbol.asyncIterator]() { + return { + next(): Promise> { + if (i < parts.length) { + return Promise.resolve({ done: false, value: parts[i++]! }); + } + return new Promise(() => {}); // silent forever + }, + return(): Promise> { + onIteratorReturn?.(); + return Promise.resolve({ done: true, value: undefined }); + }, + }; + }, + } as unknown as StreamedMessage; +} + +function healthyStream(parts: StreamedMessagePart[]): StreamedMessage { + let i = 0; + return { + id: 'msg-1', + usage: null, + finishReason: 'completed', + rawFinishReason: 'stop', + [Symbol.asyncIterator]() { + return { + next(): Promise> { + if (i < parts.length) { + return Promise.resolve({ done: false, value: parts[i++]! }); + } + return Promise.resolve({ done: true, value: undefined }); + }, + }; + }, + } as unknown as StreamedMessage; +} + +// All three cases below share the same 200ms watchdog window. The env var is +// read once at module init; setting it before the first import is enough. +async function importGenerateWithTimeout(ms: number) { + process.env['KIMI_STREAM_IDLE_TIMEOUT_MS'] = String(ms); + const mod = await import('#/generate'); + return mod; +} + +describe('kosong.generate() stream idle watchdog', () => { + it('throws StreamIdleTimeoutError when the stream goes silent after first part', async () => { + const { generate, StreamIdleTimeoutError } = await importGenerateWithTimeout(200); + let cancelled = false; + const stream = stalledStream( + [{ type: 'think', think: 'partial thinking' }], + () => { + cancelled = true; + }, + ); + await expect( + generate(makeProvider(stream), 'system', [], [ + { role: 'user', content: [{ type: 'text', text: 'hi' }], toolCalls: [] }, + ]), + ).rejects.toBeInstanceOf(StreamIdleTimeoutError); + expect(cancelled).toBe(true); + }); + + it('classifies the timeout error as a retryable APITimeoutError so step-retry drives recovery', async () => { + const { generate, StreamIdleTimeoutError } = await importGenerateWithTimeout(200); + const stream = stalledStream( + [{ type: 'think', think: 'partial thinking' }], + undefined, + 'trace-abc123', + ); + let captured: unknown; + try { + await generate(makeProvider(stream), 'system', [], [ + { role: 'user', content: [{ type: 'text', text: 'hi' }], toolCalls: [] }, + ]); + } catch (error) { + captured = error; + } + expect(captured).toBeInstanceOf(StreamIdleTimeoutError); + expect(captured).toBeInstanceOf(APITimeoutError); + expect(isRetryableGenerateError(captured)).toBe(true); + const err = captured as InstanceType; + expect(err.traceId).toBe('trace-abc123'); + expect(err.idleMs).toBe(200); + expect(err.elapsedMs).toBeGreaterThanOrEqual(0); + expect(err.message).toContain('traceId=trace-abc123'); + }); + + it('does not interfere with a healthy stream', async () => { + const { generate } = await importGenerateWithTimeout(200); + const stream = healthyStream([ + { type: 'text', text: 'hello ' }, + { type: 'text', text: 'world' }, + ]); + const result = await generate(makeProvider(stream), 'system', [], [ + { role: 'user', content: [{ type: 'text', text: 'hi' }], toolCalls: [] }, + ]); + expect(result.message.content.some((p) => p.type === 'text')).toBe(true); + }); + + it('aborts the provider request and closes the iterator when the watchdog fires', async () => { + const { generate, StreamIdleTimeoutError } = await importGenerateWithTimeout(200); + let providerSignal: AbortSignal | undefined; + let iteratorReturned = false; + const stream = stalledStream( + [{ type: 'think', think: 'partial thinking' }], + undefined, + null, + () => { + iteratorReturned = true; + }, + ); + const provider = makeProvider(stream, (options) => { + providerSignal = options?.signal; + }); + await expect( + generate(provider, 'system', [], [ + { role: 'user', content: [{ type: 'text', text: 'hi' }], toolCalls: [] }, + ]), + ).rejects.toBeInstanceOf(StreamIdleTimeoutError); + expect(providerSignal?.aborted).toBe(true); + expect(providerSignal?.reason).toBeInstanceOf(StreamIdleTimeoutError); + expect(iteratorReturned).toBe(true); + }); + + it('still delivers caller aborts to the provider through the merged signal', async () => { + const { generate } = await importGenerateWithTimeout(200); + let providerSignal: AbortSignal | undefined; + const stream = healthyStream([{ type: 'text', text: 'hello' }]); + const provider = makeProvider(stream, (options) => { + providerSignal = options?.signal; + }); + const caller = new AbortController(); + await generate( + provider, + 'system', + [], + [{ role: 'user', content: [{ type: 'text', text: 'hi' }], toolCalls: [] }], + undefined, + { signal: caller.signal }, + ); + expect(providerSignal).toBeDefined(); + expect(providerSignal?.aborted).toBe(false); + caller.abort(); + expect(providerSignal?.aborted).toBe(true); + }); +});