Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions packages/types/src/vscode-extension-host.ts
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,7 @@ export interface ExtensionMessage {
| "claudeCodeRateLimits"
| "customToolsResult"
| "modes"
| "taskWithAggregatedCosts"
text?: string
payload?: any // eslint-disable-line @typescript-eslint/no-explicit-any
checkpointWarning?: {
Expand Down Expand Up @@ -182,6 +183,13 @@ export interface ExtensionMessage {
stepIndex?: number // For browserSessionNavigate: the target step index to display
tools?: SerializedCustomToolDefinition[] // For customToolsResult
modes?: { slug: string; name: string }[] // For modes response
aggregatedCosts?: {
// For taskWithAggregatedCosts response
totalCost: number
ownCost: number
childrenCost: number
}
historyItem?: HistoryItem
}

export type ExtensionState = Pick<
Expand Down Expand Up @@ -498,6 +506,7 @@ export interface WebviewMessage {
| "getDismissedUpsells"
| "updateSettings"
| "allowedCommands"
| "getTaskWithAggregatedCosts"
| "deniedCommands"
| "killBrowserSession"
| "openBrowserSessionPanel"
Expand Down
15 changes: 15 additions & 0 deletions src/core/webview/ClineProvider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ import {
DEFAULT_CHECKPOINT_TIMEOUT_SECONDS,
getModelId,
} from "@roo-code/types"
import { aggregateTaskCostsRecursive, type AggregatedCosts } from "./aggregateTaskCosts"
import { TelemetryService } from "@roo-code/telemetry"
import { CloudService, BridgeOrchestrator, getRooCodeApiUrl } from "@roo-code/cloud"

Expand Down Expand Up @@ -1705,6 +1706,20 @@ export class ClineProvider
throw new Error("Task not found")
}

async getTaskWithAggregatedCosts(taskId: string): Promise<{
historyItem: HistoryItem
aggregatedCosts: AggregatedCosts
}> {
const { historyItem } = await this.getTaskWithId(taskId)

const aggregatedCosts = await aggregateTaskCostsRecursive(taskId, async (id: string) => {
const result = await this.getTaskWithId(id)
return result.historyItem
})

return { historyItem, aggregatedCosts }
}

async showTaskWithId(id: string) {
if (id !== this.getCurrentTask()?.taskId) {
// Non-current task.
Expand Down
326 changes: 326 additions & 0 deletions src/core/webview/__tests__/aggregateTaskCosts.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,326 @@
import { describe, it, expect, vi, beforeEach } from "vitest"
import { aggregateTaskCostsRecursive } from "../aggregateTaskCosts.js"
import type { HistoryItem } from "@roo-code/types"

describe("aggregateTaskCostsRecursive", () => {
let consoleWarnSpy: ReturnType<typeof vi.spyOn>

beforeEach(() => {
consoleWarnSpy = vi.spyOn(console, "warn").mockImplementation(() => {})
})

it("should calculate cost for task with no children", async () => {
const mockHistory: Record<string, HistoryItem> = {
"task-1": {
id: "task-1",
totalCost: 1.5,
childIds: [],
} as unknown as HistoryItem,
}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("task-1", getTaskHistory)

expect(result.ownCost).toBe(1.5)
expect(result.childrenCost).toBe(0)
expect(result.totalCost).toBe(1.5)
expect(result.childBreakdown).toEqual({})
})

it("should calculate cost for task with undefined childIds", async () => {
const mockHistory: Record<string, HistoryItem> = {
"task-1": {
id: "task-1",
totalCost: 2.0,
// childIds is undefined
} as unknown as HistoryItem,
}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("task-1", getTaskHistory)

expect(result.ownCost).toBe(2.0)
expect(result.childrenCost).toBe(0)
expect(result.totalCost).toBe(2.0)
expect(result.childBreakdown).toEqual({})
})

it("should aggregate parent with one child", async () => {
const mockHistory: Record<string, HistoryItem> = {
parent: {
id: "parent",
totalCost: 1.0,
childIds: ["child-1"],
} as unknown as HistoryItem,
"child-1": {
id: "child-1",
totalCost: 0.5,
childIds: [],
} as unknown as HistoryItem,
}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("parent", getTaskHistory)

expect(result.ownCost).toBe(1.0)
expect(result.childrenCost).toBe(0.5)
expect(result.totalCost).toBe(1.5)
expect(result.childBreakdown).toHaveProperty("child-1")
const child1 = result.childBreakdown?.["child-1"]
expect(child1).toBeDefined()
expect(child1!.totalCost).toBe(0.5)
})

it("should aggregate parent with multiple children", async () => {
const mockHistory: Record<string, HistoryItem> = {
parent: {
id: "parent",
totalCost: 1.0,
childIds: ["child-1", "child-2", "child-3"],
} as unknown as HistoryItem,
"child-1": {
id: "child-1",
totalCost: 0.5,
childIds: [],
} as unknown as HistoryItem,
"child-2": {
id: "child-2",
totalCost: 0.75,
childIds: [],
} as unknown as HistoryItem,
"child-3": {
id: "child-3",
totalCost: 0.25,
childIds: [],
} as unknown as HistoryItem,
}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("parent", getTaskHistory)

expect(result.ownCost).toBe(1.0)
expect(result.childrenCost).toBe(1.5) // 0.5 + 0.75 + 0.25
expect(result.totalCost).toBe(2.5)
expect(Object.keys(result.childBreakdown || {})).toHaveLength(3)
})

it("should recursively aggregate multi-level hierarchy", async () => {
const mockHistory: Record<string, HistoryItem> = {
parent: {
id: "parent",
totalCost: 1.0,
childIds: ["child"],
} as unknown as HistoryItem,
child: {
id: "child",
totalCost: 0.5,
childIds: ["grandchild"],
} as unknown as HistoryItem,
grandchild: {
id: "grandchild",
totalCost: 0.25,
childIds: [],
} as unknown as HistoryItem,
}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("parent", getTaskHistory)

expect(result.ownCost).toBe(1.0)
expect(result.childrenCost).toBe(0.75) // child (0.5) + grandchild (0.25)
expect(result.totalCost).toBe(1.75)

// Verify child breakdown
const child = result.childBreakdown?.["child"]
expect(child).toBeDefined()
expect(child!.ownCost).toBe(0.5)
expect(child!.childrenCost).toBe(0.25)
expect(child!.totalCost).toBe(0.75)

// Verify grandchild breakdown
const grandchild = child!.childBreakdown?.["grandchild"]
expect(grandchild).toBeDefined()
expect(grandchild!.ownCost).toBe(0.25)
expect(grandchild!.childrenCost).toBe(0)
expect(grandchild!.totalCost).toBe(0.25)
})

it("should detect and prevent circular references", async () => {
const mockHistory: Record<string, HistoryItem> = {
"task-a": {
id: "task-a",
totalCost: 1.0,
childIds: ["task-b"],
} as unknown as HistoryItem,
"task-b": {
id: "task-b",
totalCost: 0.5,
childIds: ["task-a"], // Circular reference back to task-a
} as unknown as HistoryItem,
}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("task-a", getTaskHistory)

// Should still process task-b but ignore the circular reference
expect(result.ownCost).toBe(1.0)
expect(result.childrenCost).toBe(0.5) // Only task-b's own cost, circular ref returns 0
expect(result.totalCost).toBe(1.5)

// Verify warning was logged
expect(consoleWarnSpy).toHaveBeenCalledWith(expect.stringContaining("Circular reference detected: task-a"))
})

it("should handle missing task gracefully", async () => {
const mockHistory: Record<string, HistoryItem> = {
parent: {
id: "parent",
totalCost: 1.0,
childIds: ["nonexistent-child"],
} as unknown as HistoryItem,
}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("parent", getTaskHistory)

expect(result.ownCost).toBe(1.0)
expect(result.childrenCost).toBe(0) // Missing child contributes 0
expect(result.totalCost).toBe(1.0)

// Verify warning was logged
expect(consoleWarnSpy).toHaveBeenCalledWith(expect.stringContaining("Task nonexistent-child not found"))
})

it("should return zero costs for completely missing task", async () => {
const mockHistory: Record<string, HistoryItem> = {}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("nonexistent", getTaskHistory)

expect(result.ownCost).toBe(0)
expect(result.childrenCost).toBe(0)
expect(result.totalCost).toBe(0)

expect(consoleWarnSpy).toHaveBeenCalledWith(expect.stringContaining("Task nonexistent not found"))
})

it("should handle task with null totalCost", async () => {
const mockHistory: Record<string, HistoryItem> = {
"task-1": {
id: "task-1",
totalCost: null as unknown as number, // Explicitly null (invalid type in prod)
childIds: [],
} as unknown as HistoryItem,
}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("task-1", getTaskHistory)

expect(result.ownCost).toBe(0)
expect(result.childrenCost).toBe(0)
expect(result.totalCost).toBe(0)
})

it("should handle task with undefined totalCost", async () => {
const mockHistory: Record<string, HistoryItem> = {
"task-1": {
id: "task-1",
// totalCost is undefined
childIds: [],
} as unknown as HistoryItem,
}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("task-1", getTaskHistory)

expect(result.ownCost).toBe(0)
expect(result.childrenCost).toBe(0)
expect(result.totalCost).toBe(0)
})

it("should handle complex hierarchy with mixed costs", async () => {
const mockHistory: Record<string, HistoryItem> = {
root: {
id: "root",
totalCost: 2.5,
childIds: ["child-1", "child-2"],
} as unknown as HistoryItem,
"child-1": {
id: "child-1",
totalCost: 1.2,
childIds: ["grandchild-1", "grandchild-2"],
} as unknown as HistoryItem,
"child-2": {
id: "child-2",
totalCost: 0.8,
childIds: [],
} as unknown as HistoryItem,
"grandchild-1": {
id: "grandchild-1",
totalCost: 0.3,
childIds: [],
} as unknown as HistoryItem,
"grandchild-2": {
id: "grandchild-2",
totalCost: 0.15,
childIds: [],
} as unknown as HistoryItem,
}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("root", getTaskHistory)

expect(result.ownCost).toBe(2.5)
// child-1: 1.2 + 0.3 + 0.15 = 1.65
// child-2: 0.8
// Total children: 2.45
expect(result.childrenCost).toBe(2.45)
expect(result.totalCost).toBe(4.95) // 2.5 + 2.45
})

it("should handle siblings without cross-contamination", async () => {
const mockHistory: Record<string, HistoryItem> = {
parent: {
id: "parent",
totalCost: 1.0,
childIds: ["sibling-1", "sibling-2"],
} as unknown as HistoryItem,
"sibling-1": {
id: "sibling-1",
totalCost: 0.5,
childIds: ["nephew"],
} as unknown as HistoryItem,
"sibling-2": {
id: "sibling-2",
totalCost: 0.3,
childIds: ["nephew"], // Same child ID as sibling-1
} as unknown as HistoryItem,
nephew: {
id: "nephew",
totalCost: 0.1,
childIds: [],
} as unknown as HistoryItem,
}

const getTaskHistory = vi.fn(async (id: string) => mockHistory[id])

const result = await aggregateTaskCostsRecursive("parent", getTaskHistory)

// Both siblings should independently count nephew
// sibling-1: 0.5 + 0.1 = 0.6
// sibling-2: 0.3 + 0.1 = 0.4
// Total: 1.0 + 0.6 + 0.4 = 2.0
expect(result.totalCost).toBe(2.0)
})
})
Loading
Loading