From e805d2f756b785776cb1f60b97a253e3ae4c837c Mon Sep 17 00:00:00 2001 From: Marius Date: Tue, 12 May 2026 10:36:56 +0200 Subject: [PATCH] fix(vscode): refresh checkpoint header stats (#10156) * fix(vscode): refresh checkpoint header stats * test(vscode): update story session mock * chore: update kilo-vscode visual regression baselines --------- Co-authored-by: github-actions[bot] --- .changeset/bright-checkpoints-graph.md | 5 ++ .../task-header-with-todos-chromium-linux.png | 4 +- .../tests/unit/session-queue.test.ts | 31 ++++++++++++ .../src/components/chat/TaskHeader.tsx | 6 +-- .../src/components/chat/TaskTimeline.tsx | 2 +- .../webview-ui/src/context/session-queue.ts | 6 +++ .../webview-ui/src/context/session.tsx | 50 ++++++++++++++----- .../webview-ui/src/stories/StoryProviders.tsx | 1 + 8 files changed, 87 insertions(+), 18 deletions(-) create mode 100644 .changeset/bright-checkpoints-graph.md diff --git a/.changeset/bright-checkpoints-graph.md b/.changeset/bright-checkpoints-graph.md new file mode 100644 index 00000000000..c9812a3ca69 --- /dev/null +++ b/.changeset/bright-checkpoints-graph.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Keep VS Code session timeline, token, context, and cost stats in sync after restoring a checkpoint. diff --git a/packages/kilo-docs/public/img/screenshot-tests/kilo-vscode/visual-regression/chat/task-header-with-todos-chromium-linux.png b/packages/kilo-docs/public/img/screenshot-tests/kilo-vscode/visual-regression/chat/task-header-with-todos-chromium-linux.png index f1e6bc8a35b..da291495f58 100644 --- a/packages/kilo-docs/public/img/screenshot-tests/kilo-vscode/visual-regression/chat/task-header-with-todos-chromium-linux.png +++ b/packages/kilo-docs/public/img/screenshot-tests/kilo-vscode/visual-regression/chat/task-header-with-todos-chromium-linux.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:d34153690f8aa4ef7e06e12385a184c151843f3ae75a9955619a6d7e88dc8c8b -size 7981 +oid sha256:22cfec1cfdfe2deb4299f62bd677eeee34b522c581504297fd356efbf26e9b44 +size 6351 diff --git a/packages/kilo-vscode/tests/unit/session-queue.test.ts b/packages/kilo-vscode/tests/unit/session-queue.test.ts index 9234a499aa5..5d15ac66282 100644 --- a/packages/kilo-vscode/tests/unit/session-queue.test.ts +++ b/packages/kilo-vscode/tests/unit/session-queue.test.ts @@ -4,6 +4,7 @@ import { messageTurns, queuedUserMessageIDs, stableMessageTurns, + visibleMessages, } from "../../webview-ui/src/context/session-queue" import type { Message } from "../../webview-ui/src/types/messages" @@ -107,6 +108,36 @@ describe("messageTurns", () => { { id: "message_6", partial: undefined, assistant: [] }, ]) }) + + it("stops at the revert boundary user turn", () => { + const messages = [ + user("message_1"), + assistant("message_2", "message_1"), + user("message_3"), + assistant("message_4", "message_3"), + ] + + expect(messageTurns(messages, "message_3").map((turn) => turn.user.id)).toEqual(["message_1"]) + }) +}) + +describe("visibleMessages", () => { + it("flattens only turns before the revert boundary", () => { + const messages = [ + user("message_1"), + assistant("message_2", "message_1"), + user("message_3"), + assistant("message_4", "message_3"), + ] + + expect(visibleMessages(messages, "message_3").map((msg) => msg.id)).toEqual(["message_1", "message_2"]) + }) + + it("keeps leading partial assistant output", () => { + const messages = [assistant("message_2", "message_1"), user("message_3")] + + expect(visibleMessages(messages).map((msg) => msg.id)).toEqual(["message_2", "message_3"]) + }) }) describe("stableMessageTurns", () => { diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/TaskHeader.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/TaskHeader.tsx index e04f51dc7d3..f15ba684ae8 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/TaskHeader.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/TaskHeader.tsx @@ -32,7 +32,7 @@ export const TaskHeader: Component = (props) => { const title = createMemo(() => session.currentSession()?.title ?? language.t("command.session.new")) const hasMessages = createMemo(() => session.messages().length > 0) const busy = createMemo(() => session.status() === "busy") - const canCompact = createMemo(() => !busy() && hasMessages() && !!session.selected()) + const canCompact = createMemo(() => !busy() && session.visibleMessages().length > 0 && !!session.selected()) const fmt = (n: number) => new Intl.NumberFormat(language.locale(), { style: "currency", currency: "USD" }).format(n) @@ -67,7 +67,7 @@ export const TaskHeader: Component = (props) => { // Token breakdown from the last assistant message — only return if at least one value is > 0 const tokens = createMemo(() => { - const msgs = session.messages() + const msgs = session.visibleMessages() for (let i = msgs.length - 1; i >= 0; i--) { const m = msgs[i] if (m.role !== "assistant" || !m.tokens) continue @@ -79,7 +79,7 @@ export const TaskHeader: Component = (props) => { }) const hasTimeline = createMemo(() => { - for (const m of session.messages()) { + for (const m of session.visibleMessages()) { if (m.role !== "assistant") continue if (session.getParts(m.id).some((p) => p.type !== "step-start")) return true } diff --git a/packages/kilo-vscode/webview-ui/src/components/chat/TaskTimeline.tsx b/packages/kilo-vscode/webview-ui/src/components/chat/TaskTimeline.tsx index bde1a2082b8..d1c5b3eb109 100644 --- a/packages/kilo-vscode/webview-ui/src/components/chat/TaskTimeline.tsx +++ b/packages/kilo-vscode/webview-ui/src/components/chat/TaskTimeline.tsx @@ -57,7 +57,7 @@ export const TaskTimeline: Component = () => { let startX = 0 let startScroll = 0 - const messages = () => session.messages() + const messages = () => session.visibleMessages() const allParts = () => { const msgs = messages() const result: Record = {} diff --git a/packages/kilo-vscode/webview-ui/src/context/session-queue.ts b/packages/kilo-vscode/webview-ui/src/context/session-queue.ts index 6d8fff78f42..2ec9079ffff 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session-queue.ts +++ b/packages/kilo-vscode/webview-ui/src/context/session-queue.ts @@ -74,6 +74,12 @@ export function messageTurns(messages: Message[], boundary?: string): MessageTur return [...partials(lead), ...result] } +export function visibleMessages(messages: Message[], boundary?: string): Message[] { + return messageTurns(messages, boundary).flatMap((turn) => + turn.partial ? turn.assistant : [turn.user, ...turn.assistant], + ) +} + function sameMessages(a: Message[], b: Message[]) { if (a.length !== b.length) return false for (let i = 0; i < a.length; i++) { diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index 48695930192..741e9e8a505 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/packages/kilo-vscode/webview-ui/src/context/session.tsx @@ -51,6 +51,7 @@ import { errorIDs } from "./session-errors" import { PartStash } from "./part-stash" import { getVariant, sessionVariantKeys, transferVariants, variantKey } from "./session-variant-store" import { KILO_AUTO, parseModelString } from "../../../src/shared/provider-model" +import { visibleMessages as filterVisibleMessages } from "./session-queue" const RECENT_LIMIT = 5 const MESSAGE_PAGE_LIMIT = 80 @@ -109,6 +110,9 @@ interface SessionContextValue { // Messages for current session messages: Accessor + // Messages for current session with soft-reverted turns hidden + visibleMessages: Accessor + // User messages for current session (role === "user") userMessages: Accessor @@ -1401,13 +1405,12 @@ export const SessionProvider: ParentComponent = (props) => { * family tree (self + subagents + sub-subagents). Reads directly from the * store so it's reactive — automatically updates when new parts arrive. */ - function sessionFamily(rootID: string): Set { - const family = new Set([rootID]) + function sessionIDs(rootID: string, source: (sessionID: string) => Message[]): Set { + const ids = new Set([rootID]) const queue = [rootID] while (queue.length > 0) { const sid = queue.pop()! - const msgs = store.messages[sid] - if (!msgs) continue + const msgs = source(sid) for (const msg of msgs) { const parts = store.parts[msg.id] if (!parts) continue @@ -1422,14 +1425,22 @@ export const SessionProvider: ParentComponent = (props) => { state?: { metadata?: { sessionId?: string } } }, ) - if (child && !family.has(child)) { - family.add(child) + if (child && !ids.has(child)) { + ids.add(child) queue.push(child) } } } } - return family + return ids + } + + function sessionFamily(rootID: string): Set { + return sessionIDs(rootID, (sid) => store.messages[sid] ?? []) + } + + function visibleFamily(rootID: string): Set { + return sessionIDs(rootID, visible) } /** Return permissions scoped to the given session's family (self + subagents). */ @@ -2138,12 +2149,21 @@ export const SessionProvider: ParentComponent = (props) => { const userMessages = createMemo(() => messages().filter((m) => m.role === "user")) + function visible(sessionID: string) { + return filterVisibleMessages(store.messages[sessionID] ?? [], store.sessions[sessionID]?.revert?.messageID) + } + const revert = createMemo(() => { const id = currentSessionID() // revert can be null (cleared by unrevert) or undefined (never set) — treat both as "no revert" return id ? (store.sessions[id]?.revert ?? undefined) : undefined }) + const visibleMessages = createMemo(() => { + const id = currentSessionID() + return id ? visible(id) : [] + }) + const revertedCount = createMemo(() => { const boundary = revert()?.messageID if (!boundary) return 0 @@ -2203,16 +2223,21 @@ export const SessionProvider: ParentComponent = (props) => { const familyCosts = createMemo>(() => { const id = currentSessionID() if (!id) return new Map() - const family = sessionFamily(id) - const parents = buildFamilyParents(family, store.messages as any, store.parts as any) - return buildFamilyCosts(family, store.messages, store.sessions, parents) + const family = visibleFamily(id) + const msgs: Record = {} + for (const sid of family) msgs[sid] = visible(sid) + const parents = buildFamilyParents(family, msgs, store.parts) + return buildFamilyCosts(family, msgs, store.sessions, parents) }) /** Child session labels — only reads store.parts (not message costs). */ const familyLabels = createMemo>(() => { const id = currentSessionID() if (!id) return new Map() - return buildFamilyLabels(sessionFamily(id), store.messages as any, store.parts as any) + const family = visibleFamily(id) + const msgs: Record = {} + for (const sid of family) msgs[sid] = visible(sid) + return buildFamilyLabels(family, msgs, store.parts) }) /** Combined cost breakdown with labels. */ @@ -2249,7 +2274,7 @@ export const SessionProvider: ParentComponent = (props) => { }) const contextUsage = createMemo(() => { - const msgs = messages() + const msgs = visibleMessages() for (let i = msgs.length - 1; i >= 0; i--) { const m = msgs[i] if (m.role !== "assistant" || !m.tokens) continue @@ -2277,6 +2302,7 @@ export const SessionProvider: ParentComponent = (props) => { hasOlderMessages, messageMutation, messages, + visibleMessages, userMessages, getParts, isErrorHidden: (messageID: string) => hiddenErrors().has(messageID), diff --git a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx b/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx index 892060c72f8..473ad5dd8ee 100644 --- a/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx +++ b/packages/kilo-vscode/webview-ui/src/stories/StoryProviders.tsx @@ -180,6 +180,7 @@ export function mockSessionValue(overrides?: { hasOlderMessages: () => false, messageMutation: () => undefined, messages: () => [], + visibleMessages: () => [], userMessages: () => [], allMessages: () => ({}), allParts: () => ({}),