diff --git a/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java b/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java index e7c07f4..ccd32c4 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java @@ -129,6 +129,52 @@ public class McpAsyncServer { notificationHandlers); } + // --------------------------------------- + // Lifecycle Management + // --------------------------------------- + private DefaultMcpSession.RequestHandler initializeRequestHandler() { + return params -> { + McpSchema.InitializeRequest initializeRequest = transport.unmarshalFrom(params, + new TypeReference() { + }); + + logger.info("Client initialize request - Protocol: {}, Capabilities: {}, Info: {}", + initializeRequest.protocolVersion(), initializeRequest.capabilities(), + initializeRequest.clientInfo()); + + if (!McpSchema.LATEST_PROTOCOL_VERSION.equals(initializeRequest.protocolVersion())) { + return Mono + .error(new McpError( + "Unsupported protocol version from client: " + initializeRequest.protocolVersion())) + .publishOn(Schedulers.boundedElastic()); + } + + return Mono + .just(new McpSchema.InitializeResult(McpSchema.LATEST_PROTOCOL_VERSION, this.serverCapabilities, + this.serverInfo, null)) + .publishOn(Schedulers.boundedElastic()); + }; + } + + /** + * Gracefully closes the server, allowing any in-progress operations to complete. + * @return A Mono that completes when the server has been closed + */ + public Mono closeGracefully() { + return this.mcpSession.closeGracefully(); + } + + /** + * Close the server immediately. + */ + public void close() { + this.mcpSession.close(); + } + + // --------------------------------------- + // Tool Management + // --------------------------------------- + /** * Add a new tool registration at runtime. * @param toolRegistration The tool registration to add @@ -185,6 +231,53 @@ public class McpAsyncServer { return Mono.error(new McpError("Tool with name '" + toolName + "' not found")); } + /** + * Notifies clients that the list of available tools has changed. + * @return A Mono that completes when all clients have been notified + */ + public Mono notifyToolsListChanged() { + return this.mcpSession.sendNotification("notifications/tools/list_changed", null); + } + + private DefaultMcpSession.RequestHandler toolsListRequestHandler() { + return params -> { + McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, + new TypeReference() { + }); + + List tools = this.tools.stream().map(toolRegistration -> { + return toolRegistration.tool(); + }).toList(); + + logger.info("Client tools list request - Cursor: {}", request.cursor()); + return Mono.just(new McpSchema.ListToolsResult(tools, null)); + }; + } + + private DefaultMcpSession.RequestHandler toolsCallRequestHandler() { + return params -> { + McpSchema.CallToolRequest callToolRequest = transport.unmarshalFrom(params, + new TypeReference() { + }); + + Optional toolRegistration = this.tools.stream() + .filter(tr -> callToolRequest.name().equals(tr.tool().name())) + .findAny(); + + if (toolRegistration.isEmpty()) { + return Mono.error(new McpError("Tool not found: " + callToolRequest.name())); + } + + CallToolResult callResponse = toolRegistration.get().call().apply(callToolRequest.arguments()); + + return Mono.just(callResponse); + }; + } + + // --------------------------------------- + // Resource Management + // --------------------------------------- + /** * Add a new resource handler at runtime. * @param resourceHandler The resource handler to add @@ -236,6 +329,43 @@ public class McpAsyncServer { return Mono.error(new McpError("Resource with URI '" + resourceUri + "' not found")); } + /** + * Notifies clients that the list of available resources has changed. + * @return A Mono that completes when all clients have been notified + */ + public Mono notifyResourcesListChanged() { + return this.mcpSession.sendNotification("notifications/resources/list_changed", null); + } + + private DefaultMcpSession.RequestHandler resourcesListRequestHandler() { + return params -> { + var resourceList = this.resources.values().stream().map(ResourceRegistration::resource).toList(); + return Mono.just(new McpSchema.ListResourcesResult(resourceList, null)); + }; + } + + private DefaultMcpSession.RequestHandler resourceTemplateListRequestHandler() { + return params -> Mono.just(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null)); + + } + + private DefaultMcpSession.RequestHandler resourcesReadRequestHandler() { + return params -> { + McpSchema.ReadResourceRequest resourceRequest = transport.unmarshalFrom(params, + new TypeReference() { + }); + var resourceUri = resourceRequest.uri(); + if (this.resources.containsKey(resourceUri)) { + return Mono.just(this.resources.get(resourceUri).readHandler().apply(resourceRequest)); + } + return Mono.error(new McpError("Resource not found: " + resourceUri)); + }; + } + + // --------------------------------------- + // Prompt Management + // --------------------------------------- + /** * Add a new prompt handler at runtime. * @param promptRegistration The prompt handler to add @@ -255,7 +385,11 @@ public class McpAsyncServer { } this.prompts.put(promptRegistration.propmpt().name(), promptRegistration); + logger.info("Added prompt handler: {}", promptRegistration.propmpt().name()); + + // Servers that declared the listChanged capability SHOULD send a notification, + // when the list of available prompts changes if (this.serverCapabilities.prompts().listChanged()) { return notifyPromptsListChanged(); } @@ -276,8 +410,11 @@ public class McpAsyncServer { } PromptRegistration removed = this.prompts.remove(promptName); + if (removed != null) { logger.info("Removed prompt handler: {}", promptName); + // Servers that declared the listChanged capability SHOULD send a + // notification, when the list of available prompts changes if (this.serverCapabilities.prompts().listChanged()) { return this.notifyPromptsListChanged(); } @@ -286,100 +423,12 @@ public class McpAsyncServer { return Mono.error(new McpError("Prompt with name '" + promptName + "' not found")); } - // --------------------------------------- - // Request Handlers - // --------------------------------------- - private DefaultMcpSession.RequestHandler initializeRequestHandler() { - return params -> { - McpSchema.InitializeRequest request = transport.unmarshalFrom(params, - new TypeReference() { - }); - - logger.info("Client initialize request - Protocol: {}, Capabilities: {}, Info: {}", - request.protocolVersion(), request.capabilities(), request.clientInfo()); - - if (!McpSchema.LATEST_PROTOCOL_VERSION.equals(request.protocolVersion())) { - return Mono - .error( - new McpError("Unsupported protocol version from client: " + request.protocolVersion())) - .publishOn(Schedulers.boundedElastic()); - } - - return Mono - .just(new McpSchema.InitializeResult(McpSchema.LATEST_PROTOCOL_VERSION, this.serverCapabilities, - this.serverInfo, null)) - .publishOn(Schedulers.boundedElastic()); - }; - } - - private DefaultMcpSession.RequestHandler toolsListRequestHandler() { - return params -> { - McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, - new TypeReference() { - }); - - List tools = this.tools.stream().map(toolRegistration -> { - return toolRegistration.tool(); - }).toList(); - - logger.info("Client tools list request - Cursor: {}", request.cursor()); - return Mono.just(new McpSchema.ListToolsResult(tools, null)); - }; - } - - private DefaultMcpSession.RequestHandler toolsCallRequestHandler() { - return params -> { - McpSchema.CallToolRequest callToolRequest = transport.unmarshalFrom(params, - new TypeReference() { - }); - - Optional toolRegistration = this.tools.stream() - .filter(tr -> callToolRequest.name().equals(tr.tool().name())) - .findAny(); - - if (toolRegistration.isEmpty()) { - return Mono.error(new McpError("Tool not found: " + callToolRequest.name())); - } - - CallToolResult callResponse = toolRegistration.get().call().apply(callToolRequest.arguments()); - - return Mono.just(callResponse); - }; - } - - private DefaultMcpSession.RequestHandler resourcesListRequestHandler() { - return params -> { - // McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, - // new TypeReference() { - // }); - - var resourceList = this.resources.values().stream().map(ResourceRegistration::resource).toList(); - - return Mono.just(new McpSchema.ListResourcesResult(resourceList, null)); - }; - } - - private DefaultMcpSession.RequestHandler resourceTemplateListRequestHandler() { - return params -> { - // McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, - // new TypeReference() { - // }); - - return Mono.just(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null)); - }; - } - - private DefaultMcpSession.RequestHandler resourcesReadRequestHandler() { - return params -> { - McpSchema.ReadResourceRequest request = transport.unmarshalFrom(params, - new TypeReference() { - }); - var resourceUri = request.uri(); - if (this.resources.containsKey(resourceUri)) { - return Mono.just(this.resources.get(resourceUri).readHandler().apply(request)); - } - return Mono.error(new McpError("Resource not found: " + resourceUri)); - }; + /** + * Notifies clients that the list of available prompts has changed. + * @return A Mono that completes when all clients have been notified + */ + public Mono notifyPromptsListChanged() { + return this.mcpSession.sendNotification("notifications/prompts/list_changed", null); } private DefaultMcpSession.RequestHandler promptsListRequestHandler() { @@ -396,56 +445,17 @@ public class McpAsyncServer { private DefaultMcpSession.RequestHandler promptsGetRequestHandler() { return params -> { - McpSchema.GetPromptRequest request = transport.unmarshalFrom(params, + McpSchema.GetPromptRequest promptRequest = transport.unmarshalFrom(params, new TypeReference() { }); // Implement prompt retrieval logic here - if (this.prompts.containsKey(request.name())) { - return Mono.just(this.prompts.get(request.name()).promptHandler().apply(request)); + if (this.prompts.containsKey(promptRequest.name())) { + return Mono.just(this.prompts.get(promptRequest.name()).promptHandler().apply(promptRequest)); } - return Mono.error(new McpError("Prompt not found: " + request.name())); + return Mono.error(new McpError("Prompt not found: " + promptRequest.name())); }; } - /** - * Notifies clients that the list of available tools has changed. - * @return A Mono that completes when all clients have been notified - */ - public Mono notifyToolsListChanged() { - return this.mcpSession.sendNotification("notifications/tools/list_changed", null); - } - - /** - * Notifies clients that the list of available resources has changed. - * @return A Mono that completes when all clients have been notified - */ - public Mono notifyResourcesListChanged() { - return this.mcpSession.sendNotification("notifications/resources/list_changed", null); - } - - /** - * Notifies clients that the list of available prompts has changed. - * @return A Mono that completes when all clients have been notified - */ - public Mono notifyPromptsListChanged() { - return this.mcpSession.sendNotification("notifications/prompts/list_changed", null); - } - - /** - * Gracefully closes the server, allowing any in-progress operations to complete. - * @return A Mono that completes when the server has been closed - */ - public Mono closeGracefully() { - return this.mcpSession.closeGracefully(); - } - - /** - * Close the server immediately. - */ - public void close() { - this.mcpSession.close(); - } - } diff --git a/mcp/src/main/java/org/springframework/ai/mcp/server/McpServer.java b/mcp/src/main/java/org/springframework/ai/mcp/server/McpServer.java index 22d6840..a92879a 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/server/McpServer.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/server/McpServer.java @@ -392,7 +392,7 @@ public interface McpServer { public Builder resources(List resourceRegsitrations) { Assert.notNull(resourceRegsitrations, "Resource handlers list must not be null"); for (ResourceRegistration resource : resourceRegsitrations) { - this.resources.put(resource.resource().name(), resource); + this.resources.put(resource.resource().uri(), resource); } return this; } @@ -417,7 +417,7 @@ public interface McpServer { public Builder resources(ResourceRegistration... resourceRegistrations) { Assert.notNull(resourceRegistrations, "Resource handlers list must not be null"); for (ResourceRegistration resource : resourceRegistrations) { - this.resources.put(resource.resource().name(), resource); + this.resources.put(resource.resource().uri(), resource); } return this; } diff --git a/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java b/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java index 675737f..c568233 100644 --- a/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java +++ b/mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java @@ -153,31 +153,23 @@ public class StdioServerTransport implements McpTransport { try { JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line); if (!this.inboundSink.tryEmitNext(message).isSuccess()) { - if (!isClosing) { - logger.error("Failed to enqueue message"); - } + logIfNotClosing("Failed to enqueue message"); break; } } catch (Exception e) { - if (!isClosing) { - logger.error("Error processing inbound message", e); - } + logIfNotClosing("Error processing inbound message", e); break; } } catch (IOException e) { - if (!isClosing) { - logger.error("Error reading from stdin", e); - } + logIfNotClosing("Error reading from stdin", e); break; } } } catch (Exception e) { - if (!isClosing) { - logger.error("Error in inbound processing", e); - } + logIfNotClosing("Error in inbound processing", e); } finally { isClosing = true; @@ -254,7 +246,6 @@ public class StdioServerTransport implements McpTransport { // Dispose schedulers with longer timeout inboundScheduler.dispose(); outboundScheduler.dispose(); - logger.info("Graceful shutdown completed"); } catch (Exception e) { @@ -268,4 +259,16 @@ public class StdioServerTransport implements McpTransport { return this.objectMapper.convertValue(data, typeRef); } + private void logIfNotClosing(String message, Exception e) { + if (!this.isClosing) { + logger.error(message, e); + } + } + + private void logIfNotClosing(String message) { + if (!this.isClosing) { + logger.error(message); + } + } + } diff --git a/spring-ai-mcp-sample/src/main/java/org/springframework/ai/mcp/sample/client/ClientStdio.java b/spring-ai-mcp-sample/src/main/java/org/springframework/ai/mcp/sample/client/ClientStdio.java index 7777359..ee84037 100644 --- a/spring-ai-mcp-sample/src/main/java/org/springframework/ai/mcp/sample/client/ClientStdio.java +++ b/spring-ai-mcp-sample/src/main/java/org/springframework/ai/mcp/sample/client/ClientStdio.java @@ -19,9 +19,9 @@ import org.springframework.ai.mcp.client.transport.ServerParameters; import org.springframework.ai.mcp.client.transport.StdioClientTransport; /** - * @author Christian Tzolov + * With stdio transport, the MCP server is automatically started by the client. But you + * have to build the server jar first: 
./mvnw clean install -DskipTests
*/ - public class ClientStdio { public static void main(String[] args) { diff --git a/spring-ai-mcp-sample/src/main/java/org/springframework/ai/mcp/sample/server/McpServerConfig.java b/spring-ai-mcp-sample/src/main/java/org/springframework/ai/mcp/sample/server/McpServerConfig.java index 968b122..06d2dc3 100644 --- a/spring-ai-mcp-sample/src/main/java/org/springframework/ai/mcp/sample/server/McpServerConfig.java +++ b/spring-ai-mcp-sample/src/main/java/org/springframework/ai/mcp/sample/server/McpServerConfig.java @@ -27,18 +27,21 @@ import org.springframework.web.reactive.function.server.RouterFunction; @Configuration public class McpServerConfig { + // STDIO transport @Bean @ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "stdio") public StdioServerTransport stdioServerTransport() { return new StdioServerTransport(); } + // SSE transport @Bean @ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "sse") public SseServerTransport sseServerTransport() { return new SseServerTransport(new ObjectMapper(), "/mcp/message"); } + // Router function for SSE transport used by Spring WebFlux to start an HTTP server. @Bean @ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "sse") public RouterFunction mcpRouterFunction(SseServerTransport transport) { @@ -59,13 +62,13 @@ public class McpServerConfig { return McpServer.using(transport) .info("MCP Demo Server", "1.0.0") .capabilities(capabilities) - .tools(List.of(weatherTool(), calculatorTool())) - .resources(Map.of("system://info", systemInfoResource())) - .prompts(List.of(greetingPrompt())) + .resources(systemInfoResourceRegistration()) + .prompts(greetingPromptRegistration()) + .tools(weatherToolRegistration(), calculatorToolRegistration()) .async(); } // @formatter:on - private static ResourceRegistration systemInfoResource() { + private static ResourceRegistration systemInfoResourceRegistration() { // Create a resource registration for system information var systemInfoResource = new McpSchema.Resource( // @formatter:off @@ -98,7 +101,26 @@ public class McpServerConfig { return resourceRegistration; } - private static ToolRegistration weatherTool() { + private static PromptRegistration greetingPromptRegistration() { + + var prompt = new McpSchema.Prompt("greeting", "A friendly greeting prompt", + List.of(new McpSchema.PromptArgument("name", "The name to greet", true))); + + return new PromptRegistration(prompt, getPromptRequest -> { + + String nameArgument = (String) getPromptRequest.arguments().get("name"); + if (nameArgument == null) { + nameArgument = "friend"; + } + + var userMessage = new PromptMessage(Role.USER, + new TextContent("Hello " + nameArgument + "! How can I assist you today?")); + + return new GetPromptResult("A personalized greeting message", List.of(userMessage)); + }); + } + + private static ToolRegistration weatherToolRegistration() { return new ToolRegistration( new McpSchema.Tool("weather", "Weather forecast tool by location", Map.of("city", "String")), (arguments) -> { @@ -108,7 +130,7 @@ public class McpServerConfig { }); } - private static ToolRegistration calculatorTool() { + private static ToolRegistration calculatorToolRegistration() { return new ToolRegistration(new McpSchema.Tool("calculator", "Performs basic arithmetic operations (add, subtract, multiply, divide)", """ { @@ -164,24 +186,4 @@ public class McpServerConfig { }); } - private static PromptRegistration greetingPrompt() { - var nameArg = new McpSchema.PromptArgument("name", "The name to greet", true // required - // parameter - ); - - var prompt = new McpSchema.Prompt("greeting", "A friendly greeting prompt", List.of(nameArg)); - - return new PromptRegistration(prompt, request -> { - String name = (String) request.arguments().get("name"); - if (name == null) { - name = "friend"; - } - - var assistantMessage = new PromptMessage(Role.ASSISTANT, - new TextContent("Hello " + name + "! How can I assist you today?")); - - return new GetPromptResult("A personalized greeting message", List.of(assistantMessage)); - }); - } - }