Skip to content

Commit 74b9288

Browse files
committed
Refactor: Extract costHistory calculation to separate utility function
- Created costUtils.ts with calculateCostHistory function - Added comprehensive unit tests for costUtils - Updated ChatView.tsx to use the new utility function - Improves separation of concerns and testability
1 parent 6caa57b commit 74b9288

File tree

3 files changed

+268
-45
lines changed

3 files changed

+268
-45
lines changed

webview-ui/src/components/chat/ChatView.tsx

Lines changed: 2 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import { useTaskSearch } from "../history/useTaskSearch"
3232
import HistoryPreview from "../history/HistoryPreview"
3333
import RooHero from "@src/components/welcome/RooHero"
3434
import RooTips from "@src/components/welcome/RooTips"
35+
import { calculateCostHistory } from "@src/utils/costUtils"
3536
import Announcement from "./Announcement"
3637
import BrowserSessionRow from "./BrowserSessionRow"
3738
import ChatRow from "./ChatRow"
@@ -110,51 +111,7 @@ const ChatViewComponent: React.ForwardRefRenderFunction<ChatViewRef, ChatViewPro
110111
// Has to be after api_req_finished are all reduced into api_req_started messages.
111112
const apiMetrics = useMemo(() => getApiMetrics(modifiedMessages), [modifiedMessages])
112113

113-
// Calculate cost history for the chart, including token and cache data
114-
const costHistory = useMemo(() => {
115-
let cumulativeCost = 0
116-
let requestIndex = 0
117-
// Define the structure for the history array, including optional fields
118-
const history: {
119-
requestIndex: number
120-
cumulativeCost: number
121-
costDelta: number
122-
tokensIn?: number
123-
tokensOut?: number
124-
cacheReads?: number
125-
cacheWrites?: number
126-
}[] = []
127-
128-
modifiedMessages.forEach((message) => {
129-
// Look for messages indicating a completed API request with cost and other metrics
130-
if (message.say === "api_req_started" && message.text) {
131-
try {
132-
const data = JSON.parse(message.text)
133-
// Check if cost is defined and not null (indicating request completion)
134-
if (data.cost !== undefined && data.cost !== null) {
135-
const costDelta = data.cost ?? 0
136-
cumulativeCost += costDelta
137-
requestIndex++
138-
139-
// Push the extended data point
140-
history.push({
141-
requestIndex,
142-
cumulativeCost,
143-
costDelta,
144-
tokensIn: data.tokensIn, // Add tokensIn if available
145-
tokensOut: data.tokensOut, // Add tokensOut if available
146-
cacheReads: data.cacheReads, // Add cacheReads if available
147-
cacheWrites: data.cacheWrites, // Add cacheWrites if available
148-
})
149-
}
150-
} catch (e) {
151-
console.error("Error parsing api_req_started text for cost history:", e, message.text)
152-
// Ignore parsing errors for robustness
153-
}
154-
}
155-
})
156-
return history
157-
}, [modifiedMessages])
114+
const costHistory = useMemo(() => calculateCostHistory(modifiedMessages), [modifiedMessages])
158115

