chore: send ADK tracking headers from Claude and from Apigee custom headers

PiperOrigin-RevId: 1003140617

Author: damianmomotgoogle

core/src/main/java/com/google/adk/models/ApigeeLlm.java

@@ -95,27 +95,21 @@ private ApigeeLlm(
               + " APIGEE_PROXY_URL.");
     }
 
-    // Chat completions clients send no language tag, so they get both tags; the genai SDK adds its
-    // own gl-java tag.
-    ImmutableMap<String, String> trackingHeaders =
-        this.apiType == ApiType.CHAT_COMPLETIONS
-            ? TrackingHeaders.HEADERS
-            : TrackingHeaders.GENAI_SDK_HEADERS;
-
     // Build the Client
+    Map<String, String> callerHeaders = customHeaders != null ? customHeaders : ImmutableMap.of();
     HttpOptions.Builder httpOptionsBuilder =
-        HttpOptions.builder().baseUrl(effectiveProxyUrl).headers(trackingHeaders);
+        HttpOptions.builder()
+            .baseUrl(effectiveProxyUrl)
+            // Chat completions clients send no language tag, so they get both tags; the genai SDK
+            // adds its own gl-java tag.
+            .headers(
+                this.apiType == ApiType.CHAT_COMPLETIONS
+                    ? TrackingHeaders.mergeInto(callerHeaders)
+                    : TrackingHeaders.mergeIntoGenAiSdkHeaders(callerHeaders));
     String apiVersion = identifyApiVersion(modelName);
     if (!apiVersion.isEmpty()) {
       httpOptionsBuilder.apiVersion(apiVersion);
     }
