Skip to content

Commit 70af1d2

Browse files
committed
fix: update context token calculation to match upstream
The context token calculation has been updated to: - Include cache tokens (cacheWrites + cacheReads) in the total context size - Use a cleaner approach to find the last valid API request - Fix issue where placeholder messages without token info were being counted - Match upstream's implementation of getTotalTokensFromMessage This fixes the issue where context size was incorrectly showing as 139 tokens without any API requests being made.
1 parent 1edb939 commit 70af1d2

File tree

1 file changed

+20
-23
lines changed

1 file changed

+20
-23
lines changed

src/shared/getApiMetrics.ts

Lines changed: 20 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ interface ApiMetrics {
66
totalCacheWrites?: number
77
totalCacheReads?: number
88
totalCost: number
9-
contextTokens: number // Total tokens in conversation (last message's tokensIn + tokensOut)
9+
contextTokens: number // Total tokens in conversation (last message's tokensIn + tokensOut + cacheWrites + cacheReads)
1010
}
1111

1212
/**
@@ -17,7 +17,7 @@ interface ApiMetrics {
1717
* It extracts and sums up the tokensIn, tokensOut, cacheWrites, cacheReads, and cost from these messages.
1818
*
1919
* @param messages - An array of ClineMessage objects to process.
20-
* @returns An ApiMetrics object containing totalTokensIn, totalTokensOut, totalCacheWrites, totalCacheReads, and totalCost.
20+
* @returns An ApiMetrics object containing totalTokensIn, totalTokensOut, totalCacheWrites, totalCacheReads, totalCost, and contextTokens.
2121
*
2222
* @example
2323
* const messages = [
@@ -36,27 +36,30 @@ export function getApiMetrics(messages: ClineMessage[]): ApiMetrics {
3636
contextTokens: 0,
3737
}
3838

39-
// Find the last api_req_started message that has valid token information
39+
// Helper function to get total tokens from a message
40+
const getTotalTokensFromMessage = (message: ClineMessage): number => {
41+
if (!message.text) return 0
42+
try {
43+
const { tokensIn, tokensOut, cacheWrites, cacheReads } = JSON.parse(message.text)
44+
return (tokensIn || 0) + (tokensOut || 0) + (cacheWrites || 0) + (cacheReads || 0)
45+
} catch {
46+
return 0
47+
}
48+
}
49+
50+
// Find the last api_req_started message that has any tokens
4051
const lastApiReq = [...messages].reverse().find((message) => {
41-
if (message.type === "say" && message.say === "api_req_started" && message.text) {
42-
try {
43-
const parsedData = JSON.parse(message.text)
44-
return typeof parsedData.tokensIn === "number" && typeof parsedData.tokensOut === "number"
45-
} catch {
46-
return false
47-
}
52+
if (message.type === "say" && message.say === "api_req_started") {
53+
return getTotalTokensFromMessage(message) > 0
4854
}
4955
return false
5056
})
5157

52-
// Keep track of the last valid context tokens
53-
let lastValidContextTokens = 0
54-
58+
// Calculate running totals
5559
messages.forEach((message) => {
5660
if (message.type === "say" && message.say === "api_req_started" && message.text) {
5761
try {
58-
const parsedData = JSON.parse(message.text)
59-
const { tokensIn, tokensOut, cacheWrites, cacheReads, cost } = parsedData
62+
const { tokensIn, tokensOut, cacheWrites, cacheReads, cost } = JSON.parse(message.text)
6063

6164
if (typeof tokensIn === "number") {
6265
result.totalTokensIn += tokensIn
@@ -74,15 +77,9 @@ export function getApiMetrics(messages: ClineMessage[]): ApiMetrics {
7477
result.totalCost += cost
7578
}
7679

77-
// Update last valid context tokens whenever we have valid input and output tokens
78-
if (tokensIn > 0 && tokensOut > 0) {
79-
lastValidContextTokens = tokensIn + tokensOut
80-
}
81-
82-
// If this is the last api request, use its tokens for context size
80+
// If this is the last api request with tokens, use its total for context size
8381
if (message === lastApiReq) {
84-
// Use the last valid context tokens if the current request doesn't have valid tokens
85-
result.contextTokens = tokensIn > 0 && tokensOut > 0 ? tokensIn + tokensOut : lastValidContextTokens
82+
result.contextTokens = getTotalTokensFromMessage(message)
8683
}
8784
} catch (error) {
8885
console.error("Error parsing JSON:", error)

0 commit comments

Comments
 (0)