diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java index cacb30522..d811862da 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java @@ -768,8 +768,19 @@ public Mono sendMessage(McpSchema.JSONRPCMessage message, String messageId } catch (Exception e) { logger.error("Failed to send message to session {}: {}", this.sessionId, e.getMessage()); - HttpServletStreamableServerTransportProvider.this.sessions.remove(this.sessionId); - this.asyncContext.complete(); + // A failed write breaks this SSE stream, not the session itself: the + // client may reopen a stream (GET with Last-Event-ID) or keep + // POSTing. + // Session removal stays with explicit DELETE handling and lifecycle + // events, mirroring close() below. + this.closed = true; + try { + this.asyncContext.complete(); + } + catch (Exception completionError) { + logger.warn("Failed to complete async context for session {}: {}", this.sessionId, + completionError.getMessage()); + } } finally { lock.unlock(); diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProviderSessionRetentionTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProviderSessionRetentionTests.java new file mode 100644 index 000000000..cb010c7fd --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProviderSessionRetentionTests.java @@ -0,0 +1,197 @@ +/* + * Copyright 2025 - 2025 the original author or authors. + */ + +package io.modelcontextprotocol.server.transport; + +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.PrintWriter; +import java.io.StringWriter; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.Map; + +import com.google.gson.Gson; +import com.google.gson.GsonBuilder; +import com.google.gson.JsonObject; +import com.google.gson.ToNumberPolicy; +import com.google.gson.JsonSerializer; +import io.modelcontextprotocol.json.McpJsonMapper; +import io.modelcontextprotocol.spec.HttpHeaders; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.McpStreamableServerSession; +import io.modelcontextprotocol.spec.json.gson.GsonMcpJsonMapper; +import jakarta.servlet.AsyncContext; +import jakarta.servlet.ReadListener; +import jakarta.servlet.ServletInputStream; +import jakarta.servlet.http.HttpServletRequest; +import jakarta.servlet.http.HttpServletResponse; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import reactor.core.publisher.Mono; + +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +/** + * Tests that a failed write on one SSE stream does not evict the whole streamable HTTP + * session, while explicit DELETE requests still do. + */ +class HttpServletStreamableServerTransportProviderSessionRetentionTests { + + private static final String INIT_BODY = """ + {"jsonrpc":"2.0","id":"1","method":"initialize","params":{"protocolVersion":"2025-06-18", + "capabilities":{},"clientInfo":{"name":"it","version":"1.0"}}}"""; + + private static final String NOTIFICATION_BODY = """ + {"jsonrpc":"2.0","method":"notifications/test","params":{}}"""; + + private HttpServletStreamableServerTransportProvider provider; + + private final StringWriter initResponseBody = new StringWriter(); + + @BeforeEach + void setUp() { + this.provider = HttpServletStreamableServerTransportProvider.builder() + .jsonMapper(jsonMapper()) + .mcpEndpoint("/mcp") + .build(); + this.provider.setSessionFactory(this::startSession); + } + + private McpJsonMapper jsonMapper() { + // The test Gson needs a Throwable adapter: responseError serializes McpError + // (a RuntimeException), and reflective serialization of Throwable fields is + // not accessible on JDK 17+. + Gson gson = new GsonBuilder().serializeNulls() + .setObjectToNumberStrategy(ToNumberPolicy.LONG_OR_DOUBLE) + .setNumberToNumberStrategy(ToNumberPolicy.LONG_OR_DOUBLE) + .registerTypeHierarchyAdapter(Throwable.class, (JsonSerializer) (src, typeOfSrc, context) -> { + JsonObject error = new JsonObject(); + error.addProperty("message", src.getMessage()); + return error; + }) + .create(); + return new GsonMcpJsonMapper(gson); + } + + private McpStreamableServerSession.McpStreamableServerSessionInit startSession( + McpSchema.InitializeRequest initializeRequest) { + McpStreamableServerSession session = new McpStreamableServerSession("test-session-id", + new McpSchema.ClientCapabilities(null, null, null, null), + new McpSchema.Implementation("it", null, "1.0", null, null, null), Duration.ofSeconds(5), Map.of(), + Map.of()); + McpSchema.InitializeResult initResult = new McpSchema.InitializeResult("2025-06-18", + McpSchema.ServerCapabilities.builder().build(), + new McpSchema.Implementation("test-server", null, "1.0", null, null, null), null, null); + return new McpStreamableServerSession.McpStreamableServerSessionInit(session, Mono.just(initResult)); + } + + @Test + void sessionSurvivesFailedWriteOnSseStream() throws Exception { + String sessionId = initializeSession(); + + openBrokenSseStream(sessionId); + + // The write failure happens inside the notification delivery + this.provider.notifyClient(sessionId, "notifications/test", Map.of()).block(); + + // The session must still be usable for subsequent client requests + HttpServletRequest post = postRequest(NOTIFICATION_BODY, sessionId); + HttpServletResponse response = mockResponse(new PrintWriter(new StringWriter())); + this.provider.doPost(post, response); + + verify(response).setStatus(HttpServletResponse.SC_ACCEPTED); + } + + @Test + void explicitDeleteStillRemovesSession() throws Exception { + String sessionId = initializeSession(); + + HttpServletRequest delete = mock(HttpServletRequest.class); + when(delete.getRequestURI()).thenReturn("/mcp"); + when(delete.getHeader(HttpHeaders.MCP_SESSION_ID)).thenReturn(sessionId); + HttpServletResponse deleteResponse = mockResponse(new PrintWriter(new StringWriter())); + this.provider.doDelete(delete, deleteResponse); + verify(deleteResponse).setStatus(HttpServletResponse.SC_OK); + + HttpServletRequest post = postRequest(NOTIFICATION_BODY, sessionId); + HttpServletResponse postResponse = mockResponse(new PrintWriter(new StringWriter())); + this.provider.doPost(post, postResponse); + + verify(postResponse).setStatus(HttpServletResponse.SC_NOT_FOUND); + } + + private String initializeSession() throws Exception { + HttpServletRequest post = postRequest(INIT_BODY, null); + HttpServletResponse response = mockResponse(new PrintWriter(this.initResponseBody)); + this.provider.doPost(post, response); + + ArgumentCaptor sessionId = ArgumentCaptor.forClass(String.class); + verify(response).setHeader(eq(HttpHeaders.MCP_SESSION_ID), sessionId.capture()); + return sessionId.getValue(); + } + + private void openBrokenSseStream(String sessionId) throws Exception { + HttpServletRequest get = mock(HttpServletRequest.class); + when(get.getRequestURI()).thenReturn("/mcp"); + when(get.getHeader(HttpHeaders.ACCEPT)).thenReturn("text/event-stream"); + when(get.getHeader(HttpHeaders.MCP_SESSION_ID)).thenReturn(sessionId); + + PrintWriter brokenWriter = mock(PrintWriter.class); + when(brokenWriter.checkError()).thenReturn(true); + + AsyncContext asyncContext = mock(AsyncContext.class); + when(get.startAsync()).thenReturn(asyncContext); + + HttpServletResponse response = mockResponse(brokenWriter); + this.provider.doGet(get, response); + } + + private HttpServletRequest postRequest(String body, String sessionId) throws IOException { + HttpServletRequest request = mock(HttpServletRequest.class); + when(request.getRequestURI()).thenReturn("/mcp"); + when(request.getHeader(HttpHeaders.ACCEPT)).thenReturn("application/json, text/event-stream"); + if (sessionId != null) { + when(request.getHeader(HttpHeaders.MCP_SESSION_ID)).thenReturn(sessionId); + } + when(request.getContentLengthLong()).thenReturn((long) body.getBytes(StandardCharsets.UTF_8).length); + when(request.getInputStream()).thenReturn(bodyStream(body)); + return request; + } + + private HttpServletResponse mockResponse(PrintWriter writer) throws IOException { + HttpServletResponse response = mock(HttpServletResponse.class); + when(response.getWriter()).thenReturn(writer); + return response; + } + + private static ServletInputStream bodyStream(String body) { + ByteArrayInputStream source = new ByteArrayInputStream(body.getBytes(StandardCharsets.UTF_8)); + return new ServletInputStream() { + @Override + public boolean isFinished() { + return source.available() == 0; + } + + @Override + public boolean isReady() { + return true; + } + + @Override + public void setReadListener(ReadListener listener) { + } + + @Override + public int read() { + return source.read(); + } + }; + } + +}