-    if (customHeaders != null) {
-      httpOptionsBuilder.headers(
-          ImmutableMap.<String, String>builder()
-              .putAll(trackingHeaders)
-              .putAll(customHeaders)
-              .buildOrThrow());
-    }
     this.httpOptions = httpOptionsBuilder.build();
 
     if (this.apiType == ApiType.CHAT_COMPLETIONS) {
@@ -289,7 +283,9 @@ public Builder proxyUrl(String proxyUrl) {
     }
 
     /**
-     * Sets a dictionary of headers to be sent with the request.
+     * Sets a dictionary of headers to be sent with the request. A {@code user-agent} or {@code
+     * x-goog-api-client} value, matched case-insensitively, is kept and extended with the ADK
+     * tracking tags it lacks.
      *
      * @param customHeaders the custom headers.
      * @return this builder.

core/src/main/java/com/google/adk/models/Claude.java

@@ -19,6 +19,8 @@
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.anthropic.client.AnthropicClient;
+import com.anthropic.core.Properties;
+import com.anthropic.core.http.Headers;
 import com.anthropic.models.messages.ContentBlock;
 import com.anthropic.models.messages.ContentBlockParam;
 import com.anthropic.models.messages.Message;
@@ -89,22 +91,52 @@ public class Claude extends BaseLlm {
   private final AnthropicClient anthropicClient;
 
   /**
-   * Constructs a new Claude instance.
+   * Constructs a new Claude instance. Requests sent through {@code anthropicClient} carry ADK's
+   * usage-tracking headers, added to the ones the client already sends.
    *
    * @param modelName The name of the Claude model to use (e.g., "claude-3-opus-20240229").
    * @param anthropicClient The Anthropic API client instance.
    */
   public Claude(String modelName, AnthropicClient anthropicClient) {
     super(modelName);
-    this.anthropicClient = anthropicClient;
+    this.anthropicClient = withTrackingHeaders(anthropicClient);
   }
 
   public Claude(String modelName, AnthropicClient anthropicClient, int maxTokens) {
     super(modelName);
-    this.anthropicClient = anthropicClient;
+    this.anthropicClient = withTrackingHeaders(anthropicClient);
     this.maxTokens = maxTokens;
   }
 
+  /**
+   * Returns a copy of {@code client} that shares its connections and sends ADK's tracking headers,
+   * extending any the client already sets. A null client, or a test double that returns no copy, is
+   * returned unchanged.
+   */
+  private static @Nullable AnthropicClient withTrackingHeaders(@Nullable AnthropicClient client) {
+    if (client == null) {
+      return null;
+    }
+    AnthropicClient tracked =
+        client.withOptions(
+            options -> {
+              Headers configured = options.build().headers();
+              Map<String, String> current = new HashMap<>();
+              // Start from the SDK's default User-Agent, which it adds only when none is set.
+              current.put(
+                  "user-agent",
+                  client.getClass().getSimpleName() + "/Java " + Properties.getPackageVersion());
+              for (String name : TrackingHeaders.HEADERS.keySet()) {
+                List<String> values = configured.values(name);
+                if (!values.isEmpty()) {
+                  current.put(name, String.join(" ", values));
+                }
+              }
+              TrackingHeaders.mergeInto(current).forEach(options::replaceHeaders);
+            });
+    return tracked != null ? tracked : client;
+  }
+
   @Override
   public Flowable<LlmResponse> generateContent(LlmRequest llmRequest, boolean stream) {
     // TODO: Switch to streaming API.

core/src/test/java/com/google/adk/models/ApigeeLlmTest.java

@@ -331,15 +331,23 @@ public void validateModelString_rejectsOpenAiWithVersion() {
   }
 
   @Test
-  public void build_withCustomHeadersOverlappingTrackingHeaders_throwsException() {
-    ImmutableMap<String, String> overlappingHeaders = ImmutableMap.of("user-agent", "custom-agent");
-    ApigeeLlm.Builder builder =
+  public void build_chatCompletions_addsTheTrackingTagsCustomHeadersLack() {
+    ImmutableMap<String, String> overlappingHeaders =
+        ImmutableMap.of(
+            "User-Agent", "custom-agent", "x-goog-api-client", "lib/1 google-adk/0.0.1");
+    ApigeeLlm llm =
         ApigeeLlm.builder()
             .modelName("apigee/openai/gpt-4o")
             .proxyUrl(PROXY_URL)
-            .customHeaders(overlappingHeaders);
-    IllegalArgumentException e = assertThrows(IllegalArgumentException.class, builder::build);
-    assertThat(e).hasMessageThat().contains("Multiple entries with same key: user-agent=");
+            .customHeaders(overlappingHeaders)
+            .build();
+
+    assertThat(llm.getHttpOptions().headers().get())
+        .containsExactly(
+            "user-agent",
+            "custom-agent " + TrackingHeaders.VALUE,
+            "x-goog-api-client",
+            "lib/1 google-adk/0.0.1 gl-java/" + JAVA_VERSION.value());
   }
 
   @Test
@@ -428,6 +436,31 @@ public void generateContent_chatCompletions_sendsTrackingHeaders() throws Except
     }
   }
 
+  @Test
+  public void generateContent_genAi_extendsACustomUserAgentWithTheAdkTag() throws Exception {
+    Headers headers =
+        headersSentToProxy(
+            "apigee/gemini/v1beta/gemini-test",
+            ImmutableMap.of("User-Agent", "custom-agent"),
+            GENERATE_CONTENT_JSON);
+
+    assertThat(getOnlyElement(headers.get("user-agent")))
+        .endsWith(" custom-agent " + TrackingHeaders.ADK_TAG);
+    assertGenAiTrackingHeaders(headers);
+  }
+
+  @Test
+  public void generateContent_chatCompletions_extendsACustomUserAgent() throws Exception {
+    Headers headers =
+        headersSentToProxy(
+            "apigee/openai/gpt-4o",
+            ImmutableMap.of("User-Agent", "custom-agent"),
+            CHAT_COMPLETION_JSON);
+
+    assertThat(headers.get("user-agent")).containsExactly("custom-agent " + TrackingHeaders.VALUE);
+    assertThat(headers.get("x-goog-api-client")).containsExactly(TrackingHeaders.VALUE);
+  }
+
   /**
    * Sends one request through an {@link ApigeeLlm} for {@code modelName} to a local proxy that
    * answers with {@code responseJson}, and returns the request headers the proxy received.

core/src/test/java/com/google/adk/models/ClaudeTest.java

@@ -16,6 +16,7 @@
 
 package com.google.adk.models;
 
+import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.common.truth.Truth.assertThat;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.junit.Assert.assertEquals;
@@ -26,7 +27,9 @@
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.anthropic.backends.Backend;
 import com.anthropic.client.AnthropicClient;
+import com.anthropic.client.okhttp.AnthropicOkHttpClient;
 import com.anthropic.core.JsonValue;
 import com.anthropic.models.messages.ContentBlock;
 import com.anthropic.models.messages.ContentBlockParam;
@@ -46,17 +49,24 @@
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.net.InetAddresses;
 import com.google.genai.types.Content;
 import com.google.genai.types.FunctionCall;
 import com.google.genai.types.FunctionDeclaration;
 import com.google.genai.types.FunctionResponse;
 import com.google.genai.types.Part;
 import com.google.genai.types.Schema;
+import com.sun.net.httpserver.Headers;
+import com.sun.net.httpserver.HttpServer;
+import java.io.OutputStream;
 import java.lang.reflect.Method;
+import java.net.InetAddress;
+import java.net.InetSocketAddress;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -704,4 +714,82 @@ public void generateContent_userTurnOfOnlyEmptyText_isSentAsPlaceholder() {
     assertThat(blocks).hasSize(1);
     assertThat(blocks.get(0).asText().text()).isEqualTo(".");
   }
+
+  @Test
+  public void generateContent_extendsTheSdkUserAgentWithAdkTrackingHeaders() throws Exception {
+    Headers headers = headersSentThrough(AnthropicOkHttpClient.builder());
+
+    String userAgent = getOnlyElement(headers.get("user-agent"));
+    assertThat(userAgent).startsWith("AnthropicClientImpl/Java ");
+    assertThat(userAgent).endsWith(" " + TrackingHeaders.VALUE);
+    assertThat(headers.get("x-goog-api-client")).containsExactly(TrackingHeaders.VALUE);
+  }
+
+  @Test
+  public void generateContent_extendsTheCallerUserAgentWithAdkTrackingHeaders() throws Exception {
+    Headers headers =
+        headersSentThrough(AnthropicOkHttpClient.builder().putHeader("User-Agent", "my-app/1.0"));
+
+    assertThat(headers.get("user-agent")).containsExactly("my-app/1.0 " + TrackingHeaders.VALUE);
+    assertThat(headers.get("x-goog-api-client")).containsExactly(TrackingHeaders.VALUE);
+  }
+
+  @Test
+  public void constructor_nullClient_doesNotThrow() {
+    assertThat(new Claude("claude-test", (AnthropicClient) null).model()).isEqualTo("claude-test");
+  }
+
+  private static final String MESSAGE_JSON =
+      "{\"id\":\"msg_1\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-test\","
+          + "\"content\":[{\"type\":\"text\",\"text\":\"pong\"}],\"stop_reason\":\"end_turn\","
+          + "\"stop_sequence\":null,\"usage\":{\"input_tokens\":3,\"output_tokens\":2}}";
+
+  /**
+   * Sends one request through Claude using a client built from {@code clientBuilder} for a local
+   * server, and returns the headers the server received.
+   */
+  private static Headers headersSentThrough(AnthropicOkHttpClient.Builder clientBuilder)
+      throws Exception {
+    AtomicReference<Headers> requestHeaders = new AtomicReference<>();
+    HttpServer server =
+        HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0);
+    server.createContext(
+        "/",
+        exchange -> {
+          requestHeaders.set(exchange.getRequestHeaders());
+          byte[] body = MESSAGE_JSON.getBytes(UTF_8);
+          exchange.getResponseHeaders().set("Content-Type", "application/json");
+          exchange.sendResponseHeaders(200, body.length);
+          try (OutputStream out = exchange.getResponseBody()) {
+            out.write(body);
+          }
+        });
+    server.start();
+    try {
+      String baseUrl =
+          "http://"
+              + InetAddresses.toUriString(server.getAddress().getAddress())
+              + ":"
+              + server.getAddress().getPort();
+      Backend backend =
+          new Backend() {
+            @Override
+            public String baseUrl() {
+              return baseUrl;
+            }
+
+            @Override
+            public void close() {}
+          };
+      Claude model = new Claude("claude-test", clientBuilder.backend(backend).build());
+
+      LlmResponse response =
+          model.generateContent(request(userText("Ping")), /* stream= */ false).blockingFirst();
+
+      assertThat(response.content().get().parts().get().get(0).text()).hasValue("pong");
+      return requestHeaders.get();
+    } finally {
+      server.stop(0);
+    }
+  }
 }