Added per-request exchange instances to route server-sent JSONRPCMessage to their related transport streams

Author: Zachary German

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

@@ -778,6 +778,7 @@ private McpServerSession.RequestHandler<Object> setLoggerRequestHandler() {
 						new TypeReference<SetLevelRequest>() {
 						});
 
+				// This will update both the exchange and session logging levels
 				exchange.setMinLoggingLevel(newMinLoggingLevel.level());
 
 				// FIXME: this field is deprecated and should be removed together

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

@@ -45,18 +45,22 @@ public class McpAsyncServerExchange {
 	public static final TypeReference<Object> OBJECT_TYPE_REF = new TypeReference<>() {
 	};
 
+	private final String transportId;
+
 	/**
 	 * Create a new asynchronous exchange with the client.
 	 * @param session The server session representing a 1-1 interaction.
 	 * @param clientCapabilities The client capabilities that define the supported
 	 * features and functionality.
 	 * @param clientInfo The client implementation information.
+	 * @param transportId The transport ID to use for outgoing messages
 	 */
 	public McpAsyncServerExchange(McpServerSession session, McpSchema.ClientCapabilities clientCapabilities,
-			McpSchema.Implementation clientInfo) {
+			McpSchema.Implementation clientInfo, String transportId) {
 		this.session = session;
 		this.clientCapabilities = clientCapabilities;
 		this.clientInfo = clientInfo;
+		this.transportId = transportId;
 	}
 
 	/**
@@ -99,7 +103,7 @@ public Mono<McpSchema.CreateMessageResult> createMessage(McpSchema.CreateMessage
 			return Mono.error(new McpError("Client must be configured with sampling capabilities"));
 		}
 		return this.session.sendRequest(McpSchema.METHOD_SAMPLING_CREATE_MESSAGE, createMessageRequest,
-				CREATE_MESSAGE_RESULT_TYPE_REF);
+				CREATE_MESSAGE_RESULT_TYPE_REF, transportId);
 	}
 
 	/**
@@ -123,8 +127,8 @@ public Mono<McpSchema.ElicitResult> createElicitation(McpSchema.ElicitRequest el
 		if (this.clientCapabilities.elicitation() == null) {
 			return Mono.error(new McpError("Client must be configured with elicitation capabilities"));
 		}
-		return this.session.sendRequest(McpSchema.METHOD_ELICITATION_CREATE, elicitRequest,
-				ELICITATION_RESULT_TYPE_REF);
+		return this.session.sendRequest(McpSchema.METHOD_ELICITATION_CREATE, elicitRequest, ELICITATION_RESULT_TYPE_REF,
+				transportId);
 	}
 
 	/**
@@ -154,7 +158,7 @@ public Mono<McpSchema.ListRootsResult> listRoots() {
 	 */
 	public Mono<McpSchema.ListRootsResult> listRoots(String cursor) {
 		return this.session.sendRequest(McpSchema.METHOD_ROOTS_LIST, new McpSchema.PaginatedRequest(cursor),
-				LIST_ROOTS_RESULT_TYPE_REF);
+				LIST_ROOTS_RESULT_TYPE_REF, transportId);
 	}
 
 	/**
@@ -171,7 +175,8 @@ public Mono<Void> loggingNotification(LoggingMessageNotification loggingMessageN
 
 		return Mono.defer(() -> {
 			if (this.isNotificationForLevelAllowed(loggingMessageNotification.level())) {
-				return this.session.sendNotification(McpSchema.METHOD_NOTIFICATION_MESSAGE, loggingMessageNotification);
+				return this.session.sendNotification(McpSchema.METHOD_NOTIFICATION_MESSAGE, loggingMessageNotification,
+						transportId);
 			}
 			return Mono.empty();
 		});
@@ -182,17 +187,19 @@ public Mono<Void> loggingNotification(LoggingMessageNotification loggingMessageN
 	 * @return A Mono that completes with clients's ping response
 	 */
 	public Mono<Object> ping() {
-		return this.session.sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF);
+		return this.session.sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF, transportId);
 	}
 
 	/**
 	 * Set the minimum logging level for the client. Messages below this level will be
 	 * filtered out.
 	 * @param minLoggingLevel The minimum logging level
 	 */
-	void setMinLoggingLevel(LoggingLevel minLoggingLevel) {
+	public void setMinLoggingLevel(LoggingLevel minLoggingLevel) {
 		Assert.notNull(minLoggingLevel, "minLoggingLevel must not be null");
 		this.minLoggingLevel = minLoggingLevel;
+		// Also update the session level for future exchanges
+		this.session.setMinLoggingLevel(minLoggingLevel);
 	}
 
 	private boolean isNotificationForLevelAllowed(LoggingLevel loggingLevel) {

mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java

@@ -14,6 +14,7 @@
 import io.modelcontextprotocol.server.McpAsyncServerExchange;
 import io.modelcontextprotocol.spec.SseEvent;
 import io.modelcontextprotocol.spec.McpSchema.McpId;
+import io.modelcontextprotocol.spec.McpError;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -53,8 +54,6 @@ public class McpServerSession implements McpSession {
 
 	private final Map<String, NotificationHandler> notificationHandlers;
 
-	private final Sinks.One<McpAsyncServerExchange> exchangeSink = Sinks.one();
-
 	private final AtomicReference<McpSchema.ClientCapabilities> clientCapabilities = new AtomicReference<>();
 
 	private final AtomicReference<McpSchema.Implementation> clientInfo = new AtomicReference<>();
@@ -73,6 +72,8 @@ public class McpServerSession implements McpSession {
 
 	private final Map<String, Map<String, SseEvent>> transportEventHistories = new ConcurrentHashMap<>();
 
+	private volatile McpSchema.LoggingLevel minLoggingLevel = McpSchema.LoggingLevel.INFO;
+
 	/**
 	 * Creates a new server session with the given parameters and the transport to use.
 	 * @param id session id
@@ -104,6 +105,16 @@ public McpServerSession(String id, Duration requestTimeout, InitRequestHandler i
 		this(id, requestTimeout, null, initHandler, initNotificationHandler, requestHandlers, notificationHandlers);
 	}
 
+	/**
+	 * Updates the session's minimum logging level for all future exchanges.
+	 */
+	public void setMinLoggingLevel(McpSchema.LoggingLevel level) {
+		if (level != null) {
+			this.minLoggingLevel = level;
+			logger.debug("Updated session {} minimum logging level to {}", id, level);
+		}
+	}
+
 	/**
 	 * Retrieve the session initialization state
 	 * @return session initialization state
@@ -240,14 +251,26 @@ public RequestHandler<?> getRequestHandler(String method) {
 
 	@Override
 	public <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
-		McpId requestId = this.generateRequestId();
+		return sendRequest(method, requestParams, typeRef, LISTENING_TRANSPORT);
+	}
+
+	public <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef, String transportId) {
+		McpServerTransport transport = getTransport(transportId);
+		if (transport == null) {
+			// Fallback to listening transport if specific transport not found
+			transport = getTransport(LISTENING_TRANSPORT);
+			if (transport == null) {
+				return Mono.error(new RuntimeException("Transport not found: " + transportId));
+			}
+		}
 
+		final McpServerTransport finalTransport = transport;
+		McpId requestId = this.generateRequestId();
 		return Mono.<McpSchema.JSONRPCResponse>create(sink -> {
 			this.pendingResponses.put(requestId, sink);
 			McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method,
 					requestId, requestParams);
-
-			Flux.from(listeningTransport.sendMessage(jsonrpcRequest)).subscribe(v -> {
+			Flux.from(finalTransport.sendMessage(jsonrpcRequest)).subscribe(v -> {
 			}, error -> {
 				this.pendingResponses.remove(requestId);
 				sink.error(error);
@@ -260,17 +283,29 @@ else if (typeRef.getType().equals(Void.class)) {
 				sink.complete();
 			}
 			else {
-				T result = listeningTransport.unmarshalFrom(jsonRpcResponse.result(), typeRef);
+				T result = finalTransport.unmarshalFrom(jsonRpcResponse.result(), typeRef);
 				sink.next(result);
 			}
 		});
 	}
 
 	@Override
 	public Mono<Void> sendNotification(String method, Object params) {
+		return sendNotification(method, params, LISTENING_TRANSPORT);
+	}
+
+	public Mono<Void> sendNotification(String method, Object params, String transportId) {
+		McpServerTransport transport = getTransport(transportId);
+		if (transport == null) {
+			// Fallback to listening transport if specific transport not found
+			transport = getTransport(LISTENING_TRANSPORT);
+			if (transport == null) {
+				return Mono.error(new RuntimeException("Transport not found: " + transportId));
+			}
+		}
 		McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
 				method, params);
-		return this.listeningTransport.sendMessage(jsonrpcNotification);
+		return transport.sendMessage(jsonrpcNotification);
 	}
 
 	/**
@@ -300,26 +335,20 @@ public Mono<Void> handle(McpSchema.JSONRPCMessage message) {
 			}
 			else if (message instanceof McpSchema.JSONRPCRequest request) {
 				logger.debug("Received request: {}", request);
-				final String transportId;
-				if (transports.isEmpty()) {
-					transportId = LISTENING_TRANSPORT;
-				}
-				else {
-					transportId = request.id().toString();
-				}
+				final String transportId = determineTransportId(request);
 				return handleIncomingRequest(request).onErrorResume(error -> {
 					var errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null,
 							new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR,
 									error.getMessage(), null));
-					McpServerTransport transport = getTransport(transportId);
+					McpServerTransport transport = getTransportWithFallback(transportId);
 					return transport != null ? transport.sendMessage(errorResponse).then(Mono.empty()) : Mono.empty();
 				}).flatMap(response -> {
-					McpServerTransport transport = getTransport(transportId);
+					McpServerTransport transport = getTransportWithFallback(transportId);
 					if (transport != null) {
 						return transport.sendMessage(response);
 					}
 					else {
-						return Mono.error(new RuntimeException("Transport not found: " + transportId));
+						return Mono.error(new RuntimeException("No transport available"));
 					}
 				});
 			}
@@ -369,10 +398,10 @@ private Mono<McpSchema.JSONRPCResponse> handleIncomingRequest(McpSchema.JSONRPCR
 									error.message(), error.data())));
 				}
 
-				// We would need to add request.id() as a parameter to handler.handle() if
-				// we want client-request-driven requests/notifications to go to the
-				// related stream
-				resultMono = this.exchangeSink.asMono().flatMap(exchange -> handler.handle(exchange, request.params()));
+				McpAsyncServerExchange requestExchange = new McpAsyncServerExchange(this, clientCapabilities.get(),
+						clientInfo.get(), determineTransportId(request));
+				requestExchange.setMinLoggingLevel(minLoggingLevel);
+				resultMono = handler.handle(requestExchange, request.params());
 			}
 			return resultMono
 				.map(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), result, null))
@@ -392,7 +421,6 @@ private Mono<Void> handleIncomingNotification(McpSchema.JSONRPCNotification noti
 		return Mono.defer(() -> {
 			if (McpSchema.METHOD_NOTIFICATION_INITIALIZED.equals(notification.method())) {
 				this.state.lazySet(STATE_INITIALIZED);
-				exchangeSink.tryEmitValue(new McpAsyncServerExchange(this, clientCapabilities.get(), clientInfo.get()));
 				return this.initNotificationHandler.handle();
 			}
 
@@ -401,7 +429,10 @@ private Mono<Void> handleIncomingNotification(McpSchema.JSONRPCNotification noti
 				logger.error("No handler registered for notification method: {}", notification.method());
 				return Mono.empty();
 			}
-			return this.exchangeSink.asMono().flatMap(exchange -> handler.handle(exchange, notification.params()));
+			McpAsyncServerExchange notificationExchange = new McpAsyncServerExchange(this, clientCapabilities.get(),
+					clientInfo.get(), LISTENING_TRANSPORT);
+			notificationExchange.setMinLoggingLevel(minLoggingLevel);
+			return handler.handle(notificationExchange, notification.params());
 		});
 	}
 
@@ -412,6 +443,32 @@ private MethodNotFoundError getMethodNotFoundError(String method) {
 		return new MethodNotFoundError(method, "Method not found: " + method, null);
 	}
 
+	/**
+	 * Determines the appropriate transport ID for a request. Uses request ID for
+	 * per-request routing only if a transport with that ID exists, otherwise falls back
+	 * to listening transport.
+	 */
+	private String determineTransportId(McpSchema.JSONRPCRequest request) {
+		String requestTransportId = request.id().toString();
+		// Check if a transport exists for this specific request ID
+		if (getTransport(requestTransportId) != null) {
+			return requestTransportId;
+		}
+		// Fallback to listening transport
+		return LISTENING_TRANSPORT;
+	}
+
+	/**
+	 * Gets a transport with fallback to listening transport.
+	 */
+	private McpServerTransport getTransportWithFallback(String transportId) {
+		McpServerTransport transport = getTransport(transportId);
+		if (transport == null) {
+			transport = getTransport(LISTENING_TRANSPORT);
+		}
+		return transport;
+	}
+
 	@Override
 	public Mono<Void> closeGracefully() {
 		return Mono.defer(() -> {