Fix race condition in subscription/listen

Under highly contentious load, there might be a race condition in subscription/listen.
Every change broadcast before subscribing to the change feed is lost.
Subscription happens on a thread, and writing to the wire happens in another thread.

The `ack` might be fully written to the transport, and received by the client,
before the changes are subscribed to.
A test harness might fire a change right after the test client receives `ack`,
but before the changes are subscribed to, and therefore the change is lost.

So we need to ensure that:
1. Changes are subscribed to BEFORE the ack
2. But the ACK goes first in any case (and only once)

This is mostly a test-harness fix rather than a real-world scenario.

Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>

Author: Kehrlann

mcp-core/src/main/java/io/modelcontextprotocol/modern/server/SubscriptionsFeature.java

@@ -104,16 +104,26 @@ public void capabilities(ServerCapabilities.Builder builder) {
 	}
 
 	private Mono<Result> listen(McpRequestContext ctx, SubscriptionFilter honoured, McpAsyncNotifier notifier) {
+		// To avoid race conditions in test harnesses, need to ensure that:
+		// 1. Changes are subscribed to BEFORE the ack
+		// 2. But the ACK goes before any notification
+		//
+		// Achieve this by:
+		// 1. Merging changes + ack
+		// 2. Forcing ack to be sent first in the changes
+		// 3. Caching the "send ack" so that it only ever goes out once
 		Map<String, Object> subscriptionMeta = Map.of(MetaKeys.SUBSCRIPTION_ID, ctx.requestId());
-		Mono<Void> ack = notifier.notify(McpSchema.METHOD_NOTIFICATION_SUBSCRIPTIONS_ACKNOWLEDGED,
-				new SubscriptionsAcknowledgedParams(honoured, subscriptionMeta));
+		Mono<Void> ack = notifier
+			.notify(McpSchema.METHOD_NOTIFICATION_SUBSCRIPTIONS_ACKNOWLEDGED,
+					new SubscriptionsAcknowledgedParams(honoured, subscriptionMeta))
+			.cache();
 
 		Flux<Void> forwardChanges = this.publisher.changes()
 			.filter(change -> matches(change, honoured))
 			.takeUntilOther(this.shutdown.asMono())
-			.concatMap(change -> notifier.notify(methodFor(change), paramsFor(change, subscriptionMeta)));
+			.concatMap(change -> ack.then(notifier.notify(methodFor(change), paramsFor(change, subscriptionMeta))));
 
-		return ack.thenMany(forwardChanges)
+		return Flux.merge(forwardChanges, ack)
 			.then(Mono.fromSupplier(() -> (Result) SubscriptionsListenResult.forSubscription(ctx.requestId())));
 	}
 

mcp-test/src/test/java/io/modelcontextprotocol/modern/server/SubscriptionsFeatureTests.java

@@ -4,13 +4,15 @@
 
 package io.modelcontextprotocol.modern.server;
 
+import java.time.Duration;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Stream;
 
 import io.modelcontextprotocol.common.McpTransportContext;
 import io.modelcontextprotocol.json.McpJsonDefaults;
+import io.modelcontextprotocol.modern.JsonRpc.JSONRPCMessage;
 import io.modelcontextprotocol.modern.JsonRpc.JSONRPCNotification;
 import io.modelcontextprotocol.modern.JsonRpc.JSONRPCRequest;
 import io.modelcontextprotocol.modern.JsonRpc.JSONRPCResponse;
@@ -32,6 +34,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
+import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
@@ -283,6 +286,54 @@ void resourceSubscribeIsNotAdvertisedWithoutSubscriptions() {
 		assertThat(resourcesCapability(server)).containsEntry("subscribe", false);
 	}
 
+	@Test
+	void changeBroadcastAsSoonAsAckIsDeliveredReachesTheStream() {
+		McpChangeFeed feed = new McpChangeFeed();
+		McpServer server = toolsServer(feed);
+
+		// Trigger a change from the subscribing thread the change is broadcast right when
+		// ack is received.
+		// This simulates high CPU loads in test harnesses, were writing the ack might
+		// happen
+		// sooner than any subsequent event subscription.
+		Flux<JSONRPCMessage> messages = listen(server, 1).messages().doOnNext(msg -> {
+			if (msg instanceof JSONRPCNotification notification
+					&& notification.method().equals(McpSchema.METHOD_NOTIFICATION_SUBSCRIPTIONS_ACKNOWLEDGED)) {
+				feed.broadcast(new ServerChange.ToolsListChanged());
+			}
+		});
+
+		// A lost change never arrives: time out instead of hanging.
+		StepVerifier.create(messages)
+			.expectNextMatches(msg -> ((JSONRPCNotification) msg).method()
+				.equals(McpSchema.METHOD_NOTIFICATION_SUBSCRIPTIONS_ACKNOWLEDGED))
+			.expectNextMatches(msg -> ((JSONRPCNotification) msg).method()
+				.equals(McpSchema.METHOD_NOTIFICATION_TOOLS_LIST_CHANGED))
+			.then(server::closeGracefully)
+			.expectNextMatches(msg -> msg instanceof JSONRPCResponse)
+			.expectComplete()
+			.verify(Duration.ofSeconds(5));
+	}
+
+	@Test
+	void changeDeliveredBeforeAckIsSentFollowsTheAck() {
+		// Delivers a change as soon as the listen stream subscribes, before the ack is
+		// out.
+		McpChangePublisher publisher = () -> Flux.<ServerChange>just(new ServerChange.ToolsListChanged())
+			.concatWith(Flux.never());
+		McpServer server = toolsServer(publisher);
+
+		StepVerifier.create(listen(server, 1).messages())
+			.expectNextMatches(msg -> ((JSONRPCNotification) msg).method()
+				.equals(McpSchema.METHOD_NOTIFICATION_SUBSCRIPTIONS_ACKNOWLEDGED))
+			.expectNextMatches(msg -> ((JSONRPCNotification) msg).method()
+				.equals(McpSchema.METHOD_NOTIFICATION_TOOLS_LIST_CHANGED))
+			.then(server::closeGracefully)
+			.expectNextMatches(msg -> msg instanceof JSONRPCResponse)
+			.expectComplete()
+			.verify(Duration.ofSeconds(5));
+	}
+
 	private static McpServer resourcesServer(boolean supportsSubscribe, McpChangePublisher publisher) {
 		McpAsyncResourceRepository repository = new McpAsyncResourceRepository() {
 			@Override