HttpClientStreamableHttpTransport: a late 404 for an old session invalidates the current, healthy session and fails its in-flight requests #1270

package repro;

import static java.nio.charset.StandardCharsets.UTF_8;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;

import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpServer;
import io.modelcontextprotocol.client.McpAsyncClient;
import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport;
import io.modelcontextprotocol.json.McpJsonDefaults;
import io.modelcontextprotocol.json.McpJsonMapper;
import io.modelcontextprotocol.spec.McpSchema;
import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;

@Timeout(30)
class StaleSessionNotFoundTest {

	@Test
	void late404ForOldSessionTearsDownTheNewSession() throws Exception {
		try (FakeServer server = new FakeServer()) {
			McpAsyncClient client = McpClient
				.async(HttpClientStreamableHttpTransport.builder(server.baseUri()).build())
				.requestTimeout(Duration.ofSeconds(20))
				.build();
			try {
				client.initialize().block(Duration.ofSeconds(5)); // session s1

				call(client, "A"); // the server holds A, which carries s1
				assertEquals("s1", server.arrived("A").get(5, SECONDS));

				server.expire("s1");
				CompletableFuture<?> b = call(client, "B"); // 404 -> client re-initializes
				assertThrows(ExecutionException.class, () -> b.get(5, SECONDS));

				CompletableFuture<?> c = call(client, "C"); // sent on the new session s2, held
				assertEquals("s2", server.arrived("C").get(5, SECONDS));

				server.release("A"); // late 404 for A, i.e. for s1

				// Actual (2.1.0): C fails and the client itself deletes the healthy s2.
				ExecutionException cFailure = assertThrows(ExecutionException.class, () -> c.get(5, SECONDS));
				assertEquals("MCP session with server terminated", cFailure.getCause().getMessage());
				server.deleted("s2").get(5, SECONDS);
			}
			finally {
				client.close();
			}
		}
	}

	private static CompletableFuture<McpSchema.CallToolResult> call(McpAsyncClient client, String tool) {
		return client.callTool(McpSchema.CallToolRequest.builder(tool).build()).toFuture();
	}

	/**
	 * Minimal stateful Streamable HTTP server. Tools A and C are answered only once
	 * released; a request whose session is unknown at answer time gets 404. The DELETE
	 * for an expired session is answered only at close: while it is outstanding the
	 * client keeps the old session's exchanges (here A) subscribed.
	 */
	static final class FakeServer implements AutoCloseable {

		private static final McpJsonMapper JSON = McpJsonDefaults.getMapper();

		private final ExecutorService pool = Executors.newCachedThreadPool();

		private final HttpServer http;

		private final AtomicInteger sessionCounter = new AtomicInteger();

		private final Set<String> live = ConcurrentHashMap.newKeySet();

		private final Map<String, CountDownLatch> held = Map.of("A", new CountDownLatch(1), "C",
				new CountDownLatch(1));

		private final Map<String, CompletableFuture<String>> arrived = new ConcurrentHashMap<>();

		private final Map<String, CompletableFuture<Void>> deleted = new ConcurrentHashMap<>();

		private final CountDownLatch closing = new CountDownLatch(1);

		FakeServer() throws IOException {
			this.http = HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0);
			this.http.setExecutor(this.pool);
			this.http.createContext("/mcp", this::handle);
			this.http.start();
		}

		String baseUri() {
			return "http://127.0.0.1:" + this.http.getAddress().getPort();
		}

		void expire(String sessionId) {
			this.live.remove(sessionId);
		}

		void release(String tool) {
			this.held.get(tool).countDown();
		}

		CompletableFuture<String> arrived(String tool) {
			return this.arrived.computeIfAbsent(tool, k -> new CompletableFuture<>());
		}

		CompletableFuture<Void> deleted(String sessionId) {
			return this.deleted.computeIfAbsent(sessionId, k -> new CompletableFuture<>());
		}

		private void handle(HttpExchange exchange) throws IOException {
			try (exchange) {
				String sessionId = exchange.getRequestHeaders().getFirst("Mcp-Session-Id");
				switch (exchange.getRequestMethod()) {
					case "GET" -> exchange.sendResponseHeaders(405, -1); // no standalone SSE
					case "DELETE" -> {
						if (!this.live.remove(sessionId)) {
							this.closing.await(20, SECONDS);
						}
						deleted(sessionId).complete(null);
						exchange.sendResponseHeaders(200, -1);
					}
					default -> post(exchange, sessionId);
				}
			}
			catch (IOException | InterruptedException ex) {
				// client went away or test is closing
			}
		}

		private void post(HttpExchange exchange, String sessionId) throws IOException, InterruptedException {
			var message = McpSchema.deserializeJsonRpcMessage(JSON,
					new String(exchange.getRequestBody().readAllBytes(), UTF_8));
			if (message instanceof McpSchema.JSONRPCRequest request && "initialize".equals(request.method())) {
				String newSession = "s" + this.sessionCounter.incrementAndGet();
				this.live.add(newSession);
				exchange.getResponseHeaders().add("Mcp-Session-Id", newSession);
				reply(exchange, request.id(), Map.of("protocolVersion", "2025-11-25", "capabilities",
						Map.of("tools", Map.of()), "serverInfo", Map.of("name", "fake", "version", "1")));
				return;
			}
			if (message instanceof McpSchema.JSONRPCRequest request && "tools/call".equals(request.method())) {
				String tool = (String) ((Map<?, ?>) request.params()).get("name");
				arrived(tool).complete(sessionId);
				CountDownLatch latch = this.held.get(tool);
				if (latch != null) {
					latch.await(20, SECONDS);
				}
				if (this.live.contains(sessionId)) {
					reply(exchange, request.id(), Map.of("content", List.of(Map.of("type", "text", "text", tool))));
					return;
				}
			}
			else if (this.live.contains(sessionId)) { // notifications/initialized
				exchange.sendResponseHeaders(202, -1);
				return;
			}
			exchange.sendResponseHeaders(404, -1);
		}

		private static void reply(HttpExchange exchange, Object id, Object result) throws IOException {
			byte[] body = JSON.writeValueAsString(McpSchema.JSONRPCResponse.result(id, result)).getBytes(UTF_8);
			exchange.getResponseHeaders().add("Content-Type", "application/json");
			exchange.sendResponseHeaders(200, body.length);
			exchange.getResponseBody().write(body);
		}

		@Override
		public void close() {
			this.closing.countDown();
			this.held.values().forEach(CountDownLatch::countDown);
			this.http.stop(0);
			this.pool.shutdownNow();
		}

	}

}