feat: add support for Streamable HTTP Connections to MCP Tools

Poggecci · copybara-github · commit bea3244c5850 · 2025-09-02T13:43:41.000-07:00
PiperOrigin-RevId: 802274756
diff --git a/core/src/main/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilder.java b/core/src/main/java/com/google/adk/tools/mcp/DefaultMcpTransportBuilder.java
@@ -2,16 +2,18 @@
 
 import com.google.common.collect.ImmutableMap;
 import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport;
+import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport;
 import io.modelcontextprotocol.client.transport.ServerParameters;
 import io.modelcontextprotocol.client.transport.StdioClientTransport;
 import io.modelcontextprotocol.spec.McpClientTransport;
 import java.util.Collection;
 import java.util.Optional;
+import reactor.core.publisher.Mono;
 
 /**
  * The default builder for creating MCP client transports. Supports StdioClientTransport based on
- * {@link ServerParameters} and the standard HttpClientSseClientTransport based on {@link
- * SseServerParameters}.
+ * {@link ServerParameters}, HttpClientSseClientTransport based on {@link SseServerParameters}, and
+ * HttpClientStreamableHttpTransport based on {@link StreamableHttpServerParameters}.
  */
 public class DefaultMcpTransportBuilder implements McpTransportBuilder {
 
@@ -37,10 +39,21 @@ public McpClientTransport build(Object connectionParams) {
                                       .map(Object::toString)
                                       .orElse(""))))
           .build();
+    } else if (connectionParams instanceof StreamableHttpServerParameters streamableParams) {
+      return HttpClientStreamableHttpTransport.builder(streamableParams.url())
+          .connectTimeout(streamableParams.timeout())
+          // HttpClientStreamableHttpTransport uses connectTimeout for general HTTP ops
+          // and sseReadTimeout for the SSE stream part.
+          .asyncHttpRequestCustomizer(
+              (builder, method, uri, body) -> {
+                streamableParams.headers().forEach((key, value) -> builder.header(key, value));
+                return Mono.just(builder);
+              })
+          .build();
     } else {
       throw new IllegalArgumentException(
-          "DefaultMcpTransportBuilder supports only ServerParameters or SseServerParameters, but"
-              + " got "
+          "DefaultMcpTransportBuilder supports only ServerParameters, SseServerParameters, or"
+              + " StreamableHttpServerParameters, but got "
               + connectionParams.getClass().getName());
     }
   }
