chore: tag Gemini requests sent through a caller-supplied genai Client

PiperOrigin-RevId: 1002229130

Author: damianmomotgoogle

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<String, String> 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 {
@@ -110,15 +97,15 @@ private ApigeeLlm(
 
     // Build the Client
     HttpOptions.Builder httpOptionsBuilder =
-        HttpOptions.builder().baseUrl(effectiveProxyUrl).headers(TRACKING_HEADERS);
+        HttpOptions.builder().baseUrl(effectiveProxyUrl).headers(TrackingHeaders.HEADERS);
     String apiVersion = identifyApiVersion(modelName);
     if (!apiVersion.isEmpty()) {
       httpOptionsBuilder.apiVersion(apiVersion);
     }
     if (customHeaders != null) {
       httpOptionsBuilder.headers(
           ImmutableMap.<String, String>builder()
-              .putAll(TRACKING_HEADERS)
+              .putAll(TrackingHeaders.HEADERS)
               .putAll(customHeaders)
               .buildOrThrow());
     }
@@ -134,7 +121,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;
     }
 

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<String, String> TRACKING_HEADERS;
 
   private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executorService) {
     OkHttpClient client =
@@ -79,7 +75,7 @@ private static Client buildApiKeyClient(
       String apiKey, @Nullable ExecutorService executorService) {
     return Client.builder()
         .apiKey(apiKey)
-        .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
+        .httpOptions(HttpOptions.builder().headers(TrackingHeaders.HEADERS).build())
         .clientOptions(
             ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build())
         .build();
@@ -89,7 +85,7 @@ private static Client buildVertexClient(
       VertexCredentials vertexCredentials, @Nullable ExecutorService executorService) {
     Client.Builder apiClientBuilder =
         Client.builder()
-            .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
+            .httpOptions(HttpOptions.builder().headers(TrackingHeaders.HEADERS).build())
             .clientOptions(
                 ClientOptions.builder()
                     .customHttpClient(prepareHttpClient(executorService))
@@ -102,34 +98,31 @@ private static Client buildVertexClient(
 
   private static Client buildDefaultClient(@Nullable ExecutorService executorService) {
     return Client.builder()
-        .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
+        .httpOptions(HttpOptions.builder().headers(TrackingHeaders.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);
-  }
-
   private final Client apiClient;
+  // True when apiClient already sends TrackingHeaders.HEADERS.
+  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 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 +132,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 +145,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 +252,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 +291,9 @@ private Flowable<LlmResponse> 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<ResponseStream<GenerateContentResponse>> streamFuture =
@@ -316,6 +318,21 @@ private Flowable<LlmResponse> sendRequest(
     }
   }
 
+  /**
+   * Returns {@code config} with ADK's tracking headers 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.mergeInto(httpOptions.headers().orElse(ImmutableMap.of())))
+                .build())
+        .build();
+  }
+
   private static LlmResponse withCacheMetadata(
       LlmResponse response, @Nullable CacheMetadata cacheMetadata) {
     return cacheMetadata == null

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

@@ -0,0 +1,73 @@
+/*
+ * 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 com.google.adk.Version;
+import com.google.common.base.Ascii;
+import com.google.common.base.Splitter;
+import com.google.common.collect.ImmutableMap;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import org.jspecify.annotations.Nullable;
+
+/** Headers that let Google APIs attribute model calls to ADK, in the format all ADK SDKs use. */
+final class TrackingHeaders {
+
+  private static final String ADK_TOKEN_PREFIX = "google-adk/";
+
+  static final String VALUE =
+      ADK_TOKEN_PREFIX + Version.JAVA_ADK_VERSION + " gl-java/" + JAVA_VERSION.value();
+
+  static final ImmutableMap<String, String> HEADERS =
+      ImmutableMap.of("x-goog-api-client", VALUE, "user-agent", VALUE);
+
+  /**
+   * Returns {@code headers} with {@link #HEADERS} merged in. Any tracking header already set,
+   * regardless of letter case, keeps its value and is extended with ADK's value unless it already
+   * has an ADK token.
+   */
+  static ImmutableMap<String, String> mergeInto(Map<String, String> headers) {
+    Map<String, String> merged = new LinkedHashMap<>();
+    Map<String, String> 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);
+          }
+        });
+    HEADERS.forEach((name, value) -> merged.put(name, extend(callerValues.get(name))));
+    return ImmutableMap.copyOf(merged);
+  }
+
+  private static String extend(@Nullable String callerValue) {
+    if (callerValue == null || callerValue.isBlank()) {
+      return VALUE;
+    }
+    boolean hasAdkToken =
+        Splitter.on(' ')
+            .splitToStream(callerValue)
+            .anyMatch(token -> token.startsWith(ADK_TOKEN_PREFIX));
+    return hasAdkToken ? callerValue : callerValue + " " + VALUE;
+  }
+
+  private TrackingHeaders() {}
+}

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

@@ -15,7 +15,9 @@
  */
 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.assertThrows;
 import static org.junit.Assume.assumeNotNull;
 import static org.junit.Assume.assumeTrue;
@@ -28,13 +30,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 +403,52 @@ public void build_malformedModelsWithTrailingSlashes_throwsException() {
       assertThat(e).hasMessageThat().contains("Invalid model string: " + modelName);
     }
   }
+
+  @Test
+  public void generateContent_genAi_sendsTrackingHeadersOnce() 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 = 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();
+    try {
+      ApigeeLlm llm =
+          ApigeeLlm.builder()
+              .modelName("apigee/gemini/v1beta/gemini-test")
+              .proxyUrl(
+                  "http://"
+                      + InetAddresses.toUriString(server.getAddress().getAddress())
+                      + ":"
+                      + server.getAddress().getPort())
+              .build();
+      LlmRequest request =
+          LlmRequest.builder()
+              .contents(ImmutableList.of(Content.builder().parts(Part.fromText("hello")).build()))
+              .build();
+
+      llm.generateContent(request, /* stream= */ false).blockingSubscribe();
+
+      for (String name : TrackingHeaders.HEADERS.keySet()) {
+        String value = getOnlyElement(requestHeaders.get().get(name));
+        assertThat(Splitter.on(' ').splitToStream(value).filter(t -> t.startsWith("google-adk/")))
+            .hasSize(1);
+      }
+    } finally {
+      server.stop(0);
+    }
+  }
+
+  private static final String GENERATE_CONTENT_JSON =
+      "{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"hi\"}]},"
+          + "\"finishReason\":\"STOP\"}]}";
 }