feat(server): clean up

Author: Aliaksie

mcp/src/main/java/io/modelcontextprotocol/server/transport/StreamableHttpServerTransportProvider.java

@@ -149,8 +149,6 @@ protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws S
 			return;
 		}
 
-		// resp
-		resp.setStatus(HttpServletResponse.SC_OK);
 		resp.setCharacterEncoding("UTF-8");
 
 		final McpServerTransport transport = new StreamableHttpServerTransport(resp.getOutputStream(), objectMapper);
@@ -168,14 +166,17 @@ protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws S
 		try {
 			messages = parseRequestBodyAsStream(req);
 		}
-		catch (Exception e) {
+		catch (final Exception e) {
 			resp.sendError(HttpServletResponse.SC_BAD_REQUEST, "Invalid JSON input");
 			return;
 		}
 
 		boolean hasRequest = messages.stream().anyMatch(m -> m instanceof McpSchema.JSONRPCRequest);
 		if (!hasRequest) {
 			resp.setStatus(HttpServletResponse.SC_ACCEPTED);
+			if ("stateless".equals(session.getId())) {
+				transport.closeGracefully().subscribe();
+			}
 			return;
 		}
 
@@ -190,24 +191,17 @@ protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws S
 			Flux.fromIterable(messages)
 				.flatMap(session::handle)
 				.doOnError(e -> sendError(resp, 500, "Streaming failed: " + e.getMessage()))
-				.then(transport.closeGracefully())
+				.then(closeIfStateless(session, transport))
 				.subscribe();
 		}
 		else if (accept.contains(APPLICATION_JSON)) {
 			// TODO: Handle traditional JSON-RPC,
 			resp.setContentType(APPLICATION_JSON);
-			Flux.fromIterable(messages).flatMap(session::handle).collectList().flatMap(responses -> {
-				try {
-					// todo: collect result if it's a response,
-					// hm handle should not be void ...
-					String json = objectMapper.writeValueAsString(responses.size() == 1 ? responses.get(0) : responses);
-					resp.getWriter().write(json);
-					return transport.closeGracefully();
-				}
-				catch (IOException e) {
-					return Mono.error(e);
-				}
-			}).doOnError(e -> sendError(resp, 500, "JSON response failed: " + e.getMessage())).subscribe();
+			Flux.fromIterable(messages)
+					.flatMap(session::handle)
+					.doOnError(e -> sendError(resp, 500, "Streaming failed: " + e.getMessage()))
+					.then(closeIfStateless(session, transport))
+					.subscribe();
 
 		}
 		else {
@@ -266,7 +260,7 @@ protected void doDelete(HttpServletRequest req, HttpServletResponse resp) throws
 		resp.setStatus(HttpServletResponse.SC_NO_CONTENT);
 	}
 
-	private List<McpSchema.JSONRPCMessage> parseRequestBodyAsStream(final HttpServletRequest req) {
+	private List<McpSchema.JSONRPCMessage> parseRequestBodyAsStream(final HttpServletRequest req) throws IOException {
 		try (final InputStream inputStream = req.getInputStream()) {
 			final JsonNode node = objectMapper.readTree(inputStream);
 			if (node.isArray()) {
@@ -288,10 +282,6 @@ else if (node.isObject()) {
 
 			}
 		}
-		catch (Exception e) {
-			throw new IllegalArgumentException("Invalid JSON-RPC request: " + e.getMessage());
-		}
-
 	}
 
 	private void sendEvent(PrintWriter writer, String eventType, String data) throws IOException {
@@ -320,6 +310,12 @@ else if (sessionFactory != null) {
 		}
 	}
 
+	Mono<Void> closeIfStateless(final McpSession session, final McpServerTransport transport) {
+		return "stateless".equals(session.getId())
+				? transport.closeGracefully()
+				: Mono.empty();
+	}
+
 	private void sendError(final HttpServletResponse resp, final int code, final String msg) {
 		try {
 			resp.sendError(code, msg);