diff --git a/packages/kilo-vscode/tests/unit/session-utils.test.ts b/packages/kilo-vscode/tests/unit/session-utils.test.ts index 3ffd26095ea..e308a576e3d 100644 --- a/packages/kilo-vscode/tests/unit/session-utils.test.ts +++ b/packages/kilo-vscode/tests/unit/session-utils.test.ts @@ -301,6 +301,53 @@ describe("buildFamilyCosts", () => { expect(costs.get("root")).toBeCloseTo(0.05) expect(costs.get("child")).toBeCloseTo(0.1) }) + + it("returns own costs for a nested propagated subagent chain", () => { + const ids = ["root", "a", "b", "c"] + const own = new Map([ + ["root", 1], + ["a", 2], + ["b", 3], + ["c", 4], + ]) + const total = ids.reduce((sum, sid) => sum + own.get(sid)!, 0) + const subtree = (index: number) => ids.slice(index).reduce((sum, sid) => sum + own.get(sid)!, 0) + const messages = Object.fromEntries(ids.map((sid, index) => [sid, [msg(`m-${sid}`, "assistant", subtree(index))]])) + const sessions = Object.fromEntries(ids.map((sid, index) => [sid, index === 0 ? {} : { parentID: ids[index - 1] }])) + + const costs = buildFamilyCosts(new Set(ids), messages, sessions) + const sum = [...costs.values()].reduce((acc, cost) => acc + cost, 0) + const bad = ids.reduce((acc, _sid, index) => acc + subtree(index), 0) + + expect(bad).toBe(30) + expect(costs.get("root")).toBe(1) + expect(costs.get("a")).toBe(2) + expect(costs.get("b")).toBe(3) + expect(costs.get("c")).toBe(4) + expect(sum).toBe(total) + }) + + it("returns own costs for a nested chain using task-derived parent links", () => { + const ids = ["root", "a", "b", "c"] + const own = { root: 1, a: 2, b: 3, c: 4 } + const subtree = (index: number) => ids.slice(index).reduce((sum, sid) => sum + own[sid as keyof typeof own], 0) + const messages = Object.fromEntries(ids.map((sid, index) => [sid, [msg(`m-${sid}`, "assistant", subtree(index))]])) + const parts = Object.fromEntries( + ids + .slice(0, -1) + .map((sid, index) => [`m-${sid}`, [toolPart("task", ids[index + 1], { subagent_type: "explore" })]]), + ) + const parents = buildFamilyParents(new Set(ids), messages, parts) + const costs = buildFamilyCosts(new Set(ids), messages, { root: {} }, parents) + const sum = [...costs.values()].reduce((acc, cost) => acc + cost, 0) + + expect(parents.size).toBe(3) + expect(costs.get("root")).toBe(1) + expect(costs.get("a")).toBe(2) + expect(costs.get("b")).toBe(3) + expect(costs.get("c")).toBe(4) + expect(sum).toBe(10) + }) }) describe("buildFamilyParents", () => { @@ -525,4 +572,17 @@ describe("collapseCostBreakdown", () => { expect(aggregated.label).toBe("12 older sessions") expect(aggregated.cost).toBeCloseTo(0.05 * 12) }) + + it("collapses older nested rows after own-cost correction, not subtree totals", () => { + const items = [ + { label: "This session", cost: 1 }, + ...Array.from({ length: 10 }, (_, i) => ({ label: "explore", cost: i + 2 })), + ] + const result = collapseCostBreakdown(items, summary) + const hidden = result.at(-1) + const shown = result.reduce((sum, item) => sum + item.cost, 0) + + expect(hidden).toEqual({ label: "2 older sessions", cost: 2 + 3 }) + expect(shown).toBe(1 + 2 + 3 + 4 + 5 + 6 + 7 + 8 + 9 + 10 + 11) + }) })