Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 15 additions & 5 deletions src/main/java/com/google/genai/ApiClient.java
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
package com.google.genai;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.collect.ImmutableMap.toImmutableMap;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
Expand All @@ -39,6 +38,7 @@
import java.net.URI;
import java.net.URISyntaxException;
import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
Expand Down Expand Up @@ -660,10 +660,20 @@ HttpOptions mergeHttpOptions(HttpOptions httpOptionsToApply) {
.entrySet()
.stream()),
httpOptionsToApply.headers().orElse(ImmutableMap.of()).entrySet().stream());
Map<String, String> mergedHeaders =
headersStream.collect(
toImmutableMap(Map.Entry::getKey, Map.Entry::getValue, (val1, val2) -> val2));
mergedHttpOptionsBuilder.headers(mergedHeaders);
final Map<String, String> mergedHeaders = new HashMap<>();
headersStream.forEach(
entry ->
mergedHeaders.merge(
entry.getKey(),
entry.getValue(),
(val1, val2) -> {
if (entry.getKey().equals("user-agent")
|| entry.getKey().equals("x-goog-api-client")) {
return val1 + " " + val2;
}
return val2;
}));
mergedHttpOptionsBuilder.headers(ImmutableMap.copyOf(mergedHeaders));
}

return mergedHttpOptionsBuilder.build();
Expand Down
4 changes: 3 additions & 1 deletion src/main/java/com/google/genai/ReplayApiClient.java
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,9 @@ private void matchRequest(ReplayRequest replayRequest, Request actualRequest) {
actualHeaders.put(entry.getKey(), String.join(" ", entry.getValue()));
}
actualHeaders = redactRequestHeaders(actualHeaders);
if (!equalsIgnoreKeyCase(replayHeaders, actualHeaders)) {
// TODO(b/491838117): Enable header checks for vertex extension modules
if (!equalsIgnoreKeyCase(replayHeaders, actualHeaders)
&& !this.replaysDirectory.contains("vertex_sdk_genai_replays")) {
throw new AssertionError(
String.format(
"Request headers mismatch:\nReplay: %s\nActual: %s", replayHeaders, actualHeaders));
Expand Down
54 changes: 54 additions & 0 deletions src/test/java/com/google/genai/HttpApiClientTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -551,6 +551,60 @@ public void testAsyncRequest_Failure() throws Exception {
assertEquals(networkError, cause.getCause());
}

@Test
public void testInitHttpClientCustomUserAgent() throws Exception {
HttpApiClient client1 =
new HttpApiClient(Optional.of(API_KEY), Optional.empty(), Optional.empty());

String userAgentRegex1 = "^google-genai-sdk/[^ ]+ gl-java/[^ ]+$";
assertTrue(
client1.httpOptions.headers().get().get("user-agent").matches(userAgentRegex1),
"The User-Agent string '"
+ client1.httpOptions.headers().get().get("user-agent")
+ "' does not match the expected format.");
assertTrue(
client1.httpOptions.headers().get().get("x-goog-api-client").matches(userAgentRegex1),
"The x-goog-api-client string '"
+ client1.httpOptions.headers().get().get("x-goog-api-client")
+ "' does not match the expected format.");

ImmutableMap<String, String> trackingHeaders =
ImmutableMap.of(
"x-goog-api-client", "library-name/1.2.3",
"user-agent", "library-name/1.2.3");
HttpOptions httpOptions = HttpOptions.builder().headers(trackingHeaders).build();
HttpApiClient client2 =
new HttpApiClient(Optional.of(API_KEY), Optional.of(httpOptions), Optional.empty());

// Ex. google-genai-sdk/1.42.0 gl-java/22.0.2 library-name/1.2.3
String userAgentRegex2 = "^google-genai-sdk/[^ ]+ gl-java/[^ ]+ library-name/1.2.3$";
assertTrue(
client2.httpOptions.headers().get().get("user-agent").matches(userAgentRegex2),
"The User-Agent string '"
+ client2.httpOptions.headers().get().get("user-agent")
+ "' does not match the expected format.");
assertTrue(
client2.httpOptions.headers().get().get("x-goog-api-client").matches(userAgentRegex2),
"The x-goog-api-client string '"
+ client2.httpOptions.headers().get().get("x-goog-api-client")
+ "' does not match the expected format.");

// Ensure that new clients still have the default tracking headers.
HttpApiClient client3 =
new HttpApiClient(Optional.of(API_KEY), Optional.empty(), Optional.empty());

assertTrue(
client3.httpOptions.headers().get().get("user-agent").matches(userAgentRegex1),
"The User-Agent string '"
+ client3.httpOptions.headers().get().get("user-agent")
+ "' does not match the expected format.");
assertTrue(
client3.httpOptions.headers().get().get("x-goog-api-client").matches(userAgentRegex1),
"The x-goog-api-client string '"
+ client3.httpOptions.headers().get().get("x-goog-api-client")
+ "' does not match the expected format.");
}

@Test
public void testInitHttpClientMldev() throws Exception {
HttpOptions httpOptions =
Expand Down
Loading