Removing streamTools in favor of exchange-driven transport upgrading. Added extensive integration testing.

Author: Zachary German

mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransport.java

@@ -29,6 +29,7 @@
 import io.modelcontextprotocol.spec.McpTransportSession;
 import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
 import io.modelcontextprotocol.spec.McpTransportStream;
+import io.modelcontextprotocol.spec.McpSchema.JSONRPCNotification;
 import io.modelcontextprotocol.util.Assert;
 import reactor.core.Disposable;
 import reactor.core.publisher.Flux;
@@ -117,10 +118,6 @@ public static Builder builder(WebClient.Builder webClientBuilder) {
 	public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
 		return Mono.deferContextual(ctx -> {
 			this.handler.set(handler);
-			if (openConnectionOnStartup) {
-				logger.debug("Eagerly opening connection on startup");
-				return this.reconnect(null).then();
-			}
 			return Mono.empty();
 		});
 	}
@@ -250,11 +247,13 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
 				})
 				.bodyValue(message)
 				.exchangeToFlux(response -> {
-					if (transportSession
-						.markInitialized(response.headers().asHttpHeaders().getFirst("mcp-session-id"))) {
-						// Once we have a session, we try to open an async stream for
-						// the server to send notifications and requests out-of-band.
-						reconnect(null).contextWrite(sink.contextView()).subscribe();
+					transportSession.markInitialized(response.headers().asHttpHeaders().getFirst("mcp-session-id"));
+					if (response.statusCode().is2xxSuccessful()
+							&& message instanceof JSONRPCNotification notification) {
+						if (notification.method().equals("notifications/initialized")) {
+							// Establish SSE stream after session is initialized
+							reconnect(null).contextWrite(sink.contextView()).subscribe();
+						}
 					}
 
 					String sessionRepresentation = sessionIdOrPlaceholder(transportSession);

mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/transport/StreamableHttpTransportIntegrationTest.java

@@ -88,16 +88,12 @@ public class McpAsyncServer {
 
 	private final McpSchema.ServerCapabilities serverCapabilities;
 
-	private final boolean isStreamableHttp;
-
 	private final McpSchema.Implementation serverInfo;
 
 	private final String instructions;
 
 	private final CopyOnWriteArrayList<McpServerFeatures.AsyncToolSpecification> tools = new CopyOnWriteArrayList<>();
 
-	private final CopyOnWriteArrayList<McpServerFeatures.AsyncStreamingToolSpecification> streamTools = new CopyOnWriteArrayList<>();
-
 	private final CopyOnWriteArrayList<McpSchema.ResourceTemplate> resourceTemplates = new CopyOnWriteArrayList<>();
 
 	private final ConcurrentHashMap<String, McpServerFeatures.AsyncResourceSpecification> resources = new ConcurrentHashMap<>();
@@ -123,7 +119,7 @@ public class McpAsyncServer {
 	 */
 	McpAsyncServer(McpServerTransportProvider mcpTransportProvider, ObjectMapper objectMapper,
 			McpServerFeatures.Async features, Duration requestTimeout,
-			McpUriTemplateManagerFactory uriTemplateManagerFactory, boolean isStreamableHttp) {
+			McpUriTemplateManagerFactory uriTemplateManagerFactory) {
 		this.mcpTransportProvider = mcpTransportProvider;
 		this.objectMapper = objectMapper;
 		this.serverInfo = features.serverInfo();
@@ -135,7 +131,6 @@ public class McpAsyncServer {
 		this.prompts.putAll(features.prompts());
 		this.completions.putAll(features.completions());
 		this.uriTemplateManagerFactory = uriTemplateManagerFactory;
-		this.isStreamableHttp = isStreamableHttp;
 
 		Map<String, McpServerSession.RequestHandler<?>> requestHandlers = new HashMap<>();
 
@@ -193,13 +188,6 @@ public class McpAsyncServer {
 				notificationHandlers));
 	}
 
-	// Alternate constructor for HTTP+SSE servers (past spec)
-	McpAsyncServer(McpServerTransportProvider mcpTransportProvider, ObjectMapper objectMapper,
-			McpServerFeatures.Async features, Duration requestTimeout,
-			McpUriTemplateManagerFactory uriTemplateManagerFactory) {
-		this(mcpTransportProvider, objectMapper, features, requestTimeout, uriTemplateManagerFactory, false);
-	}
-
 	// ---------------------------------------
 	// Lifecycle Management
 	// ---------------------------------------
@@ -341,69 +329,6 @@ public Mono<Void> removeTool(String toolName) {
 		});
 	}
 
-	/**
-	 * Add a new tool specification at runtime.
-	 * @param toolSpecification The tool specification to add
-	 * @return Mono that completes when clients have been notified of the change
-	 */
-	public Mono<Void> addStreamTool(McpServerFeatures.AsyncStreamingToolSpecification toolSpecification) {
-		if (toolSpecification == null) {
-			return Mono.error(new McpError("Tool specification must not be null"));
-		}
-		if (toolSpecification.tool() == null) {
-			return Mono.error(new McpError("Tool must not be null"));
-		}
-		if (toolSpecification.call() == null) {
-			return Mono.error(new McpError("Tool call handler must not be null"));
-		}
-		if (this.serverCapabilities.tools() == null) {
-			return Mono.error(new McpError("Server must be configured with tool capabilities"));
-		}
-
-		return Mono.defer(() -> {
-			// Check for duplicate tool names
-			if (this.streamTools.stream().anyMatch(th -> th.tool().name().equals(toolSpecification.tool().name()))) {
-				return Mono
-					.error(new McpError("Tool with name '" + toolSpecification.tool().name() + "' already exists"));
-			}
-
-			this.streamTools.add(toolSpecification);
-			logger.debug("Added tool handler: {}", toolSpecification.tool().name());
-
-			if (this.serverCapabilities.tools().listChanged()) {
-				return notifyToolsListChanged();
-			}
-			return Mono.empty();
-		});
-	}
-
-	/**
-	 * Remove a tool handler at runtime.
-	 * @param toolName The name of the tool handler to remove
-	 * @return Mono that completes when clients have been notified of the change
-	 */
-	public Mono<Void> removeStreamTool(String toolName) {
-		if (toolName == null) {
-			return Mono.error(new McpError("Tool name must not be null"));
-		}
-		if (this.serverCapabilities.tools() == null) {
-			return Mono.error(new McpError("Server must be configured with tool capabilities"));
-		}
-
-		return Mono.defer(() -> {
-			boolean removed = this.tools
-				.removeIf(toolSpecification -> toolSpecification.tool().name().equals(toolName));
-			if (removed) {
-				logger.debug("Removed tool handler: {}", toolName);
-				if (this.serverCapabilities.tools().listChanged()) {
-					return notifyToolsListChanged();
-				}
-				return Mono.empty();
-			}
-			return Mono.error(new McpError("Tool with name '" + toolName + "' not found"));
-		});
-	}
-
 	/**
 	 * Notifies clients that the list of available tools has changed.
 	 * @return A Mono that completes when all clients have been notified
@@ -416,95 +341,27 @@ private McpServerSession.RequestHandler<McpSchema.ListToolsResult> toolsListRequ
 		return (exchange, params) -> {
 			List<Tool> tools = new ArrayList<>();
 			tools.addAll(this.tools.stream().map(McpServerFeatures.AsyncToolSpecification::tool).toList());
-			tools.addAll(
-					this.streamTools.stream().map(McpServerFeatures.AsyncStreamingToolSpecification::tool).toList());
 
 			return Mono.just(new McpSchema.ListToolsResult(tools, null));
 		};
 	}
 
 	private McpServerSession.RequestHandler<CallToolResult> toolsCallRequestHandler() {
-		if (isStreamableHttp) {
-			return new McpServerSession.StreamingRequestHandler<CallToolResult>() {
-				@Override
-				public Mono<CallToolResult> handle(McpAsyncServerExchange exchange, Object params) {
-					var callToolRequest = objectMapper.convertValue(params, McpSchema.CallToolRequest.class);
-
-					// Check regular tools first
-					var regularTool = tools.stream()
-						.filter(tool -> callToolRequest.name().equals(tool.tool().name()))
-						.findFirst();
-
-					if (regularTool.isPresent()) {
-						return regularTool.get().call().apply(exchange, callToolRequest.arguments());
-					}
-
-					// Check streaming tools (take first result)
-					var streamingTool = streamTools.stream()
-						.filter(tool -> callToolRequest.name().equals(tool.tool().name()))
-						.findFirst();
-
-					if (streamingTool.isPresent()) {
-						return streamingTool.get().call().apply(exchange, callToolRequest.arguments()).next();
-					}
-
-					return Mono.error(new McpError("Tool not found: " + callToolRequest.name()));
-				}
-
-				@Override
-				public Flux<CallToolResult> handleStreaming(McpAsyncServerExchange exchange, Object params) {
-					var callToolRequest = objectMapper.convertValue(params, McpSchema.CallToolRequest.class);
-
-					// Check streaming tools first (preferred for streaming)
-					var streamingTool = streamTools.stream()
-						.filter(tool -> callToolRequest.name().equals(tool.tool().name()))
-						.findFirst();
-
-					if (streamingTool.isPresent()) {
-						return streamingTool.get().call().apply(exchange, callToolRequest.arguments());
-					}
-
-					// Fallback to regular tools (convert Mono to Flux)
-					var regularTool = tools.stream()
-						.filter(tool -> callToolRequest.name().equals(tool.tool().name()))
-						.findFirst();
-
-					if (regularTool.isPresent()) {
-						return regularTool.get().call().apply(exchange, callToolRequest.arguments()).flux();
-					}
-
-					return Flux.error(new McpError("Tool not found: " + callToolRequest.name()));
-				}
-			};
-		}
-		else {
-			return (exchange, params) -> {
-				McpSchema.CallToolRequest callToolRequest = objectMapper.convertValue(params,
-						new TypeReference<McpSchema.CallToolRequest>() {
-						});
-
-				// Check regular tools first
-				Optional<McpServerFeatures.AsyncToolSpecification> toolSpecification = this.tools.stream()
-					.filter(tr -> callToolRequest.name().equals(tr.tool().name()))
-					.findAny();
-
-				if (toolSpecification.isPresent()) {
-					return toolSpecification.get().call().apply(exchange, callToolRequest.arguments());
-				}
+		return (exchange, params) -> {
+			McpSchema.CallToolRequest callToolRequest = objectMapper.convertValue(params,
+					new TypeReference<McpSchema.CallToolRequest>() {
+					});
 
-				// Check streaming tools (take first result)
-				Optional<McpServerFeatures.AsyncStreamingToolSpecification> streamToolSpecification = this.streamTools
-					.stream()
-					.filter(tr -> callToolRequest.name().equals(tr.tool().name()))
-					.findAny();
+			Optional<McpServerFeatures.AsyncToolSpecification> toolSpecification = this.tools.stream()
+				.filter(tr -> callToolRequest.name().equals(tr.tool().name()))
+				.findAny();
 
-				if (streamToolSpecification.isPresent()) {
-					return streamToolSpecification.get().call().apply(exchange, callToolRequest.arguments()).next();
-				}
+			if (toolSpecification.isPresent()) {
+				return toolSpecification.get().call().apply(exchange, callToolRequest.arguments());
+			}
 
-				return Mono.error(new McpError("Tool not found: " + callToolRequest.name()));
-			};
-		}
+			return Mono.error(new McpError("Tool not found: " + callToolRequest.name()));
+		};
 	}
 
 	// ---------------------------------------

mcp/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java

@@ -8,11 +8,14 @@
 import java.util.Collections;
 
 import com.fasterxml.jackson.core.type.TypeReference;
+
+import io.modelcontextprotocol.server.transport.StreamableHttpServerTransportProvider;
 import io.modelcontextprotocol.spec.McpError;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpSchema.LoggingLevel;
 import io.modelcontextprotocol.spec.McpSchema.LoggingMessageNotification;
 import io.modelcontextprotocol.spec.McpServerSession;
+import io.modelcontextprotocol.spec.McpServerTransport;
 import io.modelcontextprotocol.util.Assert;
 import reactor.core.publisher.Mono;
 
@@ -79,6 +82,28 @@ public McpSchema.Implementation getClientInfo() {
 		return this.clientInfo;
 	}
 
+	/**
+	 * If the exchange's session is using StreamableHttp: Upgrades the transport
+	 * referenced by this exchange's transportId to an SSE stream if it isn't already one.
+	 */
+	private void establishSseStream() {
+		final McpServerTransport currentTransport = session.getTransport(transportId);
+		if (session.isStreamableHttp()
+				&& currentTransport instanceof StreamableHttpServerTransportProvider.HttpTransport transport) {
+			session.registerTransport(transportId,
+					new StreamableHttpServerTransportProvider.SseTransport(transport.getObjectMapper(),
+							transport.getResponse(), transport.getAsyncContext(), null, transportId, session.getId()));
+		}
+	}
+
+	/**
+	 * This is for tool writers to use if they want to send their tool response over an
+	 * SSE stream without using any other McpAsyncServerExchange methods
+	 */
+	public void upgradeTransport() {
+		establishSseStream();
+	}
+
 	/**
 	 * Create a new message using the sampling capabilities of the client. The Model
 	 * Context Protocol (MCP) provides a standardized way for servers to request LLM
@@ -96,6 +121,9 @@ public McpSchema.Implementation getClientInfo() {
 	 * Specification</a>
 	 */
 	public Mono<McpSchema.CreateMessageResult> createMessage(McpSchema.CreateMessageRequest createMessageRequest) {
+
+		establishSseStream();
+
 		if (this.clientCapabilities == null) {
 			return Mono.error(new McpError("Client must be initialized. Call the initialize method first!"));
 		}
@@ -121,6 +149,9 @@ public Mono<McpSchema.CreateMessageResult> createMessage(McpSchema.CreateMessage
 	 * Specification</a>
 	 */
 	public Mono<McpSchema.ElicitResult> createElicitation(McpSchema.ElicitRequest elicitRequest) {
+
+		establishSseStream();
+
 		if (this.clientCapabilities == null) {
 			return Mono.error(new McpError("Client must be initialized. Call the initialize method first!"));
 		}
@@ -157,6 +188,9 @@ public Mono<McpSchema.ListRootsResult> listRoots() {
 	 * @return A Mono that emits the list of roots result containing
 	 */
 	public Mono<McpSchema.ListRootsResult> listRoots(String cursor) {
+
+		establishSseStream();
+
 		return this.session.sendRequest(McpSchema.METHOD_ROOTS_LIST, new McpSchema.PaginatedRequest(cursor),
 				LIST_ROOTS_RESULT_TYPE_REF, transportId);
 	}
@@ -169,6 +203,8 @@ public Mono<McpSchema.ListRootsResult> listRoots(String cursor) {
 	 */
 	public Mono<Void> loggingNotification(LoggingMessageNotification loggingMessageNotification) {
 
+		establishSseStream();
+
 		if (loggingMessageNotification == null) {
 			return Mono.error(new McpError("Logging message must not be null"));
 		}
@@ -187,6 +223,9 @@ public Mono<Void> loggingNotification(LoggingMessageNotification loggingMessageN
 	 * @return A Mono that completes with clients's ping response
 	 */
 	public Mono<Object> ping() {
+
+		establishSseStream();
+
 		return this.session.sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF, transportId);
 	}