159116
const [inputValue, setInputValue] = useState("")
160117
const textAreaRef = useRef<HTMLTextAreaElement>(null)
Lines changed: 223 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,223 @@
1+
import { ClineMessage } from "@roo/shared/ExtensionMessage"
2+
import { calculateCostHistory, CostHistoryDataPoint } from "../costUtils"
3+
4+
describe("calculateCostHistory", () => {
5+
const addTs = (msg: Omit<ClineMessage, "ts">, index: number): ClineMessage => ({
6+
...msg,
7+
ts: Date.now() + index,
8+
})
9+
10+
it("should return an empty array for empty input", () => {
11+
const messages: ClineMessage[] = []
12+
expect(calculateCostHistory(messages)).toEqual([])
13+
})
14+
15+
it("should return an empty array if no api_req_started messages exist", () => {
16+
const messages: ClineMessage[] = [
17+
addTs({ type: "say", say: "text", text: "Hello" }, 0),
18+
addTs({ type: "ask", ask: "followup", text: "How are you?" }, 1),
19+
]
20+
expect(calculateCostHistory(messages)).toEqual([])
21+
})
22+
23+
it("should return an empty array if api_req_started messages have undefined text", () => {
24+
const messages: ClineMessage[] = [addTs({ type: "say", say: "api_req_started", text: undefined }, 0)]
25+
expect(calculateCostHistory(messages)).toEqual([])
26+
})
27+
28+
it("should return an empty array if api_req_started messages have invalid JSON", () => {
29+
const messages: ClineMessage[] = [addTs({ type: "say", say: "api_req_started", text: "not json" }, 0)]
30+
const consoleErrorSpy = jest.spyOn(console, "error").mockImplementation(() => {})
31+
expect(calculateCostHistory(messages)).toEqual([])
32+
expect(consoleErrorSpy).toHaveBeenCalled()
33+
consoleErrorSpy.mockRestore()
34+
})
35+
36+
it("should return an empty array if api_req_started messages have no cost (incomplete)", () => {
37+
const messages: ClineMessage[] = [
38+
addTs({ type: "say", say: "api_req_started", text: JSON.stringify({ tokensIn: 10 }) }, 0),
39+
]
40+
expect(calculateCostHistory(messages)).toEqual([])
41+
})
42+
43+
it("should return an empty array if api_req_started messages have cost as null (incomplete)", () => {
44+
const messages: ClineMessage[] = [
45+
addTs({ type: "say", say: "api_req_started", text: JSON.stringify({ cost: null, tokensIn: 10 }) }, 0),
46+
]
47+
expect(calculateCostHistory(messages)).toEqual([])
48+
})
49+
50+
it("should correctly calculate history for a single completed request", () => {
51+
const messages: ClineMessage[] = [
52+
addTs(
53+
{
54+
type: "say",
55+
say: "api_req_started",
56+
text: JSON.stringify({ cost: 0.01, tokensIn: 100, tokensOut: 50 }),
57+
},
58+
0,
59+
),
60+
]
61+
const expected: CostHistoryDataPoint[] = [
62+
{
63+
requestIndex: 1,
64+
cumulativeCost: 0.01,
65+
costDelta: 0.01,
66+
tokensIn: 100,
67+
tokensOut: 50,
68+
cacheReads: undefined,
69+
cacheWrites: undefined,
70+
},
71+
]
72+
expect(calculateCostHistory(messages)).toEqual(expected)
73+
})
74+
75+
it("should correctly calculate history for a single completed request with cost 0", () => {
76+
const messages: ClineMessage[] = [
77+
addTs(
78+
{ type: "say", say: "api_req_started", text: JSON.stringify({ cost: 0, tokensIn: 10, tokensOut: 5 }) },
79+
0,
80+
),
81+
]
82+
const expected: CostHistoryDataPoint[] = [
83+
{
84+
requestIndex: 1,
85+
cumulativeCost: 0,
86+
costDelta: 0,
87+
tokensIn: 10,
88+
tokensOut: 5,
89+
cacheReads: undefined,
90+
cacheWrites: undefined,
91+
},
92+
]
93+
expect(calculateCostHistory(messages)).toEqual(expected)
94+
})
95+
96+
it("should correctly calculate history for multiple completed requests", () => {
97+
const messages: ClineMessage[] = [
98+
addTs(
99+
{
100+
type: "say",
101+
say: "api_req_started",
102+
text: JSON.stringify({ cost: 0.01, tokensIn: 100, tokensOut: 50 }),
103+
},
104+
0,
105+
),
106+
addTs({ type: "say", say: "text", text: "Some response" }, 1),
107+
addTs(
108+
{
109+
type: "say",
110+
say: "api_req_started",
111+
text: JSON.stringify({ cost: 0.025, tokensIn: 200, tokensOut: 150, cacheReads: 1 }),
112+
},
113+
2,
114+
),
115+
addTs(
116+
{
117+
type: "say",
118+
say: "api_req_started",
119+
text: JSON.stringify({ cost: 0.005, tokensIn: 50, tokensOut: 20, cacheWrites: 1 }),
120+
},
121+
3,
122+
),
123+
]
124+
const expected: CostHistoryDataPoint[] = [
125+
{
126+
requestIndex: 1,
127+
cumulativeCost: 0.01,
128+
costDelta: 0.01,
129+
tokensIn: 100,
130+
tokensOut: 50,
131+
cacheReads: undefined,
132+
cacheWrites: undefined,
133+
},
134+
{
135+
requestIndex: 2,
136+
cumulativeCost: 0.035,
137+
costDelta: 0.025,
138+
tokensIn: 200,
139+
tokensOut: 150,
140+
cacheReads: 1,
141+
cacheWrites: undefined,
142+
},
143+
{
144+
requestIndex: 3,
145+
cumulativeCost: 0.04,
146+
costDelta: 0.005,
147+
tokensIn: 50,
148+
tokensOut: 20,
149+
cacheReads: undefined,
150+
cacheWrites: 1,
151+
},
152+
]
153+
expect(calculateCostHistory(messages)).toEqual(expected)
154+
})
155+
156+
it("should ignore incomplete requests mixed with complete ones", () => {
157+
const messages: ClineMessage[] = [
158+
addTs({ type: "say", say: "api_req_started", text: JSON.stringify({ cost: 0.01, tokensIn: 100 }) }, 0),
159+
addTs({ type: "say", say: "api_req_started", text: JSON.stringify({ tokensIn: 200 }) }, 1),
160+
addTs({ type: "say", say: "api_req_started", text: JSON.stringify({ cost: null, tokensIn: 50 }) }, 2),
161+
addTs({ type: "say", say: "api_req_started", text: JSON.stringify({ cost: 0.03, tokensOut: 150 }) }, 3),
162+
]
163+
const expected: CostHistoryDataPoint[] = [
164+
{
165+
requestIndex: 1,
166+
cumulativeCost: 0.01,
167+
costDelta: 0.01,
168+
tokensIn: 100,
169+
tokensOut: undefined,
170+
cacheReads: undefined,
171+
cacheWrites: undefined,
172+
},
173+
{
174+
requestIndex: 2,
175+
cumulativeCost: 0.04,
176+
costDelta: 0.03,
177+
tokensIn: undefined,
178+
tokensOut: 150,
179+
cacheReads: undefined,
180+
cacheWrites: undefined,
181+
},
182+
]
183+
expect(calculateCostHistory(messages)).toEqual(expected)
184+
})
185+
186+
it("should handle missing optional fields gracefully", () => {
187+
const messages: ClineMessage[] = [
188+
addTs({ type: "say", say: "api_req_started", text: JSON.stringify({ cost: 0.01 }) }, 0),
189+
addTs({ type: "say", say: "api_req_started", text: JSON.stringify({ cost: 0.02, tokensIn: 100 }) }, 1),
190+
addTs({ type: "say", say: "api_req_started", text: JSON.stringify({ cost: 0.03, cacheReads: 5 }) }, 2),
191+
]
192+
const expected: CostHistoryDataPoint[] = [
193+
{
194+
requestIndex: 1,
195+
cumulativeCost: 0.01,
196+
costDelta: 0.01,
197+
tokensIn: undefined,
198+
tokensOut: undefined,
199+
cacheReads: undefined,
200+
cacheWrites: undefined,
201+
},
202+
{
203+
requestIndex: 2,
204+
cumulativeCost: 0.03,
205+
costDelta: 0.02,
206+
tokensIn: 100,
207+
tokensOut: undefined,
208+
cacheReads: undefined,
209+
cacheWrites: undefined,
210+
},
211+
{
212+
requestIndex: 3,
213+
cumulativeCost: 0.06,
214+
costDelta: 0.03,
215+
tokensIn: undefined,
216+
tokensOut: undefined,
217+
cacheReads: 5,
218+
cacheWrites: undefined,
219+
},
220+
]
221+
expect(calculateCostHistory(messages)).toEqual(expected)
222+
})
223+
})

