feat: Allowing McpAsycToolset Builder to take in a McpSessionManager

PiperOrigin-RevId: 897886449

Author: google-genai-bot

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;
   }