Yield the cost in usage data for Anthropic (#4849)

Author: cte

src/api/providers/anthropic.ts

@@ -17,6 +17,7 @@ import { getModelParams } from "../transform/model-params"
 
 import { BaseProvider } from "./base-provider"
 import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
+import { calculateApiCostAnthropic } from "../../shared/cost"
 
 export class AnthropicHandler extends BaseProvider implements SingleCompletionHandler {
 	private options: ApiHandlerOptions
@@ -132,20 +133,35 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa
 			}
 		}
 
+		let inputTokens = 0
+		let outputTokens = 0
+		let cacheWriteTokens = 0
+		let cacheReadTokens = 0
+
 		for await (const chunk of stream) {
 			switch (chunk.type) {
 				case "message_start": {
 					// Tells us cache reads/writes/input/output.
-					const usage = chunk.message.usage
+					const {
+						input_tokens = 0,
+						output_tokens = 0,
+						cache_creation_input_tokens,
+						cache_read_input_tokens,
+					} = chunk.message.usage
 
 					yield {
 						type: "usage",
-						inputTokens: usage.input_tokens || 0,
-						outputTokens: usage.output_tokens || 0,
-						cacheWriteTokens: usage.cache_creation_input_tokens || undefined,
-						cacheReadTokens: usage.cache_read_input_tokens || undefined,
+						inputTokens: input_tokens,
+						outputTokens: output_tokens,
+						cacheWriteTokens: cache_creation_input_tokens || undefined,
+						cacheReadTokens: cache_read_input_tokens || undefined,
 					}
 
+					inputTokens += input_tokens
+					outputTokens += output_tokens
+					cacheWriteTokens += cache_creation_input_tokens || 0
+					cacheReadTokens += cache_read_input_tokens || 0
+
 					break
 				}
 				case "message_delta":
@@ -198,6 +214,21 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa
 					break
 			}
 		}
+
+		if (inputTokens > 0 || outputTokens > 0 || cacheWriteTokens > 0 || cacheReadTokens > 0) {
+			yield {
+				type: "usage",
+				inputTokens: 0,
+				outputTokens: 0,
+				totalCost: calculateApiCostAnthropic(
+					this.getModel().info,
+					inputTokens,
+					outputTokens,
+					cacheWriteTokens,
+					cacheReadTokens,
+				),
+			}
+		}
 	}
 
 	getModel() {

src/core/task/Task.ts

@@ -1439,19 +1439,22 @@ export class Task extends EventEmitter<ClineEvents> {
 			} finally {
 				this.isStreaming = false
 			}
-			if (
-				inputTokens > 0 ||
-				outputTokens > 0 ||
-				cacheWriteTokens > 0 ||
-				cacheReadTokens > 0 ||
-				typeof totalCost !== "undefined"
-			) {
+
+			if (inputTokens > 0 || outputTokens > 0 || cacheWriteTokens > 0 || cacheReadTokens > 0) {
 				TelemetryService.instance.captureLlmCompletion(this.taskId, {
 					inputTokens,
 					outputTokens,
 					cacheWriteTokens,
 					cacheReadTokens,
-					cost: totalCost,
+					cost:
+						totalCost ??
+						calculateApiCostAnthropic(
+							this.api.getModel().info,
+							inputTokens,
+							outputTokens,
+							cacheWriteTokens,
+							cacheReadTokens,
+						),
 				})
 			}
 

src/services/checkpoints/__tests__/ShadowCheckpointService.spec.ts

@@ -12,7 +12,7 @@ import * as fileSearch from "../../../services/search/file-search"
 
 import { RepoPerTaskCheckpointService } from "../RepoPerTaskCheckpointService"
 
-vitest.setConfig({ testTimeout: 10_000 })
+vitest.setConfig({ testTimeout: 20_000 })
 
 const tmpDir = path.join(os.tmpdir(), "CheckpointService")
 

src/shared/cost.ts

@@ -15,26 +15,25 @@ function calculateApiCostInternal(
 	return totalCost
 }
 
-// For Anthropic compliant usage, the input tokens count does NOT include the cached tokens
+// For Anthropic compliant usage, the input tokens count does NOT include the
+// cached tokens.
 export function calculateApiCostAnthropic(
 	modelInfo: ModelInfo,
 	inputTokens: number,
 	outputTokens: number,
 	cacheCreationInputTokens?: number,
 	cacheReadInputTokens?: number,
 ): number {
-	const cacheCreationInputTokensNum = cacheCreationInputTokens || 0
-	const cacheReadInputTokensNum = cacheReadInputTokens || 0
 	return calculateApiCostInternal(
 		modelInfo,
 		inputTokens,
 		outputTokens,
-		cacheCreationInputTokensNum,
-		cacheReadInputTokensNum,
+		cacheCreationInputTokens || 0,
+		cacheReadInputTokens || 0,
 	)
 }
 
-// For OpenAI compliant usage, the input tokens count INCLUDES the cached tokens
+// For OpenAI compliant usage, the input tokens count INCLUDES the cached tokens.
 export function calculateApiCostOpenAI(
 	modelInfo: ModelInfo,
 	inputTokens: number,
@@ -45,6 +44,7 @@ export function calculateApiCostOpenAI(
 	const cacheCreationInputTokensNum = cacheCreationInputTokens || 0
 	const cacheReadInputTokensNum = cacheReadInputTokens || 0
 	const nonCachedInputTokens = Math.max(0, inputTokens - cacheCreationInputTokensNum - cacheReadInputTokensNum)
+
 	return calculateApiCostInternal(
 		modelInfo,
 		nonCachedInputTokens,