Skip to content

Commit de731dd

Browse files
pantanurag555pantanurag555
committed
feat: Add client validation for structuredContent
1 parent a0afdcd commit de731dd

File tree

4 files changed

+333
-22
lines changed

4 files changed

+333
-22
lines changed

‎mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java‎

Lines changed: 88 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,18 @@
1212
import java.util.HashMap;
1313
import java.util.List;
1414
import java.util.Map;
15+
import java.util.Optional;
1516
import java.util.concurrent.ConcurrentHashMap;
1617
import java.util.function.Function;
1718

1819
import org.slf4j.Logger;
1920
import org.slf4j.LoggerFactory;
2021

2122
import io.modelcontextprotocol.json.TypeRef;
22-
23+
import io.modelcontextprotocol.json.schema.JsonSchemaValidator;
2324
import io.modelcontextprotocol.spec.McpClientSession;
2425
import io.modelcontextprotocol.spec.McpClientTransport;
26+
import io.modelcontextprotocol.spec.McpError;
2527
import io.modelcontextprotocol.spec.McpSchema;
2628
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
2729
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
@@ -75,6 +77,7 @@
7577
* @author Dariusz Jędrzejczyk
7678
* @author Christian Tzolov
7779
* @author Jihoon Kim
80+
* @author Anurag Pant
7881
* @see McpClient
7982
* @see McpSchema
8083
* @see McpClientSession
@@ -152,16 +155,33 @@ public class McpAsyncClient {
152155
*/
153156
private final LifecycleInitializer initializer;
154157

158+
/**
159+
* JSON schema validator to use for validating tool responses against output schemas.
160+
*/
161+
private final JsonSchemaValidator jsonSchemaValidator;
162+
163+
/**
164+
* Cached tool output schemas.
165+
*/
166+
private final ConcurrentHashMap<String, Optional<Map<String, Object>>> toolsOutputSchemaCache;
167+
168+
/**
169+
* Whether to enable automatic schema caching during callTool operations.
170+
*/
171+
private final boolean enableCallToolSchemaCaching;
172+
155173
/**
156174
* Create a new McpAsyncClient with the given transport and session request-response
157175
* timeout.
158176
* @param transport the transport to use.
159177
* @param requestTimeout the session request-response timeout.
160178
* @param initializationTimeout the max timeout to await for the client-server
161-
* @param features the MCP Client supported features.
179+
* @param jsonSchemaValidator the JSON schema validator to use for validating tool
180+
* @param features the MCP Client supported features. responses against output
181+
* schemas.
162182
*/
163183
McpAsyncClient(McpClientTransport transport, Duration requestTimeout, Duration initializationTimeout,
164-
McpClientFeatures.Async features) {
184+
JsonSchemaValidator jsonSchemaValidator, McpClientFeatures.Async features) {
165185

166186
Assert.notNull(transport, "Transport must not be null");
167187
Assert.notNull(requestTimeout, "Request timeout must not be null");
@@ -171,6 +191,9 @@ public class McpAsyncClient {
171191
this.clientCapabilities = features.clientCapabilities();
172192
this.transport = transport;
173193
this.roots = new ConcurrentHashMap<>(features.roots());
194+
this.jsonSchemaValidator = jsonSchemaValidator;
195+
this.toolsOutputSchemaCache = new ConcurrentHashMap<>();
196+
this.enableCallToolSchemaCaching = features.enableCallToolSchemaCaching();
174197

175198
// Request Handlers
176199
Map<String, RequestHandler<?>> requestHandlers = new HashMap<>();
@@ -539,15 +562,61 @@ private RequestHandler<ElicitResult> elicitationCreateHandler() {
539562
* @see #listTools()
540563
*/
541564
public Mono<McpSchema.CallToolResult> callTool(McpSchema.CallToolRequest callToolRequest) {
542-
return this.initializer.withIntitialization("calling tools", init -> {
543-
if (init.initializeResult().capabilities().tools() == null) {
544-
return Mono.error(new IllegalStateException("Server does not provide tools capability"));
545-
}
546-
return init.mcpSession()
547-
.sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF);
565+
return Mono.defer(() -> {
566+
// Conditionally cache schemas if needed, otherwise return empty Mono
567+
Mono<Void> cachingStep = (this.enableCallToolSchemaCaching
568+
&& !this.toolsOutputSchemaCache.containsKey(callToolRequest.name())) ? this.listTools().then()
569+
: Mono.empty();
570+
571+
return cachingStep.then(this.initializer.withIntitialization("calling tool", init -> {
572+
if (init.initializeResult().capabilities().tools() == null) {
573+
return Mono.error(new IllegalStateException("Server does not provide tools capability"));
574+
}
575+
576+
return init.mcpSession()
577+
.sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF)
578+
.flatMap(result -> validateToolResult(callToolRequest.name(), result));
579+
}));
548580
});
549581
}
550582

583+
/**
584+
* Calls a tool provided by the server and validates the result against the cached
585+
* output schema.
586+
* @param toolName The name of the tool to call
587+
* @param result The result of the tool call
588+
* @return A Mono that emits the validated tool result
589+
*/
590+
private Mono<McpSchema.CallToolResult> validateToolResult(String toolName, McpSchema.CallToolResult result) {
591+
Optional<Map<String, Object>> optOutputSchema = toolsOutputSchemaCache.get(toolName);
592+
593+
if (result != null && result.isError() != null && !result.isError()) {
594+
if (optOutputSchema == null) {
595+
// Tool not found in cache - skip validation and proceed
596+
logger.debug("Tool '{}' not found in cache, skipping validation", toolName);
597+
return Mono.just(result);
598+
}
599+
else {
600+
if (optOutputSchema.isPresent()) {
601+
// Validate the tool output against the cached output schema
602+
var validation = this.jsonSchemaValidator.validate(optOutputSchema.get(),
603+
result.structuredContent());
604+
if (!validation.valid()) {
605+
logger.warn("Tool call result validation failed: {}", validation.errorMessage());
606+
return Mono.just(new McpSchema.CallToolResult(validation.errorMessage(), true));
607+
}
608+
}
609+
else if (result.structuredContent() != null) {
610+
logger.warn(
611+
"Calling a tool with no outputSchema is not expected to return result with structured content, but got: {}",
612+
result.structuredContent());
613+
}
614+
}
615+
}
616+
617+
return Mono.just(result);
618+
}
619+
551620
/**
552621
* Retrieves the list of all tools provided by the server.
553622
* @return A Mono that emits the list of all tools result
@@ -574,7 +643,16 @@ public Mono<McpSchema.ListToolsResult> listTools(String cursor) {
574643
}
575644
return init.mcpSession()
576645
.sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor),
577-
LIST_TOOLS_RESULT_TYPE_REF);
646+
LIST_TOOLS_RESULT_TYPE_REF)
647+
.map(result -> {
648+
if (result.tools() != null) {
649+
// Cache tools output schema
650+
result.tools()
651+
.forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(),
652+
Optional.ofNullable(tool.outputSchema())));
653+
}
654+
return result;
655+
});
578656
});
579657
}
580658

‎mcp-core/src/main/java/io/modelcontextprotocol/client/McpClient.java‎

Lines changed: 69 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import java.util.function.Function;
1414
import java.util.function.Supplier;
1515

16+
import io.modelcontextprotocol.json.schema.JsonSchemaValidator;
1617
import io.modelcontextprotocol.common.McpTransportContext;
1718
import io.modelcontextprotocol.spec.McpClientTransport;
1819
import io.modelcontextprotocol.spec.McpSchema;
@@ -99,6 +100,7 @@
99100
*
100101
* @author Christian Tzolov
101102
* @author Dariusz Jędrzejczyk
103+
* @author Anurag Pant
102104
* @see McpAsyncClient
103105
* @see McpSyncClient
104106
* @see McpTransport
@@ -187,6 +189,10 @@ class SyncSpec {
187189

188190
private Supplier<McpTransportContext> contextProvider = () -> McpTransportContext.EMPTY;
189191

192+
private JsonSchemaValidator jsonSchemaValidator;
193+
194+
private boolean enableCallToolSchemaCaching = false; // Default to false
195+
190196
private SyncSpec(McpClientTransport transport) {
191197
Assert.notNull(transport, "Transport must not be null");
192198
this.transport = transport;
@@ -429,6 +435,32 @@ public SyncSpec transportContextProvider(Supplier<McpTransportContext> contextPr
429435
return this;
430436
}
431437

438+
/**
439+
* Add a {@link JsonSchemaValidator} to validate the JSON structure of the
440+
* structured output.
441+
* @param jsonSchemaValidator A validator to validate the JSON structure of the
442+
* structured output. Must not be null.
443+
* @return This builder for method chaining
444+
* @throws IllegalArgumentException if jsonSchemaValidator is null
445+
*/
446+
public SyncSpec jsonSchemaValidator(JsonSchemaValidator jsonSchemaValidator) {
447+
Assert.notNull(jsonSchemaValidator, "JsonSchemaValidator must not be null");
448+
this.jsonSchemaValidator = jsonSchemaValidator;
449+
return this;
450+
}
451+
452+
/**
453+
* Enables automatic schema caching during callTool operations. When a tool's
454+
* output schema is not found in the cache, callTool will automatically fetch and
455+
* cache all tool schemas via listTools.
456+
* @param enableCallToolSchemaCaching true to enable, false to disable
457+
* @return This builder instance for method chaining
458+
*/
459+
public SyncSpec enableCallToolSchemaCaching(boolean enableCallToolSchemaCaching) {
460+
this.enableCallToolSchemaCaching = enableCallToolSchemaCaching;
461+
return this;
462+
}
463+
432464
/**
433465
* Create an instance of {@link McpSyncClient} with the provided configurations or
434466
* sensible defaults.
@@ -438,13 +470,13 @@ public McpSyncClient build() {
438470
McpClientFeatures.Sync syncFeatures = new McpClientFeatures.Sync(this.clientInfo, this.capabilities,
439471
this.roots, this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
440472
this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers, this.samplingHandler,
441-
this.elicitationHandler);
473+
this.elicitationHandler, this.enableCallToolSchemaCaching);
442474

443475
McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);
444476

445-
return new McpSyncClient(
446-
new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures),
447-
this.contextProvider);
477+
return new McpSyncClient(new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout,
478+
jsonSchemaValidator != null ? jsonSchemaValidator : JsonSchemaValidator.getDefault(),
479+
asyncFeatures), this.contextProvider);
448480
}
449481

450482
}
@@ -495,6 +527,10 @@ class AsyncSpec {
495527

496528
private Function<ElicitRequest, Mono<ElicitResult>> elicitationHandler;
497529

530+
private JsonSchemaValidator jsonSchemaValidator;
531+
532+
private boolean enableCallToolSchemaCaching = false; // Default to false
533+
498534
private AsyncSpec(McpClientTransport transport) {
499535
Assert.notNull(transport, "Transport must not be null");
500536
this.transport = transport;
@@ -741,17 +777,45 @@ public AsyncSpec progressConsumers(
741777
return this;
742778
}
743779

780+
/**
781+
* Sets the JSON schema validator to use for validating tool responses against
782+
* output schemas.
783+
* @param jsonSchemaValidator The validator to use. Must not be null.
784+
* @return This builder instance for method chaining
785+
* @throws IllegalArgumentException if jsonSchemaValidator is null
786+
*/
787+
public AsyncSpec jsonSchemaValidator(JsonSchemaValidator jsonSchemaValidator) {
788+
Assert.notNull(jsonSchemaValidator, "JsonSchemaValidator must not be null");
789+
this.jsonSchemaValidator = jsonSchemaValidator;
790+
return this;
791+
}
792+
793+
/**
794+
* Enables automatic schema caching during callTool operations. When a tool's
795+
* output schema is not found in the cache, callTool will automatically fetch and
796+
* cache all tool schemas via listTools.
797+
* @param enableCallToolSchemaCaching true to enable, false to disable
798+
* @return This builder instance for method chaining
799+
*/
800+
public AsyncSpec enableCallToolSchemaCaching(boolean enableCallToolSchemaCaching) {
801+
this.enableCallToolSchemaCaching = enableCallToolSchemaCaching;
802+
return this;
803+
}
804+
744805
/**
745806
* Create an instance of {@link McpAsyncClient} with the provided configurations
746807
* or sensible defaults.
747808
* @return a new instance of {@link McpAsyncClient}.
748809
*/
749810
public McpAsyncClient build() {
811+
var jsonSchemaValidator = (this.jsonSchemaValidator != null) ? this.jsonSchemaValidator
812+
: JsonSchemaValidator.getDefault();
750813
return new McpAsyncClient(this.transport, this.requestTimeout, this.initializationTimeout,
814+
jsonSchemaValidator,
751815
new McpClientFeatures.Async(this.clientInfo, this.capabilities, this.roots,
752816
this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
753817
this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers,
754-
this.samplingHandler, this.elicitationHandler));
818+
this.samplingHandler, this.elicitationHandler, this.enableCallToolSchemaCaching));
755819
}
756820

757821
}

0 commit comments

Comments
 (0)