From eca73954abbf335b0c1979a4fc7390b6c747064c Mon Sep 17 00:00:00 2001
From: Hemasekhar Puchuginjala <hemasekhar@google.com>
Date: Fri, 12 Jun 2026 16:02:05 +0530
Subject: [PATCH] fix: map token usage metadata for Anthropic Claude model

---
 .../java/com/google/adk/models/Claude.java    | 15 ++++++++-----
 .../com/google/adk/models/ClaudeTest.java     | 22 +++++++++++++++++++
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/core/src/main/java/com/google/adk/models/Claude.java b/core/src/main/java/com/google/adk/models/Claude.java
index 79201c261..d3f65a40c 100644
--- a/core/src/main/java/com/google/adk/models/Claude.java
+++ b/core/src/main/java/com/google/adk/models/Claude.java
@@ -34,11 +34,7 @@
 import com.google.adk.JsonBaseModel;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.genai.types.Content;
-import com.google.genai.types.FunctionCall;
-import com.google.genai.types.FunctionDeclaration;
-import com.google.genai.types.GenerateContentConfig;
-import com.google.genai.types.Part;
+import com.google.genai.types.*;
 import io.reactivex.rxjava3.core.Flowable;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -268,6 +264,15 @@ private LlmResponse convertAnthropicResponseToLlmResponse(Message message) {
       responseBuilder.content(
           Content.builder().role("model").parts(ImmutableList.copyOf(parts)).build());
     }
+    if (message.usage() != null) {
+      responseBuilder.usageMetadata(
+          GenerateContentResponseUsageMetadata.builder()
+              .promptTokenCount((int) message.usage().inputTokens())
+              .candidatesTokenCount((int) message.usage().outputTokens())
+              .totalTokenCount(
+                  (int) (message.usage().inputTokens() + message.usage().outputTokens()))
+              .build());
+    }
     return responseBuilder.build();
   }
 
diff --git a/core/src/test/java/com/google/adk/models/ClaudeTest.java b/core/src/test/java/com/google/adk/models/ClaudeTest.java
index febcaf4be..c2f8f2d56 100644
--- a/core/src/test/java/com/google/adk/models/ClaudeTest.java
+++ b/core/src/test/java/com/google/adk/models/ClaudeTest.java
@@ -17,10 +17,15 @@
 package com.google.adk.models;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 import com.anthropic.client.AnthropicClient;
 import com.anthropic.models.messages.ContentBlockParam;
+import com.anthropic.models.messages.Message;
 import com.anthropic.models.messages.ToolResultBlockParam;
+import com.anthropic.models.messages.Usage;
 import com.google.common.collect.ImmutableMap;
 import com.google.genai.types.FunctionResponse;
 import com.google.genai.types.Part;
@@ -78,4 +83,21 @@ public void testPartToAnthropicMessageBlock_jsonFallback() throws Exception {
     ToolResultBlockParam toolResult = result.asToolResult();
     assertThat(toolResult.content().get().asString()).contains("\"custom_key\":\"custom_value\"");
   }
+
+  @Test
+  public void testClaudeUsageMapping_ShouldFailWhenMappingIsMissing() throws Exception {
+    Usage mockUsage = mock(Usage.class);
+    when(mockUsage.inputTokens()).thenReturn(10L);
+    when(mockUsage.outputTokens()).thenReturn(20L);
+
+    Message mockMessage = mock(Message.class);
+    when(mockMessage.usage()).thenReturn(mockUsage);
+    when(mockMessage.content()).thenReturn(java.util.Collections.emptyList());
+
+    Method convertMethod =
+        Claude.class.getDeclaredMethod("convertAnthropicResponseToLlmResponse", Message.class);
+    convertMethod.setAccessible(true);
+    LlmResponse result = (LlmResponse) convertMethod.invoke(claude, mockMessage);
+    assertTrue(result.usageMetadata().isPresent());
+  }
 }