From 482a3b19dd984fb08afc6b6fb58cbb600e4b2e07 Mon Sep 17 00:00:00 2001 From: Marius Date: Mon, 11 May 2026 12:19:03 +0200 Subject: [PATCH] fix(vscode): correct subagent cost breakdown after sync (#10134) --- .changeset/fix-session-cost-breakdown.md | 5 +++ .../tests/unit/session-utils.test.ts | 31 +++++++++++++++++++ .../webview-ui/src/context/session-utils.ts | 29 ++++++++++++++++- .../webview-ui/src/context/session.tsx | 7 +++-- 4 files changed, 69 insertions(+), 3 deletions(-) create mode 100644 .changeset/fix-session-cost-breakdown.md diff --git a/.changeset/fix-session-cost-breakdown.md b/.changeset/fix-session-cost-breakdown.md new file mode 100644 index 00000000000..d81c0e55589 --- /dev/null +++ b/.changeset/fix-session-cost-breakdown.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Fix session cost breakdowns after reloading conversations with subagents. diff --git a/packages/kilo-vscode/tests/unit/session-utils.test.ts b/packages/kilo-vscode/tests/unit/session-utils.test.ts index 64454c9e51e..3ffd26095ea 100644 --- a/packages/kilo-vscode/tests/unit/session-utils.test.ts +++ b/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", () => { diff --git a/packages/kilo-vscode/webview-ui/src/context/session-utils.ts b/packages/kilo-vscode/webview-ui/src/context/session-utils.ts index 3506d2d07b9..2544246964a 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session-utils.ts +++ b/packages/kilo-vscode/webview-ui/src/context/session-utils.ts @@ -121,13 +121,14 @@ export function buildFamilyCosts( family: Set, messages: Record>, sessions: Record, + parents: Map = new Map(), ): Map { const totals = new Map() for (const sid of family) totals.set(sid, calcTotalCost(messages[sid] ?? [])) const own = new Map(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, + messages: Record, + parts: Record, +): Map { + const parents = new Map() + 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 /** diff --git a/packages/kilo-vscode/webview-ui/src/context/session.tsx b/packages/kilo-vscode/webview-ui/src/context/session.tsx index 2da1f032c68..48695930192 100644 --- a/packages/kilo-vscode/webview-ui/src/context/session.tsx +++ b/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>(() => { 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). */