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] <github-actions[bot]@users.noreply.github.com>

.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.

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

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", () => {

packages/kilo-vscode/webview-ui/src/components/chat/TaskHeader.tsx

@@ -32,7 +32,7 @@ export const TaskHeader: Component<TaskHeaderProps> = (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<TaskHeaderProps> = (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<TaskHeaderProps> = (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
     }

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<string, Part[]> = {}

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++) {

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<Message[]>
 
+  // Messages for current session with soft-reverted turns hidden
+  visibleMessages: Accessor<Message[]>
+
   // User messages for current session (role === "user")
   userMessages: Accessor<Message[]>
 
@@ -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<string> {
-    const family = new Set<string>([rootID])
+  function sessionIDs(rootID: string, source: (sessionID: string) => Message[]): Set<string> {
+    const ids = new Set<string>([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<string> {
+    return sessionIDs(rootID, (sid) => store.messages[sid] ?? [])
+  }
+
+  function visibleFamily(rootID: string): Set<string> {
+    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<Map<string, number>>(() => {
     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<string, Message[]> = {}
+    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<Map<string, string>>(() => {
     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<string, Message[]> = {}
+    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<ContextUsage | undefined>(() => {
-    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),

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: () => ({}),