diff --git a/electron/ai-edition/chat-compaction.test.ts b/electron/ai-edition/chat-compaction.test.ts index 63a61f379..5c9c17bd3 100644 --- a/electron/ai-edition/chat-compaction.test.ts +++ b/electron/ai-edition/chat-compaction.test.ts @@ -4,6 +4,7 @@ import { applyCompaction, budgetSnapshot, buildCompactionPrompt, + compactionReducesHistory, estimateHistoryTokens, shouldCompact, } from "./chat-compaction"; @@ -78,18 +79,33 @@ describe("shouldCompact", () => { }); describe("applyCompaction", () => { - it("inserts the summary message at the split point", () => { + it("replaces the summarized prefix with the summary message", () => { const msgs = [ - msg("user", "1", "u1"), - msg("assistant", "2", "a1"), + msg("user", "x".repeat(4_000), "u1"), + msg("assistant", "y".repeat(4_000), "a1"), msg("user", "3", "u2"), msg("assistant", "4", "a2"), ]; const out = applyCompaction(msgs, 2, "summary text", "2026-02-01T00:00:00.000Z"); - expect(out.length).toBe(5); - expect(out[2]?.content).toBe("summary text"); - expect(out[0]?.id).toBe("u1"); + expect(out).toHaveLength(3); + expect(out[0]?.content).toBe("summary text"); + expect(out[0]?.id).toMatch(/^summary_\d+$/); + expect(out[1]?.id).toBe("u2"); expect(out.at(-1)?.id).toBe("a2"); + expect(estimateHistoryTokens(out)).toBeLessThan(estimateHistoryTokens(msgs)); + expect(compactionReducesHistory(msgs, out)).toBe(true); + }); + + it("rejects a summary that does not reduce estimated context use", () => { + const msgs = [ + msg("user", "small", "u1"), + msg("assistant", "reply", "a1"), + msg("user", "recent", "u2"), + msg("assistant", "tail", "a2"), + ]; + const oversized = applyCompaction(msgs, 2, "x".repeat(1_000), "2026-02-01T00:00:00.000Z"); + + expect(compactionReducesHistory(msgs, oversized)).toBe(false); }); }); diff --git a/electron/ai-edition/chat-compaction.ts b/electron/ai-edition/chat-compaction.ts index e9faf4706..1b11b5e32 100644 --- a/electron/ai-edition/chat-compaction.ts +++ b/electron/ai-edition/chat-compaction.ts @@ -83,7 +83,6 @@ export function applyCompaction( summary: string, summaryAt: string, ): AiEditionChatMessage[] { - const head = messages.slice(0, splitIndex); const tail = messages.slice(splitIndex); const summaryMessage: AiEditionChatMessage = { id: `summary_${Date.now()}`, @@ -91,7 +90,15 @@ export function applyCompaction( content: summary, createdAt: summaryAt, }; - return [...head, summaryMessage, ...tail]; + return [summaryMessage, ...tail]; +} + +/** True only when a proposed compaction strictly reduces estimated context use. */ +export function compactionReducesHistory( + original: AiEditionChatMessage[], + compacted: AiEditionChatMessage[], +): boolean { + return estimateHistoryTokens(compacted) < estimateHistoryTokens(original); } /** System-prompt addendum to ask the LLM to compact its own history. */ diff --git a/electron/ai-edition/chat-service.compaction.test.ts b/electron/ai-edition/chat-service.compaction.test.ts new file mode 100644 index 000000000..ee9427844 --- /dev/null +++ b/electron/ai-edition/chat-service.compaction.test.ts @@ -0,0 +1,169 @@ +// Compaction as seen from chat-service: what the user keeps versus what the +// model is handed. Both seams are mocked — `invokeOpenScreenAgent` for the +// turn itself (so we can read the history it was given) and the chat model +// behind the summarizer, so no test here needs a provider or a key. + +import { beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("./deep-agent/service", () => ({ + invokeOpenScreenAgent: vi.fn(), +})); + +vi.mock("./deep-agent/chat-model", () => ({ + createOpenScreenChatModel: vi.fn(), + messageContentToText: (content: unknown) => String(content), +})); + +import { + compactSessionNow, + createSession, + getSessionContextUsage, + runChat, + selectSession, +} from "./chat-service"; +import { createOpenScreenChatModel } from "./deep-agent/chat-model"; +import { invokeOpenScreenAgent } from "./deep-agent/service"; +import type { LlmConfigStore } from "./llm-config-store"; + +const invokeMock = vi.mocked(invokeOpenScreenAgent); +const chatModelMock = vi.mocked(createOpenScreenChatModel); + +type ModelHistory = Array<{ role: "user" | "assistant" | "system"; content: string }>; + +let histories: ModelHistory[] = []; + +function stubConfig(): LlmConfigStore { + return { + getConfig: () => ({ provider: "openai", model: "gpt-4o" }), + getApiKey: () => "sk-test", + getCredential: () => ({ value: "sk-test", entry: { kind: "api-key", apiKey: "sk-test" } }), + } as unknown as LlmConfigStore; +} + +/** Point the summarizer at a fixed reply and return its call spy. */ +function stubSummarizer(reply: string) { + const invoke = vi.fn(async () => ({ content: reply })); + chatModelMock.mockImplementation( + async () => ({ invoke }) as unknown as Awaited>, + ); + return invoke; +} + +// Long enough that four of them clear the 70%-of-80k-tokens trip point, short +// enough that three of them do not. +const LONG = "x".repeat(60_000); + +beforeEach(() => { + histories = []; + invokeMock.mockReset(); + chatModelMock.mockReset(); + invokeMock.mockImplementation(async (args) => { + histories.push([...args.history]); + return { text: "ok", document: args.document, mutated: false }; + }); +}); + +describe("auto-compaction", () => { + it("leaves the transcript whole and compacts only what the model is given", async () => { + stubSummarizer("EARLIER CONTEXT"); + const session = createSession("proj_compact_transcript"); + for (let i = 0; i < 4; i += 1) { + await runChat("proj_compact_transcript", session.id, `${LONG}#${i}`, stubConfig()); + } + + // Four user turns, four replies, nothing deleted: this array is what the + // renderer shows, and the user never asked for half of it to go away. + const transcript = selectSession("proj_compact_transcript", session.id)?.messages ?? []; + expect(transcript).toHaveLength(8); + expect(transcript[0]?.content).toBe(`${LONG}#0`); + expect(transcript.filter((m) => m.role === "user")).toHaveLength(4); + + // The fourth turn is the one that tripped the budget: the model got the + // summary in place of the older half, not the whole conversation. + const history = histories.at(-1) ?? []; + expect(history[0]?.content).toBe("EARLIER CONTEXT"); + expect(history).toHaveLength(4); + expect(history.some((m) => m.content === `${LONG}#0`)).toBe(false); + expect(history.at(-1)?.content).toBe(`${LONG}#3`); + + // The context pill measures the payload, so compaction actually shows up: + // the whole transcript estimates at ~60k tokens, the payload at half. + const usage = getSessionContextUsage("proj_compact_transcript", session.id); + expect(usage?.usedTokens).toBeLessThan(40_000); + }); + + it("keeps the summary in the payload when the tail is longer than the window", async () => { + stubSummarizer("EARLIER CONTEXT"); + const session = createSession("proj_compact_window"); + for (let i = 0; i < 30; i += 1) { + await runChat("proj_compact_window", session.id, `turn ${i}`, stubConfig()); + } + const huge = LONG.repeat(4); + await runChat("proj_compact_window", session.id, huge, stubConfig()); + + // 31 messages survive the boundary — a plain slice(-20) would drop the + // summary we just paid a model call to produce. + const history = histories.at(-1) ?? []; + expect(history).toHaveLength(20); + expect(history[0]?.content).toBe("EARLIER CONTEXT"); + expect(history.at(-1)?.content).toBe(huge); + }); + + it("stops retrying after a summary that does not shrink the payload", async () => { + const oversized = stubSummarizer("z".repeat(400_000)); + const session = createSession("proj_compact_blocked"); + for (let i = 0; i < 5; i += 1) { + await runChat("proj_compact_blocked", session.id, `${LONG}#${i}`, stubConfig()); + } + + // Two more turns tripped the heuristic after the failure; neither paid + // for another summarizer call. + expect(oversized).toHaveBeenCalledTimes(1); + const history = histories.at(-1) ?? []; + expect(history.some((m) => m.content === "EARLIER CONTEXT")).toBe(false); + expect(selectSession("proj_compact_blocked", session.id)?.messages).toHaveLength(10); + + // The Compact button is an explicit request, so it tries again — and a + // success unblocks the automatic path. + const usable = stubSummarizer("EARLIER CONTEXT"); + const manual = await compactSessionNow("proj_compact_blocked", session.id, stubConfig()); + expect(usable).toHaveBeenCalledTimes(1); + expect(manual?.summary).toBe("EARLIER CONTEXT"); + expect(manual?.session.messages).toHaveLength(10); + + await runChat("proj_compact_blocked", session.id, "and then?", stubConfig()); + expect(histories.at(-1)?.[0]?.content).toBe("EARLIER CONTEXT"); + }); + + // The regression this guards is `planCompaction` measuring the wrong list. + // `splitIndex` comes back from `shouldCompact` as an index INTO WHAT IT WAS + // GIVEN, and it is then applied to the payload. Measure the transcript + // instead — which never shrinks, so it keeps tripping — and the index runs + // off the end of the much shorter payload, so `payload.slice(0, splitIndex)` + // swallows the whole thing, current user turn included. The model is then + // asked to answer a question it was never shown. + // + // Three turns is not enough to see it: the collapse needs a payload that has + // already been compacted at least once, so the two lists have diverged. + it("never summarizes away the turn the user just sent", async () => { + stubSummarizer("EARLIER CONTEXT"); + const session = createSession("proj_compact_current_turn"); + for (let i = 0; i < 10; i += 1) { + await runChat("proj_compact_current_turn", session.id, `${LONG}#${i}`, stubConfig()); + } + + // Every turn, not just the last: the collapse is intermittent, so a + // spot-check on `histories.at(-1)` walks straight past it. When it bites, + // the payload is `[summary]` alone, so the last entry is the summary + // rather than the message the user just typed — which is exactly what + // this asserts. (Turn 0 is legitimately a one-message payload, so length + // is the wrong thing to check.) + expect(histories).toHaveLength(10); + histories.forEach((history, turn) => { + expect( + history.at(-1)?.content, + `turn ${turn} was handed a payload that did not end with the user's message`, + ).toBe(`${LONG}#${turn}`); + }); + }); +}); diff --git a/electron/ai-edition/chat-service.ts b/electron/ai-edition/chat-service.ts index 180f8e0bc..eec200364 100644 --- a/electron/ai-edition/chat-service.ts +++ b/electron/ai-edition/chat-service.ts @@ -26,6 +26,7 @@ import { budgetSnapshot, buildCompactionPrompt, COMPACTION_SYSTEM_PROMPT, + compactionReducesHistory, DEFAULT_BUDGET_TOKENS, shouldCompact, } from "./chat-compaction"; @@ -112,12 +113,55 @@ function dropCheckpointsFrom( ); } +// ponytail: what compaction leaves behind. `coveredCount` counts the leading +// transcript messages the summary stands in for — the transcript itself is +// never rewritten, so the user keeps every message they wrote while the model +// gets the shortened list. Sessions are in-memory only; deleting the messages +// the renderer shows would be unrecoverable, and nothing in the UI would say +// it happened. +interface SessionCompaction { + summary: AiEditionChatMessage; + coveredCount: number; +} + export interface ChatSession { id: string; projectId: string; title: string; createdAt: string; messages: AiEditionChatMessage[]; + /** Main-process bookkeeping: the compaction boundary, not part of the + * transcript. See `modelMessages`. */ + compaction?: SessionCompaction; + /** Set when a summarize call came back no smaller than what it replaced. + * Auto-compaction then stops trying — otherwise every following turn pays + * for the same useless summarizer call. */ + compactionBlocked?: boolean; +} + +/** The message list compaction hands the model: summary first, then everything + * after the boundary. Identical to the transcript until a compaction lands. */ +function modelMessages(session: ChatSession): AiEditionChatMessage[] { + const state = session.compaction; + if (!state) return session.messages; + return [state.summary, ...session.messages.slice(state.coveredCount)]; +} + +// The model sees a sliding window of recent turns, not the whole session. +const MODEL_HISTORY_WINDOW = 20; + +/** + * Window `modelMessages` down to what we send. The summary is pinned to the + * front rather than left to the window: the tail after a compaction is roughly + * half the session, so at the token counts that trip compaction in the first + * place a plain `slice(-N)` drops the very summary we just paid to produce. + */ +function modelHistory(session: ChatSession): AiEditionChatMessage[] { + const messages = modelMessages(session); + if (messages.length <= MODEL_HISTORY_WINDOW) return messages; + const summary = session.compaction?.summary; + if (!summary) return messages.slice(-MODEL_HISTORY_WINDOW); + return [summary, ...messages.slice(1).slice(-(MODEL_HISTORY_WINDOW - 1))]; } export interface ChatSessionSummary { @@ -179,7 +223,15 @@ export function selectSession(projectId: string, sessionId: string): ChatSession const s = m?.get(sessionId); if (!s) return null; // ponytail: shallow-copy messages so the caller can't mutate the live array. - return { ...s, messages: [...s.messages] }; + // The compaction boundary stays behind: it is main-process bookkeeping, and + // every caller of this wants the transcript as the user sees it. + return { + id: s.id, + projectId: s.projectId, + title: s.title, + createdAt: s.createdAt, + messages: [...s.messages], + }; } export function renameSession( @@ -316,15 +368,15 @@ export async function runChat( const editsAllowed = config.allowAgentEdits !== false; - // P3.7 — context compaction: when the session grows past the heuristic - // budget, summarize the older half into a single "Earlier context" - // assistant message. The current user turn stays uncompacted, so the - // model still sees the request verbatim. - const decision = shouldCompact(session.messages); - if (decision && decision.compact) { + // P3.7 — context compaction: when what we send the model grows past the + // heuristic budget, summarize the older half into a single "Earlier + // context" assistant message. The current user turn stays uncompacted, so + // the model still sees the request verbatim. + const plan = session.compactionBlocked ? null : planCompaction(session); + if (plan) { await tryCompactSession({ session, - splitIndex: decision.splitIndex, + plan, apiKey: apiKey ?? "", provider: config.provider, model: config.model, @@ -333,9 +385,10 @@ export async function runChat( }); } - const history = session.messages - .slice(-20) - .map((m) => ({ role: m.role as "user" | "assistant" | "system", content: m.content })); + const history = modelHistory(session).map((m) => ({ + role: m.role as "user" | "assistant" | "system", + content: m.content, + })); const appliedToolCalls: AiEditionToolCallSummary[] = []; @@ -453,6 +506,12 @@ export function rewindToMessage( toolCalls: m.toolCalls ? [...m.toolCalls] : undefined, })); session.messages = survived; + // ponytail: a rewind that cuts back past the compaction boundary leaves a + // summary standing for messages this lineage no longer has. Drop it and let + // the next turn re-derive one, same reasoning as the checkpoints below. + if (session.compaction && session.compaction.coveredCount > survived.length) { + session.compaction = undefined; + } dropCheckpointsFrom(projectId, sessionId, target.checkpointId); return { @@ -549,12 +608,15 @@ export async function compactSessionNow( const credential = llmConfig.getCredential(def.id, def.envKeys); const apiKey = credential?.value ?? ""; - const decision = shouldCompact(session.messages); - if (!decision || !decision.compact) return null; + // The button is an explicit request, so it ignores `compactionBlocked` — + // the user knows they are spending a summarizer call, and a success clears + // the flag for the automatic path too. + const plan = planCompaction(session); + if (!plan) return null; const ok = await tryCompactSession({ session, - splitIndex: decision.splitIndex, + plan, apiKey, provider: config.provider, model: config.model, @@ -580,7 +642,10 @@ export function getSessionContextUsage( ): { usedTokens: number; budgetTokens: number; ratio: number; fillPercent: number } | null { const session = sessionsByProject.get(projectId)?.get(sessionId); if (!session) return null; - const snap = budgetSnapshot(session.messages, budgetTokens); + // 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); const fillPercent = Math.min(100, Math.round(snap.ratio * 100)); return { usedTokens: snap.usedTokens, @@ -613,7 +678,7 @@ export function getSessionBudget( ): SessionBudgetSnapshot | null { const s = sessionsByProject.get(projectId)?.get(sessionId); if (!s) return null; - const snap = budgetSnapshot(s.messages, budgetTokens); + const snap = budgetSnapshot(modelMessages(s), budgetTokens); return { usedTokens: snap.usedTokens, budgetTokens: snap.budgetTokens, @@ -639,12 +704,12 @@ export async function compactSession( const credential = def ? llmConfig.getCredential(def.id, def.envKeys) : null; const apiKey = credential?.value ?? ""; - const decision = shouldCompact(session.messages); - if (!decision) return { summaryMessageId: null, summary: "" }; + const plan = planCompaction(session); + if (!plan) return { summaryMessageId: null, summary: "" }; const ok = await tryCompactSession({ session, - splitIndex: decision.splitIndex, + plan, apiKey, provider: config.provider, model: config.model, @@ -654,17 +719,53 @@ export async function compactSession( return ok; } +interface CompactionPlan { + /** What the model currently gets — the input compaction shortens. */ + payload: AiEditionChatMessage[]; + /** Boundary inside `payload`. */ + splitIndex: number; + /** The same boundary expressed as a count of transcript messages. */ + coveredCount: number; +} + +/** + * Decide whether the next turn should compact, measuring the payload rather + * than the transcript. Measuring the transcript would re-trip on every turn + * for the rest of the session, since compaction no longer shrinks it. + * + * A second compaction folds the previous summary into the new one: it sits at + * `payload[0]`, so it is part of the prefix being summarized. + */ +function planCompaction(session: ChatSession): CompactionPlan | null { + const payload = modelMessages(session); + const decision = shouldCompact(payload); + if (!decision?.compact || decision.splitIndex <= 0) return null; + // Payload index → transcript index. With a summary in front, payload[i] + // is transcript message `coveredCount + i - 1`. + const offset = session.compaction ? session.compaction.coveredCount - 1 : 0; + return { + payload, + splitIndex: decision.splitIndex, + coveredCount: decision.splitIndex + offset, + }; +} + +/** + * Summarize the older half of the model payload and record the boundary on the + * session. The transcript is left alone — `session.messages` is what the + * renderer shows, and compaction is a fact about the model's input. + */ async function tryCompactSession(opts: { session: ChatSession; - splitIndex: number; + plan: CompactionPlan; apiKey: string; provider: string; model: string; baseUrl?: string; reasoningEffort?: string; }): Promise<{ summaryMessageId: string | null; summary: string } | null> { - const { session, splitIndex, apiKey, provider, model, baseUrl, reasoningEffort } = opts; - const oldMessages = session.messages.slice(0, splitIndex); + const { session, plan, apiKey, provider, model, baseUrl, reasoningEffort } = opts; + const oldMessages = plan.payload.slice(0, plan.splitIndex); if (oldMessages.length === 0) return null; const prompt = buildCompactionPrompt(oldMessages); @@ -695,15 +796,25 @@ async function tryCompactSession(opts: { } const compacted = applyCompaction( - session.messages, - splitIndex, + plan.payload, + plan.splitIndex, summary, new Date().toISOString(), ); - const inserted = compacted[splitIndex]; - session.messages = compacted; + const summaryMessage = compacted[0]; + if (!summaryMessage) return null; + if (!compactionReducesHistory(plan.payload, compacted)) { + // ponytail: the model handed back a summary at least as long as the + // messages it replaced. Adopting it would grow the payload, and + // retrying next turn just buys the same answer again — so stop asking + // until the user compacts by hand. + session.compactionBlocked = true; + return null; + } + session.compaction = { summary: summaryMessage, coveredCount: plan.coveredCount }; + session.compactionBlocked = false; return { - summaryMessageId: inserted?.id ?? null, + summaryMessageId: summaryMessage.id, summary, }; }