From 13b10bb54dcbed52f0c11ac99defe6b2766e0e54 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Tue, 20 Oct 2026 08:53:17 -0700 Subject: [PATCH] chore: send ADK tracking headers from Claude and from Apigee custom headers PiperOrigin-RevId: 1003140617 --- .../java/com/google/adk/models/ApigeeLlm.java | 28 +++--- .../java/com/google/adk/models/Claude.java | 38 +++++++- .../com/google/adk/models/ApigeeLlmTest.java | 45 ++++++++-- .../com/google/adk/models/ClaudeTest.java | 88 +++++++++++++++++++ 4 files changed, 174 insertions(+), 25 deletions(-) diff --git a/core/src/main/java/com/google/adk/models/ApigeeLlm.java b/core/src/main/java/com/google/adk/models/ApigeeLlm.java index 419a8ff7f..9f357662d 100644 --- a/core/src/main/java/com/google/adk/models/ApigeeLlm.java +++ b/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 trackingHeaders = - this.apiType == ApiType.CHAT_COMPLETIONS - ? TrackingHeaders.HEADERS - : TrackingHeaders.GENAI_SDK_HEADERS; - // Build the Client + Map 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.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. 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 f20f8f40c..3227c32a8 100644 --- a/core/src/main/java/com/google/adk/models/Claude.java +++ b/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 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 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 generateContent(LlmRequest llmRequest, boolean stream) { // TODO: Switch to streaming API. diff --git a/core/src/test/java/com/google/adk/models/ApigeeLlmTest.java b/core/src/test/java/com/google/adk/models/ApigeeLlmTest.java index 3eb986fc8..92f9611e3 100644 --- a/core/src/test/java/com/google/adk/models/ApigeeLlmTest.java +++ b/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 overlappingHeaders = ImmutableMap.of("user-agent", "custom-agent"); - ApigeeLlm.Builder builder = + public void build_chatCompletions_addsTheTrackingTagsCustomHeadersLack() { + ImmutableMap 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. 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 ff3103d73..0eac85a2f 100644 --- a/core/src/test/java/com/google/adk/models/ClaudeTest.java +++ b/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 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); + } + } }