From 9312f948285d3f6107f250e13762f85601f3286f Mon Sep 17 00:00:00 2001 From: gcf Date: Fri, 17 Jul 2026 10:23:43 +0800 Subject: [PATCH 1/3] fix(agent-core-v2): watchdog and retry silent LLM stream stalls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The OpenAI-compatible client used for streaming completions clears its request timeout as soon as response headers arrive (`openai/src/client.ts` wraps the fetch in `try { … } finally { clearTimeout(timeout) }` and the fetch promise resolves on header receipt). Once headers are back the streaming body has no read timeout of its own, so a mid-stream stall blocks `for await` forever and hangs the whole turn — including cases where the assistant has already streamed a reasoning delta and then falls silent, and where the process holds pooled TCP connections that a NAT/LB has half-closed. Add an inactivity watchdog around the streamed message iterator: after `KIMI_STREAM_IDLE_TIMEOUT_MS` (default 180 s) without a new chunk, cancel the underlying stream and throw `StreamIdleTimeoutError`. The error now extends `APITimeoutError`, so `isRetryableGenerateError` recognises it and the loop's `stepRetry` plugin re-drives the failed step automatically; users see the same recoverable behaviour as for any transient network timeout instead of a permanent hang. The error message carries the elapsed stream time and, when available, the Kimi `x-trace-id` for server-side triage. --- .changeset/stream-idle-watchdog.md | 5 + .../src/app/llmProtocol/generate.ts | 90 ++++++++++++- .../llmProtocol/stream-idle-timeout.test.ts | 127 ++++++++++++++++++ 3 files changed, 220 insertions(+), 2 deletions(-) create mode 100644 .changeset/stream-idle-watchdog.md create mode 100644 packages/agent-core-v2/test/app/llmProtocol/stream-idle-timeout.test.ts diff --git a/.changeset/stream-idle-watchdog.md b/.changeset/stream-idle-watchdog.md new file mode 100644 index 0000000000..fd6e1039f5 --- /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 idle watchdog (default 180 s, tunable via `KIMI_STREAM_IDLE_TIMEOUT_MS`) cancels the stream and raises an `APITimeoutError` subclass carrying elapsed time and the Kimi `x-trace-id`, which the loop's step-retry plugin picks up to re-drive the failed step. diff --git a/packages/agent-core-v2/src/app/llmProtocol/generate.ts b/packages/agent-core-v2/src/app/llmProtocol/generate.ts index cbe70260eb..6433fe6fa3 100644 --- a/packages/agent-core-v2/src/app/llmProtocol/generate.ts +++ b/packages/agent-core-v2/src/app/llmProtocol/generate.ts @@ -1,4 +1,4 @@ -import { APIEmptyResponseError } from './errors'; +import { APIEmptyResponseError, APITimeoutError } from './errors'; import { isContentPart, isToolCall, @@ -62,7 +62,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)) { const arrivedAt = Date.now(); if (firstPartAt === undefined) { firstPartAt = arrivedAt; @@ -175,6 +175,92 @@ 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 HTTP client's + * request timeout is cleared as soon as response headers arrive (see + * `openai/src/client.ts` `fetchWithTimeout`, which wraps the fetch call in + * `try { … } finally { clearTimeout(timeout) }` and resolves once headers + * are received) — so a connection that goes silent mid-stream would + * otherwise block `for await` forever and leave the whole turn hanging. + * 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, +): AsyncGenerator { + const iterator = stream[Symbol.asyncIterator](); + const startedAt = Date.now(); + 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 stream is cancelled below; + // swallow it so it never surfaces as an unhandled rejection. + next.catch(() => {}); + await cancelStream(stream); + } + throw error; + } finally { + clearTimeout(timer); + } + if (result.done === true) return; + yield result.value; + } +} + 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..4ee94496b5 --- /dev/null +++ b/packages/agent-core-v2/test/app/llmProtocol/stream-idle-timeout.test.ts @@ -0,0 +1,127 @@ +/** + * 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): ChatProvider { + return { + name: 'fake', + modelName: 'fake-model', + generate: async () => stream, + } as unknown as ChatProvider; +} + +function stalledStream( + parts: StreamedMessagePart[], + onCancel?: () => void, + traceId: string | null = null, +): 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 + }, + }; + }, + } 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); + }); +}); From 7cc81d3bf6317cacfb2f4b194e25175c35228e29 Mon Sep 17 00:00:00 2001 From: gcf Date: Fri, 17 Jul 2026 10:43:13 +0800 Subject: [PATCH 2/3] fix(kosong): apply stream idle watchdog on the ACP hot path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ACP path (kimi-code TUI + `kimi acp` → @moonshot-ai/kimi-code-sdk → @moonshot-ai/agent-core → agent-core/turn/kosong-llm.ts) drives its LLM stream through kosong.generate(), not agent-core-v2.generate(). The earlier watchdog on agent-core-v2 fired only on the kap-server path (@moonshot-ai/agent-core-v2), so the ACP session that hung mid-stream still had no read-side timeout. Live proof: session 6871de3a-3499-4e94-94f7-39585a347f1b silently stalled after the first `think` chunk at 02:13:47 on a bundle that contained the agent-core-v2 watchdog symbols — no StreamIdleTimeoutError ever surfaced. Port the identical inter-chunk deadline guard into kosong: wrap the `for await (const part of stream)` loop with `withStreamIdleTimeout()` that races each iterator.next() against a `setTimeout` reject; on timeout, cancel the stream and throw a `StreamIdleTimeoutError` (extends kosong `APITimeoutError`) carrying the elapsed time and the `x-trace-id`. The existing `isRetryableGenerateError()` classifies it as retryable, so agent-core's `chatWithRetry()` re-drives the failed step via the standard `step.retrying` event — same recovery path the kap-server watchdog already exercises. Watchdog window shares the same `KIMI_STREAM_IDLE_TIMEOUT_MS` knob (default 180000 ms) so operators tune both stacks in lockstep. Test: packages/kosong/test/stream-idle-timeout.test.ts covers cancel, retryable classification with traceId propagation, and healthy-stream pass-through. --- .changeset/stream-idle-watchdog.md | 2 +- packages/kosong/src/generate.ts | 94 +++++++++++- packages/kosong/src/index.ts | 2 +- .../kosong/test/stream-idle-timeout.test.ts | 144 ++++++++++++++++++ 4 files changed, 238 insertions(+), 4 deletions(-) create mode 100644 packages/kosong/test/stream-idle-timeout.test.ts diff --git a/.changeset/stream-idle-watchdog.md b/.changeset/stream-idle-watchdog.md index fd6e1039f5..25c9e37774 100644 --- a/.changeset/stream-idle-watchdog.md +++ b/.changeset/stream-idle-watchdog.md @@ -2,4 +2,4 @@ "@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 idle watchdog (default 180 s, tunable via `KIMI_STREAM_IDLE_TIMEOUT_MS`) cancels the stream and raises an `APITimeoutError` subclass carrying elapsed time and the Kimi `x-trace-id`, which the loop's step-retry plugin picks up to re-drive the failed step. +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/kosong/src/generate.ts b/packages/kosong/src/generate.ts index 626dc30736..a4688197b8 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, @@ -142,7 +142,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)) { const arrivedAt = Date.now(); if (firstPartAt === undefined) { firstPartAt = arrivedAt; @@ -270,6 +270,96 @@ 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, +): AsyncGenerator { + const iterator = stream[Symbol.asyncIterator](); + const startedAt = Date.now(); + 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 stream is cancelled below; + // swallow it so it never surfaces as an unhandled rejection. + next.catch(() => {}); + await cancelStream(stream); + } + throw error; + } finally { + clearTimeout(timer); + } + if (result.done === true) return; + yield result.value; + } +} + 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..613f43f08e --- /dev/null +++ b/packages/kosong/test/stream-idle-timeout.test.ts @@ -0,0 +1,144 @@ +/** + * 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 { StreamedMessagePart } from '#/message'; +import type { ChatProvider, StreamedMessage, ThinkingEffort } from '#/provider'; +import type { Tool } from '#/tool'; +import type { Message } from '#/message'; +import { describe, expect, it } from 'vitest'; + +function makeProvider(stream: StreamedMessage): ChatProvider { + return { + name: 'fake', + modelName: 'fake-model', + thinkingEffort: null, + generate: async ( + _systemPrompt: string, + _tools: Tool[], + _history: Message[], + ): Promise => stream, + withThinking(_effort: ThinkingEffort): ChatProvider { + return this; + }, + }; +} + +function stalledStream( + parts: StreamedMessagePart[], + onCancel?: () => void, + traceId: string | null = null, +): 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 + }, + }; + }, + } 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); + }); +}); From cd11470025560d7d3242a67a4a6ce7b293f24331 Mon Sep 17 00:00:00 2001 From: gcf Date: Fri, 17 Jul 2026 22:54:51 +0800 Subject: [PATCH 3/3] fix(kosong,agent-core-v2): abort the transport when the stream idle watchdog fires Provider stream classes expose no cancellation handle, so the previous cancelStream() call was a no-op on every real provider and each stall leaked its HTTP connection until the server gave up. The watchdog now owns an AbortController merged with the caller's signal via AbortSignal.any and aborts it on timeout, so the stalled request is actually torn down. The provider iterator is additionally closed from a finally block (fire-and-forget: awaiting return() on a stalled stream would re-introduce the hang) so generator cleanup also runs on consumer throws and early exits. In agent-core-v2, comments move to the top-of-file header per that package's convention. --- .../src/app/llmProtocol/generate.ts | 103 ++++++++++-------- .../llmProtocol/stream-idle-timeout.test.ts | 67 +++++++++++- packages/kosong/src/generate.ts | 88 +++++++++------ .../kosong/test/stream-idle-timeout.test.ts | 73 ++++++++++++- 4 files changed, 246 insertions(+), 85 deletions(-) diff --git a/packages/agent-core-v2/src/app/llmProtocol/generate.ts b/packages/agent-core-v2/src/app/llmProtocol/generate.ts index 6433fe6fa3..1341c8f531 100644 --- a/packages/agent-core-v2/src/app/llmProtocol/generate.ts +++ b/packages/agent-core-v2/src/app/llmProtocol/generate.ts @@ -1,3 +1,16 @@ +/** + * 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, @@ -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 withStreamIdleTimeout(stream, provider)) { + for await (const part of withStreamIdleTimeout(stream, provider, watchdog)) { const arrivedAt = Date.now(); if (firstPartAt === undefined) { firstPartAt = arrivedAt; @@ -175,20 +196,6 @@ 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 HTTP client's - * request timeout is cleared as soon as response headers arrive (see - * `openai/src/client.ts` `fetchWithTimeout`, which wraps the fetch call in - * `try { … } finally { clearTimeout(timeout) }` and resolves once headers - * are received) — so a connection that goes silent mid-stream would - * otherwise block `for await` forever and leave the whole turn hanging. - * 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; @@ -223,41 +230,45 @@ export class StreamIdleTimeoutError extends APITimeoutError { async function* withStreamIdleTimeout( stream: StreamedMessage, provider: ChatProvider, + watchdog: AbortController, ): AsyncGenerator { const iterator = stream[Symbol.asyncIterator](); const startedAt = Date.now(); - 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 stream is cancelled below; - // swallow it so it never surfaces as an unhandled rejection. - next.catch(() => {}); - await cancelStream(stream); + 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); } - throw error; - } finally { - clearTimeout(timer); + if (result.done === true) return; + yield result.value; } - if (result.done === true) return; - yield result.value; + } finally { + void Promise.resolve(iterator.return?.()).catch(() => {}); } } 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 index 4ee94496b5..74a5cbe2cd 100644 --- 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 @@ -12,11 +12,22 @@ import { } from '#/app/llmProtocol/errors'; import { describe, expect, it } from 'vitest'; -function makeProvider(stream: StreamedMessage): ChatProvider { +function makeProvider( + stream: StreamedMessage, + onGenerate?: (options?: { signal?: AbortSignal }) => void, +): ChatProvider { return { name: 'fake', modelName: 'fake-model', - generate: async () => stream, + generate: async ( + _systemPrompt: unknown, + _tools: unknown, + _history: unknown, + options?: { signal?: AbortSignal }, + ) => { + onGenerate?.(options); + return stream; + }, } as unknown as ChatProvider; } @@ -24,6 +35,7 @@ function stalledStream( parts: StreamedMessagePart[], onCancel?: () => void, traceId: string | null = null, + onIteratorReturn?: () => void, ): StreamedMessage { let i = 0; return { @@ -41,6 +53,10 @@ function stalledStream( } return new Promise(() => {}); // silent forever }, + return(): Promise> { + onIteratorReturn?.(); + return Promise.resolve({ done: true, value: undefined }); + }, }; }, } as StreamedMessage; @@ -124,4 +140,51 @@ describe('generate() stream idle watchdog', () => { ]); 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 a4688197b8..5925842400 100644 --- a/packages/kosong/src/generate.ts +++ b/packages/kosong/src/generate.ts @@ -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 withStreamIdleTimeout(stream, provider)) { + for await (const part of withStreamIdleTimeout(stream, provider, watchdog)) { const arrivedAt = Date.now(); if (firstPartAt === undefined) { firstPartAt = arrivedAt; @@ -322,41 +334,53 @@ export class StreamIdleTimeoutError extends APITimeoutError { async function* withStreamIdleTimeout( stream: StreamedMessage, provider: ChatProvider, + watchdog: AbortController, ): AsyncGenerator { const iterator = stream[Symbol.asyncIterator](); const startedAt = Date.now(); - 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 stream is cancelled below; - // swallow it so it never surfaces as an unhandled rejection. - next.catch(() => {}); - await cancelStream(stream); + 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); } - throw error; - } finally { - clearTimeout(timer); + if (result.done === true) return; + yield result.value; } - 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(() => {}); } } diff --git a/packages/kosong/test/stream-idle-timeout.test.ts b/packages/kosong/test/stream-idle-timeout.test.ts index 613f43f08e..e2e2310483 100644 --- a/packages/kosong/test/stream-idle-timeout.test.ts +++ b/packages/kosong/test/stream-idle-timeout.test.ts @@ -11,13 +11,20 @@ * kap-server path — the same failure signature must fail loudly on both. */ import { APITimeoutError, isRetryableGenerateError } from '#/errors'; -import type { StreamedMessagePart } from '#/message'; -import type { ChatProvider, StreamedMessage, ThinkingEffort } from '#/provider'; +import type { Message, StreamedMessagePart } from '#/message'; +import type { + ChatProvider, + GenerateOptions, + StreamedMessage, + ThinkingEffort, +} from '#/provider'; import type { Tool } from '#/tool'; -import type { Message } from '#/message'; import { describe, expect, it } from 'vitest'; -function makeProvider(stream: StreamedMessage): ChatProvider { +function makeProvider( + stream: StreamedMessage, + onGenerate?: (options?: GenerateOptions) => void, +): ChatProvider { return { name: 'fake', modelName: 'fake-model', @@ -26,7 +33,11 @@ function makeProvider(stream: StreamedMessage): ChatProvider { _systemPrompt: string, _tools: Tool[], _history: Message[], - ): Promise => stream, + options?: GenerateOptions, + ): Promise => { + onGenerate?.(options); + return stream; + }, withThinking(_effort: ThinkingEffort): ChatProvider { return this; }, @@ -37,6 +48,7 @@ function stalledStream( parts: StreamedMessagePart[], onCancel?: () => void, traceId: string | null = null, + onIteratorReturn?: () => void, ): StreamedMessage { let i = 0; return { @@ -54,6 +66,10 @@ function stalledStream( } return new Promise(() => {}); // silent forever }, + return(): Promise> { + onIteratorReturn?.(); + return Promise.resolve({ done: true, value: undefined }); + }, }; }, } as unknown as StreamedMessage; @@ -141,4 +157,51 @@ describe('kosong.generate() stream idle watchdog', () => { ]); 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); + }); });