Skip to content

Commit 7011142

Browse files
committed
feat: Add Tool.outputSchema and CallToolResult.structuredContent
- Implement support for outputSchema in Tool and structuredContent in CallToolResult. - Add a client-side cache that maintains track of all server tools and their output schemas (if any). - Implement a mechanism to refresh the client side cache when encountering a tool call request for an uncached tool. - Implement json validation between tool's output schema and result's structured content, when an output schema is specified for the tool.
1 parent 07e7b8f commit 7011142

File tree

8 files changed

+318
-42
lines changed

8 files changed

+318
-42
lines changed
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
/local/home/pantanur/workspace/java-sdk/mcp-test/src/main/java/io/modelcontextprotocol/server/AbstractMcpSyncServerTests.java

mcp-test/src/main/java/io/modelcontextprotocol/server/AbstractMcpSyncServerTests.java

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
package io.modelcontextprotocol.server;
66

7+
import java.util.HashMap;
78
import java.util.List;
89

910
import io.modelcontextprotocol.spec.McpError;
@@ -17,6 +18,7 @@
1718
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
1819
import io.modelcontextprotocol.spec.McpSchema.Tool;
1920
import io.modelcontextprotocol.spec.McpServerTransportProvider;
21+
2022
import org.junit.jupiter.api.AfterEach;
2123
import org.junit.jupiter.api.BeforeEach;
2224
import org.junit.jupiter.api.Test;
@@ -125,7 +127,7 @@ void testAddTool() {
125127

126128
@Test
127129
void testAddDuplicateTool() {
128-
Tool duplicateTool = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", emptyJsonSchema);
130+
Tool duplicateTool = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", emptyJsonSchema, emptyJsonSchema);
129131

130132
var mcpSyncServer = McpServer.sync(createMcpTransportProvider())
131133
.serverInfo("test-server", "1.0.0")
@@ -134,7 +136,7 @@ void testAddDuplicateTool() {
134136
.build();
135137

136138
assertThatThrownBy(() -> mcpSyncServer.addTool(new McpServerFeatures.SyncToolSpecification(duplicateTool,
137-
(exchange, args) -> new CallToolResult(List.of(), false))))
139+
(exchange, args) -> new CallToolResult(List.of(), false, new HashMap<String, Object>()))))
138140
.isInstanceOf(McpError.class)
139141
.hasMessage("Tool with name '" + TEST_TOOL_NAME + "' already exists");
140142

mcp/pom.xml

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,13 @@
202202
<scope>test</scope>
203203
</dependency>
204204

205+
<!-- Json validator dependency -->
206+
<dependency>
207+
<groupId>com.networknt</groupId>
208+
<artifactId>json-schema-validator</artifactId>
209+
<version>1.5.7</version>
210+
</dependency>
211+
205212
</dependencies>
206213

207214

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

Lines changed: 53 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import java.util.function.Function;
1515

1616
import com.fasterxml.jackson.core.type.TypeReference;
17+
1718
import io.modelcontextprotocol.spec.McpClientSession;
1819
import io.modelcontextprotocol.spec.McpClientSession.NotificationHandler;
1920
import io.modelcontextprotocol.spec.McpClientSession.RequestHandler;
@@ -33,8 +34,11 @@
3334
import io.modelcontextprotocol.spec.McpTransport;
3435
import io.modelcontextprotocol.util.Assert;
3536
import io.modelcontextprotocol.util.Utils;
37+
3638
import org.slf4j.Logger;
3739
import org.slf4j.LoggerFactory;
40+
41+
import io.modelcontextprotocol.server.McpServerFeatures;
3842
import reactor.core.publisher.Flux;
3943
import reactor.core.publisher.Mono;
4044
import reactor.core.publisher.Sinks;
@@ -123,6 +127,11 @@ public class McpAsyncClient {
123127
*/
124128
private McpSchema.Implementation serverInfo;
125129

130+
/**
131+
* Cached tool output schemas.
132+
*/
133+
private final HashMap<String, McpSchema.JsonSchema> toolsOutputSchemaCache;
134+
126135
/**
127136
* Roots define the boundaries of where servers can operate within the filesystem,
128137
* allowing them to understand which directories and files they have access to.
@@ -171,6 +180,7 @@ public class McpAsyncClient {
171180
this.transport = transport;
172181
this.roots = new ConcurrentHashMap<>(features.roots());
173182
this.initializationTimeout = initializationTimeout;
183+
this.toolsOutputSchemaCache = new HashMap<>();
174184

175185
// Request Handlers
176186
Map<String, RequestHandler<?>> requestHandlers = new HashMap<>();
@@ -287,6 +297,14 @@ public McpSchema.Implementation getClientInfo() {
287297
return this.clientInfo;
288298
}
289299

300+
/**
301+
* Get the cached tool output schemas.
302+
* @return The cached tool output schemas
303+
*/
304+
public HashMap<String, McpSchema.JsonSchema> getToolsOutputSchemaCache() {
305+
return this.toolsOutputSchemaCache;
306+
}
307+
290308
/**
291309
* Closes the client connection immediately.
292310
*/
@@ -525,7 +543,13 @@ public Mono<McpSchema.CallToolResult> callTool(McpSchema.CallToolRequest callToo
525543
if (this.serverCapabilities.tools() == null) {
526544
return Mono.error(new McpError("Server does not provide tools capability"));
527545
}
528-
return this.mcpSession.sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF);
546+
// Refresh tool output schema cache, if necessary, prior to making tool call
547+
Mono<Void> refreshCacheMono = Mono.empty();
548+
if (!this.toolsOutputSchemaCache.containsKey(callToolRequest.name())) {
549+
refreshCacheMono = refreshToolOutputSchemaCache();
550+
}
551+
return refreshCacheMono.then(this.mcpSession.sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest,
552+
CALL_TOOL_RESULT_TYPE_REF));
529553
});
530554
}
531555

@@ -547,8 +571,34 @@ public Mono<McpSchema.ListToolsResult> listTools(String cursor) {
547571
if (this.serverCapabilities.tools() == null) {
548572
return Mono.error(new McpError("Server does not provide tools capability"));
549573
}
550-
return this.mcpSession.sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor),
551-
LIST_TOOLS_RESULT_TYPE_REF);
574+
return this.mcpSession
575+
.sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor),
576+
LIST_TOOLS_RESULT_TYPE_REF)
577+
.doOnNext(result -> {
578+
// Cache tools output schema
579+
if (result.tools() != null) {
580+
// Cache tools output schema
581+
result.tools()
582+
.forEach(tool -> this.toolsOutputSchemaCache.put(tool.name(), tool.outputSchema()));
583+
}
584+
});
585+
});
586+
}
587+
588+
/**
589+
* Refreshes the tool output schema cache by fetching all tools from the server.
590+
* @return A Mono that completes when all tool output schemas have been cached
591+
*/
592+
private Mono<Void> refreshToolOutputSchemaCache() {
593+
return this.withInitializationCheck("refreshing tool output schema cache", initializedResult -> {
594+
595+
// Use expand operator to handle pagination in a reactive way
596+
return this.listTools(null).expand(result -> {
597+
if (result.nextCursor() != null) {
598+
return this.listTools(result.nextCursor());
599+
}
600+
return Mono.empty();
601+
}).then();
552602
});
553603
}
554604

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

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

77
import java.time.Duration;
8+
import java.util.HashMap;
9+
import java.util.Set;
810

11+
import org.slf4j.Logger;
12+
import org.slf4j.LoggerFactory;
13+
14+
import com.fasterxml.jackson.core.JsonProcessingException;
15+
import com.fasterxml.jackson.databind.JsonNode;
16+
import com.fasterxml.jackson.databind.ObjectMapper;
17+
import com.fasterxml.jackson.databind.node.ObjectNode;
18+
import com.networknt.schema.JsonSchema;
19+
import com.networknt.schema.JsonSchemaFactory;
20+
import com.networknt.schema.SpecVersion;
21+
import com.networknt.schema.ValidationMessage;
22+
23+
import io.modelcontextprotocol.spec.McpError;
924
import io.modelcontextprotocol.spec.McpSchema;
1025
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
1126
import io.modelcontextprotocol.spec.McpSchema.GetPromptRequest;
1227
import io.modelcontextprotocol.spec.McpSchema.GetPromptResult;
1328
import io.modelcontextprotocol.spec.McpSchema.ListPromptsResult;
1429
import io.modelcontextprotocol.util.Assert;
15-
import org.slf4j.Logger;
16-
import org.slf4j.LoggerFactory;
1730

