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

feat(mcp): Add sampling capabilities to MCP server

Browse source 
Implement client-side sampling features in the Model Context Protocol (MCP) server:
- Add createMessage method to McpAsyncServer and McpSyncServer for LLM sampling
- Move MockMcpTransport to separate file for reusability
- Add integration tests for sampling functionality
- Test error handling for uninitialized clients and missing capabilities

The sampling implementation allows servers to request LLM completions from
clients while maintaining client-side control over model access and permissions.

Resolves #42
This commit is contained in:
Christian Tzolov 2025-01-02 15:45:17 +01:00
commit 07769a651d
7 changed files with 331 additions and 153 deletions
Download patch file Download diff file Expand all files Collapse all files

103
mcp/src/test/java/org/springframework/ai/mcp/MockMcpTransport.java Normal file
View file

@ -0,0 +1,103 @@
/*
* 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;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.scheduler.Schedulers;
import org.springframework.ai.mcp.spec.McpSchema;
import org.springframework.ai.mcp.spec.McpTransport;
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCMessage;
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCNotification;
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
@SuppressWarnings("unused")
public class MockMcpTransport implements McpTransport {
private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
private final Sinks.Many<McpSchema.JSONRPCMessage> outgoing = Sinks.many().multicast().onBackpressureBuffer();
private final Sinks.Many<McpSchema.JSONRPCMessage> inbound = Sinks.many().unicast().onBackpressureBuffer();
private final Flux<McpSchema.JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) {
if (inbound.tryEmitNext(message).isFailure()) {
throw new RuntimeException("Failed to emit message " + message);
}
inboundMessageCount.incrementAndGet();
}
@Override
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
if (outgoing.tryEmitNext(message).isFailure()) {
return Mono.error(new RuntimeException("Can't emit outgoing message " + message));
}
return Mono.empty();
}
public McpSchema.JSONRPCRequest getLastSentMessageAsRequest() {
return (JSONRPCRequest) outboundView.blockFirst();
}
public McpSchema.JSONRPCNotification getLastSentMessageAsNotifiation() {
return (JSONRPCNotification) outboundView.blockFirst();
}
public McpSchema.JSONRPCMessage getLastSentMessage() {
return outboundView.blockFirst();
}
private volatile boolean connected = false;
@Override
public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
if (connected) {
return Mono.error(new IllegalStateException("Already connected"));
}
connected = true;
return inbound.asFlux()
.publishOn(Schedulers.boundedElastic())
.flatMap(message -> Mono.just(message).transform(handler))
.doFinally(signal -> connected = false)
.then();
}
@Override
public Mono<Void> closeGracefully() {
return Mono.defer(() -> {
connected = false;
outgoing.tryEmitComplete();
inbound.tryEmitComplete();
return Mono.empty();
});
}
@Override
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
return new ObjectMapper().convertValue(data, typeRef);
}
}

80
mcp/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientResponseHandlerTests.java
View file

@ -20,100 +20,22 @@ import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.function.Function;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.scheduler.Schedulers;
import org.springframework.ai.mcp.MockMcpTransport;
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 {
@SuppressWarnings("unused")
private static class MockMcpTransport implements McpTransport {
private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
private final Sinks.Many<McpSchema.JSONRPCMessage> outgoing = Sinks.many().multicast().onBackpressureBuffer();
private final Sinks.Many<McpSchema.JSONRPCMessage> inbound = Sinks.many().unicast().onBackpressureBuffer();
private final Flux<McpSchema.JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) {
if (inbound.tryEmitNext(message).isFailure()) {
throw new RuntimeException("Failed to emit message " + message);
}
inboundMessageCount.incrementAndGet();
}
@Override
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
if (outgoing.tryEmitNext(message).isFailure()) {
return Mono.error(new RuntimeException("Can't emit outgoing message " + message));
}
return Mono.empty();
}
public McpSchema.JSONRPCRequest getLastSentMessageAsRequest() {
return (JSONRPCRequest) outboundView.blockFirst();
}
public McpSchema.JSONRPCNotification getLastSentMessageAsNotifiation() {
return (JSONRPCNotification) outboundView.blockFirst();
}
public McpSchema.JSONRPCMessage getLastSentMessage() {
return outboundView.blockFirst();
}
private volatile boolean connected = false;
@Override
public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
if (connected) {
return Mono.error(new IllegalStateException("Already connected"));
}
connected = true;
return inbound.asFlux()
.publishOn(Schedulers.boundedElastic())
.flatMap(message -> Mono.just(message).transform(handler))
.doFinally(signal -> connected = false)
.then();
}
@Override
public Mono<Void> closeGracefully() {
return Mono.defer(() -> {
connected = false;
outgoing.tryEmitComplete();
inbound.tryEmitComplete();
return Mono.empty();
});
}
@Override
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
return new ObjectMapper().convertValue(data, typeRef);
}
}
@Test
void testToolsChangeNotificationHandling() {
MockMcpTransport transport = new MockMcpTransport();

2
mcp/src/test/java/org/springframework/ai/mcp/server/AbstractMcpAsyncServerTests.java
View file

@ -25,11 +25,11 @@ 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;
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
import org.springframework.ai.mcp.spec.McpSchema.GetPromptResult;
import org.springframework.ai.mcp.spec.McpSchema.Prompt;

159
mcp/src/test/java/org/springframework/ai/mcp/server/SseAsyncIntegrationTests.java Normal file
View file

@ -0,0 +1,159 @@
/*
* 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 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 reactor.netty.DisposableServer;
import reactor.netty.http.server.HttpServer;
import reactor.test.StepVerifier;
import org.springframework.ai.mcp.client.McpClient;
import org.springframework.ai.mcp.client.transport.SseClientTransport;
import org.springframework.ai.mcp.server.transport.SseServerTransport;
import org.springframework.ai.mcp.spec.McpError;
import org.springframework.ai.mcp.spec.McpSchema;
import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities;
import org.springframework.ai.mcp.spec.McpSchema.CreateMessageRequest;
import org.springframework.ai.mcp.spec.McpSchema.CreateMessageResult;
import org.springframework.ai.mcp.spec.McpSchema.InitializeResult;
import org.springframework.ai.mcp.spec.McpSchema.Role;
import org.springframework.http.server.reactive.HttpHandler;
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.server.RouterFunctions;
import static org.assertj.core.api.Assertions.assertThat;
public class SseAsyncIntegrationTests {
private static final int PORT = 8181;
private static final String MESSAGE_ENDPOINT = "/mcp/message";
private DisposableServer httpServer;
private SseServerTransport mcpServerTransport;
McpClient.Builder clientBuilder;
@BeforeEach
public void before() {
this.mcpServerTransport = new SseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpServerTransport.getRouterFunction());
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
this.httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
this.clientBuilder = McpClient
.using(new SseClientTransport(WebClient.builder().baseUrl("http://localhost:" + PORT)));
}
@AfterEach
public void after() {
if (httpServer != null) {
httpServer.disposeNow();
}
}
// ---------------------------------------
// Sampling Tests
// ---------------------------------------
@Test
void testCreateMessageWithoutInitialization() {
var mcpAsyncServer = McpServer.using(mcpServerTransport).serverInfo("test-server", "1.0.0").async();
var messages = List
.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message")));
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of());
StepVerifier.create(mcpAsyncServer.createMessage(request)).verifyErrorSatisfies(error -> {
assertThat(error).isInstanceOf(McpError.class)
.hasMessage("Client must be initialized. Call the initialize method first!");
});
}
@Test
void testCreateMessageWithoutSamplingCapabilities() {
var mcpAsyncServer = McpServer.using(mcpServerTransport).serverInfo("test-server", "1.0.0").async();
var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).sync();
InitializeResult initResult = client.initialize();
assertThat(initResult).isNotNull();
var messages = List
.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message")));
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of());
StepVerifier.create(mcpAsyncServer.createMessage(request)).verifyErrorSatisfies(error -> {
assertThat(error).isInstanceOf(McpError.class)
.hasMessage("Client must be configured with sampling capabilities");
});
}
@Test
void testCreateMessageSuccess() throws InterruptedException {
var mcpAsyncServer = McpServer.using(mcpServerTransport).serverInfo("test-server", "1.0.0").async();
Function<CreateMessageRequest, CreateMessageResult> samplingHandler = request -> {
assertThat(request.messages()).hasSize(1);
assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class);
return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName",
CreateMessageResult.StopReason.STOP_SEQUENCE);
};
var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
.capabilities(ClientCapabilities.builder().sampling().build())
.sampling(samplingHandler)
.sync();
InitializeResult initResult = client.initialize();
assertThat(initResult).isNotNull();
var messages = List
.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message")));
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of());
StepVerifier.create(mcpAsyncServer.createMessage(request)).consumeNextWith(result -> {
assertThat(result).isNotNull();
assertThat(result.role()).isEqualTo(Role.USER);
assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class);
assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message");
assertThat(result.model()).isEqualTo("MockModelName");
assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE);
}).verifyComplete();
}
}

7
mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpAsyncServerTests.java
View file

@ -18,15 +18,14 @@ package org.springframework.ai.mcp.server;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Timeout;
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;
import org.springframework.http.server.reactive.HttpHandler;
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
import org.springframework.web.reactive.function.server.RouterFunctions;
/**
* Tests for {@link McpAsyncServer} using {@link SseServerTransport}.

7
mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpSyncServerTests.java
View file

@ -18,15 +18,14 @@ package org.springframework.ai.mcp.server;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Timeout;
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;
import org.springframework.http.server.reactive.HttpHandler;
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
import org.springframework.web.reactive.function.server.RouterFunctions;
/**
* Tests for {@link McpSyncServer} using {@link SseServerTransport}.

66
mcp/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpSessionTests.java
View file

@ -18,24 +18,18 @@ package org.springframework.ai.mcp.spec;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import com.fasterxml.jackson.core.type.TypeReference;
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.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.scheduler.Schedulers;
import reactor.test.StepVerifier;
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCNotification;
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
import org.springframework.ai.mcp.MockMcpTransport;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@ -62,64 +56,6 @@ class DefaultMcpSessionTests {
private MockMcpTransport transport;
@SuppressWarnings("unused")
private static class MockMcpTransport implements McpTransport {
private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
private Sinks.Many<McpSchema.JSONRPCMessage> outgoing = Sinks.many().multicast().onBackpressureBuffer();
private Sinks.Many<McpSchema.JSONRPCMessage> inbound = Sinks.many().unicast().onBackpressureBuffer();
private Flux<McpSchema.JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) {
if (inbound.tryEmitNext(message).isFailure()) {
throw new RuntimeException("Failed to emit message " + message);
}
inboundMessageCount.incrementAndGet();
}
@Override
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
if (outgoing.tryEmitNext(message).isFailure()) {
return Mono.error(new RuntimeException("Can't emit outgoing message " + message));
}
return Mono.empty();
}
public McpSchema.JSONRPCRequest getLastSentMessageAsRequest() {
return (JSONRPCRequest) outboundView.blockFirst();
}
public McpSchema.JSONRPCNotification getLastSentMessageAsNotifiation() {
return (JSONRPCNotification) outboundView.blockFirst();
}
public McpSchema.JSONRPCMessage getLastSentMessage() {
return outboundView.blockFirst();
}
@Override
public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
return inbound.asFlux()
.publishOn(Schedulers.boundedElastic())
.flatMap(message -> Mono.just(message).transform(handler))
.then();
}
@Override
public Mono<Void> closeGracefully() {
return Mono.empty();
}
@Override
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
return new ObjectMapper().convertValue(data, typeRef);
}
}
@BeforeEach
void setUp() {
transport = new MockMcpTransport();