Skip to content

fix: calculate contextTokens as cumulative sum of all conversation tokens #6007

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
roomote wants to merge 1 commit into from
Closed

fix: calculate contextTokens as cumulative sum of all conversation tokens #6007

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
55 changes: 34 additions & 21 deletions src/shared/__tests__/getApiMetrics.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ describe("getApiMetrics", () => {
expect(result.totalCacheWrites).toBe(5)
expect(result.totalCacheReads).toBe(10)
expect(result.totalCost).toBe(0.005)
expect(result.contextTokens).toBe(300) // 100 + 200 (OpenAI default, no cache tokens)
expect(result.contextTokens).toBe(300) // 100 + 200 (cumulative tokens)
})

it("should calculate metrics from multiple api_req_started messages", () => {
Expand All @@ -83,7 +83,7 @@ describe("getApiMetrics", () => {
expect(result.totalCacheWrites).toBe(8) // 5 + 3
expect(result.totalCacheReads).toBe(17) // 10 + 7
expect(result.totalCost).toBe(0.008) // 0.005 + 0.003
expect(result.contextTokens).toBe(200) // 50 + 150 (OpenAI default, no cache tokens)
expect(result.contextTokens).toBe(500) // (100 + 200) + (50 + 150) - cumulative
})

it("should calculate metrics from condense_context messages", () => {
Expand Down Expand Up @@ -123,7 +123,7 @@ describe("getApiMetrics", () => {
expect(result.totalCacheWrites).toBe(8) // 5 + 3
expect(result.totalCacheReads).toBe(17) // 10 + 7
expect(result.totalCost).toBe(0.01) // 0.005 + 0.002 + 0.003
expect(result.contextTokens).toBe(200) // 50 + 150 (OpenAI default, no cache tokens)
expect(result.contextTokens).toBe(700) // 500 (from condense) + 50 + 150
})
})

Expand Down Expand Up @@ -242,9 +242,8 @@ describe("getApiMetrics", () => {
expect(result.totalCacheReads).toBe(10)
expect(result.totalCost).toBe(0.005)

// The implementation will use the last message that has any tokens
// In this case, it's the message with tokensOut:200 (since the last few messages have no tokensIn/Out)
expect(result.contextTokens).toBe(200) // 0 + 200 (from the tokensOut message)
// The cumulative context should be the sum of all tokens
expect(result.contextTokens).toBe(300) // 100 + 0 + 0 + 0 + 200 (cumulative)
})

it("should handle non-number values in api_req_started message", () => {
Expand All @@ -264,48 +263,62 @@ describe("getApiMetrics", () => {
expect(result.totalCacheReads).toBeUndefined()
expect(result.totalCost).toBe(0)

// The implementation concatenates all token values including cache tokens
expect(result.contextTokens).toBe("not-a-numbernot-a-number") // tokensIn + tokensOut (OpenAI default)
// Non-number values should result in 0 context tokens
expect(result.contextTokens).toBe(0)
})
})

describe("Context tokens calculation", () => {
it("should calculate contextTokens from the last api_req_started message", () => {
it("should calculate cumulative contextTokens from all api_req_started messages", () => {
const messages: ClineMessage[] = [
createApiReqStartedMessage('{"tokensIn":100,"tokensOut":200,"cacheWrites":5,"cacheReads":10}', 1000),
createApiReqStartedMessage('{"tokensIn":50,"tokensOut":150,"cacheWrites":3,"cacheReads":7}', 2000),
]

const result = getApiMetrics(messages)

// Should use the values from the last api_req_started message
expect(result.contextTokens).toBe(200) // 50 + 150 (OpenAI default, no cache tokens)
// Should sum all tokens from all messages
expect(result.contextTokens).toBe(500) // (100 + 200) + (50 + 150)
})

it("should calculate contextTokens from the last condense_context message", () => {
it("should reset contextTokens after condense_context message", () => {
const messages: ClineMessage[] = [
createApiReqStartedMessage('{"tokensIn":100,"tokensOut":200,"cacheWrites":5,"cacheReads":10}', 1000),
createCondenseContextMessage(0.002, 500, 1000, 2000),
]

const result = getApiMetrics(messages)

// Should use newContextTokens from the last condense_context message
// Should use newContextTokens from the condense_context message
expect(result.contextTokens).toBe(500)
})

it("should prioritize the last message for contextTokens calculation", () => {
it("should accumulate tokens after condense_context", () => {
const messages: ClineMessage[] = [
createCondenseContextMessage(0.002, 500, 1000, 1000),
createApiReqStartedMessage('{"tokensIn":100,"tokensOut":200,"cacheWrites":5,"cacheReads":10}', 2000),
createCondenseContextMessage(0.003, 400, 800, 3000),
createApiReqStartedMessage('{"tokensIn":50,"tokensOut":150,"cacheWrites":3,"cacheReads":7}', 4000),
createApiReqStartedMessage('{"tokensIn":100,"tokensOut":200,"cacheWrites":5,"cacheReads":10}', 1000),
createCondenseContextMessage(0.002, 500, 1000, 2000),
createApiReqStartedMessage('{"tokensIn":50,"tokensOut":150,"cacheWrites":3,"cacheReads":7}', 3000),
]

const result = getApiMetrics(messages)

// Should use the values from the last api_req_started message
expect(result.contextTokens).toBe(200) // 50 + 150 (OpenAI default, no cache tokens)
// Should use condense tokens + new tokens after condense
expect(result.contextTokens).toBe(700) // 500 + (50 + 150)
})

it("should handle multiple condense_context messages correctly", () => {
const messages: ClineMessage[] = [
createApiReqStartedMessage('{"tokensIn":100,"tokensOut":200}', 1000),
createCondenseContextMessage(0.002, 500, 1000, 2000),
createApiReqStartedMessage('{"tokensIn":50,"tokensOut":150}', 3000),
createCondenseContextMessage(0.003, 400, 800, 4000),
createApiReqStartedMessage('{"tokensIn":25,"tokensOut":75}', 5000),
]

const result = getApiMetrics(messages)

// Should use the last condense tokens + tokens after it
expect(result.contextTokens).toBe(500) // 400 + (25 + 75)
})

it("should handle missing values when calculating contextTokens", () => {
Expand All @@ -320,7 +333,7 @@ describe("getApiMetrics", () => {
const result = getApiMetrics(messages)

// Should handle missing or invalid values
expect(result.contextTokens).toBe(0) // 0 + 0 (OpenAI default, no cache tokens)
expect(result.contextTokens).toBe(0)

// Restore console.error
console.error = originalConsoleError
Expand Down
63 changes: 34 additions & 29 deletions src/shared/getApiMetrics.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,20 @@ export function getApiMetrics(messages: ClineMessage[]) {
contextTokens: 0,
}

// Calculate running totals
messages.forEach((message) => {
// Track cumulative context tokens
let cumulativeContextTokens = 0
let lastCondenseIndex = -1

// Find the last condense_context message if any
for (let i = messages.length - 1; i >= 0; i--) {
if (messages[i].type === "say" && messages[i].say === "condense_context") {
lastCondenseIndex = i
break
}
}

// Calculate running totals and context tokens
messages.forEach((message, index) => {
if (message.type === "say" && message.say === "api_req_started" && message.text) {
try {
const parsedText: ParsedApiReqStartedTextType = JSON.parse(message.text)
Expand All @@ -58,41 +70,34 @@ export function getApiMetrics(messages: ClineMessage[]) {
if (typeof cost === "number") {
result.totalCost += cost
}

// Add to cumulative context tokens if this message is after the last condense
if (index > lastCondenseIndex) {
// For context calculation, we count input and output tokens
// Cache reads represent tokens that were already in context, so we don't add them
// Cache writes are new tokens being added to context
if (typeof tokensIn === "number") {
cumulativeContextTokens += tokensIn
}
if (typeof tokensOut === "number") {
cumulativeContextTokens += tokensOut
}
}
} catch (error) {
console.error("Error parsing JSON:", error)
}
} else if (message.type === "say" && message.say === "condense_context") {
result.totalCost += message.contextCondense?.cost ?? 0
}
})

// Calculate context tokens, from the last API request started or condense context message
result.contextTokens = 0
for (let i = messages.length - 1; i >= 0; i--) {
const message = messages[i]
if (message.type === "say" && message.say === "api_req_started" && message.text) {
try {
const parsedText: ParsedApiReqStartedTextType = JSON.parse(message.text)
const { tokensIn, tokensOut, cacheWrites, cacheReads, apiProtocol } = parsedText

// Calculate context tokens based on API protocol
if (apiProtocol === "anthropic") {
result.contextTokens = (tokensIn || 0) + (tokensOut || 0) + (cacheWrites || 0) + (cacheReads || 0)
} else {
// For OpenAI (or when protocol is not specified)
result.contextTokens = (tokensIn || 0) + (tokensOut || 0)
}
} catch (error) {
console.error("Error parsing JSON:", error)
continue
// When we hit a condense_context, reset the cumulative tokens to the new context size
if (index === lastCondenseIndex && message.contextCondense?.newContextTokens !== undefined) {
cumulativeContextTokens = message.contextCondense.newContextTokens
}
} else if (message.type === "say" && message.say === "condense_context") {
result.contextTokens = message.contextCondense?.newContextTokens ?? 0
}
if (result.contextTokens) {
break
}
}
})

// Set the final context tokens
result.contextTokens = cumulativeContextTokens

return result
}