diff --git a/core/src/main/java/com/google/adk/tools/mcp/McpSessionManager.java b/core/src/main/java/com/google/adk/tools/mcp/McpSessionManager.java
@@ -65,6 +65,9 @@ public static McpSyncClient initializeSession(
     if (connectionParams instanceof SseServerParameters sseServerParams) {
       initializationTimeout = sseServerParams.timeout();
       requestTimeout = sseServerParams.sseReadTimeout();
+    } else if (connectionParams instanceof StreamableHttpServerParameters streamableParams) {
+      initializationTimeout = streamableParams.timeout();
+      requestTimeout = streamableParams.sseReadTimeout();
     }
     McpSyncClient client =
         McpClient.sync(transport)
@@ -95,6 +98,9 @@ public static McpAsyncClient initializeAsyncSession(
     if (connectionParams instanceof SseServerParameters sseServerParams) {
       initializationTimeout = sseServerParams.timeout();
       requestTimeout = sseServerParams.sseReadTimeout();
+    } else if (connectionParams instanceof StreamableHttpServerParameters streamableParams) {
+      initializationTimeout = streamableParams.timeout();
+      requestTimeout = streamableParams.sseReadTimeout();
     }
     return McpClient.async(transport)
         .initializationTimeout(
diff --git a/core/src/main/java/com/google/adk/tools/mcp/McpToolset.java b/core/src/main/java/com/google/adk/tools/mcp/McpToolset.java
@@ -165,6 +165,34 @@ public McpToolset(
     this.toolFilter = toolFilter;
   }
 
+  /**
+   * Initializes the McpToolset with Steamable HTTP server parameters.
+   *
+   * @param connectionParams The Streamable HTTP connection parameters to the MCP server.
+   * @param objectMapper An ObjectMapper instance for parsing schemas.
+   * @param toolFilter An Optional containing either a ToolPredicate or a List of tool names.
+   */
+  public McpToolset(
+      StreamableHttpServerParameters connectionParams,
+      ObjectMapper objectMapper,
+      Optional<Object> toolFilter) {
+    Objects.requireNonNull(connectionParams);
+    Objects.requireNonNull(objectMapper);
+    this.objectMapper = objectMapper;
+    this.mcpSessionManager = new McpSessionManager(connectionParams);
+    this.toolFilter = toolFilter;
+  }
+
+  /**
+   * Initializes the McpToolset with Streamable HTTP server parameters, using the ObjectMapper used
+   * across the ADK and no tool filter.
+   *
+   * @param connectionParams The Streamable HTTP connection parameters to the MCP server.
+   */
+  public McpToolset(StreamableHttpServerParameters connectionParams) {
+    this(connectionParams, JsonBaseModel.getMapper(), Optional.empty());
+  }
+
   @Override
   public Flowable<BaseTool> getTools(ReadonlyContext readonlyContext) {
     return Flowable.fromCallable(
diff --git a/core/src/main/java/com/google/adk/tools/mcp/StreamableHttpServerParameters.java b/core/src/main/java/com/google/adk/tools/mcp/StreamableHttpServerParameters.java
@@ -0,0 +1,94 @@
+/*
+ * Copyright 2025 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.tools.mcp;
+
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
+import io.modelcontextprotocol.util.Assert;
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Map;
+
+/**
+ * Server parameters for Streamable HTTP client transport.
+ *
+ * @param url The base URL for the MCP Streamable HTTP server.
+ * @param headers Optional headers to include in requests.
+ * @param timeout Timeout for HTTP operations (default: 30 seconds).
+ * @param sseReadTimeout Timeout for reading data from the SSE stream (default: 5 minutes).
+ * @param terminateOnClose Whether to terminate the session on close (default: true).
+ */
+public record StreamableHttpServerParameters(
+    String url,
+    Map<String, String> headers,
+    Duration timeout,
+    Duration sseReadTimeout,
+    boolean terminateOnClose) {
+
+  public StreamableHttpServerParameters {
+    Assert.hasText(url, "url must not be empty");
+    headers = headers == null ? Collections.emptyMap() : headers;
+    timeout = timeout == null ? Duration.ofSeconds(30) : timeout;
+    sseReadTimeout = sseReadTimeout == null ? Duration.ofMinutes(5) : sseReadTimeout;
+  }
+
+  public static Builder builder(String url) {
+    return new Builder(url);
+  }
+
+  /** Builder for {@link StreamableHttpServerParameters}. */
+  public static class Builder {
+    private final String url;
+    private Map<String, String> headers;
+    private Duration timeout = Duration.ofSeconds(30);
+    private Duration sseReadTimeout = Duration.ofMinutes(5);
+    private boolean terminateOnClose = true;
+
+    private Builder(String url) {
+      Assert.hasText(url, "url must not be empty");
+      this.url = url;
+    }
+
+    @CanIgnoreReturnValue
+    public Builder headers(Map<String, String> headers) {
+      this.headers = headers;
+      return this;
+    }
+
+    @CanIgnoreReturnValue
+    public Builder timeout(Duration timeout) {
+      this.timeout = timeout;
+      return this;
+    }
+
+    @CanIgnoreReturnValue
+    public Builder sseReadTimeout(Duration sseReadTimeout) {
+      this.sseReadTimeout = sseReadTimeout;
+      return this;
+    }
+
+    @CanIgnoreReturnValue
+    public Builder terminateOnClose(boolean terminateOnClose) {
+      this.terminateOnClose = terminateOnClose;
+      return this;
+    }
+
+    public StreamableHttpServerParameters build() {
+      return new StreamableHttpServerParameters(
+          url, headers, timeout, sseReadTimeout, terminateOnClose);
+    }
+  }
+}
