refactor: extract MCP client initialization logic into LifecycleInitializer (#370)
- Create new LifecycleInitializer class to handle protocol initialization phase - Move initialization logic from McpAsyncClient to dedicated initializer - Add javadocs for MCP initialization process - Implement protocol version negotiation and capability exchange - Add exception handling for transport session recovery - Include test suite for LifecyleInitializer - Simplify McpAsyncClient by delegating initialization responsibilities This refactoring improves separation of concerns and makes the initialization process more maintainable and testable. Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
This commit is contained in:
Christian Tzolov
parent
9977a064d5
commit
b37c32cdba
2 changed files with 795 additions and 187 deletions
|
|
@ -0,0 +1,348 @@
|
|||
package io.modelcontextprotocol.client;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
import io.modelcontextprotocol.spec.McpClientSession;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import io.modelcontextprotocol.util.Assert;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.Sinks;
|
||||
import reactor.util.context.ContextView;
|
||||
|
||||
/**
|
||||
* <b>Handles the protocol initialization phase between client and server</b>
|
||||
*
|
||||
* <p>
|
||||
* The initialization phase MUST be the first interaction between client and server.
|
||||
* During this phase, the client and server perform the following operations:
|
||||
* <ul>
|
||||
* <li>Establish protocol version compatibility</li>
|
||||
* <li>Exchange and negotiate capabilities</li>
|
||||
* <li>Share implementation details</li>
|
||||
* </ul>
|
||||
*
|
||||
* <b>Client Initialization Process</b>
|
||||
* <p>
|
||||
* The client MUST initiate this phase by sending an initialize request containing:
|
||||
* <ul>
|
||||
* <li>Protocol version supported</li>
|
||||
* <li>Client capabilities</li>
|
||||
* <li>Client implementation information</li>
|
||||
* </ul>
|
||||
*
|
||||
* <p>
|
||||
* After successful initialization, the client MUST send an initialized notification to
|
||||
* indicate it is ready to begin normal operations.
|
||||
*
|
||||
* <b>Server Response</b>
|
||||
* <p>
|
||||
* The server MUST respond with its own capabilities and information.
|
||||
*
|
||||
* <b>Protocol Version Negotiation</b>
|
||||
* <p>
|
||||
* In the initialize request, the client MUST send a protocol version it supports. This
|
||||
* SHOULD be the latest version supported by the client.
|
||||
*
|
||||
* <p>
|
||||
* If the server supports the requested protocol version, it MUST respond with the same
|
||||
* version. Otherwise, the server MUST respond with another protocol version it supports.
|
||||
* This SHOULD be the latest version supported by the server.
|
||||
*
|
||||
* <p>
|
||||
* If the client does not support the version in the server's response, it SHOULD
|
||||
* disconnect.
|
||||
*
|
||||
* <b>Request Restrictions</b>
|
||||
* <p>
|
||||
* <strong>Important:</strong> The following restrictions apply during initialization:
|
||||
* <ul>
|
||||
* <li>The client SHOULD NOT send requests other than pings before the server has
|
||||
* responded to the initialize request</li>
|
||||
* <li>The server SHOULD NOT send requests other than pings and logging before receiving
|
||||
* the initialized notification</li>
|
||||
* </ul>
|
||||
*/
|
||||
class LifecycleInitializer {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(LifecycleInitializer.class);
|
||||
|
||||
/**
|
||||
* The MCP session supplier that manages bidirectional JSON-RPC communication between
|
||||
* clients and servers.
|
||||
*/
|
||||
private final Function<ContextView, McpClientSession> sessionSupplier;
|
||||
|
||||
private final McpSchema.ClientCapabilities clientCapabilities;
|
||||
|
||||
private final McpSchema.Implementation clientInfo;
|
||||
|
||||
private List<String> protocolVersions;
|
||||
|
||||
private final AtomicReference<DefaultInitialization> initializationRef = new AtomicReference<>();
|
||||
|
||||
/**
|
||||
* The max timeout to await for the client-server connection to be initialized.
|
||||
*/
|
||||
private final Duration initializationTimeout;
|
||||
|
||||
public LifecycleInitializer(McpSchema.ClientCapabilities clientCapabilities, McpSchema.Implementation clientInfo,
|
||||
List<String> protocolVersions, Duration initializationTimeout,
|
||||
Function<ContextView, McpClientSession> sessionSupplier) {
|
||||
|
||||
Assert.notNull(sessionSupplier, "Session supplier must not be null");
|
||||
Assert.notNull(clientCapabilities, "Client capabilities must not be null");
|
||||
Assert.notNull(clientInfo, "Client info must not be null");
|
||||
Assert.notEmpty(protocolVersions, "Protocol versions must not be empty");
|
||||
Assert.notNull(initializationTimeout, "Initialization timeout must not be null");
|
||||
|
||||
this.sessionSupplier = sessionSupplier;
|
||||
this.clientCapabilities = clientCapabilities;
|
||||
this.clientInfo = clientInfo;
|
||||
this.protocolVersions = Collections.unmodifiableList(new ArrayList<>(protocolVersions));
|
||||
this.initializationTimeout = initializationTimeout;
|
||||
}
|
||||
|
||||
/**
|
||||
* This method is package-private and used for test only. Should not be called by user
|
||||
* code.
|
||||
* @param protocolVersions the Client supported protocol versions.
|
||||
*/
|
||||
void setProtocolVersions(List<String> protocolVersions) {
|
||||
this.protocolVersions = protocolVersions;
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents the initialization state of the MCP client.
|
||||
*/
|
||||
interface Initialization {
|
||||
|
||||
/**
|
||||
* Returns the MCP client session that is used to communicate with the server.
|
||||
* This session is established during the initialization process and is used for
|
||||
* sending requests and notifications.
|
||||
* @return The MCP client session
|
||||
*/
|
||||
McpClientSession mcpSession();
|
||||
|
||||
/**
|
||||
* Returns the result of the MCP initialization process. This result contains
|
||||
* information about the protocol version, capabilities, server info, and
|
||||
* instructions provided by the server during the initialization phase.
|
||||
* @return The result of the MCP initialization process
|
||||
*/
|
||||
McpSchema.InitializeResult initializeResult();
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Default implementation of the {@link Initialization} interface that manages the MCP
|
||||
* client initialization process.
|
||||
*/
|
||||
private static class DefaultInitialization implements Initialization {
|
||||
|
||||
/**
|
||||
* A sink that emits the result of the MCP initialization process. It allows
|
||||
* subscribers to wait for the initialization to complete.
|
||||
*/
|
||||
private final Sinks.One<McpSchema.InitializeResult> initSink;
|
||||
|
||||
/**
|
||||
* Holds the result of the MCP initialization process. It is used to cache the
|
||||
* result for future requests.
|
||||
*/
|
||||
private final AtomicReference<McpSchema.InitializeResult> result;
|
||||
|
||||
/**
|
||||
* Holds the MCP client session that is used to communicate with the server. It is
|
||||
* set during the initialization process and used for sending requests and
|
||||
* notifications.
|
||||
*/
|
||||
private final AtomicReference<McpClientSession> mcpClientSession;
|
||||
|
||||
private DefaultInitialization() {
|
||||
this.initSink = Sinks.one();
|
||||
this.result = new AtomicReference<>();
|
||||
this.mcpClientSession = new AtomicReference<>();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------
|
||||
// Public access for mcpSession and initializeResult because they are
|
||||
// used in by the McpAsyncClient.
|
||||
// ----------------------------------------------------
|
||||
public McpClientSession mcpSession() {
|
||||
return this.mcpClientSession.get();
|
||||
}
|
||||
|
||||
public McpSchema.InitializeResult initializeResult() {
|
||||
return this.result.get();
|
||||
}
|
||||
|
||||
// ---------------------------------------------------
|
||||
// Private accessors used internally by the LifecycleInitializer to set the MCP
|
||||
// client session and complete the initialization process.
|
||||
// ---------------------------------------------------
|
||||
private void setMcpClientSession(McpClientSession mcpClientSession) {
|
||||
this.mcpClientSession.set(mcpClientSession);
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns a Mono that completes when the MCP client initialization is complete.
|
||||
* This allows subscribers to wait for the initialization to finish before
|
||||
* proceeding with further operations.
|
||||
* @return A Mono that emits the result of the MCP initialization process
|
||||
*/
|
||||
private Mono<McpSchema.InitializeResult> await() {
|
||||
return this.initSink.asMono();
|
||||
}
|
||||
|
||||
/**
|
||||
* Completes the initialization process with the given result. It caches the
|
||||
* result and emits it to all subscribers waiting for the initialization to
|
||||
* complete.
|
||||
* @param initializeResult The result of the MCP initialization process
|
||||
*/
|
||||
private void complete(McpSchema.InitializeResult initializeResult) {
|
||||
// first ensure the result is cached
|
||||
this.result.set(initializeResult);
|
||||
// inform all the subscribers waiting for the initialization
|
||||
this.initSink.emitValue(initializeResult, Sinks.EmitFailureHandler.FAIL_FAST);
|
||||
}
|
||||
|
||||
private void error(Throwable t) {
|
||||
this.initSink.emitError(t, Sinks.EmitFailureHandler.FAIL_FAST);
|
||||
}
|
||||
|
||||
private void close() {
|
||||
this.mcpSession().close();
|
||||
}
|
||||
|
||||
private Mono<Void> closeGracefully() {
|
||||
return this.mcpSession().closeGracefully();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
public boolean isInitialized() {
|
||||
return this.currentInitializationResult() != null;
|
||||
}
|
||||
|
||||
public McpSchema.InitializeResult currentInitializationResult() {
|
||||
DefaultInitialization current = this.initializationRef.get();
|
||||
McpSchema.InitializeResult initializeResult = current != null ? current.result.get() : null;
|
||||
return initializeResult;
|
||||
}
|
||||
|
||||
/**
|
||||
* Hook to handle exceptions that occur during the MCP transport session.
|
||||
* <p>
|
||||
* If the exception is a {@link McpTransportSessionNotFoundException}, it indicates
|
||||
* that the session was not found, and we should re-initialize the client.
|
||||
* </p>
|
||||
* @param t The exception to handle
|
||||
*/
|
||||
public void handleException(Throwable t) {
|
||||
logger.warn("Handling exception", t);
|
||||
if (t instanceof McpTransportSessionNotFoundException) {
|
||||
DefaultInitialization previous = this.initializationRef.getAndSet(null);
|
||||
if (previous != null) {
|
||||
previous.close();
|
||||
}
|
||||
// Providing an empty operation since we are only interested in triggering
|
||||
// the implicit initialization step.
|
||||
withIntitialization("re-initializing", result -> Mono.empty()).subscribe();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Utility method to ensure the initialization is established before executing an
|
||||
* operation.
|
||||
* @param <T> The type of the result Mono
|
||||
* @param actionName The action to perform when the client is initialized
|
||||
* @param operation The operation to execute when the client is initialized
|
||||
* @return A Mono that completes with the result of the operation
|
||||
*/
|
||||
public <T> Mono<T> withIntitialization(String actionName, Function<Initialization, Mono<T>> operation) {
|
||||
return Mono.deferContextual(ctx -> {
|
||||
DefaultInitialization newInit = new DefaultInitialization();
|
||||
DefaultInitialization previous = this.initializationRef.compareAndExchange(null, newInit);
|
||||
|
||||
boolean needsToInitialize = previous == null;
|
||||
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
|
||||
|
||||
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize ? doInitialize(newInit, ctx)
|
||||
: previous.await();
|
||||
|
||||
return initializationJob.map(initializeResult -> this.initializationRef.get())
|
||||
.timeout(this.initializationTimeout)
|
||||
.onErrorResume(ex -> {
|
||||
logger.warn("Failed to initialize", ex);
|
||||
return Mono.error(new McpError("Client failed to initialize " + actionName));
|
||||
})
|
||||
.flatMap(operation);
|
||||
});
|
||||
}
|
||||
|
||||
private Mono<McpSchema.InitializeResult> doInitialize(DefaultInitialization initialization, ContextView ctx) {
|
||||
initialization.setMcpClientSession(this.sessionSupplier.apply(ctx));
|
||||
|
||||
McpClientSession mcpClientSession = initialization.mcpSession();
|
||||
|
||||
String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);
|
||||
|
||||
McpSchema.InitializeRequest initializeRequest = new McpSchema.InitializeRequest(latestVersion,
|
||||
this.clientCapabilities, this.clientInfo);
|
||||
|
||||
Mono<McpSchema.InitializeResult> result = mcpClientSession.sendRequest(McpSchema.METHOD_INITIALIZE,
|
||||
initializeRequest, McpAsyncClient.INITIALIZE_RESULT_TYPE_REF);
|
||||
|
||||
return result.flatMap(initializeResult -> {
|
||||
logger.info("Server response with Protocol: {}, Capabilities: {}, Info: {} and Instructions {}",
|
||||
initializeResult.protocolVersion(), initializeResult.capabilities(), initializeResult.serverInfo(),
|
||||
initializeResult.instructions());
|
||||
|
||||
if (!this.protocolVersions.contains(initializeResult.protocolVersion())) {
|
||||
return Mono.error(new McpError(
|
||||
"Unsupported protocol version from the server: " + initializeResult.protocolVersion()));
|
||||
}
|
||||
|
||||
return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null)
|
||||
.thenReturn(initializeResult);
|
||||
}).doOnNext(initialization::complete).onErrorResume(ex -> {
|
||||
initialization.error(ex);
|
||||
return Mono.error(ex);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Closes the current initialization if it exists.
|
||||
*/
|
||||
public void close() {
|
||||
DefaultInitialization current = this.initializationRef.getAndSet(null);
|
||||
if (current != null) {
|
||||
current.close();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Gracefully closes the current initialization if it exists.
|
||||
* @return A Mono that completes when the connection is closed
|
||||
*/
|
||||
public Mono<?> closeGracefully() {
|
||||
return Mono.defer(() -> {
|
||||
DefaultInitialization current = this.initializationRef.getAndSet(null);
|
||||
Mono<?> sessionClose = current != null ? current.closeGracefully() : Mono.empty();
|
||||
return sessionClose;
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -10,7 +10,6 @@ import java.util.HashMap;
|
|||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
|
|
@ -19,8 +18,6 @@ import org.slf4j.LoggerFactory;
|
|||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
|
||||
import io.modelcontextprotocol.spec.McpClientSession;
|
||||
import io.modelcontextprotocol.spec.McpClientSession.NotificationHandler;
|
||||
import io.modelcontextprotocol.spec.McpClientSession.RequestHandler;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
|
|
@ -36,13 +33,12 @@ import io.modelcontextprotocol.spec.McpSchema.LoggingLevel;
|
|||
import io.modelcontextprotocol.spec.McpSchema.LoggingMessageNotification;
|
||||
import io.modelcontextprotocol.spec.McpSchema.PaginatedRequest;
|
||||
import io.modelcontextprotocol.spec.McpSchema.Root;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import io.modelcontextprotocol.spec.McpClientSession.NotificationHandler;
|
||||
import io.modelcontextprotocol.spec.McpClientSession.RequestHandler;
|
||||
import io.modelcontextprotocol.util.Assert;
|
||||
import io.modelcontextprotocol.util.Utils;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.Sinks;
|
||||
import reactor.util.context.ContextView;
|
||||
|
||||
/**
|
||||
* The Model Context Protocol (MCP) client implementation that provides asynchronous
|
||||
|
|
@ -104,13 +100,6 @@ public class McpAsyncClient {
|
|||
public static final TypeReference<LoggingMessageNotification> LOGGING_MESSAGE_NOTIFICATION_TYPE_REF = new TypeReference<>() {
|
||||
};
|
||||
|
||||
private final AtomicReference<Initialization> initializationRef = new AtomicReference<>();
|
||||
|
||||
/**
|
||||
* The max timeout to await for the client-server connection to be initialized.
|
||||
*/
|
||||
private final Duration initializationTimeout;
|
||||
|
||||
/**
|
||||
* Client capabilities.
|
||||
*/
|
||||
|
|
@ -154,15 +143,9 @@ public class McpAsyncClient {
|
|||
private final McpClientTransport transport;
|
||||
|
||||
/**
|
||||
* Supported protocol versions.
|
||||
* The lifecycle initializer that manages the client-server connection initialization.
|
||||
*/
|
||||
private List<String> protocolVersions = List.of(McpSchema.LATEST_PROTOCOL_VERSION);
|
||||
|
||||
/**
|
||||
* The MCP session supplier that manages bidirectional JSON-RPC communication between
|
||||
* clients and servers.
|
||||
*/
|
||||
private final Function<ContextView, McpClientSession> sessionSupplier;
|
||||
private final LifecycleInitializer initializer;
|
||||
|
||||
/**
|
||||
* Create a new McpAsyncClient with the given transport and session request-response
|
||||
|
|
@ -183,7 +166,6 @@ public class McpAsyncClient {
|
|||
this.clientCapabilities = features.clientCapabilities();
|
||||
this.transport = transport;
|
||||
this.roots = new ConcurrentHashMap<>(features.roots());
|
||||
this.initializationTimeout = initializationTimeout;
|
||||
|
||||
// Request Handlers
|
||||
Map<String, RequestHandler<?>> requestHandlers = new HashMap<>();
|
||||
|
|
@ -271,28 +253,11 @@ public class McpAsyncClient {
|
|||
notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_MESSAGE,
|
||||
asyncLoggingNotificationHandler(loggingConsumersFinal));
|
||||
|
||||
this.transport.setExceptionHandler(this::handleException);
|
||||
this.sessionSupplier = ctx -> new McpClientSession(requestTimeout, transport, requestHandlers,
|
||||
notificationHandlers, con -> con.contextWrite(ctx));
|
||||
}
|
||||
|
||||
private void handleException(Throwable t) {
|
||||
logger.warn("Handling exception", t);
|
||||
if (t instanceof McpTransportSessionNotFoundException) {
|
||||
Initialization previous = this.initializationRef.getAndSet(null);
|
||||
if (previous != null) {
|
||||
previous.close();
|
||||
}
|
||||
// Providing an empty operation since we are only interested in triggering the
|
||||
// implicit initialization step.
|
||||
withSession("re-initializing", result -> Mono.empty()).subscribe();
|
||||
}
|
||||
}
|
||||
|
||||
private McpSchema.InitializeResult currentInitializationResult() {
|
||||
Initialization current = this.initializationRef.get();
|
||||
McpSchema.InitializeResult initializeResult = current != null ? current.result.get() : null;
|
||||
return initializeResult;
|
||||
this.initializer = new LifecycleInitializer(clientCapabilities, clientInfo,
|
||||
List.of(McpSchema.LATEST_PROTOCOL_VERSION), initializationTimeout,
|
||||
ctx -> new McpClientSession(requestTimeout, transport, requestHandlers, notificationHandlers,
|
||||
con -> con.contextWrite(ctx)));
|
||||
this.transport.setExceptionHandler(this.initializer::handleException);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -300,7 +265,7 @@ public class McpAsyncClient {
|
|||
* @return The server capabilities
|
||||
*/
|
||||
public McpSchema.ServerCapabilities getServerCapabilities() {
|
||||
McpSchema.InitializeResult initializeResult = currentInitializationResult();
|
||||
McpSchema.InitializeResult initializeResult = this.initializer.currentInitializationResult();
|
||||
return initializeResult != null ? initializeResult.capabilities() : null;
|
||||
}
|
||||
|
||||
|
|
@ -310,7 +275,7 @@ public class McpAsyncClient {
|
|||
* @return The server instructions
|
||||
*/
|
||||
public String getServerInstructions() {
|
||||
McpSchema.InitializeResult initializeResult = currentInitializationResult();
|
||||
McpSchema.InitializeResult initializeResult = this.initializer.currentInitializationResult();
|
||||
return initializeResult != null ? initializeResult.instructions() : null;
|
||||
}
|
||||
|
||||
|
|
@ -319,7 +284,7 @@ public class McpAsyncClient {
|
|||
* @return The server implementation details
|
||||
*/
|
||||
public McpSchema.Implementation getServerInfo() {
|
||||
McpSchema.InitializeResult initializeResult = currentInitializationResult();
|
||||
McpSchema.InitializeResult initializeResult = this.initializer.currentInitializationResult();
|
||||
return initializeResult != null ? initializeResult.serverInfo() : null;
|
||||
}
|
||||
|
||||
|
|
@ -328,8 +293,7 @@ public class McpAsyncClient {
|
|||
* @return true if the client-server connection is initialized
|
||||
*/
|
||||
public boolean isInitialized() {
|
||||
Initialization current = this.initializationRef.get();
|
||||
return current != null && (current.result.get() != null);
|
||||
return this.initializer.isInitialized();
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -352,10 +316,7 @@ public class McpAsyncClient {
|
|||
* Closes the client connection immediately.
|
||||
*/
|
||||
public void close() {
|
||||
Initialization current = this.initializationRef.getAndSet(null);
|
||||
if (current != null) {
|
||||
current.close();
|
||||
}
|
||||
this.initializer.close();
|
||||
this.transport.close();
|
||||
}
|
||||
|
||||
|
|
@ -365,9 +326,7 @@ public class McpAsyncClient {
|
|||
*/
|
||||
public Mono<Void> closeGracefully() {
|
||||
return Mono.defer(() -> {
|
||||
Initialization current = this.initializationRef.getAndSet(null);
|
||||
Mono<?> sessionClose = current != null ? current.closeGracefully() : Mono.empty();
|
||||
return sessionClose.then(transport.closeGracefully());
|
||||
return this.initializer.closeGracefully().then(transport.closeGracefully());
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -401,118 +360,7 @@ public class McpAsyncClient {
|
|||
* </p>
|
||||
*/
|
||||
public Mono<McpSchema.InitializeResult> initialize() {
|
||||
return withSession("by explicit API call", init -> Mono.just(init.get()));
|
||||
}
|
||||
|
||||
private Mono<McpSchema.InitializeResult> doInitialize(Initialization initialization, ContextView ctx) {
|
||||
initialization.setMcpClientSession(this.sessionSupplier.apply(ctx));
|
||||
|
||||
McpClientSession mcpClientSession = initialization.mcpSession();
|
||||
|
||||
String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);
|
||||
|
||||
McpSchema.InitializeRequest initializeRequest = new McpSchema.InitializeRequest(// @formatter:off
|
||||
latestVersion,
|
||||
this.clientCapabilities,
|
||||
this.clientInfo); // @formatter:on
|
||||
|
||||
Mono<McpSchema.InitializeResult> result = mcpClientSession.sendRequest(McpSchema.METHOD_INITIALIZE,
|
||||
initializeRequest, INITIALIZE_RESULT_TYPE_REF);
|
||||
|
||||
return result.flatMap(initializeResult -> {
|
||||
logger.info("Server response with Protocol: {}, Capabilities: {}, Info: {} and Instructions {}",
|
||||
initializeResult.protocolVersion(), initializeResult.capabilities(), initializeResult.serverInfo(),
|
||||
initializeResult.instructions());
|
||||
|
||||
if (!this.protocolVersions.contains(initializeResult.protocolVersion())) {
|
||||
return Mono.error(new McpError(
|
||||
"Unsupported protocol version from the server: " + initializeResult.protocolVersion()));
|
||||
}
|
||||
|
||||
return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null)
|
||||
.thenReturn(initializeResult);
|
||||
}).doOnNext(initialization::complete).onErrorResume(ex -> {
|
||||
initialization.error(ex);
|
||||
return Mono.error(ex);
|
||||
});
|
||||
}
|
||||
|
||||
private static class Initialization {
|
||||
|
||||
private final Sinks.One<McpSchema.InitializeResult> initSink = Sinks.one();
|
||||
|
||||
private final AtomicReference<McpSchema.InitializeResult> result = new AtomicReference<>();
|
||||
|
||||
private final AtomicReference<McpClientSession> mcpClientSession = new AtomicReference<>();
|
||||
|
||||
static Initialization create() {
|
||||
return new Initialization();
|
||||
}
|
||||
|
||||
void setMcpClientSession(McpClientSession mcpClientSession) {
|
||||
this.mcpClientSession.set(mcpClientSession);
|
||||
}
|
||||
|
||||
McpClientSession mcpSession() {
|
||||
return this.mcpClientSession.get();
|
||||
}
|
||||
|
||||
McpSchema.InitializeResult get() {
|
||||
return this.result.get();
|
||||
}
|
||||
|
||||
Mono<McpSchema.InitializeResult> await() {
|
||||
return this.initSink.asMono();
|
||||
}
|
||||
|
||||
void complete(McpSchema.InitializeResult initializeResult) {
|
||||
// first ensure the result is cached
|
||||
this.result.set(initializeResult);
|
||||
// inform all the subscribers waiting for the initialization
|
||||
this.initSink.emitValue(initializeResult, Sinks.EmitFailureHandler.FAIL_FAST);
|
||||
}
|
||||
|
||||
void error(Throwable t) {
|
||||
this.initSink.emitError(t, Sinks.EmitFailureHandler.FAIL_FAST);
|
||||
}
|
||||
|
||||
void close() {
|
||||
this.mcpSession().close();
|
||||
}
|
||||
|
||||
Mono<Void> closeGracefully() {
|
||||
return this.mcpSession().closeGracefully();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Utility method to handle the common pattern of ensuring initialization before
|
||||
* executing an operation.
|
||||
* @param <T> The type of the result Mono
|
||||
* @param actionName The action to perform when the client is initialized
|
||||
* @param operation The operation to execute when the client is initialized
|
||||
* @return A Mono that completes with the result of the operation
|
||||
*/
|
||||
private <T> Mono<T> withSession(String actionName, Function<Initialization, Mono<T>> operation) {
|
||||
return Mono.deferContextual(ctx -> {
|
||||
Initialization newInit = Initialization.create();
|
||||
Initialization previous = this.initializationRef.compareAndExchange(null, newInit);
|
||||
|
||||
boolean needsToInitialize = previous == null;
|
||||
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
|
||||
|
||||
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize ? doInitialize(newInit, ctx)
|
||||
: previous.await();
|
||||
|
||||
return initializationJob.map(initializeResult -> this.initializationRef.get())
|
||||
.timeout(this.initializationTimeout)
|
||||
.onErrorResume(ex -> {
|
||||
logger.warn("Failed to initialize", ex);
|
||||
return Mono.error(new McpError("Client failed to initialize " + actionName));
|
||||
})
|
||||
.flatMap(operation);
|
||||
});
|
||||
return this.initializer.withIntitialization("by explicit API call", init -> Mono.just(init.initializeResult()));
|
||||
}
|
||||
|
||||
// --------------------------
|
||||
|
|
@ -524,7 +372,7 @@ public class McpAsyncClient {
|
|||
* @return A Mono that completes with the server's ping response
|
||||
*/
|
||||
public Mono<Object> ping() {
|
||||
return this.withSession("pinging the server",
|
||||
return this.initializer.withIntitialization("pinging the server",
|
||||
init -> init.mcpSession().sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF));
|
||||
}
|
||||
|
||||
|
|
@ -605,7 +453,7 @@ public class McpAsyncClient {
|
|||
* @return A Mono that completes when the notification is sent.
|
||||
*/
|
||||
public Mono<Void> rootsListChangedNotification() {
|
||||
return this.withSession("sending roots list changed notification",
|
||||
return this.initializer.withIntitialization("sending roots list changed notification",
|
||||
init -> init.mcpSession().sendNotification(McpSchema.METHOD_NOTIFICATION_ROOTS_LIST_CHANGED));
|
||||
}
|
||||
|
||||
|
|
@ -664,8 +512,8 @@ public class McpAsyncClient {
|
|||
* @see #listTools()
|
||||
*/
|
||||
public Mono<McpSchema.CallToolResult> callTool(McpSchema.CallToolRequest callToolRequest) {
|
||||
return this.withSession("calling tools", init -> {
|
||||
if (init.get().capabilities().tools() == null) {
|
||||
return this.initializer.withIntitialization("calling tools", init -> {
|
||||
if (init.initializeResult().capabilities().tools() == null) {
|
||||
return Mono.error(new McpError("Server does not provide tools capability"));
|
||||
}
|
||||
return init.mcpSession()
|
||||
|
|
@ -693,8 +541,8 @@ public class McpAsyncClient {
|
|||
* @return A Mono that emits the list of tools result
|
||||
*/
|
||||
public Mono<McpSchema.ListToolsResult> listTools(String cursor) {
|
||||
return this.withSession("listing tools", init -> {
|
||||
if (init.get().capabilities().tools() == null) {
|
||||
return this.initializer.withIntitialization("listing tools", init -> {
|
||||
if (init.initializeResult().capabilities().tools() == null) {
|
||||
return Mono.error(new McpError("Server does not provide tools capability"));
|
||||
}
|
||||
return init.mcpSession()
|
||||
|
|
@ -757,8 +605,8 @@ public class McpAsyncClient {
|
|||
* @see #readResource(McpSchema.Resource)
|
||||
*/
|
||||
public Mono<McpSchema.ListResourcesResult> listResources(String cursor) {
|
||||
return this.withSession("listing resources", init -> {
|
||||
if (init.get().capabilities().resources() == null) {
|
||||
return this.initializer.withIntitialization("listing resources", init -> {
|
||||
if (init.initializeResult().capabilities().resources() == null) {
|
||||
return Mono.error(new McpError("Server does not provide the resources capability"));
|
||||
}
|
||||
return init.mcpSession()
|
||||
|
|
@ -789,8 +637,8 @@ public class McpAsyncClient {
|
|||
* @see McpSchema.ReadResourceResult
|
||||
*/
|
||||
public Mono<McpSchema.ReadResourceResult> readResource(McpSchema.ReadResourceRequest readResourceRequest) {
|
||||
return this.withSession("reading resources", init -> {
|
||||
if (init.get().capabilities().resources() == null) {
|
||||
return this.initializer.withIntitialization("reading resources", init -> {
|
||||
if (init.initializeResult().capabilities().resources() == null) {
|
||||
return Mono.error(new McpError("Server does not provide the resources capability"));
|
||||
}
|
||||
return init.mcpSession()
|
||||
|
|
@ -827,8 +675,8 @@ public class McpAsyncClient {
|
|||
* @see McpSchema.ListResourceTemplatesResult
|
||||
*/
|
||||
public Mono<McpSchema.ListResourceTemplatesResult> listResourceTemplates(String cursor) {
|
||||
return this.withSession("listing resource templates", init -> {
|
||||
if (init.get().capabilities().resources() == null) {
|
||||
return this.initializer.withIntitialization("listing resource templates", init -> {
|
||||
if (init.initializeResult().capabilities().resources() == null) {
|
||||
return Mono.error(new McpError("Server does not provide the resources capability"));
|
||||
}
|
||||
return init.mcpSession()
|
||||
|
|
@ -847,7 +695,7 @@ public class McpAsyncClient {
|
|||
* @see #unsubscribeResource(McpSchema.UnsubscribeRequest)
|
||||
*/
|
||||
public Mono<Void> subscribeResource(McpSchema.SubscribeRequest subscribeRequest) {
|
||||
return this.withSession("subscribing to resources", init -> init.mcpSession()
|
||||
return this.initializer.withIntitialization("subscribing to resources", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_RESOURCES_SUBSCRIBE, subscribeRequest, VOID_TYPE_REFERENCE));
|
||||
}
|
||||
|
||||
|
|
@ -861,7 +709,7 @@ public class McpAsyncClient {
|
|||
* @see #subscribeResource(McpSchema.SubscribeRequest)
|
||||
*/
|
||||
public Mono<Void> unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) {
|
||||
return this.withSession("unsubscribing from resources", init -> init.mcpSession()
|
||||
return this.initializer.withIntitialization("unsubscribing from resources", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_RESOURCES_UNSUBSCRIBE, unsubscribeRequest, VOID_TYPE_REFERENCE));
|
||||
}
|
||||
|
||||
|
|
@ -927,7 +775,7 @@ public class McpAsyncClient {
|
|||
* @see #getPrompt(GetPromptRequest)
|
||||
*/
|
||||
public Mono<ListPromptsResult> listPrompts(String cursor) {
|
||||
return this.withSession("listing prompts", init -> init.mcpSession()
|
||||
return this.initializer.withIntitialization("listing prompts", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_PROMPT_LIST, new PaginatedRequest(cursor), LIST_PROMPTS_RESULT_TYPE_REF));
|
||||
}
|
||||
|
||||
|
|
@ -941,7 +789,7 @@ public class McpAsyncClient {
|
|||
* @see #listPrompts()
|
||||
*/
|
||||
public Mono<GetPromptResult> getPrompt(GetPromptRequest getPromptRequest) {
|
||||
return this.withSession("getting prompts", init -> init.mcpSession()
|
||||
return this.initializer.withIntitialization("getting prompts", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_PROMPT_GET, getPromptRequest, GET_PROMPT_RESULT_TYPE_REF));
|
||||
}
|
||||
|
||||
|
|
@ -992,7 +840,7 @@ public class McpAsyncClient {
|
|||
return Mono.error(new McpError("Logging level must not be null"));
|
||||
}
|
||||
|
||||
return this.withSession("setting logging level", init -> {
|
||||
return this.initializer.withIntitialization("setting logging level", init -> {
|
||||
var params = new McpSchema.SetLevelRequest(loggingLevel);
|
||||
return init.mcpSession().sendRequest(McpSchema.METHOD_LOGGING_SET_LEVEL, params, OBJECT_TYPE_REF).then();
|
||||
});
|
||||
|
|
@ -1004,7 +852,7 @@ public class McpAsyncClient {
|
|||
* @param protocolVersions the Client supported protocol versions.
|
||||
*/
|
||||
void setProtocolVersions(List<String> protocolVersions) {
|
||||
this.protocolVersions = protocolVersions;
|
||||
this.initializer.setProtocolVersions(protocolVersions);
|
||||
}
|
||||
|
||||
// --------------------------
|
||||
|
|
@ -1024,7 +872,7 @@ public class McpAsyncClient {
|
|||
* @see McpSchema.CompleteResult
|
||||
*/
|
||||
public Mono<McpSchema.CompleteResult> completeCompletion(McpSchema.CompleteRequest completeRequest) {
|
||||
return this.withSession("complete completions", init -> init.mcpSession()
|
||||
return this.initializer.withIntitialization("complete completions", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_COMPLETION_COMPLETE, completeRequest, COMPLETION_COMPLETE_RESULT_TYPE_REF));
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue