From 70254f4e6dc67a09dec3d2640e23cd9a4ca0f00f Mon Sep 17 00:00:00 2001
From: Matthew Tang <tangmatthew@google.com>
Date: Wed, 13 May 2026 10:35:27 -0700
Subject: [PATCH] chore: Update java-genai header merging logic to handle
 vertex-genai-modules case

PiperOrigin-RevId: 914939877
---
 src/main/java/com/google/genai/ApiClient.java | 32 ++++++++++++++++++-
 .../com/google/genai/ReplayApiClient.java     |  4 +--
 .../com/google/genai/HttpApiClientTest.java   | 20 ++++++++++++
 3 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/src/main/java/com/google/genai/ApiClient.java b/src/main/java/com/google/genai/ApiClient.java
index a59c505c677..f763457f626 100644
--- a/src/main/java/com/google/genai/ApiClient.java
+++ b/src/main/java/com/google/genai/ApiClient.java
@@ -714,6 +714,36 @@ private Optional<Map<String, String>> getTimeoutHeader(HttpOptions httpOptionsTo
     return Optional.empty();
   }
 
+  /**
+   * Merges two user agent values, handling a special case for vertex-genai-modules.
+   *
+   * <p>Example:
+   * <li>val1 = google-genai-sdk/1.42.0 gl-java/22.0.2
+   * <li>val2 = vertex-genai-modules/1.2.3
+   * <li>Result: google-genai-sdk/1.42.0+vertex-genai-modules/1.2.3 gl-java/22.0.2
+   *
+   * @param val1 The first value.
+   * @param val2 The second value.
+   * @return The merged user agent value.
+   */
+  private String mergeUserAgentValues(String val1, String val2) {
+    String extensionRegex = "^vertex-genai-modules/[^ ]+.*";
+    boolean val1Matches = val1.matches(extensionRegex);
+    boolean val2Matches = val2.matches(extensionRegex);
+
+    if (val1Matches || val2Matches) {
+      String extensionHeader = val1Matches ? val1 : val2;
+      String baseHeader = val1Matches ? val2 : val1;
+
+      // Find google-genai-sdk/x.y.z in baseHeader
+      if (baseHeader.contains("google-genai-sdk/")) {
+        // $1 refers to the matched (google-genai-sdk/[^ ]+) group
+        return baseHeader.replaceFirst("(google-genai-sdk/[^ ]+)", "$1+" + extensionHeader);
+      }
+    }
+    return val1 + " " + val2;
+  }
+
   /**
    * Merges the http options to the client's http options.
    *
@@ -752,7 +782,7 @@ HttpOptions mergeHttpOptions(HttpOptions httpOptionsToApply) {
                   (val1, val2) -> {
                     if (entry.getKey().equals("user-agent")
                         || entry.getKey().equals("x-goog-api-client")) {
-                      return val1 + " " + val2;
+                      return mergeUserAgentValues(val1, val2);
                     }
                     return val2;
                   }));
diff --git a/src/main/java/com/google/genai/ReplayApiClient.java b/src/main/java/com/google/genai/ReplayApiClient.java
index d1a3622ed2e..57f3abfc486 100644
--- a/src/main/java/com/google/genai/ReplayApiClient.java
+++ b/src/main/java/com/google/genai/ReplayApiClient.java
@@ -206,9 +206,7 @@ private void matchRequest(ReplayRequest replayRequest, Request actualRequest) {
       actualHeaders.put(entry.getKey(), String.join(" ", entry.getValue()));
     }
     actualHeaders = redactRequestHeaders(actualHeaders);
-    // TODO(b/491838117): Enable header checks for vertex extension modules
-    if (!equalsIgnoreKeyCase(replayHeaders, actualHeaders)
-        && !this.replaysDirectory.contains("vertex_sdk_genai_replays")) {
+    if (!equalsIgnoreKeyCase(replayHeaders, actualHeaders)) {
       throw new AssertionError(
           String.format(
               "Request headers mismatch:\nReplay: %s\nActual: %s", replayHeaders, actualHeaders));
diff --git a/src/test/java/com/google/genai/HttpApiClientTest.java b/src/test/java/com/google/genai/HttpApiClientTest.java
index 1b801fd0b73..baa45871b5c 100644
--- a/src/test/java/com/google/genai/HttpApiClientTest.java
+++ b/src/test/java/com/google/genai/HttpApiClientTest.java
@@ -605,6 +605,26 @@ public void testInitHttpClientCustomUserAgent() throws Exception {
             + "' does not match the expected format.");
   }
 
+  @Test
+  public void testInitHttpClientExtensionsUserAgentMerge() throws Exception {
+    String vertexHeader = "vertex-genai-modules/1.2.3";
+    ImmutableMap<String, String> trackingHeaders =
+        ImmutableMap.of(
+            "x-goog-api-client", vertexHeader,
+            "user-agent", vertexHeader);
+    HttpOptions httpOptions = HttpOptions.builder().headers(trackingHeaders).build();
+    HttpApiClient client =
+        new HttpApiClient(Optional.of(API_KEY), Optional.of(httpOptions), Optional.empty());
+
+    String userAgent = client.httpOptions.headers().get().get("user-agent");
+    String apiClient = client.httpOptions.headers().get().get("x-goog-api-client");
+
+    String expectedRegex = "^google-genai-sdk/[^ ]+vertex-genai-modules/1.2.3 gl-java/[^ ]+$";
+
+    assertTrue(userAgent.matches(expectedRegex), "User-Agent did not match: " + userAgent);
+    assertTrue(apiClient.matches(expectedRegex), "x-goog-api-client did not match: " + apiClient);
+  }
+
   @Test
   public void testInitHttpClientMldev() throws Exception {
     HttpOptions httpOptions =