mirror of
https://github.com/Kilo-Org/kilocode.git
synced 2026-09-01 15:32:11 +08:00
Revert "fix(cli): prevent stalled agent streams (#12249)"
This reverts commit cd205d857a.
This commit is contained in:
@@ -1,5 +0,0 @@
|
||||
---
|
||||
"@kilocode/cli": patch
|
||||
---
|
||||
|
||||
Fix CLI agent stream stalls caused by the idle watchdog firing during time-to-first-content. The per-chunk idle watchdog armed on the AI SDK's synthetic `start` part and raced the wait for the first content-bearing part (prompt processing / time-to-first-token) against the same idle window as inter-chunk gaps, so a healthy but slow first response was falsely aborted and immediately retried (regression #12467). The watchdog now bounds the pre-content phase by the provider's configured request `timeout` (falling back to a 5-minute default) and only applies the per-chunk idle window after the first content/tool part. The same first-content-aware split is applied to the lower-level SSE fetch watchdog (`wrapSSE`) so provider-level `chunkTimeout` no longer aborts slow first content, while mid-response stall detection is preserved at both layers.
|
||||
@@ -115,11 +115,9 @@ export const Info = Schema.Struct({
|
||||
description:
|
||||
"Timeout in milliseconds to wait for response headers. Provider integrations may set defaults. Set to false to disable timeout.",
|
||||
}),
|
||||
// kilocode_change: accept `false` so internal callers can disable the
|
||||
// watchdog. PositiveInt already excludes 0, so a public zero stays invalid.
|
||||
chunkTimeout: Schema.optional(Schema.Union([PositiveInt, Schema.Literal(false)])).annotate({
|
||||
chunkTimeout: Schema.optional(PositiveInt).annotate({
|
||||
description:
|
||||
"Timeout in milliseconds between streamed SSE chunks for this provider. If no chunk arrives within this window, the request is aborted. Set to false to disable the idle watchdog. The pre-content bound is only shape-aware for OpenAI-compatible SSE and otherwise uses the request `timeout` budget.",
|
||||
"Timeout in milliseconds between streamed SSE chunks for this provider. If no chunk arrives within this window, the request is aborted.",
|
||||
}),
|
||||
}),
|
||||
[Schema.Record(Schema.String, Schema.Any)],
|
||||
|
||||
@@ -6,7 +6,6 @@
|
||||
// This module exports patch functions and data that the upstream provider.ts
|
||||
// calls at well-defined injection points (each marked with kilocode_change).
|
||||
|
||||
import { ProviderError } from "@/provider/error"
|
||||
import { createKilo, type KiloProvider, AI_SDK_PROVIDERS, PROMPTS } from "@kilocode/kilo-gateway"
|
||||
import { DEFAULT_HEADERS } from "@/kilocode/const"
|
||||
import { ProviderV2 } from "@opencode-ai/core/provider"
|
||||
@@ -19,173 +18,6 @@ import { mapValues, omit, pickBy } from "remeda"
|
||||
/** Default timeout (ms) for provider HTTP requests (connection phase). */
|
||||
export const REQUEST_TIMEOUT_MS = 300_000 // 5 minutes
|
||||
|
||||
/**
|
||||
* Pre-content (time-to-first-content) budget for raw SSE streams. Mirrors the
|
||||
* value of `KiloLLM.DEFAULT_FIRST_TOKEN_MS` in `src/kilocode/session/llm.ts` —
|
||||
* keep these two constants in sync if either ever changes.
|
||||
*/
|
||||
export const SSE_FIRST_TOKEN_MS = 300_000 // 5 minutes
|
||||
|
||||
/**
|
||||
* Resolves the pre-content timeout budget for `wrapSSEFirstContent`, mirroring
|
||||
* `KiloLLM.resolveFirstTokenMs` semantics: a positive finite `options.timeout`
|
||||
* wins; otherwise falls back to a positive finite provider `timeout`; otherwise
|
||||
* `SSE_FIRST_TOKEN_MS`. `false` / `0` / unset / invalid / non-finite all map to
|
||||
* the default, because a never-first-content hang must remain bounded.
|
||||
*/
|
||||
export function resolveSseFirstTokenMs(options: Record<string, any>, fallback: Record<string, any> = {}): number {
|
||||
const val = options["timeout"]
|
||||
if (typeof val === "number" && Number.isFinite(val) && val > 0) return val
|
||||
const fb = fallback["timeout"]
|
||||
if (typeof fb === "number" && Number.isFinite(fb) && fb > 0) return fb
|
||||
return SSE_FIRST_TOKEN_MS
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// SSE first-content detection
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type ContentEventResult = { found: boolean; carry: string }
|
||||
|
||||
/**
|
||||
* Returns `found: true` iff any COMPLETE SSE event in `text` has a `data:` line
|
||||
* whose JSON `choices[0].delta` carries `content`, `reasoning_content`, or
|
||||
* `tool_calls`. The last (incomplete) event is skipped — it will be re-checked
|
||||
* as the head of the next read's buffer.
|
||||
*
|
||||
* To bound memory usage, callers should replace `decoderBuf` with `carry`
|
||||
* after each call. `carry` is the trailing incomplete event fragment (the text
|
||||
* after the last `\n\n` boundary), so completed events are never retained or
|
||||
* re-scanned.
|
||||
*/
|
||||
export function looksLikeContentEvent(text: string): ContentEventResult {
|
||||
const events = text.split(/\r?\n\r?\n/)
|
||||
if (events.length <= 1) return { found: false, carry: text }
|
||||
|
||||
const complete = events.slice(0, -1)
|
||||
for (const evt of complete) {
|
||||
const dataLines: string[] = []
|
||||
for (const line of evt.split(/\r?\n/)) {
|
||||
if (line.startsWith(":")) continue
|
||||
if (line.startsWith("data:")) {
|
||||
dataLines.push(line.slice(5).replace(/^ /, ""))
|
||||
}
|
||||
}
|
||||
if (dataLines.length === 0) continue
|
||||
const payload = dataLines.join("\n")
|
||||
if (payload === "[DONE]") continue
|
||||
let parsed: unknown
|
||||
try {
|
||||
parsed = JSON.parse(payload)
|
||||
} catch {
|
||||
continue
|
||||
}
|
||||
if (!parsed || typeof parsed !== "object") continue
|
||||
const choices = (parsed as { choices?: unknown }).choices
|
||||
if (!Array.isArray(choices) || choices.length === 0) continue
|
||||
const delta = (choices[0] as { delta?: unknown })?.delta
|
||||
if (!delta || typeof delta !== "object") continue
|
||||
const d = delta as { content?: unknown; reasoning_content?: unknown; tool_calls?: unknown }
|
||||
if (typeof d.content === "string" && d.content.length > 0) return { found: true, carry: "" }
|
||||
if (typeof d.reasoning_content === "string" && d.reasoning_content.length > 0) return { found: true, carry: "" }
|
||||
if (Array.isArray(d.tool_calls) && d.tool_calls.length > 0) return { found: true, carry: "" }
|
||||
}
|
||||
|
||||
return { found: false, carry: events[events.length - 1] }
|
||||
}
|
||||
|
||||
/**
|
||||
* Wraps an SSE `Response` body so that reads before the first content-bearing
|
||||
* `data:` event are bounded by `firstTokenMs`, while reads after content are
|
||||
* bounded by the per-chunk `chunkTimeout`. Original bytes are enqueued
|
||||
* untouched — the text decoder is observational only.
|
||||
*
|
||||
* This is the Kilo-specific first-content-aware layer; the upstream provider
|
||||
* calls it at the single `wrapSSE` injection point.
|
||||
*
|
||||
* Scope note: this is installed for every provider SDK built by `resolveSDK`,
|
||||
* but `looksLikeContentEvent` only recognizes the OpenAI chat-completions SSE
|
||||
* shape (`choices[0].delta` with `content` / `reasoning_content` /
|
||||
* `tool_calls`). For other provider-native shapes (Anthropic
|
||||
* `content_block_delta`, OpenAI Responses `response.output_text.delta`, Google
|
||||
* `candidates`, etc.) `seenContent` never flips, so pre-content reads remain
|
||||
* bounded by `firstTokenMs` (the request `timeout` budget) rather than by
|
||||
* `chunkTimeout` once content starts. This is intentional:
|
||||
* (a) the stream is still bounded (no hang);
|
||||
* (b) the provider-agnostic session-layer `KiloLLM.watchIterator` guards the
|
||||
* main session path for all providers using normalized AI SDK parts;
|
||||
* (c) treating the first arbitrary `data:` event as content would arm on
|
||||
* OpenAI's immediate role-delta and reintroduce the #12467 false-positive;
|
||||
* (d) this only arms when a positive provider-level `chunkTimeout` is
|
||||
* configured (no built-in default).
|
||||
*/
|
||||
export function wrapSSEFirstContent(
|
||||
res: Response,
|
||||
chunkTimeout: number,
|
||||
ctl: AbortController,
|
||||
firstTokenMs: number,
|
||||
): Response {
|
||||
if (typeof chunkTimeout !== "number" || chunkTimeout <= 0) return res
|
||||
if (!res.body) return res
|
||||
if (!res.headers.get("content-type")?.includes("text/event-stream")) return res
|
||||
|
||||
const reader = res.body.getReader()
|
||||
const decoder = new TextDecoder("utf-8")
|
||||
let decoderBuf = ""
|
||||
let seenContent = false
|
||||
|
||||
const body = new ReadableStream<Uint8Array>({
|
||||
async pull(ctrl) {
|
||||
const budget = seenContent ? chunkTimeout : firstTokenMs
|
||||
const part = await new Promise<Awaited<ReturnType<typeof reader.read>>>((resolve, reject) => {
|
||||
const id = setTimeout(() => {
|
||||
const err = new ProviderError.ResponseStreamError("SSE read timed out")
|
||||
ctl.abort(err)
|
||||
void reader.cancel(err)
|
||||
reject(err)
|
||||
}, budget)
|
||||
|
||||
reader.read().then(
|
||||
(part) => {
|
||||
clearTimeout(id)
|
||||
resolve(part)
|
||||
},
|
||||
(err) => {
|
||||
clearTimeout(id)
|
||||
reject(err)
|
||||
},
|
||||
)
|
||||
})
|
||||
|
||||
if (part.done) {
|
||||
ctrl.close()
|
||||
return
|
||||
}
|
||||
|
||||
if (!seenContent && part.value) {
|
||||
decoderBuf += decoder.decode(part.value, { stream: true })
|
||||
const result = looksLikeContentEvent(decoderBuf)
|
||||
if (result.found) {
|
||||
seenContent = true
|
||||
}
|
||||
decoderBuf = result.carry
|
||||
}
|
||||
|
||||
ctrl.enqueue(part.value)
|
||||
},
|
||||
async cancel(reason) {
|
||||
ctl.abort(reason)
|
||||
await reader.cancel(reason)
|
||||
},
|
||||
})
|
||||
|
||||
return new Response(body, {
|
||||
headers: new Headers(res.headers),
|
||||
status: res.status,
|
||||
statusText: res.statusText,
|
||||
})
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Bundled providers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -1,24 +1,12 @@
|
||||
import type { LanguageModelV2StreamPart } from "@ai-sdk/provider"
|
||||
import * as Stream from "effect/Stream"
|
||||
import { ProviderError } from "@/provider/error"
|
||||
import type { LLMEvent } from "@opencode-ai/llm"
|
||||
import type { ModelMessage } from "ai"
|
||||
import * as Stream from "effect/Stream"
|
||||
import type { LLMEvent } from "@opencode-ai/llm"
|
||||
import type { Logger } from "@opencode-ai/core/util/log"
|
||||
import type { Provider } from "@/provider/provider"
|
||||
import { KiloSessionOverflow } from "./overflow"
|
||||
|
||||
const SAFETY = 2048
|
||||
const MIN_OUTPUT = 1024
|
||||
const DEFAULT_CHUNK_IDLE_MS = 300_000
|
||||
// Kilo default for the pre-content (time-to-first-content / prompt-processing)
|
||||
// bound. The provider's own request `timeout` signal is cleared once response
|
||||
// headers arrive (see `buildTimeoutSignal` in src/kilocode/provider/provider.ts),
|
||||
// so the watchdog is the only thing that bounds the pre-content phase, and it
|
||||
// uses the configured `timeout` value (or this default) as its budget. 5 min is
|
||||
// generous for slow prompt processing on large-context / local models while
|
||||
// still bounding a genuine never-first-content hang.
|
||||
const DEFAULT_FIRST_TOKEN_MS = 300_000
|
||||
|
||||
type FullStreamPart = LanguageModelV2StreamPart
|
||||
|
||||
export namespace KiloLLM {
|
||||
// Stream failures and interruptions propagate while text deltas are collected.
|
||||
@@ -29,307 +17,20 @@ export namespace KiloLLM {
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves the configured chunk idle timeout in milliseconds, or `undefined`
|
||||
* when the watchdog should be disabled.
|
||||
*
|
||||
* Precedence:
|
||||
* 1. prepared `options.chunkTimeout`
|
||||
* 2. provider `fallback.chunkTimeout`
|
||||
* 3. DEFAULT_CHUNK_IDLE_MS
|
||||
*
|
||||
* Rules:
|
||||
* - positive finite number wins.
|
||||
* - public `false` or internal `0` disables (returns undefined).
|
||||
* - invalid prepared values (non-number, negative, non-finite, strings, ...)
|
||||
* fall through to the provider fallback. The same rules apply at every
|
||||
* layer.
|
||||
*/
|
||||
export function resolveIdleMs(input: {
|
||||
export function timeout(input: {
|
||||
options: Record<string, unknown>
|
||||
fallback?: Record<string, unknown>
|
||||
}): number | undefined {
|
||||
const prepared = resolve(input.options["chunkTimeout"])
|
||||
if (prepared.disabled) return undefined
|
||||
if (prepared.value !== undefined) return prepared.value
|
||||
const fallback = resolve(input.fallback?.["chunkTimeout"])
|
||||
if (fallback.disabled) return undefined
|
||||
if (fallback.value !== undefined) return fallback.value
|
||||
return DEFAULT_CHUNK_IDLE_MS
|
||||
}
|
||||
|
||||
/**
|
||||
* Resolves the pre-content (time-to-first-content / prompt-processing) budget
|
||||
* in milliseconds. Unlike `resolveIdleMs`, this value is NEVER `undefined`:
|
||||
* the pre-content phase is disabled only when the whole watchdog is disabled
|
||||
* (`chunkTimeout: false` → `idleMs === undefined` → `watchdogStream` returns
|
||||
* the stream unchanged), which is decided at the watchdog entry point, not
|
||||
* here.
|
||||
*
|
||||
* Rationale: the provider's own request `timeout` signal is cleared once HTTP
|
||||
* response headers arrive (`buildTimeoutSignal`,
|
||||
* `src/kilocode/provider/provider.ts:254-272`), so it does NOT bound the
|
||||
* pre-content phase (which is post-header). The watchdog is therefore the
|
||||
* only enforcement mechanism for time-to-first-content, and using the
|
||||
* configured `timeout` VALUE as its budget is the only way to give a
|
||||
* never-first-content hang a finite bound. Mapping `timeout: false`/`0`/
|
||||
* invalid/unset to "disabled" would leave such a hang with no bound at all
|
||||
* (idleMs only arms AFTER the first content part).
|
||||
*
|
||||
* Precedence:
|
||||
* 1. prepared `options.timeout` (positive finite number)
|
||||
* 2. provider `fallback.timeout` (positive finite number)
|
||||
* 3. DEFAULT_FIRST_TOKEN_MS (5 min)
|
||||
*/
|
||||
export function resolveFirstTokenMs(input: {
|
||||
options: Record<string, unknown>
|
||||
fallback?: Record<string, unknown>
|
||||
}): number {
|
||||
const prepared = input.options["timeout"]
|
||||
if (isPositiveFiniteMs(prepared)) return prepared
|
||||
const fallback = input.fallback?.["timeout"]
|
||||
if (isPositiveFiniteMs(fallback)) return fallback
|
||||
return DEFAULT_FIRST_TOKEN_MS
|
||||
}
|
||||
|
||||
function isPositiveFiniteMs(value: unknown): value is number {
|
||||
return typeof value === "number" && Number.isFinite(value) && value > 0
|
||||
}
|
||||
|
||||
// Tri-state: `disabled` means "explicitly off"; `value` is a usable ms count.
|
||||
// `null`/`undefined`/invalid numeric values are treated as not-configured.
|
||||
function resolve(value: unknown): { value: number | undefined; disabled: boolean } {
|
||||
if (value === false || value === 0) return { value: undefined, disabled: true }
|
||||
if (value == null) return { value: undefined, disabled: false }
|
||||
if (typeof value !== "number") return { value: undefined, disabled: false }
|
||||
if (!Number.isFinite(value)) return { value: undefined, disabled: false }
|
||||
if (value <= 0) return { value: undefined, disabled: false }
|
||||
return { value, disabled: false }
|
||||
}
|
||||
|
||||
/**
|
||||
* Wraps an AI SDK `fullStream` with a Kilo-owned per-event idle watchdog.
|
||||
*
|
||||
* Behavior:
|
||||
* - `idleMs === undefined` returns the stream unchanged (disabled). In that
|
||||
* case `firstTokenMs` is also unused; the whole watchdog is off.
|
||||
* - The watchdog arms in two phases, keyed off the first content-bearing
|
||||
* part (the AI SDK emits a synthetic `start` immediately on stream open,
|
||||
* well before any provider byte, so arming on the first raw event would
|
||||
* race the time-to-first-content / prompt-processing phase — see issue
|
||||
* #12467):
|
||||
* * Pre-content: every pull before the first content/tool part is
|
||||
* raced against `firstTokenMs` (the request-timeout budget, see
|
||||
* `resolveFirstTokenMs`).
|
||||
* * Post-content: every pull after the first content/tool part is raced
|
||||
* against `idleMs` (the per-event inter-chunk idle).
|
||||
* - Content-bearing part types (verified against
|
||||
* `src/session/llm/ai-sdk.ts:140-223`): `text-delta`, `reasoning-delta`,
|
||||
* `tool-call`, `tool-input-start`, `tool-input-delta`. Structural parts
|
||||
* (`start`, `start-step`, `text-start`, `reasoning-start`, `stream-start`)
|
||||
* do NOT arm the post-content phase.
|
||||
* - non-provider-executed `tool-call` adds an active tool id; matching
|
||||
* `tool-result` / `tool-error` removes it. While any local tool id is
|
||||
* active, the watchdog is suspended (long-running tool work is not a
|
||||
* stall).
|
||||
* - provider-executed `tool-call` does not suspend the watchdog and no id
|
||||
* is tracked — those are settled server-side and a missing result is a
|
||||
* real stall.
|
||||
* - parallel local tool calls remain suspended until the last one settles.
|
||||
* - the wrapper fails the stream with `ProviderError.ResponseStreamError`
|
||||
* on stall. Existing `MessageV2` retry mapping handles that error.
|
||||
*
|
||||
* `firstTokenMs` defaults to `idleMs` so legacy callers (and the many
|
||||
* direct-unit tests that omit it) keep their exact current behavior
|
||||
* (pre- and post-content budgets both equal to `idleMs`).
|
||||
*
|
||||
* The wrapper is implemented against `AsyncIterable` so it composes with
|
||||
* any stream the AI SDK exposes, including its native `fullStream`. The
|
||||
* outer `Stream` is rebuilt from the wrapped iterable, which keeps the
|
||||
* contract simple: one pull = one raw event.
|
||||
*/
|
||||
export function watchdogStream(
|
||||
stream: Stream.Stream<FullStreamPart, unknown>,
|
||||
idleMs: number | undefined,
|
||||
abort?: AbortController,
|
||||
firstTokenMs: number = idleMs as number,
|
||||
): Stream.Stream<FullStreamPart, unknown> {
|
||||
if (idleMs === undefined) return stream
|
||||
const source = Stream.toAsyncIterable(stream)
|
||||
return Stream.fromAsyncIterable(watchdogAsyncIterable(source, idleMs, abort, firstTokenMs), (e) =>
|
||||
e instanceof Error ? e : new Error(String(e)),
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Wraps an `AsyncIterable` of raw AI SDK `fullStream` parts with the same
|
||||
* Kilo-owned per-event idle watchdog. Use this when the upstream is already
|
||||
* an `AsyncIterable` (e.g. the AI SDK's `fullStream`) so we avoid a
|
||||
* Stream → AsyncIterable → Stream round-trip. See `watchdogStream` for the
|
||||
* pre-content / post-content phase semantics; `firstTokenMs` defaults to
|
||||
* `idleMs` at this boundary for legacy callers.
|
||||
*/
|
||||
export function watchdogAsyncIterable(
|
||||
source: AsyncIterable<FullStreamPart>,
|
||||
idleMs: number | undefined,
|
||||
abort?: AbortController,
|
||||
firstTokenMs: number = idleMs as number,
|
||||
): AsyncIterable<FullStreamPart> {
|
||||
if (idleMs === undefined) return source
|
||||
return { [Symbol.asyncIterator]: () => watchIterator(source, idleMs, abort, firstTokenMs) }
|
||||
}
|
||||
|
||||
/**
|
||||
* Implemented as a hand-rolled `AsyncIterator` rather than an `async
|
||||
* function*` generator. An async generator's `.return()` cannot preempt an
|
||||
* in-flight internal `await`: per spec, when the generator is suspended
|
||||
* mid-`await` (as opposed to suspended at a `yield`), a `.return()` call
|
||||
* only takes effect once that `await` settles on its own. When the source
|
||||
* is genuinely stalled — the exact case this watchdog exists to catch —
|
||||
* that `await` never settles, so a caller that wants to cancel promptly
|
||||
* (e.g. Effect interrupting the consuming Stream) would hang forever
|
||||
* waiting for cleanup instead. A plain iterator object's `return()` runs
|
||||
* immediately and forwards to the underlying source's `return()` without
|
||||
* waiting on any outstanding pull, matching how interruption already
|
||||
* behaves for the unwrapped upstream iterator.
|
||||
*/
|
||||
function watchIterator(
|
||||
source: AsyncIterable<FullStreamPart>,
|
||||
idleMs: number,
|
||||
abort?: AbortController,
|
||||
firstTokenMs: number = idleMs,
|
||||
): AsyncIterator<FullStreamPart> {
|
||||
const local = new Set<string>()
|
||||
const iter = source[Symbol.asyncIterator]()
|
||||
let suspended = false
|
||||
let closed = false
|
||||
// The AI SDK emits a synthetic `start` part at t+2ms, well before any
|
||||
// provider byte. Arming the per-event idle timer on the first pull would
|
||||
// therefore race the time-to-first-content / prompt-processing phase and
|
||||
// falsely abort healthy slow streams (issue #12467). Instead, we arm the
|
||||
// post-content budget (`idleMs`) only after the first content-bearing
|
||||
// part (text-delta / reasoning-delta / tool-call / tool-input-*) has
|
||||
// been observed. The pre-content wait is bounded by `firstTokenMs` (the
|
||||
// request-timeout budget, see `resolveFirstTokenMs`).
|
||||
let seenContent = false
|
||||
return {
|
||||
async next(): Promise<IteratorResult<FullStreamPart>> {
|
||||
if (closed) return { done: true, value: undefined }
|
||||
try {
|
||||
// Decide BEFORE pulling whether the next event is allowed to take as
|
||||
// long as upstream needs. Local tool work in flight must not be timed
|
||||
// out — the AI SDK only emits a tool-result / tool-error once the
|
||||
// client-side tool has actually finished. The pre-content phase uses
|
||||
// the larger firstTokenMs budget; the post-content phase uses
|
||||
// idleMs.
|
||||
const budget = suspended ? undefined : seenContent ? idleMs : firstTokenMs
|
||||
const pull = budget === undefined ? iter.next() : raceWithTimeout(iter.next(), budget, abort)
|
||||
const value = await pull
|
||||
suspended = false
|
||||
if (value.done) {
|
||||
closed = true
|
||||
await safeClose(iter)
|
||||
return value
|
||||
}
|
||||
const part = value.value
|
||||
const isContent = trackPart(local, part)
|
||||
if (isContent) seenContent = true
|
||||
suspended = local.size > 0
|
||||
return { done: false, value: part }
|
||||
} catch (e) {
|
||||
closed = true
|
||||
await safeClose(iter)
|
||||
throw e
|
||||
}
|
||||
},
|
||||
async return(value?: unknown): Promise<IteratorResult<FullStreamPart>> {
|
||||
if (!closed) {
|
||||
closed = true
|
||||
await safeClose(iter)
|
||||
}
|
||||
return { done: true, value: value as FullStreamPart }
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Inspects a raw AI SDK `fullStream` part and updates the local-tool
|
||||
* suspension set. Returns `true` iff the part is content-bearing (i.e.
|
||||
* arms the post-content `seenContent` gate in `watchIterator`), so the
|
||||
* pre-content `firstTokenMs` budget only applies until the first
|
||||
* content/tool part is observed.
|
||||
*
|
||||
* Content-bearing part types (verified against
|
||||
* `src/session/llm/ai-sdk.ts:140-223`): `text-delta`, `reasoning-delta`,
|
||||
* `tool-call`, `tool-input-start`, `tool-input-delta`. Structural parts
|
||||
* (`start`, `start-step`, `text-start`, `reasoning-start`, `stream-start`)
|
||||
* do NOT arm the post-content phase.
|
||||
*
|
||||
* Suspension logic (unchanged): non-provider-executed `tool-call` adds
|
||||
* the id to the local set; matching `tool-result` / `tool-error` removes
|
||||
* it. Provider-executed `tool-call` is content-bearing but does NOT add
|
||||
* to the local set (those calls are settled server-side, so a missing
|
||||
* result is a real stall).
|
||||
*/
|
||||
function trackPart(local: Set<string>, part: FullStreamPart): boolean {
|
||||
if (!part || typeof part !== "object") return false
|
||||
const t = (part as { type?: unknown }).type
|
||||
if (t === "tool-call") {
|
||||
const call = part as unknown as {
|
||||
toolCallId?: unknown
|
||||
providerExecuted?: unknown
|
||||
}
|
||||
if (call.providerExecuted === true) {
|
||||
// Provider-executed tool-call: content-bearing (carries the call
|
||||
// payload) but does NOT suspend the watchdog — a missing server-side
|
||||
// result is a genuine stall.
|
||||
return true
|
||||
}
|
||||
if (typeof call.toolCallId !== "string") return false
|
||||
local.add(call.toolCallId)
|
||||
return true
|
||||
}
|
||||
if (t === "tool-result" || t === "tool-error") {
|
||||
const call = part as unknown as { toolCallId?: unknown }
|
||||
if (typeof call.toolCallId !== "string") return false
|
||||
local.delete(call.toolCallId)
|
||||
return false
|
||||
}
|
||||
if (
|
||||
t === "text-delta" ||
|
||||
t === "reasoning-delta" ||
|
||||
t === "tool-input-start" ||
|
||||
t === "tool-input-delta"
|
||||
) {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
function raceWithTimeout<T>(promise: Promise<T>, ms: number, abort?: AbortController): Promise<T> {
|
||||
return new Promise<T>((resolve, reject) => {
|
||||
const timer = setTimeout(() => {
|
||||
const err = new ProviderError.ResponseStreamError(`AI SDK stream stalled: no event for ${ms}ms`)
|
||||
if (abort && !abort.signal.aborted) {
|
||||
abort.abort(err)
|
||||
}
|
||||
reject(err)
|
||||
}, ms)
|
||||
promise.then(
|
||||
(v) => {
|
||||
clearTimeout(timer)
|
||||
resolve(v)
|
||||
},
|
||||
(e) => {
|
||||
clearTimeout(timer)
|
||||
reject(e)
|
||||
},
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
async function safeClose<T>(iter: AsyncIterator<T>) {
|
||||
if (typeof iter.return === "function") await iter.return()
|
||||
log?: Pick<Logger, "debug">
|
||||
}): { timeout?: { chunkMs: number } } {
|
||||
const value =
|
||||
typeof input.options["chunkTimeout"] === "number"
|
||||
? input.options["chunkTimeout"]
|
||||
: typeof input.fallback?.["chunkTimeout"] === "number"
|
||||
? input.fallback["chunkTimeout"]
|
||||
: undefined
|
||||
if (!value) return {}
|
||||
input.log?.debug("chunk idle timeout configured", { chunkTimeout: value })
|
||||
return { timeout: { chunkMs: value } }
|
||||
}
|
||||
|
||||
export function needsEstimate(input: { model: Provider.Model; configured: number | undefined }) {
|
||||
|
||||
@@ -41,8 +41,6 @@ import {
|
||||
patchKiloProviderPrivacy,
|
||||
kiloSmallModelPriority,
|
||||
buildTimeoutSignal,
|
||||
resolveSseFirstTokenMs,
|
||||
wrapSSEFirstContent,
|
||||
} from "@/kilocode/provider/provider"
|
||||
import * as ModelsRefresh from "@/kilocode/provider/models-refresh"
|
||||
// kilocode_change end
|
||||
@@ -50,13 +48,53 @@ import { ProviderError } from "./error"
|
||||
|
||||
const OPENAI_HEADER_TIMEOUT_DEFAULT = 10_000
|
||||
|
||||
// kilocode_change start
|
||||
// Kilo-specific SSE first-content-aware chunk-idle wrapper. The implementation
|
||||
// lives in `src/kilocode/provider/provider.ts` so upstream diffs stay minimal.
|
||||
function wrapSSE(res: Response, chunkTimeout: number, ctl: AbortController, firstTokenMs: number) {
|
||||
return wrapSSEFirstContent(res, chunkTimeout, ctl, firstTokenMs)
|
||||
function wrapSSE(res: Response, ms: number, ctl: AbortController) {
|
||||
if (typeof ms !== "number" || ms <= 0) return res
|
||||
if (!res.body) return res
|
||||
if (!res.headers.get("content-type")?.includes("text/event-stream")) return res
|
||||
|
||||
const reader = res.body.getReader()
|
||||
const body = new ReadableStream<Uint8Array>({
|
||||
async pull(ctrl) {
|
||||
const part = await new Promise<Awaited<ReturnType<typeof reader.read>>>((resolve, reject) => {
|
||||
const id = setTimeout(() => {
|
||||
const err = new ProviderError.ResponseStreamError("SSE read timed out")
|
||||
ctl.abort(err)
|
||||
void reader.cancel(err)
|
||||
reject(err)
|
||||
}, ms)
|
||||
|
||||
reader.read().then(
|
||||
(part) => {
|
||||
clearTimeout(id)
|
||||
resolve(part)
|
||||
},
|
||||
(err) => {
|
||||
clearTimeout(id)
|
||||
reject(err)
|
||||
},
|
||||
)
|
||||
})
|
||||
|
||||
if (part.done) {
|
||||
ctrl.close()
|
||||
return
|
||||
}
|
||||
|
||||
ctrl.enqueue(part.value)
|
||||
},
|
||||
async cancel(reason) {
|
||||
ctl.abort(reason)
|
||||
await reader.cancel(reason)
|
||||
},
|
||||
})
|
||||
|
||||
return new Response(body, {
|
||||
headers: new Headers(res.headers),
|
||||
status: res.status,
|
||||
statusText: res.statusText,
|
||||
})
|
||||
}
|
||||
// kilocode_change end
|
||||
|
||||
function timeoutController(ms: number) {
|
||||
const ctl = new AbortController()
|
||||
@@ -1717,13 +1755,6 @@ export const layer = Layer.effect(
|
||||
const customFetch = options["fetch"]
|
||||
const chunkTimeout = options["chunkTimeout"]
|
||||
const headerTimeout = options["headerTimeout"]
|
||||
// kilocode_change start
|
||||
// Pre-content (time-to-first-content) budget for `wrapSSE`. Mirrors
|
||||
// `KiloLLM.resolveFirstTokenMs` semantics from S2: a positive finite
|
||||
// `options["timeout"]` wins; `false`/`0`/unset/invalid fall back to
|
||||
// `SSE_FIRST_TOKEN_MS` in the Kilo mirror. See `resolveSseFirstTokenMs`.
|
||||
const firstTokenMs = resolveSseFirstTokenMs(options)
|
||||
// kilocode_change end
|
||||
delete options["chunkTimeout"]
|
||||
delete options["headerTimeout"]
|
||||
|
||||
@@ -1753,7 +1784,7 @@ export const layer = Layer.effect(
|
||||
}).finally(() => headerTimeoutCtl?.clear())
|
||||
timeout.clear()
|
||||
if (!chunkAbortCtl) return res
|
||||
return wrapSSE(res, chunkTimeout, chunkAbortCtl, firstTokenMs)
|
||||
return wrapSSE(res, chunkTimeout, chunkAbortCtl)
|
||||
} catch (err) {
|
||||
timeout.clear()
|
||||
throw err
|
||||
|
||||
@@ -392,9 +392,7 @@ const live: Layer.Layer<
|
||||
toolChoice: input.toolChoice,
|
||||
maxOutputTokens: prepared.params.maxOutputTokens,
|
||||
abortSignal: input.abort,
|
||||
// kilocode_change: AI SDK's built-in chunk timeout is removed in favor
|
||||
// of a Kilo-owned per-event watchdog applied to the raw fullStream
|
||||
// before LLMAISDK.toLLMEvents normalization (see below).
|
||||
...KiloLLM.timeout({ options: prepared.params.options, fallback: item.options, log: l }), // kilocode_change
|
||||
headers: prepared.headers,
|
||||
maxRetries: input.retries ?? 0,
|
||||
messages: prepared.messages,
|
||||
@@ -422,27 +420,7 @@ const live: Layer.Layer<
|
||||
})
|
||||
// kilocode_change end
|
||||
// kilocode_change start - capture eligible session export request completion off the stream path
|
||||
// kilocode_change: resolve per-subscription idle watchdog so concurrent
|
||||
// sessions each get their own timer. Computed here (not at the stream
|
||||
// consumer) so the resolved value travels with the returned fullStream.
|
||||
const idleMs = KiloLLM.resolveIdleMs({
|
||||
options: prepared.params.options,
|
||||
fallback: item.options,
|
||||
})
|
||||
// kilocode_change: also resolve the pre-content (time-to-first-content /
|
||||
// prompt-processing) budget. The provider's own request `timeout` signal
|
||||
// is cleared once response headers arrive (buildTimeoutSignal,
|
||||
// src/kilocode/provider/provider.ts:254-272), so it does NOT bound the
|
||||
// pre-content phase — the watchdog is the enforcement mechanism here,
|
||||
// using the configured `timeout` VALUE as its budget. unset / false / 0
|
||||
// / invalid maps to DEFAULT_FIRST_TOKEN_MS (5 min), never to "disabled".
|
||||
// The same two option bags as `idleMs` are read so a model-level
|
||||
// `timeout` always wins over a provider-level `timeout`.
|
||||
const firstTokenMs = KiloLLM.resolveFirstTokenMs({
|
||||
options: prepared.params.options,
|
||||
fallback: item.options,
|
||||
})
|
||||
if (!exportable) return { type: "ai-sdk" as const, result, idleMs, firstTokenMs }
|
||||
if (!exportable) return { type: "ai-sdk" as const, result }
|
||||
return {
|
||||
type: "ai-sdk" as const,
|
||||
result: {
|
||||
@@ -456,8 +434,6 @@ const live: Layer.Layer<
|
||||
retries: input.retries ?? 0,
|
||||
}),
|
||||
},
|
||||
idleMs,
|
||||
firstTokenMs,
|
||||
}
|
||||
// kilocode_change end
|
||||
})
|
||||
@@ -478,25 +454,10 @@ const live: Layer.Layer<
|
||||
// Adapter seam: both runtimes expose the same LLMEvent stream. Native
|
||||
// already returns one; AI SDK streams are converted here.
|
||||
const state = LLMAISDK.adapterState()
|
||||
// kilocode_change start: wrap the raw AI SDK fullStream with the Kilo
|
||||
// idle watchdog before normalization. Per-subscription timers
|
||||
// (post-content `idleMs` and pre-content `firstTokenMs`) were
|
||||
// resolved inside `run` and travel with the result. Pass the
|
||||
// scoped controller so the watchdog can abort a stalled source
|
||||
// and avoid hanging cleanup.
|
||||
const watched = KiloLLM.watchdogAsyncIterable(
|
||||
result.result.fullStream as AsyncIterable<import("@ai-sdk/provider").LanguageModelV2StreamPart>,
|
||||
result.idleMs,
|
||||
ctrl,
|
||||
result.firstTokenMs,
|
||||
)
|
||||
// kilocode_change end
|
||||
return Stream.fromAsyncIterable(watched, (e) => (e instanceof Error ? e : new Error(String(e)))).pipe(
|
||||
// kilocode_change: the watchdog consumes raw LanguageModelV2 parts;
|
||||
// cast back to the TextStreamPart shape LLMAISDK.toLLMEvents expects.
|
||||
Stream.mapEffect((event) =>
|
||||
LLMAISDK.toLLMEvents(state, event as Parameters<typeof LLMAISDK.toLLMEvents>[1]),
|
||||
),
|
||||
return Stream.fromAsyncIterable(result.result.fullStream, (e) =>
|
||||
e instanceof Error ? e : new Error(String(e)),
|
||||
).pipe(
|
||||
Stream.mapEffect((event) => LLMAISDK.toLLMEvents(state, event)),
|
||||
Stream.flatMap((events) => Stream.fromIterable(events)),
|
||||
)
|
||||
}),
|
||||
|
||||
@@ -1,168 +0,0 @@
|
||||
import { describe, expect, test } from "bun:test"
|
||||
import { looksLikeContentEvent, resolveSseFirstTokenMs, SSE_FIRST_TOKEN_MS, wrapSSEFirstContent } from "@/kilocode/provider/provider"
|
||||
import { ProviderError } from "@/provider/error"
|
||||
|
||||
describe("kilocode.provider.resolveSseFirstTokenMs", () => {
|
||||
test("returns a positive finite timeout as-is", () => {
|
||||
expect(resolveSseFirstTokenMs({ timeout: 90_000 })).toBe(90_000)
|
||||
})
|
||||
|
||||
test("falls back to a positive finite provider timeout", () => {
|
||||
expect(resolveSseFirstTokenMs({ timeout: "90_000" }, { timeout: 30_000 })).toBe(30_000)
|
||||
})
|
||||
|
||||
test("maps false to the default", () => {
|
||||
expect(resolveSseFirstTokenMs({ timeout: false })).toBe(SSE_FIRST_TOKEN_MS)
|
||||
})
|
||||
|
||||
test("maps 0 to the default", () => {
|
||||
expect(resolveSseFirstTokenMs({ timeout: 0 })).toBe(SSE_FIRST_TOKEN_MS)
|
||||
})
|
||||
|
||||
test("maps negative / non-finite / invalid to the default", () => {
|
||||
expect(resolveSseFirstTokenMs({ timeout: -1 })).toBe(SSE_FIRST_TOKEN_MS)
|
||||
expect(resolveSseFirstTokenMs({ timeout: Number.POSITIVE_INFINITY })).toBe(SSE_FIRST_TOKEN_MS)
|
||||
expect(resolveSseFirstTokenMs({ timeout: Number.NaN })).toBe(SSE_FIRST_TOKEN_MS)
|
||||
expect(resolveSseFirstTokenMs({ timeout: "x" })).toBe(SSE_FIRST_TOKEN_MS)
|
||||
expect(resolveSseFirstTokenMs({})).toBe(SSE_FIRST_TOKEN_MS)
|
||||
})
|
||||
})
|
||||
|
||||
describe("kilocode.provider.looksLikeContentEvent", () => {
|
||||
test("returns false + full carry when no event boundary is present", () => {
|
||||
const text = 'data: {"choices":[{"delta":{"content":"incomplete'
|
||||
const result = looksLikeContentEvent(text)
|
||||
expect(result.found).toBe(false)
|
||||
expect(result.carry).toBe(text)
|
||||
})
|
||||
|
||||
test("role-only delta is not content", () => {
|
||||
const text = 'data: {"choices":[{"delta":{"role":"assistant"}}]}\n\n'
|
||||
const result = looksLikeContentEvent(text)
|
||||
expect(result.found).toBe(false)
|
||||
expect(result.carry).toBe("")
|
||||
})
|
||||
|
||||
test("empty delta is not content", () => {
|
||||
const text = 'data: {"choices":[{"delta":{}}]}\n\n'
|
||||
const result = looksLikeContentEvent(text)
|
||||
expect(result.found).toBe(false)
|
||||
expect(result.carry).toBe("")
|
||||
})
|
||||
|
||||
test(":heartbeat comment is not content", () => {
|
||||
const text = ":heartbeat\n\n"
|
||||
const result = looksLikeContentEvent(text)
|
||||
expect(result.found).toBe(false)
|
||||
expect(result.carry).toBe("")
|
||||
})
|
||||
|
||||
test("data: [DONE] is not content", () => {
|
||||
const text = "data: [DONE]\n\n"
|
||||
const result = looksLikeContentEvent(text)
|
||||
expect(result.found).toBe(false)
|
||||
expect(result.carry).toBe("")
|
||||
})
|
||||
|
||||
test("content delta is content", () => {
|
||||
const text = 'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n'
|
||||
const result = looksLikeContentEvent(text)
|
||||
expect(result.found).toBe(true)
|
||||
expect(result.carry).toBe("")
|
||||
})
|
||||
|
||||
test("reasoning_content delta is content", () => {
|
||||
const text = 'data: {"choices":[{"delta":{"reasoning_content":"thinking"}}]}\n\n'
|
||||
const result = looksLikeContentEvent(text)
|
||||
expect(result.found).toBe(true)
|
||||
expect(result.carry).toBe("")
|
||||
})
|
||||
|
||||
test("tool_calls delta is content", () => {
|
||||
const text = 'data: {"choices":[{"delta":{"tool_calls":[{"id":"1"}]}}]}\n\n'
|
||||
const result = looksLikeContentEvent(text)
|
||||
expect(result.found).toBe(true)
|
||||
expect(result.carry).toBe("")
|
||||
})
|
||||
|
||||
test("content event split across two reads is detected once the boundary arrives", () => {
|
||||
const head = 'data: {"choices":[{"delta":{"content":"h'
|
||||
const tail = 'i"}}]}\n\n'
|
||||
|
||||
const r1 = looksLikeContentEvent(head)
|
||||
expect(r1.found).toBe(false)
|
||||
|
||||
const r2 = looksLikeContentEvent(r1.carry + tail)
|
||||
expect(r2.found).toBe(true)
|
||||
})
|
||||
|
||||
test("keeps only the trailing fragment after a complete non-content event (buffer trimming)", () => {
|
||||
const text =
|
||||
'data: {"choices":[{"delta":{"role":"assistant"}}]}\n\ndata: {"choices":[{"delta":{"con'
|
||||
const result = looksLikeContentEvent(text)
|
||||
expect(result.found).toBe(false)
|
||||
expect(result.carry).toBe('data: {"choices":[{"delta":{"con')
|
||||
})
|
||||
|
||||
test("drops many complete heartbeats and keeps only the trailing fragment (Repair B)", () => {
|
||||
const text = Array.from({ length: 20 }, (_, i) => `:heartbeat${i}`).join("\n\n") + "\n\n:incomplete"
|
||||
const result = looksLikeContentEvent(text)
|
||||
expect(result.found).toBe(false)
|
||||
expect(result.carry).toBe(":incomplete")
|
||||
})
|
||||
})
|
||||
|
||||
describe("kilocode.provider.wrapSSEFirstContent", () => {
|
||||
test("returns the original response when chunkTimeout is not positive", () => {
|
||||
const res = new Response("body")
|
||||
expect(wrapSSEFirstContent(res, 0, new AbortController(), 60_000)).toBe(res)
|
||||
})
|
||||
|
||||
test("returns the original response when there is no body", () => {
|
||||
const res = new Response(null, { headers: { "content-type": "text/event-stream" } })
|
||||
expect(wrapSSEFirstContent(res, 50, new AbortController(), 60_000)).toBe(res)
|
||||
})
|
||||
|
||||
test("returns the original response when content-type is not SSE", () => {
|
||||
const res = new Response(new ReadableStream(), { headers: { "content-type": "application/json" } })
|
||||
expect(wrapSSEFirstContent(res, 50, new AbortController(), 60_000)).toBe(res)
|
||||
})
|
||||
|
||||
test("enqueues original bytes untouched and uses pre-content budget until first content", async () => {
|
||||
const stream = new ReadableStream<Uint8Array>({
|
||||
start(ctrl) {
|
||||
// Feed a non-content event first, then a content event.
|
||||
ctrl.enqueue(new TextEncoder().encode('data: {"choices":[{"delta":{"role":"assistant"}}]}\n\n'))
|
||||
ctrl.enqueue(new TextEncoder().encode('data: {"choices":[{"delta":{"content":"ok"}}]}\n\n'))
|
||||
ctrl.close()
|
||||
},
|
||||
})
|
||||
const res = new Response(stream, { headers: { "content-type": "text/event-stream" } })
|
||||
const ctl = new AbortController()
|
||||
const wrapped = wrapSSEFirstContent(res, 1_000, ctl, 1_000)
|
||||
const reader = wrapped.body!.getReader()
|
||||
const chunks: string[] = []
|
||||
let done = false
|
||||
while (!done) {
|
||||
const { value, done: d } = await reader.read()
|
||||
done = d
|
||||
if (value) chunks.push(new TextDecoder().decode(value))
|
||||
}
|
||||
expect(chunks.join("")).toContain('"role":"assistant"')
|
||||
expect(chunks.join("")).toContain('"content":"ok"')
|
||||
expect(ctl.signal.aborted).toBe(false)
|
||||
})
|
||||
|
||||
test("aborts a stalled pre-content SSE stream within firstTokenMs", async () => {
|
||||
const stream = new ReadableStream<Uint8Array>({
|
||||
start() {
|
||||
// Never enqueues — simulates a headers-only SSE connection.
|
||||
},
|
||||
})
|
||||
const res = new Response(stream, { headers: { "content-type": "text/event-stream" } })
|
||||
const ctl = new AbortController()
|
||||
const wrapped = wrapSSEFirstContent(res, 60_000, ctl, 50)
|
||||
const reader = wrapped.body!.getReader()
|
||||
await expect(reader.read()).rejects.toBeInstanceOf(ProviderError.ResponseStreamError)
|
||||
})
|
||||
})
|
||||
@@ -1,626 +0,0 @@
|
||||
import { NodeFileSystem } from "@effect/platform-node"
|
||||
import { afterEach, describe, expect } from "bun:test"
|
||||
import { Effect, Exit, Fiber, Layer } from "effect"
|
||||
import { FetchHttpClient } from "effect/unstable/http"
|
||||
import fs from "fs/promises"
|
||||
import { SessionV1 } from "@opencode-ai/core/v1/session"
|
||||
import { Database } from "@opencode-ai/core/database/database"
|
||||
import type { SessionID } from "../../src/session/schema"
|
||||
import path from "path"
|
||||
import { FSUtil } from "@opencode-ai/core/fs-util"
|
||||
import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner"
|
||||
import * as Log from "@opencode-ai/core/util/log"
|
||||
import { Agent as AgentSvc } from "../../src/agent/agent"
|
||||
import { BackgroundJob } from "../../src/background/job"
|
||||
import { Bus } from "../../src/bus"
|
||||
import { Command } from "../../src/command"
|
||||
import { Auth } from "../../src/auth"
|
||||
import { Config } from "../../src/config/config"
|
||||
import { RuntimeFlags } from "../../src/effect/runtime-flags"
|
||||
import { EventV2Bridge } from "../../src/event-v2-bridge"
|
||||
import { Env } from "../../src/env"
|
||||
import { Format } from "../../src/format"
|
||||
import { Git } from "../../src/git"
|
||||
import { Image } from "../../src/image/image"
|
||||
import { KiloSessions } from "../../src/kilo-sessions/kilo-sessions"
|
||||
import { LSP } from "../../src/lsp/lsp"
|
||||
import { MCP } from "../../src/mcp"
|
||||
import { Permission } from "../../src/permission"
|
||||
import { Plugin } from "../../src/plugin"
|
||||
import { Provider as ProviderSvc } from "../../src/provider/provider"
|
||||
import { Question } from "../../src/question"
|
||||
import { RepositoryCache } from "@opencode-ai/core/repository-cache"
|
||||
import { SessionCompaction } from "../../src/session/compaction"
|
||||
import { Instruction } from "../../src/session/instruction"
|
||||
import { LLM } from "../../src/session/llm"
|
||||
import { SessionProcessor } from "../../src/session/processor"
|
||||
import { SessionPrompt } from "../../src/session/prompt"
|
||||
import { SessionRevert } from "../../src/session/revert"
|
||||
import { SessionRunState } from "../../src/session/run-state"
|
||||
import { Session } from "../../src/session/session"
|
||||
import { SessionStatus } from "../../src/session/status"
|
||||
import { SystemPrompt } from "../../src/session/system"
|
||||
import { SessionSummary } from "../../src/session/summary"
|
||||
import { Todo } from "../../src/session/todo"
|
||||
import { Skill } from "../../src/skill"
|
||||
import { Snapshot } from "../../src/snapshot"
|
||||
import { Storage } from "../../src/storage/storage"
|
||||
import { SyncEvent } from "../../src/sync"
|
||||
import { Ripgrep } from "@opencode-ai/core/ripgrep"
|
||||
import { ToolRegistry } from "../../src/tool/registry"
|
||||
import { Truncate } from "../../src/tool/truncate"
|
||||
import { MemoryService } from "@kilocode/kilo-memory/effect/service"
|
||||
import { provideTmpdirServer } from "../fixture/fixture"
|
||||
import { awaitWithTimeout, pollWithTimeout, testEffect } from "../lib/effect"
|
||||
import { reply, TestLLMServer } from "../lib/llm-server"
|
||||
|
||||
void Log.init({ print: false })
|
||||
|
||||
afterEach(async () => {
|
||||
// Dispose all test instances between integration scenarios.
|
||||
const { disposeAllInstances } = await import("../fixture/fixture")
|
||||
await disposeAllInstances()
|
||||
})
|
||||
|
||||
const summary = Layer.succeed(
|
||||
SessionSummary.Service,
|
||||
SessionSummary.Service.of({
|
||||
summarize: () => Effect.void,
|
||||
diff: () => Effect.succeed([]),
|
||||
computeDiff: () => Effect.succeed([]),
|
||||
}),
|
||||
)
|
||||
|
||||
const mcp = Layer.succeed(
|
||||
MCP.Service,
|
||||
MCP.Service.of({
|
||||
status: () => Effect.succeed({}),
|
||||
clients: () => Effect.succeed({}),
|
||||
tools: () => Effect.succeed({}),
|
||||
prompts: () => Effect.succeed({}),
|
||||
resources: () => Effect.succeed({}),
|
||||
add: () => Effect.succeed({ status: { status: "disabled" as const } }),
|
||||
connect: () => Effect.void,
|
||||
disconnect: () => Effect.void,
|
||||
getPrompt: () => Effect.succeed(undefined),
|
||||
readResource: () => Effect.succeed(undefined),
|
||||
startAuth: () => Effect.die("unexpected MCP auth in watchdog tests"),
|
||||
authenticate: () => Effect.die("unexpected MCP auth in watchdog tests"),
|
||||
finishAuth: () => Effect.die("unexpected MCP auth in watchdog tests"),
|
||||
removeAuth: () => Effect.void,
|
||||
supportsOAuth: () => Effect.succeed(false),
|
||||
hasStoredTokens: () => Effect.succeed(false),
|
||||
getAuthStatus: () => Effect.succeed("not_authenticated" as const),
|
||||
}),
|
||||
)
|
||||
|
||||
const lsp = Layer.succeed(
|
||||
LSP.Service,
|
||||
LSP.Service.of({
|
||||
init: () => Effect.void,
|
||||
status: () => Effect.succeed([]),
|
||||
hasClients: () => Effect.succeed(false),
|
||||
touchFile: () => Effect.void,
|
||||
diagnostics: () => Effect.succeed({}),
|
||||
hover: () => Effect.succeed(undefined),
|
||||
definition: () => Effect.succeed([]),
|
||||
references: () => Effect.succeed([]),
|
||||
implementation: () => Effect.succeed([]),
|
||||
documentSymbol: () => Effect.succeed([]),
|
||||
workspaceSymbol: () => Effect.succeed([]),
|
||||
prepareCallHierarchy: () => Effect.succeed([]),
|
||||
incomingCalls: () => Effect.succeed([]),
|
||||
outgoingCalls: () => Effect.succeed([]),
|
||||
}),
|
||||
)
|
||||
|
||||
const status = Layer.mergeAll(SessionStatus.defaultLayer, Bus.layer)
|
||||
const run = SessionRunState.layer.pipe(Layer.provide(status))
|
||||
const infra = Layer.mergeAll(NodeFileSystem.layer, CrossSpawnSpawner.defaultLayer)
|
||||
|
||||
function makeHttp() {
|
||||
const deps = Layer.mergeAll(
|
||||
Session.defaultLayer,
|
||||
BackgroundJob.defaultLayer,
|
||||
Snapshot.defaultLayer,
|
||||
LLM.defaultLayer,
|
||||
Env.defaultLayer,
|
||||
AgentSvc.defaultLayer,
|
||||
Command.defaultLayer,
|
||||
Permission.defaultLayer,
|
||||
Plugin.defaultLayer,
|
||||
Config.defaultLayer,
|
||||
RuntimeFlags.layer(),
|
||||
ProviderSvc.defaultLayer,
|
||||
lsp,
|
||||
mcp,
|
||||
FSUtil.defaultLayer,
|
||||
SyncEvent.defaultLayer,
|
||||
EventV2Bridge.defaultLayer,
|
||||
Database.defaultLayer,
|
||||
status,
|
||||
MemoryService.layer,
|
||||
).pipe(Layer.provideMerge(infra))
|
||||
const question = Question.layer.pipe(Layer.provideMerge(deps))
|
||||
const todo = Todo.layer.pipe(Layer.provideMerge(deps))
|
||||
const registry = ToolRegistry.layer.pipe(
|
||||
Layer.provide(KiloSessions.testLayer),
|
||||
Layer.provide(Skill.defaultLayer),
|
||||
Layer.provide(FetchHttpClient.layer),
|
||||
Layer.provide(CrossSpawnSpawner.defaultLayer),
|
||||
Layer.provide(RepositoryCache.defaultLayer),
|
||||
Layer.provide(Ripgrep.defaultLayer),
|
||||
Layer.provide(Format.defaultLayer),
|
||||
Layer.provide(Git.defaultLayer),
|
||||
Layer.provide(Command.defaultLayer),
|
||||
Layer.provide(Auth.defaultLayer),
|
||||
Layer.provideMerge(todo),
|
||||
Layer.provideMerge(question),
|
||||
Layer.provideMerge(deps),
|
||||
)
|
||||
const trunc = Truncate.layer.pipe(Layer.provideMerge(deps))
|
||||
const proc = SessionProcessor.layer.pipe(
|
||||
Layer.provide(summary),
|
||||
Layer.provide(Image.defaultLayer),
|
||||
Layer.provideMerge(deps),
|
||||
)
|
||||
const compact = SessionCompaction.layer.pipe(Layer.provideMerge(proc), Layer.provideMerge(deps))
|
||||
return Layer.mergeAll(
|
||||
TestLLMServer.layer,
|
||||
SessionPrompt.layer.pipe(
|
||||
Layer.provide(SessionRevert.defaultLayer),
|
||||
Layer.provide(Image.defaultLayer),
|
||||
Layer.provide(summary),
|
||||
Layer.provideMerge(run),
|
||||
Layer.provideMerge(compact),
|
||||
Layer.provideMerge(proc),
|
||||
Layer.provideMerge(registry),
|
||||
Layer.provideMerge(trunc),
|
||||
Layer.provideMerge(question),
|
||||
Layer.provide(Instruction.defaultLayer),
|
||||
Layer.provide(SystemPrompt.defaultLayer),
|
||||
Layer.provideMerge(deps),
|
||||
),
|
||||
).pipe(
|
||||
Layer.provide(
|
||||
Layer.mergeAll(
|
||||
summary,
|
||||
deps,
|
||||
Config.defaultLayer,
|
||||
RuntimeFlags.layer(),
|
||||
BackgroundJob.defaultLayer,
|
||||
Bus.layer,
|
||||
infra,
|
||||
Storage.defaultLayer,
|
||||
),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
const it = testEffect(makeHttp())
|
||||
|
||||
const cfg = {
|
||||
provider: {
|
||||
test: {
|
||||
name: "Test",
|
||||
id: "test",
|
||||
env: [],
|
||||
npm: "@ai-sdk/openai-compatible",
|
||||
models: {
|
||||
"test-model": {
|
||||
id: "test-model",
|
||||
name: "Test Model",
|
||||
attachment: false,
|
||||
reasoning: false,
|
||||
temperature: false,
|
||||
tool_call: true,
|
||||
release_date: "2025-01-01",
|
||||
limit: { context: 100000, output: 10000 },
|
||||
cost: { input: 0, output: 0 },
|
||||
options: { chunkTimeout: 1_000 },
|
||||
},
|
||||
},
|
||||
options: {
|
||||
apiKey: "test-key",
|
||||
baseURL: "http://localhost:1/v1",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
function providerCfg(url: string) {
|
||||
return {
|
||||
...cfg,
|
||||
provider: {
|
||||
...cfg.provider,
|
||||
test: {
|
||||
...cfg.provider.test,
|
||||
options: {
|
||||
...cfg.provider.test.options,
|
||||
baseURL: url,
|
||||
chunkTimeout: false as const,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Model-level `chunkTimeout: 1_000` (from the base `cfg`'s `test-model`
|
||||
// options) is what arms the session watchdog because `resolveIdleMs` reads model
|
||||
// options first. The provider-level option is intentionally left unset so this
|
||||
// test exercises only the session-layer `watchIterator` path (not `wrapSSE`,
|
||||
// which arms only on a positive provider-level `chunkTimeout`).
|
||||
function enabledWatchdogCfg(url: string) {
|
||||
return {
|
||||
...cfg,
|
||||
provider: {
|
||||
...cfg.provider,
|
||||
test: {
|
||||
...cfg.provider.test,
|
||||
options: {
|
||||
...cfg.provider.test.options,
|
||||
baseURL: url,
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
const worktreeFile = (dir: string, name: string) => path.join(dir, name)
|
||||
|
||||
const exists = (file: string) =>
|
||||
Effect.promise(() =>
|
||||
fs
|
||||
.access(file)
|
||||
.then(() => true)
|
||||
.catch(() => false),
|
||||
)
|
||||
|
||||
// The production bash tool runs every command through a *login* shell
|
||||
// (`bash -l -c ...`, see src/shell/shell.ts) so `~/.bashrc` and shell
|
||||
// aliases behave the same as an interactive terminal. Git for Windows'
|
||||
// login-shell startup rescans the full Windows `PATH`, which is slower
|
||||
// than the Unix shells used elsewhere in this file. Give the marker file
|
||||
// these tests poll for a little extra headroom there, on top of the tests
|
||||
// A/C `config.shell: "bash"` override that makes the bash tool actually
|
||||
// use git-bash instead of cmd.exe on Windows (see those config comments).
|
||||
const waitForFile = (file: string, label: string, duration = process.platform === "win32" ? 15_000 : 5_000) =>
|
||||
pollWithTimeout(
|
||||
Effect.gen(function* () {
|
||||
const ok = yield* exists(file)
|
||||
return ok ? true : undefined
|
||||
}),
|
||||
label,
|
||||
duration,
|
||||
)
|
||||
|
||||
const touch = (file: string) => Effect.promise(() => fs.writeFile(file, ""))
|
||||
|
||||
const waitForRunningTool = (sessionID: SessionID, sessions: Session.Interface, label: string, duration = 15_000) =>
|
||||
pollWithTimeout(
|
||||
Effect.gen(function* () {
|
||||
const msgs = yield* sessions.messages({ sessionID })
|
||||
const running = msgs
|
||||
.flatMap((msg) => msg.parts)
|
||||
.find((part) => part.type === "tool" && part.state.status === "running")
|
||||
return running ? running : undefined
|
||||
}),
|
||||
label,
|
||||
duration,
|
||||
)
|
||||
|
||||
const waitForRequestHit = (llm: TestLLMServer["Service"], needle: string, label: string) =>
|
||||
pollWithTimeout(
|
||||
Effect.gen(function* () {
|
||||
const hits = yield* llm.hits
|
||||
const matched = hits.filter((hit) => JSON.stringify(hit.body).includes(needle))
|
||||
return matched.length > 0 ? matched : undefined
|
||||
}),
|
||||
label,
|
||||
5_000,
|
||||
)
|
||||
|
||||
const matchContains = (needle: string) => (hit: { body: Record<string, unknown> }) =>
|
||||
JSON.stringify(hit.body).includes(needle)
|
||||
|
||||
const assertNotInterrupted = (parts: SessionV1.WithParts["parts"]) => {
|
||||
for (const part of parts) {
|
||||
if (part.type === "tool") {
|
||||
expect(part.state.status).toBe("completed")
|
||||
if (part.state.status === "completed") {
|
||||
expect(part.state.metadata?.interrupted).not.toBe(true)
|
||||
expect(part.state.output).not.toContain("Tool execution aborted")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// kilocode_change: normalize to forward slashes before embedding in the
|
||||
// shell script. `ready`/`release` come from `path.join`, which yields
|
||||
// backslash-separated paths on Windows; inside a double-quoted git-bash
|
||||
// string a literal backslash is an escape character, so a Windows path can
|
||||
// silently mangle into the wrong filename (or a path bash's `[ -f ... ]`
|
||||
// test can't resolve) rather than throwing. Git-bash/MSYS accept
|
||||
// forward-slash paths natively, so this is safe on every platform this
|
||||
// suite runs on.
|
||||
const posixPath = (p: string) => p.replaceAll("\\", "/")
|
||||
|
||||
const bashGate = (dir: string, ready: string, release: string) =>
|
||||
`touch ${JSON.stringify(posixPath(ready))} && while [ ! -f ${JSON.stringify(posixPath(release))} ]; do sleep 0.05; done && echo done`
|
||||
|
||||
describe("session stream watchdog integration", () => {
|
||||
it.live(
|
||||
"A: root session long-running Bash is not interrupted by the idle watchdog",
|
||||
() =>
|
||||
provideTmpdirServer(
|
||||
Effect.fnUntraced(function* ({ dir, llm }) {
|
||||
const prompt = yield* SessionPrompt.Service
|
||||
const sessions = yield* Session.Service
|
||||
const chat = yield* sessions.create({ title: "Root long bash" })
|
||||
const ready = worktreeFile(dir, "bash-ready")
|
||||
const release = worktreeFile(dir, "bash-release")
|
||||
|
||||
yield* llm.tool("bash", {
|
||||
command: bashGate(dir, ready, release),
|
||||
description: "Long running bash command",
|
||||
timeout: 60_000,
|
||||
workdir: dir,
|
||||
})
|
||||
yield* llm.text("bash complete")
|
||||
|
||||
yield* prompt.prompt({
|
||||
sessionID: chat.id,
|
||||
agent: "build",
|
||||
noReply: true,
|
||||
parts: [{ type: "text", text: "run a long bash command" }],
|
||||
})
|
||||
|
||||
const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
|
||||
|
||||
yield* llm.wait(1)
|
||||
yield* waitForRunningTool(chat.id, sessions, "root bash tool never started")
|
||||
yield* waitForFile(ready, "root bash readiness marker never appeared")
|
||||
yield* Effect.sleep("1500 millis")
|
||||
|
||||
yield* touch(release)
|
||||
|
||||
const exit = yield* awaitWithTimeout(Fiber.await(fiber), "root bash loop did not finish", "15 seconds")
|
||||
expect(Exit.isSuccess(exit)).toBe(true)
|
||||
|
||||
// Check all messages in the session for interrupted tools
|
||||
const allMessages = yield* sessions.messages({ sessionID: chat.id })
|
||||
for (const msg of allMessages) {
|
||||
assertNotInterrupted(msg.parts)
|
||||
}
|
||||
}),
|
||||
{
|
||||
git: true,
|
||||
// kilocode_change: without an explicit `shell`, the bash tool's
|
||||
// `defaultShell()` falls back to cmd.exe on Windows (see
|
||||
// packages/core/src/tool/bash.ts), which cannot run bashGate's
|
||||
// POSIX syntax (`touch`, `[ -f ... ]`, `while ... done`). That
|
||||
// made `touch` fail immediately and silently, so the readiness
|
||||
// marker never appeared regardless of how long the test waited.
|
||||
config: (url) => ({ ...providerCfg(url), shell: "bash", permission: { bash: "allow" } }),
|
||||
},
|
||||
),
|
||||
{ timeout: 30_000 },
|
||||
)
|
||||
|
||||
it.live(
|
||||
"B: root foreground TaskTool child with held LLM response is not interrupted",
|
||||
() =>
|
||||
provideTmpdirServer(
|
||||
Effect.fnUntraced(function* ({ dir, llm }) {
|
||||
const prompt = yield* SessionPrompt.Service
|
||||
const sessions = yield* Session.Service
|
||||
const chat = yield* sessions.create({ title: "Root foreground child" })
|
||||
const gate = Promise.withResolvers<void>()
|
||||
|
||||
yield* llm.tool("task", {
|
||||
description: "Foreground child task",
|
||||
prompt: "child task: say hello",
|
||||
subagent_type: "child",
|
||||
})
|
||||
yield* llm.pushMatch(
|
||||
matchContains("child task: say hello"),
|
||||
reply().wait(gate.promise).text("child done").stop(),
|
||||
)
|
||||
yield* llm.text("parent done")
|
||||
|
||||
yield* prompt.prompt({
|
||||
sessionID: chat.id,
|
||||
agent: "build",
|
||||
noReply: true,
|
||||
parts: [{ type: "text", text: "run a foreground child" }],
|
||||
})
|
||||
|
||||
const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
|
||||
|
||||
yield* llm.wait(1)
|
||||
yield* waitForRunningTool(chat.id, sessions, "root task tool never started")
|
||||
yield* waitForRequestHit(llm, "child task: say hello", "child LLM request never hit server")
|
||||
yield* Effect.sleep("1500 millis")
|
||||
|
||||
gate.resolve(undefined)
|
||||
|
||||
const exit = yield* awaitWithTimeout(
|
||||
Fiber.await(fiber),
|
||||
"root foreground child loop did not finish",
|
||||
"15 seconds",
|
||||
)
|
||||
expect(Exit.isSuccess(exit)).toBe(true)
|
||||
|
||||
// Check all messages in root and child sessions for interrupted tools
|
||||
const allMessages = yield* sessions.messages({ sessionID: chat.id })
|
||||
for (const msg of allMessages) {
|
||||
assertNotInterrupted(msg.parts)
|
||||
}
|
||||
|
||||
const children = yield* sessions.children(chat.id)
|
||||
expect(children).toHaveLength(1)
|
||||
const childMessages = yield* sessions.messages({ sessionID: children[0]!.id })
|
||||
for (const msg of childMessages) {
|
||||
assertNotInterrupted(msg.parts)
|
||||
}
|
||||
}),
|
||||
{
|
||||
git: true,
|
||||
config: (url) => ({
|
||||
...providerCfg(url),
|
||||
permission: { bash: "allow", task: "allow" },
|
||||
agent: {
|
||||
child: {
|
||||
model: "test/test-model",
|
||||
mode: "subagent",
|
||||
options: { chunkTimeout: false },
|
||||
permission: { bash: "allow", task: "allow" },
|
||||
},
|
||||
},
|
||||
}),
|
||||
},
|
||||
),
|
||||
{ timeout: 30_000 },
|
||||
)
|
||||
|
||||
it.live(
|
||||
"C: child session long-running Bash while root awaits it is not interrupted",
|
||||
() =>
|
||||
provideTmpdirServer(
|
||||
Effect.fnUntraced(function* ({ dir, llm }) {
|
||||
const prompt = yield* SessionPrompt.Service
|
||||
const sessions = yield* Session.Service
|
||||
const chat = yield* sessions.create({ title: "Nested long bash" })
|
||||
const ready = worktreeFile(dir, "child-bash-ready")
|
||||
const release = worktreeFile(dir, "child-bash-release")
|
||||
|
||||
yield* llm.tool("task", {
|
||||
description: "Nested child task",
|
||||
prompt: "child task: run a long bash command",
|
||||
subagent_type: "child",
|
||||
})
|
||||
yield* llm.pushMatch(
|
||||
matchContains("child task: run a long bash command"),
|
||||
reply().tool("bash", {
|
||||
command: bashGate(dir, ready, release),
|
||||
description: "Long running child bash command",
|
||||
timeout: 60_000,
|
||||
workdir: dir,
|
||||
}),
|
||||
)
|
||||
yield* llm.text("child done")
|
||||
yield* llm.text("parent done")
|
||||
|
||||
yield* prompt.prompt({
|
||||
sessionID: chat.id,
|
||||
agent: "build",
|
||||
noReply: true,
|
||||
parts: [{ type: "text", text: "run a nested child" }],
|
||||
})
|
||||
|
||||
const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
|
||||
|
||||
yield* llm.wait(1)
|
||||
yield* waitForRunningTool(chat.id, sessions, "root task tool never started")
|
||||
yield* waitForRequestHit(llm, "child task: run a long bash command", "child LLM request never hit server")
|
||||
|
||||
const children = yield* sessions.children(chat.id)
|
||||
expect(children).toHaveLength(1)
|
||||
const childID = children[0]!.id
|
||||
|
||||
yield* waitForRunningTool(childID, sessions, "child bash tool never started")
|
||||
yield* waitForFile(ready, "child bash readiness marker never appeared")
|
||||
yield* Effect.sleep("1500 millis")
|
||||
|
||||
yield* touch(release)
|
||||
|
||||
const exit = yield* awaitWithTimeout(
|
||||
Fiber.await(fiber),
|
||||
"nested child bash loop did not finish",
|
||||
"15 seconds",
|
||||
)
|
||||
expect(Exit.isSuccess(exit)).toBe(true)
|
||||
|
||||
// Check all messages in root and child sessions for interrupted tools
|
||||
const rootMessages = yield* sessions.messages({ sessionID: chat.id })
|
||||
for (const msg of rootMessages) {
|
||||
assertNotInterrupted(msg.parts)
|
||||
}
|
||||
|
||||
// Reuse childID captured above
|
||||
const childMessages = yield* sessions.messages({ sessionID: childID })
|
||||
for (const msg of childMessages) {
|
||||
assertNotInterrupted(msg.parts)
|
||||
}
|
||||
}),
|
||||
{
|
||||
git: true,
|
||||
// kilocode_change: see the matching comment on test A — without
|
||||
// this, the nested child's bash tool falls back to cmd.exe on
|
||||
// Windows and the readiness marker never appears.
|
||||
config: (url) => ({
|
||||
...providerCfg(url),
|
||||
shell: "bash",
|
||||
permission: { bash: "allow", task: "allow" },
|
||||
agent: {
|
||||
child: {
|
||||
model: "test/test-model",
|
||||
mode: "subagent",
|
||||
permission: { bash: "allow", task: "allow" },
|
||||
},
|
||||
},
|
||||
}),
|
||||
},
|
||||
),
|
||||
{ timeout: 30_000 },
|
||||
)
|
||||
|
||||
it.live(
|
||||
"AC1: slow time-to-first-content is not aborted or retried",
|
||||
() =>
|
||||
provideTmpdirServer(
|
||||
Effect.fnUntraced(function* ({ llm }) {
|
||||
const prompt = yield* SessionPrompt.Service
|
||||
const sessions = yield* Session.Service
|
||||
const chat = yield* sessions.create({ title: "AC1 slow first content" })
|
||||
const gate = Promise.withResolvers<void>()
|
||||
|
||||
yield* llm.hold("done", gate.promise)
|
||||
|
||||
yield* prompt.prompt({
|
||||
sessionID: chat.id,
|
||||
agent: "build",
|
||||
noReply: true,
|
||||
parts: [{ type: "text", text: "say something after a long think" }],
|
||||
})
|
||||
|
||||
const fiber = yield* prompt.loop({ sessionID: chat.id }).pipe(Effect.forkChild)
|
||||
|
||||
yield* llm.wait(1)
|
||||
yield* Effect.sleep("2500 millis")
|
||||
gate.resolve(undefined)
|
||||
|
||||
const exit = yield* awaitWithTimeout(
|
||||
Fiber.await(fiber),
|
||||
"AC1 loop did not finish",
|
||||
"15 seconds",
|
||||
)
|
||||
expect(Exit.isSuccess(exit)).toBe(true)
|
||||
|
||||
const calls = yield* llm.calls
|
||||
expect(calls).toBe(1)
|
||||
|
||||
const messages = yield* sessions.messages({ sessionID: chat.id })
|
||||
const assistantText = messages
|
||||
.filter((msg) => msg.info.role === "assistant")
|
||||
.flatMap((msg) => msg.parts)
|
||||
.filter((part) => part.type === "text")
|
||||
.map((part) => part.text)
|
||||
.join("")
|
||||
expect(assistantText).toContain("done")
|
||||
}),
|
||||
{ git: true, config: (url) => enabledWatchdogCfg(url) },
|
||||
),
|
||||
{ timeout: 30_000 },
|
||||
)
|
||||
})
|
||||
@@ -3,84 +3,36 @@ import { Effect, Stream } from "effect"
|
||||
import { LLMEvent } from "@opencode-ai/llm"
|
||||
import { KiloLLM } from "@/kilocode/session/llm"
|
||||
|
||||
describe("kilocode.session.llm.resolveIdleMs", () => {
|
||||
describe("kilocode.session.llm.timeout", () => {
|
||||
test("uses prepared options before the provider fallback", () => {
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({
|
||||
options: { chunkTimeout: 15_000 },
|
||||
fallback: { chunkTimeout: 30_000 },
|
||||
}),
|
||||
).toBe(15_000)
|
||||
const result = KiloLLM.timeout({
|
||||
options: { chunkTimeout: 15_000 },
|
||||
fallback: { chunkTimeout: 30_000 },
|
||||
})
|
||||
|
||||
expect(result).toEqual({ timeout: { chunkMs: 15_000 } })
|
||||
})
|
||||
|
||||
test("uses the provider fallback when prepared options omit the timeout", () => {
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({
|
||||
options: {},
|
||||
fallback: { chunkTimeout: 30_000 },
|
||||
}),
|
||||
).toBe(30_000)
|
||||
const result = KiloLLM.timeout({
|
||||
options: {},
|
||||
fallback: { chunkTimeout: 30_000 },
|
||||
})
|
||||
|
||||
expect(result).toEqual({ timeout: { chunkMs: 30_000 } })
|
||||
})
|
||||
|
||||
test("uses the provider fallback when the prepared value is not a number", () => {
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({
|
||||
options: { chunkTimeout: "15_000" },
|
||||
fallback: { chunkTimeout: 30_000 },
|
||||
}),
|
||||
).toBe(30_000)
|
||||
const result = KiloLLM.timeout({
|
||||
options: { chunkTimeout: "15_000" },
|
||||
fallback: { chunkTimeout: 30_000 },
|
||||
})
|
||||
|
||||
expect(result).toEqual({ timeout: { chunkMs: 30_000 } })
|
||||
})
|
||||
|
||||
test("defaults the chunk idle timeout to 300_000 ms when no override is configured", () => {
|
||||
expect(KiloLLM.resolveIdleMs({ options: {} })).toBe(300_000)
|
||||
})
|
||||
|
||||
test("returns undefined when prepared is false (disabled)", () => {
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({
|
||||
options: { chunkTimeout: false },
|
||||
fallback: { chunkTimeout: 30_000 },
|
||||
}),
|
||||
).toBeUndefined()
|
||||
})
|
||||
|
||||
test("returns undefined when prepared is 0 (internal disable)", () => {
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({
|
||||
options: { chunkTimeout: 0 },
|
||||
fallback: { chunkTimeout: 30_000 },
|
||||
}),
|
||||
).toBeUndefined()
|
||||
})
|
||||
|
||||
test("returns undefined when provider fallback is false", () => {
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({
|
||||
options: {},
|
||||
fallback: { chunkTimeout: false },
|
||||
}),
|
||||
).toBeUndefined()
|
||||
})
|
||||
|
||||
test("falls through invalid prepared values to provider fallback", () => {
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({
|
||||
options: { chunkTimeout: -1 },
|
||||
fallback: { chunkTimeout: 5_000 },
|
||||
}),
|
||||
).toBe(5_000)
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({
|
||||
options: { chunkTimeout: Number.POSITIVE_INFINITY },
|
||||
fallback: { chunkTimeout: 5_000 },
|
||||
}),
|
||||
).toBe(5_000)
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({
|
||||
options: { chunkTimeout: Number.NaN },
|
||||
fallback: { chunkTimeout: 5_000 },
|
||||
}),
|
||||
).toBe(5_000)
|
||||
test("omits the timeout when it is not configured", () => {
|
||||
expect(KiloLLM.timeout({ options: {} })).toEqual({})
|
||||
})
|
||||
})
|
||||
|
||||
|
||||
@@ -1,499 +0,0 @@
|
||||
import { describe, expect, test } from "bun:test"
|
||||
import { Effect, Stream } from "effect"
|
||||
import type { LanguageModelV2CallWarning, LanguageModelV2StreamPart } from "@ai-sdk/provider"
|
||||
import { KiloLLM } from "@/kilocode/session/llm"
|
||||
import { ProviderError } from "@/provider/error"
|
||||
|
||||
type FullStreamPart = LanguageModelV2StreamPart
|
||||
|
||||
function part(type: string, extra: Record<string, unknown> = {}): FullStreamPart {
|
||||
return { type, ...extra } as unknown as FullStreamPart
|
||||
}
|
||||
|
||||
async function run<T>(eff: Effect.Effect<T, unknown>) {
|
||||
return await Effect.runPromise(eff)
|
||||
}
|
||||
|
||||
function fromSchedule(events: Array<[number, FullStreamPart]>, end: number): Stream.Stream<FullStreamPart, never> {
|
||||
// Each `[at, value]` is an ABSOLUTE time in milliseconds from stream start.
|
||||
// This matches how the tests are written: post-tool events (finish-step,
|
||||
// finish) are scheduled within a short idle window of the tool-result so
|
||||
// the watchdog, after the local active set drains, still receives the next
|
||||
// event in time.
|
||||
return Stream.fromAsyncIterable(
|
||||
(async function* () {
|
||||
const start = Date.now()
|
||||
for (const [at, value] of events) {
|
||||
const wait = at - (Date.now() - start)
|
||||
if (wait > 0) await new Promise((r) => setTimeout(r, wait))
|
||||
yield value
|
||||
}
|
||||
await new Promise((r) => setTimeout(r, end))
|
||||
})(),
|
||||
(e) => e as never,
|
||||
)
|
||||
}
|
||||
|
||||
// kilocode_change: async-generator version of `fromSchedule` that yields an
|
||||
// `AsyncIterable` for tests driving `watchdogAsyncIterable` directly. The
|
||||
// production consumer in `src/session/llm.ts` wraps the AI SDK's native
|
||||
// `fullStream` AsyncIterable and uses `watchdogAsyncIterable`, not
|
||||
// `watchdogStream`, so the direct-unit AC1/AC3 tests target that same entry
|
||||
// point. This helper is fine for tests that COMPLETE normally (e.g. AC1). The
|
||||
// AC3 test, which must abort during a stall, instead uses a separate
|
||||
// hand-rolled `AsyncIterable` (not this generator) because an async
|
||||
// generator's `return()` cannot preempt an in-flight internal `await` — see
|
||||
// the AC3 test comment for that distinction.
|
||||
function iterableFromSchedule(
|
||||
events: Array<[number, FullStreamPart]>,
|
||||
end: number,
|
||||
): AsyncIterable<FullStreamPart> {
|
||||
return (async function* () {
|
||||
const start = Date.now()
|
||||
for (const [at, value] of events) {
|
||||
const wait = at - (Date.now() - start)
|
||||
if (wait > 0) await new Promise((r) => setTimeout(r, wait))
|
||||
yield value
|
||||
}
|
||||
await new Promise((r) => setTimeout(r, end))
|
||||
})()
|
||||
}
|
||||
|
||||
describe("kilocode.session.llm.resolveIdleMs", () => {
|
||||
test("returns prepared positive finite value as-is", () => {
|
||||
const out = KiloLLM.resolveIdleMs({ options: { chunkTimeout: 15_000 }, fallback: { chunkTimeout: 30_000 } })
|
||||
expect(out).toBe(15_000)
|
||||
})
|
||||
|
||||
test("falls back to provider value when prepared is missing", () => {
|
||||
const out = KiloLLM.resolveIdleMs({ options: {}, fallback: { chunkTimeout: 30_000 } })
|
||||
expect(out).toBe(30_000)
|
||||
})
|
||||
|
||||
test("falls back to provider value when prepared is a non-number string", () => {
|
||||
const out = KiloLLM.resolveIdleMs({
|
||||
options: { chunkTimeout: "15_000" },
|
||||
fallback: { chunkTimeout: 30_000 },
|
||||
})
|
||||
expect(out).toBe(30_000)
|
||||
})
|
||||
|
||||
test("falls back to provider value when prepared is negative", () => {
|
||||
const out = KiloLLM.resolveIdleMs({
|
||||
options: { chunkTimeout: -1 },
|
||||
fallback: { chunkTimeout: 30_000 },
|
||||
})
|
||||
expect(out).toBe(30_000)
|
||||
})
|
||||
|
||||
test("falls back to provider value when prepared is non-finite (Infinity, NaN)", () => {
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({ options: { chunkTimeout: Number.POSITIVE_INFINITY }, fallback: { chunkTimeout: 5_000 } }),
|
||||
).toBe(5_000)
|
||||
expect(KiloLLM.resolveIdleMs({ options: { chunkTimeout: Number.NaN }, fallback: { chunkTimeout: 5_000 } })).toBe(
|
||||
5_000,
|
||||
)
|
||||
})
|
||||
|
||||
test("treats boolean false as a request to disable the watchdog", () => {
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({ options: { chunkTimeout: false }, fallback: { chunkTimeout: 30_000 } }),
|
||||
).toBeUndefined()
|
||||
})
|
||||
|
||||
test("treats internal 0 as a request to disable the watchdog", () => {
|
||||
expect(KiloLLM.resolveIdleMs({ options: { chunkTimeout: 0 }, fallback: { chunkTimeout: 30_000 } })).toBeUndefined()
|
||||
})
|
||||
|
||||
test("provider fallback false also disables", () => {
|
||||
expect(KiloLLM.resolveIdleMs({ options: {}, fallback: { chunkTimeout: false } })).toBeUndefined()
|
||||
})
|
||||
|
||||
test("uses 300_000 default when nothing valid is configured", () => {
|
||||
expect(KiloLLM.resolveIdleMs({ options: {} })).toBe(300_000)
|
||||
})
|
||||
|
||||
test("uses 300_000 default when both prepared and fallback are invalid", () => {
|
||||
expect(
|
||||
KiloLLM.resolveIdleMs({
|
||||
options: { chunkTimeout: "x" },
|
||||
fallback: { chunkTimeout: -5 },
|
||||
}),
|
||||
).toBe(300_000)
|
||||
})
|
||||
})
|
||||
|
||||
// kilocode_change: resolveFirstTokenMs is the new S2 helper for the
|
||||
// pre-content (time-to-first-content) budget. Unlike resolveIdleMs it
|
||||
// NEVER returns undefined: the pre-content bound is disabled only when
|
||||
// the whole watchdog is disabled (idleMs === undefined at the entry
|
||||
// point). unset / false / 0 / invalid / non-finite / negative MUST map
|
||||
// to DEFAULT_FIRST_TOKEN_MS — leaving a never-first-content hang
|
||||
// unbounded when `timeout` is configured as "off" would re-introduce a
|
||||
// worse version of #12467 (silent infinite hang instead of a 60s false
|
||||
// positive).
|
||||
describe("kilocode.session.llm.resolveFirstTokenMs", () => {
|
||||
test("returns prepared positive finite value as-is", () => {
|
||||
const out = KiloLLM.resolveFirstTokenMs({
|
||||
options: { timeout: 90_000 },
|
||||
fallback: { timeout: 30_000 },
|
||||
})
|
||||
expect(out).toBe(90_000)
|
||||
})
|
||||
|
||||
test("falls back to provider value when prepared is missing", () => {
|
||||
const out = KiloLLM.resolveFirstTokenMs({ options: {}, fallback: { timeout: 30_000 } })
|
||||
expect(out).toBe(30_000)
|
||||
})
|
||||
|
||||
test("falls back to provider value when prepared is a non-number string", () => {
|
||||
const out = KiloLLM.resolveFirstTokenMs({
|
||||
options: { timeout: "90_000" },
|
||||
fallback: { timeout: 30_000 },
|
||||
})
|
||||
expect(out).toBe(30_000)
|
||||
})
|
||||
|
||||
test("falls back to provider value when prepared is negative", () => {
|
||||
const out = KiloLLM.resolveFirstTokenMs({
|
||||
options: { timeout: -1 },
|
||||
fallback: { timeout: 30_000 },
|
||||
})
|
||||
expect(out).toBe(30_000)
|
||||
})
|
||||
|
||||
test("falls back to provider value when prepared is non-finite (Infinity, NaN)", () => {
|
||||
expect(
|
||||
KiloLLM.resolveFirstTokenMs({
|
||||
options: { timeout: Number.POSITIVE_INFINITY },
|
||||
fallback: { timeout: 5_000 },
|
||||
}),
|
||||
).toBe(5_000)
|
||||
expect(
|
||||
KiloLLM.resolveFirstTokenMs({ options: { timeout: Number.NaN }, fallback: { timeout: 5_000 } }),
|
||||
).toBe(5_000)
|
||||
})
|
||||
|
||||
// Crucial divergence from resolveIdleMs: boolean false must NOT disable
|
||||
// the pre-content bound — it falls back / uses the Kilo default instead.
|
||||
test("treats boolean false as not-configured (falls back, does NOT disable)", () => {
|
||||
expect(
|
||||
KiloLLM.resolveFirstTokenMs({ options: { timeout: false }, fallback: { timeout: 30_000 } }),
|
||||
).toBe(30_000)
|
||||
expect(KiloLLM.resolveFirstTokenMs({ options: { timeout: false } })).toBe(300_000)
|
||||
})
|
||||
|
||||
test("treats internal 0 as not-configured (falls back, does NOT disable)", () => {
|
||||
expect(
|
||||
KiloLLM.resolveFirstTokenMs({ options: { timeout: 0 }, fallback: { timeout: 30_000 } }),
|
||||
).toBe(30_000)
|
||||
expect(KiloLLM.resolveFirstTokenMs({ options: { timeout: 0 } })).toBe(300_000)
|
||||
})
|
||||
|
||||
test("uses 300_000 default when nothing valid is configured", () => {
|
||||
expect(KiloLLM.resolveFirstTokenMs({ options: {} })).toBe(300_000)
|
||||
})
|
||||
|
||||
test("uses 300_000 default when both prepared and fallback are invalid", () => {
|
||||
expect(
|
||||
KiloLLM.resolveFirstTokenMs({
|
||||
options: { timeout: "x" },
|
||||
fallback: { timeout: -5 },
|
||||
}),
|
||||
).toBe(300_000)
|
||||
})
|
||||
})
|
||||
|
||||
describe("kilocode.session.llm.watchdogStream", () => {
|
||||
test("returns the stream unchanged when idle is undefined (disabled)", async () => {
|
||||
const events: FullStreamPart[] = [
|
||||
part("stream-start", { warnings: [] as LanguageModelV2CallWarning[] }),
|
||||
part("text-delta", { id: "t1", delta: "ok" }),
|
||||
]
|
||||
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(Stream.fromIterable(events), undefined)))
|
||||
expect(out.length).toBe(2)
|
||||
})
|
||||
|
||||
test("emits events and completes when the stream delivers them within the idle window", async () => {
|
||||
const events: FullStreamPart[] = [
|
||||
part("stream-start", { warnings: [] as LanguageModelV2CallWarning[] }),
|
||||
part("text-delta", { id: "t1", delta: "hi" }),
|
||||
part("text-delta", { id: "t1", delta: "!" }),
|
||||
]
|
||||
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(Stream.fromIterable(events), 1_000)))
|
||||
expect(out.length).toBe(3)
|
||||
})
|
||||
|
||||
test("fails with ProviderError.ResponseStreamError when the stream stalls", async () => {
|
||||
const slow = Stream.fromEffect(
|
||||
Effect.flatMap(Effect.sleep("5 seconds"), () => Effect.succeed(part("text-delta", { id: "t1", delta: "x" }))),
|
||||
)
|
||||
const wrapped = KiloLLM.watchdogStream(slow, 100)
|
||||
const err = await run(Effect.flip(Stream.runCollect(wrapped)))
|
||||
expect(err).toBeInstanceOf(ProviderError.ResponseStreamError)
|
||||
})
|
||||
|
||||
test("every raw AI SDK event resets the idle timer", async () => {
|
||||
// idle 200ms; emit text-delta at 0 and 60ms (a single 260ms pull would time out without reset).
|
||||
const stream = fromSchedule(
|
||||
[
|
||||
[0, part("text-delta", { id: "t1", delta: "a" })],
|
||||
[60, part("text-delta", { id: "t1", delta: "b" })],
|
||||
],
|
||||
10,
|
||||
)
|
||||
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(stream, 200)))
|
||||
expect(out.length).toBe(2)
|
||||
})
|
||||
|
||||
test("AC2: mid-response stall after first content is still aborted", async () => {
|
||||
// A content part at t=0, then a long quiet gap that exceeds the idle window.
|
||||
const stream = fromSchedule([[0, part("text-delta", { id: "t1", delta: "hi" })]], 400)
|
||||
const err = await run(Effect.flip(Stream.runCollect(KiloLLM.watchdogStream(stream, 200))))
|
||||
expect(err).toBeInstanceOf(ProviderError.ResponseStreamError)
|
||||
})
|
||||
|
||||
test("pending local tool calls suspend the idle timeout until they settle", async () => {
|
||||
// tool-call at t=0 (local). 250ms quiet gap then tool-result. A healthy AI
|
||||
// SDK run also emits finish-step + finish right after the tool-result, so
|
||||
// the watchdog sees another event within idleMs and resets.
|
||||
const stream = fromSchedule(
|
||||
[
|
||||
[0, part("tool-call", { toolCallId: "c1", toolName: "bash" })],
|
||||
[250, part("tool-result", { toolCallId: "c1", toolName: "bash", output: "ok" })],
|
||||
[260, part("finish-step", { finishReason: "tool-calls" })],
|
||||
[270, part("finish", { finishReason: "stop" })],
|
||||
],
|
||||
10,
|
||||
)
|
||||
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(stream, 200)))
|
||||
expect(out.length).toBe(4)
|
||||
})
|
||||
|
||||
test("provider-executed tool calls do not suspend the watchdog", async () => {
|
||||
const stream = fromSchedule(
|
||||
[[0, part("tool-call", { toolCallId: "c1", toolName: "web", providerExecuted: true })]],
|
||||
400,
|
||||
)
|
||||
const err = await run(Effect.flip(Stream.runCollect(KiloLLM.watchdogStream(stream, 200))))
|
||||
expect(err).toBeInstanceOf(ProviderError.ResponseStreamError)
|
||||
})
|
||||
|
||||
test("parallel local tool calls remain suspended until the last settles", async () => {
|
||||
const stream = fromSchedule(
|
||||
[
|
||||
[0, part("tool-call", { toolCallId: "a", toolName: "bash" })],
|
||||
[10, part("tool-call", { toolCallId: "b", toolName: "bash" })],
|
||||
[200, part("tool-result", { toolCallId: "a", toolName: "bash", output: "x" })],
|
||||
[350, part("tool-result", { toolCallId: "b", toolName: "bash", output: "y" })],
|
||||
[360, part("finish-step", { finishReason: "tool-calls" })],
|
||||
[370, part("finish", { finishReason: "stop" })],
|
||||
],
|
||||
10,
|
||||
)
|
||||
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(stream, 200)))
|
||||
expect(out.length).toBe(6)
|
||||
})
|
||||
|
||||
test("tool-error for a local tool id also releases the suspension", async () => {
|
||||
const stream = fromSchedule(
|
||||
[
|
||||
[0, part("tool-call", { toolCallId: "c1", toolName: "bash" })],
|
||||
[200, part("tool-error", { toolCallId: "c1", toolName: "bash", error: new Error("nope") })],
|
||||
[210, part("finish-step", { finishReason: "tool-calls" })],
|
||||
[220, part("finish", { finishReason: "stop" })],
|
||||
],
|
||||
10,
|
||||
)
|
||||
const out = await run(Stream.runCollect(KiloLLM.watchdogStream(stream, 200)))
|
||||
expect(out.length).toBe(4)
|
||||
})
|
||||
|
||||
test("aborts the underlying source on timeout so cleanup does not hang", async () => {
|
||||
const ctrl = new AbortController()
|
||||
let abortReason: unknown
|
||||
let nextResolved = false
|
||||
const source: AsyncIterable<FullStreamPart> = {
|
||||
[Symbol.asyncIterator]() {
|
||||
let nextPromise: Promise<IteratorResult<FullStreamPart>> | undefined
|
||||
let resolveNext: ((value: IteratorResult<FullStreamPart>) => void) | undefined
|
||||
ctrl.signal.addEventListener("abort", () => {
|
||||
abortReason = ctrl.signal.reason
|
||||
if (resolveNext) {
|
||||
resolveNext({ done: true, value: undefined })
|
||||
nextResolved = true
|
||||
}
|
||||
})
|
||||
return {
|
||||
next() {
|
||||
nextPromise = new Promise((resolve) => {
|
||||
resolveNext = resolve
|
||||
})
|
||||
return nextPromise
|
||||
},
|
||||
async return() {
|
||||
if (nextPromise) await nextPromise
|
||||
return { done: true, value: undefined }
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
const wrapped = KiloLLM.watchdogAsyncIterable(source, 100, ctrl)
|
||||
const err = await run(Effect.flip(Stream.runCollect(Stream.fromAsyncIterable(wrapped, (e) => e as never))))
|
||||
expect(err).toBeInstanceOf(ProviderError.ResponseStreamError)
|
||||
expect(nextResolved).toBe(true)
|
||||
expect(abortReason).toBeInstanceOf(ProviderError.ResponseStreamError)
|
||||
})
|
||||
|
||||
test("propagates upstream stream errors without false timeout", async () => {
|
||||
const stream = Stream.fail(new Error("upstream broken"))
|
||||
const err = await run(Effect.flip(Stream.runCollect(KiloLLM.watchdogStream(stream, 1_000))))
|
||||
expect((err as Error).message).toBe("upstream broken")
|
||||
})
|
||||
|
||||
test("return() closes the source immediately without waiting on a stalled pull", async () => {
|
||||
// Regression test: a hand-rolled async generator's `.return()` cannot
|
||||
// preempt an in-flight internal `await` — it only takes effect once that
|
||||
// await settles on its own, which never happens for a genuinely stalled
|
||||
// source. `watchdogAsyncIterable` must instead expose a `return()` that
|
||||
// runs immediately and forwards to the source's `return()` without
|
||||
// waiting for the outstanding `next()` to resolve.
|
||||
let sourceReturnCalled = false
|
||||
let neverResolvingNextCalled = false
|
||||
const source: AsyncIterable<FullStreamPart> = {
|
||||
[Symbol.asyncIterator]() {
|
||||
return {
|
||||
next() {
|
||||
neverResolvingNextCalled = true
|
||||
return new Promise<IteratorResult<FullStreamPart>>(() => {
|
||||
// Never resolves — simulates a fully stalled source (e.g. a
|
||||
// hung fetch response) whose pending pull is abandoned once
|
||||
// the consumer decides to stop.
|
||||
})
|
||||
},
|
||||
async return() {
|
||||
sourceReturnCalled = true
|
||||
return { done: true, value: undefined }
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
const wrapped = KiloLLM.watchdogAsyncIterable(source, 60_000)
|
||||
const it = wrapped[Symbol.asyncIterator]()
|
||||
const pending = it.next()
|
||||
expect(neverResolvingNextCalled).toBe(true)
|
||||
|
||||
const returned = await Promise.race([
|
||||
it.return!(),
|
||||
new Promise((_, reject) => setTimeout(() => reject(new Error("return() hung")), 500)),
|
||||
])
|
||||
expect(returned).toMatchObject({ done: true })
|
||||
expect(sourceReturnCalled).toBe(true)
|
||||
// The abandoned pull is left unresolved; only return() is asserted here.
|
||||
void pending
|
||||
})
|
||||
|
||||
// kilocode_change: AC1 (parameterised) — slow time-to-first-content is
|
||||
// bounded by `firstTokenMs`, not `idleMs`. The first content part
|
||||
// (`text-delta`) is delayed 300ms, which is past the 100ms `idleMs` but
|
||||
// within the 1000ms `firstTokenMs` (the request-timeout budget). Without
|
||||
// the fix the per-event idle would fire at t≈100ms and abort a healthy
|
||||
// slow stream (issue #12467). With the fix the pre-content phase is
|
||||
// bounded by the larger `firstTokenMs` and the stream completes.
|
||||
//
|
||||
// Uses `watchdogAsyncIterable` directly (with a raw AsyncIterable) to match
|
||||
// the production call site (`session/llm.ts`), which wraps the AI SDK's
|
||||
// native `fullStream` AsyncIterable. `watchdogStream` would round-trip
|
||||
// through `Stream.toAsyncIterable`, which blocks the event loop on long
|
||||
// stalls and prevents the watchdog's setTimeout from firing.
|
||||
test("AC1 (parameterised): slow time-to-first-content completes when firstTokenMs > idleMs", async () => {
|
||||
// Synthetic structural parts arrive instantly (mirroring the AI SDK's
|
||||
// t+2ms synthetic `start`); the first content part is gated 300ms.
|
||||
const source = iterableFromSchedule(
|
||||
[
|
||||
[0, part("stream-start", { warnings: [] as LanguageModelV2CallWarning[] })],
|
||||
[0, part("start-step", { request: {}, warnings: [] as LanguageModelV2CallWarning[] })],
|
||||
[0, part("text-start", { id: "t1", providerMetadata: undefined })],
|
||||
[300, part("text-delta", { id: "t1", delta: "hi", providerMetadata: undefined })],
|
||||
[310, part("text-delta", { id: "t1", delta: "!", providerMetadata: undefined })],
|
||||
[320, part("finish-step", { finishReason: "stop", usage: undefined, providerMetadata: undefined })],
|
||||
[330, part("finish", { finishReason: "stop", usage: undefined, providerMetadata: undefined })],
|
||||
],
|
||||
10,
|
||||
)
|
||||
const wrapped = KiloLLM.watchdogAsyncIterable(source, 100, undefined, 1_000)
|
||||
const out = await run(
|
||||
Stream.runCollect(Stream.fromAsyncIterable(wrapped, (e) => (e instanceof Error ? e : new Error(String(e))))),
|
||||
)
|
||||
// 7 raw events must be collected end-to-end — the 300ms wait for the
|
||||
// first content part exceeds idleMs=100 but is well under
|
||||
// firstTokenMs=1000, so the watchdog must NOT abort.
|
||||
expect(out.length).toBe(7)
|
||||
})
|
||||
|
||||
// kilocode_change: AC3 — a never-first-content hang is bounded by
|
||||
// `firstTokenMs` (the request-timeout budget), not by the post-content
|
||||
// `idleMs`. The schedule emits ONLY structural parts and then hangs; the
|
||||
// watchdog must fire at ~firstTokenMs (500ms), not at the much-larger
|
||||
// `idleMs` (5000ms) and not at the Kilo default (300s). This is the
|
||||
// primary direct-unit guard that a `timeout: false` / unset config
|
||||
// still gets a finite bound on time-to-first-content, since the
|
||||
// provider's own request `timeout` signal is cleared once response
|
||||
// headers arrive.
|
||||
//
|
||||
// Uses a hand-rolled `AsyncIterable` rather than an async generator
|
||||
// because the watchdog's catch path calls `safeClose(source.return())`
|
||||
// to clean up the source. On a hand-rolled generator that's suspended
|
||||
// mid-`setTimeout`, `return()` blocks until the pending `next()`
|
||||
// settles (per the async-generator spec), so the timeout would not
|
||||
// appear to fire until the source's setTimeout elapsed. A custom
|
||||
// `AsyncIterable` whose `return()` resolves immediately matches the
|
||||
// behavior of the AI SDK's native `fullStream` consumer in production
|
||||
// and is the same pattern used in the existing "aborts the underlying
|
||||
// source" test in this file.
|
||||
test("AC3: never-first-content hang is bounded by firstTokenMs, not idleMs", async () => {
|
||||
// Yield 3 structural parts instantly, then hang.
|
||||
const parts: FullStreamPart[] = [
|
||||
part("stream-start", { warnings: [] as LanguageModelV2CallWarning[] }),
|
||||
part("start-step", { request: {}, warnings: [] as LanguageModelV2CallWarning[] }),
|
||||
part("text-start", { id: "t1", providerMetadata: undefined }),
|
||||
]
|
||||
let index = 0
|
||||
let resolveHanging: ((v: IteratorResult<FullStreamPart>) => void) | undefined
|
||||
const source: AsyncIterable<FullStreamPart> = {
|
||||
[Symbol.asyncIterator]() {
|
||||
return {
|
||||
next() {
|
||||
if (index < parts.length) {
|
||||
return Promise.resolve({ done: false, value: parts[index++]! })
|
||||
}
|
||||
return new Promise<IteratorResult<FullStreamPart>>((resolve) => {
|
||||
resolveHanging = resolve
|
||||
})
|
||||
},
|
||||
return() {
|
||||
if (resolveHanging) resolveHanging({ done: true, value: undefined })
|
||||
return Promise.resolve({ done: true, value: undefined })
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
const start = Date.now()
|
||||
const wrapped = KiloLLM.watchdogAsyncIterable(source, 5_000, undefined, 500)
|
||||
const err = await run(
|
||||
Effect.flip(
|
||||
Stream.runCollect(
|
||||
Stream.fromAsyncIterable(wrapped, (e) => (e instanceof Error ? e : new Error(String(e)))),
|
||||
),
|
||||
),
|
||||
)
|
||||
const elapsed = Date.now() - start
|
||||
expect(err).toBeInstanceOf(ProviderError.ResponseStreamError)
|
||||
// Must abort at ~firstTokenMs (500ms), not at ~idleMs (5000ms). The
|
||||
// lower bound allows for setTimeout drift; the upper bound is well
|
||||
// below any reasonable post-content idle window.
|
||||
expect(elapsed).toBeGreaterThanOrEqual(400)
|
||||
expect(elapsed).toBeLessThan(2_000)
|
||||
})
|
||||
})
|
||||
@@ -45,59 +45,13 @@ it.live("headerTimeout does not abort delayed SSE body after headers arrive", ()
|
||||
}),
|
||||
)
|
||||
|
||||
// kilocode_change start - S2b: under the first-content-aware fix, a slow FIRST
|
||||
// content chunk is no longer raced against `chunkTimeout`; it is bounded by
|
||||
// the request `timeout` (or the Kilo default of 5 min) instead. With no
|
||||
// `timeout` configured here, a 250ms first-content delay is comfortably
|
||||
// within that budget, so the stream completes without error. The previous
|
||||
// assertion (that this would raise `ProviderError.ResponseStreamError`) is
|
||||
// what the fix removes.
|
||||
it.live("chunkTimeout does NOT abort a slow first content chunk (bounded by timeout instead)", () =>
|
||||
it.live("chunkTimeout raises a response stream error when SSE body stalls", () =>
|
||||
Effect.gen(function* () {
|
||||
const server = yield* Effect.acquireRelease(
|
||||
Effect.promise(() => delayedBodyServer(250)),
|
||||
(server) => Effect.sync(() => server.server.close()),
|
||||
)
|
||||
|
||||
yield* provideTmpdirInstance(
|
||||
() =>
|
||||
Effect.gen(function* () {
|
||||
const provider = yield* Provider.Service
|
||||
const model = yield* provider.getModel(ProviderV2.ID.make("test"), ModelV2.ID.make("test-model"))
|
||||
const result = streamText({
|
||||
model: yield* provider.getLanguage(model),
|
||||
onError() {},
|
||||
messages: [{ role: "user", content: "hello" }],
|
||||
})
|
||||
|
||||
const { error, text } = yield* Effect.promise(async () => {
|
||||
try {
|
||||
const text = await result.text
|
||||
return { error: undefined, text }
|
||||
} catch (error) {
|
||||
return { error, text: "" }
|
||||
}
|
||||
})
|
||||
expect(error).toBeUndefined()
|
||||
expect(text).toBe("late")
|
||||
}),
|
||||
{ config: providerConfig(server.url, { chunkTimeout: 50 }) },
|
||||
)
|
||||
}),
|
||||
)
|
||||
|
||||
// kilocode_change - S2b AC9b: mid-content stall protection is preserved. After
|
||||
// the first content-bearing SSE `data:` chunk is observed, `wrapSSE` reverts
|
||||
// to racing subsequent reads against `chunkTimeout`; a stall past that window
|
||||
// must still raise `ProviderError.ResponseStreamError` (this is the protection
|
||||
// the raw-`fullStream` memory/agent consumers rely on — see plan AC9).
|
||||
it.live("chunkTimeout still aborts a mid-content stall after the first content chunk", () =>
|
||||
Effect.gen(function* () {
|
||||
const server = yield* Effect.acquireRelease(
|
||||
Effect.promise(() => contentThenStallServer(5000)),
|
||||
(server) => Effect.sync(() => server.server.close()),
|
||||
)
|
||||
|
||||
yield* provideTmpdirInstance(
|
||||
() =>
|
||||
Effect.gen(function* () {
|
||||
@@ -124,7 +78,6 @@ it.live("chunkTimeout still aborts a mid-content stall after the first content c
|
||||
)
|
||||
}),
|
||||
)
|
||||
// kilocode_change end
|
||||
|
||||
it.live("headerTimeout aborts when response headers do not arrive", () =>
|
||||
Effect.gen(function* () {
|
||||
@@ -258,25 +211,6 @@ async function delayedBodyServer(delay: number, prelude = ""): Promise<{ server:
|
||||
if (!address || typeof address === "string") throw new Error("server did not bind to a TCP port")
|
||||
return { server, url: `http://127.0.0.1:${address.port}` }
|
||||
}
|
||||
|
||||
// Writes a content-bearing SSE `data:` chunk immediately, flushes, then holds
|
||||
// the connection open without sending any further bytes. Used to exercise the
|
||||
// post-first-content branch of `wrapSSE` (mid-content stall must still be
|
||||
// caught by `chunkTimeout`).
|
||||
async function contentThenStallServer(stallMs: number): Promise<{ server: Server; url: string }> {
|
||||
const server = createServer((_, res) => {
|
||||
res.writeHead(200, { "content-type": "text/event-stream" })
|
||||
res.flushHeaders()
|
||||
res.write('data: {"choices":[{"delta":{"content":"hi"}}]}\n\n')
|
||||
setTimeout(() => {
|
||||
res.end('data: {"choices":[{"delta":{"content":"there"}}]}\n\ndata: [DONE]\n\n')
|
||||
}, stallMs)
|
||||
})
|
||||
await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve))
|
||||
const address = server.address()
|
||||
if (!address || typeof address === "string") throw new Error("server did not bind to a TCP port")
|
||||
return { server, url: `http://127.0.0.1:${address.port}` }
|
||||
}
|
||||
// kilocode_change end
|
||||
|
||||
function withAuthContent<A, E, R>(self: Effect.Effect<A, E, R>, value: Record<string, unknown> = defaultAuthContent()) {
|
||||
|
||||
@@ -1407,11 +1407,8 @@ export type ProviderConfig = {
|
||||
* Timeout in milliseconds to wait for response headers. Provider integrations may set defaults. Set to false to disable timeout.
|
||||
*/
|
||||
headerTimeout?: number | false
|
||||
/**
|
||||
* Timeout in milliseconds between streamed SSE chunks for this provider. If no chunk arrives within this window, the request is aborted. Set to false to disable the idle watchdog.
|
||||
*/
|
||||
chunkTimeout?: number | false
|
||||
[key: string]: unknown | string | boolean | number | false | number | false | number | false | undefined
|
||||
chunkTimeout?: number
|
||||
[key: string]: unknown | string | boolean | number | false | number | false | number | undefined
|
||||
}
|
||||
models?: {
|
||||
[key: string]: {
|
||||
|
||||
@@ -27837,17 +27837,8 @@
|
||||
"description": "Timeout in milliseconds to wait for response headers. Provider integrations may set defaults. Set to false to disable timeout."
|
||||
},
|
||||
"chunkTimeout": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "integer",
|
||||
"exclusiveMinimum": 0
|
||||
},
|
||||
{
|
||||
"type": "boolean",
|
||||
"enum": [false]
|
||||
}
|
||||
],
|
||||
"description": "Timeout in milliseconds between streamed SSE chunks for this provider. If no chunk arrives within this window, the request is aborted. Set to false to disable the idle watchdog."
|
||||
"type": "integer",
|
||||
"exclusiveMinimum": 0
|
||||
}
|
||||
},
|
||||
"additionalProperties": {}
|
||||
|
||||
Reference in New Issue
Block a user