fix(vscode): correct subagent cost breakdown after sync (#10134)

.changeset/fix-session-cost-breakdown.md

@@ -0,0 +1,5 @@
+---
+"kilo-code": patch
+---
+
+Fix session cost breakdowns after reloading conversations with subagents.

packages/kilo-vscode/tests/unit/session-utils.test.ts

@@ -4,6 +4,7 @@ import {
   calcTotalCost,
   calcContextUsage,
   buildFamilyCosts,
+  buildFamilyParents,
   buildFamilyLabels,
   buildCostBreakdown,
   collapseCostBreakdown,
@@ -287,6 +288,36 @@ describe("buildFamilyCosts", () => {
     expect(costs.size).toBe(1)
     expect(costs.get("s1")).toBeCloseTo(0.01)
   })
+
+  it("subtracts child totals using task parent links when session metadata is missing", () => {
+    const family = new Set(["root", "child"])
+    const messages = {
+      root: [msg("m1", "assistant", 0.15)],
+      child: [msg("m2", "assistant", 0.1)],
+    }
+    const parts = { m1: [toolPart("task", "child", { subagent_type: "explore" })] }
+    const parents = buildFamilyParents(family, messages, parts)
+    const costs = buildFamilyCosts(family, messages, { root: {} }, parents)
+    expect(costs.get("root")).toBeCloseTo(0.05)
+    expect(costs.get("child")).toBeCloseTo(0.1)
+  })
+})
+
+describe("buildFamilyParents", () => {
+  it("derives child-to-parent links from task tool parts", () => {
+    const family = new Set(["root", "child"])
+    const messages = { root: [msg("m1", "assistant")], child: [msg("m2", "assistant")] }
+    const parts = { m1: [toolPart("task", "child", { subagent_type: "general" })] }
+    const parents = buildFamilyParents(family, messages, parts)
+    expect(parents.get("child")).toBe("root")
+  })
+
+  it("ignores task parts that point outside the family", () => {
+    const family = new Set(["root"])
+    const messages = { root: [msg("m1", "assistant")] }
+    const parts = { m1: [toolPart("task", "orphan", { subagent_type: "general" })] }
+    expect(buildFamilyParents(family, messages, parts).size).toBe(0)
+  })
 })
 
 describe("buildFamilyLabels", () => {

packages/kilo-vscode/webview-ui/src/context/session-utils.ts

@@ -121,13 +121,14 @@ export function buildFamilyCosts(
   family: Set<string>,
   messages: Record<string, Array<{ role: string; cost?: number }>>,
   sessions: Record<string, { parentID?: string | null } | undefined>,
+  parents: Map<string, string> = new Map(),
 ): Map<string, number> {
   const totals = new Map<string, number>()
   for (const sid of family) totals.set(sid, calcTotalCost(messages[sid] ?? []))
 
   const own = new Map<string, number>(totals)
   for (const sid of family) {
-    const parent = sessions[sid]?.parentID
+    const parent = sessions[sid]?.parentID ?? parents.get(sid)
     if (!parent || !own.has(parent)) continue
     own.set(parent, (own.get(parent) ?? 0) - (totals.get(sid) ?? 0))
   }
@@ -139,6 +140,32 @@ export function buildFamilyCosts(
   return costs
 }
 
+/**
+ * Build child session ID -> parent session ID links from task tool metadata.
+ * This fills the gap when child messages are synced before their SessionInfo.
+ */
+export function buildFamilyParents(
+  family: Set<string>,
+  messages: Record<string, CostMessage[]>,
+  parts: Record<string, TaskPart[]>,
+): Map<string, string> {
+  const parents = new Map<string, string>()
+  for (const sid of family) {
+    const msgs = messages[sid]
+    if (!msgs) continue
+    for (const msg of msgs) {
+      const list = parts[msg.id]
+      if (!list) continue
+      for (const p of list) {
+        const child = childID(p)
+        if (!child || !family.has(child) || parents.has(child)) continue
+        parents.set(child, sid)
+      }
+    }
+  }
+  return parents
+}
+
 const LABEL_CAP = 24
 
 /**

packages/kilo-vscode/webview-ui/src/context/session.tsx

@@ -39,6 +39,7 @@ import {
   computeStatus,
   calcContextUsage,
   buildFamilyCosts,
+  buildFamilyParents,
   buildFamilyLabels,
   buildCostBreakdown,
   childID,
@@ -2195,14 +2196,16 @@ export const SessionProvider: ParentComponent = (props) => {
 
   /**
    * Per-session **own cost** — reads `store.messages` for per-session
-   * propagated totals and `store.sessions` for parent links so each
+   * propagated totals and task metadata as a fallback for parent links so each
    * session's entry excludes the cost already propagated up from its
    * descendants by the CLI backend.
    */
   const familyCosts = createMemo<Map<string, number>>(() => {
     const id = currentSessionID()
     if (!id) return new Map()
-    return buildFamilyCosts(sessionFamily(id), store.messages, store.sessions)
+    const family = sessionFamily(id)
+    const parents = buildFamilyParents(family, store.messages as any, store.parts as any)
+    return buildFamilyCosts(family, store.messages, store.sessions, parents)
   })
 
   /** Child session labels — only reads store.parts (not message costs). */