From ca4b1c31af774da04e954b2e3a83be4e66b3cb75 Mon Sep 17 00:00:00 2001 From: Valera Brizhatiuk <19537764+plombeer31@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:07:20 +0300 Subject: [PATCH 1/3] fix(llm): a streaming llama response is bounded by idle time, not total time MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `localModels.requestTimeoutMs` (default 300s) was applied as a wall-clock deadline over an entire streaming generation. `createRequestController` armed one `setTimeout` when the request was sent and nothing refreshed it, so `completeStream` aborted its own healthy stream 300s in — a reasoning model on CPU, or a llama-server across a LAN, is killed mid-token with every byte already produced discarded. Nothing downstream recovers it: `isRetryableLlamaError` refuses to replay a `timedOut` error, and `timedOutOf` in `should-advance` deliberately keeps a self-inflicted timeout off the immediate-fallover path. The turn just dies. The cloud path never behaved this way. `openAiFetch` clears its identical timer in `finally` when the fetch promise settles, i.e. at response headers, so for OpenAI-compatible providers the same knob bounds only connect. Local and cloud read one config key two incompatible ways. `createRequestController` now returns `keepAlive()`, which re-arms the deadline and marks it an idle budget; `completeStream` calls it once at headers and again on every chunk. `complete()` is untouched — a unary request has exactly one event to wait for and no idle signal to refresh against, so a total budget is the only one it can have. `timedOut()` now reports which deadline fired so the two carry honest, opposite advice: "lower completionMaxTokens" is meaningless for a stall where nothing arrived at all. An idle stall still sets `LlamaServerError.timedOut`, leaving retry and fallover exactly where they were — llama.cpp sends headers before it evaluates the prompt, so minutes of silence during a long CPU prompt-eval is not evidence the provider is dead. --- src/llm/llama-server-client.test.ts | 240 +++++++++++++++++++++++++++- src/llm/llama-server-client.ts | 119 ++++++++++++-- 2 files changed, 341 insertions(+), 18 deletions(-) diff --git a/src/llm/llama-server-client.test.ts b/src/llm/llama-server-client.test.ts index 0bfcc8d8..a1f1fe10 100644 --- a/src/llm/llama-server-client.test.ts +++ b/src/llm/llama-server-client.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect } from "vitest"; +import { afterEach, describe, it, expect, vi } from "vitest"; import { LlamaServerClient, LlamaServerError, @@ -474,6 +474,244 @@ describe("LlamaServerClient.completeStream", () => { }); }); +/** + * The streaming deadline is an *idle* deadline: `requestTimeoutMs` bounds + * how long the server may stay silent, not how long the answer may be. + * It used to bound the whole generation, so a healthy reasoning model on + * CPU — or any llama-server on the far side of a LAN — was killed at + * exactly the budget with every token already produced thrown away, and + * neither the retry policy (`timedOut` is not retryable) nor the fallback + * chain (a self-inflicted timeout is not an immediate signal) recovered + * it. The cloud path never had this problem: `openAiFetch` clears its + * timer as soon as the fetch promise settles, i.e. at response headers. + */ +describe("LlamaServerClient.completeStream deadlines", () => { + interface PushableStream { + response: Response; + push: (text: string) => void; + close: () => void; + } + + /** + * An SSE body the test drives by hand. Aborting the request signal + * errors the body mid-read, which is what undici does when the + * controller fires while the response is still streaming — the + * behaviour the production bug depends on. + */ + function pushableSse(signal: AbortSignal | null | undefined): PushableStream { + const encoder = new TextEncoder(); + let ctrl!: ReadableStreamDefaultController; + const stream = new ReadableStream({ + start(c) { + ctrl = c; + }, + }); + let finished = false; + signal?.addEventListener("abort", () => { + if (finished) return; + finished = true; + ctrl.error( + Object.assign(new Error("The operation was aborted"), { + name: "AbortError", + }), + ); + }); + return { + response: new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }), + push: (text: string) => { + if (!finished) ctrl.enqueue(encoder.encode(text)); + }, + close: () => { + if (finished) return; + finished = true; + ctrl.close(); + }, + }; + } + + function streamingClient(requestTimeoutMs: number): { + client: LlamaServerClient; + opened: () => PushableStream; + } { + let handle: PushableStream | null = null; + const client = new LlamaServerClient({ + baseUrl: "http://127.0.0.1:9999", + requestTimeoutMs, + fetchImpl: createMockFetch(async (_url, init) => { + handle = pushableSse(init.signal); + return handle.response; + }), + completionRetries: 1, + completionRetryBackoffMs: 0, + sleep: async () => {}, + }); + return { + client, + opened: () => { + if (!handle) throw new Error("stream not opened yet"); + return handle; + }, + }; + } + + afterEach(() => { + vi.useRealTimers(); + }); + + it("keeps streaming past requestTimeoutMs while chunks keep arriving", async () => { + // The regression test. Six chunks 999ms apart is 5,994ms of healthy + // generation under a 1,000ms budget — six times over the old + // wall-clock cap, and every one of those gaps is under it. + vi.useFakeTimers(); + const { client, opened } = streamingClient(1_000); + const iterator = client.completeStream({ prompt: "hi" }); + const deltas: string[] = []; + let final: { content: string } | null = null; + const consumed = (async () => { + while (true) { + const next = await iterator.next(); + if (next.done) { + final = next.value; + return; + } + if (next.value.delta) deltas.push(next.value.delta); + } + })(); + + // Let the generator open the request and park on its first read(). + await vi.advanceTimersByTimeAsync(0); + for (let i = 0; i < 6; i += 1) { + opened().push(`data: {"content":"t${i}","stop":false}\n\n`); + await vi.advanceTimersByTimeAsync(999); + } + opened().push('data: {"content":"","stop":true}\n\n'); + await vi.advanceTimersByTimeAsync(0); + opened().close(); + await vi.advanceTimersByTimeAsync(0); + await consumed; + + expect(deltas.join("")).toBe("t0t1t2t3t4t5"); + expect(final).not.toBeNull(); + expect(final!.content).toBe("t0t1t2t3t4t5"); + }); + + it("aborts a stream that goes silent for longer than requestTimeoutMs", async () => { + vi.useFakeTimers(); + const { client, opened } = streamingClient(1_000); + const iterator = client.completeStream({ prompt: "hi" }); + const deltas: string[] = []; + const failure = (async (): Promise => { + try { + while (true) { + const next = await iterator.next(); + if (next.done) return null; + if (next.value.delta) deltas.push(next.value.delta); + } + } catch (err) { + return err; + } + })(); + + await vi.advanceTimersByTimeAsync(0); + opened().push('data: {"content":"partial","stop":false}\n\n'); + await vi.advanceTimersByTimeAsync(500); + // …and then the server goes quiet for a full budget. + await vi.advanceTimersByTimeAsync(1_001); + const err = await failure; + + expect(deltas.join("")).toBe("partial"); + expect(err).toBeInstanceOf(LlamaServerError); + const llamaErr = err as LlamaServerError; + expect(llamaErr.status).toBeNull(); + // Still `timedOut` — see the field's doc comment. llama-server sends + // headers before it evaluates the prompt, so silence is not proof the + // provider is dead, and flipping this would turn a slow local model + // into an immediate fallover. + expect(llamaErr.timedOut).toBe(true); + expect(llamaErr.message).toContain("sent no data for 1000ms"); + // The old advice is wrong for a stall: nothing was too long. + expect(llamaErr.message).not.toContain("lower completionMaxTokens"); + }); + + it("still enforces a total deadline on the unary complete() path", async () => { + // Pinned deliberately. A non-streaming request has exactly one event + // to wait for, so it has no idle signal to refresh against — the + // wall-clock budget is all it can have. + vi.useFakeTimers(); + const client = new LlamaServerClient({ + baseUrl: "http://127.0.0.1:9999", + requestTimeoutMs: 1_000, + fetchImpl: createMockFetch( + (_url, init) => + new Promise((_resolve, reject) => { + init.signal?.addEventListener("abort", () => { + reject( + Object.assign(new Error("aborted"), { name: "AbortError" }), + ); + }); + }), + ), + completionRetries: 1, + completionRetryBackoffMs: 0, + sleep: async () => {}, + }); + const failure = client.complete({ prompt: "hi" }).then( + () => null, + (err: unknown) => err, + ); + await vi.advanceTimersByTimeAsync(1_001); + const err = (await failure) as LlamaServerError; + + expect(err).toBeInstanceOf(LlamaServerError); + expect(err.timedOut).toBe(true); + expect(err.message).toContain("exceeded requestTimeoutMs (1000ms)"); + }); + + it("lets an external abort cancel mid-stream without reporting a timeout", async () => { + // Esc in the TUI. The abort must not be laundered into our own + // idle-timeout error: `timedOut` stays false, so the fallback chain + // and `toLlmFailure` (which reads `ctx.signal.aborted`) still see a + // cancellation rather than a provider failure. + vi.useFakeTimers(); + const { client, opened } = streamingClient(60_000); + const abort = new AbortController(); + const iterator = client.completeStream({ + prompt: "hi", + signal: abort.signal, + }); + const deltas: string[] = []; + const failure = (async (): Promise => { + try { + while (true) { + const next = await iterator.next(); + if (next.done) return null; + if (next.value.delta) deltas.push(next.value.delta); + } + } catch (err) { + return err; + } + })(); + + await vi.advanceTimersByTimeAsync(0); + opened().push('data: {"content":"half","stop":false}\n\n'); + await vi.advanceTimersByTimeAsync(10); + abort.abort(); + await vi.advanceTimersByTimeAsync(0); + const err = await failure; + + expect(deltas.join("")).toBe("half"); + expect(err).toBeInstanceOf(LlamaServerError); + const llamaErr = err as LlamaServerError; + expect(llamaErr.timedOut).toBe(false); + expect(llamaErr.message).toMatch(/abort/i); + expect(llamaErr.message).not.toContain("requestTimeoutMs"); + expect(llamaErr.message).not.toContain("sent no data"); + }); +}); + describe("extractLlamaErrorDetail", () => { it("pulls the message from { error: { message } }", () => { expect( diff --git a/src/llm/llama-server-client.ts b/src/llm/llama-server-client.ts index 3abe229a..34460e03 100644 --- a/src/llm/llama-server-client.ts +++ b/src/llm/llama-server-client.ts @@ -39,6 +39,17 @@ const ENV_TOP_P = parseFloatEnv(process.env.ATOMIC_AGENT_LLAMA_TOP_P); const ENV_TOP_K = parseIntEnv(process.env.ATOMIC_AGENT_LLAMA_TOP_K); const ENV_SEED = parseIntEnv(process.env.ATOMIC_AGENT_LLAMA_SEED); +/** + * Which of our own deadlines fired. + * + * - `total` — the whole request was given `requestTimeoutMs` and never + * produced a response. The only signal a unary request has. + * - `idle` — a *stream* went `requestTimeoutMs` without sending a byte. + * A healthy generation refreshes this budget on every chunk, so it + * means the server went quiet, not that the answer was long. + */ +export type LlamaTimeoutKind = "total" | "idle"; + export class LlamaServerError extends Error { constructor( message: string, @@ -46,11 +57,21 @@ export class LlamaServerError extends Error { public readonly url: string, /** * True when *our own* `requestTimeoutMs` controller fired rather than - * the transport failing. Both surface as `status === null`, but a + * the transport failing — for either deadline, `total` or `idle` + * (see `LlamaTimeoutKind`). Both surface as `status === null`, but a * timeout is a "the model is slower than the budget" signal, not a * transient blip — replaying it just burns another full timeout of * GPU time (3 attempts x 300s = 15 silent minutes). See * `isRetryableLlamaError`. + * + * An idle stall stays flagged here on purpose. It is tempting to read + * "the server went quiet" as harder evidence of a dead provider than + * "the server is slow", and so let `shouldAdvance` fall over on the + * first occurrence — but llama.cpp streams response headers before it + * evaluates the prompt, so a long CPU prompt-eval is genuinely silent + * for minutes while nothing is wrong. Keeping the flag leaves the + * fallover behaviour exactly where it was: advance on the + * consecutive-failure threshold, never immediately. */ public readonly timedOut = false, /** @@ -254,13 +275,14 @@ export class LlamaServerClient { response: Response; controller: AbortController; cleanup: () => void; - timedOut: () => boolean; + timedOut: () => LlamaTimeoutKind | null; + keepAlive: () => void; }; try { opened = await this.runWithRetry( url, async () => { - const { controller, cleanup, timedOut } = + const { controller, cleanup, timedOut, keepAlive } = this.createRequestController(request.signal); try { const response = await this.fetchImpl(url, { @@ -272,7 +294,7 @@ export class LlamaServerClient { if (!response.ok || !response.body) { throw await buildHttpError(response, url); } - return { response, controller, cleanup, timedOut }; + return { response, controller, cleanup, timedOut, keepAlive }; } catch (err) { cleanup(); throw this.wrapTransportError(err, url, timedOut()); @@ -287,7 +309,7 @@ export class LlamaServerClient { cause: err, }); } - const { response, cleanup, timedOut } = opened; + const { response, cleanup, timedOut, keepAlive } = opened; let finalResult: CompletionResult = { content: "", reasoningContent: "", @@ -314,12 +336,23 @@ export class LlamaServerClient { const reader = response.body .pipeThrough(new TextDecoderStream()) .getReader(); + // Headers are in; from here the deadline bounds *silence*, not the + // length of the answer. Re-arming once here also hands the body a + // full budget rather than whatever the connect phase left over — + // llama.cpp answers with headers immediately and only then evaluates + // the prompt, so the first token can legitimately be minutes away. + keepAlive(); let buffer = ""; let accumulated = ""; let accumulatedReasoning = ""; while (true) { const { value, done } = await reader.read(); if (done) break; + // A byte arrived: the server is alive, so start the clock over. + // Deliberately not called on `done` — that breaks straight out of + // the loop into `finally { cleanup() }` with nothing awaited in + // between, so there is no window left for the timer to fire. + keepAlive(); buffer += value; let eventEnd = buffer.indexOf("\n\n"); while (eventEnd !== -1) { @@ -367,32 +400,69 @@ export class LlamaServerClient { * The returned `cleanup` clears the timeout and detaches the external * listener — call it in `finally` so a long-lived stream does not leak * the listener. + * + * The deadline starts as a **total** budget, which is all a unary + * request can be given: it has exactly one event to wait for. A + * streaming caller converts it into an **idle** budget by calling + * `keepAlive()` on every byte it receives — see `completeStream`. + * Without that, `requestTimeoutMs` was a wall-clock cap on the whole + * generation and killed healthy long answers at exactly the budget, + * discarding every token already produced. */ private createRequestController(externalSignal?: AbortSignal): { controller: AbortController; cleanup: () => void; - /** True once the per-request timeout (not the caller) fired the abort. */ - timedOut: () => boolean; + /** + * Which of our own deadlines fired the abort, or `null` when the + * abort came from the caller / nothing fired at all. + */ + timedOut: () => LlamaTimeoutKind | null; + /** + * Restart the deadline and mark it an idle budget. A no-op once the + * request is already aborted, so a late call cannot resurrect a + * controller the caller or the timer has finished with. + */ + keepAlive: () => void; } { const controller = new AbortController(); - let expired = false; - const timer = setTimeout(() => { - expired = true; - controller.abort(); - }, this.requestTimeoutMs); - const timedOut = (): boolean => expired; + let expired: LlamaTimeoutKind | null = null; + let kind: LlamaTimeoutKind = "total"; + const arm = (): ReturnType => + setTimeout(() => { + expired = kind; + controller.abort(); + }, this.requestTimeoutMs); + let timer = arm(); + const timedOut = (): LlamaTimeoutKind | null => expired; + const keepAlive = (): void => { + if (expired !== null || controller.signal.aborted) return; + clearTimeout(timer); + kind = "idle"; + timer = arm(); + }; if (!externalSignal) { - return { controller, cleanup: () => clearTimeout(timer), timedOut }; + return { + controller, + cleanup: () => clearTimeout(timer), + timedOut, + keepAlive, + }; } if (externalSignal.aborted) { controller.abort(); - return { controller, cleanup: () => clearTimeout(timer), timedOut }; + return { + controller, + cleanup: () => clearTimeout(timer), + timedOut, + keepAlive, + }; } const onAbort = (): void => controller.abort(); externalSignal.addEventListener("abort", onAbort, { once: true }); return { controller, timedOut, + keepAlive, cleanup: () => { clearTimeout(timer); externalSignal.removeEventListener("abort", onAbort); @@ -408,10 +478,25 @@ export class LlamaServerClient { private wrapTransportError( err: unknown, url: string, - timedOut: boolean, + timedOut: LlamaTimeoutKind | null, ): LlamaServerError { if (err instanceof LlamaServerError) return err; - if (timedOut) { + // An idle stall and a blown total budget need opposite advice. + // "Lower completionMaxTokens" is meaningless when the server sent + // nothing at all — the answer was not too long, it never came. + if (timedOut === "idle") { + return new LlamaServerError( + `llama-server sent no data for ${this.requestTimeoutMs}ms mid-stream — ` + + `the server stopped responding after starting the reply; check that ` + + `llama-server is still running, or raise localModels.requestTimeoutMs`, + null, + url, + true, + undefined, + { cause: err }, + ); + } + if (timedOut === "total") { return new LlamaServerError( `llama-server request exceeded requestTimeoutMs (${this.requestTimeoutMs}ms) — ` + `raise localModels.requestTimeoutMs or lower completionMaxTokens`, From 9d2bb40247243e1ed1b27cc055338c0c844ee9fa Mon Sep 17 00:00:00 2001 From: Valera Brizhatiuk <19537764+plombeer31@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:49:02 +0300 Subject: [PATCH 2/3] fix(llm): put an upper bound back on a single streaming response MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Making `requestTimeoutMs` an idle deadline removed the last cap on one local completion: a server emitting one byte every (budget - 1)ms refreshes the deadline forever, and nothing else on the turn path stops it — `src/agent` arms no timers, there is no `AbortSignal.timeout` anywhere on the path, and `ctx.signal` is user-driven only. A wedged or hostile llama-server could pin a slot, a session and, under headless `run`, the process indefinitely. Adds `localModels.streamTotalTimeoutMs` (env `ATOMIC_AGENT_LLAMA_STREAM_TOTAL_TIMEOUT_MS`, default 6 h) alongside the existing knob. It is a backstop, not a budget: 72x `REQUEST_TIMEOUT_MS`, and well clear of the worst honest local generation — the default `completionMaxTokens` of 8192 decoded at 0.4 tok/s is ~5.7 h. --- src/config/config-schema.ts | 31 +++++++++++++++++++++++++++++++ src/config/load-config.ts | 4 ++++ 2 files changed, 35 insertions(+) diff --git a/src/config/config-schema.ts b/src/config/config-schema.ts index bf541c99..1de6b5ee 100644 --- a/src/config/config-schema.ts +++ b/src/config/config-schema.ts @@ -175,7 +175,23 @@ export interface AtomicAgentConfig { /** Upper bound on `n_predict` for each completion when the caller omits `maxTokens`. */ completionMaxTokens: number; healthTimeoutMs: number; + /** + * For a unary `complete()`, the whole-request budget. For + * `completeStream()`, an **idle** budget: how long llama-server may + * stay silent between bytes. A healthy generation refreshes it on + * every chunk, so it never caps how long an answer may be — see + * `streamTotalTimeoutMs` for that. + */ requestTimeoutMs: number; + /** + * Absolute cap on one streaming response, measured from the moment + * response headers arrive. `requestTimeoutMs` only bounds silence, + * so without this a server dribbling one byte just under the idle + * budget would pin a slot, a session and — in headless `run` — a + * process forever. Deliberately far above any honest local + * generation; it is a backstop, not a budget. + */ + streamTotalTimeoutMs: number; healthRetries: number; healthRetryBackoffMs: number; /** @@ -2103,6 +2119,21 @@ export const ENV_DEFAULTS = { STATE_DIR: "~/.atomic-agent", HEALTH_TIMEOUT_MS: 3000, REQUEST_TIMEOUT_MS: 300_000, + /** + * 6 hours. The backstop on a single streaming response — see + * `AtomicAgentConfig.localModels.streamTotalTimeoutMs`. + * + * Chosen to clear the worst *honest* local generation by a wide + * margin: the default `completionMaxTokens` of 8 192 tokens decoded at + * 0.4 tok/s — slower than any CPU setup people actually sit through — + * is about 5.7 h. It is also 72x `REQUEST_TIMEOUT_MS`, so the idle + * deadline gets dozens of chances to fire first; if this one fires, + * the server was streaming continuously for six hours without + * finishing, which no local model this project targets does by + * accident. Raise it with `ATOMIC_AGENT_LLAMA_STREAM_TOTAL_TIMEOUT_MS` + * if you really do run a 131 072-token completion on a slow box. + */ + STREAM_TOTAL_TIMEOUT_MS: 6 * 60 * 60 * 1_000, HEALTH_RETRIES: 5, HEALTH_BACKOFF_MS: 500, COMPLETION_RETRIES: 3, diff --git a/src/config/load-config.ts b/src/config/load-config.ts index 43dec648..a602d16b 100644 --- a/src/config/load-config.ts +++ b/src/config/load-config.ts @@ -162,6 +162,10 @@ export function loadConfig(): AtomicAgentConfig { "ATOMIC_AGENT_LLAMA_REQUEST_TIMEOUT_MS", ENV_DEFAULTS.REQUEST_TIMEOUT_MS, ), + streamTotalTimeoutMs: readInt( + "ATOMIC_AGENT_LLAMA_STREAM_TOTAL_TIMEOUT_MS", + ENV_DEFAULTS.STREAM_TOTAL_TIMEOUT_MS, + ), healthRetries: readInt( "ATOMIC_AGENT_LLAMA_HEALTH_RETRIES", ENV_DEFAULTS.HEALTH_RETRIES, From f3f103b4fb4d9fb28b7e9a521274d1ad6d95b8a3 Mon Sep 17 00:00:00 2001 From: Valera Brizhatiuk <19537764+plombeer31@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:49:11 +0300 Subject: [PATCH 3/3] fix(llm): tell a pre-first-token stall apart from a mid-reply stall MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three follow-ups to the idle deadline, in one file: 1. Enforce `streamTotalTimeoutMs`. `createRequestController` now also returns `startStreamDeadline()`, a second timer armed once at response headers and never refreshed, reported as its own `stream-total` kind with its own wording ("data kept arriving, so this is the absolute cap … not a stall"). A live stream therefore holds two pending timers; `cleanup()` clears both. 2. A stall *before the first token* no longer claims the server "stopped responding after starting the reply". llama.cpp sends headers and only then evaluates the prompt, so that silence is the ordinary look of a long CPU prompt eval — exactly the population this change exists to protect. `keepAlive()` now takes the kind to record: `first-token` at headers, `idle` once a byte has actually arrived. The new wording says the request was accepted but no first token came, and points at `requestTimeoutMs` or the prompt/context size. 3. Tests for both, plus the two behaviours that were load-bearing but unpinned: the `keepAlive()` at headers (deleting it now fails two tests instead of none) and its no-op-once-aborted guard (a byte still in the decode pipe when the abort lands must not re-arm the timer and rewrite which deadline gets reported). --- src/llm/llama-server-client.test.ts | 151 +++++++++++++++++++++++++++- src/llm/llama-server-client.ts | 148 ++++++++++++++++++++++----- 2 files changed, 269 insertions(+), 30 deletions(-) diff --git a/src/llm/llama-server-client.test.ts b/src/llm/llama-server-client.test.ts index a1f1fe10..53bfcfb6 100644 --- a/src/llm/llama-server-client.test.ts +++ b/src/llm/llama-server-client.test.ts @@ -490,6 +490,8 @@ describe("LlamaServerClient.completeStream deadlines", () => { response: Response; push: (text: string) => void; close: () => void; + /** Error the body by hand — for streams that ignore the abort. */ + fail: () => void; } /** @@ -497,8 +499,15 @@ describe("LlamaServerClient.completeStream deadlines", () => { * errors the body mid-read, which is what undici does when the * controller fires while the response is still streaming — the * behaviour the production bug depends on. + * + * `errorOnAbort: false` models the narrow window in which the abort + * has landed but bytes already sitting in the decode pipe are still + * delivered; the test then errors the body itself with `fail()`. */ - function pushableSse(signal: AbortSignal | null | undefined): PushableStream { + function pushableSse( + signal: AbortSignal | null | undefined, + errorOnAbort = true, + ): PushableStream { const encoder = new TextEncoder(); let ctrl!: ReadableStreamDefaultController; const stream = new ReadableStream({ @@ -507,7 +516,7 @@ describe("LlamaServerClient.completeStream deadlines", () => { }, }); let finished = false; - signal?.addEventListener("abort", () => { + const fail = (): void => { if (finished) return; finished = true; ctrl.error( @@ -515,7 +524,8 @@ describe("LlamaServerClient.completeStream deadlines", () => { name: "AbortError", }), ); - }); + }; + if (errorOnAbort) signal?.addEventListener("abort", fail); return { response: new Response(stream, { status: 200, @@ -529,10 +539,14 @@ describe("LlamaServerClient.completeStream deadlines", () => { finished = true; ctrl.close(); }, + fail, }; } - function streamingClient(requestTimeoutMs: number): { + function streamingClient( + requestTimeoutMs: number, + options: { streamTotalTimeoutMs?: number; errorOnAbort?: boolean } = {}, + ): { client: LlamaServerClient; opened: () => PushableStream; } { @@ -540,8 +554,11 @@ describe("LlamaServerClient.completeStream deadlines", () => { const client = new LlamaServerClient({ baseUrl: "http://127.0.0.1:9999", requestTimeoutMs, + ...(options.streamTotalTimeoutMs === undefined + ? {} + : { streamTotalTimeoutMs: options.streamTotalTimeoutMs }), fetchImpl: createMockFetch(async (_url, init) => { - handle = pushableSse(init.signal); + handle = pushableSse(init.signal, options.errorOnAbort ?? true); return handle.response; }), completionRetries: 1, @@ -710,6 +727,130 @@ describe("LlamaServerClient.completeStream deadlines", () => { expect(llamaErr.message).not.toContain("requestTimeoutMs"); expect(llamaErr.message).not.toContain("sent no data"); }); + + it("reports a stall before the first token as a prompt eval, not a dead server", async () => { + // llama.cpp sends response headers and *then* evaluates the prompt, + // so this is the exact shape of the population this change exists to + // protect: a healthy server grinding a long context on CPU. Telling + // that user the server "stopped responding after starting the reply" + // would just be a different piece of wrong advice. + // + // This is also the test that covers the `keepAlive()` call at + // headers: delete it and the deadline is still the connect-phase + // `total` budget, so the error comes back with the unary wording. + vi.useFakeTimers(); + const { client } = streamingClient(1_000); + const iterator = client.completeStream({ prompt: "hi" }); + const failure = (async (): Promise => { + try { + while (true) { + const next = await iterator.next(); + if (next.done) return null; + } + } catch (err) { + return err; + } + })(); + + // Headers land, and then the body sends nothing at all. + await vi.advanceTimersByTimeAsync(0); + await vi.advanceTimersByTimeAsync(1_001); + const err = (await failure) as LlamaServerError; + + expect(err).toBeInstanceOf(LlamaServerError); + expect(err.status).toBeNull(); + expect(err.timedOut).toBe(true); + expect(err.message).toContain("sent no first token within 1000ms"); + expect(err.message).toContain("still be evaluating the prompt"); + // The two wordings this one must not be confused with. + expect(err.message).not.toContain("stopped responding"); + expect(err.message).not.toContain("after starting the reply"); + expect(err.message).not.toContain("exceeded requestTimeoutMs"); + }); + + it("caps one streaming response with streamTotalTimeoutMs even while chunks keep arriving", async () => { + // The idle deadline is not an upper bound: a server emitting one + // byte every (budget - 1)ms refreshes it forever. Without this cap a + // wedged or hostile llama-server pins a slot, a session and — under + // headless `run` — the process, with nothing else on the turn path + // to stop it (`ctx.signal` is user-driven only). + vi.useFakeTimers(); + const { client, opened } = streamingClient(1_000, { + streamTotalTimeoutMs: 5_000, + }); + const iterator = client.completeStream({ prompt: "hi" }); + const deltas: string[] = []; + const failure = (async (): Promise => { + try { + while (true) { + const next = await iterator.next(); + if (next.done) return null; + if (next.value.delta) deltas.push(next.value.delta); + } + } catch (err) { + return err; + } + })(); + + await vi.advanceTimersByTimeAsync(0); + // 900ms apart: every gap is inside the 1,000ms idle budget, so the + // idle deadline can never fire. Only the cap can. + for (let i = 0; i < 20; i += 1) { + opened().push('data: {"content":"t","stop":false}\n\n'); + await vi.advanceTimersByTimeAsync(900); + } + const err = (await failure) as LlamaServerError; + + expect(err).toBeInstanceOf(LlamaServerError); + expect(err.status).toBeNull(); + expect(err.timedOut).toBe(true); + // It streamed healthily right up to the cap. + expect(deltas.length).toBeGreaterThanOrEqual(5); + expect(err.message).toContain("streamTotalTimeoutMs (5000ms)"); + expect(err.message).toContain("ATOMIC_AGENT_LLAMA_STREAM_TOTAL_TIMEOUT_MS"); + // Not a stall, and the user must not be sent looking for one. + expect(err.message).not.toContain("sent no data for"); + expect(err.message).not.toContain("sent no first token"); + }); + + it("does not let a byte still in flight rewrite which deadline fired", async () => { + // `keepAlive()` is a no-op once a deadline has fired or the caller + // has aborted. The window is narrow but real: the abort lands while + // bytes already sitting in the decode pipe are still delivered, and + // the read loop calls `keepAlive()` on each of them. Without the + // guard those late bytes re-arm the timer, which fires a second time + // and overwrites the recorded reason — so the user is told the + // server stalled mid-reply when what actually happened is that it + // never produced a first token. + vi.useFakeTimers(); + const { client, opened } = streamingClient(1_000, { errorOnAbort: false }); + const iterator = client.completeStream({ prompt: "hi" }); + const failure = (async (): Promise => { + try { + while (true) { + const next = await iterator.next(); + if (next.done) return null; + } + } catch (err) { + return err; + } + })(); + + await vi.advanceTimersByTimeAsync(0); + // Silence past the budget: the first-token deadline fires and aborts. + await vi.advanceTimersByTimeAsync(1_001); + // …and only now does the byte that was already in flight land. + opened().push('data: {"content":"late","stop":false}\n\n'); + await vi.advanceTimersByTimeAsync(0); + // Long enough for a re-armed deadline to fire a second time. + await vi.advanceTimersByTimeAsync(2_000); + opened().fail(); + const err = (await failure) as LlamaServerError; + + expect(err).toBeInstanceOf(LlamaServerError); + expect(err.message).toContain("sent no first token within 1000ms"); + expect(err.message).not.toContain("sent no data for"); + }); }); describe("extractLlamaErrorDetail", () => { diff --git a/src/llm/llama-server-client.ts b/src/llm/llama-server-client.ts index 34460e03..d0eedcf5 100644 --- a/src/llm/llama-server-client.ts +++ b/src/llm/llama-server-client.ts @@ -44,11 +44,24 @@ const ENV_SEED = parseIntEnv(process.env.ATOMIC_AGENT_LLAMA_SEED); * * - `total` — the whole request was given `requestTimeoutMs` and never * produced a response. The only signal a unary request has. - * - `idle` — a *stream* went `requestTimeoutMs` without sending a byte. - * A healthy generation refreshes this budget on every chunk, so it - * means the server went quiet, not that the answer was long. + * - `first-token` — a *stream*'s headers arrived and then nothing did, + * for `requestTimeoutMs`. llama.cpp answers with headers immediately + * and only then evaluates the prompt, so this usually means the + * prompt eval is still running, **not** that the server is broken. + * - `idle` — a *stream* that had already sent at least one byte went + * `requestTimeoutMs` without sending another. A healthy generation + * refreshes this budget on every chunk, so it means the server went + * quiet mid-reply, not that the answer was long. + * - `stream-total` — a stream kept sending but never finished within + * `streamTotalTimeoutMs`. The backstop that keeps a wedged or + * hostile server from pinning a slot forever by dribbling one byte + * just under the idle budget. */ -export type LlamaTimeoutKind = "total" | "idle"; +export type LlamaTimeoutKind = + | "total" + | "first-token" + | "idle" + | "stream-total"; export class LlamaServerError extends Error { constructor( @@ -57,8 +70,8 @@ export class LlamaServerError extends Error { public readonly url: string, /** * True when *our own* `requestTimeoutMs` controller fired rather than - * the transport failing — for either deadline, `total` or `idle` - * (see `LlamaTimeoutKind`). Both surface as `status === null`, but a + * the transport failing — for any of our deadlines (see + * `LlamaTimeoutKind`). Both surface as `status === null`, but a * timeout is a "the model is slower than the budget" signal, not a * transient blip — replaying it just burns another full timeout of * GPU time (3 attempts x 300s = 15 silent minutes). See @@ -158,6 +171,12 @@ export interface LlamaServerClientOptions { baseUrl?: string; apiKey?: string | null; requestTimeoutMs?: number; + /** + * Overrides `config.localModels.streamTotalTimeoutMs`, the absolute + * cap on a single streaming response. Streaming only — a unary + * request is already bounded by `requestTimeoutMs`. + */ + streamTotalTimeoutMs?: number; fetchImpl?: typeof fetch; /** * Overrides the retry budget for `complete()` and the initial fetch @@ -187,6 +206,7 @@ export class LlamaServerClient { private readonly baseUrlOverride: string | undefined; private readonly apiKey: string | null; private readonly requestTimeoutMs: number; + private readonly streamTotalTimeoutMs: number; private readonly fetchImpl: typeof fetch; private readonly completionRetriesOverride: number | undefined; private readonly completionRetryBackoffMsOverride: number | undefined; @@ -198,6 +218,8 @@ export class LlamaServerClient { this.apiKey = options.apiKey ?? config.localModels.apiKey; this.requestTimeoutMs = options.requestTimeoutMs ?? config.localModels.requestTimeoutMs; + this.streamTotalTimeoutMs = + options.streamTotalTimeoutMs ?? config.localModels.streamTotalTimeoutMs; this.fetchImpl = options.fetchImpl ?? fetch; this.completionRetriesOverride = options.completionRetries; this.completionRetryBackoffMsOverride = options.completionRetryBackoffMs; @@ -276,13 +298,14 @@ export class LlamaServerClient { controller: AbortController; cleanup: () => void; timedOut: () => LlamaTimeoutKind | null; - keepAlive: () => void; + keepAlive: (next: Exclude) => void; + startStreamDeadline: () => void; }; try { opened = await this.runWithRetry( url, async () => { - const { controller, cleanup, timedOut, keepAlive } = + const { controller, cleanup, timedOut, keepAlive, startStreamDeadline } = this.createRequestController(request.signal); try { const response = await this.fetchImpl(url, { @@ -294,7 +317,14 @@ export class LlamaServerClient { if (!response.ok || !response.body) { throw await buildHttpError(response, url); } - return { response, controller, cleanup, timedOut, keepAlive }; + return { + response, + controller, + cleanup, + timedOut, + keepAlive, + startStreamDeadline, + }; } catch (err) { cleanup(); throw this.wrapTransportError(err, url, timedOut()); @@ -309,7 +339,8 @@ export class LlamaServerClient { cause: err, }); } - const { response, cleanup, timedOut, keepAlive } = opened; + const { response, cleanup, timedOut, keepAlive, startStreamDeadline } = + opened; let finalResult: CompletionResult = { content: "", reasoningContent: "", @@ -341,18 +372,27 @@ export class LlamaServerClient { // full budget rather than whatever the connect phase left over — // llama.cpp answers with headers immediately and only then evaluates // the prompt, so the first token can legitimately be minutes away. - keepAlive(); + // Until a byte actually arrives the deadline reports `first-token`: + // a silence *before* the reply starts is most likely a long prompt + // eval, and telling that user their server "stopped responding" is + // the same bad advice this change exists to remove. + keepAlive("first-token"); + // And an idle budget alone is not an upper bound — arm the absolute + // cap so a server dribbling one byte per (budget - 1)ms cannot pin + // this slot, session and process forever. + startStreamDeadline(); let buffer = ""; let accumulated = ""; let accumulatedReasoning = ""; while (true) { const { value, done } = await reader.read(); if (done) break; - // A byte arrived: the server is alive, so start the clock over. + // A byte arrived: the server is alive, so start the clock over — + // and from now on a silence really is a mid-reply stall. // Deliberately not called on `done` — that breaks straight out of // the loop into `finally { cleanup() }` with nothing awaited in // between, so there is no window left for the timer to fire. - keepAlive(); + keepAlive("idle"); buffer += value; let eventEnd = buffer.indexOf("\n\n"); while (eventEnd !== -1) { @@ -408,6 +448,12 @@ export class LlamaServerClient { * Without that, `requestTimeoutMs` was a wall-clock cap on the whole * generation and killed healthy long answers at exactly the budget, * discarding every token already produced. + * + * An idle budget alone is not an upper bound: a server emitting one + * byte just under it streams forever. `startStreamDeadline()` arms the + * second, never-refreshed timer that puts a ceiling back on — see + * `streamTotalTimeoutMs`. So a live stream holds two pending timers, + * and `cleanup` clears both. */ private createRequestController(externalSignal?: AbortSignal): { controller: AbortController; @@ -418,11 +464,20 @@ export class LlamaServerClient { */ timedOut: () => LlamaTimeoutKind | null; /** - * Restart the deadline and mark it an idle budget. A no-op once the - * request is already aborted, so a late call cannot resurrect a - * controller the caller or the timer has finished with. + * Restart the deadline and record what a subsequent expiry means: + * `first-token` once headers are in, `idle` once the body has + * actually produced something. A no-op once the request is already + * aborted or a deadline has already fired, so a byte that was still + * in the decode pipe when the abort landed cannot re-arm the timer + * or rewrite which deadline gets reported. */ - keepAlive: () => void; + keepAlive: (next: Exclude) => void; + /** + * Arm the absolute streaming cap. Idempotent, and a no-op once the + * request is aborted. Called once, at response headers, so the cap + * measures the body and not the connect phase. + */ + startStreamDeadline: () => void; } { const controller = new AbortController(); let expired: LlamaTimeoutKind | null = null; @@ -433,28 +488,43 @@ export class LlamaServerClient { controller.abort(); }, this.requestTimeoutMs); let timer = arm(); + let streamTimer: ReturnType | null = null; const timedOut = (): LlamaTimeoutKind | null => expired; - const keepAlive = (): void => { + const keepAlive = (next: Exclude): void => { if (expired !== null || controller.signal.aborted) return; clearTimeout(timer); - kind = "idle"; + kind = next; timer = arm(); }; + const startStreamDeadline = (): void => { + if (expired !== null || controller.signal.aborted) return; + if (streamTimer !== null) return; + streamTimer = setTimeout(() => { + expired = "stream-total"; + controller.abort(); + }, this.streamTotalTimeoutMs); + }; + const clearTimers = (): void => { + clearTimeout(timer); + if (streamTimer !== null) clearTimeout(streamTimer); + }; if (!externalSignal) { return { controller, - cleanup: () => clearTimeout(timer), + cleanup: clearTimers, timedOut, keepAlive, + startStreamDeadline, }; } if (externalSignal.aborted) { controller.abort(); return { controller, - cleanup: () => clearTimeout(timer), + cleanup: clearTimers, timedOut, keepAlive, + startStreamDeadline, }; } const onAbort = (): void => controller.abort(); @@ -463,8 +533,9 @@ export class LlamaServerClient { controller, timedOut, keepAlive, + startStreamDeadline, cleanup: () => { - clearTimeout(timer); + clearTimers(); externalSignal.removeEventListener("abort", onAbort); }, }; @@ -481,9 +552,36 @@ export class LlamaServerClient { timedOut: LlamaTimeoutKind | null, ): LlamaServerError { if (err instanceof LlamaServerError) return err; - // An idle stall and a blown total budget need opposite advice. - // "Lower completionMaxTokens" is meaningless when the server sent - // nothing at all — the answer was not too long, it never came. + // Each deadline needs its own advice. "Lower completionMaxTokens" + // is meaningless when the server sent nothing at all — the answer + // was not too long, it never came — and "the server stopped + // responding" is wrong when it never started, which for llama.cpp + // is the ordinary look of a long prompt eval. + if (timedOut === "first-token") { + return new LlamaServerError( + `llama-server accepted the request but sent no first token within ${this.requestTimeoutMs}ms — ` + + `it may still be evaluating the prompt; raise localModels.requestTimeoutMs, ` + + `or shorten the prompt/context if it is too large for this machine to evaluate in time`, + null, + url, + true, + undefined, + { cause: err }, + ); + } + if (timedOut === "stream-total") { + return new LlamaServerError( + `llama-server streamed for longer than streamTotalTimeoutMs (${this.streamTotalTimeoutMs}ms) without finishing — ` + + `data kept arriving, so this is the absolute cap on one streaming reply, not a stall; ` + + `raise localModels.streamTotalTimeoutMs (ATOMIC_AGENT_LLAMA_STREAM_TOTAL_TIMEOUT_MS) ` + + `or lower completionMaxTokens`, + null, + url, + true, + undefined, + { cause: err }, + ); + } if (timedOut === "idle") { return new LlamaServerError( `llama-server sent no data for ${this.requestTimeoutMs}ms mid-stream — ` +