Add resource subscribe/unsubscribe support to McpAsyncServer #838
+269
−11
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,217 @@ | ||
| /* | ||
| * Copyright 2024-2024 the original author or authors. | ||
| */ | ||
|
|
||
| package io.modelcontextprotocol.server; | ||
|
|
||
| import java.time.Duration; | ||
| import java.util.List; | ||
|
|
||
| import io.modelcontextprotocol.MockMcpServerTransport; | ||
| import io.modelcontextprotocol.MockMcpServerTransportProvider; | ||
| import io.modelcontextprotocol.spec.McpSchema; | ||
| import io.modelcontextprotocol.spec.McpSchema.ReadResourceResult; | ||
| import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; | ||
| import io.modelcontextprotocol.spec.ProtocolVersions; | ||
| import org.junit.jupiter.api.AfterEach; | ||
| import org.junit.jupiter.api.BeforeEach; | ||
| import org.junit.jupiter.api.Test; | ||
| import reactor.core.publisher.Mono; | ||
|
|
||
| import static org.assertj.core.api.Assertions.assertThat; | ||
| import static org.assertj.core.api.Assertions.assertThatCode; | ||
|
|
||
| /** | ||
| * Tests for resource subscribe/unsubscribe support in {@link McpAsyncServer}. | ||
| * | ||
| * Covers the handlers registered under {@code resources/subscribe} and | ||
| * {@code resources/unsubscribe}, the capability guard that keeps them absent when | ||
| * {@code subscribe=false}, and the {@code notifications/resources/updated} notification | ||
| * emitted by {@link McpAsyncServer#addResource} when a subscribed URI is (re-)added. | ||
| */ | ||
| class ResourceSubscriptionTests { | ||
|
|
||
| private static final String RESOURCE_URI = "test://resource/item"; | ||
|
|
||
| private MockMcpServerTransport mockTransport; | ||
|
|
||
| private MockMcpServerTransportProvider mockTransportProvider; | ||
|
|
||
| private McpAsyncServer mcpAsyncServer; | ||
|
|
||
| @BeforeEach | ||
| void setUp() { | ||
| mockTransport = new MockMcpServerTransport(); | ||
| mockTransportProvider = new MockMcpServerTransportProvider(mockTransport); | ||
| } | ||
|
|
||
| @AfterEach | ||
| void tearDown() { | ||
| if (mcpAsyncServer != null) { | ||
| assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))) | ||
| .doesNotThrowAnyException(); | ||
| } | ||
| } | ||
|
|
||
| // ------------------------------------------------------------------------- | ||
| // Capability guard | ||
| // ------------------------------------------------------------------------- | ||
|
|
||
| @Test | ||
| void subscribeHandlerNotRegisteredWhenSubscribeCapabilityDisabled() { | ||
| mcpAsyncServer = McpServer.async(mockTransportProvider) | ||
| .serverInfo("test-server", "1.0.0") | ||
| .capabilities(ServerCapabilities.builder().resources(false, false).build()) | ||
| .build(); | ||
|
|
||
| initializeSession(); | ||
| mockTransportProvider.simulateIncomingMessage(subscribeRequest("req-1", RESOURCE_URI)); | ||
|
|
||
| McpSchema.JSONRPCMessage response = mockTransport.getLastSentMessage(); | ||
| assertThat(response).isInstanceOf(McpSchema.JSONRPCResponse.class); | ||
| McpSchema.JSONRPCResponse jsonResponse = (McpSchema.JSONRPCResponse) response; | ||
| assertThat(jsonResponse.error()).isNotNull(); | ||
| assertThat(jsonResponse.error().code()).isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND); | ||
| } | ||
|
|
||
| // ------------------------------------------------------------------------- | ||
| // Subscribe / unsubscribe happy paths | ||
| // ------------------------------------------------------------------------- | ||
|
|
||
| @Test | ||
| void subscribeRequestReturnsEmptyResult() { | ||
| mcpAsyncServer = McpServer.async(mockTransportProvider) | ||
| .serverInfo("test-server", "1.0.0") | ||
| .capabilities(ServerCapabilities.builder().resources(true, false).build()) | ||
| .build(); | ||
|
|
||
| initializeSession(); | ||
| mockTransportProvider.simulateIncomingMessage(subscribeRequest("req-1", RESOURCE_URI)); | ||
|
|
||
| McpSchema.JSONRPCMessage response = mockTransport.getLastSentMessage(); | ||
| assertThat(response).isInstanceOf(McpSchema.JSONRPCResponse.class); | ||
| McpSchema.JSONRPCResponse jsonResponse = (McpSchema.JSONRPCResponse) response; | ||
| assertThat(jsonResponse.id()).isEqualTo("req-1"); | ||
| assertThat(jsonResponse.error()).isNull(); | ||
| } | ||
|
|
||
| // ------------------------------------------------------------------------- | ||
| // addResource notification behaviour | ||
| // ------------------------------------------------------------------------- | ||
|
|
||
| @Test | ||
| void addSubscribedResourceSendsResourcesUpdatedNotification() { | ||
| mcpAsyncServer = McpServer.async(mockTransportProvider) | ||
| .serverInfo("test-server", "1.0.0") | ||
| .capabilities(ServerCapabilities.builder().resources(true, false).build()) | ||
| .build(); | ||
|
|
||
| initializeSession(); | ||
| // Client subscribes first | ||
| mockTransportProvider.simulateIncomingMessage(subscribeRequest("req-1", RESOURCE_URI)); | ||
| mockTransport.clearSentMessages(); | ||
|
|
||
| mcpAsyncServer.addResource(resourceSpec(RESOURCE_URI)).block(Duration.ofSeconds(5)); | ||
|
|
||
| List<McpSchema.JSONRPCMessage> sent = mockTransport.getAllSentMessages(); | ||
| assertThat(sent).hasSize(1); | ||
| assertThat(sent.get(0)).isInstanceOf(McpSchema.JSONRPCNotification.class); | ||
| McpSchema.JSONRPCNotification notification = (McpSchema.JSONRPCNotification) sent.get(0); | ||
| assertThat(notification.method()).isEqualTo(McpSchema.METHOD_NOTIFICATION_RESOURCES_UPDATED); | ||
| } | ||
|
|
||
| @Test | ||
| void addSubscribedResourceWithListChangedSendsBothNotifications() { | ||
| mcpAsyncServer = McpServer.async(mockTransportProvider) | ||
| .serverInfo("test-server", "1.0.0") | ||
| .capabilities(ServerCapabilities.builder().resources(true, true).build()) | ||
| .build(); | ||
|
|
||
| initializeSession(); | ||
| mockTransportProvider.simulateIncomingMessage(subscribeRequest("req-1", RESOURCE_URI)); | ||
| mockTransport.clearSentMessages(); | ||
|
|
||
| mcpAsyncServer.addResource(resourceSpec(RESOURCE_URI)).block(Duration.ofSeconds(5)); | ||
|
|
||
| List<McpSchema.JSONRPCMessage> sent = mockTransport.getAllSentMessages(); | ||
| assertThat(sent).hasSize(2); | ||
|
|
||
| List<String> methods = sent.stream() | ||
| .filter(m -> m instanceof McpSchema.JSONRPCNotification) | ||
| .map(m -> ((McpSchema.JSONRPCNotification) m).method()) | ||
| .toList(); | ||
| assertThat(methods).containsExactly(McpSchema.METHOD_NOTIFICATION_RESOURCES_UPDATED, | ||
| McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED); | ||
| } | ||
|
|
||
| @Test | ||
| void addUnsubscribedResourceDoesNotSendResourcesUpdatedNotification() { | ||
| mcpAsyncServer = McpServer.async(mockTransportProvider) | ||
| .serverInfo("test-server", "1.0.0") | ||
| .capabilities(ServerCapabilities.builder().resources(true, true).build()) | ||
| .build(); | ||
|
|
||
| initializeSession(); | ||
| // No subscribe call — resource URI is not in the subscription map | ||
| mcpAsyncServer.addResource(resourceSpec(RESOURCE_URI)).block(Duration.ofSeconds(5)); | ||
|
|
||
| List<McpSchema.JSONRPCMessage> sent = mockTransport.getAllSentMessages(); | ||
| assertThat(sent).hasSize(1); | ||
| McpSchema.JSONRPCNotification notification = (McpSchema.JSONRPCNotification) sent.get(0); | ||
| assertThat(notification.method()).isEqualTo(McpSchema.METHOD_NOTIFICATION_RESOURCES_LIST_CHANGED); | ||
| } | ||
|
|
||
| @Test | ||
| void addResourceAfterUnsubscribeDoesNotSendResourcesUpdatedNotification() { | ||
| mcpAsyncServer = McpServer.async(mockTransportProvider) | ||
| .serverInfo("test-server", "1.0.0") | ||
| .capabilities(ServerCapabilities.builder().resources(true, false).build()) | ||
| .build(); | ||
|
|
||
| initializeSession(); | ||
| mockTransportProvider.simulateIncomingMessage(subscribeRequest("req-1", RESOURCE_URI)); | ||
| mockTransportProvider.simulateIncomingMessage(unsubscribeRequest("req-2", RESOURCE_URI)); | ||
| mockTransport.clearSentMessages(); | ||
|
|
||
| mcpAsyncServer.addResource(resourceSpec(RESOURCE_URI)).block(Duration.ofSeconds(5)); | ||
|
|
||
| // No notifications expected: not subscribed, listChanged=false | ||
| assertThat(mockTransport.getAllSentMessages()).isEmpty(); | ||
| } | ||
|
|
||
| // ------------------------------------------------------------------------- | ||
| // Helpers | ||
| // ------------------------------------------------------------------------- | ||
|
|
||
| /** | ||
| * Performs the MCP initialization handshake so that the session's exchangeSink is | ||
| * populated and subsequent request handlers can be invoked. | ||
| */ | ||
| private void initializeSession() { | ||
| mockTransportProvider.simulateIncomingMessage(new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, | ||
| McpSchema.METHOD_INITIALIZE, "init-req", new McpSchema.InitializeRequest( | ||
| ProtocolVersions.MCP_2025_11_25, null, new McpSchema.Implementation("test-client", "1.0.0")))); | ||
|
|
||
| mockTransportProvider.simulateIncomingMessage(new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, | ||
| McpSchema.METHOD_NOTIFICATION_INITIALIZED, null)); | ||
|
|
||
| mockTransport.clearSentMessages(); | ||
| } | ||
|
|
||
| private static McpSchema.JSONRPCRequest subscribeRequest(String id, String uri) { | ||
| return new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, McpSchema.METHOD_RESOURCES_SUBSCRIBE, id, | ||
| new McpSchema.SubscribeRequest(uri)); | ||
| } | ||
|
|
||
| private static McpSchema.JSONRPCRequest unsubscribeRequest(String id, String uri) { | ||
| return new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, McpSchema.METHOD_RESOURCES_UNSUBSCRIBE, id, | ||
| new McpSchema.SubscribeRequest(uri)); | ||
| } | ||
|
|
||
| private static McpServerFeatures.AsyncResourceSpecification resourceSpec(String uri) { | ||
| McpSchema.Resource resource = McpSchema.Resource.builder().uri(uri).name("test-resource").build(); | ||
| return new McpServerFeatures.AsyncResourceSpecification(resource, | ||
| (exchange, req) -> Mono.just(new ReadResourceResult(List.of()))); | ||
| } | ||
|
|
||
| } |
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
SubscribeRequestis not needed right now, but perhaps the meta-information can be of use for later. For now I'd keep the field asMap<String, Set<String>>of uri -> set(session_id).io.modelcontextprotocol.server.McpAsyncServer#notifyResourcesUpdatedto filter the clients to notify about a resource update (look at uri fromio.modelcontextprotocol.spec.McpSchema.ResourcesUpdatedNotification#uri)There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I realized this requires the
McpTransportProviderto expose an API that accepts the sessionId or a set of them. For subscription related notifications, it has to be the listening stream (because the act of subscribing which has to receive a response, so even if the request upgrades to an SSE stream, it has to be closed after the response). So it makes sense to expose it at this level.