Merge pull request #1268 from svetanis:fix/graceful-jvm-shutdown

PiperOrigin-RevId: 941709147

Author: copybara-github

core/src/main/java/com/google/adk/internal/http/HttpClientFactory.java

@@ -0,0 +1,56 @@
+/*
+ * 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 java.util.concurrent.Executors;
+import java.util.concurrent.ThreadFactory;
+import okhttp3.Dispatcher;
+import okhttp3.OkHttpClient;
+
+/** Utility class for common HTTP client configuration across the ADK. */
+public final class HttpClientFactory {
+
+  private HttpClientFactory() {}
+
+  /**
+   * Configures a custom OkHttp dispatcher that uses daemon threads. By default, OkHttp uses
+   * non-daemon threads for its async call thread pool, which prevents the JVM from shutting down
+   * for 60 seconds (the default keep-alive) after the last streaming request completes.
+   *
+   * @param name The prefix name to use for the dispatcher threads.
+   * @return A pre-configured Dispatcher using daemon threads.
+   */
+  private static Dispatcher createDaemonDispatcher(String name) {
+    ThreadFactory daemonThreadFactory =
+        r -> {
+          Thread t = new Thread(r, name + "-Dispatcher");
+          t.setDaemon(true);
+          return t;
+        };
+    return new Dispatcher(Executors.newCachedThreadPool(daemonThreadFactory));
+  }
+
+  /**
+   * Creates a shared OkHttpClient instance equipped with a daemon thread dispatcher.
+   *
+   * @param threadName The prefix name to use for the dispatcher threads.
+   * @return A pre-configured OkHttpClient.
+   */
+  public static OkHttpClient createSharedHttpClient(String threadName) {
+    return new OkHttpClient.Builder().dispatcher(createDaemonDispatcher(threadName)).build();
+  }
+}

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

@@ -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;
@@ -42,6 +44,7 @@
 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;
@@ -56,6 +59,8 @@ public class Gemini extends BaseLlm {
 
   private static final Logger logger = LoggerFactory.getLogger(Gemini.class);
   private static final ImmutableMap<String, String> TRACKING_HEADERS;
+  private static final OkHttpClient SHARED_HTTP_CLIENT =
+      HttpClientFactory.createSharedHttpClient("GeminiApiClient");
 
   static {
     String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION;
@@ -94,6 +99,7 @@ public Gemini(String modelName, String apiKey) {
         Client.builder()
             .apiKey(apiKey)
             .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
+            .clientOptions(ClientOptions.builder().customHttpClient(SHARED_HTTP_CLIENT).build())
             .build();
   }
 
@@ -107,7 +113,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(SHARED_HTTP_CLIENT).build());
     vertexCredentials.project().ifPresent(apiClientBuilder::project);
     vertexCredentials.location().ifPresent(apiClientBuilder::location);
     vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials);
@@ -206,6 +214,7 @@ public Gemini build() {
             modelName,
             Client.builder()
                 .httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
+                .clientOptions(ClientOptions.builder().customHttpClient(SHARED_HTTP_CLIENT).build())
                 .build());
       }
     }

core/src/main/java/com/google/adk/models/chat/ChatCompletionsHttpClient.java

@@ -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;
@@ -72,7 +73,8 @@ 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 final OkHttpClient SHARED_POOL_CLIENT =
+      HttpClientFactory.createSharedHttpClient("ChatCompletionsHttpClient");
 
   private final OkHttpClient client;
   private final HttpUrl completionsUrl;

core/src/main/java/com/google/adk/sessions/ApiClient.java

@@ -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;
@@ -102,8 +103,22 @@ abstract class ApiClient {
     this.httpClient = createHttpClient(httpOptions.timeout().orElse(null));
   }
 
+  /**
+   * Shared OkHttpClient instance whose connection pool and thread dispatcher are reused across all
+   * {@link ApiClient} subclass instances.
+   *
+   * <p>Even though {@link ApiClient} is abstract, every instantiation of a concrete subclass (e.g.,
+   * GoogleAiClient) triggers this class's constructor. Making this static prevents a severe
+   * resource leak where every new LLM Agent would otherwise spin up a brand new thread pool and TCP
+   * connection pool. Instead, all clients share this pool and fork it via {@link
+   * OkHttpClient#newBuilder()}.
+   */
+  private static final OkHttpClient SHARED_POOL_CLIENT =
+      HttpClientFactory.createSharedHttpClient("ApiClient");
+
   private OkHttpClient createHttpClient(@Nullable Integer timeout) {
-    OkHttpClient.Builder builder = new OkHttpClient().newBuilder();
+    OkHttpClient.Builder builder = SHARED_POOL_CLIENT.newBuilder();
+
     if (timeout != null) {
       builder.connectTimeout(Duration.ofMillis(timeout));
     }

core/src/test/java/com/google/adk/internal/http/HttpClientFactoryTest.java

@@ -0,0 +1,37 @@
+/*
+ * 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.jupiter.api.Assertions.assertTrue;
+
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.ThreadPoolExecutor;
+import org.junit.jupiter.api.Test;
+
+public class HttpClientFactoryTest {
+
+  @Test
+  public void testSharedHttpClientDaemonThreads() {
+    ThreadFactory tf =
+        ((ThreadPoolExecutor)
+                HttpClientFactory.createSharedHttpClient("Test").dispatcher().executorService())
+            .getThreadFactory();
+    // Usually OkHttp uses a specific thread factory. We passed our own DaemonThreadFactory
+    Thread t = tf.newThread(() -> {});
+    assertTrue(t.isDaemon(), "HttpClientFactory thread factory should produce daemon threads");
+  }
+}