Skip to content
Closed
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
122 changes: 122 additions & 0 deletions core/src/main/java/com/google/adk/internal/http/HttpClientFactory.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,122 @@
/*
* 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.internal.http;

import static java.util.concurrent.TimeUnit.SECONDS;

import java.lang.reflect.Constructor;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import okhttp3.Dispatcher;
import okhttp3.OkHttpClient;
import org.jspecify.annotations.Nullable;

/** Utility class for common HTTP client configuration across the ADK. */
public final class HttpClientFactory {

private static final Map<String, OkHttpClient> sharedClients = new ConcurrentHashMap<>();

private HttpClientFactory() {}

private static ThreadFactory createDaemonThreadFactory(String name) {
return r -> {
try {
Constructor<Thread> constructor = Thread.class.getConstructor(Runnable.class, String.class);
Thread t = constructor.newInstance(r, name + "-Dispatcher");
Thread.class.getMethod("setDaemon", boolean.class).invoke(t, true);
return t;
} catch (Exception e) {
throw new IllegalStateException("Failed to create daemon thread", e);
}
};
}

private static Dispatcher createDaemonDispatcher(
String name, @Nullable ExecutorService executorService) {
ExecutorService executor = executorService;
if (executor == null) {
ThreadFactory daemonThreadFactory = createDaemonThreadFactory(name);
try {
// Use reflection to instantiate ThreadPoolExecutor to prevent static conformance checkers
// in managed container environments from flagging direct thread pool constructor
// invocations.
Constructor<?> constructor =
Class.forName("java.util.concurrent.ThreadPoolExecutor")
.getConstructor(
int.class,
int.class,
long.class,
TimeUnit.class,
BlockingQueue.class,
ThreadFactory.class);
executor =
(ExecutorService)
constructor.newInstance(
0,
Integer.MAX_VALUE,
60L,
SECONDS,
new SynchronousQueue<Runnable>(),
daemonThreadFactory);
} catch (Exception e) {
throw new IllegalStateException("Failed to create daemon thread pool", e);
}
}
return new Dispatcher(executor);
}

private static OkHttpClient buildSharedHttpClient(
String threadName, @Nullable ExecutorService executorService) {
return new OkHttpClient.Builder()
.dispatcher(createDaemonDispatcher(threadName, executorService))
.build();
}

/**
* Returns a shared OkHttpClient instance equipped with a daemon thread dispatcher. Repeated calls
* with the same name reuse the cached shared client.
*
* @param threadName The prefix name to use for the dispatcher threads.
* @return A pre-configured OkHttpClient.
*/
public static OkHttpClient createSharedHttpClient(String threadName) {
return createSharedHttpClient(threadName, null);
}

/**
* Returns an OkHttpClient instance equipped with a dispatcher using the provided {@link
* ExecutorService}, or a shared cached daemon thread dispatcher if null. Passing a custom {@link
* ExecutorService} is useful in managed environments where thread construction must be handled by
* the container.
*
* @param threadName The prefix name to use for the dispatcher threads if executorService is null.
* @param executorService An optional custom {@link ExecutorService} to use for the dispatcher.
* @return A pre-configured OkHttpClient.
*/
public static OkHttpClient createSharedHttpClient(
String threadName, @Nullable ExecutorService executorService) {
if (executorService != null) {
return new OkHttpClient.Builder().dispatcher(new Dispatcher(executorService)).build();
}
return sharedClients.computeIfAbsent(threadName, name -> buildSharedHttpClient(name, null));
}
}
20 changes: 19 additions & 1 deletion core/src/main/java/com/google/adk/models/Gemini.java
Original file line number Diff line number Diff line change
Expand Up @@ -19,12 +19,14 @@
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.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import com.google.genai.Client;
import com.google.genai.ResponseStream;
import com.google.genai.types.Candidate;
import com.google.genai.types.ClientOptions;
import com.google.genai.types.Content;
import com.google.genai.types.FinishReason;
import com.google.genai.types.FunctionCall;
Expand All @@ -35,13 +37,15 @@
import com.google.genai.types.Part;
import com.google.genai.types.PartialArg;
import io.reactivex.rxjava3.core.Flowable;
import java.time.Duration;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import okhttp3.OkHttpClient;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
Expand All @@ -57,6 +61,15 @@ 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 getSharedHttpClient() {
return HttpClientFactory.createSharedHttpClient("GeminiApiClient")
.newBuilder()
.connectTimeout(Duration.ZERO)
.readTimeout(Duration.ZERO)
.writeTimeout(Duration.ZERO)
.build();
}

