mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
fix: implement persistent multi-model cost tracking (#7755)
- Add CostLedger class with Write-Ahead Logging (WAL) for crash-safe persistence - Integrate cost tracking into Task streaming loop to capture costs after each API call - Preserve cost ledger instance when switching models in ClineProvider - Update UI components to display ledger data instead of recalculated metrics - Add comprehensive unit tests for CostLedger functionality This ensures accurate cumulative cost tracking across model switches, fixing the issue where costs were incorrectly recalculated based on the current model pricing instead of preserving actual historical costs.
This commit is contained in:
parent
247da38b02
commit
2b5146611e
6 changed files with 909 additions and 5 deletions
335
src/core/cost-ledger/CostLedger.ts
Normal file
335
src/core/cost-ledger/CostLedger.ts
Normal file
|
|
@ -0,0 +1,335 @@
|
|||
import * as fs from "fs/promises"
|
||||
import * as path from "path"
|
||||
import { v4 as uuidv4 } from "uuid"
|
||||
import { safeWriteJson } from "../../utils/safeWriteJson"
|
||||
|
||||
/**
|
||||
* Represents a single cost entry in the ledger
|
||||
*/
|
||||
export interface CostEntry {
|
||||
entry_id: string
|
||||
task_id: string
|
||||
origin_task_id?: string
|
||||
root_task_id?: string
|
||||
provider: string
|
||||
model_id: string
|
||||
feature: string
|
||||
tokens_in: number
|
||||
tokens_out: number
|
||||
cache_writes?: number
|
||||
cache_reads?: number
|
||||
cost: number
|
||||
timestamp: string
|
||||
}
|
||||
|
||||
/**
|
||||
* Model breakdown for cost reporting
|
||||
*/
|
||||
export interface ModelCostBreakdown {
|
||||
provider: string
|
||||
model_id: string
|
||||
total_cost: number
|
||||
total_tokens_in: number
|
||||
total_tokens_out: number
|
||||
total_cache_writes: number
|
||||
total_cache_reads: number
|
||||
entry_count: number
|
||||
}
|
||||
|
||||
/**
|
||||
* CostLedger manages persistent cost tracking across model switches
|
||||
* Uses Write-Ahead Logging (WAL) for crash safety
|
||||
*/
|
||||
export class CostLedger {
|
||||
private entries: CostEntry[] = []
|
||||
private walPath: string
|
||||
private snapshotPath: string
|
||||
private walFileHandle: fs.FileHandle | null = null
|
||||
private snapshotInterval = 100 // Snapshot every 100 entries
|
||||
private isInitialized = false
|
||||
|
||||
constructor(private storagePath: string) {
|
||||
this.walPath = path.join(storagePath, "cost-ledger-wal.jsonl")
|
||||
this.snapshotPath = path.join(storagePath, "cost-ledger.json")
|
||||
}
|
||||
|
||||
/**
|
||||
* Initialize the ledger by loading existing data
|
||||
*/
|
||||
async initialize(): Promise<void> {
|
||||
if (this.isInitialized) {
|
||||
return
|
||||
}
|
||||
|
||||
try {
|
||||
// Ensure storage directory exists
|
||||
await fs.mkdir(this.storagePath, { recursive: true })
|
||||
|
||||
// Load snapshot if exists
|
||||
await this.loadSnapshot()
|
||||
|
||||
// Replay WAL entries after snapshot
|
||||
await this.replayWAL()
|
||||
|
||||
// Open WAL file for appending
|
||||
try {
|
||||
this.walFileHandle = await fs.open(this.walPath, "a")
|
||||
} catch (error: any) {
|
||||
// If file doesn't exist, create it
|
||||
if (error.code === "ENOENT") {
|
||||
await fs.writeFile(this.walPath, "")
|
||||
this.walFileHandle = await fs.open(this.walPath, "a")
|
||||
} else {
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
this.isInitialized = true
|
||||
} catch (error) {
|
||||
console.error("Failed to initialize CostLedger:", error)
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Append a new cost entry to the ledger
|
||||
*/
|
||||
async appendEntry(params: {
|
||||
task_id: string
|
||||
origin_task_id?: string
|
||||
root_task_id?: string
|
||||
provider: string
|
||||
model_id: string
|
||||
feature: string
|
||||
tokens_in: number
|
||||
tokens_out: number
|
||||
cache_writes?: number
|
||||
cache_reads?: number
|
||||
cost: number
|
||||
}): Promise<void> {
|
||||
if (!this.isInitialized) {
|
||||
await this.initialize()
|
||||
}
|
||||
|
||||
const entry: CostEntry = {
|
||||
entry_id: uuidv4(),
|
||||
timestamp: new Date().toISOString(),
|
||||
...params,
|
||||
}
|
||||
|
||||
// Append to WAL first (for durability)
|
||||
await this.appendToWAL(entry)
|
||||
|
||||
// Add to in-memory entries
|
||||
this.entries.push(entry)
|
||||
|
||||
// Check if we need to create a snapshot
|
||||
if (this.entries.length % this.snapshotInterval === 0) {
|
||||
await this.createSnapshot()
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get cumulative total cost across all models
|
||||
*/
|
||||
getCumulativeTotal(): number {
|
||||
return this.entries.reduce((total, entry) => total + entry.cost, 0)
|
||||
}
|
||||
|
||||
/**
|
||||
* Get breakdown of costs by model
|
||||
*/
|
||||
getBreakdownByModel(): Record<
|
||||
string,
|
||||
{
|
||||
provider: string
|
||||
tokens_in: number
|
||||
tokens_out: number
|
||||
cache_writes: number
|
||||
cache_reads: number
|
||||
cost: number
|
||||
count: number
|
||||
}
|
||||
> {
|
||||
const breakdown: Record<string, any> = {}
|
||||
|
||||
for (const entry of this.entries) {
|
||||
const key = entry.model_id
|
||||
if (!breakdown[key]) {
|
||||
breakdown[key] = {
|
||||
provider: entry.provider,
|
||||
tokens_in: 0,
|
||||
tokens_out: 0,
|
||||
cache_writes: 0,
|
||||
cache_reads: 0,
|
||||
cost: 0,
|
||||
count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
breakdown[key].tokens_in += entry.tokens_in
|
||||
breakdown[key].tokens_out += entry.tokens_out
|
||||
breakdown[key].cache_writes += entry.cache_writes || 0
|
||||
breakdown[key].cache_reads += entry.cache_reads || 0
|
||||
breakdown[key].cost += entry.cost
|
||||
breakdown[key].count += 1
|
||||
}
|
||||
|
||||
return breakdown
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all entries for a specific task
|
||||
*/
|
||||
getEntriesForTask(taskId: string): CostEntry[] {
|
||||
return this.entries.filter(
|
||||
(entry) => entry.task_id === taskId || entry.origin_task_id === taskId || entry.root_task_id === taskId,
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Get total metrics (for UI display)
|
||||
*/
|
||||
getTotalMetrics(): {
|
||||
totalTokensIn: number
|
||||
totalTokensOut: number
|
||||
totalCacheWrites: number
|
||||
totalCacheReads: number
|
||||
totalCost: number
|
||||
} {
|
||||
const metrics = {
|
||||
totalTokensIn: 0,
|
||||
totalTokensOut: 0,
|
||||
totalCacheWrites: 0,
|
||||
totalCacheReads: 0,
|
||||
totalCost: 0,
|
||||
}
|
||||
|
||||
for (const entry of this.entries) {
|
||||
metrics.totalTokensIn += entry.tokens_in
|
||||
metrics.totalTokensOut += entry.tokens_out
|
||||
metrics.totalCacheWrites += entry.cache_writes || 0
|
||||
metrics.totalCacheReads += entry.cache_reads || 0
|
||||
metrics.totalCost += entry.cost
|
||||
}
|
||||
|
||||
return metrics
|
||||
}
|
||||
|
||||
/**
|
||||
* Clear the ledger (for new tasks)
|
||||
*/
|
||||
async clear(): Promise<void> {
|
||||
this.entries = []
|
||||
|
||||
// Close and truncate WAL
|
||||
if (this.walFileHandle) {
|
||||
await this.walFileHandle.close()
|
||||
}
|
||||
await fs.writeFile(this.walPath, "")
|
||||
this.walFileHandle = await fs.open(this.walPath, "a")
|
||||
|
||||
// Remove snapshot
|
||||
try {
|
||||
await fs.unlink(this.snapshotPath)
|
||||
} catch (error) {
|
||||
// Ignore if file doesn't exist
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Close the ledger (cleanup)
|
||||
*/
|
||||
async close(): Promise<void> {
|
||||
// Save a final snapshot before closing
|
||||
if (this.entries.length > 0) {
|
||||
await this.createSnapshot()
|
||||
}
|
||||
|
||||
if (this.walFileHandle) {
|
||||
await this.walFileHandle.close()
|
||||
this.walFileHandle = null
|
||||
}
|
||||
this.isInitialized = false
|
||||
}
|
||||
|
||||
/**
|
||||
* Append entry to WAL file
|
||||
*/
|
||||
private async appendToWAL(entry: CostEntry): Promise<void> {
|
||||
if (!this.walFileHandle) {
|
||||
throw new Error("WAL file handle not initialized")
|
||||
}
|
||||
|
||||
const line = JSON.stringify(entry) + "\n"
|
||||
await this.walFileHandle.write(line)
|
||||
}
|
||||
|
||||
/**
|
||||
* Load snapshot from disk
|
||||
*/
|
||||
private async loadSnapshot(): Promise<void> {
|
||||
try {
|
||||
const data = await fs.readFile(this.snapshotPath, "utf-8")
|
||||
const snapshot = JSON.parse(data)
|
||||
if (Array.isArray(snapshot)) {
|
||||
this.entries = snapshot
|
||||
}
|
||||
} catch (error) {
|
||||
// Snapshot doesn't exist or is corrupted, start fresh
|
||||
this.entries = []
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Replay WAL entries after snapshot
|
||||
*/
|
||||
private async replayWAL(): Promise<void> {
|
||||
try {
|
||||
const walContent = await fs.readFile(this.walPath, "utf-8")
|
||||
const lines = walContent.split("\n").filter((line) => line.trim())
|
||||
|
||||
// Get the last entry ID from snapshot
|
||||
const lastSnapshotEntryId = this.entries.length > 0 ? this.entries[this.entries.length - 1].entry_id : null
|
||||
|
||||
let foundSnapshot = !lastSnapshotEntryId
|
||||
for (const line of lines) {
|
||||
try {
|
||||
const entry = JSON.parse(line) as CostEntry
|
||||
|
||||
// Skip entries until we find the one after snapshot
|
||||
if (!foundSnapshot) {
|
||||
if (entry.entry_id === lastSnapshotEntryId) {
|
||||
foundSnapshot = true
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// Add entries after snapshot
|
||||
if (!this.entries.find((e) => e.entry_id === entry.entry_id)) {
|
||||
this.entries.push(entry)
|
||||
}
|
||||
} catch (error) {
|
||||
// Skip malformed lines
|
||||
console.warn("Skipping malformed WAL entry:", line)
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
// WAL doesn't exist, that's fine
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a snapshot of current entries
|
||||
*/
|
||||
private async createSnapshot(): Promise<void> {
|
||||
await safeWriteJson(this.snapshotPath, this.entries)
|
||||
|
||||
// Truncate WAL after successful snapshot
|
||||
if (this.walFileHandle) {
|
||||
await this.walFileHandle.close()
|
||||
}
|
||||
await fs.writeFile(this.walPath, "")
|
||||
this.walFileHandle = await fs.open(this.walPath, "a")
|
||||
}
|
||||
}
|
||||
480
src/core/cost-ledger/__tests__/CostLedger.test.ts
Normal file
480
src/core/cost-ledger/__tests__/CostLedger.test.ts
Normal file
|
|
@ -0,0 +1,480 @@
|
|||
import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"
|
||||
import * as path from "path"
|
||||
import { CostLedger } from "../CostLedger"
|
||||
import type { CostEntry } from "../CostLedger"
|
||||
|
||||
// Mock fs/promises
|
||||
vi.mock("fs/promises", () => ({
|
||||
mkdir: vi.fn().mockResolvedValue(undefined),
|
||||
readFile: vi.fn().mockRejectedValue(new Error("ENOENT")),
|
||||
writeFile: vi.fn().mockResolvedValue(undefined),
|
||||
appendFile: vi.fn().mockResolvedValue(undefined),
|
||||
unlink: vi.fn().mockResolvedValue(undefined),
|
||||
access: vi.fn().mockResolvedValue(undefined),
|
||||
open: vi.fn().mockResolvedValue({
|
||||
write: vi.fn().mockResolvedValue({ bytesWritten: 0 }),
|
||||
close: vi.fn().mockResolvedValue(undefined),
|
||||
}),
|
||||
}))
|
||||
|
||||
// Mock safeWriteJson
|
||||
vi.mock("../../../utils/safeWriteJson", () => ({
|
||||
safeWriteJson: vi.fn().mockResolvedValue(undefined),
|
||||
}))
|
||||
|
||||
import * as fs from "fs/promises"
|
||||
import { safeWriteJson } from "../../../utils/safeWriteJson"
|
||||
|
||||
describe("CostLedger", () => {
|
||||
let ledger: CostLedger
|
||||
const testDir = "/test/cost-ledger"
|
||||
const walPath = path.join(testDir, "cost-ledger-wal.jsonl")
|
||||
const snapshotPath = path.join(testDir, "cost-ledger.json")
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
// Reset mocks to default behavior
|
||||
vi.mocked(fs.mkdir).mockResolvedValue(undefined)
|
||||
vi.mocked(fs.readFile).mockRejectedValue(new Error("ENOENT"))
|
||||
vi.mocked(fs.writeFile).mockResolvedValue(undefined)
|
||||
vi.mocked(fs.appendFile).mockResolvedValue(undefined)
|
||||
vi.mocked(fs.unlink).mockResolvedValue(undefined)
|
||||
vi.mocked(fs.open).mockResolvedValue({
|
||||
write: vi.fn().mockResolvedValue({ bytesWritten: 0 }),
|
||||
close: vi.fn().mockResolvedValue(undefined),
|
||||
} as any)
|
||||
|
||||
ledger = new CostLedger(testDir)
|
||||
})
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks()
|
||||
})
|
||||
|
||||
describe("appendEntry", () => {
|
||||
it("should append a cost entry and update totals", async () => {
|
||||
await ledger.initialize()
|
||||
|
||||
const entry: Omit<CostEntry, "entry_id"> = {
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "anthropic",
|
||||
model_id: "claude-3-opus",
|
||||
feature: "chat",
|
||||
tokens_in: 100,
|
||||
tokens_out: 50,
|
||||
cache_writes: 10,
|
||||
cache_reads: 5,
|
||||
cost: 0.015,
|
||||
timestamp: new Date().toISOString(),
|
||||
}
|
||||
|
||||
await ledger.appendEntry(entry)
|
||||
|
||||
const metrics = ledger.getTotalMetrics()
|
||||
expect(metrics.totalTokensIn).toBe(100)
|
||||
expect(metrics.totalTokensOut).toBe(50)
|
||||
expect(metrics.totalCacheWrites).toBe(10)
|
||||
expect(metrics.totalCacheReads).toBe(5)
|
||||
expect(metrics.totalCost).toBe(0.015)
|
||||
})
|
||||
|
||||
it("should accumulate multiple entries", async () => {
|
||||
await ledger.initialize()
|
||||
|
||||
const entry1: Omit<CostEntry, "entry_id"> = {
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "anthropic",
|
||||
model_id: "claude-3-opus",
|
||||
feature: "chat",
|
||||
tokens_in: 100,
|
||||
tokens_out: 50,
|
||||
cache_writes: 10,
|
||||
cache_reads: 5,
|
||||
cost: 0.015,
|
||||
timestamp: new Date().toISOString(),
|
||||
}
|
||||
|
||||
const entry2: Omit<CostEntry, "entry_id"> = {
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "openai",
|
||||
model_id: "gpt-4",
|
||||
feature: "chat",
|
||||
tokens_in: 200,
|
||||
tokens_out: 100,
|
||||
cache_writes: 0,
|
||||
cache_reads: 0,
|
||||
cost: 0.025,
|
||||
timestamp: new Date().toISOString(),
|
||||
}
|
||||
|
||||
await ledger.appendEntry(entry1)
|
||||
await ledger.appendEntry(entry2)
|
||||
|
||||
const metrics = ledger.getTotalMetrics()
|
||||
expect(metrics.totalTokensIn).toBe(300)
|
||||
expect(metrics.totalTokensOut).toBe(150)
|
||||
expect(metrics.totalCacheWrites).toBe(10)
|
||||
expect(metrics.totalCacheReads).toBe(5)
|
||||
expect(metrics.totalCost).toBe(0.04)
|
||||
})
|
||||
|
||||
it("should write to WAL file", async () => {
|
||||
await ledger.initialize()
|
||||
|
||||
const entry: Omit<CostEntry, "entry_id"> = {
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "anthropic",
|
||||
model_id: "claude-3-opus",
|
||||
feature: "chat",
|
||||
tokens_in: 100,
|
||||
tokens_out: 50,
|
||||
cache_writes: 10,
|
||||
cache_reads: 5,
|
||||
cost: 0.015,
|
||||
timestamp: new Date().toISOString(),
|
||||
}
|
||||
|
||||
await ledger.appendEntry(entry)
|
||||
|
||||
// Check that the WAL file handle was used
|
||||
const mockFileHandle = await vi.mocked(fs.open).mock.results[0]?.value
|
||||
expect(mockFileHandle?.write).toHaveBeenCalled()
|
||||
})
|
||||
})
|
||||
|
||||
describe("getCumulativeTotal", () => {
|
||||
it("should return cumulative total cost", async () => {
|
||||
await ledger.initialize()
|
||||
|
||||
await ledger.appendEntry({
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "anthropic",
|
||||
model_id: "claude-3-opus",
|
||||
feature: "chat",
|
||||
tokens_in: 100,
|
||||
tokens_out: 50,
|
||||
cache_writes: 10,
|
||||
cache_reads: 5,
|
||||
cost: 0.015,
|
||||
})
|
||||
|
||||
await ledger.appendEntry({
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "openai",
|
||||
model_id: "gpt-4",
|
||||
feature: "chat",
|
||||
tokens_in: 200,
|
||||
tokens_out: 100,
|
||||
cache_writes: 0,
|
||||
cache_reads: 0,
|
||||
cost: 0.025,
|
||||
})
|
||||
|
||||
expect(ledger.getCumulativeTotal()).toBe(0.04)
|
||||
})
|
||||
|
||||
it("should return 0 when no entries", async () => {
|
||||
await ledger.initialize()
|
||||
expect(ledger.getCumulativeTotal()).toBe(0)
|
||||
})
|
||||
})
|
||||
|
||||
describe("getBreakdownByModel", () => {
|
||||
it("should return breakdown by model", async () => {
|
||||
await ledger.initialize()
|
||||
|
||||
await ledger.appendEntry({
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "anthropic",
|
||||
model_id: "claude-3-opus",
|
||||
feature: "chat",
|
||||
tokens_in: 100,
|
||||
tokens_out: 50,
|
||||
cache_writes: 10,
|
||||
cache_reads: 5,
|
||||
cost: 0.015,
|
||||
})
|
||||
|
||||
await ledger.appendEntry({
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "anthropic",
|
||||
model_id: "claude-3-opus",
|
||||
feature: "chat",
|
||||
tokens_in: 50,
|
||||
tokens_out: 25,
|
||||
cache_writes: 5,
|
||||
cache_reads: 2,
|
||||
cost: 0.008,
|
||||
})
|
||||
|
||||
await ledger.appendEntry({
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "openai",
|
||||
model_id: "gpt-4",
|
||||
feature: "chat",
|
||||
tokens_in: 200,
|
||||
tokens_out: 100,
|
||||
cache_writes: 0,
|
||||
cache_reads: 0,
|
||||
cost: 0.025,
|
||||
})
|
||||
|
||||
const breakdown = ledger.getBreakdownByModel()
|
||||
|
||||
expect((breakdown as any)["claude-3-opus"]).toEqual({
|
||||
provider: "anthropic",
|
||||
tokens_in: 150,
|
||||
tokens_out: 75,
|
||||
cache_writes: 15,
|
||||
cache_reads: 7,
|
||||
cost: 0.023,
|
||||
count: 2,
|
||||
})
|
||||
|
||||
expect((breakdown as any)["gpt-4"]).toEqual({
|
||||
provider: "openai",
|
||||
tokens_in: 200,
|
||||
tokens_out: 100,
|
||||
cache_writes: 0,
|
||||
cache_reads: 0,
|
||||
cost: 0.025,
|
||||
count: 1,
|
||||
})
|
||||
})
|
||||
|
||||
it("should return empty object when no entries", async () => {
|
||||
await ledger.initialize()
|
||||
const breakdown = ledger.getBreakdownByModel()
|
||||
expect(breakdown).toEqual({})
|
||||
})
|
||||
})
|
||||
|
||||
describe("getTotalMetrics", () => {
|
||||
it("should return total metrics", async () => {
|
||||
await ledger.initialize()
|
||||
|
||||
await ledger.appendEntry({
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "anthropic",
|
||||
model_id: "claude-3-opus",
|
||||
feature: "chat",
|
||||
tokens_in: 100,
|
||||
tokens_out: 50,
|
||||
cache_writes: 10,
|
||||
cache_reads: 5,
|
||||
cost: 0.015,
|
||||
})
|
||||
|
||||
await ledger.appendEntry({
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "openai",
|
||||
model_id: "gpt-4",
|
||||
feature: "chat",
|
||||
tokens_in: 200,
|
||||
tokens_out: 100,
|
||||
cache_writes: 20,
|
||||
cache_reads: 10,
|
||||
cost: 0.025,
|
||||
})
|
||||
|
||||
const metrics = ledger.getTotalMetrics()
|
||||
|
||||
expect(metrics).toEqual({
|
||||
totalTokensIn: 300,
|
||||
totalTokensOut: 150,
|
||||
totalCacheWrites: 30,
|
||||
totalCacheReads: 15,
|
||||
totalCost: 0.04,
|
||||
})
|
||||
})
|
||||
|
||||
it("should return zeros when no entries", async () => {
|
||||
await ledger.initialize()
|
||||
const metrics = ledger.getTotalMetrics()
|
||||
|
||||
expect(metrics).toEqual({
|
||||
totalTokensIn: 0,
|
||||
totalTokensOut: 0,
|
||||
totalCacheWrites: 0,
|
||||
totalCacheReads: 0,
|
||||
totalCost: 0,
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
describe("persistence", () => {
|
||||
it("should load from snapshot on initialization", async () => {
|
||||
const existingData = [
|
||||
{
|
||||
entry_id: "entry-1",
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "anthropic",
|
||||
model_id: "claude-3-opus",
|
||||
feature: "chat",
|
||||
tokens_in: 100,
|
||||
tokens_out: 50,
|
||||
cache_writes: 10,
|
||||
cache_reads: 5,
|
||||
cost: 0.015,
|
||||
timestamp: new Date().toISOString(),
|
||||
},
|
||||
]
|
||||
|
||||
vi.mocked(fs.readFile).mockImplementation((filePath) => {
|
||||
if (filePath === snapshotPath) {
|
||||
return Promise.resolve(JSON.stringify(existingData))
|
||||
}
|
||||
return Promise.reject(new Error("ENOENT"))
|
||||
})
|
||||
|
||||
const newLedger = new CostLedger(testDir)
|
||||
await newLedger.initialize()
|
||||
|
||||
const metrics = newLedger.getTotalMetrics()
|
||||
expect(metrics.totalTokensIn).toBe(100)
|
||||
expect(metrics.totalTokensOut).toBe(50)
|
||||
expect(metrics.totalCost).toBe(0.015)
|
||||
})
|
||||
|
||||
it("should recover from WAL if snapshot is missing", async () => {
|
||||
const walEntries = [
|
||||
JSON.stringify({
|
||||
entry_id: "entry-1",
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "anthropic",
|
||||
model_id: "claude-3-opus",
|
||||
feature: "chat",
|
||||
tokens_in: 100,
|
||||
tokens_out: 50,
|
||||
cache_writes: 10,
|
||||
cache_reads: 5,
|
||||
cost: 0.015,
|
||||
timestamp: new Date().toISOString(),
|
||||
}),
|
||||
JSON.stringify({
|
||||
entry_id: "entry-2",
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "openai",
|
||||
model_id: "gpt-4",
|
||||
feature: "chat",
|
||||
tokens_in: 200,
|
||||
tokens_out: 100,
|
||||
cache_writes: 0,
|
||||
cache_reads: 10,
|
||||
cost: 0.025,
|
||||
timestamp: new Date().toISOString(),
|
||||
}),
|
||||
].join("\n")
|
||||
|
||||
vi.mocked(fs.readFile).mockImplementation((filePath) => {
|
||||
if (filePath === walPath) {
|
||||
return Promise.resolve(walEntries)
|
||||
}
|
||||
return Promise.reject(new Error("ENOENT"))
|
||||
})
|
||||
|
||||
const newLedger = new CostLedger(testDir)
|
||||
await newLedger.initialize()
|
||||
|
||||
const metrics = newLedger.getTotalMetrics()
|
||||
expect(metrics.totalTokensIn).toBe(300)
|
||||
expect(metrics.totalTokensOut).toBe(150)
|
||||
expect(metrics.totalCacheReads).toBe(15)
|
||||
expect(metrics.totalCost).toBe(0.04)
|
||||
})
|
||||
|
||||
it("should create snapshot after 100 entries", async () => {
|
||||
await ledger.initialize()
|
||||
|
||||
// Add 101 entries to trigger snapshot
|
||||
for (let i = 0; i < 101; i++) {
|
||||
await ledger.appendEntry({
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "anthropic",
|
||||
model_id: "claude-3-opus",
|
||||
feature: "chat",
|
||||
tokens_in: 10,
|
||||
tokens_out: 5,
|
||||
cache_writes: 1,
|
||||
cache_reads: 1,
|
||||
cost: 0.001,
|
||||
})
|
||||
}
|
||||
|
||||
// Wait for async operations
|
||||
await new Promise((resolve) => setTimeout(resolve, 50))
|
||||
|
||||
// Should have created a snapshot
|
||||
expect(safeWriteJson).toHaveBeenCalledWith(
|
||||
snapshotPath,
|
||||
expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
task_id: "task-123",
|
||||
}),
|
||||
]),
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("close", () => {
|
||||
it("should save snapshot when closing", async () => {
|
||||
await ledger.initialize()
|
||||
|
||||
await ledger.appendEntry({
|
||||
task_id: "task-123",
|
||||
origin_task_id: "task-123",
|
||||
root_task_id: "task-123",
|
||||
provider: "anthropic",
|
||||
model_id: "claude-3-opus",
|
||||
feature: "chat",
|
||||
tokens_in: 100,
|
||||
tokens_out: 50,
|
||||
cache_writes: 10,
|
||||
cache_reads: 5,
|
||||
cost: 0.015,
|
||||
})
|
||||
|
||||
await ledger.close()
|
||||
|
||||
// Should have saved a snapshot
|
||||
expect(safeWriteJson).toHaveBeenCalledWith(
|
||||
snapshotPath,
|
||||
expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
task_id: "task-123",
|
||||
}),
|
||||
]),
|
||||
)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
@ -73,6 +73,7 @@ import { TerminalRegistry } from "../../integrations/terminal/TerminalRegistry"
|
|||
// utils
|
||||
import { calculateApiCostAnthropic } from "../../shared/cost"
|
||||
import { getWorkspacePath } from "../../utils/path"
|
||||
import { CostLedger } from "../cost-ledger/CostLedger"
|
||||
|
||||
// prompts
|
||||
import { formatResponse } from "../prompts/responses"
|
||||
|
|
@ -222,6 +223,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
|
|||
api: ApiHandler
|
||||
private static lastGlobalApiRequestTime?: number
|
||||
private autoApprovalHandler: AutoApprovalHandler
|
||||
private costLedger: CostLedger
|
||||
|
||||
/**
|
||||
* Reset the global API request timestamp. This should only be used for testing.
|
||||
|
|
@ -359,6 +361,10 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
|
|||
this.consecutiveMistakeLimit = consecutiveMistakeLimit ?? DEFAULT_CONSECUTIVE_MISTAKE_LIMIT
|
||||
this.providerRef = new WeakRef(provider)
|
||||
this.globalStoragePath = provider.context.globalStorageUri.fsPath
|
||||
|
||||
// Initialize cost ledger for persistent cost tracking (after globalStoragePath is set)
|
||||
const costLedgerPath = path.join(this.globalStoragePath, "cost-ledgers", this.taskId)
|
||||
this.costLedger = new CostLedger(costLedgerPath)
|
||||
this.diffViewProvider = new DiffViewProvider(this.cwd, this)
|
||||
this.enableCheckpoints = enableCheckpoints
|
||||
this.enableBridge = enableBridge
|
||||
|
|
@ -554,6 +560,21 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
|
|||
return this._taskMode
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the cost ledger metrics for this task
|
||||
*/
|
||||
public get costLedgerMetrics() {
|
||||
return (
|
||||
this.costLedger?.getTotalMetrics() || {
|
||||
totalTokensIn: 0,
|
||||
totalTokensOut: 0,
|
||||
totalCacheWrites: 0,
|
||||
totalCacheReads: 0,
|
||||
totalCost: 0,
|
||||
}
|
||||
)
|
||||
}
|
||||
|
||||
static create(options: TaskOptions): [Task, Promise<void>] {
|
||||
const instance = new Task({ ...options, startTask: false })
|
||||
const { images, task, historyItem } = options
|
||||
|
|
@ -1177,6 +1198,8 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
|
|||
// Start / Resume / Abort / Dispose
|
||||
|
||||
private async startTask(task?: string, images?: string[]): Promise<void> {
|
||||
// Initialize the cost ledger for this new task
|
||||
await this.costLedger.initialize()
|
||||
if (this.enableBridge) {
|
||||
try {
|
||||
await BridgeOrchestrator.subscribeToTask(this)
|
||||
|
|
@ -1218,6 +1241,8 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
|
|||
}
|
||||
|
||||
private async resumeTaskFromHistory() {
|
||||
// Initialize the cost ledger for resumed task
|
||||
await this.costLedger.initialize()
|
||||
if (this.enableBridge) {
|
||||
try {
|
||||
await BridgeOrchestrator.subscribeToTask(this)
|
||||
|
|
@ -1506,6 +1531,15 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
|
|||
public dispose(): void {
|
||||
console.log(`[Task#dispose] disposing task ${this.taskId}.${this.instanceId}`)
|
||||
|
||||
// Close the cost ledger
|
||||
try {
|
||||
this.costLedger?.close().catch((error) => {
|
||||
console.error("Error closing cost ledger:", error)
|
||||
})
|
||||
} catch (error) {
|
||||
console.error("Error closing cost ledger:", error)
|
||||
}
|
||||
|
||||
// Dispose message queue and remove event listeners.
|
||||
try {
|
||||
if (this.messageQueueStateChangedHandler) {
|
||||
|
|
@ -1859,6 +1893,40 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
|
|||
cancelReason,
|
||||
streamingFailedMessage,
|
||||
} satisfies ClineApiReqInfo)
|
||||
|
||||
// Track cost in the ledger if we have cost data
|
||||
const finalCost =
|
||||
totalCost ??
|
||||
calculateApiCostAnthropic(
|
||||
this.api.getModel().info,
|
||||
inputTokens,
|
||||
outputTokens,
|
||||
cacheWriteTokens,
|
||||
cacheReadTokens,
|
||||
)
|
||||
|
||||
if (finalCost > 0 && !cancelReason) {
|
||||
// Only track successful API calls in the ledger
|
||||
// Get the current mode from the task
|
||||
const currentMode = this._taskMode || "default"
|
||||
this.costLedger
|
||||
.appendEntry({
|
||||
task_id: this.taskId,
|
||||
origin_task_id: this.parentTaskId,
|
||||
root_task_id: this.rootTaskId,
|
||||
provider: this.apiConfiguration.apiProvider || "unknown",
|
||||
model_id: this.api.getModel().id,
|
||||
feature: currentMode,
|
||||
tokens_in: inputTokens,
|
||||
tokens_out: outputTokens,
|
||||
cache_writes: cacheWriteTokens,
|
||||
cache_reads: cacheReadTokens,
|
||||
cost: finalCost,
|
||||
})
|
||||
.catch((error) => {
|
||||
console.error("Failed to append cost entry to ledger:", error)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
const abortStream = async (cancelReason: ClineApiReqCancelReason, streamingFailedMessage?: string) => {
|
||||
|
|
|
|||
|
|
@ -1247,7 +1247,13 @@ export class ClineProvider
|
|||
const task = this.getCurrentTask()
|
||||
|
||||
if (task) {
|
||||
// Preserve the cost ledger when switching providers
|
||||
const previousLedger = (task as any).costLedger
|
||||
task.api = buildApiHandler(providerSettings)
|
||||
// Ensure the cost ledger is preserved across provider switches
|
||||
if (previousLedger) {
|
||||
;(task as any).costLedger = previousLedger
|
||||
}
|
||||
}
|
||||
} else {
|
||||
await this.updateGlobalState("listApiConfigMeta", await this.providerSettingsManager.listConfig())
|
||||
|
|
@ -1308,7 +1314,13 @@ export class ClineProvider
|
|||
const task = this.getCurrentTask()
|
||||
|
||||
if (task) {
|
||||
// Preserve the cost ledger when switching providers
|
||||
const previousLedger = (task as any).costLedger
|
||||
task.api = buildApiHandler(providerSettings)
|
||||
// Ensure the cost ledger is preserved across provider switches
|
||||
if (previousLedger) {
|
||||
;(task as any).costLedger = previousLedger
|
||||
}
|
||||
}
|
||||
|
||||
await this.postStateToWebview()
|
||||
|
|
@ -1818,6 +1830,7 @@ export class ClineProvider
|
|||
: undefined,
|
||||
clineMessages: this.getCurrentTask()?.clineMessages || [],
|
||||
currentTaskTodos: this.getCurrentTask()?.todoList || [],
|
||||
costLedgerMetrics: this.getCurrentTask()?.costLedgerMetrics,
|
||||
messageQueue: this.getCurrentTask()?.messageQueueService?.messages,
|
||||
taskHistory: (taskHistory || [])
|
||||
.filter((item: HistoryItem) => item.ts && item.task)
|
||||
|
|
|
|||
|
|
@ -286,6 +286,13 @@ export type ExtensionState = Pick<
|
|||
clineMessages: ClineMessage[]
|
||||
currentTaskItem?: HistoryItem
|
||||
currentTaskTodos?: TodoItem[] // Initial todos for the current task
|
||||
costLedgerMetrics?: {
|
||||
totalTokensIn: number
|
||||
totalTokensOut: number
|
||||
totalCacheWrites: number
|
||||
totalCacheReads: number
|
||||
totalCost: number
|
||||
}
|
||||
apiConfiguration: ProviderSettings
|
||||
uriScheme?: string
|
||||
shouldShowAnnouncement: boolean
|
||||
|
|
|
|||
|
|
@ -121,6 +121,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
|
|||
soundVolume,
|
||||
cloudIsAuthenticated,
|
||||
messageQueue = [],
|
||||
costLedgerMetrics,
|
||||
} = useExtensionState()
|
||||
|
||||
const messagesRef = useRef(messages)
|
||||
|
|
@ -1781,11 +1782,11 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
|
|||
<>
|
||||
<TaskHeader
|
||||
task={task}
|
||||
tokensIn={apiMetrics.totalTokensIn}
|
||||
tokensOut={apiMetrics.totalTokensOut}
|
||||
cacheWrites={apiMetrics.totalCacheWrites}
|
||||
cacheReads={apiMetrics.totalCacheReads}
|
||||
totalCost={apiMetrics.totalCost}
|
||||
tokensIn={costLedgerMetrics?.totalTokensIn || apiMetrics.totalTokensIn}
|
||||
tokensOut={costLedgerMetrics?.totalTokensOut || apiMetrics.totalTokensOut}
|
||||
cacheWrites={costLedgerMetrics?.totalCacheWrites || apiMetrics.totalCacheWrites}
|
||||
cacheReads={costLedgerMetrics?.totalCacheReads || apiMetrics.totalCacheReads}
|
||||
totalCost={costLedgerMetrics?.totalCost || apiMetrics.totalCost}
|
||||
contextTokens={apiMetrics.contextTokens}
|
||||
buttonsDisabled={sendingDisabled}
|
||||
handleCondenseContext={handleCondenseContext}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue