Skip to content
This repository was archived by the owner on Feb 14, 2025. It is now read-only.

Commit aa26cfa

Browse files
tzolovtzolov
committed
feat: Improve Tool schema handling
- Make McpSchema final and add private constructor - Add JsonSchema record for structured schema representation - Change Tool.inputSchema type from Map<String, Object> to JsonSchema. - Use the String JSON Schema constructor when creating the Tool programaticaly. - Enhance StdioClientTransport error logging with more context - Add ToolHelper utility class for: - Converting Spring AI FunctionCallbacks to MCP tools - Generating JSON schemas for tool input validation - Facilitating integration between Spring AI functions and MCP tools - Update tests and samples to use the new Tool schema format Resolves #52
1 parent 8334e2e commit aa26cfa

File tree

13 files changed

+392
-51
lines changed

13 files changed

+392
-51
lines changed

‎mcp/src/main/java/org/springframework/ai/mcp/client/transport/StdioClientTransport.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -275,14 +275,14 @@ private void startInboundProcessing() {
275275
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line);
276276
if (!this.inboundSink.tryEmitNext(message).isSuccess()) {
277277
if (!isClosing) {
278-
logger.error("Failed to enqueue inbound message");
278+
logger.error("Failed to enqueue inbound message: {}", message);
279279
}
280280
break;
281281
}
282282
}
283283
catch (Exception e) {
284284
if (!isClosing) {
285-
logger.error("Error processing inbound message", e);
285+
logger.error("Error processing inbound message for line: " + line, e);
286286
}
287287
break;
288288
}

‎mcp/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java‎

