diff --git a/electron/ai-edition/chat-compaction.test.ts b/electron/ai-edition/chat-compaction.test.ts index 838458dd..758136c6 100644 --- a/electron/ai-edition/chat-compaction.test.ts +++ b/electron/ai-edition/chat-compaction.test.ts @@ -27,19 +27,21 @@ describe("estimateHistoryTokens", () => { expect(tokens).toBe(100); }); - it("adds 4 tokens per tool call", () => { - const base = estimateHistoryTokens([msg("user", "hi")]); + it("does not bill tool calls, which are never sent", () => { + // `chat-service` maps the history to `{role, content}` before handing it to the + // provider, so a message's tool-call names and summaries are rendered in the + // chat and then dropped. Counting them inflated the context pill with text + // nobody pays for. const withTool = estimateHistoryTokens([ { id: "a", role: "assistant", - content: "done", + content: "x".repeat(400), createdAt: "2026-01-01T00:00:00.000Z", - toolCalls: [{ name: "addTrim", summary: "skip 5-8s" }], + toolCalls: [{ name: "addTrim", summary: "y".repeat(400) }], }, ]); - // tool adds roughly: 16 + 4-chars-per-token of name+summary - expect(withTool).toBeGreaterThan(base); + expect(withTool).toBe(100); }); }); diff --git a/electron/ai-edition/chat-compaction.ts b/electron/ai-edition/chat-compaction.ts index 5be23d87..87554d69 100644 --- a/electron/ai-edition/chat-compaction.ts +++ b/electron/ai-edition/chat-compaction.ts @@ -39,11 +39,13 @@ export const DEFAULT_BUDGET_TOKENS = 80_000; export function estimateHistoryTokens(messages: AiEditionChatMessage[]): number { let chars = 0; for (const m of messages) { - // 4 chars per token + 4 tokens per message overhead (rough). + // Content only. `chat-service` maps the history to `{role, content}` before it + // sends it, so a message's tool-call names and summaries never reach the model -- + // billing them here inflated the pill with text nobody pays for. The two + // compaction comparisons are unaffected: both sides were counted the same way. chars += m.content.length; - for (const tc of m.toolCalls ?? []) - chars += (tc.name?.length ?? 0) + (tc.summary?.length ?? 0) + 16; } + // 4 chars per token (rough). return Math.ceil(chars / CHARS_PER_TOKEN); } diff --git a/electron/ai-edition/chat-service.compaction.test.ts b/electron/ai-edition/chat-service.compaction.test.ts index 2489398c..37f715b3 100644 --- a/electron/ai-edition/chat-service.compaction.test.ts +++ b/electron/ai-edition/chat-service.compaction.test.ts @@ -114,6 +114,33 @@ describe("compaction", () => { expect(usage?.usedTokens).toBeLessThan(40_000); }); + it("measures the WINDOW the model is sent, not the whole post-compaction list", async () => { + // The model gets `MODEL_HISTORY_WINDOW` (20) messages, not everything after the + // compaction boundary. Measuring the wider list made the pill report roughly + // double on a long session, and made pressing Compact halve the number while + // what actually reaches the provider barely moved -- the same class of wrong + // number this pill exists to avoid. + const session = createSession("proj_window"); + // 12 turns = 24 messages, so 4 fall outside the window. + for (let i = 0; i < 12; i += 1) { + await runChat("proj_window", session.id, `${LONG}#${i}`, stubConfig()); + } + + const transcript = selectSession("proj_window", session.id)?.messages ?? []; + expect(transcript).toHaveLength(24); + const sentToTheModel = histories.at(-1) ?? []; + // The last turn's payload is what the pill has to describe. + const sentChars = sentToTheModel.reduce((acc, m) => acc + m.content.length, 0); + const transcriptChars = transcript.reduce((acc, m) => acc + m.content.length, 0); + expect(sentChars).toBeLessThan(transcriptChars); + + const usage = getSessionContextUsage("proj_window", session.id); + // Within one turn of the payload (the pill is read after the reply lands, the + // payload was built before it), and nowhere near the whole transcript. + expect(usage?.usedTokens).toBeLessThan(Math.ceil(transcriptChars / 4)); + expect(usage?.usedTokens).toBeGreaterThanOrEqual(Math.ceil(sentChars / 4)); + }); + it("compacts an ORDINARY conversation — the button is not gated by a budget", async () => { // The same guessed budget gated the manual path: `compactSessionNow` // went through the same heuristic, so below 70% of 80k the button did diff --git a/electron/ai-edition/chat-service.ts b/electron/ai-edition/chat-service.ts index 89c1d78e..4dcac342 100644 --- a/electron/ai-edition/chat-service.ts +++ b/electron/ai-edition/chat-service.ts @@ -630,10 +630,15 @@ export function getSessionContextUsage( ): { usedTokens: number; budgetTokens: number; ratio: number; fillPercent: number } | null { const session = sessionsByProject.get(projectId)?.get(sessionId); if (!session) return null; - // Measured on what the model is given, not on the transcript — after a + // Measured on what the model is given, not on the transcript -- after a // compaction those differ, and the number that matters is the one that // fills the context window. - const snap = budgetSnapshot(modelMessages(session), budgetTokens); + // + // `modelHistory`, NOT `modelMessages`: the model gets a 20-message window, so on + // a 40-message session `modelMessages` reports double what is sent, and pressing + // Compact halves the reported number while what actually reaches the provider + // barely moves. That is the same class of wrong number this pill exists to avoid. + const snap = budgetSnapshot(modelHistory(session), budgetTokens); const fillPercent = Math.min(100, Math.round(snap.ratio * 100)); return { usedTokens: snap.usedTokens, @@ -666,7 +671,8 @@ export function getSessionBudget( ): SessionBudgetSnapshot | null { const s = sessionsByProject.get(projectId)?.get(sessionId); if (!s) return null; - const snap = budgetSnapshot(modelMessages(s), budgetTokens); + // Same window the model is actually sent — see `getSessionContextUsage`. + const snap = budgetSnapshot(modelHistory(s), budgetTokens); return { usedTokens: snap.usedTokens, budgetTokens: snap.budgetTokens, diff --git a/src/components/ai-edition/LeftPanel.tsx b/src/components/ai-edition/LeftPanel.tsx index 1942d1da..f7ea5a95 100644 --- a/src/components/ai-edition/LeftPanel.tsx +++ b/src/components/ai-edition/LeftPanel.tsx @@ -27,11 +27,11 @@ import { } from "../../../electron/ai-edition/provider-registry"; import { ChatWelcome } from "./ChatWelcome"; import { canSendChat } from "./chatAvailability"; -import { computeBudget } from "./chatBudget"; import { ChatHistoryModal, SourceTranscriptModal } from "./Modals"; import styles from "./NewEditorShell.module.css"; import { ProviderSettings } from "./ProviderSettings"; import { TranscriptionStatusDot } from "./TranscriptionStatus"; +import { useChatBudget } from "./useChatBudget"; export type LeftTab = "chat" | "media"; @@ -1141,10 +1141,12 @@ function ChatStripPanel() { }); }, [llmConfig]); - // Real context usage — feeds the badge in the chat strip and gates the - // auto-compact heuristic on the main side. Recomputed on every messages - // change so the % tracks the live history. - const budget = computeBudget(messages); + // Prefer the main process's estimate of the windowed history it actually sends, so + // manual compaction can shrink this meter while the complete transcript remains + // visible. The renderer estimate is only shown until the first answer arrives -- + // including in web builds, where the shim answers with its own transcript estimate + // rather than leaving the fallback in place. + const budget = useChatBudget({ projectId, sessionId: activeSessionId, messages }); const [compactNowPending, setCompactNowPending] = useState(false); const compactNow = useCallback(async () => { diff --git a/src/components/ai-edition/chatBudget.ts b/src/components/ai-edition/chatBudget.ts index 912c9187..91834762 100644 --- a/src/components/ai-edition/chatBudget.ts +++ b/src/components/ai-edition/chatBudget.ts @@ -1,12 +1,17 @@ // Renderer-side budget helper. Mirrors `electron/ai-edition/chat-compaction.ts` // but inline so we don't drag electron/ into the renderer bundle. // -// This feeds the context pill and NOTHING else — no code decides anything from -// it. `DEFAULT_CHAT_BUDGET_TOKENS` is a made-up denominator (the app has no way -// to ask a provider how big its context window is), which is exactly why the -// automatic compaction that used to branch on the main-process twin is gone. -// Read the pill as "the conversation is about this big", never as "you are this -// close to a limit", and do not let this number regain a decision. +// This is the renderer fallback for the context pill while native usage is loading +// or has never arrived. Desktop builds replace it with the main process's estimate +// of the windowed history it actually sends, which understands compaction. +// +// Either way it feeds the pill and NOTHING else -- no code decides anything from it. +// `DEFAULT_CHAT_BUDGET_TOKENS` is a made-up denominator (the app has no way to ask a +// provider how big its context window is), which is exactly why the automatic +// compaction that used to branch on the main-process twin is gone. Read the pill as +// "the conversation is about this big", never as "you are this close to a limit", +// and do not let this number regain a decision. Understanding compaction does not +// make the denominator any less invented. const CHARS_PER_TOKEN = 4; @@ -18,24 +23,23 @@ export interface ChatBudget { const DEFAULT_CHAT_BUDGET_TOKENS = 80_000; -interface RenderableChatMessage { +export interface RenderableChatMessage { content: string; - toolCalls?: Array<{ name?: string; summary?: string }>; } -function estimateTokens(messages: RenderableChatMessage[]): number { +// Content only, matching `estimateHistoryTokens` in the main process: tool-call +// names and summaries are rendered in the chat but stripped before the history is +// sent, so counting them here would bill text the model never sees. +function estimateTokens(messages: readonly RenderableChatMessage[]): number { let chars = 0; for (const m of messages) { chars += m.content.length; - for (const tc of m.toolCalls ?? []) { - chars += (tc.name?.length ?? 0) + (tc.summary?.length ?? 0) + 16; - } } return Math.ceil(chars / CHARS_PER_TOKEN); } export function computeBudget( - messages: RenderableChatMessage[], + messages: readonly RenderableChatMessage[], budgetTokens: number = DEFAULT_CHAT_BUDGET_TOKENS, ): ChatBudget { const used = estimateTokens(messages); diff --git a/src/components/ai-edition/useChatBudget.test.ts b/src/components/ai-edition/useChatBudget.test.ts new file mode 100644 index 00000000..83944076 --- /dev/null +++ b/src/components/ai-edition/useChatBudget.test.ts @@ -0,0 +1,191 @@ +// @vitest-environment jsdom +import { act, renderHook, waitFor } from "@testing-library/react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { useChatBudget } from "./useChatBudget"; + +const chatBudgetMock = vi.hoisted(() => vi.fn()); + +vi.mock("@/native/client", () => ({ + nativeBridgeClient: { + aiEdition: { chatBudget: chatBudgetMock }, + }, +})); + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (reason?: unknown) => void; + const promise = new Promise((done, fail) => { + resolve = done; + reject = fail; + }); + return { promise, resolve, reject }; +} + +const message = (content: string) => [{ content }]; + +describe("useChatBudget", () => { + beforeEach(() => chatBudgetMock.mockReset()); + + it("uses the transcript estimate until native model-context usage arrives", async () => { + const native = deferred<{ + usedTokens: number; + budgetTokens: number; + ratio: number; + fillPercent: number; + }>(); + chatBudgetMock.mockReturnValue(native.promise); + const visibleMessages = message("x".repeat(400)); + + const { result } = renderHook(() => + useChatBudget({ projectId: "project_1", sessionId: "session_1", messages: visibleMessages }), + ); + expect(result.current.usedTokens).toBe(100); + + act(() => + native.resolve({ usedTokens: 12, budgetTokens: 80_000, ratio: 0.00015, fillPercent: 0.015 }), + ); + await waitFor(() => expect(result.current.usedTokens).toBe(12)); + }); + + it("keeps the transcript estimate when no session is selected", () => { + const visibleMessages = message("x".repeat(400)); + const { result } = renderHook(() => + useChatBudget({ projectId: "project_1", sessionId: null, messages: visibleMessages }), + ); + + expect(result.current.usedTokens).toBe(100); + expect(chatBudgetMock).not.toHaveBeenCalled(); + }); + + it("tracks the transcript while native usage has never arrived", async () => { + // The transcript has to keep MOVING, not just happen to match: asserting 100 + // before and after a null answer cannot tell "fell back correctly" apart from + // "the effect never ran". + chatBudgetMock.mockResolvedValue(undefined); + const { result, rerender } = renderHook( + ({ messages }) => useChatBudget({ projectId: "project_1", sessionId: "session_1", messages }), + { initialProps: { messages: message("x".repeat(400)) } }, + ); + await waitFor(() => expect(chatBudgetMock).toHaveBeenCalledTimes(1)); + expect(result.current.usedTokens).toBe(100); + + rerender({ messages: message("x".repeat(800)) }); + await waitFor(() => expect(result.current.usedTokens).toBe(200)); + expect(chatBudgetMock).toHaveBeenCalledTimes(2); + }); + + it("keeps the last native number when a refresh fails, instead of swapping quantity", async () => { + // The two numbers are not interchangeable. After a compaction the native one is + // roughly half the transcript estimate, so silently substituting the transcript + // on a failed refresh made the compaction visibly un-happen -- the PR's own + // premise turned against it. A stale native number is the honest answer here. + const warn = vi.spyOn(console, "warn").mockImplementation(() => { + // swallowed: the point is that it is CALLED, not what it prints + }); + chatBudgetMock + .mockResolvedValueOnce({ + usedTokens: 12, + budgetTokens: 80_000, + ratio: 0.00015, + fillPercent: 0.015, + }) + .mockRejectedValueOnce(new Error("native budget unavailable")); + const { result, rerender } = renderHook( + ({ messages }) => useChatBudget({ projectId: "project_1", sessionId: "session_1", messages }), + { initialProps: { messages: message("x".repeat(400)) } }, + ); + await waitFor(() => expect(result.current.usedTokens).toBe(12)); + + rerender({ messages: message("x".repeat(800)) }); + await waitFor(() => expect(chatBudgetMock).toHaveBeenCalledTimes(2)); + await act(async () => Promise.resolve()); + expect(result.current.usedTokens).toBe(12); + // And it is not silent: renaming the bridge action must leave a trace. + expect(warn).toHaveBeenCalled(); + warn.mockRestore(); + }); + + it("drops a native number when the session changes", async () => { + // The kept-on-failure number belongs to one conversation. Session 2's first + // refresh never answering must not leave session 1's number on screen. + const pending = deferred(); + chatBudgetMock + .mockResolvedValueOnce({ + usedTokens: 12, + budgetTokens: 80_000, + ratio: 0.00015, + fillPercent: 0.015, + }) + .mockReturnValueOnce(pending.promise); + const visibleMessages = message("x".repeat(400)); + const { result, rerender } = renderHook( + ({ sessionId }) => + useChatBudget({ projectId: "project_1", sessionId, messages: visibleMessages }), + { initialProps: { sessionId: "session_1" } }, + ); + await waitFor(() => expect(result.current.usedTokens).toBe(12)); + + rerender({ sessionId: "session_2" }); + expect(result.current.usedTokens).toBe(100); + }); + + it("ignores a late response from the previously selected session", async () => { + const first = deferred<{ + usedTokens: number; + budgetTokens: number; + ratio: number; + fillPercent: number; + }>(); + const second = deferred<{ + usedTokens: number; + budgetTokens: number; + ratio: number; + fillPercent: number; + }>(); + chatBudgetMock.mockReturnValueOnce(first.promise).mockReturnValueOnce(second.promise); + const visibleMessages = message("visible transcript"); + + const { result, rerender } = renderHook( + ({ sessionId }) => + useChatBudget({ projectId: "project_1", sessionId, messages: visibleMessages }), + { initialProps: { sessionId: "session_1" } }, + ); + rerender({ sessionId: "session_2" }); + act(() => + second.resolve({ usedTokens: 20, budgetTokens: 80_000, ratio: 0.00025, fillPercent: 0.025 }), + ); + await waitFor(() => expect(result.current.usedTokens).toBe(20)); + + act(() => + first.resolve({ usedTokens: 999, budgetTokens: 80_000, ratio: 0.012, fillPercent: 1.2 }), + ); + await act(async () => Promise.resolve()); + expect(result.current.usedTokens).toBe(20); + }); + + it("refreshes native usage when compaction returns a new transcript array", async () => { + chatBudgetMock + .mockResolvedValueOnce({ + usedTokens: 500, + budgetTokens: 80_000, + ratio: 0.00625, + fillPercent: 0.625, + }) + .mockResolvedValueOnce({ + usedTokens: 40, + budgetTokens: 80_000, + ratio: 0.0005, + fillPercent: 0.05, + }); + const visibleMessages = message("the transcript remains visible"); + const { result, rerender } = renderHook( + ({ messages }) => useChatBudget({ projectId: "project_1", sessionId: "session_1", messages }), + { initialProps: { messages: visibleMessages } }, + ); + await waitFor(() => expect(result.current.usedTokens).toBe(500)); + + rerender({ messages: [...visibleMessages] }); + await waitFor(() => expect(result.current.usedTokens).toBe(40)); + expect(chatBudgetMock).toHaveBeenCalledTimes(2); + }); +}); diff --git a/src/components/ai-edition/useChatBudget.ts b/src/components/ai-edition/useChatBudget.ts new file mode 100644 index 00000000..2fd64657 Binary files /dev/null and b/src/components/ai-edition/useChatBudget.ts differ