fix(kosong,agent-core-v2): watchdog and retry silent LLM stream stalls

.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.

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<StreamedMessagePart> {
+  const iterator = stream[Symbol.asyncIterator]();
+  const startedAt = Date.now();
+  try {
+    while (true) {
+      let timer: ReturnType<typeof setTimeout> | undefined;
+      const timeout = new Promise<never>((_, 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<StreamedMessagePart>;
+      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');
 }

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<IteratorResult<StreamedMessagePart>> {
+          if (i < parts.length) {
+            return Promise.resolve({ done: false, value: parts[i++]! });
+          }
+          return new Promise<never>(() => {}); // silent forever
+        },
+        return(): Promise<IteratorResult<StreamedMessagePart>> {
+          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<IteratorResult<StreamedMessagePart>> {
+          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<typeof StreamIdleTimeoutError>;
+    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);
+  });
+});