diff --git a/assets/ai-models/qwen3.5-chat-template.jinja b/assets/ai-models/qwen3.5-chat-template.jinja deleted file mode 100644 index c2c2eef..0000000 --- a/assets/ai-models/qwen3.5-chat-template.jinja +++ /dev/null @@ -1,57 +0,0 @@ -{%- if messages[0]["role"] == "system" %} - {%- set system_message = messages[0]["content"] %} - {%- set loop_messages = messages[1:] %} -{%- else %} - {%- set loop_messages = messages %} -{%- endif %} - -{%- if not tools is defined %} - {%- set tools = [] %} -{%- endif %} - -{{- "<|im_start|>system\n" }} -{%- if system_message is defined %} - {{- system_message }} -{%- else %} - {{- "You are Qwen, a helpful assistant." }} -{%- endif %} -{%- if tools | length > 0 %} - {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within XML tags:\n" }} - {%- for tool in tools %} - {{- "\n" ~ tool | tojson }} - {%- endfor %} - {{- "\n\n\nFor each function call, return a json object with function name and arguments within XML tags:\n\n{\"name\": , \"arguments\": }\n" }} -{%- endif %} -{{- "<|im_end|>\n" }} -{%- for message in loop_messages %} - {%- if message.role == "assistant" and message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %} - {{- "<|im_start|>assistant" }} - {%- if message.content is defined and message.content is string and message.content | trim | length > 0 %} - {{- "\n" ~ message.content | trim }} - {%- endif %} - {%- for tool_call in message.tool_calls %} - {%- if tool_call.function is defined %} - {%- set tc = tool_call.function %} - {%- else %} - {%- set tc = tool_call %} - {%- endif %} - {{- "\n\n" ~ {"name": tc.name, "arguments": tc.arguments} | tojson ~ "\n" }} - {%- endfor %} - {{- "<|im_end|>\n" }} - {%- elif message.role == "tool" %} - {%- if loop.previtem is defined and loop.previtem.role != "tool" %} - {{- "<|im_start|>user\n" }} - {%- endif %} - {{- "\n" ~ message.content ~ "\n\n" }} - {%- if not loop.last and loop.nextitem is defined and loop.nextitem.role != "tool" %} - {{- "<|im_end|>\n" }} - {%- elif loop.last %} - {{- "<|im_end|>\n" }} - {%- endif %} - {%- else %} - {{- "<|im_start|>" ~ message.role ~ "\n" ~ message.content ~ "<|im_end|>\n" }} - {%- endif %} -{%- endfor %} -{%- if add_generation_prompt %} - {{- "<|im_start|>assistant\n" }} -{%- endif %} diff --git a/grammars/tool-call.gbnf b/grammars/tool-call.gbnf index 2f3869e..a3c60cf 100644 --- a/grammars/tool-call.gbnf +++ b/grammars/tool-call.gbnf @@ -43,4 +43,13 @@ exponent ::= ("e" | "E") ("+" | "-")? [0-9]+ boolean ::= "true" | "false" null-lit ::= "null" -ws ::= [ \t\n\r]* +# Bounded, deliberately. Whitespace is legal at ~10 seams inside the JSON +# body and GBNF masks EOG until the whole root is satisfied, so an +# unbounded `[ \t\n\r]*` gives a stuck sampler an infinite legal move: it +# emits newlines until `n_predict` while the stream parser (still waiting +# on `"tool":`) surfaces nothing to the UI — a silent multi-minute hang. +# 64 is far past any realistic newline+indent run (deep pretty-printed +# JSON tops out well under it); past that the mask simply forces the next +# real token, which is the desired recovery. Same grouped `{0,n}` shape +# already used for the reasoning-seam whitespace rule in build-grammar.ts. +ws ::= ( [ \t\n\r] ){0,64} diff --git a/src/llm/grammar/build-grammar.test.ts b/src/llm/grammar/build-grammar.test.ts index f3cbd6f..b45d545 100644 --- a/src/llm/grammar/build-grammar.test.ts +++ b/src/llm/grammar/build-grammar.test.ts @@ -15,6 +15,24 @@ describe("buildGrammar", () => { expect(grammar).not.toContain("think-prelude ::= think-body"); }); + // GBNF masks EOG until the whole root is satisfied, so any unbounded + // rule is an infinite legal move for a stuck sampler. `ws` appears at + // ~10 seams inside the JSON body; unbounded it produced a silent + // newline-until-n_predict hang that surfaced nothing to the UI. + it("bounds the shared whitespace rule so it cannot loop forever", async () => { + for (const profile of [ + PLAIN_INSTRUCT_PROFILE, + QWEN_THINK_PROFILE, + GEMMA4_THINK_PROFILE, + ]) { + const grammar = await buildGrammar(profile); + const ws = grammar.match(/^ws ::= .*$/m); + expect(ws, profile.id).not.toBeNull(); + expect(ws![0], profile.id).toMatch(/\{0,\d+\}\s*$/); + expect(ws![0], profile.id).not.toMatch(/\*\s*$/); + } + }); + it("builds a qwen think grammar with a think prelude routed into the array", async () => { const grammar = await buildGrammar(QWEN_THINK_PROFILE); expect(grammar).toContain("root ::= think-prelude tool-call-array"); diff --git a/src/llm/index.ts b/src/llm/index.ts index eda9128..e2349a8 100644 --- a/src/llm/index.ts +++ b/src/llm/index.ts @@ -9,10 +9,11 @@ export type { } from "./llama-server-client.js"; export { checkLlamaServer } from "./llama-server-health.js"; export type { HealthCheckOptions, HealthResult } from "./llama-server-health.js"; -export { SlotManager, hashPrefix } from "./slot-manager.js"; +export { SlotManager, hashPrefix, DEFAULT_SLOT_COUNT } from "./slot-manager.js"; export type { SlotAssignment } from "./slot-manager.js"; export { detectModelProfile, + extractTotalSlots, GEMMA4_THINK_PROFILE, PLAIN_INSTRUCT_PROFILE, QWEN_THINK_PROFILE, diff --git a/src/llm/llama-server-client.test.ts b/src/llm/llama-server-client.test.ts index 174a491..192b210 100644 --- a/src/llm/llama-server-client.test.ts +++ b/src/llm/llama-server-client.test.ts @@ -159,6 +159,55 @@ describe("LlamaServerClient.complete", () => { expect(result.content).toBe("ok"); }); + // A slow model that blows `requestTimeoutMs` surfaces as an abort with + // `status === null`, structurally identical to a dropped socket. Retrying + // it burns another full timeout of GPU time per attempt and cannot + // succeed, so it must short-circuit like a 4xx does. + it("does not retry when its own requestTimeoutMs fires", async () => { + let calls = 0; + const client = new LlamaServerClient({ + baseUrl: "http://127.0.0.1:9999", + requestTimeoutMs: 5, + fetchImpl: createMockFetch( + (_url, init) => + new Promise((_resolve, reject) => { + calls += 1; + init.signal?.addEventListener("abort", () => { + reject(Object.assign(new Error("aborted"), { name: "AbortError" })); + }); + }), + ), + completionRetries: 5, + completionRetryBackoffMs: 0, + sleep: async () => {}, + }); + await expect(client.complete({ prompt: "hi" })).rejects.toMatchObject({ + name: "LlamaServerError", + status: null, + timedOut: true, + }); + expect(calls).toBe(1); + }); + + it("still retries genuine transport failures", async () => { + let calls = 0; + const client = new LlamaServerClient({ + baseUrl: "http://127.0.0.1:9999", + requestTimeoutMs: 60_000, + fetchImpl: createMockFetch(async () => { + calls += 1; + throw new Error("ECONNRESET"); + }), + completionRetries: 3, + completionRetryBackoffMs: 0, + sleep: async () => {}, + }); + await expect(client.complete({ prompt: "hi" })).rejects.toMatchObject({ + timedOut: false, + }); + expect(calls).toBe(3); + }); + it("does not retry on 4xx grammar/validation errors", async () => { let calls = 0; const client = new LlamaServerClient({ diff --git a/src/llm/llama-server-client.ts b/src/llm/llama-server-client.ts index 5c1ea33..d2dfa01 100644 --- a/src/llm/llama-server-client.ts +++ b/src/llm/llama-server-client.ts @@ -42,6 +42,15 @@ export class LlamaServerError extends Error { message: string, public readonly status: number | null, public readonly url: string, + /** + * True when *our own* `requestTimeoutMs` controller fired rather than + * the transport failing. 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`. + */ + public readonly timedOut = false, ) { super(message); this.name = "LlamaServerError"; @@ -193,7 +202,7 @@ export class LlamaServerClient { return this.runWithRetry( url, async () => { - const { controller, cleanup } = this.createRequestController( + const { controller, cleanup, timedOut } = this.createRequestController( request.signal, ); try { @@ -209,9 +218,7 @@ export class LlamaServerClient { const json = (await response.json()) as Record; return normaliseCompletionResponse(json); } catch (err) { - if (err instanceof LlamaServerError) throw err; - const message = err instanceof Error ? err.message : String(err); - throw new LlamaServerError(message, null, url); + throw this.wrapTransportError(err, url, timedOut()); } finally { cleanup(); } @@ -231,14 +238,14 @@ export class LlamaServerClient { response: Response; controller: AbortController; cleanup: () => void; + timedOut: () => boolean; }; try { opened = await this.runWithRetry( url, async () => { - const { controller, cleanup } = this.createRequestController( - request.signal, - ); + const { controller, cleanup, timedOut } = + this.createRequestController(request.signal); try { const response = await this.fetchImpl(url, { method: "POST", @@ -249,12 +256,10 @@ export class LlamaServerClient { if (!response.ok || !response.body) { throw await buildHttpError(response, url); } - return { response, controller, cleanup }; + return { response, controller, cleanup, timedOut }; } catch (err) { cleanup(); - if (err instanceof LlamaServerError) throw err; - const message = err instanceof Error ? err.message : String(err); - throw new LlamaServerError(message, null, url); + throw this.wrapTransportError(err, url, timedOut()); } }, request.signal, @@ -264,7 +269,7 @@ export class LlamaServerClient { const message = err instanceof Error ? err.message : String(err); throw new LlamaServerError(message, null, url); } - const { response, cleanup } = opened; + const { response, cleanup, timedOut } = opened; let finalResult: CompletionResult = { content: "", reasoningContent: "", @@ -331,9 +336,7 @@ export class LlamaServerClient { } return finalResult; } catch (err) { - if (err instanceof LlamaServerError) throw err; - const message = err instanceof Error ? err.message : String(err); - throw new LlamaServerError(message, null, url); + throw this.wrapTransportError(err, url, timedOut()); } finally { cleanup(); } @@ -350,20 +353,28 @@ export class LlamaServerClient { private createRequestController(externalSignal?: AbortSignal): { controller: AbortController; cleanup: () => void; + /** True once the per-request timeout (not the caller) fired the abort. */ + timedOut: () => boolean; } { const controller = new AbortController(); - const timer = setTimeout(() => controller.abort(), this.requestTimeoutMs); + let expired = false; + const timer = setTimeout(() => { + expired = true; + controller.abort(); + }, this.requestTimeoutMs); + const timedOut = (): boolean => expired; if (!externalSignal) { - return { controller, cleanup: () => clearTimeout(timer) }; + return { controller, cleanup: () => clearTimeout(timer), timedOut }; } if (externalSignal.aborted) { controller.abort(); - return { controller, cleanup: () => clearTimeout(timer) }; + return { controller, cleanup: () => clearTimeout(timer), timedOut }; } const onAbort = (): void => controller.abort(); externalSignal.addEventListener("abort", onAbort, { once: true }); return { controller, + timedOut, cleanup: () => { clearTimeout(timer); externalSignal.removeEventListener("abort", onAbort); @@ -371,6 +382,30 @@ export class LlamaServerClient { }; } + /** + * Normalise a caught transport failure into a `LlamaServerError`, + * preserving whether it was our own request-timeout so the retry policy + * can refuse to replay it. Pass-through for errors already of that type. + */ + private wrapTransportError( + err: unknown, + url: string, + timedOut: boolean, + ): LlamaServerError { + if (err instanceof LlamaServerError) return err; + if (timedOut) { + return new LlamaServerError( + `llama-server request exceeded requestTimeoutMs (${this.requestTimeoutMs}ms) — ` + + `raise localModels.requestTimeoutMs or lower completionMaxTokens`, + null, + url, + true, + ); + } + const message = err instanceof Error ? err.message : String(err); + return new LlamaServerError(message, null, url); + } + private prepareRequest( request: CompletionRequest, stream: boolean, @@ -528,6 +563,10 @@ function toNumber(value: unknown, fallback = 0): number { */ function isRetryableLlamaError(err: unknown): boolean { if (!(err instanceof LlamaServerError)) return false; + // Our own request-timeout, not a transport failure. Replaying it costs + // another full `requestTimeoutMs` of GPU time and cannot succeed if the + // model simply needs longer than the budget. + if (err.timedOut) return false; if (err.status === null) return true; if (err.status >= 500 && err.status < 600) return true; return false; diff --git a/src/llm/model-profile-manager.test.ts b/src/llm/model-profile-manager.test.ts index 12f8c15..cb6c083 100644 --- a/src/llm/model-profile-manager.test.ts +++ b/src/llm/model-profile-manager.test.ts @@ -34,6 +34,56 @@ function makeLlamaStub( } as { client: Pick; calls: number }; } +describe("ModelProfileManager slot discovery", () => { + // Managed mode defers the boot health check, so bootstrap builds the + // SlotManager on a one-slot default. This refresh hook is the only path + // that ever widens it to the daemon's real `--parallel` count. + it("reports /props.total_slots on a successful refresh", async () => { + const stub = makeLlamaStub([{ ...QWEN3_PROPS, total_slots: 2 }]); + const seen: number[] = []; + const mgr = new ModelProfileManager({ + llama: stub.client as LlamaServerClient, + initialProfile: PLAIN_INSTRUCT_PROFILE, + initialGrammar: await buildGrammar(PLAIN_INSTRUCT_PROFILE), + initialModelId: null, + onTotalSlots: (n) => seen.push(n), + }); + await mgr.refresh(); + expect(seen).toEqual([2]); + }); + + it("stays silent when the server omits total_slots", async () => { + const stub = makeLlamaStub([QWEN3_PROPS]); + const seen: number[] = []; + const mgr = new ModelProfileManager({ + llama: stub.client as LlamaServerClient, + initialProfile: PLAIN_INSTRUCT_PROFILE, + initialGrammar: await buildGrammar(PLAIN_INSTRUCT_PROFILE), + initialModelId: null, + onTotalSlots: (n) => seen.push(n), + }); + await mgr.refresh(); + expect(seen).toEqual([]); + }); + + it("reports slots even when the probe body fails profile detection", async () => { + // Slot discovery must not be hostage to the grammar rebuild path — an + // oversized pool thrashes the KV cache on every turn regardless of + // whether the profile changed. + const stub = makeLlamaStub([{ total_slots: 4 }]); + const seen: number[] = []; + const mgr = new ModelProfileManager({ + llama: stub.client as LlamaServerClient, + initialProfile: PLAIN_INSTRUCT_PROFILE, + initialGrammar: await buildGrammar(PLAIN_INSTRUCT_PROFILE), + initialModelId: null, + onTotalSlots: (n) => seen.push(n), + }); + await mgr.refresh(); + expect(seen).toEqual([4]); + }); +}); + describe("ModelProfileManager", () => { it("reports the initial profile/grammar until a mismatch is observed", async () => { const grammar = await buildGrammar(GEMMA4_THINK_PROFILE); diff --git a/src/llm/model-profile-manager.ts b/src/llm/model-profile-manager.ts index dc2243d..78f1c00 100644 --- a/src/llm/model-profile-manager.ts +++ b/src/llm/model-profile-manager.ts @@ -3,6 +3,7 @@ import { buildGrammar } from "./grammar/build-grammar.js"; import type { LlamaServerClient } from "./llama-server-client.js"; import { detectModelProfile, + extractTotalSlots, type ModelProfile, } from "./model-profile.js"; import { checkProfileGrammarAligned } from "./profile-invariants.js"; @@ -26,6 +27,15 @@ export interface ModelProfileManagerOptions { * `browser.*` into the `tool-name` rule. Default `true`. */ browserEnabled?: boolean; + /** + * Invoked with `/props.total_slots` after every successful probe. + * Managed mode defers the boot health check (the daemon may not be + * running yet), so bootstrap builds the `SlotManager` on a conservative + * one-slot default; this is the hook that widens it to the server's + * real `--parallel` count on the first refresh. Omit to ignore slot + * discovery entirely (tests with a stubbed HTTP layer). + */ + onTotalSlots?: (totalSlots: number) => void; logger?: StructuredLogger; } @@ -65,6 +75,9 @@ export class ModelProfileManager { private readonly llama: LlamaServerClient; private readonly grammarsDir: string | undefined; private readonly browserEnabled: boolean; + private readonly onTotalSlots: + | ((totalSlots: number) => void) + | undefined; private readonly logger: StructuredLogger | undefined; constructor(options: ModelProfileManagerOptions) { @@ -74,6 +87,7 @@ export class ModelProfileManager { this.llama = options.llama; this.grammarsDir = options.grammarsDir; this.browserEnabled = options.browserEnabled ?? true; + this.onTotalSlots = options.onTotalSlots; this.logger = options.logger; } @@ -131,6 +145,14 @@ export class ModelProfileManager { async refresh(): Promise { try { const props = await this.llama.fetchProps(); + // Slot discovery rides this probe. Done before the profile work so a + // grammar rebuild failure cannot swallow it — an oversized slot pool + // silently thrashes the server's KV cache and is worth correcting + // even on a turn where nothing else changed. + const totalSlots = extractTotalSlots(props); + if (totalSlots !== null) { + this.onTotalSlots?.(totalSlots); + } const nextProfile = detectModelProfile(props); const nextModelId = normaliseId( typeof props.model_alias === "string" ? props.model_alias : null, diff --git a/src/llm/model-profile.test.ts b/src/llm/model-profile.test.ts index 37e692d..172306d 100644 --- a/src/llm/model-profile.test.ts +++ b/src/llm/model-profile.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest"; import { detectModelProfile, + extractTotalSlots, detectVisionSupport, GEMMA4_THINK_PROFILE, PLAIN_INSTRUCT_PROFILE, @@ -14,6 +15,27 @@ import { QWEN3_PROPS, } from "./model-profile.fixtures.js"; +describe("extractTotalSlots", () => { + it("reads a positive integer slot count", () => { + expect(extractTotalSlots({ total_slots: 2 })).toBe(2); + }); + + // Falling back to `null` (and thus the conservative SlotManager default) + // is deliberate: llama.cpp wraps an out-of-range id into another + // session's slot rather than erroring, so guessing high corrupts cache + // affinity silently. + it("collapses missing, non-numeric, and sub-1 values to null", () => { + for (const raw of [undefined, null, "2", 0, -1, Number.NaN, Infinity]) { + expect(extractTotalSlots({ total_slots: raw })).toBeNull(); + } + expect(extractTotalSlots({})).toBeNull(); + }); + + it("truncates a fractional count", () => { + expect(extractTotalSlots({ total_slots: 2.9 })).toBe(2); + }); +}); + describe("detectModelProfile", () => { it("detects qwen think profile from props", () => { expect(detectModelProfile(QWEN3_PROPS)).toEqual(QWEN_THINK_PROFILE); diff --git a/src/llm/model-profile.ts b/src/llm/model-profile.ts index e5a0e7d..1829f6a 100644 --- a/src/llm/model-profile.ts +++ b/src/llm/model-profile.ts @@ -237,6 +237,24 @@ function selectBaseProfile( * also fall back to a root-level `n_ctx` for older builds and custom * forks. Returns `null` when neither is a positive number. */ +/** + * Extract `total_slots` from a `/props` payload — the number of parallel + * slots llama-server was started with (`--parallel`). Anything missing, + * non-finite, or below 1 collapses to `null` so callers fall back to the + * conservative `SlotManager` default rather than guessing high: an id + * above the real count is wrapped by llama.cpp (`id_slot % n_slots`) into + * a *different session's* slot, silently thrashing the KV cache. + */ +export function extractTotalSlots( + props: Record, +): number | null { + const raw = props.total_slots; + if (typeof raw !== "number" || !Number.isFinite(raw) || raw < 1) { + return null; + } + return Math.trunc(raw); +} + function readContextWindow(props: Record): number | null { const defaults = readObject(props.default_generation_settings); const nested = toPositiveInt(defaults.n_ctx); diff --git a/src/llm/slot-manager.test.ts b/src/llm/slot-manager.test.ts index 2d1411d..40cdb97 100644 --- a/src/llm/slot-manager.test.ts +++ b/src/llm/slot-manager.test.ts @@ -1,7 +1,92 @@ import { describe, it, expect } from "vitest"; -import { SlotManager, hashPrefix } from "./slot-manager.js"; +import { SlotManager, hashPrefix, DEFAULT_SLOT_COUNT } from "./slot-manager.js"; describe("SlotManager", () => { + // Guessing high is not free: llama.cpp wraps an out-of-range `id_slot` + // (`id_slot % n_slots`) into another session's slot instead of erroring, + // so an oversized pool silently evicts KV cache. Slot 0 is the only id + // every llama-server is guaranteed to have. + it("defaults to a single slot so no id can exceed the server's count", () => { + expect(DEFAULT_SLOT_COUNT).toBe(1); + const mgr = new SlotManager(); + expect(mgr.getSlotCount()).toBe(1); + for (let i = 0; i < 10; i++) { + expect(mgr.acquire(`sess-${i}`, "prefix").slotId).toBe(0); + } + }); + + describe("resize", () => { + it("widens the pool to the discovered slot count", () => { + const mgr = new SlotManager(); + mgr.resize(2); + expect(mgr.getSlotCount()).toBe(2); + const ids = ["a", "b", "c"].map((s) => mgr.acquire(s, "prefix").slotId); + expect(ids).toEqual([0, 1, 0]); + }); + + it("is a no-op when the count is unchanged, preserving affinity", () => { + const mgr = new SlotManager(2); + const before = mgr.acquire("sess", "prefix"); + mgr.resize(2); + const after = mgr.acquire("sess", "prefix"); + expect(after.slotId).toBe(before.slotId); + expect(after.cacheReused).toBe(true); + }); + + it("drops stale assignments whose slot may no longer exist", () => { + const mgr = new SlotManager(4); + const before = mgr.acquire("sess", "prefix"); + mgr.resize(2); + const after = mgr.acquire("sess", "prefix"); + expect(after.cacheReused).toBe(false); + expect(after.slotId).toBeLessThan(2); + void before; + }); + + it("never hands out an id at or above the new slot count", () => { + const mgr = new SlotManager(8); + mgr.reserveReflectionSlot(); + mgr.resize(2); + for (let i = 0; i < 20; i++) { + expect(mgr.acquire(`sess-${i}`, "prefix").slotId).toBeLessThan(2); + } + }); + + it("releases a reflection reservation that fell out of range", () => { + const mgr = new SlotManager(4); + expect(mgr.reserveReflectionSlot()).toBe(3); + mgr.resize(2); + // 3 no longer exists, so the reservation must be re-taken in range. + const reserved = mgr.reserveReflectionSlot(); + expect(reserved).not.toBeNull(); + expect(reserved).toBeLessThan(2); + }); + + it("keeps an in-range reflection reservation off the acquire pool", () => { + const mgr = new SlotManager(4); + const reserved = mgr.reserveReflectionSlot(); + expect(reserved).toBe(3); + // Widening keeps slot 3 valid, so the reservation must survive. + mgr.resize(5); + expect(mgr.reserveReflectionSlot()).toBe(reserved); + for (let i = 0; i < 10; i++) { + expect(mgr.acquire(`sess-${i}`, "prefix").slotId).not.toBe(reserved); + } + }); + + it("gives the sole slot back to the agent when shrinking to one", () => { + const mgr = new SlotManager(4); + mgr.reserveReflectionSlot(); + mgr.resize(1); + expect(mgr.acquire("sess", "prefix").slotId).toBe(0); + expect(mgr.reserveReflectionSlot()).toBeNull(); + }); + + it("rejects a non-positive slot count", () => { + expect(() => new SlotManager(2).resize(0)).toThrow(/must be positive/); + }); + }); + it("returns the same slot for an unchanged prefix within a session", () => { const mgr = new SlotManager(4); const a = mgr.acquire("sess-1", "stable prefix v1"); diff --git a/src/llm/slot-manager.ts b/src/llm/slot-manager.ts index 7ba97fc..302c67a 100644 --- a/src/llm/slot-manager.ts +++ b/src/llm/slot-manager.ts @@ -33,14 +33,27 @@ export interface SlotAssignment { * controller-contract violation and would race the prefix swap. Do * not call `acquire` outside an `AgentLoop.runTurn` frame. */ +/** + * Slot count used when the `/props` probe has not answered yet. Every + * llama-server has slot 0; anything above that is a guess. The previous + * default of 4 was one such guess, and it outlived the probe in managed + * mode (which always defers the boot health check), so sessions were + * handed slot ids 1-3 against a `--parallel 2` daemon. llama.cpp wraps + * out-of-range ids (`id_slot % n_slots`), so nothing errored — the ids + * silently collided instead, evicting the KV cache and forcing a full + * prompt reprocess on every rotation. One slot is always correct; + * `resize()` widens the pool once the real count is known. + */ +export const DEFAULT_SLOT_COUNT = 1; + export class SlotManager { private readonly assignments = new Map(); - private readonly slotCount: number; + private slotCount: number; private slotPool: number[]; private nextRoundRobin = 0; private reservedReflectionSlot: number | null = null; - constructor(slotCount = 4) { + constructor(slotCount = DEFAULT_SLOT_COUNT) { if (slotCount <= 0) { throw new Error("slotCount must be positive"); } @@ -48,6 +61,51 @@ export class SlotManager { this.slotPool = Array.from({ length: slotCount }, (_, i) => i); } + /** Slot count the pool is currently sized for. */ + getSlotCount(): number { + return this.slotCount; + } + + /** + * Re-size the pool to the server's actual slot count, discovered from a + * later `/props` probe. No-op when the count is unchanged, so the common + * refresh path costs nothing and never disturbs live cache affinity. + * + * On a real change every existing assignment is dropped: a session's + * slot id may no longer exist, and the server-side KV for it is not + * where we think it is either way. The next `acquire` re-assigns with + * `cacheReused: false`, which is honest — one cold prefix rebuild beats + * silently pointing at a stranger's cache. + * + * A reflection reservation that is still in range is preserved; one that + * fell outside is released back so `reserveReflectionSlot()` can re-take + * a valid slot. Call this between turns — `acquire` has no locking and + * the single-active-turn-per-session invariant is what keeps it safe. + */ + resize(slotCount: number): void { + if (slotCount <= 0) { + throw new Error("slotCount must be positive"); + } + if (slotCount === this.slotCount) return; + const reserved = + this.reservedReflectionSlot !== null && + this.reservedReflectionSlot < slotCount + ? this.reservedReflectionSlot + : null; + this.slotCount = slotCount; + this.assignments.clear(); + this.nextRoundRobin = 0; + this.reservedReflectionSlot = reserved; + this.slotPool = Array.from({ length: slotCount }, (_, i) => i).filter( + (id) => id !== reserved, + ); + // Reserving the only slot would starve the agent loop; hand it back. + if (this.slotPool.length === 0) { + this.reservedReflectionSlot = null; + this.slotPool = [0]; + } + } + acquire(sessionId: string, stablePrefix: string): SlotAssignment { const prefixHash = hashPrefix(stablePrefix); const existing = this.assignments.get(sessionId); diff --git a/src/local-llm/chat-templates.test.ts b/src/local-llm/chat-templates.test.ts index 0b36715..219d5e3 100644 --- a/src/local-llm/chat-templates.test.ts +++ b/src/local-llm/chat-templates.test.ts @@ -1,16 +1,31 @@ import { describe, expect, it } from "vitest"; import { resolveChatTemplatePath } from "./chat-templates.js"; -import { getLocalModelDef } from "./models-catalog.js"; +import { getLocalModelDef, LOCAL_MODELS_CATALOG } from "./models-catalog.js"; describe("chat-templates", () => { - it("resolves qwen3.5 jinja when asset exists in repo", () => { - const path = resolveChatTemplatePath(getLocalModelDef("qwen-3.5-4b")); - expect(path).not.toBeNull(); - expect(path!.endsWith("qwen3.5-chat-template.jinja")).toBe(true); - }); - it("returns null when model has no template asset", () => { expect(resolveChatTemplatePath(getLocalModelDef("gemma-4-e4b"))).toBeNull(); }); + + // Regression guard for the managed-mode hang: `--chat-template-file` + // overrides what `llama-server` reports at `/props.chat_template`, which + // is the only signal `detectModelProfile` has. Shipping a template that + // omits the model's reasoning markers silently demotes a thinking model + // to `plain-instruct`, which strips the reasoning prelude out of the + // GBNF root — the sampler is then forced to open with `[` and stalls in + // the unbounded `ws` rule until `n_predict`. No catalog model should + // override its GGUF's own template unless the override preserves those + // markers. + it("ships no chat-template override for any catalog model", () => { + const overriding = LOCAL_MODELS_CATALOG.filter((m) => m.chatTemplateAsset); + expect(overriding.map((m) => m.id)).toEqual([]); + }); + + it("resolves every declared template asset to a file on disk", () => { + for (const model of LOCAL_MODELS_CATALOG) { + if (!model.chatTemplateAsset) continue; + expect(resolveChatTemplatePath(model), model.id).not.toBeNull(); + } + }); }); diff --git a/src/local-llm/context-size.test.ts b/src/local-llm/context-size.test.ts index 86d6dfd..bf5253d 100644 --- a/src/local-llm/context-size.test.ts +++ b/src/local-llm/context-size.test.ts @@ -8,6 +8,27 @@ import { resolveDeviceFreeVramMiB, } from "./context-size.js"; import type { GpuDevice } from "./gpu-devices.js"; +import { minUsableContextWindow } from "../prompt/token-budget.js"; +import { USER_CONFIG_DEFAULTS } from "../config/config-schema.js"; + +// Drift guard. The floor exists so a model that does not fit VRAM still +// gets a *workable* context rather than a merely non-zero one. When it sat +// at 8192 it was equal to `completionMaxTokens` and smaller than the fixed +// prompt plus that budget, so every step on a floored model came back +// `truncated` and the agent never emitted a tool call. Either constant may +// move; they must not cross. +describe("MIN_AUTO_CONTEXT", () => { + it("clears the agent's fixed prompt plus a full generation budget", () => { + const required = minUsableContextWindow( + USER_CONFIG_DEFAULTS.localModels.completionMaxTokens, + ); + expect(MIN_AUTO_CONTEXT).toBeGreaterThanOrEqual(required); + }); + + it("is not itself capped away by the auto-size ceiling", () => { + expect(MIN_AUTO_CONTEXT).toBeLessThanOrEqual(MAX_AUTO_CONTEXT); + }); +}); const QWEN_9B = { modelSizeGb: 5.3, diff --git a/src/local-llm/context-size.ts b/src/local-llm/context-size.ts index 8d513cd..c471b95 100644 --- a/src/local-llm/context-size.ts +++ b/src/local-llm/context-size.ts @@ -32,14 +32,21 @@ const COMPUTE_OVERHEAD_MIB = 768; const KV_BYTES_PER_TOKEN_PER_GB = 32_000; /** - * Lower bound for the auto-sized context. The agent's stable prefix - * (persona + tool catalog + capabilities + instructions) alone is - * several thousand tokens; below this floor the agent cannot hold even a - * single turn, so we force it even when VRAM is tight. On small GPUs - * this makes llama.cpp's `-fit` step spill some layers to the CPU - * (slower) rather than reject every request with HTTP 400. + * Lower bound for the auto-sized context. On small GPUs this makes + * llama.cpp's `-fit` step spill some layers to the CPU (slower) rather + * than reject every request with HTTP 400. + * + * Derived, not arbitrary: the agent's stable prefix (persona + tool + * catalog + capabilities + instructions) measures ~5.2k tokens on its + * own, and `localModels.completionMaxTokens` defaults to 8192. The old + * 8192 floor was therefore *equal to* the generation budget and smaller + * than prefix + budget combined — a model that landed on it had ~2-3k + * tokens of room, could not finish a reasoning block plus a tool-call + * array, and hit llama.cpp's context ceiling with `truncated: true` on + * every step. Any model >= ~14 GB on a 16 GB card lands on this floor, + * so the value has to clear 5.2k + 8192 + margin. */ -export const MIN_AUTO_CONTEXT = 8192; +export const MIN_AUTO_CONTEXT = 16_384; /** * Upper bound for the auto-sized context. Caps KV growth on large GPUs diff --git a/src/local-llm/models-catalog.test.ts b/src/local-llm/models-catalog.test.ts index 4e6c8ce..8e5792f 100644 --- a/src/local-llm/models-catalog.test.ts +++ b/src/local-llm/models-catalog.test.ts @@ -21,10 +21,13 @@ describe("models-catalog", () => { expect(DEFAULT_LLAMACPP_MODEL_ID).toBe("qwen-3.5-4b"); }); - it("resolves qwen-3.5-4b chat template asset", () => { - expect(getLocalModelDef("qwen-3.5-4b").chatTemplateAsset).toBe( - "qwen3.5-chat-template.jinja", - ); + // The bundled Qwen 3.5 override was a Qwen2.5-style ChatML template with + // no `` / `enable_thinking` markers. Because `--chat-template-file` + // is what `/props.chat_template` reports back, it demoted the profile to + // `plain-instruct` and deadlocked the grammar. See chat-templates.test.ts. + it("does not override the Qwen 3.5 chat template", () => { + expect(getLocalModelDef("qwen-3.5-4b").chatTemplateAsset).toBeUndefined(); + expect(getLocalModelDef("qwen-3.5-35b").chatTemplateAsset).toBeUndefined(); }); it("throws on unknown id", () => { diff --git a/src/local-llm/models-catalog.ts b/src/local-llm/models-catalog.ts index 5055205..22d8e68 100644 --- a/src/local-llm/models-catalog.ts +++ b/src/local-llm/models-catalog.ts @@ -43,6 +43,19 @@ export interface LocalModelDef { minRamGb: number; recommendedRamGb: number; family: "qwen" | "gemma"; + /** + * Jinja file under `assets/ai-models/` passed to llama-server as + * `--chat-template-file`, overriding the template baked into the GGUF. + * + * **Invariant:** an override MUST preserve every reasoning marker + * `detectModelProfile` keys on (`` + `enable_thinking` for Qwen, + * `<|channel>thought` + `<|turn>` for Gemma 4). llama-server echoes the + * override back at `/props.chat_template`, which is the runtime's only + * profile signal — a template without those markers demotes a thinking + * model to `plain-instruct`, `buildGrammar` then drops the reasoning + * prelude from the GBNF root, and the sampler is forced to open with `[` + * and stalls in `ws` until `n_predict`. Prefer leaving this unset. + */ chatTemplateAsset?: string; tag?: string; /** @@ -199,7 +212,6 @@ export const LOCAL_MODELS_CATALOG: readonly LocalModelDef[] = [ minRamGb: 6, recommendedRamGb: 8, family: "qwen", - chatTemplateAsset: "qwen3.5-chat-template.jinja", supportsVision: true, mmprojUrl: "https://huggingface.co/unsloth/Qwen3.5-4B-GGUF/resolve/main/mmproj-F16.gguf", @@ -241,7 +253,6 @@ export const LOCAL_MODELS_CATALOG: readonly LocalModelDef[] = [ minRamGb: 24, recommendedRamGb: 36, family: "qwen", - chatTemplateAsset: "qwen3.5-chat-template.jinja", tag: "High Performance", supportsVision: true, mmprojUrl: diff --git a/src/prompt/index.ts b/src/prompt/index.ts index 1093181..3240667 100644 --- a/src/prompt/index.ts +++ b/src/prompt/index.ts @@ -18,6 +18,8 @@ export { checkBudget, truncateToTokens, computeEffectiveConversationCap, + minUsableContextWindow, + AGENT_FIXED_PROMPT_TOKENS, CONVERSATION_CAP_SAFETY_MARGIN, CONVERSATION_CAP_FLOOR, } from "./token-budget.js"; diff --git a/src/prompt/token-budget.ts b/src/prompt/token-budget.ts index 3408be2..2220ebd 100644 --- a/src/prompt/token-budget.ts +++ b/src/prompt/token-budget.ts @@ -113,6 +113,29 @@ export interface EffectiveConversationCapInput { */ export const CONVERSATION_CAP_SAFETY_MARGIN = 512; +/** + * Approximate token cost of the agent's fixed prompt scaffolding — the + * stable prefix (persona + tool catalog + capabilities + instructions), + * measured at ~5.2k and rounded up for drift. Only used for the startup + * sanity check on a model's context window; nothing depends on it being + * exact, and the real per-build figure comes from `checkBudget`. + */ +export const AGENT_FIXED_PROMPT_TOKENS = 6000; + +/** + * Smallest context window in which the agent can actually complete a + * step: fixed scaffolding + a full generation budget + boundary margin. + * `contextWindow` below this means every step will hit llama.cpp's + * context ceiling and come back `truncated`. + */ +export function minUsableContextWindow(completionMaxTokens: number): number { + return ( + AGENT_FIXED_PROMPT_TOKENS + + completionMaxTokens + + CONVERSATION_CAP_SAFETY_MARGIN + ); +} + /** * Hard minimum for the effective conversation cap. Even on tiny-context * models we keep at least this many tokens so the last user turn and a diff --git a/src/runtime/bootstrap.ts b/src/runtime/bootstrap.ts index 3f5afde..f0bb25e 100644 --- a/src/runtime/bootstrap.ts +++ b/src/runtime/bootstrap.ts @@ -34,11 +34,12 @@ import type { import { buildGrammar, detectModelProfile, + extractTotalSlots, ModelProfileManager, PLAIN_INSTRUCT_PROFILE, } from "../llm/index.js"; import { checkProfileGrammarAligned } from "../llm/profile-invariants.js"; -import { SlotManager } from "../llm/slot-manager.js"; +import { DEFAULT_SLOT_COUNT, SlotManager } from "../llm/slot-manager.js"; import { checkLlamaServer } from "../llm/llama-server-health.js"; import { ApprovalGate } from "../approval/approval-gate.js"; @@ -123,6 +124,7 @@ import { seedStarterSkillsIfMissing } from "../skills/seed-starter-skills.js"; import { DEFAULT_TOOL_DESCRIPTORS } from "../prompt/tool-descriptors.js"; import { filterToolDescriptorsByConfig } from "./filter-disabled-tools.js"; import { buildCapabilities } from "../prompt/capabilities.js"; +import { minUsableContextWindow } from "../prompt/token-budget.js"; import type { CapabilitiesSummary, SkillCatalogEntry, @@ -680,8 +682,32 @@ export async function createAgentRuntime( url: config.localModels.url, }); } else { - logger.info("slot manager using default slot count (probe miss)", { - slotCount: 4, + // Managed mode always defers the boot probe (the daemon may not be up + // yet), so this is the normal path there. `ModelProfileManager` calls + // `slotManager.resize()` on its first successful `/props` refresh at + // turn start; until then the pool is the single slot every + // llama-server is guaranteed to have. + logger.info("slot manager using conservative default (probe deferred)", { + slotCount: DEFAULT_SLOT_COUNT, + }); + } + + // A context window that cannot hold the fixed prompt plus a full + // generation budget makes every step come back `truncated` — the model + // burns its remaining tokens and never closes a tool-call array. Loud + // at startup because the failure mode downstream is silent. + const minUsableCtx = minUsableContextWindow( + config.localModels.completionMaxTokens, + ); + if (profile.contextWindow && profile.contextWindow < minUsableCtx) { + logger.warn("context window too small for the agent prompt", { + contextWindow: profile.contextWindow, + required: minUsableCtx, + completionMaxTokens: config.localModels.completionMaxTokens, + hint: + config.localModels.mode === "managed" + ? "raise localModels.managed.contextSize, lower localModels.completionMaxTokens, or pick a model that fits VRAM" + : "start llama-server with a larger --ctx-size, or lower localModels.completionMaxTokens", }); } @@ -961,6 +987,14 @@ export async function createAgentRuntime( initialModelId: modelAlias, grammarsDir: config.paths.grammarsDir, browserEnabled: config.browser.enabled, + onTotalSlots: (discovered) => { + if (discovered === slotManager.getSlotCount()) return; + logger.info("slot pool resized from /props", { + from: slotManager.getSlotCount(), + to: discovered, + }); + slotManager.resize(discovered); + }, logger, }) : undefined; @@ -2363,20 +2397,6 @@ function logResolvedProfile( return { profile: resolved, modelAlias: alias, totalSlots }; } -/** - * Extract `total_slots` from a `/props` payload. `llama-server` reports - * the number as an integer at the top level; anything else (missing, - * non-finite, non-positive) collapses to `null` so the caller falls back - * to the SlotManager default. - */ -function extractTotalSlots(props: Record): number | null { - const raw = props.total_slots; - if (typeof raw !== "number" || !Number.isFinite(raw) || raw < 1) { - return null; - } - return Math.trunc(raw); -} - /** * Only wire the hot-swap manager when the runtime actually talks to a * real llama-server. Any test override that replaces the HTTP layer