Lines changed: 30 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,10 @@
3838
*
3939
* @author Christian Tzolov
4040
*/
41-
public class McpSchema {
41+
public final class McpSchema {
42+
43+
private McpSchema() {
44+
}
4245

4346
public static final String LATEST_PROTOCOL_VERSION = "2024-11-05";
4447

@@ -95,6 +98,8 @@ public class McpSchema {
9598
// Sampling Methods
9699
public static final String METHOD_SAMPLING_CREATE_MESSAGE = "sampling/createMessage";
97100

101+
private static ObjectMapper OBJECT_MAPPER = new ObjectMapper();
102+
98103
// ---------------------------
99104
// JSON-RPC Error Codes
100105
// ---------------------------
@@ -659,6 +664,15 @@ public record ListToolsResult( // @formatter:off
659664
@JsonProperty("nextCursor") String nextCursor) {
660665
}// @formatter:on
661666

667+
@JsonInclude(JsonInclude.Include.NON_ABSENT)
668+
@JsonIgnoreProperties(ignoreUnknown = true)
669+
record JsonSchema( // @formatter:off
670+
@JsonProperty("type") String type,
671+
@JsonProperty("properties") Map<String, Object> properties,
672+
@JsonProperty("required") List<String> required,
673+
@JsonProperty("additionalProperties") Boolean additionalProperties) {
674+
} // @formatter:on
675+
662676
/**
663677
* Represents a tool that the server provides. Tools enable servers to expose
664678
* executable functionality to the system. Through these tools, you can interact with
@@ -676,17 +690,17 @@ public record ListToolsResult( // @formatter:off
676690
public record Tool( // @formatter:off
677691
@JsonProperty("name") String name,
678692
@JsonProperty("description") String description,
679-
@JsonProperty("inputSchema") Map<String, Object> inputSchema) {
693+
@JsonProperty("inputSchema") JsonSchema inputSchema) {
680694

681695
public Tool(String name, String description, String schema) {
682696
this(name, description, parseSchema(schema));
683697
}
684698

685699
} // @formatter:on
686700

687-
public static Map<String, Object> parseSchema(String schema) {
701+
public static JsonSchema parseSchema(String schema) {
688702
try {
689-
return new ObjectMapper().readValue(schema, MAP_TYPE_REF);
703+
return OBJECT_MAPPER.readValue(schema, JsonSchema.class);
690704
}
691705
catch (IOException e) {
692706
throw new IllegalArgumentException("Invalid schema: " + schema, e);
@@ -724,6 +738,18 @@ public record CallToolResult( // @formatter:off
724738
// ---------------------------
725739
// Sampling Interfaces
726740
// ---------------------------
741+
@JsonInclude(JsonInclude.Include.NON_ABSENT)
742+
public record ModelPreferences(// @formatter:off
743+
@JsonProperty("hints") List<ModelHint> hints,
744+
@JsonProperty("costPriority") Double costPriority,
745+
@JsonProperty("speedPriority") Double speedPriority,
746+
@JsonProperty("intelligencePriority") Double intelligencePriority) {
747+
} // @formatter:on
748+
749+
@JsonInclude(JsonInclude.Include.NON_ABSENT)
750+
public record ModelHint(@JsonProperty("name") String name) {
751+
}
752+
727753
@JsonInclude(JsonInclude.Include.NON_ABSENT)
728754
public record SamplingMessage(// @formatter:off
729755
@JsonProperty("role") Role role,
@@ -803,18 +829,6 @@ public CreateMessageResult build() {
803829
}
804830
}// @formatter:on
805831

806-
// ---------------------------
807-
// Model Preferences
808-
// ---------------------------
809-
@JsonInclude(JsonInclude.Include.NON_ABSENT)
810-
public record ModelPreferences(List<ModelHint> hints, Double costPriority, Double speedPriority,
811-
Double intelligencePriority) {
812-
}
813-
814-
@JsonInclude(JsonInclude.Include.NON_ABSENT)
815-
public record ModelHint(String name) {
816-
}
817-
818832
// ---------------------------
819833
// Pagination Interfaces
820834
// ---------------------------

‎mcp/src/test/java/org/springframework/ai/mcp/attic/DemoServer.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ public static void main(String[] args) {
4444

4545
var mcpServer = McpServer.using(transport)
4646
.serverInfo("Weather Forecast", "1.0.0")
47-
.tool(new McpSchema.Tool("weather", "Weather forecast tool by location", Map.of("city", "String")),
47+
.tool(new McpSchema.Tool("weather", "Weather forecast tool by location", "{ \"type\": \"object\" }"),
4848
(arguments) -> {
4949
String city = (String) arguments.get("city");
5050
return new CallToolResult(List.of(), false);

‎mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,9 @@
2323
import java.util.function.Consumer;
2424
import java.util.function.Function;
2525

26+
import com.fasterxml.jackson.core.JsonProcessingException;
2627
import com.fasterxml.jackson.core.type.TypeReference;
28+
import com.fasterxml.jackson.databind.ObjectMapper;
2729
import org.junit.jupiter.api.Test;
2830

2931
import org.springframework.ai.mcp.MockMcpTransport;
@@ -39,7 +41,7 @@
3941
class McpAsyncClientResponseHandlerTests {
4042

4143
@Test
42-
void testToolsChangeNotificationHandling() {
44+
void testToolsChangeNotificationHandling() throws JsonProcessingException {
4345
MockMcpTransport transport = new MockMcpTransport();
4446

4547
// Create a list to store received tools for verification
@@ -55,7 +57,8 @@ void testToolsChangeNotificationHandling() {
5557

5658
// Create a mock tools list that the server will return
5759
Map<String, Object> inputSchema = Map.of("type", "object", "properties", Map.of(), "required", List.of());
58-
McpSchema.Tool mockTool = new McpSchema.Tool("test-tool", "Test Tool Description", inputSchema);
60+
McpSchema.Tool mockTool = new McpSchema.Tool("test-tool", "Test Tool Description",
61+
new ObjectMapper().writeValueAsString(inputSchema));
5962
McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(mockTool), null);
6063

6164
// Simulate server sending tools/list_changed notification

‎mcp/src/test/java/org/springframework/ai/mcp/server/AbstractMcpAsyncServerTests.java‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -107,10 +107,17 @@ void testImmediateClose() {
107107
// ---------------------------------------
108108
// Tools Tests
109109
// ---------------------------------------
110+
String emptyJsonSchema = """
111+
{
112+
"$schema": "http://json-schema.org/draft-07/schema#",
113+
"type": "object",
114+
"properties": {}
115+
}
116+
""";
110117

111118
@Test
112119
void testAddTool() {
113-
Tool newTool = new McpSchema.Tool("new-tool", "New test tool", Map.of("input", "string"));
120+
Tool newTool = new McpSchema.Tool("new-tool", "New test tool", emptyJsonSchema);
114121
var mcpAsyncServer = McpServer.using(createMcpTransport())
115122
.serverInfo("test-server", "1.0.0")
116123
.capabilities(ServerCapabilities.builder().tools(true).build())
@@ -125,7 +132,7 @@ void testAddTool() {
125132

126133
@Test
127134
void testAddDuplicateTool() {
128-
Tool duplicateTool = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", Map.of("input", "string"));
135+
Tool duplicateTool = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", emptyJsonSchema);
129136

130137
var mcpAsyncServer = McpServer.using(createMcpTransport())
131138
.serverInfo("test-server", "1.0.0")
@@ -146,7 +153,7 @@ void testAddDuplicateTool() {
146153

147154
@Test
148155
void testRemoveTool() {
149-
Tool too = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", Map.of("input", "string"));
156+
Tool too = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", emptyJsonSchema);
150157

151158
var mcpAsyncServer = McpServer.using(createMcpTransport())
152159
.serverInfo("test-server", "1.0.0")
@@ -175,7 +182,7 @@ void testRemoveNonexistentTool() {
175182

176183
@Test
177184
void testNotifyToolsListChanged() {
178-
Tool too = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", Map.of("input", "string"));
185+
Tool too = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", emptyJsonSchema);
179186

180187
var mcpAsyncServer = McpServer.using(createMcpTransport())
181188
.serverInfo("test-server", "1.0.0")

‎mcp/src/test/java/org/springframework/ai/mcp/server/AbstractMcpSyncServerTests.java‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -116,14 +116,22 @@ void testGetAsyncServer() {
116116
// Tools Tests
117117
// ---------------------------------------
118118

119+
String emptyJsonSchema = """
120+
{
121+
"$schema": "http://json-schema.org/draft-07/schema#",
122+
"type": "object",
123+
"properties": {}
124+
}
125+
""";
126+
119127
@Test
120128
void testAddTool() {
121129
var mcpSyncServer = McpServer.using(createMcpTransport())
122130
.serverInfo("test-server", "1.0.0")
123131
.capabilities(ServerCapabilities.builder().tools(true).build())
124132
.sync();
125133

126-
Tool newTool = new McpSchema.Tool("new-tool", "New test tool", Map.of("input", "string"));
134+
Tool newTool = new McpSchema.Tool("new-tool", "New test tool", emptyJsonSchema);
127135
assertThatCode(() -> mcpSyncServer
128136
.addTool(new ToolRegistration(newTool, args -> new CallToolResult(List.of(), false))))
129137
.doesNotThrowAnyException();
@@ -133,7 +141,7 @@ void testAddTool() {
133141

134142
@Test
135143
void testAddDuplicateTool() {
136-
Tool duplicateTool = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", Map.of("input", "string"));
144+
Tool duplicateTool = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", emptyJsonSchema);
137145

138146
var mcpSyncServer = McpServer.using(createMcpTransport())
139147
.serverInfo("test-server", "1.0.0")
@@ -151,7 +159,7 @@ void testAddDuplicateTool() {
151159

152160
@Test
153161
void testRemoveTool() {
154-
Tool tool = new McpSchema.Tool(TEST_TOOL_NAME, "Test tool", Map.of("input", "string"));
162+
Tool tool = new McpSchema.Tool(TEST_TOOL_NAME, "Test tool", emptyJsonSchema);
155163

156164
var mcpSyncServer = McpServer.using(createMcpTransport())
157165
.serverInfo("test-server", "1.0.0")

‎mcp/src/test/java/org/springframework/ai/mcp/server/SseAsyncIntegrationTests.java‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -322,12 +322,22 @@ void testRootsServerCloseWithActiveSubscription() {
322322
// ---------------------------------------
323323
// Tools Tests
324324
// ---------------------------------------
325+
326+
String emptyJsonSchema = """
327+
{
328+
"$schema": "http://json-schema.org/draft-07/schema#",
329+
"type": "object",
330+
"properties": {}
331+
}
332+
""";
333+
325334
@Test
335+
326336
void testToolCallSuccess() {
327337

328338
var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null);
329-
ToolRegistration tool1 = new ToolRegistration(
330-
new McpSchema.Tool("tool1", "tool1 description", Map.of("city", "String")), request -> {
339+
ToolRegistration tool1 = new ToolRegistration(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema),
340+
request -> {
331341
// perform a blocking call to a remote service
332342
String response = RestClient.create()
333343
.get()
@@ -363,8 +373,8 @@ void testToolCallSuccess() {
363373
void testToolListChangeHandlingSuccess() {
364374

365375
var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null);
366-
ToolRegistration tool1 = new ToolRegistration(
367-
new McpSchema.Tool("tool1", "tool1 description", Map.of("city", "String")), request -> {
376+
ToolRegistration tool1 = new ToolRegistration(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema),
377+
request -> {
368378
// perform a blocking call to a remote service
369379
String response = RestClient.create()
370380
.get()
@@ -413,8 +423,8 @@ void testToolListChangeHandlingSuccess() {
413423
});
414424

415425
// Add a new tool
416-
ToolRegistration tool2 = new ToolRegistration(
417-
new McpSchema.Tool("tool2", "tool2 description", Map.of("city", "String")), request -> callResponse);
426+
ToolRegistration tool2 = new ToolRegistration(new McpSchema.Tool("tool2", "tool2 description", emptyJsonSchema),
427+
request -> callResponse);
418428

419429
mcpServer.addTool(tool2);
420430

‎spring-ai-mcp-sample/pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
<dependencies>
2121
<dependency>
2222
<groupId>org.springframework.experimental</groupId>
23-
<artifactId>mcp</artifactId>
23+
<artifactId>spring-ai-mcp</artifactId>
2424
<version>${project.version}</version>
2525
</dependency>
2626

‎spring-ai-mcp-sample/src/main/java/org/springframework/ai/mcp/sample/client/SampleClient.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,13 @@ public void run() {
5858
.callTool(new CallToolRequest("calculator", Map.of("operation", "multiply", "a", 2.0, "b", 3.0)));
5959
System.out.println("Calculator Response = " + calcResponse);
6060

61+
CallToolResult paymentStatus = client.callTool(
62+
new CallToolRequest("paymentTransactionStatus", Map.of("transactionId", "006", "accountName", "John")));
63+
System.out.println("Payment Status Response = " + paymentStatus);
64+
65+
CallToolResult parks = client.callTool(new CallToolRequest("getBooks", Map.of("title", "Spring Framework")));
66+
System.out.println("Books Response = " + parks);
67+
6168
// List and demonstrate resources
6269
var resourcesList = client.listResources();
6370
System.out.println("\nAvailable Resources = " + resourcesList);

0 commit comments

Comments
 (0)