1831
/**
1932
* A synchronous client implementation for the Model Context Protocol (MCP) that wraps an
@@ -55,6 +68,8 @@ public class McpSyncClient implements AutoCloseable {
5568

5669
private static final Logger logger = LoggerFactory.getLogger(McpSyncClient.class);
5770

71+
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
72+
5873
// TODO: Consider providing a client config to set this properly
5974
// this is currently a concern only because AutoCloseable is used - perhaps it
6075
// is not a requirement?
@@ -206,7 +221,8 @@ public Object ping() {
206221
/**
207222
* Calls a tool provided by the server. Tools enable servers to expose executable
208223
* functionality that can interact with external systems, perform computations, and
209-
* take actions in the real world.
224+
* take actions in the real world. If tool contains an output schema, validates the
225+
* tool result structured content against the output schema.
210226
* @param callToolRequest The request containing: - name: The name of the tool to call
211227
* (must match a tool name from tools/list) - arguments: Arguments that conform to the
212228
* tool's input schema
@@ -215,7 +231,53 @@ public Object ping() {
215231
* Boolean indicating if the execution failed (true) or succeeded (false/absent)
216232
*/
217233
public McpSchema.CallToolResult callTool(McpSchema.CallToolRequest callToolRequest) {
218-
return this.delegate.callTool(callToolRequest).block();
234+
McpSchema.CallToolResult result = this.delegate.callTool(callToolRequest).block();
235+
HashMap<String, McpSchema.JsonSchema> toolsOutputSchemaCache = this.delegate.getToolsOutputSchemaCache();
236+
// Should not be triggered but added for completeness
237+
if (!toolsOutputSchemaCache.containsKey(callToolRequest.name())) {
238+
throw new McpError("Tool with name '" + callToolRequest.name() + "' not found");
239+
}
240+
if (result != null && toolsOutputSchemaCache.get(callToolRequest.name()) != null) {
241+
if (result.structuredContent() == null) {
242+
throw new McpError("CallToolResult validation failed: structuredContent is null and "
243+
+ "does not match tool outputSchema.");
244+
}
245+
246+
McpSchema.JsonSchema outputSchema = toolsOutputSchemaCache.get(callToolRequest.name());
247+
248+
try {
249+
// Convert outputSchema to string
250+
String outputSchemaString = OBJECT_MAPPER.writeValueAsString(outputSchema);
251+
252+
// Create JsonSchema validator
253+
ObjectNode schemaNode = (ObjectNode) OBJECT_MAPPER.readTree(outputSchemaString);
254+
// Set additional properties to false if not specified in output schema
255+
if (!schemaNode.has("additionalProperties")) {
256+
schemaNode.put("additionalProperties", false);
257+
}
258+
JsonSchema schema = JsonSchemaFactory.getInstance(SpecVersion.VersionFlag.V202012)
259+
.getSchema(schemaNode);
260+
261+
// Convert structured content in reult to JsonNode
262+
JsonNode jsonNode = OBJECT_MAPPER.valueToTree(result.structuredContent());
263+
264+
// Validate outputSchema against structuredContent
265+
Set<ValidationMessage> validationResult = schema.validate(jsonNode);
266+
267+
// Check if validation passed
268+
if (!validationResult.isEmpty()) {
269+
// Handle validation errors
270+
throw new McpError(
271+
"CallToolResult validation failed: structuredContent does not match tool outputSchema.");
272+
}
273+
}
274+
catch (JsonProcessingException e) {
275+
// Log error if output schema can't be parsed to prevent erroring out for
276+
// successful call tool request
277+
logger.error("Encountered exception when parsing outputSchema: {}", e);
278+
}
279+
}
280+
return result;
219281
}
220282

221283
/**
@@ -353,4 +415,8 @@ public McpSchema.CompleteResult completeCompletion(McpSchema.CompleteRequest com
353415
return this.delegate.completeCompletion(completeRequest).block();
354416
}
355417

418+
private void isStrict(boolean b) {
419+
throw new UnsupportedOperationException("Not supported yet.");
420+
}
421+
356422
}

0 commit comments

Comments
 (0)