Skip to content

Commit 1911386

Browse files
committed
feat: Add client validation for structuredContent
1 parent 4937fc1 commit 1911386

File tree

3 files changed

+400
-5
lines changed

3 files changed

+400
-5
lines changed

mcp/src/main/java/io/modelcontextprotocol/client/McpClient.java

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,11 @@
1212
import java.util.function.Consumer;
1313
import java.util.function.Function;
1414

15+
import com.fasterxml.jackson.databind.ObjectMapper;
16+
17+
import io.modelcontextprotocol.server.McpServer.AsyncSpecification;
18+
import io.modelcontextprotocol.spec.DefaultJsonSchemaValidator;
19+
import io.modelcontextprotocol.spec.JsonSchemaValidator;
1520
import io.modelcontextprotocol.spec.McpClientTransport;
1621
import io.modelcontextprotocol.spec.McpSchema;
1722
import io.modelcontextprotocol.spec.McpTransport;
@@ -97,6 +102,7 @@
97102
*
98103
* @author Christian Tzolov
99104
* @author Dariusz Jędrzejczyk
105+
* @author Anurag Pant
100106
* @see McpAsyncClient
101107
* @see McpSyncClient
102108
* @see McpTransport
@@ -183,6 +189,8 @@ class SyncSpec {
183189

184190
private Function<ElicitRequest, ElicitResult> elicitationHandler;
185191

192+
private JsonSchemaValidator jsonSchemaValidator;
193+
186194
private SyncSpec(McpClientTransport transport) {
187195
Assert.notNull(transport, "Transport must not be null");
188196
this.transport = transport;
@@ -409,12 +417,27 @@ public SyncSpec progressConsumers(List<Consumer<McpSchema.ProgressNotification>>
409417
return this;
410418
}
411419

420+
/**
421+
* Sets the JSON schema validator to use for validating tool responses against
422+
* output schemas.
423+
* @param jsonSchemaValidator The validator to use. Must not be null.
424+
* @return This builder instance for method chaining
425+
* @throws IllegalArgumentException if jsonSchemaValidator is null
426+
*/
427+
public SyncSpec jsonSchemaValidator(JsonSchemaValidator jsonSchemaValidator) {
428+
Assert.notNull(jsonSchemaValidator, "JsonSchemaValidator must not be null");
429+
this.jsonSchemaValidator = jsonSchemaValidator;
430+
return this;
431+
}
432+
412433
/**
413434
* Create an instance of {@link McpSyncClient} with the provided configurations or
414435
* sensible defaults.
415436
* @return a new instance of {@link McpSyncClient}.
416437
*/
417438
public McpSyncClient build() {
439+
var jsonSchemaValidator = this.jsonSchemaValidator != null ? this.jsonSchemaValidator
440+
: new DefaultJsonSchemaValidator();
418441
McpClientFeatures.Sync syncFeatures = new McpClientFeatures.Sync(this.clientInfo, this.capabilities,
419442
this.roots, this.toolsChangeConsumers, this.resourcesChangeConsumers, this.resourcesUpdateConsumers,
420443
this.promptsChangeConsumers, this.loggingConsumers, this.progressConsumers, this.samplingHandler,
@@ -423,7 +446,8 @@ public McpSyncClient build() {
423446
McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);
424447

425448
return new McpSyncClient(
426-
new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures));
449+
new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures),
450+
jsonSchemaValidator);
427451
}
428452

429453
}

mcp/src/main/java/io/modelcontextprotocol/client/McpSyncClient.java

Lines changed: 63 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,15 @@
55
package io.modelcontextprotocol.client;
66

77
import java.time.Duration;
8+
import java.util.Map;
9+
import java.util.Optional;
10+
import java.util.concurrent.ConcurrentHashMap;
811

912
import org.slf4j.Logger;
1013
import org.slf4j.LoggerFactory;
1114

