diff --git a/src/main/java/com/google/genai/ApiClient.java b/src/main/java/com/google/genai/ApiClient.java index 0e5bbc384df..22448bce54e 100644 --- a/src/main/java/com/google/genai/ApiClient.java +++ b/src/main/java/com/google/genai/ApiClient.java @@ -438,6 +438,7 @@ protected Request buildRequest( requestHttpOptions.ifPresent( httpOptions -> { + requestBuilder.tag(HttpOptions.class, httpOptions); if (httpOptions.retryOptions().isPresent()) { requestBuilder.tag(HttpRetryOptions.class, mergedHttpOptions.retryOptions().get()); } diff --git a/src/main/java/com/google/genai/HttpApiClient.java b/src/main/java/com/google/genai/HttpApiClient.java index 1d253ff24fb..fcc8a887f11 100644 --- a/src/main/java/com/google/genai/HttpApiClient.java +++ b/src/main/java/com/google/genai/HttpApiClient.java @@ -22,10 +22,12 @@ import com.google.genai.types.ClientOptions; import com.google.genai.types.HttpOptions; import java.io.IOException; +import java.time.Duration; import java.util.Optional; import java.util.concurrent.CompletableFuture; import okhttp3.Call; import okhttp3.Callback; +import okhttp3.OkHttpClient; import okhttp3.Request; import okhttp3.Response; @@ -78,7 +80,16 @@ public HttpApiResponse request( /** Executes the given HTTP request. */ private HttpApiResponse executeRequest(Request request) { try { - return new HttpApiResponse(httpClient.newCall(request).execute()); + OkHttpClient client = httpClient; + HttpOptions requestOptions = request.tag(HttpOptions.class); + if (requestOptions != null && requestOptions.timeout().isPresent()) { + client = + httpClient + .newBuilder() + .callTimeout(Duration.ofMillis(requestOptions.timeout().get())) + .build(); + } + return new HttpApiResponse(client.newCall(request).execute()); } catch (IOException e) { throw new GenAiIOException("Failed to execute HTTP request.", e); } @@ -114,7 +125,17 @@ public CompletableFuture asyncRequest( private CompletableFuture asyncExecuteRequest(Request request) { CompletableFuture future = new CompletableFuture<>(); - httpClient + OkHttpClient client = httpClient; + HttpOptions requestOptions = request.tag(HttpOptions.class); + if (requestOptions != null && requestOptions.timeout().isPresent()) { + client = + httpClient + .newBuilder() + .callTimeout(Duration.ofMillis(requestOptions.timeout().get())) + .build(); + } + + client .newCall(request) .enqueue( new Callback() { diff --git a/src/test/java/com/google/genai/HttpApiClientTest.java b/src/test/java/com/google/genai/HttpApiClientTest.java index 8ca839655e0..5b28bf0fa2b 100644 --- a/src/test/java/com/google/genai/HttpApiClientTest.java +++ b/src/test/java/com/google/genai/HttpApiClientTest.java @@ -222,6 +222,31 @@ public void testRequestWithHttpOptions() throws Exception { assertEquals("header", capturedRequest.header("test")); } + @Test + public void testRequestWithHttpOptions_overrideTimeout() throws Exception { + // Arrange + HttpApiClient client = + new HttpApiClient(Optional.of(API_KEY), Optional.empty(), Optional.empty()); + setMockClient(client); + + okhttp3.OkHttpClient.Builder mockBuilder = Mockito.mock(okhttp3.OkHttpClient.Builder.class); + when(mockHttpClient.newBuilder()).thenReturn(mockBuilder); + when(mockBuilder.callTimeout(any(java.time.Duration.class))).thenReturn(mockBuilder); + when(mockBuilder.build()).thenReturn(mockHttpClient); + + Optional requestOptions = + Optional.of(HttpOptions.builder().timeout(30000).build()); + + // Act + client.request("POST", TEST_PATH, TEST_REQUEST_JSON, requestOptions); + + // Assert + verify(mockHttpClient).newBuilder(); + verify(mockBuilder).callTimeout(java.time.Duration.ofMillis(30000)); + verify(mockBuilder).build(); + verify(mockHttpClient).newCall(any(okhttp3.Request.class)); + } + @Test public void testRequestWithHttpOptions_extraBody_addNewKey() throws Exception { // Arrange