From 48e9ae7a6f32ed678404a5cafc9b5f48118b0adb Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Tue, 20 Oct 2026 06:23:37 -0700 Subject: [PATCH] chore: send one ADK tag and one Java tag on every Gemini request PiperOrigin-RevId: 1003055754 --- core/pom.xml | 14 + .../java/com/google/adk/models/ApigeeLlm.java | 27 +- .../java/com/google/adk/models/Gemini.java | 102 ++++--- .../google/adk/models/TrackingHeaders.java | 109 +++++++ .../com/google/adk/models/ApigeeLlmTest.java | 103 +++++++ .../adk/models/GeminiTrackingHeadersTest.java | 269 ++++++++++++++++++ .../adk/models/TrackingHeadersTest.java | 50 ++++ 7 files changed, 614 insertions(+), 60 deletions(-) create mode 100644 core/src/main/java/com/google/adk/models/TrackingHeaders.java create mode 100644 core/src/test/java/com/google/adk/models/GeminiTrackingHeadersTest.java create mode 100644 core/src/test/java/com/google/adk/models/TrackingHeadersTest.java diff --git a/core/pom.xml b/core/pom.xml index 4feb5adde..7e7b973df 100644 --- a/core/pom.xml +++ b/core/pom.xml @@ -319,6 +319,20 @@ + + gemini-tracking-headers + + test + + + GeminiTrackingHeadersTest#generateContent_adkBuiltRegistryClient_sendsTrackingHeaders, GeminiTrackingHeadersTest#generateContent_adkBuiltVertexCredentialsClient_sendsTrackingHeaders + + + api-key + false + + + 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 9e68a56bd..419a8ff7f 100644 --- a/core/src/main/java/com/google/adk/models/ApigeeLlm.java +++ b/core/src/main/java/com/google/adk/models/ApigeeLlm.java @@ -15,10 +15,8 @@ */ package com.google.adk.models; -import static com.google.common.base.StandardSystemProperty.JAVA_VERSION; import static com.google.common.base.Strings.isNullOrEmpty; -import com.google.adk.Version; import com.google.adk.models.chat.ChatCompletionsClient; import com.google.adk.models.chat.ChatCompletionsHttpClient; import com.google.common.annotations.VisibleForTesting; @@ -46,17 +44,6 @@ public class ApigeeLlm extends BaseLlm { private static final String GOOGLE_GENAI_USE_VERTEXAI_ENV_VARIABLE_NAME = "GOOGLE_GENAI_USE_VERTEXAI"; private static final String APIGEE_PROXY_URL_ENV_VARIABLE_NAME = "APIGEE_PROXY_URL"; - private static final ImmutableMap TRACKING_HEADERS; - - static { - final String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION; - final String languageLabel = "gl-java/" + JAVA_VERSION.value(); - final String versionHeaderValue = String.format("%s %s", frameworkLabel, languageLabel); - TRACKING_HEADERS = - ImmutableMap.of( - "x-goog-api-client", versionHeaderValue, - "user-agent", versionHeaderValue); - } /** Defines the type of API to be used by the Apigee proxy. */ public enum ApiType { @@ -108,9 +95,16 @@ 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 HttpOptions.Builder httpOptionsBuilder = - HttpOptions.builder().baseUrl(effectiveProxyUrl).headers(TRACKING_HEADERS); + HttpOptions.builder().baseUrl(effectiveProxyUrl).headers(trackingHeaders); String apiVersion = identifyApiVersion(modelName); if (!apiVersion.isEmpty()) { httpOptionsBuilder.apiVersion(apiVersion); @@ -118,7 +112,7 @@ private ApigeeLlm( if (customHeaders != null) { httpOptionsBuilder.headers( ImmutableMap.builder() - .putAll(TRACKING_HEADERS) + .putAll(trackingHeaders) .putAll(customHeaders) .buildOrThrow()); } @@ -134,7 +128,8 @@ private ApigeeLlm( apiClientBuilder.vertexAI(true); } this.apiClient = apiClientBuilder.build(); - this.geminiDelegate = new Gemini(modelName, apiClient); + this.geminiDelegate = + new Gemini(modelName, apiClient, /* clientSendsTrackingHeaders= */ true); this.chatCompletionsClient = null; } diff --git a/core/src/main/java/com/google/adk/models/Gemini.java b/core/src/main/java/com/google/adk/models/Gemini.java index 09402c2b5..7bd524500 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -16,9 +16,6 @@ package com.google.adk.models; -import static com.google.common.base.StandardSystemProperty.JAVA_VERSION; - -import com.google.adk.Version; import com.google.adk.internal.http.HttpClientFactory; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; @@ -60,7 +57,6 @@ public class Gemini extends BaseLlm { private static final Logger logger = LoggerFactory.getLogger(Gemini.class); - private static final ImmutableMap TRACKING_HEADERS; private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executorService) { OkHttpClient client = @@ -75,25 +71,22 @@ private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executor .build(); } - private static Client buildApiKeyClient( - String apiKey, @Nullable ExecutorService executorService) { + // Every client ADK builds starts here, so each sends ADK's tracking tag. + private static Client.Builder adkClientBuilder(@Nullable ExecutorService executorService) { return Client.builder() - .apiKey(apiKey) - .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build()) + .httpOptions(HttpOptions.builder().headers(TrackingHeaders.GENAI_SDK_HEADERS).build()) .clientOptions( - ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build()) - .build(); + ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build()); + } + + private static Client buildApiKeyClient( + String apiKey, @Nullable ExecutorService executorService) { + return adkClientBuilder(executorService).apiKey(apiKey).build(); } private static Client buildVertexClient( VertexCredentials vertexCredentials, @Nullable ExecutorService executorService) { - Client.Builder apiClientBuilder = - Client.builder() - .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build()) - .clientOptions( - ClientOptions.builder() - .customHttpClient(prepareHttpClient(executorService)) - .build()); + Client.Builder apiClientBuilder = adkClientBuilder(executorService); vertexCredentials.project().ifPresent(apiClientBuilder::project); vertexCredentials.location().ifPresent(apiClientBuilder::location); vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials); @@ -101,35 +94,30 @@ private static Client buildVertexClient( } private static Client buildDefaultClient(@Nullable ExecutorService executorService) { - return Client.builder() - .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build()) - .clientOptions( - ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build()) - .build(); - } - - static { - String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION; - String languageLabel = "gl-java/" + JAVA_VERSION.value(); - String versionHeaderValue = String.format("%s %s", frameworkLabel, languageLabel); - - TRACKING_HEADERS = - ImmutableMap.of( - "x-goog-api-client", versionHeaderValue, - "user-agent", versionHeaderValue); + return adkClientBuilder(executorService).build(); } private final Client apiClient; + // True when apiClient already sends ADK's tracking tag. + private final boolean clientSendsTrackingHeaders; /** - * Constructs a new Gemini instance. + * Constructs a new Gemini instance. Each generate request sent through {@code apiClient} carries + * ADK's tracking tag, which the genai SDK appends to the client's tracking headers. If the + * client's own headers already carry an ADK tag, requests send two because the SDK does not + * expose a client's headers. * * @param modelName The name of the Gemini model to use (e.g., "gemini-2.0-flash"). * @param apiClient The genai {@link com.google.genai.Client} instance for making API calls. */ public Gemini(String modelName, Client apiClient) { + this(modelName, apiClient, /* clientSendsTrackingHeaders= */ false); + } + + Gemini(String modelName, Client apiClient, boolean clientSendsTrackingHeaders) { super(modelName); this.apiClient = Objects.requireNonNull(apiClient, "apiClient cannot be null"); + this.clientSendsTrackingHeaders = clientSendsTrackingHeaders; } /** @@ -139,9 +127,10 @@ public Gemini(String modelName, Client apiClient) { * @param apiKey The Google Gemini API key. */ public Gemini(String modelName, String apiKey) { - super(modelName); - Objects.requireNonNull(apiKey, "apiKey cannot be null"); - this.apiClient = buildApiKeyClient(apiKey, null); + this( + modelName, + buildApiKeyClient(Objects.requireNonNull(apiKey, "apiKey cannot be null"), null), + /* clientSendsTrackingHeaders= */ true); } /** @@ -151,9 +140,11 @@ public Gemini(String modelName, String apiKey) { * @param vertexCredentials The Vertex AI credentials to access the Gemini model. */ public Gemini(String modelName, VertexCredentials vertexCredentials) { - super(modelName); - Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null"); - this.apiClient = buildVertexClient(vertexCredentials, null); + this( + modelName, + buildVertexClient( + Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null"), null), + /* clientSendsTrackingHeaders= */ true); } /** @@ -256,13 +247,16 @@ public Gemini build() { if (apiClient != null) { return new Gemini(modelName, apiClient); - } else if (apiKey != null) { - return new Gemini(modelName, buildApiKeyClient(apiKey, httpExecutorService)); + } + Client client; + if (apiKey != null) { + client = buildApiKeyClient(apiKey, httpExecutorService); } else if (vertexCredentials != null) { - return new Gemini(modelName, buildVertexClient(vertexCredentials, httpExecutorService)); + client = buildVertexClient(vertexCredentials, httpExecutorService); } else { - return new Gemini(modelName, buildDefaultClient(httpExecutorService)); + client = buildDefaultClient(httpExecutorService); } + return new Gemini(modelName, client, /* clientSendsTrackingHeaders= */ true); } } @@ -292,6 +286,9 @@ private Flowable sendRequest( logger.trace("Request Contents: {}", llmRequest.contents()); logger.trace("Request Config: {}", config); + if (!clientSendsTrackingHeaders) { + config = withTrackingHeaders(config); + } if (stream) { logger.debug("Sending streaming generateContent request to model {}", effectiveModelName); CompletableFuture> streamFuture = @@ -316,6 +313,23 @@ private Flowable sendRequest( } } + /** + * Returns {@code config} with ADK's tracking tag merged into its request headers, which the genai + * SDK appends to the client's own. + */ + private static GenerateContentConfig withTrackingHeaders(@Nullable GenerateContentConfig config) { + GenerateContentConfig base = config != null ? config : GenerateContentConfig.builder().build(); + HttpOptions httpOptions = base.httpOptions().orElseGet(() -> HttpOptions.builder().build()); + return base.toBuilder() + .httpOptions( + httpOptions.toBuilder() + .headers( + TrackingHeaders.mergeIntoGenAiSdkHeaders( + httpOptions.headers().orElse(ImmutableMap.of()))) + .build()) + .build(); + } + private static LlmResponse withCacheMetadata( LlmResponse response, @Nullable CacheMetadata cacheMetadata) { return cacheMetadata == null diff --git a/core/src/main/java/com/google/adk/models/TrackingHeaders.java b/core/src/main/java/com/google/adk/models/TrackingHeaders.java new file mode 100644 index 000000000..52e0521a9 --- /dev/null +++ b/core/src/main/java/com/google/adk/models/TrackingHeaders.java @@ -0,0 +1,109 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.models; + +import static com.google.common.base.StandardSystemProperty.JAVA_VERSION; +import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static java.util.stream.Collectors.joining; + +import com.google.adk.Version; +import com.google.common.base.Ascii; +import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.ImmutableSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Stream; + +/** + * Headers that let Google APIs attribute model calls to ADK, in the format all ADK SDKs use. A + * header value is a space-separated list of {@code name/version} tags. + */ +final class TrackingHeaders { + + private static final Splitter TAG_SPLITTER = Splitter.on(' ').omitEmptyStrings(); + + /** ADK's own tag. */ + static final String ADK_TAG = "google-adk/" + Version.JAVA_ADK_VERSION; + + /** ADK's tag followed by the Java version tag, for clients that send no language tag. */ + static final String VALUE = ADK_TAG + " gl-java/" + JAVA_VERSION.value(); + + static final ImmutableMap HEADERS = + ImmutableMap.of("x-goog-api-client", VALUE, "user-agent", VALUE); + + /** + * Headers for a genai SDK client, which puts its own {@code gl-java} tag first. Both carry only + * {@link #ADK_TAG}. + */ + static final ImmutableMap GENAI_SDK_HEADERS = + ImmutableMap.of("x-goog-api-client", ADK_TAG, "user-agent", ADK_TAG); + + /** + * Like {@link #mergeInto}, but adds only {@link #ADK_TAG}, for headers that the genai SDK appends + * to its own label. + */ + static ImmutableMap mergeIntoGenAiSdkHeaders(Map headers) { + return mergeInto(headers, ADK_TAG); + } + + /** + * Returns {@code headers} with each tracking header, matched case-insensitively, followed by + * every tag of {@link #VALUE} whose name it lacks. + */ + static ImmutableMap mergeInto(Map headers) { + return mergeInto(headers, VALUE); + } + + private static ImmutableMap mergeInto( + Map headers, String adkTags) { + Map merged = new LinkedHashMap<>(); + Map callerValues = new LinkedHashMap<>(); + headers.forEach( + (name, value) -> { + String trackingName = Ascii.toLowerCase(name); + if (HEADERS.containsKey(trackingName)) { + callerValues.merge(trackingName, value, (first, second) -> first + " " + second); + } else { + merged.put(name, value); + } + }); + for (String name : HEADERS.keySet()) { + merged.put(name, extend(callerValues.getOrDefault(name, ""), adkTags)); + } + return ImmutableMap.copyOf(merged); + } + + /** Returns {@code callerValue} followed by each tag of {@code adkTags} whose name it lacks. */ + private static String extend(String callerValue, String adkTags) { + List tags = TAG_SPLITTER.splitToList(callerValue); + ImmutableSet names = + tags.stream().map(TrackingHeaders::tagName).collect(toImmutableSet()); + return Stream.concat( + tags.stream(), + TAG_SPLITTER.splitToStream(adkTags).filter(tag -> !names.contains(tagName(tag)))) + .collect(joining(" ")); + } + + private static String tagName(String tag) { + int slash = tag.indexOf('/'); + return slash < 0 ? tag : tag.substring(0, slash); + } + + private TrackingHeaders() {} +} 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 ede4243c1..3eb986fc8 100644 --- a/core/src/test/java/com/google/adk/models/ApigeeLlmTest.java +++ b/core/src/test/java/com/google/adk/models/ApigeeLlmTest.java @@ -15,7 +15,11 @@ */ package com.google.adk.models; +import static com.google.common.base.StandardSystemProperty.JAVA_VERSION; +import static com.google.common.collect.ImmutableList.toImmutableList; +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.assertThrows; import static org.junit.Assume.assumeNotNull; import static org.junit.Assume.assumeTrue; @@ -28,13 +32,21 @@ import com.google.adk.models.ApigeeLlm.ApiType; import com.google.adk.models.chat.ChatCompletionsClient; +import com.google.common.base.Splitter; 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.Part; +import com.sun.net.httpserver.Headers; +import com.sun.net.httpserver.HttpServer; import io.reactivex.rxjava3.core.Flowable; +import java.io.OutputStream; +import java.net.InetAddress; +import java.net.InetSocketAddress; import java.util.Map; import java.util.Objects; +import java.util.concurrent.atomic.AtomicReference; import org.junit.Before; import org.junit.Rule; import org.junit.Test; @@ -393,4 +405,95 @@ public void build_malformedModelsWithTrailingSlashes_throwsException() { assertThat(e).hasMessageThat().contains("Invalid model string: " + modelName); } } + + @Test + public void generateContent_genAi_sendsTrackingHeadersOnce() throws Exception { + Headers headers = + headersSentToProxy( + "apigee/gemini/v1beta/gemini-test", + /* customHeaders= */ ImmutableMap.of(), + GENERATE_CONTENT_JSON); + + assertGenAiTrackingHeaders(headers); + } + + @Test + public void generateContent_chatCompletions_sendsTrackingHeaders() throws Exception { + Headers headers = + headersSentToProxy( + "apigee/openai/gpt-4o", /* customHeaders= */ ImmutableMap.of(), CHAT_COMPLETION_JSON); + + for (String name : TrackingHeaders.HEADERS.keySet()) { + assertThat(headers.get(name)).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. + */ + private static Headers headersSentToProxy( + String modelName, Map customHeaders, String responseJson) 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 = responseJson.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 { + ApigeeLlm.Builder builder = + ApigeeLlm.builder() + .modelName(modelName) + .proxyUrl( + "http://" + + InetAddresses.toUriString(server.getAddress().getAddress()) + + ":" + + server.getAddress().getPort()); + builder.customHeaders(customHeaders); + LlmRequest request = + LlmRequest.builder() + .contents( + ImmutableList.of( + Content.builder().role("user").parts(Part.fromText("hello")).build())) + .build(); + + builder.build().generateContent(request, /* stream= */ false).blockingSubscribe(); + + return requestHeaders.get(); + } finally { + server.stop(0); + } + } + + /** + * Asserts both tracking headers as sent by the genai SDK: each starts with the SDK label, + * includes the Java version tag and ADK's tag, and repeats no tag name. + */ + private static void assertGenAiTrackingHeaders(Headers headers) { + for (String name : TrackingHeaders.HEADERS.keySet()) { + ImmutableList tags = + ImmutableList.copyOf(Splitter.on(' ').split(getOnlyElement(headers.get(name)))); + assertThat(tags.get(0)).startsWith("google-genai-sdk/"); + assertThat(tags).containsAtLeast("gl-java/" + JAVA_VERSION.value(), TrackingHeaders.ADK_TAG); + assertThat(tags.stream().map(tag -> tag.split("/", 2)[0]).collect(toImmutableList())) + .containsNoDuplicates(); + } + } + + private static final String GENERATE_CONTENT_JSON = + "{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"hi\"}]}," + + "\"finishReason\":\"STOP\"}]}"; + + private static final String CHAT_COMPLETION_JSON = + "{\"choices\":[{\"message\":{\"role\":\"assistant\",\"content\":\"hi\"}," + + "\"finish_reason\":\"stop\"}]}"; } diff --git a/core/src/test/java/com/google/adk/models/GeminiTrackingHeadersTest.java b/core/src/test/java/com/google/adk/models/GeminiTrackingHeadersTest.java new file mode 100644 index 000000000..c325e29d8 --- /dev/null +++ b/core/src/test/java/com/google/adk/models/GeminiTrackingHeadersTest.java @@ -0,0 +1,269 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.models; + +import static com.google.common.base.StandardSystemProperty.JAVA_VERSION; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.truth.Truth.assertThat; +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.Assume.assumeNotNull; + +import com.google.adk.models.FakeGeminiApi.Kind; +import com.google.common.base.Splitter; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.net.InetAddresses; +import com.google.genai.Client; +import com.google.genai.types.ClientOptions; +import com.google.genai.types.Content; +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.HttpOptions; +import com.google.genai.types.Part; +import com.sun.net.httpserver.HttpServer; +import java.io.OutputStream; +import java.net.InetAddress; +import java.net.InetSocketAddress; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Function; +import java.util.function.Supplier; +import okhttp3.Headers; +import okhttp3.OkHttpClient; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** + * Tests the tracking headers sent on Gemini requests after the genai SDK merges them. Both {@code + * user-agent} and {@code x-goog-api-client} carry one ADK tag and the Java version tag, whoever + * built the client. + */ +@RunWith(JUnit4.class) +public final class GeminiTrackingHeadersTest { + + private final FakeGeminiApi api = new FakeGeminiApi(); + + @Test + public void generateContent_callerClient_sendsTrackingHeaders() { + generate(new Gemini("gemini-test", api.client()), request(), /* stream= */ false); + + assertTrackingHeaders(onlyGenerateHeaders()::get, TrackingHeaders.ADK_TAG); + } + + @Test + public void generateContentStream_callerClient_sendsTrackingHeaders() { + generate(new Gemini("gemini-test", api.client()), request(), /* stream= */ true); + + assertTrackingHeaders(onlyGenerateHeaders()::get, TrackingHeaders.ADK_TAG); + } + + @Test + public void generateContent_adkBuiltApiKeyClient_sendsTrackingHeaders() throws Exception { + assertTrackingHeaders( + headersSentByAdkBuiltClient(() -> new Gemini("gemini-test", "test-api-key")), + TrackingHeaders.ADK_TAG); + } + + @Test + public void generateContent_adkBuiltBuilderClient_sendsTrackingHeaders() throws Exception { + assertTrackingHeaders( + headersSentByAdkBuiltClient( + () -> Gemini.builder().modelName("gemini-test").apiKey("test-api-key").build()), + TrackingHeaders.ADK_TAG); + } + + @Test + public void generateContent_adkBuiltRegistryClient_sendsTrackingHeaders() throws Exception { + assumeNotNull(System.getenv("GOOGLE_API_KEY")); + + assertTrackingHeaders( + headersSentByAdkBuiltClient(() -> LlmRegistry.getLlm("gemini-tracking-headers-test")), + TrackingHeaders.ADK_TAG); + } + + @Test + public void generateContent_adkBuiltVertexCredentialsClient_sendsTrackingHeaders() + throws Exception { + assumeNotNull(System.getenv("GOOGLE_API_KEY")); + + // Empty credentials let the environment configure the backend, here the Gemini API key. + assertTrackingHeaders( + headersSentByAdkBuiltClient( + () -> new Gemini("gemini-test", VertexCredentials.builder().build())), + TrackingHeaders.ADK_TAG); + } + + @Test + public void generateContent_callerClientUserAgent_isExtendedWithAdkTag() { + generate( + new Gemini("gemini-test", clientWithHeaders(ImmutableMap.of("user-agent", "my-app/1.0"))), + request(), + /* stream= */ false); + + Headers headers = onlyGenerateHeaders(); + assertThat(headers.get("user-agent")).endsWith(" my-app/1.0 " + TrackingHeaders.ADK_TAG); + assertTrackingHeaders(headers::get, TrackingHeaders.ADK_TAG); + } + + @Test + public void generateContent_callerRequestHeaders_areExtendedWithAdkTag() { + generate( + new Gemini("gemini-test", api.client()), + requestWithHeaders(ImmutableMap.of("User-Agent", "my-app/1.0", "X-Custom", "v")), + /* stream= */ false); + + Headers headers = onlyGenerateHeaders(); + assertThat(headers.get("user-agent")).endsWith(" my-app/1.0 " + TrackingHeaders.ADK_TAG); + assertThat(headers.get("X-Custom")).isEqualTo("v"); + assertTrackingHeaders(headers::get, TrackingHeaders.ADK_TAG); + } + + @Test + public void generateContent_callerClient_keepsOtherRequestOptions() { + LlmRequest request = + request().toBuilder() + .config( + GenerateContentConfig.builder() + .temperature(0.5f) + .httpOptions(HttpOptions.builder().timeout(30_000).build()) + .build()) + .build(); + + generate(new Gemini("gemini-test", api.client()), request, /* stream= */ false); + + Headers headers = onlyGenerateHeaders(); + assertThat(headers.get("X-Server-Timeout")).isEqualTo("30"); + assertThat(api.bodies(Kind.GENERATE).get(0).at("/generationConfig/temperature").asDouble()) + .isEqualTo(0.5); + assertTrackingHeaders(headers::get, TrackingHeaders.ADK_TAG); + } + + @Test + public void generateContent_callerRequestHeaderWithAdkTag_isNotExtended() { + generate( + new Gemini("gemini-test", api.client()), + requestWithHeaders(ImmutableMap.of("user-agent", "my-app/1.0 google-adk/0.0.1")), + /* stream= */ false); + + Headers headers = onlyGenerateHeaders(); + assertThat(headers.get("user-agent")).endsWith(" my-app/1.0 google-adk/0.0.1"); + assertTrackingHeaders(headers::get, "google-adk/0.0.1"); + } + + /** + * Sends one request through a {@link Gemini} that builds its own client pointed at a local server + * with {@link Client#setDefaultBaseUrls}, and returns the headers the server received. + */ + private static Function headersSentByAdkBuiltClient( + Supplier model) throws Exception { + AtomicReference received = new AtomicReference<>(); + HttpServer server = + HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + server.createContext( + "/", + exchange -> { + received.set(exchange.getRequestHeaders()); + byte[] body = GENERATE_CONTENT_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(); + Optional serverUrl = + Optional.of( + "http://" + + InetAddresses.toUriString(server.getAddress().getAddress()) + + ":" + + server.getAddress().getPort()); + // The environment can select Vertex AI instead of the Gemini API, so redirect both backends. + Client.setDefaultBaseUrls(serverUrl, serverUrl); + try { + generate(model.get(), request(), /* stream= */ false); + return name -> received.get().getFirst(name); + } finally { + Client.setDefaultBaseUrls(Optional.empty(), Optional.empty()); + server.stop(0); + } + } + + private Client clientWithHeaders(Map headers) { + return Client.builder() + .apiKey("test-api-key") + .vertexAI(false) + .httpOptions(HttpOptions.builder().headers(headers).build()) + .clientOptions( + ClientOptions.builder() + .customHttpClient(new OkHttpClient.Builder().addInterceptor(api).build()) + .build()) + .build(); + } + + private static LlmRequest request() { + return LlmRequest.builder() + .contents( + ImmutableList.of( + Content.builder().role("user").parts(Part.fromText("What is the answer?")).build())) + .build(); + } + + private static LlmRequest requestWithHeaders(ImmutableMap headers) { + return request().toBuilder() + .config( + GenerateContentConfig.builder() + .httpOptions(HttpOptions.builder().headers(headers).build()) + .build()) + .build(); + } + + private static void generate(BaseLlm model, LlmRequest request, boolean stream) { + assertThat(model.generateContent(request, stream).toList().blockingGet()).isNotEmpty(); + } + + private Headers onlyGenerateHeaders() { + ImmutableList headers = api.headers(Kind.GENERATE); + assertThat(headers).hasSize(1); + return headers.get(0); + } + + /** + * Asserts both tracking headers as sent: each starts with the genai SDK label, includes the Java + * version tag and one ADK tag, and repeats no tag name. The {@code x-goog-api-client} ADK tag is + * always {@link TrackingHeaders#ADK_TAG}; the {@code user-agent} tag is {@code userAgentAdkTag}. + */ + private static void assertTrackingHeaders( + Function headers, String userAgentAdkTag) { + ImmutableMap adkTags = + ImmutableMap.of( + "user-agent", userAgentAdkTag, "x-goog-api-client", TrackingHeaders.ADK_TAG); + adkTags.forEach( + (name, adkTag) -> { + ImmutableList tags = + ImmutableList.copyOf(Splitter.on(' ').split(headers.apply(name))); + assertThat(tags.get(0)).startsWith("google-genai-sdk/"); + assertThat(tags).containsAtLeast("gl-java/" + JAVA_VERSION.value(), adkTag); + assertThat(tags.stream().map(tag -> tag.split("/", 2)[0]).collect(toImmutableList())) + .containsNoDuplicates(); + }); + } + + private static final String GENERATE_CONTENT_JSON = + "{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"hi\"}]}," + + "\"finishReason\":\"STOP\"}]}"; +} diff --git a/core/src/test/java/com/google/adk/models/TrackingHeadersTest.java b/core/src/test/java/com/google/adk/models/TrackingHeadersTest.java new file mode 100644 index 000000000..9a934739c --- /dev/null +++ b/core/src/test/java/com/google/adk/models/TrackingHeadersTest.java @@ -0,0 +1,50 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.models; + +import static com.google.common.base.StandardSystemProperty.JAVA_VERSION; +import static com.google.common.truth.Truth.assertThat; + +import com.google.common.collect.ImmutableMap; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class TrackingHeadersTest { + + @Test + public void mergeInto_addsOnlyTheTagsACallerValueLacks() { + ImmutableMap merged = + TrackingHeaders.mergeInto( + ImmutableMap.of( + "User-Agent", "my-app/1.0 gl-java/11", + "X-Goog-Api-Client", "my-lib/2.0 google-adk/0.0.1", + "X-Custom", "v")); + + assertThat(merged) + .containsExactly( + "X-Custom", "v", + "x-goog-api-client", "my-lib/2.0 google-adk/0.0.1 gl-java/" + JAVA_VERSION.value(), + "user-agent", "my-app/1.0 gl-java/11 " + TrackingHeaders.ADK_TAG); + } + + @Test + public void mergeInto_noCallerValues_addsTrackingHeaders() { + assertThat(TrackingHeaders.mergeInto(ImmutableMap.of())).isEqualTo(TrackingHeaders.HEADERS); + } +}