feat: Add MCP server with Server stdio/sse transport improve MCP client
- Implement MCP server with async/sync support and builder pattern configuration - Add Server stdio/SSE transport layers with robust error handling - Convert McpClient and McpServer to interfaces for better extensibility - Add sampling capabilities for LLM interactions (text/image) - Implement dynamic root management and runtime tool management - Add comprehensive prompt and resource capabilities - Add support for dynamic root management (add/remove roots at runtime) Architecture Changes - Introduce registration records for tools, resources, and prompts - Improve thread safety with concurrent collections - Enhance resource template handling - Implement builder patterns for Server/Client capabilities - Add support for server capabilities and implementation info Transport Layer: - Implement SseServerTransport with SSE endpoints - Add StdioServerTransport for process-based communication - Reorganize transport packages for better organization - Improve transport error handling and shutdown processes Documentation: - Add detailed MCP reference documentation - Include architecture diagrams and API documentation - Add comprehensive test coverage for all components - Include sample projects demonstrating stdio/SSE transports BREAKING CHANGES: Transport Package Reorganization - Move StdioClientTransport to o.s.ai.mcp.client.transport - Move SseClientTransport to o.s.ai.mcp.client.transport - Move ServerParameters to o.s.ai.mcp.client.transport - StdioClientTransport.errorHandler changed to private with setter method Migration: Use setErrorHandler() method instead of direct field access API Changes - Replace static factory methods with builder pattern - Update ServerCapabilities to use builder pattern - Make StdioClientTransport.errorHandler private with setter - Replace rootsListProvider with direct root management - Remove deprecated static factory methods Resolves #27,#25,#24,#11 Related to #12
This commit is contained in:
parent
c83a45f880
commit
abd0facf4e
21 changed files with 5727 additions and 471 deletions
|
|
@ -20,8 +20,8 @@ import java.time.Duration;
|
|||
|
||||
import org.springframework.ai.mcp.client.McpClient;
|
||||
import org.springframework.ai.mcp.client.McpSyncClient;
|
||||
import org.springframework.ai.mcp.client.stdio.ServerParameters;
|
||||
import org.springframework.ai.mcp.client.stdio.StdioClientTransport;
|
||||
import org.springframework.ai.mcp.client.transport.ServerParameters;
|
||||
import org.springframework.ai.mcp.client.transport.StdioClientTransport;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ListResourcesResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ListToolsResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Resource;
|
||||
|
|
|
|||
|
|
@ -20,8 +20,8 @@ import java.time.Duration;
|
|||
|
||||
import org.springframework.ai.mcp.client.McpClient;
|
||||
import org.springframework.ai.mcp.client.McpSyncClient;
|
||||
import org.springframework.ai.mcp.client.stdio.ServerParameters;
|
||||
import org.springframework.ai.mcp.client.stdio.StdioClientTransport;
|
||||
import org.springframework.ai.mcp.client.transport.ServerParameters;
|
||||
import org.springframework.ai.mcp.client.transport.StdioClientTransport;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ListResourcesResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ListToolsResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Resource;
|
||||
|
|
|
|||
|
|
@ -0,0 +1,61 @@
|
|||
/*
|
||||
* Copyright 2024 - 2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.springframework.ai.mcp.attic;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
|
||||
import org.springframework.ai.mcp.server.McpServer;
|
||||
import org.springframework.ai.mcp.server.transport.SseServerTransport;
|
||||
import org.springframework.ai.mcp.spec.McpSchema;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
|
||||
/**
|
||||
*
|
||||
* https://docs.spring.io/spring-framework/reference/web/webflux-functional.html#webflux-fn-running
|
||||
* https://docs.spring.io/spring-framework/reference/web/webflux/reactive-spring.html#webflux-httphandler
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
|
||||
public class DemoServer {
|
||||
|
||||
public static void main(String[] args) {
|
||||
SseServerTransport transport = new SseServerTransport(new ObjectMapper(), "/mcp/message");
|
||||
|
||||
var mcpServer = McpServer.using(transport)
|
||||
.serverInfo("Weather Forecast", "1.0.0")
|
||||
.tool(new McpSchema.Tool("weather", "Weather forecast tool by location", Map.of("city", "String")),
|
||||
(arguments) -> {
|
||||
String city = (String) arguments.get("city");
|
||||
return new CallToolResult(List.of(), false);
|
||||
})
|
||||
.async();
|
||||
|
||||
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transport.getRouterFunction());
|
||||
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
||||
HttpServer httpServer = HttpServer.create().port(8080).handle(adapter);
|
||||
httpServer.bindNow().onDispose().block();
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -20,7 +20,7 @@ import java.util.Map;
|
|||
|
||||
import org.springframework.ai.mcp.client.McpClient;
|
||||
import org.springframework.ai.mcp.client.McpSyncClient;
|
||||
import org.springframework.ai.mcp.client.sse.SseClientTransport;
|
||||
import org.springframework.ai.mcp.client.transport.SseClientTransport;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolRequest;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ListPromptsResult;
|
||||
|
|
|
|||
|
|
@ -17,9 +17,8 @@
|
|||
package org.springframework.ai.mcp.client;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.function.Supplier;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
|
|
@ -27,11 +26,14 @@ import org.junit.jupiter.api.Test;
|
|||
import reactor.test.StepVerifier;
|
||||
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolRequest;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.GetPromptRequest;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Prompt;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Resource;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Root;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.SubscribeRequest;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Tool;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.UnsubscribeRequest;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
|
@ -57,9 +59,11 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
abstract protected McpTransport createMcpTransport();
|
||||
|
||||
abstract protected void onStart();
|
||||
protected void onStart() {
|
||||
}
|
||||
|
||||
abstract protected void onClose();
|
||||
protected void onClose() {
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
|
|
@ -67,7 +71,10 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
this.mcpTransport = createMcpTransport();
|
||||
|
||||
assertThatCode(() -> {
|
||||
mcpAsyncClient = McpClient.using(mcpTransport).requestTimeout(TIMEOUT).async();
|
||||
mcpAsyncClient = McpClient.using(mcpTransport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.capabilities(ClientCapabilities.builder().roots(true).build())
|
||||
.async();
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
|
@ -175,19 +182,142 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
@Test
|
||||
void testRootsListChanged() {
|
||||
assertThatCode(() -> mcpAsyncClient.sendRootsListChanged().block()).doesNotThrowAnyException();
|
||||
assertThatCode(() -> mcpAsyncClient.rootsListChangedNotification().block()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testInitializeWithRootsListProviders() {
|
||||
var transport = createMcpTransport();
|
||||
List<Supplier<List<Root>>> providers = List.of(() -> List.of(new Root("file:///test/path", "test-root")));
|
||||
|
||||
var client = McpClient.using(transport).requestTimeout(TIMEOUT).rootsListProvider(providers.get(0)).async();
|
||||
var client = McpClient.using(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.roots(new Root("file:///test/path", "test-root"))
|
||||
.async();
|
||||
|
||||
assertThatCode(() -> client.initialize().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
|
||||
|
||||
assertThatCode(() -> client.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddRoot() {
|
||||
Root newRoot = new Root("file:///new/test/path", "new-test-root");
|
||||
assertThatCode(() -> mcpAsyncClient.addRoot(newRoot).block()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddRootWithNullValue() {
|
||||
assertThatThrownBy(() -> mcpAsyncClient.addRoot(null).block()).hasMessageContaining("Root must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveRoot() {
|
||||
Root root = new Root("file:///test/path/to/remove", "root-to-remove");
|
||||
assertThatCode(() -> {
|
||||
mcpAsyncClient.addRoot(root).block();
|
||||
mcpAsyncClient.removeRoot(root.uri()).block();
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveNonExistentRoot() {
|
||||
assertThatThrownBy(() -> mcpAsyncClient.removeRoot("nonexistent-uri").block())
|
||||
.hasMessageContaining("Root with uri 'nonexistent-uri' not found");
|
||||
}
|
||||
|
||||
// @Test
|
||||
void testReadResource() {
|
||||
StepVerifier.create(mcpAsyncClient.listResources()).consumeNextWith(resources -> {
|
||||
if (!resources.resources().isEmpty()) {
|
||||
Resource firstResource = resources.resources().get(0);
|
||||
StepVerifier.create(mcpAsyncClient.readResource(firstResource)).consumeNextWith(result -> {
|
||||
assertThat(result).isNotNull();
|
||||
assertThat(result.contents()).isNotNull();
|
||||
}).verifyComplete();
|
||||
}
|
||||
}).verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListResourceTemplates() {
|
||||
StepVerifier.create(mcpAsyncClient.listResourceTemplates()).consumeNextWith(result -> {
|
||||
assertThat(result).isNotNull();
|
||||
assertThat(result.resourceTemplates()).isNotNull();
|
||||
}).verifyComplete();
|
||||
}
|
||||
|
||||
// @Test
|
||||
void testResourceSubscription() {
|
||||
StepVerifier.create(mcpAsyncClient.listResources()).consumeNextWith(resources -> {
|
||||
if (!resources.resources().isEmpty()) {
|
||||
Resource firstResource = resources.resources().get(0);
|
||||
|
||||
// Test subscribe
|
||||
StepVerifier.create(mcpAsyncClient.subscribeResource(new SubscribeRequest(firstResource.uri())))
|
||||
.verifyComplete();
|
||||
|
||||
// Test unsubscribe
|
||||
StepVerifier.create(mcpAsyncClient.unsubscribeResource(new UnsubscribeRequest(firstResource.uri())))
|
||||
.verifyComplete();
|
||||
}
|
||||
}).verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNotificationHandlers() {
|
||||
AtomicBoolean toolsNotificationReceived = new AtomicBoolean(false);
|
||||
AtomicBoolean resourcesNotificationReceived = new AtomicBoolean(false);
|
||||
AtomicBoolean promptsNotificationReceived = new AtomicBoolean(false);
|
||||
|
||||
var transport = createMcpTransport();
|
||||
var client = McpClient.using(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.toolsChangeConsumer(tools -> toolsNotificationReceived.set(true))
|
||||
.resourcesChangeConsumer(resources -> resourcesNotificationReceived.set(true))
|
||||
.promptsChangeConsumer(prompts -> promptsNotificationReceived.set(true))
|
||||
.async();
|
||||
|
||||
assertThatCode(() -> {
|
||||
client.initialize().block();
|
||||
// Trigger notifications
|
||||
client.sendResourcesListChanged().block();
|
||||
client.promptListChangedNotification().block();
|
||||
client.closeGracefully().block();
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testInitializeWithSamplingCapability() {
|
||||
var transport = createMcpTransport();
|
||||
|
||||
var capabilities = ClientCapabilities.builder().sampling().build();
|
||||
|
||||
var client = McpClient.using(transport).requestTimeout(TIMEOUT).capabilities(capabilities).async();
|
||||
|
||||
assertThatCode(() -> {
|
||||
client.initialize().block(Duration.ofSeconds(10));
|
||||
client.closeGracefully().block(Duration.ofSeconds(10));
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testInitializeWithAllCapabilities() {
|
||||
var transport = createMcpTransport();
|
||||
|
||||
var capabilities = ClientCapabilities.builder()
|
||||
.experimental(Map.of("feature", "test"))
|
||||
.roots(true)
|
||||
.sampling()
|
||||
.build();
|
||||
|
||||
var client = McpClient.using(transport).requestTimeout(TIMEOUT).capabilities(capabilities).async();
|
||||
|
||||
assertThatCode(() -> {
|
||||
var result = client.initialize().block(Duration.ofSeconds(10));
|
||||
assertThat(result).isNotNull();
|
||||
assertThat(result.capabilities()).isNotNull();
|
||||
client.closeGracefully().block(Duration.ofSeconds(10));
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ package org.springframework.ai.mcp.client;
|
|||
|
||||
import java.time.Duration;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
|
|
@ -25,11 +26,17 @@ import org.junit.jupiter.api.Test;
|
|||
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolRequest;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ListResourceTemplatesResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ListResourcesResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ListToolsResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ReadResourceResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Resource;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Root;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.SubscribeRequest;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.TextContent;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Tool;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.UnsubscribeRequest;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
|
@ -46,7 +53,7 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
private McpSyncClient mcpSyncClient;
|
||||
|
||||
private static final Duration TIMEOUT = Duration.ofSeconds(20);
|
||||
private static final Duration TIMEOUT = Duration.ofSeconds(10);
|
||||
|
||||
private static final String TEST_MESSAGE = "Hello MCP Spring AI!";
|
||||
|
||||
|
|
@ -64,7 +71,10 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
this.mcpTransport = createMcpTransport();
|
||||
|
||||
assertThatCode(() -> {
|
||||
mcpSyncClient = McpClient.using(mcpTransport).requestTimeout(TIMEOUT).sync();
|
||||
mcpSyncClient = McpClient.using(mcpTransport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.capabilities(ClientCapabilities.builder().roots(true).build())
|
||||
.sync();
|
||||
mcpSyncClient.initialize();
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
|
@ -142,7 +152,7 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
@Test
|
||||
void testRootsListChanged() {
|
||||
assertThatCode(() -> mcpSyncClient.sendRootsListChanged()).doesNotThrowAnyException();
|
||||
assertThatCode(() -> mcpSyncClient.rootsListChangedNotification()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -165,4 +175,106 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
assertThat(mcpSyncClient).isNotNull();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testInitializeWithRootsListProviders() {
|
||||
var transport = createMcpTransport();
|
||||
|
||||
var client = McpClient.using(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.roots(new Root("file:///test/path", "test-root"))
|
||||
.sync();
|
||||
|
||||
assertThatCode(() -> {
|
||||
client.initialize();
|
||||
client.close();
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddRoot() {
|
||||
Root newRoot = new Root("file:///new/test/path", "new-test-root");
|
||||
assertThatCode(() -> mcpSyncClient.addRoot(newRoot)).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddRootWithNullValue() {
|
||||
assertThatThrownBy(() -> mcpSyncClient.addRoot(null)).hasMessageContaining("Root must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveRoot() {
|
||||
Root root = new Root("file:///test/path/to/remove", "root-to-remove");
|
||||
assertThatCode(() -> {
|
||||
mcpSyncClient.addRoot(root);
|
||||
mcpSyncClient.removeRoot(root.uri());
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveNonExistentRoot() {
|
||||
assertThatThrownBy(() -> mcpSyncClient.removeRoot("nonexistent-uri"))
|
||||
.hasMessageContaining("Root with uri 'nonexistent-uri' not found");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testReadResource() {
|
||||
ListResourcesResult resources = mcpSyncClient.listResources(null);
|
||||
|
||||
if (!resources.resources().isEmpty()) {
|
||||
Resource firstResource = resources.resources().get(0);
|
||||
ReadResourceResult result = mcpSyncClient.readResource(firstResource);
|
||||
|
||||
assertThat(result).isNotNull();
|
||||
assertThat(result.contents()).isNotNull();
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListResourceTemplates() {
|
||||
ListResourceTemplatesResult result = mcpSyncClient.listResourceTemplates(null);
|
||||
|
||||
assertThat(result).isNotNull();
|
||||
assertThat(result.resourceTemplates()).isNotNull();
|
||||
}
|
||||
|
||||
// @Test
|
||||
void testResourceSubscription() {
|
||||
ListResourcesResult resources = mcpSyncClient.listResources(null);
|
||||
|
||||
if (!resources.resources().isEmpty()) {
|
||||
Resource firstResource = resources.resources().get(0);
|
||||
|
||||
// Test subscribe
|
||||
assertThatCode(() -> mcpSyncClient.subscribeResource(new SubscribeRequest(firstResource.uri())))
|
||||
.doesNotThrowAnyException();
|
||||
|
||||
// Test unsubscribe
|
||||
assertThatCode(() -> mcpSyncClient.unsubscribeResource(new UnsubscribeRequest(firstResource.uri())))
|
||||
.doesNotThrowAnyException();
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNotificationHandlers() {
|
||||
AtomicBoolean toolsNotificationReceived = new AtomicBoolean(false);
|
||||
AtomicBoolean resourcesNotificationReceived = new AtomicBoolean(false);
|
||||
AtomicBoolean promptsNotificationReceived = new AtomicBoolean(false);
|
||||
|
||||
var transport = createMcpTransport();
|
||||
var client = McpClient.using(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.toolsChangeConsumer(tools -> toolsNotificationReceived.set(true))
|
||||
.resourcesChangeConsumer(resources -> resourcesNotificationReceived.set(true))
|
||||
.promptsChangeConsumer(prompts -> promptsNotificationReceived.set(true))
|
||||
.sync();
|
||||
|
||||
assertThatCode(() -> {
|
||||
client.initialize();
|
||||
// Trigger notifications
|
||||
client.sendResourcesListChanged();
|
||||
client.promptListChangedNotification();
|
||||
client.close();
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -23,11 +23,8 @@ import java.util.Map;
|
|||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.function.Consumer;
|
||||
import java.util.function.Function;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
|
||||
import static org.awaitility.Awaitility.await;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import reactor.core.publisher.Flux;
|
||||
|
|
@ -36,12 +33,14 @@ import reactor.core.publisher.Sinks;
|
|||
import reactor.core.scheduler.Schedulers;
|
||||
|
||||
import org.springframework.ai.mcp.spec.McpSchema;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCNotification;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Root;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.awaitility.Awaitility.await;
|
||||
|
||||
class McpAsyncClientResponseHandlerTests {
|
||||
|
||||
|
|
@ -150,9 +149,9 @@ class McpAsyncClientResponseHandlerTests {
|
|||
void testRootsListRequestHandling() {
|
||||
MockMcpTransport transport = new MockMcpTransport();
|
||||
|
||||
Supplier<List<Root>> rootsListProvider = () -> List.of(new Root("file:///test/path", "test-root"));
|
||||
|
||||
McpAsyncClient asyncMcpClient = McpClient.using(transport).rootsListProvider(rootsListProvider).async();
|
||||
McpAsyncClient asyncMcpClient = McpClient.using(transport)
|
||||
.roots(new Root("file:///test/path", "test-root"))
|
||||
.async();
|
||||
|
||||
// Simulate incoming request
|
||||
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "roots/list",
|
||||
|
|
@ -263,4 +262,121 @@ class McpAsyncClientResponseHandlerTests {
|
|||
asyncMcpClient.closeGracefully();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSamplingCreateMessageRequestHandling() {
|
||||
MockMcpTransport transport = new MockMcpTransport();
|
||||
|
||||
// Create a test sampling handler that echoes back the input
|
||||
Function<McpSchema.CreateMessageRequest, McpSchema.CreateMessageResult> samplingHandler = request -> {
|
||||
var content = request.messages().get(0).content();
|
||||
return new McpSchema.CreateMessageResult(McpSchema.Role.ASSISTANT, content, "test-model",
|
||||
McpSchema.CreateMessageResult.StopReason.END_TURN);
|
||||
};
|
||||
|
||||
// Create client with sampling capability and handler
|
||||
McpAsyncClient asyncMcpClient = McpClient.using(transport)
|
||||
.capabilities(ClientCapabilities.builder().sampling().build())
|
||||
.sampling(samplingHandler)
|
||||
.async();
|
||||
|
||||
// Create a mock create message request
|
||||
var messageRequest = new McpSchema.CreateMessageRequest(
|
||||
List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message"))),
|
||||
null, // modelPreferences
|
||||
"Test system prompt", McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, 0.7, // temperature
|
||||
100, // maxTokens
|
||||
null, // stopSequences
|
||||
null // metadata
|
||||
);
|
||||
|
||||
// Simulate incoming request
|
||||
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION,
|
||||
"sampling/createMessage", "test-id", messageRequest);
|
||||
transport.simulateIncomingMessage(request);
|
||||
|
||||
// Verify response
|
||||
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
|
||||
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
|
||||
|
||||
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
|
||||
assertThat(response.id()).isEqualTo("test-id");
|
||||
assertThat(response.error()).isNull();
|
||||
|
||||
McpSchema.CreateMessageResult result = transport.unmarshalFrom(response.result(),
|
||||
new TypeReference<McpSchema.CreateMessageResult>() {
|
||||
});
|
||||
assertThat(result).isNotNull();
|
||||
assertThat(result.role()).isEqualTo(McpSchema.Role.ASSISTANT);
|
||||
assertThat(result.content()).isNotNull();
|
||||
assertThat(result.model()).isEqualTo("test-model");
|
||||
assertThat(result.stopReason()).isEqualTo(McpSchema.CreateMessageResult.StopReason.END_TURN);
|
||||
|
||||
asyncMcpClient.closeGracefully();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSamplingCreateMessageRequestHandlingWithoutCapability() {
|
||||
MockMcpTransport transport = new MockMcpTransport();
|
||||
|
||||
// Create client without sampling capability
|
||||
McpAsyncClient asyncMcpClient = McpClient.using(transport)
|
||||
.capabilities(ClientCapabilities.builder().build()) // No sampling capability
|
||||
.async();
|
||||
|
||||
// Create a mock create message request
|
||||
var messageRequest = new McpSchema.CreateMessageRequest(
|
||||
List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message"))),
|
||||
null, null, null, null, 0, null, null);
|
||||
|
||||
// Simulate incoming request
|
||||
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION,
|
||||
"sampling/createMessage", "test-id", messageRequest);
|
||||
transport.simulateIncomingMessage(request);
|
||||
|
||||
// Verify error response
|
||||
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
|
||||
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
|
||||
|
||||
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
|
||||
assertThat(response.id()).isEqualTo("test-id");
|
||||
assertThat(response.result()).isNull();
|
||||
assertThat(response.error()).isNotNull();
|
||||
assertThat(response.error().message()).contains("Method not found: sampling/createMessage");
|
||||
|
||||
asyncMcpClient.closeGracefully();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSamplingCreateMessageRequestHandlingWithNullHandler() {
|
||||
MockMcpTransport transport = new MockMcpTransport();
|
||||
|
||||
// Create client with sampling capability but null handler
|
||||
McpAsyncClient asyncMcpClient = McpClient.using(transport)
|
||||
.capabilities(ClientCapabilities.builder().sampling().build())
|
||||
.sampling(null) // Null handler
|
||||
.async();
|
||||
|
||||
// Create a mock create message request
|
||||
var messageRequest = new McpSchema.CreateMessageRequest(
|
||||
List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message"))),
|
||||
null, null, null, null, 0, null, null);
|
||||
|
||||
// Simulate incoming request
|
||||
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION,
|
||||
"sampling/createMessage", "test-id", messageRequest);
|
||||
transport.simulateIncomingMessage(request);
|
||||
|
||||
// Verify error response
|
||||
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
|
||||
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
|
||||
|
||||
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
|
||||
assertThat(response.id()).isEqualTo("test-id");
|
||||
assertThat(response.result()).isNull();
|
||||
assertThat(response.error()).isNotNull();
|
||||
assertThat(response.error().message()).contains("Method not found: sampling/createMessage");
|
||||
|
||||
asyncMcpClient.closeGracefully();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,42 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.mcp.client;
|
||||
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
|
||||
import org.springframework.ai.mcp.client.transport.ServerParameters;
|
||||
import org.springframework.ai.mcp.client.transport.StdioClientTransport;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
|
||||
/**
|
||||
* Tests for the {@link McpAsyncClient} with {@link StdioClientTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Dariusz Jędrzejczyk
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class McpAsyncClientTests extends AbstractMcpAsyncClientTests {
|
||||
|
||||
@Override
|
||||
protected McpTransport createMcpTransport() {
|
||||
ServerParameters stdioParams = ServerParameters.builder("npx")
|
||||
.args("-y", "@modelcontextprotocol/server-everything", "dir")
|
||||
.build();
|
||||
return new StdioClientTransport(stdioParams);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,68 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.mcp.client;
|
||||
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
|
||||
import org.springframework.ai.mcp.client.transport.ServerParameters;
|
||||
import org.springframework.ai.mcp.client.transport.StdioClientTransport;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Tests for the {@link McpSyncClient} with {@link StdioClientTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Dariusz Jędrzejczyk
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class McpSyncClientTests extends AbstractMcpSyncClientTests {
|
||||
|
||||
@Override
|
||||
protected McpTransport createMcpTransport() {
|
||||
ServerParameters stdioParams = ServerParameters.builder("npx")
|
||||
.args("-y", "@modelcontextprotocol/server-everything", "dir")
|
||||
.build();
|
||||
|
||||
return new StdioClientTransport(stdioParams);
|
||||
}
|
||||
|
||||
@Test
|
||||
void customErrorHandlerShouldReceiveErrors() {
|
||||
AtomicReference<String> receivedError = new AtomicReference<>();
|
||||
|
||||
((StdioClientTransport) mcpTransport).setErrorHandler(error -> receivedError.set(error));
|
||||
|
||||
String errorMessage = "Test error";
|
||||
((StdioClientTransport) mcpTransport).getErrorSink().tryEmitNext(errorMessage);
|
||||
|
||||
assertThat(receivedError.get()).isNotNull().isEqualTo(errorMessage);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onClose() {
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,61 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.mcp.client;
|
||||
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
|
||||
import org.springframework.ai.mcp.client.transport.SseClientTransport;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
/**
|
||||
* Tests for the {@link McpAsyncClient} with {@link SseClientTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class SseMcpAsyncClientTests extends AbstractMcpAsyncClientTests {
|
||||
|
||||
static String host = "http://localhost:3001";
|
||||
|
||||
// Uses the https://github.com/tzolov/mcp-everything-server-docker-image
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
@Override
|
||||
protected McpTransport createMcpTransport() {
|
||||
return new SseClientTransport(WebClient.builder().baseUrl(host));
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void onClose() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,61 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.mcp.client;
|
||||
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
|
||||
import org.springframework.ai.mcp.client.transport.SseClientTransport;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
/**
|
||||
* Tests for the {@link McpSyncClient} with {@link SseClientTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15) // Giving extra time beyond the client timeout
|
||||
class SseMcpSyncClientTests extends AbstractMcpSyncClientTests {
|
||||
|
||||
static String host = "http://localhost:3001";
|
||||
|
||||
// Uses the https://github.com/tzolov/mcp-everything-server-docker-image
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
@Override
|
||||
protected McpTransport createMcpTransport() {
|
||||
return new SseClientTransport(WebClient.builder().baseUrl(host));
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onClose() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,333 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.mcp.client.transport;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.function.Function;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.Sinks;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import org.springframework.ai.mcp.spec.McpSchema;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
|
||||
import org.springframework.http.codec.ServerSentEvent;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatCode;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
/**
|
||||
* Tests for the {@link SseClientTransport} class.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15)
|
||||
class SseClientTransportTests {
|
||||
|
||||
static String host = "http://localhost:3001";
|
||||
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
private TestSseServerTransport transport;
|
||||
|
||||
private WebClient.Builder webClientBuilder;
|
||||
|
||||
private ObjectMapper objectMapper;
|
||||
|
||||
// Test class to access protected methods
|
||||
static class TestSseServerTransport extends SseClientTransport {
|
||||
|
||||
private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
|
||||
|
||||
private Sinks.Many<ServerSentEvent<String>> events = Sinks.many().unicast().onBackpressureBuffer();
|
||||
|
||||
public TestSseServerTransport(WebClient.Builder webClientBuilder, ObjectMapper objectMapper) {
|
||||
super(webClientBuilder, objectMapper);
|
||||
}
|
||||
|
||||
// @Override
|
||||
// public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>,
|
||||
// Mono<McpSchema.JSONRPCMessage>> handler) {
|
||||
// simulateEndpointEvent("https://localhost:3001");
|
||||
// return super.connect(handler);
|
||||
// }
|
||||
|
||||
@Override
|
||||
protected Flux<ServerSentEvent<String>> eventStream() {
|
||||
return super.eventStream().mergeWith(events.asFlux());
|
||||
}
|
||||
|
||||
public String getLastEndpoint() {
|
||||
return messageEndpointSink.asMono().block();
|
||||
}
|
||||
|
||||
public int getInboundMessageCount() {
|
||||
return inboundMessageCount.get();
|
||||
}
|
||||
|
||||
public void simulateEndpointEvent(String jsonMessage) {
|
||||
events.tryEmitNext(ServerSentEvent.<String>builder().event("endpoint").data(jsonMessage).build());
|
||||
inboundMessageCount.incrementAndGet();
|
||||
}
|
||||
|
||||
public void simulateMessageEvent(String jsonMessage) {
|
||||
events.tryEmitNext(ServerSentEvent.<String>builder().event("message").data(jsonMessage).build());
|
||||
inboundMessageCount.incrementAndGet();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
void startContainer() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
startContainer();
|
||||
webClientBuilder = WebClient.builder().baseUrl(host);
|
||||
objectMapper = new ObjectMapper();
|
||||
transport = new TestSseServerTransport(webClientBuilder, objectMapper);
|
||||
transport.connect(Function.identity()).block();
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void afterEach() {
|
||||
if (transport != null) {
|
||||
assertThatCode(() -> transport.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
|
||||
}
|
||||
cleanup();
|
||||
}
|
||||
|
||||
void cleanup() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testEndpointEventHandling() {
|
||||
assertThat(transport.getLastEndpoint()).startsWith("/message?");
|
||||
}
|
||||
|
||||
@Test
|
||||
void constructorValidation() {
|
||||
assertThatThrownBy(() -> new SseClientTransport(null)).isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("WebClient.Builder must not be null");
|
||||
|
||||
assertThatThrownBy(() -> new SseClientTransport(webClientBuilder, null))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("ObjectMapper must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMessageProcessing() {
|
||||
// Create a test message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Simulate receiving the message
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "test-method",
|
||||
"id": "test-id",
|
||||
"params": {"key": "value"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Subscribe to messages and verify
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testResponseMessageProcessing() {
|
||||
// Simulate receiving a response message
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "test-id",
|
||||
"result": {"status": "success"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Create and send a request message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Verify message handling
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testErrorMessageProcessing() {
|
||||
// Simulate receiving an error message
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": "test-id",
|
||||
"error": {
|
||||
"code": -32600,
|
||||
"message": "Invalid Request"
|
||||
}
|
||||
}
|
||||
""");
|
||||
|
||||
// Create and send a request message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Verify message handling
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNotificationMessageProcessing() {
|
||||
// Simulate receiving a notification message (no id)
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "update",
|
||||
"params": {"status": "processing"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Verify the notification was processed
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGracefulShutdown() {
|
||||
// Test graceful shutdown
|
||||
StepVerifier.create(transport.closeGracefully()).verifyComplete();
|
||||
|
||||
// Create a test message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Verify message is not processed after shutdown
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
// Message count should remain 0 after shutdown
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(0);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRetryBehavior() {
|
||||
// Create a WebClient that simulates connection failures
|
||||
WebClient.Builder failingWebClientBuilder = WebClient.builder().baseUrl("http://non-existent-host");
|
||||
|
||||
SseClientTransport failingTransport = new SseClientTransport(failingWebClientBuilder);
|
||||
|
||||
// Verify that the transport attempts to reconnect
|
||||
StepVerifier.create(Mono.delay(Duration.ofSeconds(2))).expectNextCount(1).verifyComplete();
|
||||
|
||||
// Clean up
|
||||
failingTransport.closeGracefully().block();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMultipleMessageProcessing() {
|
||||
// Simulate receiving multiple messages in sequence
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "method1",
|
||||
"id": "id1",
|
||||
"params": {"key": "value1"}
|
||||
}
|
||||
""");
|
||||
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "method2",
|
||||
"id": "id2",
|
||||
"params": {"key": "value2"}
|
||||
}
|
||||
""");
|
||||
|
||||
// Create and send corresponding messages
|
||||
JSONRPCRequest message1 = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "method1", "id1",
|
||||
Map.of("key", "value1"));
|
||||
|
||||
JSONRPCRequest message2 = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "method2", "id2",
|
||||
Map.of("key", "value2"));
|
||||
|
||||
// Verify both messages are processed
|
||||
StepVerifier.create(transport.sendMessage(message1).then(transport.sendMessage(message2))).verifyComplete();
|
||||
|
||||
// Verify message count
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(2);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMessageOrderPreservation() {
|
||||
// Simulate receiving messages in a specific order
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "first",
|
||||
"id": "1",
|
||||
"params": {"sequence": 1}
|
||||
}
|
||||
""");
|
||||
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "second",
|
||||
"id": "2",
|
||||
"params": {"sequence": 2}
|
||||
}
|
||||
""");
|
||||
|
||||
transport.simulateMessageEvent("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "third",
|
||||
"id": "3",
|
||||
"params": {"sequence": 3}
|
||||
}
|
||||
""");
|
||||
|
||||
// Verify message count and order
|
||||
assertThat(transport.getInboundMessageCount()).isEqualTo(3);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,342 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.mcp.server;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import org.springframework.ai.mcp.spec.McpSchema;
|
||||
import org.springframework.ai.mcp.server.McpServer.PromptRegistration;
|
||||
import org.springframework.ai.mcp.server.McpServer.ResourceRegistration;
|
||||
import org.springframework.ai.mcp.server.McpServer.ToolRegistration;
|
||||
import org.springframework.ai.mcp.spec.McpError;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.GetPromptResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Prompt;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.PromptMessage;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ReadResourceResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Resource;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ServerCapabilities;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Tool;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatCode;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
/**
|
||||
* Test suite for the {@link McpAsyncServer} that can be used with different
|
||||
* {@link McpTransport} implementations.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
public abstract class AbstractMcpAsyncServerTests {
|
||||
|
||||
private static final String TEST_TOOL_NAME = "test-tool";
|
||||
|
||||
private static final String TEST_RESOURCE_URI = "test://resource";
|
||||
|
||||
private static final String TEST_PROMPT_NAME = "test-prompt";
|
||||
|
||||
abstract protected McpTransport createMcpTransport();
|
||||
|
||||
protected void onStart() {
|
||||
}
|
||||
|
||||
protected void onClose() {
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void tearDown() {
|
||||
onClose();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testConstructorWithInvalidArguments() {
|
||||
assertThatThrownBy(() -> McpServer.using(null)).isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessage("Transport must not be null");
|
||||
|
||||
assertThatThrownBy(() -> McpServer.using(createMcpTransport()).info((McpSchema.Implementation) null))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessage("Server info must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddTool() {
|
||||
Tool newTool = new McpSchema.Tool("new-tool", "New test tool", Map.of("input", "string"));
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.async();
|
||||
|
||||
StepVerifier
|
||||
.create(mcpAsyncServer.addTool(new ToolRegistration(newTool, args -> new CallToolResult(List.of(), false))))
|
||||
.verifyComplete();
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10)));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddDuplicateTool() {
|
||||
Tool duplicateTool = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", Map.of("input", "string"));
|
||||
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.tool(duplicateTool, args -> new CallToolResult(List.of(), false))
|
||||
.async();
|
||||
|
||||
StepVerifier
|
||||
.create(mcpAsyncServer
|
||||
.addTool(new ToolRegistration(duplicateTool, args -> new CallToolResult(List.of(), false))))
|
||||
.verifyErrorSatisfies(error -> {
|
||||
assertThat(error).isInstanceOf(McpError.class)
|
||||
.hasMessage("Tool with name '" + TEST_TOOL_NAME + "' already exists");
|
||||
});
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10)));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveTool() {
|
||||
Tool too = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", Map.of("input", "string"));
|
||||
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.tool(too, args -> new CallToolResult(List.of(), false))
|
||||
.async();
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.removeTool(TEST_TOOL_NAME)).verifyComplete();
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10)));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveNonexistentTool() {
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.async();
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.removeTool("nonexistent-tool")).verifyErrorSatisfies(error -> {
|
||||
assertThat(error).isInstanceOf(McpError.class).hasMessage("Tool with name 'nonexistent-tool' not found");
|
||||
});
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10)));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNotifyToolsListChanged() {
|
||||
Tool too = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", Map.of("input", "string"));
|
||||
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.tool(too, args -> new CallToolResult(List.of(), false))
|
||||
.async();
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.notifyToolsListChanged()).verifyComplete();
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10)));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNotifyResourcesListChanged() {
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").async();
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.notifyResourcesListChanged()).verifyComplete();
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10)));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNotifyPromptsListChanged() {
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").async();
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.notifyPromptsListChanged()).verifyComplete();
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10)));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGracefulShutdown() {
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").async();
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.closeGracefully()).verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testImmediateClose() {
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").async();
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer.close()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddResource() {
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().resources(true, false).build())
|
||||
.async();
|
||||
|
||||
Resource resource = new Resource(TEST_RESOURCE_URI, "Test Resource", "text/plain", "Test resource description",
|
||||
null);
|
||||
ResourceRegistration registration = new ResourceRegistration(resource,
|
||||
req -> new ReadResourceResult(List.of()));
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.addResource(registration)).verifyComplete();
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddResourceWithNullRegistration() {
|
||||
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().resources(true, false).build())
|
||||
.async();
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.addResource(null)).verifyErrorSatisfies(error -> {
|
||||
assertThat(error).isInstanceOf(McpError.class).hasMessage("Resource must not be null");
|
||||
});
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddPromptWithNullRegistration() {
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().prompts(false).build())
|
||||
.async();
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.addPrompt(null)).verifyErrorSatisfies(error -> {
|
||||
assertThat(error).isInstanceOf(McpError.class).hasMessage("Prompt registration must not be null");
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddResourceWithoutCapability() {
|
||||
// Create a server without resource capabilities
|
||||
McpAsyncServer serverWithoutResources = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.async();
|
||||
|
||||
Resource resource = new Resource(TEST_RESOURCE_URI, "Test Resource", "text/plain", "Test resource description",
|
||||
null);
|
||||
ResourceRegistration registration = new ResourceRegistration(resource,
|
||||
req -> new ReadResourceResult(List.of()));
|
||||
|
||||
StepVerifier.create(serverWithoutResources.addResource(registration)).verifyErrorSatisfies(error -> {
|
||||
assertThat(error).isInstanceOf(McpError.class)
|
||||
.hasMessage("Server must be configured with resource capabilities");
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddPromptWithoutCapability() {
|
||||
// Create a server without prompt capabilities
|
||||
McpAsyncServer serverWithoutPrompts = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.async();
|
||||
|
||||
Prompt prompt = new Prompt(TEST_PROMPT_NAME, "Test Prompt", List.of());
|
||||
PromptRegistration registration = new PromptRegistration(prompt, req -> new GetPromptResult(
|
||||
"Test prompt description",
|
||||
List.of(new PromptMessage(McpSchema.Role.ASSISTANT, new McpSchema.TextContent("Test content")))));
|
||||
|
||||
StepVerifier.create(serverWithoutPrompts.addPrompt(registration)).verifyErrorSatisfies(error -> {
|
||||
assertThat(error).isInstanceOf(McpError.class)
|
||||
.hasMessage("Server must be configured with prompt capabilities");
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveResourceWithoutCapability() {
|
||||
// Create a server without resource capabilities
|
||||
McpAsyncServer serverWithoutResources = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.async();
|
||||
|
||||
StepVerifier.create(serverWithoutResources.removeResource(TEST_RESOURCE_URI)).verifyErrorSatisfies(error -> {
|
||||
assertThat(error).isInstanceOf(McpError.class)
|
||||
.hasMessage("Server must be configured with resource capabilities");
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemovePromptWithoutCapability() {
|
||||
// Create a server without prompt capabilities
|
||||
McpAsyncServer serverWithoutPrompts = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.async();
|
||||
|
||||
StepVerifier.create(serverWithoutPrompts.removePrompt(TEST_PROMPT_NAME)).verifyErrorSatisfies(error -> {
|
||||
assertThat(error).isInstanceOf(McpError.class)
|
||||
.hasMessage("Server must be configured with prompt capabilities");
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemovePrompt() {
|
||||
String TEST_PROMPT_NAME_TO_REMOVE = "TEST_PROMPT_NAME678";
|
||||
|
||||
Prompt prompt = new Prompt(TEST_PROMPT_NAME_TO_REMOVE, "Test Prompt", List.of());
|
||||
PromptRegistration registration = new PromptRegistration(prompt, req -> new GetPromptResult(
|
||||
"Test prompt description",
|
||||
List.of(new PromptMessage(McpSchema.Role.ASSISTANT, new McpSchema.TextContent("Test content")))));
|
||||
|
||||
var mcpAsyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().prompts(true).build())
|
||||
.prompts(registration)
|
||||
.async();
|
||||
|
||||
StepVerifier.create(mcpAsyncServer.removePrompt(TEST_PROMPT_NAME_TO_REMOVE)).verifyComplete();
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveNonexistentPrompt() {
|
||||
var mcpAsyncServer2 = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().prompts(true).build())
|
||||
.async();
|
||||
|
||||
StepVerifier.create(mcpAsyncServer2.removePrompt("nonexistent-prompt")).verifyErrorSatisfies(error -> {
|
||||
assertThat(error).isInstanceOf(McpError.class)
|
||||
.hasMessage("Prompt with name 'nonexistent-prompt' not found");
|
||||
});
|
||||
|
||||
assertThatCode(() -> mcpAsyncServer2.closeGracefully().block(Duration.ofSeconds(10)))
|
||||
.doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,314 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.mcp.server;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import org.springframework.ai.mcp.spec.McpSchema;
|
||||
import org.springframework.ai.mcp.server.McpServer.PromptRegistration;
|
||||
import org.springframework.ai.mcp.server.McpServer.ResourceRegistration;
|
||||
import org.springframework.ai.mcp.server.McpServer.ToolRegistration;
|
||||
import org.springframework.ai.mcp.spec.McpError;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.GetPromptResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Prompt;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.PromptMessage;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ReadResourceResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Resource;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ServerCapabilities;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Tool;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatCode;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
/**
|
||||
* Test suite for the {@link McpSyncServer} that can be used with different
|
||||
* {@link McpTransport} implementations.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
public abstract class AbstractMcpSyncServerTests {
|
||||
|
||||
private static final String TEST_TOOL_NAME = "test-tool";
|
||||
|
||||
private static final String TEST_RESOURCE_URI = "test://resource";
|
||||
|
||||
private static final String TEST_PROMPT_NAME = "test-prompt";
|
||||
|
||||
abstract protected McpTransport createMcpTransport();
|
||||
|
||||
protected void onStart() {
|
||||
}
|
||||
|
||||
protected void onClose() {
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void tearDown() {
|
||||
onClose();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testConstructorWithInvalidArguments() {
|
||||
assertThatThrownBy(() -> McpServer.using(null)).isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessage("Transport must not be null");
|
||||
|
||||
assertThatThrownBy(() -> McpServer.using(createMcpTransport()).info((McpSchema.Implementation) null))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessage("Server info must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddTool() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.sync();
|
||||
|
||||
Tool newTool = new McpSchema.Tool("new-tool", "New test tool", Map.of("input", "string"));
|
||||
assertThatCode(() -> mcpSyncServer
|
||||
.addTool(new ToolRegistration(newTool, args -> new CallToolResult(List.of(), false))))
|
||||
.doesNotThrowAnyException();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddDuplicateTool() {
|
||||
Tool duplicateTool = new McpSchema.Tool(TEST_TOOL_NAME, "Duplicate tool", Map.of("input", "string"));
|
||||
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.tool(duplicateTool, args -> new CallToolResult(List.of(), false))
|
||||
.sync();
|
||||
|
||||
assertThatThrownBy(() -> mcpSyncServer
|
||||
.addTool(new ToolRegistration(duplicateTool, args -> new CallToolResult(List.of(), false))))
|
||||
.isInstanceOf(McpError.class)
|
||||
.hasMessage("Tool with name '" + TEST_TOOL_NAME + "' already exists");
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveTool() {
|
||||
Tool tool = new McpSchema.Tool(TEST_TOOL_NAME, "Test tool", Map.of("input", "string"));
|
||||
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.tool(tool, args -> new CallToolResult(List.of(), false))
|
||||
.sync();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.removeTool(TEST_TOOL_NAME)).doesNotThrowAnyException();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveNonexistentTool() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.sync();
|
||||
|
||||
assertThatThrownBy(() -> mcpSyncServer.removeTool("nonexistent-tool")).isInstanceOf(McpError.class)
|
||||
.hasMessage("Tool with name 'nonexistent-tool' not found");
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNotifyToolsListChanged() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").sync();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.notifyToolsListChanged()).doesNotThrowAnyException();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNotifyResourcesListChanged() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").sync();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.notifyResourcesListChanged()).doesNotThrowAnyException();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNotifyPromptsListChanged() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").sync();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.notifyPromptsListChanged()).doesNotThrowAnyException();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGracefulShutdown() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").sync();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testImmediateClose() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").sync();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.close()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddResource() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().resources(true, false).build())
|
||||
.sync();
|
||||
|
||||
Resource resource = new Resource(TEST_RESOURCE_URI, "Test Resource", "text/plain", "Test resource description",
|
||||
null);
|
||||
ResourceRegistration registration = new ResourceRegistration(resource,
|
||||
req -> new ReadResourceResult(List.of()));
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.addResource(registration)).doesNotThrowAnyException();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddResourceWithNullRegistration() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().resources(true, false).build())
|
||||
.sync();
|
||||
|
||||
assertThatThrownBy(() -> mcpSyncServer.addResource(null)).isInstanceOf(McpError.class)
|
||||
.hasMessage("Resource must not be null");
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddPromptWithNullRegistration() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().prompts(false).build())
|
||||
.sync();
|
||||
|
||||
assertThatThrownBy(() -> mcpSyncServer.addPrompt(null)).isInstanceOf(McpError.class)
|
||||
.hasMessage("Prompt registration must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddResourceWithoutCapability() {
|
||||
var serverWithoutResources = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").sync();
|
||||
|
||||
Resource resource = new Resource(TEST_RESOURCE_URI, "Test Resource", "text/plain", "Test resource description",
|
||||
null);
|
||||
ResourceRegistration registration = new ResourceRegistration(resource,
|
||||
req -> new ReadResourceResult(List.of()));
|
||||
|
||||
assertThatThrownBy(() -> serverWithoutResources.addResource(registration)).isInstanceOf(McpError.class)
|
||||
.hasMessage("Server must be configured with resource capabilities");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAddPromptWithoutCapability() {
|
||||
var serverWithoutPrompts = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").sync();
|
||||
|
||||
Prompt prompt = new Prompt(TEST_PROMPT_NAME, "Test Prompt", List.of());
|
||||
PromptRegistration registration = new PromptRegistration(prompt, req -> new GetPromptResult(
|
||||
"Test prompt description",
|
||||
List.of(new PromptMessage(McpSchema.Role.ASSISTANT, new McpSchema.TextContent("Test content")))));
|
||||
|
||||
assertThatThrownBy(() -> serverWithoutPrompts.addPrompt(registration)).isInstanceOf(McpError.class)
|
||||
.hasMessage("Server must be configured with prompt capabilities");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveResourceWithoutCapability() {
|
||||
var serverWithoutResources = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").sync();
|
||||
|
||||
assertThatThrownBy(() -> serverWithoutResources.removeResource(TEST_RESOURCE_URI)).isInstanceOf(McpError.class)
|
||||
.hasMessage("Server must be configured with resource capabilities");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemovePromptWithoutCapability() {
|
||||
var serverWithoutPrompts = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").sync();
|
||||
|
||||
assertThatThrownBy(() -> serverWithoutPrompts.removePrompt(TEST_PROMPT_NAME)).isInstanceOf(McpError.class)
|
||||
.hasMessage("Server must be configured with prompt capabilities");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemovePrompt() {
|
||||
|
||||
Prompt prompt = new Prompt(TEST_PROMPT_NAME, "Test Prompt", List.of());
|
||||
PromptRegistration registration = new PromptRegistration(prompt, req -> new GetPromptResult(
|
||||
"Test prompt description",
|
||||
List.of(new PromptMessage(McpSchema.Role.ASSISTANT, new McpSchema.TextContent("Test content")))));
|
||||
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().prompts(true).build())
|
||||
.prompts(registration)
|
||||
.sync();
|
||||
// assertThatCode(() ->
|
||||
// mcpSyncServer.addPrompt(registration)).doesNotThrowAnyException();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.removePrompt(TEST_PROMPT_NAME)).doesNotThrowAnyException();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRemoveNonexistentPrompt() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport())
|
||||
.serverInfo("test-server", "1.0.0")
|
||||
.capabilities(ServerCapabilities.builder().prompts(true).build())
|
||||
.sync();
|
||||
|
||||
assertThatThrownBy(() -> mcpSyncServer.removePrompt("nonexistent-prompt")).isInstanceOf(McpError.class)
|
||||
.hasMessage("Prompt with name 'nonexistent-prompt' not found");
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGetAsyncServer() {
|
||||
var mcpSyncServer = McpServer.using(createMcpTransport()).serverInfo("test-server", "1.0.0").sync();
|
||||
|
||||
assertThat(mcpSyncServer.getAsyncServer()).isNotNull();
|
||||
|
||||
assertThatCode(() -> mcpSyncServer.closeGracefully()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.mcp.server;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
import reactor.netty.DisposableServer;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
|
||||
import org.springframework.ai.mcp.server.transport.SseServerTransport;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
|
||||
/**
|
||||
* Tests for {@link McpAsyncServer} using {@link SseServerTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
class SseMcpAsyncServerTests extends AbstractMcpAsyncServerTests {
|
||||
|
||||
private static final int PORT = 8181;
|
||||
|
||||
private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
||||
|
||||
private DisposableServer httpServer;
|
||||
|
||||
@Override
|
||||
protected McpTransport createMcpTransport() {
|
||||
var transport = new SseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
|
||||
|
||||
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transport.getRouterFunction());
|
||||
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
||||
httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
|
||||
|
||||
return transport;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onClose() {
|
||||
if (httpServer != null) {
|
||||
httpServer.disposeNow();
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,64 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.mcp.server;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
import reactor.netty.DisposableServer;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
|
||||
import org.springframework.ai.mcp.server.transport.SseServerTransport;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
|
||||
/**
|
||||
* Tests for {@link McpSyncServer} using {@link SseServerTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
class SseMcpSyncServerTests extends AbstractMcpSyncServerTests {
|
||||
|
||||
private static final int PORT = 8182;
|
||||
|
||||
private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
||||
|
||||
private DisposableServer httpServer;
|
||||
|
||||
private SseServerTransport transport;
|
||||
|
||||
@Override
|
||||
protected McpTransport createMcpTransport() {
|
||||
transport = new SseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
|
||||
return transport;
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transport.getRouterFunction());
|
||||
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
|
||||
httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onClose() {
|
||||
if (httpServer != null) {
|
||||
httpServer.disposeNow();
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,36 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.mcp.server;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
|
||||
import org.springframework.ai.mcp.server.transport.StdioServerTransport;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
|
||||
/**
|
||||
* Tests for {@link McpAsyncServer} using {@link StdioServerTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
class StdioMcpAsyncServerTests extends AbstractMcpAsyncServerTests {
|
||||
|
||||
@Override
|
||||
protected McpTransport createMcpTransport() {
|
||||
return new StdioServerTransport(new ObjectMapper());
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,83 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
package org.springframework.ai.mcp.server;
|
||||
|
||||
import java.io.PipedInputStream;
|
||||
import java.io.PipedOutputStream;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
|
||||
import org.springframework.ai.mcp.server.transport.StdioServerTransport;
|
||||
import org.springframework.ai.mcp.spec.McpTransport;
|
||||
|
||||
/**
|
||||
* Tests for {@link McpSyncServer} using {@link StdioServerTransport}.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
class StdioMcpSyncServerTests extends AbstractMcpSyncServerTests {
|
||||
|
||||
private PipedInputStream testInput;
|
||||
|
||||
private PipedOutputStream testOutput;
|
||||
|
||||
private PipedOutputStream writeToInput;
|
||||
|
||||
private PipedInputStream readFromOutput;
|
||||
|
||||
@Override
|
||||
protected McpTransport createMcpTransport() {
|
||||
try {
|
||||
testInput = new PipedInputStream();
|
||||
writeToInput = new PipedOutputStream(testInput);
|
||||
testOutput = new PipedOutputStream();
|
||||
readFromOutput = new PipedInputStream(testOutput);
|
||||
return new StdioServerTransport(new ObjectMapper(), testInput, testOutput);
|
||||
}
|
||||
catch (Exception e) {
|
||||
throw new RuntimeException("Failed to initialize test streams", e);
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onStart() {
|
||||
// No special setup needed for stdio transport
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void onClose() {
|
||||
try {
|
||||
if (testInput != null) {
|
||||
testInput.close();
|
||||
}
|
||||
if (testOutput != null) {
|
||||
testOutput.close();
|
||||
}
|
||||
if (writeToInput != null) {
|
||||
writeToInput.close();
|
||||
}
|
||||
if (readFromOutput != null) {
|
||||
readFromOutput.close();
|
||||
}
|
||||
}
|
||||
catch (Exception e) {
|
||||
// Log but don't throw since this is cleanup
|
||||
e.printStackTrace();
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,81 @@
|
|||
/*
|
||||
* Copyright 2024 - 2024 the original author or authors.
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* https://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
package org.springframework.ai.mcp.server.transport;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.io.InputStream;
|
||||
import java.util.concurrent.BlockingQueue;
|
||||
import java.util.concurrent.LinkedBlockingQueue;
|
||||
|
||||
public class BlockingInputStream extends InputStream {
|
||||
|
||||
private final BlockingQueue<Integer> queue = new LinkedBlockingQueue<>();
|
||||
|
||||
private volatile boolean completed = false;
|
||||
|
||||
private volatile boolean closed = false;
|
||||
|
||||
@Override
|
||||
public int read() throws IOException {
|
||||
if (closed) {
|
||||
throw new IOException("Stream is closed");
|
||||
}
|
||||
|
||||
try {
|
||||
Integer value = queue.poll();
|
||||
if (value == null) {
|
||||
if (completed) {
|
||||
return -1;
|
||||
}
|
||||
value = queue.take(); // Blocks until data is available
|
||||
if (value == null && completed) {
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
return value;
|
||||
}
|
||||
catch (InterruptedException e) {
|
||||
Thread.currentThread().interrupt();
|
||||
throw new IOException("Read interrupted", e);
|
||||
}
|
||||
}
|
||||
|
||||
public void write(int b) {
|
||||
if (!closed && !completed) {
|
||||
queue.offer(b);
|
||||
}
|
||||
}
|
||||
|
||||
public void write(byte[] data) {
|
||||
if (!closed && !completed) {
|
||||
for (byte b : data) {
|
||||
queue.offer((int) b & 0xFF);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public void complete() {
|
||||
this.completed = true;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() {
|
||||
this.closed = true;
|
||||
this.completed = true;
|
||||
this.queue.clear();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,283 @@
|
|||
package org.springframework.ai.mcp.server.transport;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Disabled;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.Disposable;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import org.springframework.ai.mcp.spec.McpSchema;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
|
||||
import org.springframework.http.HttpStatus;
|
||||
import org.springframework.http.MediaType;
|
||||
import org.springframework.http.codec.ServerSentEvent;
|
||||
import org.springframework.test.web.reactive.server.WebTestClient;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
|
||||
/**
|
||||
* Tests for the {@link SseServerTransport} class.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15)
|
||||
class SseServerTransportTests {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(SseServerTransportTests.class);
|
||||
|
||||
private ObjectMapper objectMapper;
|
||||
|
||||
private String messageEndpoint;
|
||||
|
||||
private SseServerTransport transport;
|
||||
|
||||
private WebTestClient webTestClient;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
objectMapper = new ObjectMapper();
|
||||
messageEndpoint = "/message";
|
||||
transport = new SseServerTransport(objectMapper, messageEndpoint);
|
||||
webTestClient = WebTestClient.bindToRouterFunction(transport.getRouterFunction()).build();
|
||||
}
|
||||
|
||||
@Test
|
||||
void constructorValidation() {
|
||||
assertThatThrownBy(() -> new SseServerTransport(null, "/message")).isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("ObjectMapper must not be null");
|
||||
|
||||
assertThatThrownBy(() -> new SseServerTransport(new ObjectMapper(), null))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Message endpoint must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSseConnectionEstablishment() {
|
||||
List<ServerSentEvent<String>> events = new ArrayList<>();
|
||||
|
||||
webTestClient.get()
|
||||
.uri("/sse")
|
||||
.accept(MediaType.TEXT_EVENT_STREAM)
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isOk()
|
||||
.expectHeader()
|
||||
.contentTypeCompatibleWith(MediaType.TEXT_EVENT_STREAM)
|
||||
.returnResult(String.class)
|
||||
.getResponseBody()
|
||||
.map(data -> ServerSentEvent.<String>builder().data(data).build())
|
||||
.take(1) // Take only the initial endpoint event
|
||||
.subscribe(events::add);
|
||||
|
||||
// Wait a bit for the event to be received
|
||||
StepVerifier.create(Mono.delay(Duration.ofMillis(500))).expectNextCount(1).verifyComplete();
|
||||
|
||||
assertThat(events).hasSize(1);
|
||||
assertThat(events.get(0).data()).isEqualTo(messageEndpoint);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMessageHandling() {
|
||||
AtomicReference<McpSchema.JSONRPCMessage> receivedMessage = new AtomicReference<>();
|
||||
|
||||
// Set up message handler
|
||||
transport.connect(message -> {
|
||||
message.doOnNext(receivedMessage::set).subscribe();
|
||||
return Mono.empty();
|
||||
}).block();
|
||||
|
||||
// Create a test message
|
||||
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Send message to endpoint
|
||||
webTestClient.post()
|
||||
.uri(messageEndpoint)
|
||||
.contentType(MediaType.APPLICATION_JSON)
|
||||
.bodyValue(testMessage)
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isOk();
|
||||
|
||||
// Verify the message was received and processed
|
||||
assertThat(receivedMessage.get()).isNotNull();
|
||||
McpSchema.JSONRPCRequest receivedRequest = (McpSchema.JSONRPCRequest) receivedMessage.get();
|
||||
assertThat(receivedRequest.id()).isEqualTo(testMessage.id());
|
||||
assertThat(receivedRequest.method()).isEqualTo(testMessage.method());
|
||||
}
|
||||
|
||||
@Test
|
||||
@Disabled("Flaky test")
|
||||
void testBroadcastMessage() {
|
||||
// Create test clients
|
||||
int clientCount = 3;
|
||||
CountDownLatch connectLatch = new CountDownLatch(clientCount);
|
||||
CountDownLatch messageLatch = new CountDownLatch(clientCount);
|
||||
|
||||
List<List<ServerSentEvent<String>>> allReceivedEvents = new ArrayList<>();
|
||||
List<Disposable> subscriptions = new ArrayList<>();
|
||||
|
||||
// Connect clients
|
||||
for (int i = 0; i < clientCount; i++) {
|
||||
List<ServerSentEvent<String>> clientEvents = new ArrayList<>();
|
||||
allReceivedEvents.add(clientEvents);
|
||||
|
||||
Flux<ServerSentEvent<String>> eventStream = webTestClient.get()
|
||||
.uri("/sse")
|
||||
.accept(MediaType.TEXT_EVENT_STREAM)
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isOk()
|
||||
.returnResult(String.class)
|
||||
.getResponseBody()
|
||||
.map(data -> ServerSentEvent.<String>builder().data(data).build());
|
||||
|
||||
Disposable subscription = eventStream.doOnNext(event -> {
|
||||
clientEvents.add(event);
|
||||
if (event.event() != null && event.event().equals(SseServerTransport.ENDPOINT_EVENT_TYPE)) {
|
||||
connectLatch.countDown();
|
||||
}
|
||||
else if (event.event() != null && event.event().equals(SseServerTransport.MESSAGE_EVENT_TYPE)) {
|
||||
messageLatch.countDown();
|
||||
}
|
||||
}).subscribe();
|
||||
|
||||
subscriptions.add(subscription);
|
||||
}
|
||||
|
||||
// Wait for all clients to connect
|
||||
try {
|
||||
assertThat(connectLatch.await(5, TimeUnit.SECONDS)).isTrue();
|
||||
}
|
||||
catch (InterruptedException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
|
||||
// Verify initial connections
|
||||
for (List<ServerSentEvent<String>> events : allReceivedEvents) {
|
||||
assertThat(events).hasSize(1);
|
||||
assertThat(events.get(0).data()).isEqualTo(messageEndpoint);
|
||||
}
|
||||
|
||||
// Give clients time to fully establish their subscriptions
|
||||
logger.debug("Waiting for subscriptions to stabilize...");
|
||||
try {
|
||||
Thread.sleep(1000);
|
||||
}
|
||||
catch (InterruptedException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
logger.debug("Sending broadcast message to {} clients", clientCount);
|
||||
|
||||
// Send broadcast message
|
||||
JSONRPCRequest broadcastMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "broadcast", "broadcast-id",
|
||||
Map.of("message", "Hello all!"));
|
||||
|
||||
// Send the message
|
||||
transport.sendMessage(broadcastMessage).block(Duration.ofSeconds(5));
|
||||
|
||||
// Wait for all clients to receive the broadcast
|
||||
try {
|
||||
assertThat(messageLatch.await(5, TimeUnit.SECONDS)).isTrue();
|
||||
}
|
||||
catch (InterruptedException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
|
||||
// Verify each client received both messages
|
||||
for (List<ServerSentEvent<String>> events : allReceivedEvents) {
|
||||
assertThat(events).hasSize(2);
|
||||
assertThat(events.get(0).data()).isEqualTo(messageEndpoint);
|
||||
assertThat(events.get(1).data()).contains("broadcast-id");
|
||||
}
|
||||
|
||||
// Cleanup
|
||||
subscriptions.forEach(Disposable::dispose);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGracefulShutdown() {
|
||||
// Connect a client
|
||||
Flux<ServerSentEvent<String>> eventStream = webTestClient.get()
|
||||
.uri("/sse")
|
||||
.accept(MediaType.TEXT_EVENT_STREAM)
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isOk()
|
||||
.returnResult(String.class)
|
||||
.getResponseBody()
|
||||
.map(data -> ServerSentEvent.<String>builder().data(data).build());
|
||||
|
||||
List<ServerSentEvent<String>> receivedEvents = new ArrayList<>();
|
||||
eventStream.subscribe(receivedEvents::add);
|
||||
|
||||
// Wait for connection
|
||||
StepVerifier.create(Mono.delay(Duration.ofMillis(500))).expectNextCount(1).verifyComplete();
|
||||
|
||||
// Initiate shutdown
|
||||
transport.closeGracefully().block(Duration.ofSeconds(5));
|
||||
|
||||
// Verify server rejects new connections with timeout
|
||||
webTestClient.get()
|
||||
.uri("/sse")
|
||||
.accept(MediaType.TEXT_EVENT_STREAM)
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isEqualTo(HttpStatus.SERVICE_UNAVAILABLE)
|
||||
.expectBody(String.class)
|
||||
.isEqualTo("Server is shutting down");
|
||||
|
||||
// Verify server rejects new messages with timeout
|
||||
webTestClient.post()
|
||||
.uri(messageEndpoint)
|
||||
.contentType(MediaType.APPLICATION_JSON)
|
||||
.bodyValue("""
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"method": "test",
|
||||
"id": "1",
|
||||
"params": {}
|
||||
}
|
||||
""")
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isEqualTo(HttpStatus.SERVICE_UNAVAILABLE)
|
||||
.expectBody(String.class)
|
||||
.isEqualTo("Server is shutting down");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testInvalidMessageHandling() {
|
||||
// Test invalid JSON
|
||||
webTestClient.post()
|
||||
.uri(messageEndpoint)
|
||||
.contentType(MediaType.APPLICATION_JSON)
|
||||
.bodyValue("invalid json")
|
||||
.exchange()
|
||||
.expectStatus()
|
||||
.isBadRequest();
|
||||
|
||||
// Test invalid message format
|
||||
webTestClient.post().uri(messageEndpoint).contentType(MediaType.APPLICATION_JSON).bodyValue("""
|
||||
{
|
||||
"invalid": "message"
|
||||
}
|
||||
""").exchange().expectStatus().isBadRequest();
|
||||
}
|
||||
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue