diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/SubscriptionsFeature.java b/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/SubscriptionsFeature.java index fa50d9344..da0b22f22 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/SubscriptionsFeature.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/modern/server/SubscriptionsFeature.java @@ -104,16 +104,26 @@ public void capabilities(ServerCapabilities.Builder builder) { } private Mono 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 subscriptionMeta = Map.of(MetaKeys.SUBSCRIPTION_ID, ctx.requestId()); - Mono ack = notifier.notify(McpSchema.METHOD_NOTIFICATION_SUBSCRIPTIONS_ACKNOWLEDGED, - new SubscriptionsAcknowledgedParams(honoured, subscriptionMeta)); + Mono ack = notifier + .notify(McpSchema.METHOD_NOTIFICATION_SUBSCRIPTIONS_ACKNOWLEDGED, + new SubscriptionsAcknowledgedParams(honoured, subscriptionMeta)) + .cache(); Flux 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()))); } diff --git a/mcp-test/src/test/java/io/modelcontextprotocol/modern/server/SubscriptionsFeatureTests.java b/mcp-test/src/test/java/io/modelcontextprotocol/modern/server/SubscriptionsFeatureTests.java index 115a7d8e5..d545868bc 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/modern/server/SubscriptionsFeatureTests.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/modern/server/SubscriptionsFeatureTests.java @@ -4,6 +4,7 @@ package io.modelcontextprotocol.modern.server; +import java.time.Duration; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -11,6 +12,7 @@ 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 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.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