webview-ui/src/utils/costUtils.ts

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
import { ClineMessage } from "@roo/shared/ExtensionMessage"
2+
3+
export interface CostHistoryDataPoint {
4+
requestIndex: number
5+
cumulativeCost: number
6+
costDelta: number
7+
tokensIn?: number
8+
tokensOut?: number
9+
cacheReads?: number
10+
cacheWrites?: number
11+
}
12+
13+
export function calculateCostHistory(messages: ClineMessage[]): CostHistoryDataPoint[] {
14+
let cumulativeCost = 0
15+
let requestIndex = 0
16+
const history: CostHistoryDataPoint[] = []
17+
18+
messages.forEach((message) => {
19+
if (message.say === "api_req_started" && message.text) {
20+
try {
21+
const data = JSON.parse(message.text)
22+
if (data.cost !== undefined && data.cost !== null) {
23+
const costDelta = data.cost ?? 0
24+
cumulativeCost += costDelta
25+
requestIndex++
26+
27+
history.push({
28+
requestIndex,
29+
cumulativeCost,
30+
costDelta,
31+
tokensIn: data.tokensIn,
32+
tokensOut: data.tokensOut,
33+
cacheReads: data.cacheReads,
34+
cacheWrites: data.cacheWrites,
35+
})
36+
}
37+
} catch (e) {
38+
console.error("Error parsing api_req_started text for cost history:", e, message.text)
39+
}
40+
}
41+
})
42+
return history
43+
}

0 commit comments

Comments
 (0)