From d8cc189341497b34d8932f18c8a28ab3b8e3213c Mon Sep 17 00:00:00 2001 From: Yanming Zhou Date: Thu, 21 Aug 2025 17:10:33 +0800 Subject: [PATCH] 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 --- .../spec/McpClientSession.java | 18 +++++++--- .../spec/McpServerSession.java | 17 ++++++--- .../spec/McpStreamableServerSession.java | 35 +++++++++++++------ 3 files changed, 49 insertions(+), 21 deletions(-) diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java index f7db3d7..6ac8def 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java @@ -35,6 +35,7 @@ import java.util.function.Function; * * @author Christian Tzolov * @author Dariusz Jędrzejczyk + * @author Yanming Zhou */ public class McpClientSession implements McpSession { @@ -146,13 +147,20 @@ public class McpClientSession implements McpSession { 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) { diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java index e6a0c8b..747b454 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java @@ -204,13 +204,20 @@ public class McpServerSession implements McpLoggableSession { // 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(); } diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java index af29ce0..53b56c7 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java @@ -21,6 +21,7 @@ import io.modelcontextprotocol.common.McpTransportContext; 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 @@ import reactor.core.publisher.MonoSink; * 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 class McpStreamableServerSession implements McpLoggableSession { */ public Mono 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(); });