feat(client): add client tool output schema validation and caching (#302)
- Add JSON schema validation for tool call results against output schemas - Implement automatic tool output schema caching during initialization - Add `enableCallToolSchemaCaching` configuration option to enable/disable schema caching - Add `JsonSchemaValidator` integration to McpClient builder APIs - Introduce post-initialization hook mechanism for performing operations after successful client initialization - Cache tool output schemas during `listTools` operations when caching is enabled - Validate tool results against cached schemas in `callTool` operations - Return error CallToolResult when validation fails - Add test coverage - Convert validateToolResult from Mono to synchronous method - Throw IllegalArgumentException on validation errors Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com> Co-authored-by: Anurag Pant <pantanurag555@github> Co-authored-by: Christian Tzolov <christian.tzolov@broadcom.com> Co-authored-by: Daniel Garnier-Moiroux <git@garnier.wf> Co-authored-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
This commit is contained in:
Anurag Pant
Christian Tzolov
parent
3f7578b3ce
commit
8e11fa82ae
10 changed files with 708 additions and 131 deletions
|
|
@ -11,14 +11,13 @@ 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 org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.Sinks;
|
||||
import reactor.util.context.ContextView;
|
||||
|
|
@ -99,21 +98,30 @@ class LifecycleInitializer {
|
|||
*/
|
||||
private final Duration initializationTimeout;
|
||||
|
||||
/**
|
||||
* Post-initialization hook to perform additional operations after every successful
|
||||
* initialization.
|
||||
*/
|
||||
private final Function<Initialization, Mono<Void>> postInitializationHook;
|
||||
|
||||
public LifecycleInitializer(McpSchema.ClientCapabilities clientCapabilities, McpSchema.Implementation clientInfo,
|
||||
List<String> protocolVersions, Duration initializationTimeout,
|
||||
Function<ContextView, McpClientSession> sessionSupplier) {
|
||||
Function<ContextView, McpClientSession> sessionSupplier,
|
||||
Function<Initialization, Mono<Void>> postInitializationHook) {
|
||||
|
||||
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");
|
||||
Assert.notNull(postInitializationHook, "Post-initialization hook must not be null");
|
||||
|
||||
this.sessionSupplier = sessionSupplier;
|
||||
this.clientCapabilities = clientCapabilities;
|
||||
this.clientInfo = clientInfo;
|
||||
this.protocolVersions = Collections.unmodifiableList(new ArrayList<>(protocolVersions));
|
||||
this.initializationTimeout = initializationTimeout;
|
||||
this.postInitializationHook = postInitializationHook;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -148,10 +156,6 @@ class LifecycleInitializer {
|
|||
|
||||
}
|
||||
|
||||
/**
|
||||
* Default implementation of the {@link Initialization} interface that manages the MCP
|
||||
* client initialization process.
|
||||
*/
|
||||
private static class DefaultInitialization implements Initialization {
|
||||
|
||||
/**
|
||||
|
|
@ -199,29 +203,20 @@ class LifecycleInitializer {
|
|||
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 cacheResult(McpSchema.InitializeResult initializeResult) {
|
||||
// first ensure the result is cached
|
||||
this.result.set(initializeResult);
|
||||
}
|
||||
|
||||
private void error(Throwable t) {
|
||||
this.initSink.emitError(t, Sinks.EmitFailureHandler.FAIL_FAST);
|
||||
}
|
||||
|
|
@ -263,7 +258,7 @@ class LifecycleInitializer {
|
|||
}
|
||||
// Providing an empty operation since we are only interested in triggering
|
||||
// the implicit initialization step.
|
||||
withIntitialization("re-initializing", result -> Mono.empty()).subscribe();
|
||||
this.withInitialization("re-initializing", result -> Mono.empty()).subscribe();
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -275,7 +270,7 @@ class LifecycleInitializer {
|
|||
* @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) {
|
||||
public <T> Mono<T> withInitialization(String actionName, Function<Initialization, Mono<T>> operation) {
|
||||
return Mono.deferContextual(ctx -> {
|
||||
DefaultInitialization newInit = new DefaultInitialization();
|
||||
DefaultInitialization previous = this.initializationRef.compareAndExchange(null, newInit);
|
||||
|
|
@ -283,8 +278,8 @@ class LifecycleInitializer {
|
|||
boolean needsToInitialize = previous == null;
|
||||
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
|
||||
|
||||
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize ? doInitialize(newInit, ctx)
|
||||
: previous.await();
|
||||
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize
|
||||
? this.doInitialize(newInit, this.postInitializationHook, ctx) : previous.await();
|
||||
|
||||
return initializationJob.map(initializeResult -> this.initializationRef.get())
|
||||
.timeout(this.initializationTimeout)
|
||||
|
|
@ -296,7 +291,9 @@ class LifecycleInitializer {
|
|||
});
|
||||
}
|
||||
|
||||
private Mono<McpSchema.InitializeResult> doInitialize(DefaultInitialization initialization, ContextView ctx) {
|
||||
private Mono<McpSchema.InitializeResult> doInitialize(DefaultInitialization initialization,
|
||||
Function<Initialization, Mono<Void>> postInitOperation, ContextView ctx) {
|
||||
|
||||
initialization.setMcpClientSession(this.sessionSupplier.apply(ctx));
|
||||
|
||||
McpClientSession mcpClientSession = initialization.mcpSession();
|
||||
|
|
@ -323,6 +320,9 @@ class LifecycleInitializer {
|
|||
|
||||
return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null)
|
||||
.thenReturn(initializeResult);
|
||||
}).flatMap(initializeResult -> {
|
||||
initialization.cacheResult(initializeResult);
|
||||
return postInitOperation.apply(initialization).thenReturn(initializeResult);
|
||||
}).doOnNext(initialization::complete).onErrorResume(ex -> {
|
||||
initialization.error(ex);
|
||||
return Mono.error(ex);
|
||||
|
|
|
|||
|
|
@ -15,12 +15,12 @@ import java.util.Map;
|
|||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.function.Function;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
import io.modelcontextprotocol.client.LifecycleInitializer.Initialization;
|
||||
import io.modelcontextprotocol.json.TypeRef;
|
||||
|
||||
import io.modelcontextprotocol.json.schema.JsonSchemaValidator;
|
||||
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.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
|
||||
|
|
@ -35,10 +35,10 @@ 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.McpClientSession.NotificationHandler;
|
||||
import io.modelcontextprotocol.spec.McpClientSession.RequestHandler;
|
||||
import io.modelcontextprotocol.util.Assert;
|
||||
import io.modelcontextprotocol.util.Utils;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
|
||||
|
|
@ -75,6 +75,7 @@ import reactor.core.publisher.Mono;
|
|||
* @author Dariusz Jędrzejczyk
|
||||
* @author Christian Tzolov
|
||||
* @author Jihoon Kim
|
||||
* @author Anurag Pant
|
||||
* @see McpClient
|
||||
* @see McpSchema
|
||||
* @see McpClientSession
|
||||
|
|
@ -152,16 +153,33 @@ public class McpAsyncClient {
|
|||
*/
|
||||
private final LifecycleInitializer initializer;
|
||||
|
||||
/**
|
||||
* JSON schema validator to use for validating tool responses against output schemas.
|
||||
*/
|
||||
private final JsonSchemaValidator jsonSchemaValidator;
|
||||
|
||||
/**
|
||||
* Cached tool output schemas.
|
||||
*/
|
||||
private final ConcurrentHashMap<String, Map<String, Object>> toolsOutputSchemaCache;
|
||||
|
||||
/**
|
||||
* Whether to enable automatic schema caching during callTool operations.
|
||||
*/
|
||||
private final boolean enableCallToolSchemaCaching;
|
||||
|
||||
/**
|
||||
* Create a new McpAsyncClient with the given transport and session request-response
|
||||
* timeout.
|
||||
* @param transport the transport to use.
|
||||
* @param requestTimeout the session request-response timeout.
|
||||
* @param initializationTimeout the max timeout to await for the client-server
|
||||
* @param features the MCP Client supported features.
|
||||
* @param jsonSchemaValidator the JSON schema validator to use for validating tool
|
||||
* @param features the MCP Client supported features. responses against output
|
||||
* schemas.
|
||||
*/
|
||||
McpAsyncClient(McpClientTransport transport, Duration requestTimeout, Duration initializationTimeout,
|
||||
McpClientFeatures.Async features) {
|
||||
JsonSchemaValidator jsonSchemaValidator, McpClientFeatures.Async features) {
|
||||
|
||||
Assert.notNull(transport, "Transport must not be null");
|
||||
Assert.notNull(requestTimeout, "Request timeout must not be null");
|
||||
|
|
@ -171,6 +189,9 @@ public class McpAsyncClient {
|
|||
this.clientCapabilities = features.clientCapabilities();
|
||||
this.transport = transport;
|
||||
this.roots = new ConcurrentHashMap<>(features.roots());
|
||||
this.jsonSchemaValidator = jsonSchemaValidator;
|
||||
this.toolsOutputSchemaCache = new ConcurrentHashMap<>();
|
||||
this.enableCallToolSchemaCaching = features.enableCallToolSchemaCaching();
|
||||
|
||||
// Request Handlers
|
||||
Map<String, RequestHandler<?>> requestHandlers = new HashMap<>();
|
||||
|
|
@ -273,9 +294,30 @@ public class McpAsyncClient {
|
|||
notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_PROGRESS,
|
||||
asyncProgressNotificationHandler(progressConsumersFinal));
|
||||
|
||||
Function<Initialization, Mono<Void>> postInitializationHook = init -> {
|
||||
|
||||
if (init.initializeResult().capabilities().tools() == null || !enableCallToolSchemaCaching) {
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
return this.listToolsInternal(init, McpSchema.FIRST_PAGE).doOnNext(listToolsResult -> {
|
||||
listToolsResult.tools()
|
||||
.forEach(tool -> logger.debug("Tool {} schema: {}", tool.name(), tool.outputSchema()));
|
||||
if (enableCallToolSchemaCaching && listToolsResult.tools() != null) {
|
||||
// Cache tools output schema
|
||||
listToolsResult.tools()
|
||||
.stream()
|
||||
.filter(tool -> tool.outputSchema() != null)
|
||||
.forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(), tool.outputSchema()));
|
||||
}
|
||||
}).then();
|
||||
};
|
||||
|
||||
this.initializer = new LifecycleInitializer(clientCapabilities, clientInfo, transport.protocolVersions(),
|
||||
initializationTimeout, ctx -> new McpClientSession(requestTimeout, transport, requestHandlers,
|
||||
notificationHandlers, con -> con.contextWrite(ctx)));
|
||||
notificationHandlers, con -> con.contextWrite(ctx)),
|
||||
postInitializationHook);
|
||||
|
||||
this.transport.setExceptionHandler(this.initializer::handleException);
|
||||
}
|
||||
|
||||
|
|
@ -387,7 +429,7 @@ public class McpAsyncClient {
|
|||
* </p>
|
||||
*/
|
||||
public Mono<McpSchema.InitializeResult> initialize() {
|
||||
return this.initializer.withIntitialization("by explicit API call", init -> Mono.just(init.initializeResult()));
|
||||
return this.initializer.withInitialization("by explicit API call", init -> Mono.just(init.initializeResult()));
|
||||
}
|
||||
|
||||
// --------------------------
|
||||
|
|
@ -399,7 +441,7 @@ public class McpAsyncClient {
|
|||
* @return A Mono that completes with the server's ping response
|
||||
*/
|
||||
public Mono<Object> ping() {
|
||||
return this.initializer.withIntitialization("pinging the server",
|
||||
return this.initializer.withInitialization("pinging the server",
|
||||
init -> init.mcpSession().sendRequest(McpSchema.METHOD_PING, null, OBJECT_TYPE_REF));
|
||||
}
|
||||
|
||||
|
|
@ -480,7 +522,7 @@ public class McpAsyncClient {
|
|||
* @return A Mono that completes when the notification is sent.
|
||||
*/
|
||||
public Mono<Void> rootsListChangedNotification() {
|
||||
return this.initializer.withIntitialization("sending roots list changed notification",
|
||||
return this.initializer.withInitialization("sending roots list changed notification",
|
||||
init -> init.mcpSession().sendNotification(McpSchema.METHOD_NOTIFICATION_ROOTS_LIST_CHANGED));
|
||||
}
|
||||
|
||||
|
|
@ -539,15 +581,45 @@ public class McpAsyncClient {
|
|||
* @see #listTools()
|
||||
*/
|
||||
public Mono<McpSchema.CallToolResult> callTool(McpSchema.CallToolRequest callToolRequest) {
|
||||
return this.initializer.withIntitialization("calling tools", init -> {
|
||||
return this.initializer.withInitialization("calling tool", init -> {
|
||||
if (init.initializeResult().capabilities().tools() == null) {
|
||||
return Mono.error(new IllegalStateException("Server does not provide tools capability"));
|
||||
}
|
||||
|
||||
return init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF);
|
||||
.sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF)
|
||||
.flatMap(result -> Mono.just(validateToolResult(callToolRequest.name(), result)));
|
||||
});
|
||||
}
|
||||
|
||||
private McpSchema.CallToolResult validateToolResult(String toolName, McpSchema.CallToolResult result) {
|
||||
|
||||
if (!this.enableCallToolSchemaCaching || result == null || result.isError() == Boolean.TRUE) {
|
||||
// if tool schema caching is disabled or tool call resulted in an error - skip
|
||||
// validation and return the result as it is
|
||||
return result;
|
||||
}
|
||||
|
||||
Map<String, Object> optOutputSchema = this.toolsOutputSchemaCache.get(toolName);
|
||||
|
||||
if (optOutputSchema == null) {
|
||||
logger.warn(
|
||||
"Calling a tool with no outputSchema is not expected to return result with structured content, but got: {}",
|
||||
result.structuredContent());
|
||||
return result;
|
||||
}
|
||||
|
||||
// Validate the tool output against the cached output schema
|
||||
var validation = this.jsonSchemaValidator.validate(optOutputSchema, result.structuredContent());
|
||||
|
||||
if (!validation.valid()) {
|
||||
logger.warn("Tool call result validation failed: {}", validation.errorMessage());
|
||||
throw new IllegalArgumentException("Tool call result validation failed: " + validation.errorMessage());
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieves the list of all tools provided by the server.
|
||||
* @return A Mono that emits the list of all tools result
|
||||
|
|
@ -568,14 +640,26 @@ public class McpAsyncClient {
|
|||
* @return A Mono that emits the list of tools result
|
||||
*/
|
||||
public Mono<McpSchema.ListToolsResult> listTools(String cursor) {
|
||||
return this.initializer.withIntitialization("listing tools", init -> {
|
||||
if (init.initializeResult().capabilities().tools() == null) {
|
||||
return Mono.error(new IllegalStateException("Server does not provide tools capability"));
|
||||
}
|
||||
return init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor),
|
||||
LIST_TOOLS_RESULT_TYPE_REF);
|
||||
});
|
||||
return this.initializer.withInitialization("listing tools", init -> this.listToolsInternal(init, cursor));
|
||||
}
|
||||
|
||||
private Mono<McpSchema.ListToolsResult> listToolsInternal(Initialization init, String cursor) {
|
||||
|
||||
if (init.initializeResult().capabilities().tools() == null) {
|
||||
return Mono.error(new IllegalStateException("Server does not provide tools capability"));
|
||||
}
|
||||
return init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor),
|
||||
LIST_TOOLS_RESULT_TYPE_REF)
|
||||
.doOnNext(result -> {
|
||||
if (this.enableCallToolSchemaCaching && result.tools() != null) {
|
||||
// Cache tools output schema
|
||||
result.tools()
|
||||
.stream()
|
||||
.filter(tool -> tool.outputSchema() != null)
|
||||
.forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(), tool.outputSchema()));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
private NotificationHandler asyncToolsChangeNotificationHandler(
|
||||
|
|
@ -632,7 +716,7 @@ public class McpAsyncClient {
|
|||
* @see #readResource(McpSchema.Resource)
|
||||
*/
|
||||
public Mono<McpSchema.ListResourcesResult> listResources(String cursor) {
|
||||
return this.initializer.withIntitialization("listing resources", init -> {
|
||||
return this.initializer.withInitialization("listing resources", init -> {
|
||||
if (init.initializeResult().capabilities().resources() == null) {
|
||||
return Mono.error(new IllegalStateException("Server does not provide the resources capability"));
|
||||
}
|
||||
|
|
@ -664,7 +748,7 @@ public class McpAsyncClient {
|
|||
* @see McpSchema.ReadResourceResult
|
||||
*/
|
||||
public Mono<McpSchema.ReadResourceResult> readResource(McpSchema.ReadResourceRequest readResourceRequest) {
|
||||
return this.initializer.withIntitialization("reading resources", init -> {
|
||||
return this.initializer.withInitialization("reading resources", init -> {
|
||||
if (init.initializeResult().capabilities().resources() == null) {
|
||||
return Mono.error(new IllegalStateException("Server does not provide the resources capability"));
|
||||
}
|
||||
|
|
@ -702,7 +786,7 @@ public class McpAsyncClient {
|
|||
* @see McpSchema.ListResourceTemplatesResult
|
||||
*/
|
||||
public Mono<McpSchema.ListResourceTemplatesResult> listResourceTemplates(String cursor) {
|
||||
return this.initializer.withIntitialization("listing resource templates", init -> {
|
||||
return this.initializer.withInitialization("listing resource templates", init -> {
|
||||
if (init.initializeResult().capabilities().resources() == null) {
|
||||
return Mono.error(new IllegalStateException("Server does not provide the resources capability"));
|
||||
}
|
||||
|
|
@ -722,7 +806,7 @@ public class McpAsyncClient {
|
|||
* @see #unsubscribeResource(McpSchema.UnsubscribeRequest)
|
||||
*/
|
||||
public Mono<Void> subscribeResource(McpSchema.SubscribeRequest subscribeRequest) {
|
||||
return this.initializer.withIntitialization("subscribing to resources", init -> init.mcpSession()
|
||||
return this.initializer.withInitialization("subscribing to resources", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_RESOURCES_SUBSCRIBE, subscribeRequest, VOID_TYPE_REFERENCE));
|
||||
}
|
||||
|
||||
|
|
@ -736,7 +820,7 @@ public class McpAsyncClient {
|
|||
* @see #subscribeResource(McpSchema.SubscribeRequest)
|
||||
*/
|
||||
public Mono<Void> unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) {
|
||||
return this.initializer.withIntitialization("unsubscribing from resources", init -> init.mcpSession()
|
||||
return this.initializer.withInitialization("unsubscribing from resources", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_RESOURCES_UNSUBSCRIBE, unsubscribeRequest, VOID_TYPE_REFERENCE));
|
||||
}
|
||||
|
||||
|
|
@ -802,7 +886,7 @@ public class McpAsyncClient {
|
|||
* @see #getPrompt(GetPromptRequest)
|
||||
*/
|
||||
public Mono<ListPromptsResult> listPrompts(String cursor) {
|
||||
return this.initializer.withIntitialization("listing prompts", init -> init.mcpSession()
|
||||
return this.initializer.withInitialization("listing prompts", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_PROMPT_LIST, new PaginatedRequest(cursor), LIST_PROMPTS_RESULT_TYPE_REF));
|
||||
}
|
||||
|
||||
|
|
@ -816,7 +900,7 @@ public class McpAsyncClient {
|
|||
* @see #listPrompts()
|
||||
*/
|
||||
public Mono<GetPromptResult> getPrompt(GetPromptRequest getPromptRequest) {
|
||||
return this.initializer.withIntitialization("getting prompts", init -> init.mcpSession()
|
||||
return this.initializer.withInitialization("getting prompts", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_PROMPT_GET, getPromptRequest, GET_PROMPT_RESULT_TYPE_REF));
|
||||
}
|
||||
|
||||
|
|
@ -834,14 +918,6 @@ public class McpAsyncClient {
|
|||
// --------------------------
|
||||
// Logging
|
||||
// --------------------------
|
||||
/**
|
||||
* Create a notification handler for logging notifications from the server. This
|
||||
* handler automatically distributes logging messages to all registered consumers.
|
||||
* @param loggingConsumers List of consumers that will be notified when a logging
|
||||
* message is received. Each consumer receives the logging message notification.
|
||||
* @return A NotificationHandler that processes log notifications by distributing the
|
||||
* message to all registered consumers
|
||||
*/
|
||||
private NotificationHandler asyncLoggingNotificationHandler(
|
||||
List<Function<LoggingMessageNotification, Mono<Void>>> loggingConsumers) {
|
||||
|
||||
|
|
@ -867,7 +943,7 @@ public class McpAsyncClient {
|
|||
return Mono.error(new IllegalArgumentException("Logging level must not be null"));
|
||||
}
|
||||
|
||||
return this.initializer.withIntitialization("setting logging level", init -> {
|
||||
return this.initializer.withInitialization("setting logging level", init -> {
|
||||
if (init.initializeResult().capabilities().logging() == null) {
|
||||
return Mono.error(new IllegalStateException("Server's Logging capabilities are not enabled!"));
|
||||
}
|
||||
|
|
@ -876,15 +952,6 @@ public class McpAsyncClient {
|
|||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a notification handler for progress notifications from the server. This
|
||||
* handler automatically distributes progress notifications to all registered
|
||||
* consumers.
|
||||
* @param progressConsumers List of consumers that will be notified when a progress
|
||||
* message is received. Each consumer receives the progress notification.
|
||||
* @return A NotificationHandler that processes progress notifications by distributing
|
||||
* the message to all registered consumers
|
||||
*/
|
||||
private NotificationHandler asyncProgressNotificationHandler(
|
||||
List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers) {
|
||||
|
||||
|
|
@ -924,7 +991,7 @@ public class McpAsyncClient {
|
|||
* @see McpSchema.CompleteResult
|
||||
*/
|
||||
public Mono<McpSchema.CompleteResult> completeCompletion(McpSchema.CompleteRequest completeRequest) {
|
||||
return this.initializer.withIntitialization("complete completions", init -> init.mcpSession()
|
||||
return this.initializer.withInitialization("complete completions", init -> init.mcpSession()
|
||||
.sendRequest(McpSchema.METHOD_COMPLETION_COMPLETE, completeRequest, COMPLETION_COMPLETE_RESULT_TYPE_REF));
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import java.util.function.Consumer;
|
|||
import java.util.function.Function;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
import io.modelcontextprotocol.json.schema.JsonSchemaValidator;
|
||||
import io.modelcontextprotocol.common.McpTransportContext;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
|
|
@ -99,6 +100,7 @@ import reactor.core.publisher.Mono;
|
|||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Dariusz Jędrzejczyk
|
||||
* @author Anurag Pant
|
||||
* @see McpAsyncClient
|
||||
* @see McpSyncClient
|
||||
* @see McpTransport
|
||||
|
|
@ -187,6 +189,10 @@ public interface McpClient {
|
|||
|
||||
private Supplier<McpTransportContext> contextProvider = () -> McpTransportContext.EMPTY;
|
||||
|
||||
private JsonSchemaValidator jsonSchemaValidator;
|
||||
|
||||
private boolean enableCallToolSchemaCaching = false; // Default to false
|
||||
|
||||
private SyncSpec(McpClientTransport transport) {
|
||||
Assert.notNull(transport, "Transport must not be null");
|
||||
this.transport = transport;
|
||||
|
|
@ -429,6 +435,32 @@ public interface McpClient {
|
|||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Add a {@link JsonSchemaValidator} to validate the JSON structure of the
|
||||
* structured output.
|
||||
* @param jsonSchemaValidator A validator to validate the JSON structure of the
|
||||
* structured output. Must not be null.
|
||||
* @return This builder for method chaining
|
||||
* @throws IllegalArgumentException if jsonSchemaValidator is null
|
||||
*/
|
||||
public SyncSpec jsonSchemaValidator(JsonSchemaValidator jsonSchemaValidator) {
|
||||
Assert.notNull(jsonSchemaValidator, "JsonSchemaValidator must not be null");
|
||||
this.jsonSchemaValidator = jsonSchemaValidator;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Enables automatic schema caching during callTool operations. When a tool's
|
||||
* output schema is not found in the cache, callTool will automatically fetch and
|
||||
* cache all tool schemas via listTools.
|
||||
* @param enableCallToolSchemaCaching true to enable, false to disable
|
||||
* @return This builder instance for method chaining
|
||||
*/
|
||||
public SyncSpec enableCallToolSchemaCaching(boolean enableCallToolSchemaCaching) {
|
||||
this.enableCallToolSchemaCaching = enableCallToolSchemaCaching;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create an instance of {@link McpSyncClient} with the provided configurations or
|
||||
* sensible defaults.
|
||||
|
|
@ -438,13 +470,13 @@ public interface McpClient {
|
|||
McpClientFeatures.Sync syncFeatures = new McpClientFeatures.Sync(this.clientInfo, this.capabilities,
|
||||
this.roots, this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
|
||||
this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers, this.samplingHandler,
|
||||
this.elicitationHandler);
|
||||
this.elicitationHandler, this.enableCallToolSchemaCaching);
|
||||
|
||||
McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);
|
||||
|
||||
return new McpSyncClient(
|
||||
new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures),
|
||||
this.contextProvider);
|
||||
return new McpSyncClient(new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout,
|
||||
jsonSchemaValidator != null ? jsonSchemaValidator : JsonSchemaValidator.getDefault(),
|
||||
asyncFeatures), this.contextProvider);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -495,6 +527,10 @@ public interface McpClient {
|
|||
|
||||
private Function<ElicitRequest, Mono<ElicitResult>> elicitationHandler;
|
||||
|
||||
private JsonSchemaValidator jsonSchemaValidator;
|
||||
|
||||
private boolean enableCallToolSchemaCaching = false; // Default to false
|
||||
|
||||
private AsyncSpec(McpClientTransport transport) {
|
||||
Assert.notNull(transport, "Transport must not be null");
|
||||
this.transport = transport;
|
||||
|
|
@ -741,17 +777,45 @@ public interface McpClient {
|
|||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the JSON schema validator to use for validating tool responses against
|
||||
* output schemas.
|
||||
* @param jsonSchemaValidator The validator to use. Must not be null.
|
||||
* @return This builder instance for method chaining
|
||||
* @throws IllegalArgumentException if jsonSchemaValidator is null
|
||||
*/
|
||||
public AsyncSpec jsonSchemaValidator(JsonSchemaValidator jsonSchemaValidator) {
|
||||
Assert.notNull(jsonSchemaValidator, "JsonSchemaValidator must not be null");
|
||||
this.jsonSchemaValidator = jsonSchemaValidator;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Enables automatic schema caching during callTool operations. When a tool's
|
||||
* output schema is not found in the cache, callTool will automatically fetch and
|
||||
* cache all tool schemas via listTools.
|
||||
* @param enableCallToolSchemaCaching true to enable, false to disable
|
||||
* @return This builder instance for method chaining
|
||||
*/
|
||||
public AsyncSpec enableCallToolSchemaCaching(boolean enableCallToolSchemaCaching) {
|
||||
this.enableCallToolSchemaCaching = enableCallToolSchemaCaching;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create an instance of {@link McpAsyncClient} with the provided configurations
|
||||
* or sensible defaults.
|
||||
* @return a new instance of {@link McpAsyncClient}.
|
||||
*/
|
||||
public McpAsyncClient build() {
|
||||
var jsonSchemaValidator = (this.jsonSchemaValidator != null) ? this.jsonSchemaValidator
|
||||
: JsonSchemaValidator.getDefault();
|
||||
return new McpAsyncClient(this.transport, this.requestTimeout, this.initializationTimeout,
|
||||
jsonSchemaValidator,
|
||||
new McpClientFeatures.Async(this.clientInfo, this.capabilities, this.roots,
|
||||
this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
|
||||
this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers,
|
||||
this.samplingHandler, this.elicitationHandler));
|
||||
this.samplingHandler, this.elicitationHandler, this.enableCallToolSchemaCaching));
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -62,6 +62,7 @@ class McpClientFeatures {
|
|||
* @param progressConsumers the progress consumers.
|
||||
* @param samplingHandler the sampling handler.
|
||||
* @param elicitationHandler the elicitation handler.
|
||||
* @param enableCallToolSchemaCaching whether to enable call tool schema caching.
|
||||
*/
|
||||
record Async(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities,
|
||||
Map<String, McpSchema.Root> roots, List<Function<List<McpSchema.Tool>, Mono<Void>>> toolsChangeConsumers,
|
||||
|
|
@ -71,7 +72,8 @@ class McpClientFeatures {
|
|||
List<Function<McpSchema.LoggingMessageNotification, Mono<Void>>> loggingConsumers,
|
||||
List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers,
|
||||
Function<McpSchema.CreateMessageRequest, Mono<McpSchema.CreateMessageResult>> samplingHandler,
|
||||
Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler) {
|
||||
Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler,
|
||||
boolean enableCallToolSchemaCaching) {
|
||||
|
||||
/**
|
||||
* Create an instance and validate the arguments.
|
||||
|
|
@ -84,6 +86,7 @@ class McpClientFeatures {
|
|||
* @param progressConsumers the progress consumers.
|
||||
* @param samplingHandler the sampling handler.
|
||||
* @param elicitationHandler the elicitation handler.
|
||||
* @param enableCallToolSchemaCaching whether to enable call tool schema caching.
|
||||
*/
|
||||
public Async(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities,
|
||||
Map<String, McpSchema.Root> roots,
|
||||
|
|
@ -94,7 +97,8 @@ class McpClientFeatures {
|
|||
List<Function<McpSchema.LoggingMessageNotification, Mono<Void>>> loggingConsumers,
|
||||
List<Function<McpSchema.ProgressNotification, Mono<Void>>> progressConsumers,
|
||||
Function<McpSchema.CreateMessageRequest, Mono<McpSchema.CreateMessageResult>> samplingHandler,
|
||||
Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler) {
|
||||
Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler,
|
||||
boolean enableCallToolSchemaCaching) {
|
||||
|
||||
Assert.notNull(clientInfo, "Client info must not be null");
|
||||
this.clientInfo = clientInfo;
|
||||
|
|
@ -113,6 +117,7 @@ class McpClientFeatures {
|
|||
this.progressConsumers = progressConsumers != null ? progressConsumers : List.of();
|
||||
this.samplingHandler = samplingHandler;
|
||||
this.elicitationHandler = elicitationHandler;
|
||||
this.enableCallToolSchemaCaching = enableCallToolSchemaCaching;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -129,7 +134,7 @@ class McpClientFeatures {
|
|||
Function<McpSchema.ElicitRequest, Mono<McpSchema.ElicitResult>> elicitationHandler) {
|
||||
this(clientInfo, clientCapabilities, roots, toolsChangeConsumers, resourcesChangeConsumers,
|
||||
resourcesUpdateConsumers, promptsChangeConsumers, loggingConsumers, List.of(), samplingHandler,
|
||||
elicitationHandler);
|
||||
elicitationHandler, false);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -187,7 +192,8 @@ class McpClientFeatures {
|
|||
|
||||
return new Async(syncSpec.clientInfo(), syncSpec.clientCapabilities(), syncSpec.roots(),
|
||||
toolsChangeConsumers, resourcesChangeConsumers, resourcesUpdateConsumers, promptsChangeConsumers,
|
||||
loggingConsumers, progressConsumers, samplingHandler, elicitationHandler);
|
||||
loggingConsumers, progressConsumers, samplingHandler, elicitationHandler,
|
||||
syncSpec.enableCallToolSchemaCaching);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -205,6 +211,7 @@ class McpClientFeatures {
|
|||
* @param progressConsumers the progress consumers.
|
||||
* @param samplingHandler the sampling handler.
|
||||
* @param elicitationHandler the elicitation handler.
|
||||
* @param enableCallToolSchemaCaching whether to enable call tool schema caching.
|
||||
*/
|
||||
public record Sync(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities,
|
||||
Map<String, McpSchema.Root> roots, List<Consumer<List<McpSchema.Tool>>> toolsChangeConsumers,
|
||||
|
|
@ -214,7 +221,8 @@ class McpClientFeatures {
|
|||
List<Consumer<McpSchema.LoggingMessageNotification>> loggingConsumers,
|
||||
List<Consumer<McpSchema.ProgressNotification>> progressConsumers,
|
||||
Function<McpSchema.CreateMessageRequest, McpSchema.CreateMessageResult> samplingHandler,
|
||||
Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler) {
|
||||
Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler,
|
||||
boolean enableCallToolSchemaCaching) {
|
||||
|
||||
/**
|
||||
* Create an instance and validate the arguments.
|
||||
|
|
@ -229,6 +237,7 @@ class McpClientFeatures {
|
|||
* @param progressConsumers the progress consumers.
|
||||
* @param samplingHandler the sampling handler.
|
||||
* @param elicitationHandler the elicitation handler.
|
||||
* @param enableCallToolSchemaCaching whether to enable call tool schema caching.
|
||||
*/
|
||||
public Sync(McpSchema.Implementation clientInfo, McpSchema.ClientCapabilities clientCapabilities,
|
||||
Map<String, McpSchema.Root> roots, List<Consumer<List<McpSchema.Tool>>> toolsChangeConsumers,
|
||||
|
|
@ -238,7 +247,8 @@ class McpClientFeatures {
|
|||
List<Consumer<McpSchema.LoggingMessageNotification>> loggingConsumers,
|
||||
List<Consumer<McpSchema.ProgressNotification>> progressConsumers,
|
||||
Function<McpSchema.CreateMessageRequest, McpSchema.CreateMessageResult> samplingHandler,
|
||||
Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler) {
|
||||
Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler,
|
||||
boolean enableCallToolSchemaCaching) {
|
||||
|
||||
Assert.notNull(clientInfo, "Client info must not be null");
|
||||
this.clientInfo = clientInfo;
|
||||
|
|
@ -257,6 +267,7 @@ class McpClientFeatures {
|
|||
this.progressConsumers = progressConsumers != null ? progressConsumers : List.of();
|
||||
this.samplingHandler = samplingHandler;
|
||||
this.elicitationHandler = elicitationHandler;
|
||||
this.enableCallToolSchemaCaching = enableCallToolSchemaCaching;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -272,7 +283,7 @@ class McpClientFeatures {
|
|||
Function<McpSchema.ElicitRequest, McpSchema.ElicitResult> elicitationHandler) {
|
||||
this(clientInfo, clientCapabilities, roots, toolsChangeConsumers, resourcesChangeConsumers,
|
||||
resourcesUpdateConsumers, promptsChangeConsumers, loggingConsumers, List.of(), samplingHandler,
|
||||
elicitationHandler);
|
||||
elicitationHandler, false);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -17,7 +17,6 @@ import com.fasterxml.jackson.annotation.JsonInclude;
|
|||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
import com.fasterxml.jackson.annotation.JsonSubTypes;
|
||||
import com.fasterxml.jackson.annotation.JsonTypeInfo;
|
||||
import com.fasterxml.jackson.annotation.JsonTypeInfo.As;
|
||||
import io.modelcontextprotocol.json.McpJsonMapper;
|
||||
import io.modelcontextprotocol.json.TypeRef;
|
||||
import io.modelcontextprotocol.util.Assert;
|
||||
|
|
|
|||
|
|
@ -4,15 +4,14 @@
|
|||
|
||||
package io.modelcontextprotocol.client;
|
||||
|
||||
import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import org.junit.jupiter.api.AfterAll;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
|
||||
import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
|
||||
@Timeout(15)
|
||||
public class HttpClientStreamableHttpAsyncClientTests extends AbstractMcpAsyncClientTests {
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,280 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
|
||||
import io.modelcontextprotocol.client.LifecycleInitializer.Initialization;
|
||||
import io.modelcontextprotocol.spec.McpClientSession;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.MockitoAnnotations;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
import reactor.test.StepVerifier;
|
||||
import reactor.util.context.ContextView;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.never;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
/**
|
||||
* Tests for {@link LifecycleInitializer} postInitializationHook functionality.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
class LifecycleInitializerPostInitializationHookTests {
|
||||
|
||||
private static final Duration INITIALIZATION_TIMEOUT = Duration.ofSeconds(5);
|
||||
|
||||
private static final McpSchema.ClientCapabilities CLIENT_CAPABILITIES = McpSchema.ClientCapabilities.builder()
|
||||
.build();
|
||||
|
||||
private static final McpSchema.Implementation CLIENT_INFO = new McpSchema.Implementation("test-client", "1.0.0");
|
||||
|
||||
private static final List<String> PROTOCOL_VERSIONS = List.of("1.0.0", "2.0.0");
|
||||
|
||||
private static final McpSchema.InitializeResult MOCK_INIT_RESULT = new McpSchema.InitializeResult("2.0.0",
|
||||
McpSchema.ServerCapabilities.builder().build(), new McpSchema.Implementation("test-server", "1.0.0"),
|
||||
"Test instructions");
|
||||
|
||||
@Mock
|
||||
private McpClientSession mockClientSession;
|
||||
|
||||
@Mock
|
||||
private Function<ContextView, McpClientSession> mockSessionSupplier;
|
||||
|
||||
@Mock
|
||||
private Function<Initialization, Mono<Void>> mockPostInitializationHook;
|
||||
|
||||
private LifecycleInitializer initializer;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
MockitoAnnotations.openMocks(this);
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.empty());
|
||||
when(mockSessionSupplier.apply(any(ContextView.class))).thenReturn(mockClientSession);
|
||||
when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
|
||||
.thenReturn(Mono.just(MOCK_INIT_RESULT));
|
||||
when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
|
||||
.thenReturn(Mono.empty());
|
||||
when(mockClientSession.closeGracefully()).thenReturn(Mono.empty());
|
||||
|
||||
initializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS,
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook);
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInvokePostInitializationHook() {
|
||||
AtomicReference<Initialization> capturedInit = new AtomicReference<>();
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
|
||||
capturedInit.set(invocation.getArgument(0));
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
// Verify hook was called
|
||||
verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
|
||||
|
||||
// Verify the hook received correct initialization data
|
||||
assertThat(capturedInit.get()).isNotNull();
|
||||
assertThat(capturedInit.get().mcpSession()).isEqualTo(mockClientSession);
|
||||
assertThat(capturedInit.get().initializeResult()).isEqualTo(MOCK_INIT_RESULT);
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInvokePostInitializationHookOnlyOnce() {
|
||||
// First initialization
|
||||
StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
|
||||
.expectNext("result1")
|
||||
.verifyComplete();
|
||||
|
||||
// Second call should reuse initialization and NOT call hook again
|
||||
StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2")))
|
||||
.expectNext("result2")
|
||||
.verifyComplete();
|
||||
|
||||
// Hook should only be called once
|
||||
verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInvokePostInitializationHookOnlyOnceWithConcurrentRequests() {
|
||||
AtomicInteger hookInvocationCount = new AtomicInteger(0);
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
|
||||
hookInvocationCount.incrementAndGet();
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
// Start multiple concurrent initializations
|
||||
Mono<String> init1 = initializer.withInitialization("test1", init -> Mono.just("result1"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
Mono<String> init2 = initializer.withInitialization("test2", init -> Mono.just("result2"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
Mono<String> init3 = initializer.withInitialization("test3", init -> Mono.just("result3"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
|
||||
// TODO: can we assume the order of results?
|
||||
StepVerifier.create(Mono.zip(init1, init2, init3)).assertNext(tuple -> {
|
||||
assertThat(tuple.getT1()).isEqualTo("result1");
|
||||
assertThat(tuple.getT2()).isEqualTo("result2");
|
||||
assertThat(tuple.getT3()).isEqualTo("result3");
|
||||
}).verifyComplete();
|
||||
|
||||
// Hook should only be called once despite concurrent requests
|
||||
assertThat(hookInvocationCount.get()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldFailInitializationWhenPostInitializationHookFails() {
|
||||
RuntimeException hookError = new RuntimeException("Post-initialization hook failed");
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.error(hookError));
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectErrorMatches(ex -> ex instanceof RuntimeException && ex.getCause() == hookError)
|
||||
.verify();
|
||||
|
||||
// Verify initialization was not completed
|
||||
assertThat(initializer.isInitialized()).isFalse();
|
||||
assertThat(initializer.currentInitializationResult()).isNull();
|
||||
|
||||
// Verify the hook was called
|
||||
verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldNotInvokePostInitializationHookWhenInitializationFails() {
|
||||
when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
|
||||
.thenReturn(Mono.error(new RuntimeException("Initialization failed")));
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectError(RuntimeException.class)
|
||||
.verify();
|
||||
|
||||
// Hook should NOT be called when initialization fails
|
||||
verify(mockPostInitializationHook, never()).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldNotInvokePostInitializationHookWhenNotificationFails() {
|
||||
when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
|
||||
.thenReturn(Mono.error(new RuntimeException("Notification failed")));
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectError(RuntimeException.class)
|
||||
.verify();
|
||||
|
||||
// Hook should NOT be called when notification fails
|
||||
verify(mockPostInitializationHook, never()).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInvokePostInitializationHookAgainAfterReinitialization() {
|
||||
AtomicInteger hookInvocationCount = new AtomicInteger(0);
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
|
||||
hookInvocationCount.incrementAndGet();
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
// First initialization
|
||||
StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
|
||||
.expectNext("result1")
|
||||
.verifyComplete();
|
||||
|
||||
assertThat(hookInvocationCount.get()).isEqualTo(1);
|
||||
|
||||
// Simulate transport session exception to trigger re-initialization
|
||||
initializer.handleException(new McpTransportSessionNotFoundException("Session lost"));
|
||||
|
||||
// Hook should be called twice (once for each initialization)
|
||||
assertThat(hookInvocationCount.get()).isEqualTo(2);
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldAllowPostInitializationHookToPerformAsyncOperations() {
|
||||
AtomicInteger operationCount = new AtomicInteger(0);
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class)))
|
||||
.thenReturn(Mono.fromRunnable(() -> operationCount.incrementAndGet()).then());
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
// Verify the async operation was executed
|
||||
assertThat(operationCount.get()).isEqualTo(1);
|
||||
verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldProvideCorrectInitializationDataToHook() {
|
||||
AtomicReference<McpClientSession> capturedSession = new AtomicReference<>();
|
||||
AtomicReference<McpSchema.InitializeResult> capturedResult = new AtomicReference<>();
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
|
||||
Initialization init = invocation.getArgument(0);
|
||||
capturedSession.set(init.mcpSession());
|
||||
capturedResult.set(init.initializeResult());
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
// Verify the hook received the correct session and result
|
||||
assertThat(capturedSession.get()).isEqualTo(mockClientSession);
|
||||
assertThat(capturedResult.get()).isEqualTo(MOCK_INIT_RESULT);
|
||||
assertThat(capturedResult.get().protocolVersion()).isEqualTo("2.0.0");
|
||||
assertThat(capturedResult.get().serverInfo().name()).isEqualTo("test-server");
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInvokePostInitializationHookAfterSuccessfulInitialization() {
|
||||
AtomicReference<Boolean> notificationSent = new AtomicReference<>(false);
|
||||
AtomicReference<Boolean> hookCalledAfterNotification = new AtomicReference<>(false);
|
||||
|
||||
when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
|
||||
.thenAnswer(invocation -> {
|
||||
notificationSent.set(true);
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
|
||||
// Due to flatMap chaining in doInitialize, if the hook is called,
|
||||
// the notification must have been sent first
|
||||
hookCalledAfterNotification.set(notificationSent.get());
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
// Verify the hook was called and notification was already sent at that point
|
||||
assertThat(hookCalledAfterNotification.get()).isTrue();
|
||||
verify(mockClientSession).sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any());
|
||||
verify(mockPostInitializationHook).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -10,14 +10,14 @@ import java.util.concurrent.atomic.AtomicInteger;
|
|||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
|
||||
import io.modelcontextprotocol.client.LifecycleInitializer.Initialization;
|
||||
import io.modelcontextprotocol.spec.McpClientSession;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.MockitoAnnotations;
|
||||
|
||||
import io.modelcontextprotocol.spec.McpClientSession;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
import reactor.test.StepVerifier;
|
||||
|
|
@ -58,12 +58,16 @@ class LifecycleInitializerTests {
|
|||
@Mock
|
||||
private Function<ContextView, McpClientSession> mockSessionSupplier;
|
||||
|
||||
@Mock
|
||||
private Function<Initialization, Mono<Void>> mockPostInitializationHook;
|
||||
|
||||
private LifecycleInitializer initializer;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
MockitoAnnotations.openMocks(this);
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.empty());
|
||||
when(mockSessionSupplier.apply(any(ContextView.class))).thenReturn(mockClientSession);
|
||||
when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
|
||||
.thenReturn(Mono.just(MOCK_INIT_RESULT));
|
||||
|
|
@ -72,45 +76,45 @@ class LifecycleInitializerTests {
|
|||
when(mockClientSession.closeGracefully()).thenReturn(Mono.empty());
|
||||
|
||||
initializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS,
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier);
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook);
|
||||
}
|
||||
|
||||
@Test
|
||||
void constructorShouldValidateParameters() {
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(null, CLIENT_INFO, PROTOCOL_VERSIONS, INITIALIZATION_TIMEOUT,
|
||||
mockSessionSupplier))
|
||||
mockSessionSupplier, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Client capabilities must not be null");
|
||||
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, null, PROTOCOL_VERSIONS,
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier))
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Client info must not be null");
|
||||
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, null,
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier))
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Protocol versions must not be empty");
|
||||
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, List.of(),
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier))
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Protocol versions must not be empty");
|
||||
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS, null,
|
||||
mockSessionSupplier))
|
||||
mockSessionSupplier, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Initialization timeout must not be null");
|
||||
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS,
|
||||
INITIALIZATION_TIMEOUT, null))
|
||||
INITIALIZATION_TIMEOUT, null, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Session supplier must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInitializeSuccessfully() {
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.assertNext(result -> {
|
||||
assertThat(result).isEqualTo(MOCK_INIT_RESULT);
|
||||
assertThat(initializer.isInitialized()).isTrue();
|
||||
|
|
@ -132,7 +136,7 @@ class LifecycleInitializerTests {
|
|||
return Mono.just(MOCK_INIT_RESULT);
|
||||
});
|
||||
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.assertNext(result -> {
|
||||
assertThat(capturedRequest.get().protocolVersion()).isEqualTo("2.0.0"); // Latest
|
||||
// version
|
||||
|
|
@ -152,7 +156,7 @@ class LifecycleInitializerTests {
|
|||
when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
|
||||
.thenReturn(Mono.just(unsupportedResult));
|
||||
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectError(RuntimeException.class)
|
||||
.verify();
|
||||
|
||||
|
|
@ -167,13 +171,13 @@ class LifecycleInitializerTests {
|
|||
Duration SLOW_RESPONSE_DELAY = Duration.ofSeconds(5);
|
||||
|
||||
LifecycleInitializer shortTimeoutInitializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO,
|
||||
PROTOCOL_VERSIONS, INITIALIZE_TIMEOUT, mockSessionSupplier);
|
||||
PROTOCOL_VERSIONS, INITIALIZE_TIMEOUT, mockSessionSupplier, mockPostInitializationHook);
|
||||
|
||||
when(mockClientSession.<McpSchema.InitializeResult>sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
|
||||
.thenReturn(Mono.just(MOCK_INIT_RESULT).delayElement(SLOW_RESPONSE_DELAY, virtualTimeScheduler));
|
||||
|
||||
StepVerifier
|
||||
.withVirtualTime(() -> shortTimeoutInitializer.withIntitialization("test",
|
||||
.withVirtualTime(() -> shortTimeoutInitializer.withInitialization("test",
|
||||
init -> Mono.just(init.initializeResult())), () -> virtualTimeScheduler, Long.MAX_VALUE)
|
||||
.expectSubscription()
|
||||
.expectNoEvent(INITIALIZE_TIMEOUT)
|
||||
|
|
@ -184,12 +188,12 @@ class LifecycleInitializerTests {
|
|||
@Test
|
||||
void shouldReuseExistingInitialization() {
|
||||
// First initialization
|
||||
StepVerifier.create(initializer.withIntitialization("test1", init -> Mono.just("result1")))
|
||||
StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
|
||||
.expectNext("result1")
|
||||
.verifyComplete();
|
||||
|
||||
// Second call should reuse the same initialization
|
||||
StepVerifier.create(initializer.withIntitialization("test2", init -> Mono.just("result2")))
|
||||
StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2")))
|
||||
.expectNext("result2")
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -209,11 +213,11 @@ class LifecycleInitializerTests {
|
|||
|
||||
// Start multiple concurrent initializations using subscribeOn with parallel
|
||||
// scheduler
|
||||
Mono<String> init1 = initializer.withIntitialization("test1", init -> Mono.just("result1"))
|
||||
Mono<String> init1 = initializer.withInitialization("test1", init -> Mono.just("result1"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
Mono<String> init2 = initializer.withIntitialization("test2", init -> Mono.just("result2"))
|
||||
Mono<String> init2 = initializer.withInitialization("test2", init -> Mono.just("result2"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
Mono<String> init3 = initializer.withIntitialization("test3", init -> Mono.just("result3"))
|
||||
Mono<String> init3 = initializer.withInitialization("test3", init -> Mono.just("result3"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
|
||||
StepVerifier.create(Mono.zip(init1, init2, init3)).assertNext(tuple -> {
|
||||
|
|
@ -235,7 +239,7 @@ class LifecycleInitializerTests {
|
|||
// succeeds on the second call
|
||||
.thenReturn(Mono.just(MOCK_INIT_RESULT));
|
||||
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectError(RuntimeException.class)
|
||||
.verify();
|
||||
|
||||
|
|
@ -244,7 +248,7 @@ class LifecycleInitializerTests {
|
|||
|
||||
// The initializer can recover from previous errors
|
||||
StepVerifier
|
||||
.create(initializer.withIntitialization("successful init", init -> Mono.just(init.initializeResult())))
|
||||
.create(initializer.withInitialization("successful init", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -255,7 +259,7 @@ class LifecycleInitializerTests {
|
|||
@Test
|
||||
void shouldHandleTransportSessionNotFoundException() {
|
||||
// successful initialization first
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -277,7 +281,7 @@ class LifecycleInitializerTests {
|
|||
@Test
|
||||
void shouldHandleOtherExceptions() {
|
||||
// Simulate a successful initialization first
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -295,7 +299,7 @@ class LifecycleInitializerTests {
|
|||
|
||||
@Test
|
||||
void shouldCloseGracefully() {
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -307,7 +311,7 @@ class LifecycleInitializerTests {
|
|||
|
||||
@Test
|
||||
void shouldCloseImmediately() {
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -342,7 +346,7 @@ class LifecycleInitializerTests {
|
|||
new McpSchema.Implementation("test-server", "1.0.0"), "Test instructions"));
|
||||
});
|
||||
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.assertNext(result -> {
|
||||
// Latest from new versions
|
||||
assertThat(capturedRequest.get().protocolVersion()).isEqualTo("4.0.0");
|
||||
|
|
@ -363,7 +367,7 @@ class LifecycleInitializerTests {
|
|||
});
|
||||
|
||||
StepVerifier
|
||||
.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))
|
||||
.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))
|
||||
.contextWrite(Context.of(contextKey, contextValue)))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
|
@ -374,7 +378,7 @@ class LifecycleInitializerTests {
|
|||
|
||||
@Test
|
||||
void shouldProvideAccessToMcpSessionAndInitializeResult() {
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> {
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> {
|
||||
assertThat(init.mcpSession()).isEqualTo(mockClientSession);
|
||||
assertThat(init.initializeResult()).isEqualTo(MOCK_INIT_RESULT);
|
||||
return Mono.just("success");
|
||||
|
|
@ -386,7 +390,7 @@ class LifecycleInitializerTests {
|
|||
when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
|
||||
.thenReturn(Mono.error(new RuntimeException("Notification failed")));
|
||||
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectError(RuntimeException.class)
|
||||
.verify();
|
||||
|
||||
|
|
@ -403,7 +407,7 @@ class LifecycleInitializerTests {
|
|||
@Test
|
||||
void shouldReinitializeAfterTransportSessionException() {
|
||||
// First initialization
|
||||
StepVerifier.create(initializer.withIntitialization("test1", init -> Mono.just("result1")))
|
||||
StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
|
||||
.expectNext("result1")
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -411,7 +415,7 @@ class LifecycleInitializerTests {
|
|||
initializer.handleException(new McpTransportSessionNotFoundException("Session lost"));
|
||||
|
||||
// Should be able to initialize again
|
||||
StepVerifier.create(initializer.withIntitialization("test2", init -> Mono.just("result2")))
|
||||
StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2")))
|
||||
.expectNext("result2")
|
||||
.verifyComplete();
|
||||
|
||||
|
|
|
|||
|
|
@ -10,12 +10,20 @@ import io.modelcontextprotocol.spec.McpSchema;
|
|||
import io.modelcontextprotocol.spec.ProtocolVersions;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import reactor.core.publisher.Mono;
|
||||
|
||||
import com.fasterxml.jackson.core.JsonProcessingException;
|
||||
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Objects;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
|
||||
import static io.modelcontextprotocol.util.McpJsonMapperUtils.JSON_MAPPER;
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatCode;
|
||||
|
||||
class McpAsyncClientTests {
|
||||
|
|
@ -24,6 +32,7 @@ class McpAsyncClientTests {
|
|||
"1.0.0");
|
||||
|
||||
public static final McpSchema.ServerCapabilities MOCK_SERVER_CAPABILITIES = McpSchema.ServerCapabilities.builder()
|
||||
.tools(true)
|
||||
.build();
|
||||
|
||||
public static final McpSchema.InitializeResult MOCK_INIT_RESULT = new McpSchema.InitializeResult(
|
||||
|
|
@ -31,6 +40,91 @@ class McpAsyncClientTests {
|
|||
|
||||
private static final String CONTEXT_KEY = "context.key";
|
||||
|
||||
private McpClientTransport createMockTransportForToolValidation(boolean hasOutputSchema, boolean invalidOutput)
|
||||
throws JsonProcessingException {
|
||||
|
||||
// Create tool with or without output schema
|
||||
Map<String, Object> inputSchemaMap = Map.of("type", "object", "properties",
|
||||
Map.of("expression", Map.of("type", "string")), "required", List.of("expression"));
|
||||
|
||||
McpSchema.JsonSchema inputSchema = new McpSchema.JsonSchema("object", inputSchemaMap, null, null, null, null);
|
||||
McpSchema.Tool.Builder toolBuilder = McpSchema.Tool.builder()
|
||||
.name("calculator")
|
||||
.description("Performs mathematical calculations")
|
||||
.inputSchema(inputSchema);
|
||||
|
||||
if (hasOutputSchema) {
|
||||
Map<String, Object> outputSchema = Map.of("type", "object", "properties",
|
||||
Map.of("result", Map.of("type", "number"), "operation", Map.of("type", "string")), "required",
|
||||
List.of("result", "operation"));
|
||||
toolBuilder.outputSchema(outputSchema);
|
||||
}
|
||||
|
||||
McpSchema.Tool calculatorTool = toolBuilder.build();
|
||||
McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(calculatorTool), null);
|
||||
|
||||
// Create call tool result - valid or invalid based on parameter
|
||||
Map<String, Object> structuredContent = invalidOutput ? Map.of("result", "5", "operation", "add")
|
||||
: Map.of("result", 5, "operation", "add");
|
||||
|
||||
McpSchema.CallToolResult mockCallToolResult = McpSchema.CallToolResult.builder()
|
||||
.addTextContent("Calculation result")
|
||||
.structuredContent(structuredContent)
|
||||
.build();
|
||||
|
||||
return new McpClientTransport() {
|
||||
Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler;
|
||||
|
||||
@Override
|
||||
public Mono<Void> connect(
|
||||
Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
|
||||
this.handler = handler;
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> closeGracefully() {
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
|
||||
if (!(message instanceof McpSchema.JSONRPCRequest request)) {
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
McpSchema.JSONRPCResponse response;
|
||||
if (McpSchema.METHOD_INITIALIZE.equals(request.method())) {
|
||||
response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), MOCK_INIT_RESULT,
|
||||
null);
|
||||
}
|
||||
else if (McpSchema.METHOD_TOOLS_LIST.equals(request.method())) {
|
||||
response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), mockToolsResult,
|
||||
null);
|
||||
}
|
||||
else if (McpSchema.METHOD_TOOLS_CALL.equals(request.method())) {
|
||||
response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
|
||||
mockCallToolResult, null);
|
||||
}
|
||||
else {
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
return handler.apply(Mono.just(response)).then();
|
||||
}
|
||||
|
||||
@Override
|
||||
public <T> T unmarshalFrom(Object data, TypeRef<T> typeRef) {
|
||||
return JSON_MAPPER.convertValue(data, new TypeRef<>() {
|
||||
@Override
|
||||
public java.lang.reflect.Type getType() {
|
||||
return typeRef.getType();
|
||||
}
|
||||
});
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
@Test
|
||||
void validateContextPassedToTransportConnect() {
|
||||
McpClientTransport transport = new McpClientTransport() {
|
||||
|
|
@ -87,4 +181,64 @@ class McpAsyncClientTests {
|
|||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallToolWithOutputSchemaValidationSuccess() throws JsonProcessingException {
|
||||
McpClientTransport transport = createMockTransportForToolValidation(true, false);
|
||||
|
||||
McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
|
||||
|
||||
StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
|
||||
|
||||
StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))))
|
||||
.expectNextMatches(response -> {
|
||||
assertThat(response).isNotNull();
|
||||
assertThat(response.isError()).isFalse();
|
||||
assertThat(response.structuredContent()).isInstanceOf(Map.class);
|
||||
assertThat((Map<?, ?>) response.structuredContent()).hasSize(2);
|
||||
assertThat(response.content()).hasSize(1);
|
||||
return true;
|
||||
})
|
||||
.verifyComplete();
|
||||
|
||||
StepVerifier.create(client.closeGracefully()).verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallToolWithNoOutputSchemaSuccess() throws JsonProcessingException {
|
||||
McpClientTransport transport = createMockTransportForToolValidation(false, false);
|
||||
|
||||
McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
|
||||
|
||||
StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
|
||||
|
||||
StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))))
|
||||
.expectNextMatches(response -> {
|
||||
assertThat(response).isNotNull();
|
||||
assertThat(response.isError()).isFalse();
|
||||
assertThat(response.structuredContent()).isInstanceOf(Map.class);
|
||||
assertThat((Map<?, ?>) response.structuredContent()).hasSize(2);
|
||||
assertThat(response.content()).hasSize(1);
|
||||
return true;
|
||||
})
|
||||
.verifyComplete();
|
||||
|
||||
StepVerifier.create(client.closeGracefully()).verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallToolWithOutputSchemaValidationFailure() throws JsonProcessingException {
|
||||
McpClientTransport transport = createMockTransportForToolValidation(true, true);
|
||||
|
||||
McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
|
||||
|
||||
StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
|
||||
|
||||
StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))))
|
||||
.expectErrorMatches(ex -> ex instanceof IllegalArgumentException
|
||||
&& ex.getMessage().contains("Tool call result validation failed"))
|
||||
.verify();
|
||||
|
||||
StepVerifier.create(client.closeGracefully()).verifyComplete();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,9 +8,9 @@ import java.time.Duration;
|
|||
import java.util.List;
|
||||
|
||||
import io.modelcontextprotocol.MockMcpClientTransport;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
|
||||
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.test.StepVerifier;
|
||||
|
|
@ -22,7 +22,7 @@ import static org.assertj.core.api.Assertions.assertThat;
|
|||
*/
|
||||
class McpClientProtocolVersionTests {
|
||||
|
||||
private static final Duration REQUEST_TIMEOUT = Duration.ofSeconds(30);
|
||||
private static final Duration REQUEST_TIMEOUT = Duration.ofSeconds(300);
|
||||
|
||||
private static final McpSchema.Implementation CLIENT_INFO = new McpSchema.Implementation("test-client", "1.0.0");
|
||||
|
||||
|
|
@ -46,13 +46,12 @@ class McpClientProtocolVersionTests {
|
|||
assertThat(initRequest.protocolVersion()).isEqualTo(transport.protocolVersions().get(0));
|
||||
|
||||
transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
|
||||
new McpSchema.InitializeResult(protocolVersion, null,
|
||||
new McpSchema.InitializeResult(protocolVersion, ServerCapabilities.builder().build(),
|
||||
new McpSchema.Implementation("test-server", "1.0.0"), null),
|
||||
null));
|
||||
}).assertNext(result -> {
|
||||
assertThat(result.protocolVersion()).isEqualTo(protocolVersion);
|
||||
}).verifyComplete();
|
||||
|
||||
}
|
||||
finally {
|
||||
// Ensure cleanup happens even if test fails
|
||||
|
|
@ -81,7 +80,7 @@ class McpClientProtocolVersionTests {
|
|||
assertThat(initRequest.protocolVersion()).isIn(List.of(oldVersion, McpSchema.LATEST_PROTOCOL_VERSION));
|
||||
|
||||
transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
|
||||
new McpSchema.InitializeResult(oldVersion, null,
|
||||
new McpSchema.InitializeResult(oldVersion, ServerCapabilities.builder().build(),
|
||||
new McpSchema.Implementation("test-server", "1.0.0"), null),
|
||||
null));
|
||||
}).assertNext(result -> {
|
||||
|
|
@ -110,7 +109,7 @@ class McpClientProtocolVersionTests {
|
|||
assertThat(request.params()).isInstanceOf(McpSchema.InitializeRequest.class);
|
||||
|
||||
transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
|
||||
new McpSchema.InitializeResult(unsupportedVersion, null,
|
||||
new McpSchema.InitializeResult(unsupportedVersion, ServerCapabilities.builder().build(),
|
||||
new McpSchema.Implementation("test-server", "1.0.0"), null),
|
||||
null));
|
||||
}).expectError(RuntimeException.class).verify();
|
||||
|
|
@ -143,7 +142,7 @@ class McpClientProtocolVersionTests {
|
|||
assertThat(initRequest.protocolVersion()).isEqualTo(latestVersion);
|
||||
|
||||
transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
|
||||
new McpSchema.InitializeResult(latestVersion, null,
|
||||
new McpSchema.InitializeResult(latestVersion, ServerCapabilities.builder().build(),
|
||||
new McpSchema.Implementation("test-server", "1.0.0"), null),
|
||||
null));
|
||||
}).assertNext(result -> {
|
||||
|
|
|
|||
Loading…
Reference in a new issue