Resource subscriptions

Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>

Author: chemicL

mcp-core/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java

@@ -5,10 +5,12 @@
 package io.modelcontextprotocol.server;
 
 import java.time.Duration;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.CopyOnWriteArrayList;
@@ -25,7 +27,6 @@
 import io.modelcontextprotocol.spec.McpSchema.CompleteResult.CompleteCompletion;
 import io.modelcontextprotocol.spec.McpSchema.ErrorCodes;
 import io.modelcontextprotocol.spec.McpSchema.LoggingLevel;
-import io.modelcontextprotocol.spec.McpSchema.LoggingMessageNotification;
 import io.modelcontextprotocol.spec.McpSchema.PromptReference;
 import io.modelcontextprotocol.spec.McpSchema.ResourceReference;
 import io.modelcontextprotocol.spec.McpSchema.SetLevelRequest;
@@ -111,12 +112,10 @@ public class McpAsyncServer {
 
 	private final ConcurrentHashMap<String, McpServerFeatures.AsyncPromptSpecification> prompts = new ConcurrentHashMap<>();
 
-	// FIXME: this field is deprecated and should be remvoed together with the
-	// broadcasting loggingNotification.
-	private LoggingLevel minLoggingLevel = LoggingLevel.DEBUG;
-
 	private final ConcurrentHashMap<McpSchema.CompleteReference, McpServerFeatures.AsyncCompletionSpecification> completions = new ConcurrentHashMap<>();
 
+	private final ConcurrentHashMap<String, Set<String>> resourceSubscriptions = new ConcurrentHashMap<>();
+
 	private List<String> protocolVersions;
 
 	private McpUriTemplateManagerFactory uriTemplateManagerFactory = new DefaultMcpUriTemplateManagerFactory();
@@ -149,8 +148,11 @@ public class McpAsyncServer {
 
 		this.protocolVersions = mcpTransportProvider.protocolVersions();
 
-		mcpTransportProvider.setSessionFactory(transport -> new McpServerSession(UUID.randomUUID().toString(),
-				requestTimeout, transport, this::asyncInitializeRequestHandler, requestHandlers, notificationHandlers));
+		mcpTransportProvider.setSessionFactory(transport -> {
+			String sessionId = UUID.randomUUID().toString();
+			return new McpServerSession(sessionId, requestTimeout, transport, this::asyncInitializeRequestHandler,
+					requestHandlers, notificationHandlers, () -> this.cleanupForSession(sessionId));
+		});
 	}
 
 	McpAsyncServer(McpStreamableServerTransportProvider mcpTransportProvider, McpJsonMapper jsonMapper,
@@ -174,8 +176,9 @@ public class McpAsyncServer {
 
 		this.protocolVersions = mcpTransportProvider.protocolVersions();
 
-		mcpTransportProvider.setSessionFactory(new DefaultMcpStreamableServerSessionFactory(requestTimeout,
-				this::asyncInitializeRequestHandler, requestHandlers, notificationHandlers));
+		mcpTransportProvider.setSessionFactory(
+				new DefaultMcpStreamableServerSessionFactory(requestTimeout, this::asyncInitializeRequestHandler,
+						requestHandlers, notificationHandlers, sessionId -> this.cleanupForSession(sessionId)));
 	}
 
 	private Map<String, McpNotificationHandler> prepareNotificationHandlers(McpServerFeatures.Async features) {
@@ -215,6 +218,10 @@ private Map<String, McpRequestHandler<?>> prepareRequestHandlers() {
 			requestHandlers.put(McpSchema.METHOD_RESOURCES_LIST, resourcesListRequestHandler());
 			requestHandlers.put(McpSchema.METHOD_RESOURCES_READ, resourcesReadRequestHandler());
 			requestHandlers.put(McpSchema.METHOD_RESOURCES_TEMPLATES_LIST, resourceTemplateListRequestHandler());
+			if (Boolean.TRUE.equals(this.serverCapabilities.resources().subscribe())) {
+				requestHandlers.put(McpSchema.METHOD_RESOURCES_SUBSCRIBE, resourcesSubscribeRequestHandler());
+				requestHandlers.put(McpSchema.METHOD_RESOURCES_UNSUBSCRIBE, resourcesUnsubscribeRequestHandler());
+			}
 		}
 
 		// Add prompts API handlers if provider exists
@@ -685,12 +692,73 @@ public Mono<Void> notifyResourcesListChanged() {
 	}
 
 	/**
-	 * Notifies clients that the resources have updated.
-	 * @return A Mono that completes when all clients have been notified
+	 * Notifies only the sessions that have subscribed to the updated resource URI.
+	 * @param resourcesUpdatedNotification the notification containing the updated
+	 * resource URI
+	 * @return A Mono that completes when all subscribed sessions have been notified
 	 */
 	public Mono<Void> notifyResourcesUpdated(McpSchema.ResourcesUpdatedNotification resourcesUpdatedNotification) {
-		return this.mcpTransportProvider.notifyClients(McpSchema.METHOD_NOTIFICATION_RESOURCES_UPDATED,
-				resourcesUpdatedNotification);
+		return Mono.defer(() -> {
+			String uri = resourcesUpdatedNotification.uri();
+			Set<String> subscribedSessions = this.resourceSubscriptions.get(uri);
+			if (subscribedSessions == null || subscribedSessions.isEmpty()) {
+				logger.debug("No sessions subscribed to resource URI: {}", uri);
+				return Mono.empty();
+			}
+			return Flux.fromIterable(subscribedSessions)
+				.flatMap(sessionId -> this.mcpTransportProvider
+					.notifyClient(sessionId, McpSchema.METHOD_NOTIFICATION_RESOURCES_UPDATED,
+							resourcesUpdatedNotification)
+					.doOnError(e -> logger.error("Failed to notify session {} of resource update for {}", sessionId,
+							uri, e))
+					.onErrorComplete())
+				.then();
+		});
+	}
+
+	private Mono<Void> cleanupForSession(String sessionId) {
+		return Mono.fromRunnable(() -> {
+			removeSessionSubscriptions(sessionId);
+		});
+	}
+
+	private void removeSessionSubscriptions(String sessionId) {
+		this.resourceSubscriptions.forEach((uri, sessions) -> sessions.remove(sessionId));
+		this.resourceSubscriptions.entrySet().removeIf(entry -> entry.getValue().isEmpty());
+	}
+
+	private McpRequestHandler<Object> resourcesSubscribeRequestHandler() {
+		return (exchange, params) -> Mono.defer(() -> {
+			McpSchema.SubscribeRequest subscribeRequest = jsonMapper.convertValue(params,
+					new TypeRef<McpSchema.SubscribeRequest>() {
+					});
+			String uri = subscribeRequest.uri();
+			String sessionId = exchange.sessionId();
+			this.resourceSubscriptions.computeIfAbsent(uri, k -> Collections.newSetFromMap(new ConcurrentHashMap<>()))
+				.add(sessionId);
+			logger.debug("Session {} subscribed to resource URI: {}", sessionId, uri);
+
+			return Mono.just(Map.of());
+		});
+	}
+
+	private McpRequestHandler<Object> resourcesUnsubscribeRequestHandler() {
+		return (exchange, params) -> Mono.defer(() -> {
+			McpSchema.UnsubscribeRequest unsubscribeRequest = jsonMapper.convertValue(params,
+					new TypeRef<McpSchema.UnsubscribeRequest>() {
+					});
+			String uri = unsubscribeRequest.uri();
+			String sessionId = exchange.sessionId();
+			Set<String> sessions = this.resourceSubscriptions.get(uri);
+			if (sessions != null) {
+				sessions.remove(sessionId);
+				if (sessions.isEmpty()) {
+					this.resourceSubscriptions.remove(uri, sessions);
+				}
+			}
+			logger.debug("Session {} unsubscribed from resource URI: {}", sessionId, uri);
+			return Mono.just(Map.of());
+		});
 	}
 
 	private McpRequestHandler<McpSchema.ListResourcesResult> resourcesListRequestHandler() {
@@ -878,10 +946,6 @@ private McpRequestHandler<Object> setLoggerRequestHandler() {
 
 				exchange.setMinLoggingLevel(newMinLoggingLevel.level());
 
-				// FIXME: this field is deprecated and should be removed together
-				// with the broadcasting loggingNotification.
-				this.minLoggingLevel = newMinLoggingLevel.level();
-
 				return Mono.just(Map.of());
 			});
 		};

mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java

@@ -228,6 +228,24 @@ public Mono<Void> notifyClients(String method, Object params) {
 			.then();
 	}
 
+	@Override
+	public Mono<Void> notifyClient(String sessionId, String method, Object params) {
+		return Mono.defer(() -> {
+			// Need to iterate in O(n) because the transport session id
+			// is different from the server-logical session id
+			McpServerSession session = sessions.values()
+				.stream()
+				.filter(s -> sessionId.equals(s.getId()))
+				.findFirst()
+				.orElse(null);
+			if (session == null) {
+				logger.debug("Session {} not found", sessionId);
+				return Mono.empty();
+			}
+			return session.sendNotification(method, params);
+		});
+	}
+
 	/**
 	 * Handles GET requests to establish SSE connections.
 	 * <p>

mcp-core/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java

@@ -206,6 +206,18 @@ public Mono<Void> notifyClients(String method, Object params) {
 		});
 	}
 
+	@Override
+	public Mono<Void> notifyClient(String sessionId, String method, Object params) {
+		return Mono.defer(() -> {
+			McpStreamableServerSession session = this.sessions.get(sessionId);
+			if (session == null) {
+				logger.debug("Session {} not found", sessionId);
+				return Mono.empty();
+			}
+			return session.sendNotification(method, params);
+		});
+	}
+
 	/**
 	 * Initiates a graceful shutdown of the transport.
 	 * @return A Mono that completes when all cleanup operations are finished

mcp-core/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java

@@ -98,12 +98,26 @@ public void setSessionFactory(McpServerSession.Factory sessionFactory) {
 	@Override
 	public Mono<Void> notifyClients(String method, Object params) {
 		if (this.session == null) {
-			return Mono.error(new IllegalStateException("No session to close"));
+			return Mono.error(new IllegalStateException("No session to notify"));
 		}
 		return this.session.sendNotification(method, params)
 			.doOnError(e -> logger.error("Failed to send notification: {}", e.getMessage()));
 	}
 
+	@Override
+	public Mono<Void> notifyClient(String sessionId, String method, Object params) {
+		return Mono.defer(() -> {
+			if (this.session == null) {
+				return Mono.error(new IllegalStateException("No session to notify"));
+			}
+			if (!this.session.getId().equals(sessionId)) {
+				return Mono.error(new IllegalStateException("Existing session id " + this.session.getId()
+						+ " doesn't match the notification target: " + sessionId));
+			}
+			return this.session.sendNotification(method, params);
+		});
+	}
+
 	@Override
 	public Mono<Void> closeGracefully() {
 		if (this.session == null) {

mcp-core/src/main/java/io/modelcontextprotocol/spec/DefaultMcpStreamableServerSessionFactory.java

@@ -10,6 +10,9 @@
 import java.time.Duration;
 import java.util.Map;
 import java.util.UUID;
+import java.util.function.Function;
+
+import reactor.core.publisher.Mono;
 
 /**
  * A default implementation of {@link McpStreamableServerSession.Factory}.
@@ -26,29 +29,53 @@ public class DefaultMcpStreamableServerSessionFactory implements McpStreamableSe
 
 	Map<String, McpNotificationHandler> notificationHandlers;
 
+	private final Function<String, Mono<Void>> onClose;
+
 	/**
-	 * Constructs an instance
+	 * Constructs an instance.
 	 * @param requestTimeout timeout for requests
 	 * @param initRequestHandler initialization request handler
 	 * @param requestHandlers map of MCP request handlers keyed by method name
 	 * @param notificationHandlers map of MCP notification handlers keyed by method name
+	 * @param onClose reactive callback invoked with the session ID when a session is
+	 * closed
 	 */
 	public DefaultMcpStreamableServerSessionFactory(Duration requestTimeout,
 			McpStreamableServerSession.InitRequestHandler initRequestHandler,
-			Map<String, McpRequestHandler<?>> requestHandlers,
-			Map<String, McpNotificationHandler> notificationHandlers) {
+			Map<String, McpRequestHandler<?>> requestHandlers, Map<String, McpNotificationHandler> notificationHandlers,
+			Function<String, Mono<Void>> onClose) {
 		this.requestTimeout = requestTimeout;
 		this.initRequestHandler = initRequestHandler;
 		this.requestHandlers = requestHandlers;
 		this.notificationHandlers = notificationHandlers;
+		this.onClose = onClose;
+	}
+
+	/**
+	 * Constructs an instance.
+	 * @param requestTimeout timeout for requests
+	 * @param initRequestHandler initialization request handler
+	 * @param requestHandlers map of MCP request handlers keyed by method name
+	 * @param notificationHandlers map of MCP notification handlers keyed by method name
+	 * @deprecated Use
+	 * {@link #DefaultMcpStreamableServerSessionFactory(Duration, McpStreamableServerSession.InitRequestHandler, Map, Map, Function)}
+	 * instead
+	 */
+	@Deprecated
+	public DefaultMcpStreamableServerSessionFactory(Duration requestTimeout,
+			McpStreamableServerSession.InitRequestHandler initRequestHandler,
+			Map<String, McpRequestHandler<?>> requestHandlers,
+			Map<String, McpNotificationHandler> notificationHandlers) {
+		this(requestTimeout, initRequestHandler, requestHandlers, notificationHandlers, sessionId -> Mono.empty());
 	}
 
 	@Override
 	public McpStreamableServerSession.McpStreamableServerSessionInit startSession(
 			McpSchema.InitializeRequest initializeRequest) {
-		return new McpStreamableServerSession.McpStreamableServerSessionInit(
-				new McpStreamableServerSession(UUID.randomUUID().toString(), initializeRequest.capabilities(),
-						initializeRequest.clientInfo(), requestTimeout, requestHandlers, notificationHandlers),
+		String sessionId = UUID.randomUUID().toString();
+		return new McpStreamableServerSession.McpStreamableServerSessionInit(new McpStreamableServerSession(sessionId,
+				initializeRequest.capabilities(), initializeRequest.clientInfo(), requestTimeout, requestHandlers,
+				notificationHandlers, () -> this.onClose.apply(sessionId)),
 				this.initRequestHandler.handle(initializeRequest));
 	}
 

mcp-core/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java

@@ -7,6 +7,7 @@
 import java.time.Duration;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Supplier;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
@@ -65,25 +66,47 @@ public class McpServerSession implements McpLoggableSession {
 
 	private volatile McpSchema.LoggingLevel minLoggingLevel = McpSchema.LoggingLevel.INFO;
 
+	private final Supplier<Mono<Void>> onClose;
+
 	/**
 	 * Creates a new server session with the given parameters and the transport to use.
 	 * @param id session id
+	 * @param requestTimeout duration to wait for request responses before timing out
 	 * @param transport the transport to use
 	 * @param initHandler called when a
 	 * {@link io.modelcontextprotocol.spec.McpSchema.InitializeRequest} is received by the
 	 * server
 	 * @param requestHandlers map of request handlers to use
 	 * @param notificationHandlers map of notification handlers to use
+	 * @param onClose supplier of a reactive callback invoked when the session is closed
 	 */
 	public McpServerSession(String id, Duration requestTimeout, McpServerTransport transport,
 			McpInitRequestHandler initHandler, Map<String, McpRequestHandler<?>> requestHandlers,
-			Map<String, McpNotificationHandler> notificationHandlers) {
+			Map<String, McpNotificationHandler> notificationHandlers, Supplier<Mono<Void>> onClose) {
 		this.id = id;
 		this.requestTimeout = requestTimeout;
 		this.transport = transport;
 		this.initRequestHandler = initHandler;
 		this.requestHandlers = requestHandlers;
 		this.notificationHandlers = notificationHandlers;
+		this.onClose = onClose;
+	}
+
+	/**
+	 * Creates a new server session with the given parameters and the transport to use.
+	 * @param id session id
+	 * @param requestTimeout duration to wait for request responses before timing out
+	 * @param transport the transport to use
+	 * @param initHandler called when a
+	 * {@link io.modelcontextprotocol.spec.McpSchema.InitializeRequest} is received by the
+	 * server
+	 * @param requestHandlers map of request handlers to use
+	 * @param notificationHandlers map of notification handlers to use
+	 */
+	public McpServerSession(String id, Duration requestTimeout, McpServerTransport transport,
+			McpInitRequestHandler initHandler, Map<String, McpRequestHandler<?>> requestHandlers,
+			Map<String, McpNotificationHandler> notificationHandlers) {
+		this(id, requestTimeout, transport, initHandler, requestHandlers, notificationHandlers, Mono::empty);
 	}
 
 	/**
@@ -318,12 +341,13 @@ private MethodNotFoundError getMethodNotFoundError(String method) {
 	@Override
 	public Mono<Void> closeGracefully() {
 		// TODO: clear pendingResponses and emit errors?
-		return this.transport.closeGracefully();
+		return this.onClose.get().onErrorComplete().then(this.transport.closeGracefully());
 	}
 
 	@Override
 	public void close() {
 		// TODO: clear pendingResponses and emit errors?
+		this.onClose.get().onErrorComplete().subscribe();
 		this.transport.close();
 	}
 

mcp-core/src/main/java/io/modelcontextprotocol/spec/McpServerTransportProviderBase.java

@@ -45,6 +45,23 @@ public interface McpServerTransportProviderBase {
 	 */
 	Mono<Void> notifyClients(String method, Object params);
 
+	/**
+	 * Sends a notification to a specific client session. Transport providers that support
+	 * resource subscriptions must override this method to enable per-session
+	 * notifications. The default implementation returns an error indicating that this
+	 * operation is not supported.
+	 * @param sessionId the id of the session to notify
+	 * @param method the name of the notification method to be called on the client
+	 * @param params parameters to be sent with the notification
+	 * @return a Mono that completes when the notification has been sent, or empty if the
+	 * session is not found
+	 */
+	default Mono<Void> notifyClient(String sessionId, String method, Object params) {
+		return Mono.error(
+				new UnsupportedOperationException("This transport provider does not support per-session notifications. "
+						+ "Override notifyClient() to enable resource subscription support."));
+	}
+
 	/**
 	 * Immediately closes all the transports with connected clients and releases any
 	 * associated resources.