fix: Add null check for response ID to prevent memory leaks

- Check if response.id() is not null before processing in MCP session classes
- Log error when MCP response lacks session ID to warn about potential memory leaks
- Improve error handling in McpStreamableServerSession with proper error codes

The missing null check could lead to memory leaks as pending requests
would never be completed when responses lack session IDs. This fix ensures
proper handling of such cases with appropriate error logging.

Resolves #506

Signed-off-by: Yanming Zhou <[email protected]>

Author: quaff

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

@@ -35,6 +35,7 @@
  *
  * @author Christian Tzolov
  * @author Dariusz Jędrzejczyk
+ * @author Yanming Zhou
  */
 public class McpClientSession implements McpSession {
 
@@ -146,13 +147,20 @@ private void dismissPendingResponses() {
 
 	private void handle(McpSchema.JSONRPCMessage message) {
 		if (message instanceof McpSchema.JSONRPCResponse response) {
-			logger.debug("Received Response: {}", response);
-			var sink = pendingResponses.remove(response.id());
-			if (sink == null) {
-				logger.warn("Unexpected response for unknown id {}", response.id());
+			logger.debug("Received response: {}", response);
+			if (response.id() != null) {
+				var sink = pendingResponses.remove(response.id());
+				if (sink == null) {
+					logger.warn("Unexpected response for unknown id {}", response.id());
+				}
+				else {
+					sink.success(response);
+				}
 			}
 			else {
-				sink.success(response);
+				logger.error("Discarded MCP request response without session id. "
+						+ "This is an indication of a bug in the request sender code that can lead to memory "
+						+ "leaks as pending requests will never be completed.");
 			}
 		}
 		else if (message instanceof McpSchema.JSONRPCRequest request) {

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

@@ -204,13 +204,20 @@ public Mono<Void> handle(McpSchema.JSONRPCMessage message) {
 			// TODO handle errors for communication to without initialization happening
 			// first
 			if (message instanceof McpSchema.JSONRPCResponse response) {
-				logger.debug("Received Response: {}", response);
-				var sink = pendingResponses.remove(response.id());
-				if (sink == null) {
-					logger.warn("Unexpected response for unknown id {}", response.id());
+				logger.debug("Received response: {}", response);
+				if (response.id() != null) {
+					var sink = pendingResponses.remove(response.id());
+					if (sink == null) {
+						logger.warn("Unexpected response for unknown id {}", response.id());
+					}
+					else {
+						sink.success(response);
+					}
 				}
 				else {
-					sink.success(response);
+					logger.error("Discarded MCP request response without session id. "
+							+ "This is an indication of a bug in the request sender code that can lead to memory "
+							+ "leaks as pending requests will never be completed.");
 				}
 				return Mono.empty();
 			}

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

@@ -21,6 +21,7 @@
 import io.modelcontextprotocol.server.McpAsyncServerExchange;
 import io.modelcontextprotocol.server.McpNotificationHandler;
 import io.modelcontextprotocol.server.McpRequestHandler;
+import io.modelcontextprotocol.spec.McpSchema.ErrorCodes;
 import io.modelcontextprotocol.util.Assert;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
@@ -33,6 +34,7 @@
  * capability without the insight into the transport-specific details of HTTP handling.
  *
  * @author Dariusz Jędrzejczyk
+ * @author Yanming Zhou
  */
 public class McpStreamableServerSession implements McpLoggableSession {
 
@@ -214,19 +216,30 @@ public Mono<Void> accept(McpSchema.JSONRPCNotification notification) {
 	 */
 	public Mono<Void> accept(McpSchema.JSONRPCResponse response) {
 		return Mono.defer(() -> {
-			var stream = this.requestIdToStream.get(response.id());
-			if (stream == null) {
-				return Mono.error(new McpError("Unexpected response for unknown id " + response.id())); // TODO
-																										// JSONize
-			}
-			// TODO: encapsulate this inside the stream itself
-			var sink = stream.pendingResponses.remove(response.id());
-			if (sink == null) {
-				return Mono.error(new McpError("Unexpected response for unknown id " + response.id())); // TODO
-																										// JSONize
+			logger.debug("Received response: {}", response);
+
+			if (response.id() != null) {
+				var stream = this.requestIdToStream.get(response.id());
+				if (stream == null) {
+					return Mono.error(McpError.builder(ErrorCodes.INTERNAL_ERROR)
+						.message("Unexpected response for unknown id " + response.id())
+						.build());
+				}
+				// TODO: encapsulate this inside the stream itself
+				var sink = stream.pendingResponses.remove(response.id());
+				if (sink == null) {
+					return Mono.error(McpError.builder(ErrorCodes.INTERNAL_ERROR)
+						.message("Unexpected response for unknown id " + response.id())
+						.build());
+				}
+				else {
+					sink.success(response);
+				}
 			}
 			else {
-				sink.success(response);
+				logger.error("Discarded MCP request response without session id. "
+						+ "This is an indication of a bug in the request sender code that can lead to memory "
+						+ "leaks as pending requests will never be completed.");
 			}
 			return Mono.empty();
 		});