Support resource subscriptions (#839)
Clients can now subscribe to specific resources and receive targeted notifications when those resources change. Previously, calling `notifyResourceUpdated` on the server would broadcast the notification to every connected client regardless of interest — now only sessions that have explicitly subscribed to a given resource URI receive the update, making resource change propagation both correct and efficient. The subscription lifecycle is fully handled: clients can subscribe and unsubscribe at any time, and the server cleans up subscription state when a session closes. Supersedes #838. Resolves #837, #776. --------- Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
This commit is contained in:
parent
ce2c747c92
commit
7f68ea5bc0
8 changed files with 506 additions and 60 deletions
|
|
@ -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.CallToolResult;
|
|||
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 @@ public class McpAsyncServer {
|
|||
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 class McpAsyncServer {
|
|||
}
|
||||
|
||||
/**
|
||||
* 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 @@ public class McpAsyncServer {
|
|||
|
||||
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());
|
||||
});
|
||||
};
|
||||
|
|
|
|||
|
|
@ -228,6 +228,25 @@ public class HttpServletSseServerTransportProvider extends HttpServlet implement
|
|||
.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 (in streamable http this
|
||||
// design issue was solved)
|
||||
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>
|
||||
|
|
|
|||
|
|
@ -206,6 +206,18 @@ public class HttpServletStreamableServerTransportProvider extends HttpServlet
|
|||
});
|
||||
}
|
||||
|
||||
@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
|
||||
|
|
|
|||
|
|
@ -98,12 +98,26 @@ public class StdioServerTransportProvider implements McpServerTransportProvider
|
|||
@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) {
|
||||
|
|
|
|||
|
|
@ -10,6 +10,9 @@ import io.modelcontextprotocol.server.McpRequestHandler;
|
|||
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));
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ package io.modelcontextprotocol.spec;
|
|||
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,9 +66,36 @@ 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, 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
|
||||
|
|
@ -78,12 +106,7 @@ public class McpServerSession implements McpLoggableSession {
|
|||
public McpServerSession(String id, Duration requestTimeout, McpServerTransport transport,
|
||||
McpInitRequestHandler initHandler, Map<String, McpRequestHandler<?>> requestHandlers,
|
||||
Map<String, McpNotificationHandler> notificationHandlers) {
|
||||
this.id = id;
|
||||
this.requestTimeout = requestTimeout;
|
||||
this.transport = transport;
|
||||
this.initRequestHandler = initHandler;
|
||||
this.requestHandlers = requestHandlers;
|
||||
this.notificationHandlers = notificationHandlers;
|
||||
this(id, requestTimeout, transport, initHandler, requestHandlers, notificationHandlers, Mono::empty);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -318,12 +341,13 @@ public class McpServerSession implements McpLoggableSession {
|
|||
@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();
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -62,6 +62,34 @@ public class McpStreamableServerSession implements McpLoggableSession {
|
|||
|
||||
private volatile McpSchema.LoggingLevel minLoggingLevel = McpSchema.LoggingLevel.INFO;
|
||||
|
||||
private final Supplier<Mono<Void>> onClose;
|
||||
|
||||
/**
|
||||
* Create an instance of the streamable session.
|
||||
* @param id session ID
|
||||
* @param clientCapabilities client capabilities
|
||||
* @param clientInfo client info
|
||||
* @param requestTimeout timeout to use for requests
|
||||
* @param requestHandlers the map of MCP request handlers keyed by method name
|
||||
* @param notificationHandlers the map of MCP notification handlers keyed by method
|
||||
* name
|
||||
* @param onClose supplier of a reactive callback invoked when the session is closed
|
||||
*/
|
||||
public McpStreamableServerSession(String id, McpSchema.ClientCapabilities clientCapabilities,
|
||||
McpSchema.Implementation clientInfo, Duration requestTimeout,
|
||||
Map<String, McpRequestHandler<?>> requestHandlers, Map<String, McpNotificationHandler> notificationHandlers,
|
||||
Supplier<Mono<Void>> onClose) {
|
||||
this.id = id;
|
||||
this.missingMcpTransportSession = new MissingMcpTransportSession(id);
|
||||
this.listeningStreamRef = new AtomicReference<>(this.missingMcpTransportSession);
|
||||
this.clientCapabilities.lazySet(clientCapabilities);
|
||||
this.clientInfo.lazySet(clientInfo);
|
||||
this.requestTimeout = requestTimeout;
|
||||
this.requestHandlers = requestHandlers;
|
||||
this.notificationHandlers = notificationHandlers;
|
||||
this.onClose = onClose;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create an instance of the streamable session.
|
||||
* @param id session ID
|
||||
|
|
@ -76,14 +104,7 @@ public class McpStreamableServerSession implements McpLoggableSession {
|
|||
McpSchema.Implementation clientInfo, Duration requestTimeout,
|
||||
Map<String, McpRequestHandler<?>> requestHandlers,
|
||||
Map<String, McpNotificationHandler> notificationHandlers) {
|
||||
this.id = id;
|
||||
this.missingMcpTransportSession = new MissingMcpTransportSession(id);
|
||||
this.listeningStreamRef = new AtomicReference<>(this.missingMcpTransportSession);
|
||||
this.clientCapabilities.lazySet(clientCapabilities);
|
||||
this.clientInfo.lazySet(clientInfo);
|
||||
this.requestTimeout = requestTimeout;
|
||||
this.requestHandlers = requestHandlers;
|
||||
this.notificationHandlers = notificationHandlers;
|
||||
this(id, clientCapabilities, clientInfo, requestTimeout, requestHandlers, notificationHandlers, Mono::empty);
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -126,6 +147,7 @@ public class McpStreamableServerSession implements McpLoggableSession {
|
|||
}
|
||||
|
||||
public Mono<Void> delete() {
|
||||
// onClose is invoked inside closeGracefully
|
||||
return this.closeGracefully().then(Mono.fromRunnable(() -> {
|
||||
// TODO: review in the context of history storage
|
||||
// delete history, etc.
|
||||
|
|
@ -258,15 +280,16 @@ public class McpStreamableServerSession implements McpLoggableSession {
|
|||
|
||||
@Override
|
||||
public Mono<Void> closeGracefully() {
|
||||
return Mono.defer(() -> {
|
||||
return this.onClose.get().onErrorComplete().then(Mono.defer(() -> {
|
||||
McpLoggableSession listeningStream = this.listeningStreamRef.getAndSet(missingMcpTransportSession);
|
||||
return listeningStream.closeGracefully();
|
||||
// TODO: Also close all the open streams
|
||||
});
|
||||
}));
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
this.onClose.get().onErrorComplete().subscribe();
|
||||
McpLoggableSession listeningStream = this.listeningStreamRef.getAndSet(missingMcpTransportSession);
|
||||
if (listeningStream != null) {
|
||||
listeningStream.close();
|
||||
|
|
|
|||
Loading…
Reference in a new issue