modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
- 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 <zhouyanming@gmail.com>
This commit is contained in:
Yanming Zhou 2025-08-21 17:10:33 +08:00 • committed by Christian Tzolov
commit e94a93a91b
3 changed files with 49 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

18
mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java
View file

@ -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) {

17
mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java
View file

@ -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();
}

35
mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java
View file

@ -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<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();
});