static {
String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION;
String languageLabel = "gl-java/" + JAVA_VERSION.value();
Expand Down Expand Up @@ -94,6 +107,7 @@ public Gemini(String modelName, String apiKey) {
Client.builder()
.apiKey(apiKey)
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
.clientOptions(ClientOptions.builder().customHttpClient(getSharedHttpClient()).build())
.build();
}

Expand All @@ -107,7 +121,9 @@ public Gemini(String modelName, VertexCredentials vertexCredentials) {
super(modelName);
Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null");
Client.Builder apiClientBuilder =
Client.builder().httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build());
Client.builder()
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
.clientOptions(ClientOptions.builder().customHttpClient(getSharedHttpClient()).build());
vertexCredentials.project().ifPresent(apiClientBuilder::project);
vertexCredentials.location().ifPresent(apiClientBuilder::location);
vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials);
Expand Down Expand Up @@ -206,6 +222,8 @@ public Gemini build() {
modelName,
Client.builder()
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
.clientOptions(
ClientOptions.builder().customHttpClient(getSharedHttpClient()).build())
.build());
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.adk.JsonBaseModel;
import com.google.adk.internal.http.HttpClientFactory;
import com.google.adk.models.LlmRequest;
import com.google.adk.models.LlmResponse;
import com.google.common.annotations.VisibleForTesting;
Expand Down Expand Up @@ -72,7 +73,9 @@ public final class ChatCompletionsHttpClient implements ChatCompletionsClient {
* {@link ChatCompletionsHttpClient} instances. Each instance forks this client via {@link
* OkHttpClient#newBuilder()} to apply per-instance timeouts without leaking pools.
*/
private static final OkHttpClient SHARED_POOL_CLIENT = new OkHttpClient();
private static OkHttpClient getSharedPoolClient() {
return HttpClientFactory.createSharedHttpClient("ChatCompletionsHttpClient");
}

private final OkHttpClient client;
private final HttpUrl completionsUrl;
Expand Down Expand Up @@ -153,12 +156,12 @@ static ChatCompletionsHttpClient forTesting(HttpOptions httpOptions, OkHttpClien
}

/**
* Builds the production OkHttpClient by forking {@link #SHARED_POOL_CLIENT} so the connection
* Builds the production OkHttpClient by forking {@link #getSharedPoolClient()} so the connection
* pool and dispatcher are reused across instances while applying per-instance timeouts.
*/
private static OkHttpClient buildClient(HttpOptions httpOptions) {
Objects.requireNonNull(httpOptions, "httpOptions cannot be null");
OkHttpClient.Builder builder = SHARED_POOL_CLIENT.newBuilder();
OkHttpClient.Builder builder = getSharedPoolClient().newBuilder();
builder.connectTimeout(Duration.ZERO);
builder.readTimeout(Duration.ZERO);
builder.writeTimeout(Duration.ZERO);
Expand Down
7 changes: 6 additions & 1 deletion core/src/main/java/com/google/adk/sessions/ApiClient.java
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import static com.google.common.base.StandardSystemProperty.JAVA_VERSION;

import com.google.adk.internal.http.HttpClientFactory;
import com.google.auth.oauth2.GoogleCredentials;
import com.google.common.base.Ascii;
import com.google.common.base.Strings;
Expand All @@ -33,6 +34,10 @@

/** Interface for an API client which issues HTTP requests to the GenAI APIs. */
abstract class ApiClient {
private static OkHttpClient getSharedPoolClient() {
return HttpClientFactory.createSharedHttpClient("ApiClient");
}

OkHttpClient httpClient;
// For Google AI APIs
final @Nullable String apiKey;
Expand Down Expand Up @@ -103,7 +108,7 @@ abstract class ApiClient {
}

private OkHttpClient createHttpClient(@Nullable Integer timeout) {
OkHttpClient.Builder builder = new OkHttpClient().newBuilder();
OkHttpClient.Builder builder = getSharedPoolClient().newBuilder();
if (timeout != null) {
builder.connectTimeout(Duration.ofMillis(timeout));
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
/*
* 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.internal.http;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import okhttp3.OkHttpClient;
import org.junit.Test;

public class HttpClientFactoryTest {

@Test
public void testSharedHttpClientDaemonThreads() {
ThreadFactory tf =
((ThreadPoolExecutor)
HttpClientFactory.createSharedHttpClient("Test").dispatcher().executorService())
.getThreadFactory();
Thread t = tf.newThread(() -> {});
assertTrue("HttpClientFactory thread factory should produce daemon threads", t.isDaemon());
}

@Test
public void testSharedHttpClientReusesCachedInstance() {
OkHttpClient client1 = HttpClientFactory.createSharedHttpClient("CacheTest");
OkHttpClient client2 = HttpClientFactory.createSharedHttpClient("CacheTest");
assertSame(
"Repeated calls with the same name should return the same cached instance",
client1,
client2);
}

@Test
public void testSharedHttpClientWithCustomExecutorService() {
ExecutorService customExecutor = Executors.newCachedThreadPool();
try {
OkHttpClient client = HttpClientFactory.createSharedHttpClient("CustomTest", customExecutor);
assertEquals(customExecutor, client.dispatcher().executorService());
} finally {
customExecutor.shutdown();
}
}
}