modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Clean the code and add more tests

Browse source
This commit is contained in:
Christian Tzolov 2024-12-09 23:05:23 +01:00
commit 68c21d2264
20 changed files with 312 additions and 732 deletions
Download patch file Download diff file Expand all files Collapse all files

7
mcp/pom.xml
View file

@ -55,6 +55,11 @@
<version>5.11.0</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>io.projectreactor</groupId>
<artifactId>reactor-test</artifactId>
<scope>test</scope>
</dependency>
<!-- <dependency>
<groupId>org.springframework.ai</groupId>
@ -114,4 +119,4 @@
</repository>
</repositories>
</project>
</project>

1
mcp/src/main/java/spring/ai/mcp/client/McpAsyncClient.java
View file

@ -5,6 +5,7 @@ import java.time.Duration;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import reactor.core.publisher.Mono;
import spring.ai.mcp.spec.McpAsyncSession;
import spring.ai.mcp.spec.McpAsyncTransport;
import spring.ai.mcp.spec.McpSchema;

21
mcp/src/main/java/spring/ai/mcp/client/McpClient.java
View file

@ -16,30 +16,9 @@
package spring.ai.mcp.client;
import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import spring.ai.mcp.spec.McpAsyncTransport;
import spring.ai.mcp.spec.McpTransport;
import spring.ai.mcp.spec.McpSchema;
import spring.ai.mcp.spec.McpSession;
import spring.ai.mcp.spec.McpTransport;
import spring.ai.mcp.spec.McpSchema.CallToolRequest;
import spring.ai.mcp.spec.McpSchema.CallToolResult;
import spring.ai.mcp.spec.McpSchema.ClientCapabilities;
import spring.ai.mcp.spec.McpSchema.Implementation;
import spring.ai.mcp.spec.McpSchema.InitializeRequest;
import spring.ai.mcp.spec.McpSchema.InitializeResult;
import spring.ai.mcp.spec.McpSchema.ListResourceTemplatesResult;
import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
import spring.ai.mcp.spec.McpSchema.ListToolsResult;
import spring.ai.mcp.spec.McpSchema.PaginatedRequest;
import spring.ai.mcp.spec.McpSchema.ReadResourceRequest;
import spring.ai.mcp.spec.McpSchema.Resource;
import spring.ai.mcp.spec.McpSchema.SubscribeRequest;
import spring.ai.mcp.spec.McpSchema.UnsubscribeRequest;
import spring.ai.mcp.spec.McpSchema.ClientCapabilities.RootCapabilities;
/**
* The MCP client is the main entry point for interacting with the Model Context Protocol

130
mcp/src/main/java/spring/ai/mcp/client/stdio/ServerParameters.java Normal file
View file

@ -0,0 +1,130 @@
package spring.ai.mcp.client.stdio;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import spring.ai.mcp.client.util.Assert;
/**
* Server parameters for stdio client.
*/
@JsonInclude(JsonInclude.Include.NON_ABSENT)
public class ServerParameters {
// Environment variables to inherit by default
private static final List<String> DEFAULT_INHERITED_ENV_VARS = System.getProperty("os.name")
.toLowerCase()
.contains("win")
? Arrays.asList("APPDATA", "HOMEDRIVE", "HOMEPATH", "LOCALAPPDATA", "PATH",
"PROCESSOR_ARCHITECTURE", "SYSTEMDRIVE", "SYSTEMROOT", "TEMP", "USERNAME", "USERPROFILE")
: Arrays.asList("HOME", "LOGNAME", "PATH", "SHELL", "TERM", "USER");
@JsonProperty("command")
private String command;
@JsonProperty("args")
private List<String> args = new ArrayList<>();
@JsonProperty("env")
private Map<String, String> env;
private ServerParameters(String command, List<String> args, Map<String, String> env) {
Assert.notNull(command, "The command can not be null");
Assert.notNull(args, "The args can not be null");
this.command = command;
this.args = args;
this.env = new HashMap<>(getDefaultEnvironment());
if (env != null && !env.isEmpty()) {
this.env.putAll(env);
}
}
public String getCommand() {
return this.command;
}
public List<String> getArgs() {
return this.args;
}
public Map<String, String> getEnv() {
return this.env;
}
public static Builder builder(String command) {
return new Builder(command);
}
public static class Builder {
private String command;
private List<String> args = new ArrayList<>();
private Map<String, String> env = new HashMap<>();
public Builder(String command) {
Assert.notNull(command, "The command can not be null");
this.command = command;
}
public Builder args(String... args) {
Assert.notNull(args, "The args can not be null");
this.args = Arrays.asList(args);
return this;
}
public Builder args(List<String> args) {
Assert.notNull(args, "The args can not be null");
this.args = new ArrayList<>(args);
return this;
}
public Builder arg(String arg) {
Assert.notNull(arg, "The arg can not be null");
this.args.add(arg);
return this;
}
public Builder env(Map<String, String> env) {
if (env != null && !env.isEmpty()) {
this.env.putAll(env);
}
return this;
}
public Builder addEnvVar(String key, String value) {
Assert.notNull(key, "The key can not be null");
Assert.notNull(value, "The value can not be null");
this.env.put(key, value);
return this;
}
public ServerParameters build() {
return new ServerParameters(command, args, env);
}
}
/**
* Returns a default environment object including only environment variables
* deemed safe to inherit.
*/
private static Map<String, String> getDefaultEnvironment() {
return System.getenv()
.entrySet()
.stream()
.filter(entry -> DEFAULT_INHERITED_ENV_VARS.contains(entry.getKey()))
.filter(entry -> entry.getValue() != null)
.filter(entry -> !entry.getValue().startsWith("()"))
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
}
}

88
mcp/src/main/java/spring/ai/mcp/client/stdio/ServerParametersParser.java Normal file
View file

@ -0,0 +1,88 @@
/*
* 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 spring.ai.mcp.client.stdio;
import java.util.Map;
import java.util.List;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.ObjectMapper;
import spring.ai.mcp.client.stdio.ServerParametersParser.McpServerConfigurations.McpServerConfiguration;
/**
* @author Christian Tzolov
* @since 1.0.0
*/
public class ServerParametersParser {
@JsonInclude(JsonInclude.Include.NON_ABSENT)
public record McpServerConfigurations(// @formatter:off
@JsonProperty("mcpServers") Map<String, McpServerConfiguration> mcpServers) {
@JsonInclude(JsonInclude.Include.NON_ABSENT)
public record McpServerConfiguration(
@JsonProperty("command") String command,
@JsonProperty("args") List<String> args,
@JsonProperty("env") Map<String, String> env) {
}
} // @formatter:on
public static void main(String[] args) throws Exception {
ObjectMapper objectMapper = new ObjectMapper();
// Assuming the JSON is in a file named "servers.json"
String jsonInput = """
{
"mcpServers": {
"filesystem": {
"command": "npx",
"args": ["-y", "@modelcontextprotocol/server-filesystem", "/path/to/allowed/files"]
},
"git": {
"command": "uvx",
"args": ["mcp-server-git", "--repository", "path/to/git/repo"]
},
"github": {
"command": "npx",
"args": ["-y", "@modelcontextprotocol/server-github"],
"env": {
"GITHUB_PERSONAL_ACCESS_TOKEN": "<YOUR_TOKEN>"
}
},
"postgres": {
"command": "npx",
"args": ["-y", "@modelcontextprotocol/server-postgres", "postgresql://localhost/mydb"]
}
}
}
""";
// Deserialize JSON into McpServers
McpServerConfigurations servers = objectMapper.readValue(jsonInput, McpServerConfigurations.class);
// Access individual servers
for (Map.Entry<String, McpServerConfiguration> entry : servers.mcpServers().entrySet()) {
System.out.println("Key: " + entry.getKey());
System.out.println("Command: " + entry.getValue().command());
System.out.println("Arguments: " + String.join(", ", entry.getValue().args()));
System.out.println("Environment Variables: " + entry.getValue().env());
System.out.println("-----------");
}
}
}

160
mcp/src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java
View file

@ -3,12 +3,9 @@ package spring.ai.mcp.client.stdio;
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.Executors;
import com.fasterxml.jackson.core.type.TypeReference;
@ -27,12 +24,13 @@ import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
*/
public class StdioServerTransport extends DefaultMcpTransport {
private final static TypeReference<HashMap<String, Object>> MAP_TYPE_REF = new TypeReference<HashMap<String, Object>>() {
};
private Process process;
private BufferedReader processErrorReader;
private BufferedReader processReader;
private BufferedWriter processWriter;
private Scheduler inboundScheduler;
@ -41,29 +39,28 @@ public class StdioServerTransport extends DefaultMcpTransport {
private volatile boolean isRunning;
private final StdioServerParameters params;
private final ServerParameters params;
public StdioServerTransport(StdioServerParameters params) {
this(params, Duration.ofMillis(100));
public StdioServerTransport(ServerParameters params) {
this(params, new ObjectMapper());
}
public StdioServerTransport(StdioServerParameters params, Duration writeTimeout) {
super(writeTimeout);
public StdioServerTransport(ServerParameters params, ObjectMapper objectMapper) {
super(objectMapper);
Assert.notNull(params, "The params can not be null");
Assert.notNull(objectMapper, "The ObjectMapper can not be null");
this.params = params;
// Start threads
this.inboundScheduler =
Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
"inbound");
this.outboundScheduler =
Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
"outbound");
this.errorScheduler =
Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
"error");
this.inboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
"inbound");
this.outboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
"outbound");
this.errorScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
"error");
}
@Override
@ -79,8 +76,7 @@ public class StdioServerTransport extends DefaultMcpTransport {
// Start the process
try {
this.process = processBuilder.start();
}
catch (IOException e) {
} catch (IOException e) {
throw new RuntimeException("Failed to start process with command: " + fullCommand, e);
}
@ -100,14 +96,12 @@ public class StdioServerTransport extends DefaultMcpTransport {
startInboundProcessing();
startOutboundProcessing();
startErrorProcessing();
}
public void awaitForExit() {
try {
this.process.waitFor();
}
catch (InterruptedException e) {
} catch (InterruptedException e) {
throw new RuntimeException("Process interrupted", e);
}
}
@ -121,56 +115,20 @@ public class StdioServerTransport extends DefaultMcpTransport {
System.out.println("Received error line: " + line);
// TODO: handle errors, etc.
this.getErrorSink().tryEmitNext(line);
}
catch (Exception e) {
} catch (Exception e) {
throw new RuntimeException(e);
}
}
}
catch (IOException e) {
} catch (IOException e) {
if (this.isRunning) {
throw new RuntimeException(e);
}
}
finally {
} finally {
this.isRunning = false;
}
});
}
private static ObjectMapper OBJECT_MAPPER = new ObjectMapper();
private static TypeReference<HashMap<String, Object>> MAP_TYPE_REF = new TypeReference<HashMap<String, Object>>() {
};
public static Map<String, Object> jsonToMap(String json) {
try {
return OBJECT_MAPPER.readValue(json, MAP_TYPE_REF);
}
catch (Exception e) {
throw new RuntimeException(e);
}
}
private JSONRPCMessage deserializeJsonRpcMessage(String jsonText) throws IOException {
var map = jsonToMap(jsonText);
// Determine message type based on specific JSON structure
if (map.containsKey("method") && map.containsKey("id")) {
return OBJECT_MAPPER.convertValue(map, JSONRPCRequest.class);
}
else if (map.containsKey("method") && !map.containsKey("id")) {
return OBJECT_MAPPER.convertValue(map, JSONRPCNotification.class);
}
else if (map.containsKey("result") || map.containsKey("error")) {
return OBJECT_MAPPER.convertValue(map, JSONRPCResponse.class);
}
throw new IllegalArgumentException("Cannot deserialize JSONRPCMessage: " + jsonText);
}
private void startInboundProcessing() {
this.inboundScheduler.schedule(() -> {
try {
@ -182,43 +140,55 @@ public class StdioServerTransport extends DefaultMcpTransport {
// TODO: Back off, reschedule, give up?
throw new RuntimeException("Failed to enqueue message");
}
}
catch (Exception e) {
} catch (Exception e) {
throw new RuntimeException(e);
}
}
}
catch (IOException e) {
} catch (IOException e) {
if (isRunning) {
throw new RuntimeException(e);
}
}
finally {
} finally {
isRunning = false;
}
});
}
private JSONRPCMessage deserializeJsonRpcMessage(String jsonText) throws IOException {
var map = this.objectMapper.readValue(jsonText, MAP_TYPE_REF);
// Determine message type based on specific JSON structure
if (map.containsKey("method") && map.containsKey("id")) {
return this.objectMapper.convertValue(map, JSONRPCRequest.class);
} else if (map.containsKey("method") && !map.containsKey("id")) {
return this.objectMapper.convertValue(map, JSONRPCNotification.class);
} else if (map.containsKey("result") || map.containsKey("error")) {
return this.objectMapper.convertValue(map, JSONRPCResponse.class);
}
throw new IllegalArgumentException("Cannot deserialize JSONRPCMessage: " + jsonText);
}
private void startOutboundProcessing() {
this.getOutboundSink()
.asFlux()
// this bit is important since writes come from user threads and we
// want to ensure that the actual writing happens on a dedicated thread
.publishOn(outboundScheduler)
.handle((message, s) -> {
if (message != null) {
try {
this.processWriter.write(OBJECT_MAPPER.writeValueAsString(message));
this.processWriter.newLine();
this.processWriter.flush();
s.next(message);
}
catch (IOException e) {
s.error(new RuntimeException(e));
}
}
})
.subscribe();
.asFlux()
// this bit is important since writes come from user threads and we
// want to ensure that the actual writing happens on a dedicated thread
.publishOn(outboundScheduler)
.handle((message, s) -> {
if (message != null) {
try {
this.processWriter.write(objectMapper.writeValueAsString(message));
this.processWriter.newLine();
this.processWriter.flush();
s.next(message);
} catch (IOException e) {
s.error(new RuntimeException(e));
}
}
})
.subscribe();
}
// TODO: provide a non-blocking variant with graceful option
@ -235,24 +205,30 @@ public class StdioServerTransport extends DefaultMcpTransport {
@Override
public void close() {
stop();
this.stop();
super.close(); // Do we need this?
// Close resources
if (this.processReader != null) {
try {
this.processReader.close();
}
catch (IOException e) {
} catch (IOException e) {
e.printStackTrace();
}
}
if (this.processWriter != null) {
try {
this.processWriter.close();
} catch (IOException e) {
e.printStackTrace();
}
catch (IOException e) {
}
if (this.processErrorReader != null) {
try {
this.processErrorReader.close();
} catch (IOException e) {
e.printStackTrace();
}
}

36
mcp/src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java
View file

@ -15,9 +15,9 @@
*/
package spring.ai.mcp.spec;
import java.time.Duration;
import java.util.function.Consumer;
import com.fasterxml.jackson.databind.ObjectMapper;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import spring.ai.mcp.client.util.Assert;
@ -29,19 +29,24 @@ import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
*/
public class DefaultMcpTransport implements McpAsyncTransport {
protected final ObjectMapper objectMapper;
private final Sinks.Many<String> errorSink;
private final Sinks.Many<JSONRPCMessage> inboundSink;
private final Sinks.Many<JSONRPCMessage> outboundSink;
private final Duration writeTimeout;
private Consumer<JSONRPCMessage> inboundMessageHandler = message -> System.out
.println("Received message: " + message);
private Consumer<String> errorHandler = error -> System.err.println("Received error: " + error);
private Consumer<JSONRPCMessage> messageReader = message -> System.out.println("Received message: " + message);
public DefaultMcpTransport() {
this(new ObjectMapper());
}
private Consumer<String> errorReader = error -> System.err.println("Received error: " + error);
public DefaultMcpTransport(ObjectMapper objectMapper) {
public DefaultMcpTransport(Duration readTimeout) {
Assert.notNull(objectMapper, "ObjectMapper must not be null");
this.objectMapper = objectMapper;
// TODO: consider the effects of buffering here -> the inter-process pipes are
// independent and the notifications can flood the client/server.
@ -52,19 +57,22 @@ public class DefaultMcpTransport implements McpAsyncTransport {
this.errorSink = Sinks.many().unicast().onBackpressureBuffer();
this.inboundSink = Sinks.many().unicast().onBackpressureBuffer();
this.outboundSink = Sinks.many().unicast().onBackpressureBuffer();
this.writeTimeout = readTimeout;
this.handleIncomingMessages();
}
public ObjectMapper getObjectMapper() {
return this.objectMapper;
}
private void handleIncomingMessages() {
this.inboundSink.asFlux()
.subscribe(message -> this.messageReader.accept(message));
.subscribe(message -> this.inboundMessageHandler.accept(message));
}
private void handleIncomingErrors() {
this.errorSink.asFlux().subscribe(e -> {
this.errorReader.accept(e);
this.errorHandler.accept(e);
System.err.println(e);
});
}
@ -81,12 +89,12 @@ public class DefaultMcpTransport implements McpAsyncTransport {
return errorSink;
}
public void setMessageHandler(Consumer<JSONRPCMessage> messageReader) {
this.messageReader = messageReader;
public void setInboudMessageHandler(Consumer<JSONRPCMessage> inboundMessageHandler) {
this.inboundMessageHandler = inboundMessageHandler;
}
public void setErrorHandler(Consumer<String> errorReader) {
this.errorReader = errorReader;
public void setErrorHandler(Consumer<String> errorHandler) {
this.errorHandler = errorHandler;
}
// Start processing incoming messages

121
mcp/src/main/java/spring/ai/mcp/spec/McpAsyncSession.java Normal file
View file

@ -0,0 +1,121 @@
package spring.ai.mcp.spec;
import java.time.Duration;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import reactor.core.publisher.Mono;
import reactor.core.publisher.MonoSink;
import spring.ai.mcp.client.util.Assert;
public class McpAsyncSession {
private final ConcurrentHashMap<Object, MonoSink<McpSchema.JSONRPCResponse>> pendingResponses = new ConcurrentHashMap<>();
private final Duration requestTimeout;
private final ObjectMapper objectMapper;
private final McpAsyncTransport transport;
public McpAsyncSession(Duration requestTimeout, ObjectMapper objectMapper, McpAsyncTransport transport) {
Assert.notNull(objectMapper, "The ObjectMapper can not be null");
Assert.notNull(requestTimeout, "The requstTimeout can not be null");
Assert.notNull(transport, "The transport can not be null");
this.requestTimeout = requestTimeout;
this.objectMapper = objectMapper;
this.transport = transport;
this.transport.setInboudMessageHandler(message -> {
switch (message) {
case McpSchema.JSONRPCResponse response -> {
var sink = pendingResponses.remove(response.id());
if (sink == null) {
System.out.println("Unexpected response for unkown id " + response.id());
} else {
sink.success(response);
}
}
case McpSchema.JSONRPCRequest request -> {
System.out.println("Client does not yet support server requests");
}
case McpSchema.JSONRPCNotification notification -> {
System.out.println("Notifications not yet supported");
}
}
});
this.transport.setErrorHandler(error -> System.out.println("Error received " + error));
this.transport.start();
}
public <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
String requestId = UUID.randomUUID().toString();
return Mono.<McpSchema.JSONRPCResponse>create(sink -> {
this.pendingResponses.put(requestId, sink);
McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method,
requestId, requestParams);
try {
// TODO: This is non-blocking, but it's actually a synchronous call,
// perhaps there's no need to make it return Mono?
this.transport.sendMessage(jsonrpcRequest)
// TODO: It's most efficient to create a dedicated
// Subscriber here
.subscribe(v -> {
}, e -> {
this.pendingResponses.remove(requestId);
sink.error(e);
});
} catch (Exception e) {
sink.error(e);
}
}).timeout(this.requestTimeout).handle((jsonRpcResponse, s) -> {
if (jsonRpcResponse.error() != null) {
s.error(new McpError(jsonRpcResponse.error()));
} else {
if (typeRef.getType().getTypeName().equals("java.lang.Void")) {
s.complete();
} else {
s.next(this.objectMapper.convertValue(jsonRpcResponse.result(), typeRef));
}
}
});
}
public static class McpError extends RuntimeException {
public McpError(Object error) {
super(error.toString());
}
}
public Mono<Void> sendNotification(String method) {
return sendNotification(method, null);
}
public Mono<Void> sendNotification(String method, Map<String, Object> params) {
McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
method, params);
try {
// TODO: make it non-blocking
this.transport.sendMessage(jsonrpcNotification);
} catch (Exception e) {
return Mono.error(new McpError(e));
}
return Mono.empty();
}
public Mono<Void> closeGracefully(Duration timeout) {
// TODO handle the timeout in transport
return Mono.fromRunnable(this.transport::close);
}
}

2
mcp/src/main/java/spring/ai/mcp/spec/McpAsyncTransport.java
View file

@ -31,7 +31,7 @@ public interface McpAsyncTransport extends AutoCloseable{
void close();
void setMessageHandler(Consumer<JSONRPCMessage> messageReader);
void setInboudMessageHandler(Consumer<JSONRPCMessage> messageReader);
void setErrorHandler(Consumer<String> errorReader);

155
mcp/src/main/java/spring/ai/mcp/spec/McpSession.java
View file

@ -1,155 +0,0 @@
/*
* 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 spring.ai.mcp.spec;
import java.time.Duration;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import spring.ai.mcp.client.util.Assert;
import spring.ai.mcp.spec.McpSchema.JSONRPCNotification;
import spring.ai.mcp.spec.McpSchema.JSONRPCRequest;
import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
public class McpSession implements AutoCloseable {
final ConcurrentHashMap<Object, CompletableFuture<JSONRPCResponse>> waitingForResponseStreams;
private final Duration readTimeout;
private final ObjectMapper objectMapper;
private final McpTransport transport;
public McpSession(McpTransport transport, Duration readTimeout, ObjectMapper objectMapper) {
Assert.notNull(transport, "The transport can not be null");
Assert.notNull(readTimeout, "The readTimeout can not be null");
Assert.notNull(objectMapper, "The objectMapper can not be null");
this.waitingForResponseStreams = new ConcurrentHashMap<>();
this.readTimeout = readTimeout;
this.objectMapper = objectMapper;
this.transport = transport;
this.transport.setMessageHandler(message -> {
if (message instanceof JSONRPCResponse jsonRpcResponse) {
CompletableFuture<JSONRPCResponse> responseFuture = this.waitingForResponseStreams
.remove(jsonRpcResponse.id());
if (responseFuture != null) {
responseFuture.complete(jsonRpcResponse);
}
} else if (message instanceof JSONRPCRequest jsonRpcRequest) {
handleRequest(new RequestResponder<>(jsonRpcRequest.id(), jsonRpcRequest));
} else if (message instanceof JSONRPCNotification jsonRpcNotification) {
handleNotification(jsonRpcNotification);
}
});
this.transport.setErrorHandler(error -> {
System.out.println("Received error: " + error);
});
this.transport.start();
}
public <T> CompletableFuture<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
String requestId = UUID.randomUUID().toString();
CompletableFuture<JSONRPCResponse> responseFuture = new CompletableFuture<>();
this.waitingForResponseStreams.put(requestId, responseFuture);
JSONRPCRequest jsonrpcRequest = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, method, requestId, requestParams);
try {
this.transport.sendMessage(jsonrpcRequest);
} catch (Exception e) {
return CompletableFuture.failedFuture(new McpError(e));
}
return responseFuture.orTimeout(this.readTimeout.toMillis(), TimeUnit.MILLISECONDS)
.thenApply(jsonRpcResponse -> {
if (jsonRpcResponse.error() != null) {
throw new McpError(jsonRpcResponse.error());
}
return this.objectMapper.convertValue(jsonRpcResponse.result(), typeRef);
});
}
public CompletableFuture<Void> sendNotification(String method) {
return sendNotification(method, null);
}
public CompletableFuture<Void> sendNotification(String method, Map<String, Object> params) {
JSONRPCNotification jsonrpcNotification = new JSONRPCNotification(McpSchema.JSONRPC_VERSION, method, params);
try {
this.transport.sendMessage(jsonrpcNotification);
} catch (Exception e) {
return CompletableFuture.failedFuture(new McpError(e));
}
return CompletableFuture.completedFuture(null);
}
protected void handleRequest(RequestResponder<JSONRPCRequest> responder) {
// Override to handle requests
System.out.println("Handling request: " + responder);
}
protected void handleNotification(JSONRPCNotification notification) {
// Override to handle notifications
System.out.println("Handling notification: " + notification);
}
@Override
public void close() {
this.transport.close();
}
public static class RequestResponder<R> {
private final Object requestId;
private final R request;
public RequestResponder(Object requestId, R request) {
this.requestId = requestId;
this.request = request;
}
public CompletableFuture<Void> respond(Object response) {
// Send response
return CompletableFuture.completedFuture(null);
}
}
public static class McpError extends RuntimeException {
public McpError(Object error) {
super(error.toString());
}
}
}

38
mcp/src/main/java/spring/ai/mcp/spec/McpTransport.java
View file

@ -1,38 +0,0 @@
/*
* 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 spring.ai.mcp.spec;
import java.util.function.Consumer;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
/**
* @author Christian Tzolov
* @since 1.0.0
*/
public interface McpTransport extends AutoCloseable{
void start();
void close();
void setMessageHandler(Consumer<JSONRPCMessage> messageReader);
void setErrorHandler(Consumer<String> errorReader);
void sendMessage(JSONRPCMessage message);
}

6
mcp/src/test/java/spring/ai/mcp/McpAsyncClientTests.java
View file

@ -25,7 +25,7 @@ import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import spring.ai.mcp.client.McpAsyncClient;
import spring.ai.mcp.client.McpClient;
import spring.ai.mcp.client.stdio.StdioServerParameters;
import spring.ai.mcp.client.stdio.ServerParameters;
import spring.ai.mcp.client.stdio.StdioServerTransport;
import spring.ai.mcp.spec.McpSchema.CallToolRequest;
import spring.ai.mcp.spec.McpSchema.Resource;
@ -45,7 +45,7 @@ class McpAsyncClientTests {
private McpAsyncClient mcpAsyncClient;
private StdioServerParameters stdioParams;
private ServerParameters stdioParams;
private static final Duration TIMEOUT = Duration.ofSeconds(10);
@ -53,7 +53,7 @@ class McpAsyncClientTests {
@BeforeEach
void setUp() {
stdioParams = StdioServerParameters.builder("npx")
stdioParams = ServerParameters.builder("npx")
.args("-y", "@modelcontextprotocol/server-everything", "dir")
.build();

28
mcp/src/test/java/spring/ai/mcp/McpSyncClientTests.java
View file

@ -16,6 +16,7 @@
package spring.ai.mcp;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.AfterEach;
@ -28,14 +29,16 @@ import static org.assertj.core.api.Assertions.*;
import com.fasterxml.jackson.databind.ObjectMapper;
import spring.ai.mcp.client.McpClient;
import spring.ai.mcp.client.McpSyncClient;
import spring.ai.mcp.client.stdio.StdioServerParameters;
import spring.ai.mcp.client.stdio.ServerParameters;
import spring.ai.mcp.client.stdio.StdioServerTransport;
import spring.ai.mcp.spec.McpSchema.CallToolRequest;
import spring.ai.mcp.spec.McpSchema.CallToolResult;
import spring.ai.mcp.spec.McpSchema.Content;
import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
import spring.ai.mcp.spec.McpSchema.ListToolsResult;
import spring.ai.mcp.spec.McpSchema.Tool;
import spring.ai.mcp.spec.McpSchema.Resource;
import spring.ai.mcp.spec.McpSchema.TextContent;
/**
* Unit tests for MCP Client Session functionality.
@ -47,7 +50,7 @@ class McpSyncClientTests {
private McpSyncClient mcpSyncClient;
private StdioServerParameters stdioParams;
private ServerParameters stdioParams;
private static final Duration TIMEOUT = Duration.ofSeconds(10);
@ -55,7 +58,7 @@ class McpSyncClientTests {
@BeforeEach
void setUp() {
stdioParams = StdioServerParameters.builder("npx")
stdioParams = ServerParameters.builder("npx")
.args("-y", "@modelcontextprotocol/server-everything", "dir")
.build();
@ -98,6 +101,25 @@ class McpSyncClientTests {
});
}
@Test
@Timeout(15) // Giving extra time beyond the client timeout
void testCallTools() {
CallToolResult toolResult = mcpSyncClient.callTool(new CallToolRequest("add", Map.of("a", 3, "b", 4)));
assertThat(toolResult).isNotNull().satisfies(result -> {
assertThat(result.content()).hasSize(1);
TextContent content = (TextContent) result.content().get(0);
assertThat(content).isNotNull();
assertThat(content.text()).isNotNull();
assertThat(content.text()).contains("7");
System.out.println(content);
});
}
@Test
@Timeout(15)
void testPing() {

4
mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests2.java
View file

@ -20,7 +20,7 @@ import java.time.Duration;
import com.fasterxml.jackson.databind.ObjectMapper;
import spring.ai.mcp.client.McpClient;
import spring.ai.mcp.client.McpSyncClient;
import spring.ai.mcp.client.stdio.StdioServerParameters;
import spring.ai.mcp.client.stdio.ServerParameters;
import spring.ai.mcp.client.stdio.StdioServerTransport;
import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
import spring.ai.mcp.spec.McpSchema.ListToolsResult;
@ -34,7 +34,7 @@ public class ClientSessionTests2 {
public static void main(String[] args) {
var stdioParams = StdioServerParameters.builder("uv")
var stdioParams = ServerParameters.builder("uv")
.args("--directory", "dir", "run", "mcp-server-sqlite",
"--db-path", "~/test.db")
.build();

4
mcp/src/test/java/spring/ai/mcp/attic/ClientSessionTests3.java
View file

@ -20,7 +20,7 @@ import java.time.Duration;
import com.fasterxml.jackson.databind.ObjectMapper;
import spring.ai.mcp.client.McpClient;
import spring.ai.mcp.client.McpSyncClient;
import spring.ai.mcp.client.stdio.StdioServerParameters;
import spring.ai.mcp.client.stdio.ServerParameters;
import spring.ai.mcp.client.stdio.StdioServerTransport;
import spring.ai.mcp.spec.McpSchema.ListResourcesResult;
import spring.ai.mcp.spec.McpSchema.ListToolsResult;
@ -34,7 +34,7 @@ public class ClientSessionTests3 {
public static void main(String[] args) {
var stdioParams = StdioServerParameters.builder("uv")
var stdioParams = ServerParameters.builder("uv")
.args("--directory", "dir", "run", "mcp-server-sqlite",
"--db-path", "~/test.db")
.build();

161
mcp/src/test/java/spring/ai/mcp/spec/DefaultMcpTransportTests.java Normal file
View file

@ -0,0 +1,161 @@
/*
* 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 spring.ai.mcp.spec;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
import spring.ai.mcp.spec.McpSchema.JSONRPCRequest;
import java.util.concurrent.atomic.AtomicReference;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
class DefaultMcpTransportTests {
private DefaultMcpTransport transport;
private ObjectMapper objectMapper;
@BeforeEach
void setUp() {
objectMapper = new ObjectMapper();
transport = new DefaultMcpTransport(objectMapper);
}
@Test
void constructorShouldValidateObjectMapper() {
assertThatThrownBy(() -> new DefaultMcpTransport(null))
.isInstanceOf(IllegalArgumentException.class)
.hasMessage("ObjectMapper must not be null");
}
@Test
void defaultConstructorShouldCreateValidInstance() {
DefaultMcpTransport defaultTransport = new DefaultMcpTransport();
assertThat(defaultTransport.getObjectMapper()).isNotNull();
}
@Test
void customObjectMapperConstructorShouldCreateValidInstance() {
assertThat(transport.getObjectMapper()).isSameAs(objectMapper);
}
@Test
void sendMessageShouldEmitToOutboundSink() {
JSONRPCRequest message = new JSONRPCRequest(
"2.0",
"test",
1,
null);
AtomicReference<JSONRPCMessage> receivedMessage = new AtomicReference<>();
transport.getOutboundSink().asFlux().subscribe(receivedMessage::set);
Mono<Void> result = transport.sendMessage(message);
StepVerifier.create(result)
.verifyComplete();
assertThat(receivedMessage.get())
.isNotNull()
.satisfies(msg -> {
assertThat(((JSONRPCRequest) msg).jsonrpc()).isEqualTo("2.0");
assertThat(((JSONRPCRequest) msg).method()).isEqualTo("test");
});
}
@Test
void customInboundMessageHandlerShouldReceiveMessages() {
AtomicReference<JSONRPCMessage> receivedMessage = new AtomicReference<>();
transport.setInboudMessageHandler(receivedMessage::set);
transport.start();
JSONRPCRequest message = new JSONRPCRequest(
"2.0",
"test",
1,
null);
transport.getInboundSink().tryEmitNext(message);
assertThat(receivedMessage.get())
.isNotNull()
.satisfies(msg -> {
assertThat(((JSONRPCRequest) msg).jsonrpc()).isEqualTo("2.0");
assertThat(((JSONRPCRequest) msg).method()).isEqualTo("test");
});
}
@Test
void customErrorHandlerShouldReceiveErrors() {
AtomicReference<String> receivedError = new AtomicReference<>();
transport.setErrorHandler(receivedError::set);
transport.start();
String errorMessage = "Test error";
transport.getErrorSink().tryEmitNext(errorMessage);
assertThat(receivedError.get())
.isNotNull()
.isEqualTo(errorMessage);
}
@Test
void startShouldInitializeMessageAndErrorHandling() {
AtomicReference<JSONRPCMessage> receivedMessage = new AtomicReference<>();
AtomicReference<String> receivedError = new AtomicReference<>();
transport.setInboudMessageHandler(receivedMessage::set);
transport.setErrorHandler(receivedError::set);
transport.start();
JSONRPCRequest message = new JSONRPCRequest(
"2.0",
"test",
1,
null);
String errorMessage = "Test error";
transport.getInboundSink().tryEmitNext(message);
transport.getErrorSink().tryEmitNext(errorMessage);
assertThat(receivedMessage.get()).isNotNull();
assertThat(receivedError.get()).isEqualTo(errorMessage);
}
@Test
void sendMessageShouldReturnErrorWhenSinkFails() {
JSONRPCRequest message = new JSONRPCRequest(
"2.0",
"test",
1,
null);
// Close the sink to simulate failure
transport.getOutboundSink().tryEmitComplete();
Mono<Void> result = transport.sendMessage(message);
StepVerifier.create(result)
.verifyErrorSatisfies(error -> assertThat(error)
.isInstanceOf(RuntimeException.class)
.hasMessage("Failed to enqueue message"));
}
}

158
mcp/src/test/java/spring/ai/mcp/spec/McpSessionJSONRPCTest.java
View file

@ -1,158 +0,0 @@
/*
* 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 spring.ai.mcp.spec;
import com.fasterxml.jackson.core.type.TypeReference;
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.mockito.ArgumentCaptor;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
import spring.ai.mcp.spec.McpSchema.JSONRPCNotification;
import spring.ai.mcp.spec.McpSchema.JSONRPCRequest;
import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.Mockito.*;
class McpSessionJSONRPCTest {
private McpTransport mockTransport;
private ObjectMapper objectMapper;
private McpSession mcpSession;
private ArgumentCaptor<Consumer<JSONRPCMessage>> messageHandlerCaptor;
@BeforeEach
void setUp() {
mockTransport = mock(McpTransport.class);
objectMapper = new ObjectMapper();
mcpSession = new McpSession(mockTransport, Duration.ofSeconds(5), objectMapper);
// Capture the message handler set during the session initialization
messageHandlerCaptor = ArgumentCaptor.forClass(Consumer.class);
verify(mockTransport).setMessageHandler(messageHandlerCaptor.capture());
}
@Test
void sendRequestShouldSendValidRequestMessage() {
String method = "testMethod";
Map<String, Object> params = Map.of("key", "value");
// Call sendRequest
CompletableFuture<String> future = mcpSession.sendRequest(method, params, new TypeReference<>() {});
// Capture the JSONRPCRequest sent to the transport
ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
verify(mockTransport).sendMessage(messageCaptor.capture());
// Assert the message details
JSONRPCRequest request = (JSONRPCRequest) messageCaptor.getValue();
assertEquals(method, request.method());
assertEquals("value", ((Map<?, ?>) request.params()).get("key"));
assertNotNull(request.id());
}
@Test
void sendNotificationShouldSendValidNotificationMessage() {
String method = "notifyMethod";
Map<String, Object> params = Map.of("param1", "value1");
// Call sendNotification
mcpSession.sendNotification(method, params);
// Capture the JSONRPCNotification sent to the transport
ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
verify(mockTransport).sendMessage(messageCaptor.capture());
// Assert the notification details
JSONRPCNotification notification = (JSONRPCNotification) messageCaptor.getValue();
assertEquals(method, notification.method());
assertEquals("value1", ((Map<?, ?>) notification.params()).get("param1"));
}
@Test
void messageHandlerShouldProcessResponseCorrectly() {
// Retrieve the captured message handler
Consumer<JSONRPCMessage> messageHandler = messageHandlerCaptor.getValue();
// Simulate a JSONRPCResponse
JSONRPCResponse response = new JSONRPCResponse("2.0", "1", "resultValue", null);
CompletableFuture<JSONRPCResponse> responseFuture = new CompletableFuture<>();
mcpSession.waitingForResponseStreams.put("1", responseFuture);
// Trigger the handler with the response
messageHandler.accept(response);
// Verify the response is completed
assertTrue(responseFuture.isDone());
assertEquals("resultValue", responseFuture.join().result());
}
@Test
@Disabled
void messageHandlerShouldInvokeHandleRequestForRequestMessage() {
// Retrieve the captured message handler
Consumer<JSONRPCMessage> messageHandler = messageHandlerCaptor.getValue();
// Mock request data
JSONRPCRequest request = new JSONRPCRequest("2.0", "testMethod", "1", Map.of("key", "value"));
// Override the handleRequest to check invocation
McpSession spySession = spy(mcpSession);
doNothing().when(spySession).handleRequest(any());
// Trigger the handler with a request
messageHandler.accept(request);
// Verify handleRequest was called
verify(spySession).handleRequest(any());
}
@Test
@Disabled
void messageHandlerShouldInvokeHandleNotificationForNotificationMessage() {
// Retrieve the captured message handler
Consumer<JSONRPCMessage> messageHandler = messageHandlerCaptor.getValue();
// Mock notification data
JSONRPCNotification notification = new JSONRPCNotification("2.0", "notifyMethod", Map.of("key", "value"));
// Override the handleNotification to check invocation
McpSession spySession = spy(mcpSession);
doNothing().when(spySession).handleNotification(any());
// Trigger the handler with a notification
messageHandler.accept(notification);
// Verify handleNotification was called
verify(spySession).handleNotification(notification);
}
@Test
void closeShouldInvokeTransportClose() {
// Call close on the session
mcpSession.close();
// Verify transport close was called
verify(mockTransport).close();
}
}

93
mcp/src/test/java/spring/ai/mcp/spec/McpSessionTest.java
View file

@ -1,93 +0,0 @@
/*
* 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 spring.ai.mcp.spec;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import com.fasterxml.jackson.core.type.TypeReference;
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 spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
class McpSessionTest {
private McpTransport mockTransport;
private ObjectMapper mockObjectMapper;
private McpSession mcpSession;
@BeforeEach
void setUp() {
mockTransport = mock(McpTransport.class);
mockObjectMapper = mock(ObjectMapper.class);
mcpSession = new McpSession(mockTransport, Duration.ofSeconds(5), mockObjectMapper);
}
@Test
void constructorShouldThrowIfArgumentsAreNull() {
assertThrows(IllegalArgumentException.class,
() -> new McpSession(null, Duration.ofSeconds(5), mockObjectMapper));
assertThrows(IllegalArgumentException.class, () -> new McpSession(mockTransport, null, mockObjectMapper));
assertThrows(IllegalArgumentException.class, () -> new McpSession(mockTransport, Duration.ofSeconds(5), null));
}
@Test
@Disabled
void sendRequestShouldSendJsonRpcRequest() throws Exception {
String method = "testMethod";
Object params = Map.of("key", "value");
TypeReference<String> typeRef = new TypeReference<>() {
};
CompletableFuture<McpSchema.JSONRPCResponse> mockResponseFuture = new CompletableFuture<>();
mockResponseFuture.complete(new JSONRPCResponse("2.0", "1", "result", null));
// when(mockTransport.sendMessage(any())).then(invocation -> null);
when(mockObjectMapper.convertValue("result", typeRef)).thenReturn("result");
CompletableFuture<String> result = mcpSession.sendRequest(method, params, typeRef);
assertEquals("result", result.get());
}
@Test
void sendRequestShouldHandleTimeout() {
String method = "timeoutMethod";
Object params = Map.of("key", "value");
TypeReference<String> typeRef = new TypeReference<>() {
};
CompletableFuture<String> result = mcpSession.sendRequest(method, params, typeRef);
assertThrows(Exception.class, result::get);
}
@Test
void closeShouldInvokeTransportClose() {
mcpSession.close();
verify(mockTransport, times(1)).close();
}
}

128
mcp/src/test/java/spring/ai/mcp/spec/McpTransportTest.java
View file

@ -1,128 +0,0 @@
/*
* 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 spring.ai.mcp.spec;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
import spring.ai.mcp.spec.McpSchema.JSONRPCNotification;
import spring.ai.mcp.spec.McpSchema.JSONRPCRequest;
import spring.ai.mcp.spec.McpSchema.JSONRPCResponse;
import java.time.Duration;
import java.util.Map;
import java.util.function.Consumer;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import static org.junit.jupiter.api.Assertions.*;
import static org.mockito.Mockito.*;
class McpTransportTest {
private McpTransport mockTransport;
private McpSession mcpSession;
@BeforeEach
void setUp() {
mockTransport = mock(McpTransport.class);
mcpSession = new McpSession(mockTransport, Duration.ofSeconds(5), new ObjectMapper());
}
@Test
void constructorShouldInitializeTransport() {
// Verify transport setup during construction
verify(mockTransport).setMessageHandler(any(Consumer.class));
verify(mockTransport).setErrorHandler(any(Consumer.class));
verify(mockTransport).start();
}
@Test
void sendRequestShouldCallTransportSendMessage() {
String method = "testMethod";
Object params = Map.of("key", "value");
// Call sendRequest
mcpSession.sendRequest(method, params, new TypeReference<String>() {});
// Capture the message sent to transport
ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
verify(mockTransport).sendMessage(messageCaptor.capture());
// Verify the message details
JSONRPCRequest request = (JSONRPCRequest) messageCaptor.getValue();
assertEquals("testMethod", request.method());
assertEquals("value", ((Map<?, ?>) request.params()).get("key"));
assertNotNull(request.id());
}
@Test
void sendNotificationShouldCallTransportSendMessage() {
String method = "notifyMethod";
Map<String, Object> params = Map.of("param1", "value1");
// Call sendNotification
mcpSession.sendNotification(method, params);
// Capture the notification sent to transport
ArgumentCaptor<JSONRPCMessage> messageCaptor = ArgumentCaptor.forClass(JSONRPCMessage.class);
verify(mockTransport).sendMessage(messageCaptor.capture());
// Verify the notification details
JSONRPCNotification notification = (JSONRPCNotification) messageCaptor.getValue();
assertEquals("notifyMethod", notification.method());
assertEquals("value1", ((Map<?, ?>) notification.params()).get("param1"));
}
@Test
void messageHandlerShouldProcessResponse() {
// Capture the message handler during initialization
ArgumentCaptor<Consumer<JSONRPCMessage>> handlerCaptor = ArgumentCaptor.forClass(Consumer.class);
verify(mockTransport).setMessageHandler(handlerCaptor.capture());
Consumer<JSONRPCMessage> handler = handlerCaptor.getValue();
// Simulate a JSONRPCResponse
JSONRPCResponse response = new JSONRPCResponse("2.0", "1", "result", null);
handler.accept(response);
// Verify the response was processed
// (Use CompletableFuture to match the response if needed in real tests)
}
@Test
void errorHandlerShouldLogErrors() {
// Capture the error handler during initialization
ArgumentCaptor<Consumer<String>> handlerCaptor = ArgumentCaptor.forClass(Consumer.class);
verify(mockTransport).setErrorHandler(handlerCaptor.capture());
Consumer<String> errorHandler = handlerCaptor.getValue();
// Simulate an error
String errorMessage = "Test error";
errorHandler.accept(errorMessage);
// Verify the error was logged or processed
// (In real tests, capture logging output or other side effects)
}
@Test
void closeShouldStopTransport() {
mcpSession.close();
verify(mockTransport).close();
}
}

7
spring/src/main/java/spring/ai/mcp/McpApplication.java
View file

@ -6,7 +6,7 @@ import java.util.List;
import com.fasterxml.jackson.databind.ObjectMapper;
import spring.ai.mcp.client.McpClient;
import spring.ai.mcp.client.McpSyncClient;
import spring.ai.mcp.client.stdio.StdioServerParameters;
import spring.ai.mcp.client.stdio.ServerParameters;
import spring.ai.mcp.client.stdio.StdioServerTransport;
import org.springframework.ai.chat.client.ChatClient;
@ -55,15 +55,14 @@ public class McpApplication {
@Bean(destroyMethod = "close")
public McpSyncClient clientSession() {
var stdioParams = StdioServerParameters.builder("npx")
var stdioParams = ServerParameters.builder("npx")
.args("-y", "@modelcontextprotocol/server-filesystem",
"/Users/christiantzolov/Dev/projects/demo/mcp-server/dir")
.build();
McpSyncClient mcpClient = null;
try {
mcpClient = McpClient.sync(new StdioServerTransport(stdioParams,
Duration.ofMillis(100)),
mcpClient = McpClient.sync(new StdioServerTransport(stdioParams),
Duration.ofSeconds(10), new ObjectMapper());
var init = mcpClient.initialize();