diff --git a/src/llm/llama-server-client.test.ts b/src/llm/llama-server-client.test.ts index 0bfcc8d8..ff0e46e9 100644 --- a/src/llm/llama-server-client.test.ts +++ b/src/llm/llama-server-client.test.ts @@ -50,7 +50,10 @@ describe("LlamaServerClient.complete", () => { expect(result.content).toBe('{"tool":"finish","args":{}}'); expect(result.reasoningContent).toBe(""); - expect(result.timing.promptTokens).toBe(40); + // 40 evaluated this request + 30 reused from the KV cache: the + // prompt the model saw was 70 tokens, and that is what occupancy + // consumers (the TUI context chip among them) need reported. + expect(result.timing.promptTokens).toBe(70); expect(result.cacheHitTokens).toBe(30); expect(result.slotId).toBe(2); expect(result.modelId).toBe("qwen-test"); @@ -65,6 +68,29 @@ describe("LlamaServerClient.complete", () => { expect(snapshot.body.repeat_last_n).toBe(256); }); + it("reports the bare evaluated count when nothing was cached", async () => { + const client = new LlamaServerClient({ + baseUrl: "http://127.0.0.1:9999", + fetchImpl: createMockFetch(async () => + new Response( + JSON.stringify({ + content: "ok", + stop: true, + truncated: false, + timings: { prompt_ms: 10, predicted_ms: 20, prompt_n: 40, predicted_n: 8 }, + slot_id: 0, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ), + }); + + const result = await client.complete({ prompt: "hello", maxTokens: 16 }); + + expect(result.timing.promptTokens).toBe(40); + expect(result.cacheHitTokens).toBe(0); + }); + it("forwards explicit repeatPenalty / repeatLastN overrides", async () => { let captured: Record | null = null; const client = new LlamaServerClient({ diff --git a/src/llm/llama-server-client.ts b/src/llm/llama-server-client.ts index 3abe229a..f2d8c2e4 100644 --- a/src/llm/llama-server-client.ts +++ b/src/llm/llama-server-client.ts @@ -551,6 +551,17 @@ function normaliseCompletionResponse( payload: Record, ): CompletionResult { const timings = (payload.timings ?? {}) as Record; + // `prompt_n` / `tokens_evaluated` count only the tokens llama-server + // actually evaluated this request — the prefix reused from the KV + // cache (`tokens_cached`) is excluded. Every consumer of + // `timing.promptTokens` treats it as "how big was the prompt" (their + // fallback is `prompt.tokens.total`, and the TUI shows it as occupied + // context), so report the whole prompt: evaluated + cached. On a warm + // cache the raw `prompt_n` is a small fraction of the prompt and the + // context readout collapsed to it, then leapt back to the estimator's + // full figure the moment anything reprojected it. + const evaluatedTokens = toNumber(timings.prompt_n ?? payload.tokens_evaluated); + const cachedTokens = toNumber(payload.tokens_cached); return { content: typeof payload.content === "string" ? payload.content : "", reasoningContent: @@ -562,10 +573,10 @@ function normaliseCompletionResponse( timing: { promptMs: toNumber(timings.prompt_ms), predictedMs: toNumber(timings.predicted_ms), - promptTokens: toNumber(timings.prompt_n ?? payload.tokens_evaluated), + promptTokens: evaluatedTokens + cachedTokens, predictedTokens: toNumber(timings.predicted_n ?? payload.tokens_predicted), }, - cacheHitTokens: toNumber(payload.tokens_cached), + cacheHitTokens: cachedTokens, slotId: toNumber(payload.slot_id ?? payload.id_slot, -1), modelId: typeof payload.model === "string" ? payload.model : null, diff --git a/src/tui/agent-event-reducer.ts b/src/tui/agent-event-reducer.ts index 2f8baaf6..57d2f01f 100644 --- a/src/tui/agent-event-reducer.ts +++ b/src/tui/agent-event-reducer.ts @@ -716,9 +716,11 @@ function reduceStepEvent( case "llm_completed": { // `prompt_built` carried an estimate (`estimateTokens` over-counts // by design); the provider just reported what its own tokenizer - // actually counted — llama.cpp from `tokens_evaluated`, an - // OpenAI-compatible cloud from `usage.prompt_tokens`. Prefer it, - // and leave the estimate standing when nothing was reported. + // actually counted — llama.cpp from `prompt_n + tokens_cached` + // (the whole prompt, not just the slice evaluated past the KV + // cache), an OpenAI-compatible cloud from `usage.prompt_tokens`. + // Prefer it, and leave the estimate standing when nothing was + // reported. const counted = event.completion.timing?.promptTokens ?? 0; if (counted <= 0) return state; return {