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
5 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;
|
||||
|
|
|
|||
Loading…
Reference in a new issue