15+
import io.modelcontextprotocol.spec.JsonSchemaValidator;
16+
import io.modelcontextprotocol.spec.McpError;
1217
import io.modelcontextprotocol.spec.McpSchema;
1318
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
1419
import io.modelcontextprotocol.spec.McpSchema.GetPromptRequest;
@@ -48,6 +53,7 @@
4853
* @author Dariusz Jędrzejczyk
4954
* @author Christian Tzolov
5055
* @author Jihoon Kim
56+
* @author Anurag Pant
5157
* @see McpClient
5258
* @see McpAsyncClient
5359
* @see McpSchema
@@ -63,14 +69,23 @@ public class McpSyncClient implements AutoCloseable {
6369

6470
private final McpAsyncClient delegate;
6571

72+
private final JsonSchemaValidator jsonSchemaValidator;
73+
74+
/**
75+
* Cached tool output schemas.
76+
*/
77+
private final ConcurrentHashMap<String, Optional<Map<String, Object>>> toolsOutputSchemaCache;
78+
6679
/**
6780
* Create a new McpSyncClient with the given delegate.
6881
* @param delegate the asynchronous kernel on top of which this synchronous client
6982
* provides a blocking API.
7083
*/
71-
McpSyncClient(McpAsyncClient delegate) {
84+
McpSyncClient(McpAsyncClient delegate, JsonSchemaValidator jsonSchemaValidator) {
7285
Assert.notNull(delegate, "The delegate can not be null");
7386
this.delegate = delegate;
87+
this.jsonSchemaValidator = jsonSchemaValidator;
88+
this.toolsOutputSchemaCache = new ConcurrentHashMap<>();
7489
}
7590

7691
/**
@@ -216,7 +231,37 @@ public Object ping() {
216231
* Boolean indicating if the execution failed (true) or succeeded (false/absent)
217232
*/
218233
public McpSchema.CallToolResult callTool(McpSchema.CallToolRequest callToolRequest) {
219-
return this.delegate.callTool(callToolRequest).block();
234+
if (!this.toolsOutputSchemaCache.containsKey(callToolRequest.name())) {
235+
listTools(); // Ensure tools are cached before calling
236+
}
237+
238+
McpSchema.CallToolResult result = this.delegate.callTool(callToolRequest).block();
239+
Optional<Map<String, Object>> optOutputSchema = toolsOutputSchemaCache.get(callToolRequest.name());
240+
241+
if (result != null && result.isError() != null && !result.isError()) {
242+
if (optOutputSchema == null) {
243+
// Should not be triggered but added for completeness
244+
throw new McpError("Tool with name '" + callToolRequest.name() + "' not found");
245+
}
246+
else {
247+
if (optOutputSchema.isPresent()) {
248+
// Validate the tool output against the cached output schema
249+
var validation = this.jsonSchemaValidator.validate(optOutputSchema.get(),
250+
result.structuredContent());
251+
if (!validation.valid()) {
252+
throw new McpError("Tool call result validation failed: " + validation.errorMessage());
253+
}
254+
}
255+
else if (result.structuredContent() != null) {
256+
logger.warn(
257+
"Calling a tool with no outputSchema is not expected to return result with structured content, but got: {}",
258+
result.structuredContent());
259+
}
260+
261+
}
262+
}
263+
264+
return result;
220265
}
221266

222267
/**
@@ -226,7 +271,14 @@ public McpSchema.CallToolResult callTool(McpSchema.CallToolRequest callToolReque
226271
* pagination if more tools are available
227272
*/
228273
public McpSchema.ListToolsResult listTools() {
229-
return this.delegate.listTools().block();
274+
return this.delegate.listTools().doOnNext(result -> {
275+
if (result.tools() != null) {
276+
// Cache tools output schema
277+
result.tools()
278+
.forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(),
279+
Optional.ofNullable(tool.outputSchema())));
280+
}
281+
}).block();
230282
}
231283

232284
/**
@@ -237,7 +289,14 @@ public McpSchema.ListToolsResult listTools() {
237289
* pagination if more tools are available
238290
*/
239291
public McpSchema.ListToolsResult listTools(String cursor) {
240-
return this.delegate.listTools(cursor).block();
292+
return this.delegate.listTools(cursor).doOnNext(result -> {
293+
if (result.tools() != null) {
294+
// Cache tools output schema
295+
result.tools()
296+
.forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(),
297+
Optional.ofNullable(tool.outputSchema())));
298+
}
299+
}).block();
241300
}
242301

243302
// --------------------------

0 commit comments

Comments
 (0)