Merge pull request #12588 from Kilo-Org/fix/provider-first-byte-timeout

fix(cli): prevent agent-loop freeze when a provider stalls after headers
2026-09-21 05:52:35 +08:00 · 2026-07-28 11:00:52 +02:00
parent e4759b75ac deddf0012f
commit f336cfee0a
8 changed files with 520 additions and 12 deletions
@@ -0,0 +1,30 @@
+// Plugin used by the issue #8656 regression tests.
+//
+// It attaches the simulated socket from stall-transport.ts to the `mock`
+// provider through the plugin `config` hook. This is the supported injection
+// point: src/provider/provider.ts loads plugins before reading `cfg.provider`
+// exactly so hooks can add options such as `fetch`. Injecting here keeps the
+// simulation scoped to the test's own instance instead of replacing
+// `globalThis.fetch` for the whole test process.
+
+import { createStallTransport } from "./stall-transport"
+
+type Options = { state?: unknown; answer?: unknown; provider?: unknown }
+
+type Draft = {
+  provider?: Record<string, { options?: Record<string, unknown> } | undefined>
+}
+
+export default async (_input: unknown, options?: Options) => ({
+  config: async (cfg: Draft) => {
+    const id = typeof options?.provider === "string" ? options.provider : "mock"
+    const provider = cfg.provider?.[id]
+    const state = typeof options?.state === "string" ? options.state : undefined
+    if (!provider || !state) return
+    provider.options ??= {}
+    provider.options["fetch"] = createStallTransport({
+      state,
+      answer: typeof options?.answer === "string" ? options.answer : undefined,
+    })
+  },
+})
@@ -0,0 +1,119 @@
+// Simulated provider socket for the issue #8656 regression tests.
+//
+// Only the socket is simulated. The transport is injected as the provider's
+// `fetch` option, so Kilo's own fetch wrapper (connection timeout, first-byte
+// guard, SSE chunk watchdog), the openai-compatible SDK, SSE parsing, the
+// session processor and the agent loop are all the production ones.
+//
+// Request script:
+//   1. title request                  -> short text answer
+//   2. no tool result in the messages -> a bash tool call
+//   3. first request carrying a tool result -> SSE headers, body never sends a
+//      byte (the stall reported in #8656)
+//   4. later requests carrying a tool result -> final text answer, so a bounded
+//      stall can recover through the normal retry path
+//
+// Progress is mirrored to a JSON file so tests can assert what the provider saw
+// without sharing module state with the plugin that loads this file.
+
+export type StallState = { calls: number; stalls: number; recovered: number }
+
+const HEAD = { id: "chatcmpl-stall", object: "chat.completion.chunk", created: 0, model: "mock-model" }
+
+const chunk = (obj: unknown) => `data: ${JSON.stringify(obj)}\n\n`
+
+const usage = () =>
+  chunk({ ...HEAD, choices: [], usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 } })
+
+function sse(body: BodyInit | null) {
+  return new Response(body, { status: 200, headers: { "content-type": "text/event-stream" } })
+}
+
+function answer(value: string) {
+  return sse(
+    [
+      chunk({ ...HEAD, choices: [{ index: 0, delta: { role: "assistant", content: value }, finish_reason: null }] }),
+      chunk({ ...HEAD, choices: [{ index: 0, delta: {}, finish_reason: "stop" }] }),
+      usage(),
+      "data: [DONE]\n\n",
+    ].join(""),
+  )
+}
+
+function toolCall(command: string) {
+  return sse(
+    [
+      chunk({
+        ...HEAD,
+        choices: [
+          {
+            index: 0,
+            delta: {
+              role: "assistant",
+              tool_calls: [
+                {
+                  index: 0,
+                  id: "call_1",
+                  type: "function",
+                  function: { name: "bash", arguments: JSON.stringify({ command }) },
+                },
+              ],
+            },
+            finish_reason: null,
+          },
+        ],
+      }),
+      chunk({ ...HEAD, choices: [{ index: 0, delta: {}, finish_reason: "tool_calls" }] }),
+      usage(),
+      "data: [DONE]\n\n",
+    ].join(""),
+  )
+}
+
+/** Response headers arrive, then the body never produces a byte. */
+function stalling() {
+  return sse(
+    new ReadableStream<Uint8Array>({
+      start() {},
+      cancel() {
+        // the first-byte guard cancels this reader when it gives up
+      },
+    }),
+  )
+}
+
+export function createStallTransport(input: { state: string; answer?: string; command?: string }) {
+  const state: StallState = { calls: 0, stalls: 0, recovered: 0 }
+  const save = () => Bun.write(input.state, JSON.stringify(state))
+
+  return async (_input: unknown, init?: { body?: unknown }) => {
+    const body = typeof init?.body === "string" ? init.body : ""
+    state.calls++
+
+    if (body.includes("Generate a title")) {
+      await save()
+      return answer("Stall repro")
+    }
+
+    if (!body.includes('"role":"tool"')) {
+      await save()
+      return toolCall(input.command ?? "echo repro-8656")
+    }
+
+    if (state.stalls === 0) {
+      state.stalls++
+      await save()
+      return stalling()
+    }
+
+    state.recovered++
+    await save()
+    return answer(input.answer ?? "recovered after the stall")
+  }
+}
+
+export async function readStallState(file: string): Promise<StallState> {
+  const handle = Bun.file(file)
+  if (!(await handle.exists())) return { calls: 0, stalls: 0, recovered: 0 }
+  return JSON.parse(await handle.text()) as StallState
+}
@@ -0,0 +1,180 @@
+import { describe, expect, test } from "bun:test"
+import path from "node:path"
+import { pathToFileURL } from "node:url"
+import { provideTestInstance, tmpdir } from "../fixture/fixture"
+import { readStallState } from "./fixture/stall-transport"
+import { Server } from "../../src/server/server"
+
+// Regression coverage for https://github.com/Kilo-Org/kilocode/issues/8656
+//
+// Reported symptom: after a tool call finished, `step-finish:tool-calls` was
+// recorded and the next `step-start` never arrived, leaving the session busy
+// with no error while the HTTP server stayed responsive.
+//
+// One transport state produces exactly that symptom: the follow-up request that
+// carries the tool result gets response headers and then never receives a byte
+// of body. The connection-phase timeout is cleared as soon as headers arrive, so
+// before the fix nothing bounded that wait.
+//
+// The simulated socket is injected as the provider's `fetch` through the plugin
+// `config` hook (see ../fixture/stall-plugin.ts), so the SDK, Kilo's fetch
+// wrapper, SSE parsing, the processor and the agent loop stay production code
+// and nothing global is patched.
+
+const PLUGIN = pathToFileURL(path.join(import.meta.dir, "fixture", "stall-plugin.ts")).href
+const ANSWER = "recovered after the stall"
+
+const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms))
+
+function settings(state: string, timeout: number | false) {
+  return {
+    $schema: "https://app.kilo.ai/config.json",
+    model: "mock/mock-model",
+    plugin: [[PLUGIN, { state, answer: ANSWER }]],
+    provider: {
+      mock: {
+        npm: "@ai-sdk/openai-compatible",
+        name: "Mock",
+        options: { baseURL: "http://127.0.0.1:1/v1", apiKey: "test", timeout },
+        models: {
+          "mock-model": {
+            name: "Mock Model",
+            tool_call: true,
+            limit: { context: 128000, output: 8192 },
+            cost: { input: 0, output: 0 },
+          },
+        },
+      },
+    },
+    permission: { bash: "allow" },
+  }
+}
+
+function project(timeout: number | false) {
+  return tmpdir<{ state: string }>({
+    init: async (dir) => {
+      const state = path.join(dir, "stall-state.json")
+      await Bun.write(path.join(dir, "opencode.json"), JSON.stringify(settings(state, timeout), null, 2))
+      return { state }
+    },
+  })
+}
+
+type Part = Record<string, any>
+type Message = { info: Record<string, any>; parts: Part[] }
+
+function session(dir: string) {
+  const app = Server.Default().app
+  const headers = { "Content-Type": "application/json", "x-kilo-directory": dir }
+  const query = `directory=${encodeURIComponent(dir)}`
+
+  const json = async (route: string, init?: RequestInit) => {
+    const res = await app.request(route, { headers, ...init })
+    return await res.json()
+  }
+
+  return {
+    create: async () => ((await json("/session", { method: "POST", body: "{}" })) as { id: string }).id,
+    prompt: (id: string, text: string) =>
+      app.request(`/session/${id}/prompt_async`, {
+        method: "POST",
+        headers,
+        body: JSON.stringify({ parts: [{ type: "text", text }] }),
+      }),
+    abort: (id: string) => app.request(`/session/${id}/abort`, { method: "POST", headers }),
+    messages: (id: string) => json(`/session/${id}/message?${query}`) as Promise<Message[]>,
+    status: async (id: string) => {
+      const all = (await json(`/session/status?${query}`)) as Record<string, { type: string }>
+      return all[id]?.type ?? "idle"
+    },
+  }
+}
+
+const timeline = (messages: Message[]) =>
+  messages
+    .flatMap((m) => m.parts)
+    .map((p) =>
+      p.type === "tool" ? `tool:${p.tool}:${p.state?.status}` : `${p.type}${p.reason ? `:${p.reason}` : ""}`,
+    )
+    .join(" | ")
+
+async function until(check: () => Promise<boolean>, budget: number) {
+  const deadline = Date.now() + budget
+  while (Date.now() < deadline) {
+    if (await check()) return true
+    await sleep(200)
+  }
+  return false
+}
+
+/** Runs a prompt and waits until the provider has stalled the follow-up request. */
+async function stalled(api: ReturnType<typeof session>, state: string) {
+  const id = await api.create()
+  await api.prompt(id, "run the echo command")
+  const ready = await until(async () => {
+    const stalls = (await readStallState(state)).stalls
+    const parts = timeline(await api.messages(id))
+    return stalls > 0 && parts.includes("step-finish:tool-calls")
+  }, 30_000)
+  expect(ready).toBe(true)
+  return id
+}
+
+describe("issue #8656: provider stalls after a tool call", () => {
+  test("recovers instead of freezing once the stall is bounded", async () => {
+    await using tmp = await project(2_000)
+    await provideTestInstance({
+      directory: tmp.path,
+      fn: async () => {
+        const api = session(tmp.path)
+        const id = await stalled(api, tmp.extra.state)
+
+        const done = await until(async () => (await api.status(id)) === "idle", 40_000)
+        const messages = await api.messages(id)
+        const assistant = messages.findLast((m) => m.info.role === "assistant")
+        const text = (assistant?.parts ?? []).find((p) => p.type === "text")?.text
+        const state = await readStallState(tmp.extra.state)
+        console.log("[repro] bounded ->", JSON.stringify({ timeline: timeline(messages), text, state }))
+
+        // the turn finishes on its own: the stalled request was aborted, retried
+        // and answered, so the agent loop never sits frozen
+        expect(done).toBe(true)
+        expect(timeline(messages)).toContain("tool:bash:completed")
+        expect(state.stalls).toBe(1)
+        expect(state.recovered).toBeGreaterThan(0)
+        expect(text).toContain(ANSWER)
+        expect(assistant?.info.error).toBeUndefined()
+      },
+    })
+  }, 120_000)
+
+  test("still hangs while the provider holds the connection open and timeout is disabled", async () => {
+    await using tmp = await project(false)
+    await provideTestInstance({
+      directory: tmp.path,
+      fn: async () => {
+        const api = session(tmp.path)
+        const id = await stalled(api, tmp.extra.state)
+        try {
+          await sleep(5_000)
+          const messages = await api.messages(id)
+          const parts = timeline(messages)
+          const assistant = messages.findLast((m) => m.info.role === "assistant")
+          console.log("[repro] timeout:false ->", JSON.stringify({ timeline: parts, status: await api.status(id) }))
+
+          // the reported freeze, kept reachable only through the documented opt-out
+          expect(parts.endsWith("step-finish:tool-calls")).toBe(true)
+          expect(await api.status(id)).toBe("busy")
+          expect(assistant?.info.error).toBeUndefined()
+
+          // the server itself stays responsive during the freeze
+          expect((await Server.Default().app.request("/global/health")).status).toBe(200)
+        } finally {
+          // never leave a wedged turn behind for fixture teardown
+          expect((await api.abort(id)).status).toBe(200)
+          expect(await until(async () => (await api.status(id)) === "idle", 15_000)).toBe(true)
+        }
+      },
+    })
+  }, 120_000)
+})
@@ -0,0 +1,84 @@
+import { describe, expect, test } from "bun:test"
+import { requestTimeout, wrapFirstByte, REQUEST_TIMEOUT_MS } from "../../../src/kilocode/provider/provider"
+import { ProviderError } from "../../../src/provider/error"
+
+const sse = (body: BodyInit | null) =>
+  new Response(body, { headers: { "content-type": "text/event-stream" }, status: 200 })
+
+// A body that delivers `head` immediately, then stays silent forever.
+const stalling = (head?: string) =>
+  new ReadableStream<Uint8Array>({
+    start(ctrl) {
+      if (head) ctrl.enqueue(new TextEncoder().encode(head))
+    },
+  })
+
+const drain = async (res: Response) => {
+  const reader = res.body!.getReader()
+  const out: string[] = []
+  while (true) {
+    const part = await reader.read()
+    if (part.done) return out.join("")
+    out.push(new TextDecoder().decode(part.value))
+  }
+}
+
+describe("requestTimeout", () => {
+  test("defaults to the shared request timeout", () => {
+    expect(requestTimeout({})).toBe(REQUEST_TIMEOUT_MS)
+  })
+
+  test("honours an explicit value, disables only on false, and bounds invalid input", () => {
+    expect(requestTimeout({ timeout: 1234 })).toBe(1234)
+    expect(requestTimeout({ timeout: false })).toBeUndefined()
+    // invalid/unset values fall back to the default so the wait is always bounded
+    expect(requestTimeout({ timeout: 0 })).toBe(REQUEST_TIMEOUT_MS)
+    expect(requestTimeout({ timeout: -1 })).toBe(REQUEST_TIMEOUT_MS)
+    expect(requestTimeout({ timeout: "nope" })).toBe(REQUEST_TIMEOUT_MS)
+    expect(requestTimeout({ timeout: null })).toBe(REQUEST_TIMEOUT_MS)
+  })
+})
+
+describe("wrapFirstByte", () => {
+  test("fails when the body never produces a byte", async () => {
+    const ctl = new AbortController()
+    const res = wrapFirstByte(sse(stalling()), 100, ctl)
+
+    await expect(drain(res)).rejects.toBeInstanceOf(ProviderError.ResponseStreamError)
+    expect(ctl.signal.aborted).toBe(true)
+    expect((ctl.signal.reason as Error).message).toContain("no response data within 100ms")
+  })
+
+  test("stays a passthrough once the first byte arrived, even if the stream then stalls", async () => {
+    const ctl = new AbortController()
+    const res = wrapFirstByte(sse(stalling("data: hello\n\n")), 100, ctl)
+    const reader = res.body!.getReader()
+
+    const first = await reader.read()
+    expect(new TextDecoder().decode(first.value)).toBe("data: hello\n\n")
+
+    // the guard must not arm again: idle gaps mid-stream stay opt-in (chunkTimeout)
+    const next = await Promise.race([
+      reader.read().then(() => "chunk"),
+      new Promise((resolve) => setTimeout(() => resolve("still-waiting"), 300)),
+    ])
+    expect(next).toBe("still-waiting")
+    expect(ctl.signal.aborted).toBe(false)
+    await reader.cancel("done")
+  })
+
+  test("passes a complete body through untouched", async () => {
+    const ctl = new AbortController()
+    const res = wrapFirstByte(sse("data: one\n\ndata: two\n\n"), 1_000, ctl)
+    expect(await drain(res)).toBe("data: one\n\ndata: two\n\n")
+    expect(ctl.signal.aborted).toBe(false)
+  })
+
+  test("is a no-op without a body or when disabled", () => {
+    const ctl = new AbortController()
+    const empty = new Response(null, { status: 204 })
+    expect(wrapFirstByte(empty, 100, ctl)).toBe(empty)
+    const res = sse("data: x\n\n")
+    expect(wrapFirstByte(res, 0, ctl)).toBe(res)
+  })
+})