From 9ae572baad719656ec5c4116b664aa5b30d2dee0 Mon Sep 17 00:00:00 2001
From: Google Team Member <genai-sdk-bot@google.com>
Date: Fri, 10 Apr 2026 14:21:54 -0700
Subject: [PATCH] feat: Allowing McpAsycToolset Builder to take in a
 McpSessionManager

PiperOrigin-RevId: 897886449
---
 .../google/adk/tools/mcp/McpAsyncToolset.java | 46 ++++++-------------
 1 file changed, 14 insertions(+), 32 deletions(-)

diff --git a/core/src/main/java/com/google/adk/tools/mcp/McpAsyncToolset.java b/core/src/main/java/com/google/adk/tools/mcp/McpAsyncToolset.java
index 5f4c2164b..f7970493a 100644
--- a/core/src/main/java/com/google/adk/tools/mcp/McpAsyncToolset.java
+++ b/core/src/main/java/com/google/adk/tools/mcp/McpAsyncToolset.java
@@ -66,19 +66,25 @@ public class McpAsyncToolset implements BaseToolset {
 
   /** Builder for McpAsyncToolset */
   public static class Builder {
-    private Object connectionParams = null;
+    private McpSessionManager mcpSessionManager = null;
     private ObjectMapper objectMapper = null;
     private @Nullable Object toolFilter = null;
 
     @CanIgnoreReturnValue
     public Builder connectionParams(ServerParameters connectionParams) {
-      this.connectionParams = connectionParams;
+      this.mcpSessionManager = new McpSessionManager(connectionParams);
       return this;
     }
 
     @CanIgnoreReturnValue
     public Builder connectionParams(SseServerParameters connectionParams) {
-      this.connectionParams = connectionParams;
+      this.mcpSessionManager = new McpSessionManager(connectionParams);
+      return this;
+    }
+
+    @CanIgnoreReturnValue
+    public Builder mcpSessionManager(McpSessionManager mcpSessionManager) {
+      this.mcpSessionManager = mcpSessionManager;
       return this;
     }
 
@@ -104,14 +110,8 @@ public McpAsyncToolset build() {
       if (objectMapper == null) {
         objectMapper = JsonBaseModel.getMapper();
       }
-      if (connectionParams instanceof ServerParameters setSelectedParams) {
-        return new McpAsyncToolset(setSelectedParams, objectMapper, toolFilter);
-      } else if (connectionParams instanceof SseServerParameters sseServerParameters) {
-        return new McpAsyncToolset(sseServerParameters, objectMapper, toolFilter);
-      } else {
-        throw new IllegalArgumentException(
-            "connectionParams must be either ServerParameters or SseServerParameters");
-      }
+      Preconditions.checkNotNull(mcpSessionManager, "Connection params must be set");
+      return new McpAsyncToolset(mcpSessionManager, objectMapper, toolFilter);
     }
   }
 
@@ -123,29 +123,11 @@ public McpAsyncToolset build() {
    * @param toolFilter Either a ToolPredicate or a List of tool names.
    */
   McpAsyncToolset(
-      SseServerParameters connectionParams,
-      ObjectMapper objectMapper,
-      @Nullable Object toolFilter) {
-    Objects.requireNonNull(connectionParams);
-    Objects.requireNonNull(objectMapper);
-    this.objectMapper = objectMapper;
-    this.mcpSessionManager = new McpSessionManager(connectionParams);
-    this.toolFilter = toolFilter;
-  }
-
-  /**
-   * Initializes the McpAsyncToolset with local server parameters.
-   *
-   * @param connectionParams The local server connection parameters to the MCP server.
-   * @param objectMapper An ObjectMapper instance for parsing schemas.
-   * @param toolFilter Either a ToolPredicate or a List of tool names or null.
-   */
-  McpAsyncToolset(
-      ServerParameters connectionParams, ObjectMapper objectMapper, @Nullable Object toolFilter) {
-    Objects.requireNonNull(connectionParams);
+      McpSessionManager mcpSessionManager, ObjectMapper objectMapper, @Nullable Object toolFilter) {
+    Objects.requireNonNull(mcpSessionManager);
     Objects.requireNonNull(objectMapper);
     this.objectMapper = objectMapper;
-    this.mcpSessionManager = new McpSessionManager(connectionParams);
+    this.mcpSessionManager = mcpSessionManager;
     this.toolFilter = toolFilter;
   }