feat: support stdio_connection_params in McpToolset config

PiperOrigin-RevId: 845809789

Author: google-genai-bot

core/src/main/java/com/google/adk/tools/mcp/McpSessionManager.java

@@ -61,14 +61,19 @@ public static McpSyncClient initializeSession(
       Object connectionParams, McpTransportBuilder transportBuilder) {
     Duration initializationTimeout = null;
     Duration requestTimeout = null;
-    McpClientTransport transport = transportBuilder.build(connectionParams);
-    if (connectionParams instanceof SseServerParameters sseServerParams) {
+    Object transportBuilderParams = connectionParams;
+    if (connectionParams instanceof StdioConnectionParameters stdioConnectionParameters) {
+      transportBuilderParams = stdioConnectionParameters.serverParams().toServerParameters();
+      requestTimeout = stdioConnectionParameters.timeoutDuration();
+    } else if (connectionParams instanceof SseServerParameters sseServerParams) {
       initializationTimeout = sseServerParams.timeout();
       requestTimeout = sseServerParams.sseReadTimeout();
     } else if (connectionParams instanceof StreamableHttpServerParameters streamableParams) {
       initializationTimeout = streamableParams.timeout();
       requestTimeout = streamableParams.readTimeout();
     }
+    McpClientTransport transport = transportBuilder.build(transportBuilderParams);
+
     McpSyncClient client =
         McpClient.sync(transport)
             .initializationTimeout(

core/src/main/java/com/google/adk/tools/mcp/McpToolset.java

@@ -24,6 +24,7 @@
 import com.google.adk.agents.ReadonlyContext;
 import com.google.adk.tools.BaseTool;
 import com.google.adk.tools.BaseToolset;
+import com.google.common.primitives.Booleans;
 import io.modelcontextprotocol.client.McpSyncClient;
 import io.modelcontextprotocol.client.transport.ServerParameters;
 import io.modelcontextprotocol.spec.McpSchema.ListToolsResult;
@@ -285,12 +286,23 @@ public void close() {
 
   /** Configuration class for MCPToolset. */
   public static class McpToolsetConfig extends JsonBaseModel {
+
+    private StdioConnectionParameters stdioConnectionParams;
+
     private StdioServerParameters stdioServerParams;
 
     private SseServerParameters sseServerParams;
 
     private List<String> toolFilter;
 
+    public StdioConnectionParameters stdioConnectionParams() {
+      return stdioConnectionParams;
+    }
+
+    public void setStdioConnectionParams(StdioConnectionParameters stdioConnectionParams) {
+      this.stdioConnectionParams = stdioConnectionParams;
+    }
+
     public StdioServerParameters stdioServerParams() {
       return stdioServerParams;
     }
@@ -337,23 +349,26 @@ public static McpToolset fromConfig(BaseTool.ToolConfig config, String configAbs
           mapper.convertValue(config.args(), McpToolsetConfig.class);
 
       // Validate that exactly one parameter type is set
-      if ((mcpToolsetConfig.stdioServerParams() != null)
-          == (mcpToolsetConfig.sseServerParams() != null)) {
+      if (Booleans.countTrue(
+              mcpToolsetConfig.stdioServerParams() != null,
+              mcpToolsetConfig.sseServerParams() != null,
+              mcpToolsetConfig.stdioConnectionParams() != null)
+          != 1) {
         throw new ConfigurationException(
-            "Exactly one of stdioServerParams or sseServerParams must be set for McpToolset");
+            "Exactly one of stdioConnectionParams, stdioServerParams or sseServerParams must be set"
+                + " for McpToolset");
       }
 
       // Convert tool filter to Optional<Object>
-      Optional<Object> toolFilter =
-          Optional.ofNullable(mcpToolsetConfig.toolFilter()).map(filter -> filter);
-
-      // Create McpToolset with appropriate connection parameters
-      if (mcpToolsetConfig.stdioServerParams() != null) {
-        return new McpToolset(
-            mcpToolsetConfig.stdioServerParams().toServerParameters(), mapper, toolFilter);
-      } else {
-        return new McpToolset(mcpToolsetConfig.sseServerParams(), mapper, toolFilter);
-      }
+      Optional<Object> toolFilter = Optional.ofNullable(mcpToolsetConfig.toolFilter());
+
+      Object connectionParameters =
+          Optional.<Object>ofNullable(mcpToolsetConfig.stdioConnectionParams())
+              .or(() -> Optional.ofNullable(mcpToolsetConfig.sseServerParams()))
+              .orElse(mcpToolsetConfig.stdioConnectionParams());
+
+      // Create McpToolset with McpSessionManager having appropriate connection parameters
+      return new McpToolset(new McpSessionManager(connectionParameters), mapper, toolFilter);
     } catch (IllegalArgumentException e) {
       throw new ConfigurationException("Failed to parse McpToolsetConfig from ToolArgsConfig", e);
     }

core/src/main/java/com/google/adk/tools/mcp/StdioConnectionParameters.java

@@ -0,0 +1,48 @@
+package com.google.adk.tools.mcp;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.databind.annotation.JsonDeserialize;
+import com.fasterxml.jackson.databind.annotation.JsonPOJOBuilder;
+import com.google.auto.value.AutoValue;
+import java.time.Duration;
+
+@AutoValue
+@JsonDeserialize(builder = StdioConnectionParameters.Builder.class)
+public abstract class StdioConnectionParameters {
+
+  private static final long MILLIS_IN_SEC = 1000L;
+
+  StdioConnectionParameters() {}
+
+  public abstract StdioServerParameters serverParams();
+
+  // Timeout in seconds
+  public abstract float timeout();
+
+  @JsonIgnore
+  public Duration timeoutDuration() {
+    return Duration.ofMillis((long) (timeout() * MILLIS_IN_SEC));
+  }
+
+  @AutoValue.Builder
+  @JsonPOJOBuilder(buildMethodName = "build", withPrefix = "")
+  public abstract static class Builder {
+
+    @JsonCreator
+    public static Builder jacksonBuilder() {
+      return StdioConnectionParameters.builder();
+    }
+
+    public abstract Builder serverParams(StdioServerParameters serverParams);
+
+    public abstract Builder timeout(float timeout);
+
+    public abstract StdioConnectionParameters build();
+  }
+
+  public static Builder builder() {
+    Builder b = new AutoValue_StdioConnectionParameters.Builder();
+    return b.timeout(5f);
+  }
+}