fix: count Gemini tool-use prompt tokens as input

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

backend/windmill-ai/src/ai_google.rs

@@ -292,6 +292,22 @@ pub struct GeminiUsageMetadata {
     /// Thinking tokens, billed as output but counted apart from `candidatesTokenCount`.
     #[serde(rename = "thoughtsTokenCount", default)]
     pub thoughts_token_count: Option<i32>,
+    /// Input tokens spent on tool-use prompts, counted apart from `promptTokenCount`
+    /// rather than within it.
+    #[serde(rename = "toolUsePromptTokenCount", default)]
+    pub tool_use_prompt_token_count: Option<i32>,
+}
+
+/// Input tokens as billed. Gemini reports tool-use prompts in their own field, and
+/// they are disjoint from `promptTokenCount`: a live tool call returns 17 prompt +
+/// 60 tool-use + 17 candidates + 52 thoughts against a `totalTokenCount` of 146, so
+/// leaving them out under-reports the input of every tool-using turn. Cached tokens
+/// are not added here, being already part of `promptTokenCount`.
+fn gemini_prompt_tokens(usage: &GeminiUsageMetadata) -> i32 {
+    usage
+        .prompt_token_count
+        .unwrap_or(0)
+        .saturating_add(usage.tool_use_prompt_token_count.unwrap_or(0))
 }
 
 /// Output tokens as billed: Gemini counts thinking apart from `candidatesTokenCount`
@@ -605,7 +621,7 @@ pub fn gemini_response_to_openai(parsed: &GeminiParsedEvent, model: &str) -> ser
 
     let usage = parsed.usage.as_ref().map(|u| {
         serde_json::json!({
-            "prompt_tokens": u.prompt_token_count.unwrap_or(0),
+            "prompt_tokens": gemini_prompt_tokens(u),
             "completion_tokens": gemini_completion_tokens(u),
             "total_tokens": u.total_token_count.unwrap_or(0),
             "prompt_tokens_details": {
@@ -700,7 +716,7 @@ pub fn gemini_event_to_openai_sse_chunks(
     // OpenAI's `stream_options.include_usage` terminal chunk (top-level `usage`,
     // empty `choices`) so the frontend's `'usage' in chunk` path records them.
     if let Some(usage) = &parsed.usage {
-        let prompt_tokens = usage.prompt_token_count.unwrap_or(0);
+        let prompt_tokens = gemini_prompt_tokens(usage);
         let completion_tokens = gemini_completion_tokens(usage);
         let total_tokens = usage
             .total_token_count
@@ -1003,6 +1019,7 @@ mod tests {
                 total_token_count: Some(1120),
                 cached_content_token_count: Some(900),
                 thoughts_token_count: Some(100),
+                ..Default::default()
             }),
             ..Default::default()
         };
@@ -1028,6 +1045,47 @@ mod tests {
         assert_eq!(usage_chunk["usage"]["completion_tokens"], 120);
     }
 
+    #[test]
+    fn gemini_usage_chunk_counts_tool_use_prompt_tokens() {
+        // The four counts a live tool call reported, against its totalTokenCount of
+        // 146 — tool-use prompts are disjoint from promptTokenCount, so the input is
+        // 17 + 60 rather than 17.
+        let parsed = GeminiParsedEvent {
+            text: Some("Canberra".to_string()),
+            usage: Some(GeminiUsageMetadata {
+                prompt_token_count: Some(17),
+                candidates_token_count: Some(17),
+                total_token_count: Some(146),
+                tool_use_prompt_token_count: Some(60),
+                thoughts_token_count: Some(52),
+                ..Default::default()
+            }),
+            ..Default::default()
+        };
+
+        let mut tool_call_index = 0;
+        let chunks = gemini_event_to_openai_sse_chunks(
+            &parsed,
+            "chatcmpl-test",
+            "gemini-2.5-flash",
+            &mut tool_call_index,
+        );
+        let usage_chunk = chunks
+            .iter()
+            .map(|c| parse_sse_chunk(c))
+            .find(|v| v.get("usage").map(|u| !u.is_null()).unwrap_or(false))
+            .expect("a chunk should carry top-level usage");
+
+        assert_eq!(usage_chunk["usage"]["prompt_tokens"], 77);
+        assert_eq!(usage_chunk["usage"]["completion_tokens"], 69);
+        // Everything Google billed is accounted for.
+        assert_eq!(
+            usage_chunk["usage"]["prompt_tokens"].as_i64().unwrap()
+                + usage_chunk["usage"]["completion_tokens"].as_i64().unwrap(),
+            146
+        );
+    }
+
     #[test]
     fn gemini_streaming_usage_total_falls_back_to_prompt_plus_completion() {
         let parsed = GeminiParsedEvent {

frontend/src/lib/components/copilot/chat/tokenUsage.ts

@@ -158,8 +158,10 @@ export function openAICompletionsUsageToChatTokenUsage(
 				total_tokens?: number | null
 				prompt_tokens_details?: {
 					cached_tokens?: number | null
-					/** OpenRouter surfaces Anthropic's cache creation here; OpenAI, whose
-					 * caching is automatic and unbilled, reports no such field. */
+					/** Cache creation, reported by the providers that bill for it: OpenRouter
+					 * passes Anthropic's through, and the Bedrock proxy folds
+					 * `cacheWriteInputTokens` in here. OpenAI, whose caching is automatic and
+					 * unbilled, reports no such field. */
 					cache_write_tokens?: number | null
 				} | null
 				/** OpenRouter reports what it actually charged when the request opts in. */