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:
Christian Tzolov 2024-12-19 19:38:38 +01:00
commit abd0facf4e
21 changed files with 5727 additions and 471 deletions

View file

@ -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;

View file

@ -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;

View file

@ -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();
}
}

View file

@ -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;

View file

@ -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();
}
}

View file

@ -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();
}
}

View file

@ -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();
}
}

View file

@ -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);
}
}

View file

@ -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() {
}
}

View file

@ -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();
}
}

View file

@ -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();
}
}

View file

@ -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);
}
}

View file

@ -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();
}
}

View file

@ -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();
}
}

View file

@ -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();
}
}
}

View file

@ -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();
}
}
}

View file

@ -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());
}
}

View file

@ -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();
}
}
}

View file

@ -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();
}
}

View file

@ -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