/** * Regression guard for the incremental per-message token cache in * `StatusLineComponent.getCachedContextBreakdown`. * * Before the cache: every call walked `session.messages` and ran * `estimateTokens` per message (~0.5 ms each native). For a 2,300-message * session this was a ~1.1 s blocking call. `updateEditorTopBorder()` is * invoked on every agent event (event-controller.ts:163), so during * streaming the UI froze for ~1.1 s every 2 s (the prior cache TTL). * * After the cache: messages are walked ONCE during warm-up; subsequent * refreshes append-only update the cache by `messages.length - cached` * (typically 0–1 new messages). The LAST message is recomputed every call * because its content may still be growing during streaming. Compaction * (messages.length shrinks) resets the cache. */ import { afterAll, beforeAll, describe, expect, it } from "bun:test"; import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import { StatusLineComponent } from "@oh-my-pi/pi-coding-agent/modes/components/status-line"; import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; import { computeNonMessageTokens, estimateToolSchemaTokens } from "@oh-my-pi/pi-coding-agent/modes/utils/context-usage"; import type { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; import { countTokens } from "@oh-my-pi/pi-natives"; beforeAll(async () => { resetSettingsForTest(); await Settings.init({ inMemory: true }); await initTheme(); }); afterAll(() => { resetSettingsForTest(); }); function makeSession(opts: { messages: unknown[]; systemPrompt?: string[]; tools?: { name: string; description: string; parameters?: unknown }[]; skills?: { name: string; description: string }[]; contextWindow?: number; modelId?: string; }): AgentSession { return { messages: opts.messages, systemPrompt: opts.systemPrompt ?? ["You are a helpful assistant."], agent: { state: { tools: opts.tools ?? [] } }, skills: opts.skills ?? [], model: { id: opts.modelId ?? "test-model", contextWindow: opts.contextWindow ?? 200_000 }, } as unknown as AgentSession; } function userMessage(text: string): unknown { return { role: "user", content: text }; } function assistantMessage(text: string): unknown { return { role: "assistant", content: [{ type: "text", text }] }; } describe("StatusLineComponent incremental context breakdown cache", () => { it("first call computes from scratch, second call returns same value", () => { const session = makeSession({ messages: Array.from({ length: 50 }, (_, i) => userMessage(`message ${i}`.repeat(10))), }); const comp = new StatusLineComponent(session); const first = comp.getCachedContextBreakdown(); const second = comp.getCachedContextBreakdown(); expect(first.usedTokens).toBeGreaterThan(0); expect(second.usedTokens).toBe(first.usedTokens); expect(second.contextWindow).toBe(200_000); }); it("appending a message increases the total by approximately the new message's tokens", () => { const session = makeSession({ messages: [userMessage("hello world"), userMessage("another message here")], }); const comp = new StatusLineComponent(session); const before = comp.getCachedContextBreakdown(); (session.messages as unknown[]).push(assistantMessage("a third message reply with more text content")); const after = comp.getCachedContextBreakdown(); expect(after.usedTokens).toBeGreaterThan(before.usedTokens); expect(after.contextWindow).toBe(before.contextWindow); }); it("compaction (messages.length shrinks) resets the cache and recomputes correctly", () => { const session = makeSession({ messages: Array.from({ length: 20 }, (_, i) => userMessage(`message ${i}`.repeat(10))), }); const comp = new StatusLineComponent(session); const before = comp.getCachedContextBreakdown(); expect(before.usedTokens).toBeGreaterThan(0); (session.messages as unknown[]).length = 0; (session.messages as unknown[]).push(userMessage("compacted summary")); const after = comp.getCachedContextBreakdown(); expect(after.usedTokens).toBeLessThan(before.usedTokens); expect(after.usedTokens).toBeGreaterThan(0); }); it("non-message inputs change → recomputes non-message portion", () => { const session = makeSession({ messages: [userMessage("hi")], systemPrompt: ["You are an assistant."], tools: [{ name: "bash", description: "Run shell commands", parameters: {} }], skills: [{ name: "code", description: "Write code" }], }); const comp = new StatusLineComponent(session); const v1 = comp.getCachedContextBreakdown(); const v2 = comp.getCachedContextBreakdown(); expect(v2.usedTokens).toBe(v1.usedTokens); (session.agent as { state: { tools: unknown[] } }).state.tools.push({ name: "edit", description: "Edit files", parameters: {}, }); const v3 = comp.getCachedContextBreakdown(); expect(v3.usedTokens).toBeGreaterThan(v2.usedTokens); }); it("non-message token shortcut matches previous category sum semantics", () => { const session = makeSession({ messages: [], systemPrompt: [ "You are an assistant.\n\n\n- code: Write code\n- review: Review code\n", "Loaded context file", "Runtime note", ], tools: [ { name: "bash", description: "Run shell commands", parameters: { type: "object", properties: { command: { type: "string" } } }, }, ], skills: [ { name: "code", description: "Write code" }, { name: "review", description: "Review code" }, ], }); const skillsTokens = countTokens(["code", "Write code", "review", "Review code"]); const previousCategorySum = Math.max(0, countTokens(session.systemPrompt?.[0] ?? "") - skillsTokens) + countTokens((session.systemPrompt ?? []).slice(1)) + estimateToolSchemaTokens(session.agent?.state?.tools ?? []) + skillsTokens; expect(new StatusLineComponent(session).getCachedContextBreakdown().usedTokens).toBe(previousCategorySum); expect(computeNonMessageTokens(session)).toBe(previousCategorySum); }); it("zero messages: produces only non-message tokens, no crash", () => { const session = makeSession({ messages: [] }); const comp = new StatusLineComponent(session); const result = comp.getCachedContextBreakdown(); expect(result.usedTokens).toBeGreaterThanOrEqual(0); expect(result.contextWindow).toBe(200_000); }); it("in-place mutation of a non-last message recomputes its tokens", () => { const original = userMessage("short") as { content: string }; const tail = userMessage("tail"); const session = makeSession({ messages: [original, tail] }); const comp = new StatusLineComponent(session); const before = comp.getCachedContextBreakdown(); // Mutate messages[0] in place — same object, larger content. original.content = "a much longer body that should tokenize to more".repeat(20); const after = comp.getCachedContextBreakdown(); expect(after.usedTokens).toBeGreaterThan(before.usedTokens); }); it("replaceMessages with same length but different shape recomputes tokens", () => { const session = makeSession({ messages: [userMessage("short a"), userMessage("short b")], }); const comp = new StatusLineComponent(session); const before = comp.getCachedContextBreakdown(); // Same-length replace: distinct message objects with larger payloads. (session as { messages: unknown[] }).messages = [ userMessage("a much longer payload".repeat(20)), userMessage("another longer payload".repeat(20)), ]; const after = comp.getCachedContextBreakdown(); expect(after.usedTokens).toBeGreaterThan(before.usedTokens); }); it("usage fetch error backs off — a failed fetch does not retrigger within the TTL window", async () => { const session = makeSession({ messages: [userMessage("hi")] }); let calls = 0; (session as { fetchUsageReports?: () => Promise }).fetchUsageReports = () => { calls++; return Promise.reject(new Error("network")); }; const comp = new StatusLineComponent(session); // First refresh → kicks off fetch #1. comp.refreshUsageInBackground(); // Let the rejected fetch settle so the .catch backoff stamp lands. await Bun.sleep(0); expect(calls).toBe(1); // Subsequent refreshes within the TTL window must not refetch. comp.refreshUsageInBackground(); comp.refreshUsageInBackground(); await Bun.sleep(0); expect(calls).toBe(1); }); });