Introduce HttpRequest.Builder customizer for HttpClient-based transport
- Minor improvement: speed up HttpClientSseClientTransportTests by reusing the MCP Server container across tests. - Minor improvement: rename "messageSink" to "deliveredSink" in HttpClientStreamableHttpTransport#sendMessage
This commit is contained in:
Daniel Garnier-Moiroux
Christian Tzolov
parent
bde1b6beec
commit
ae2fc8683b
2 changed files with 539 additions and 150 deletions
|
|
@ -17,10 +17,13 @@ import java.util.function.Function;
|
|||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.JSONRPCRequest;
|
||||
import org.junit.jupiter.api.AfterAll;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
import reactor.core.publisher.Mono;
|
||||
|
|
@ -28,9 +31,17 @@ import reactor.core.publisher.Sinks;
|
|||
import reactor.test.StepVerifier;
|
||||
|
||||
import org.springframework.http.codec.ServerSentEvent;
|
||||
import org.springframework.web.util.UriComponentsBuilder;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatCode;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.ArgumentMatchers.isNull;
|
||||
import static org.mockito.Mockito.clearInvocations;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
/**
|
||||
* Tests for the {@link HttpClientSseClientTransport} class.
|
||||
|
|
@ -43,7 +54,7 @@ class HttpClientSseClientTransportTests {
|
|||
static String host = "http://localhost:3001";
|
||||
|
||||
@SuppressWarnings("resource")
|
||||
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v2")
|
||||
static GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v2")
|
||||
.withCommand("node dist/index.js sse")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
|
|
@ -61,7 +72,7 @@ class HttpClientSseClientTransportTests {
|
|||
public TestHttpClientSseClientTransport(final String baseUri) {
|
||||
super(HttpClient.newBuilder().version(HttpClient.Version.HTTP_1_1).build(),
|
||||
HttpRequest.newBuilder().header("Content-Type", "application/json"), baseUri, "/sse",
|
||||
new ObjectMapper());
|
||||
new ObjectMapper(), AsyncHttpRequestCustomizer.NOOP);
|
||||
}
|
||||
|
||||
public int getInboundMessageCount() {
|
||||
|
|
@ -80,15 +91,21 @@ class HttpClientSseClientTransportTests {
|
|||
|
||||
}
|
||||
|
||||
void startContainer() {
|
||||
@BeforeAll
|
||||
static void startContainer() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
|
||||
}
|
||||
|
||||
@AfterAll
|
||||
static void stopContainer() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
startContainer();
|
||||
transport = new TestHttpClientSseClientTransport(host);
|
||||
transport.connect(Function.identity()).block();
|
||||
}
|
||||
|
|
@ -98,11 +115,6 @@ class HttpClientSseClientTransportTests {
|
|||
if (transport != null) {
|
||||
assertThatCode(() -> transport.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
|
||||
}
|
||||
cleanup();
|
||||
}
|
||||
|
||||
void cleanup() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
@Test
|
||||
|
|
@ -375,4 +387,74 @@ class HttpClientSseClientTransportTests {
|
|||
customizedTransport.closeGracefully().block();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRequestCustomizer() {
|
||||
var mockCustomizer = mock(SyncHttpRequestCustomizer.class);
|
||||
|
||||
// Create a transport with the customizer
|
||||
var customizedTransport = HttpClientSseClientTransport.builder(host)
|
||||
.httpRequestCustomizer(mockCustomizer)
|
||||
.build();
|
||||
|
||||
// Connect
|
||||
StepVerifier.create(customizedTransport.connect(Function.identity())).verifyComplete();
|
||||
|
||||
// Verify the customizer was called
|
||||
verify(mockCustomizer).customize(any(), eq("GET"),
|
||||
eq(UriComponentsBuilder.fromUriString(host).path("/sse").build().toUri()), isNull());
|
||||
clearInvocations(mockCustomizer);
|
||||
|
||||
// Send test message
|
||||
var testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Subscribe to messages and verify
|
||||
StepVerifier.create(customizedTransport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
// Verify the customizer was called
|
||||
var uriArgumentCaptor = ArgumentCaptor.forClass(URI.class);
|
||||
verify(mockCustomizer).customize(any(), eq("POST"), uriArgumentCaptor.capture(), eq(
|
||||
"{\"jsonrpc\":\"2.0\",\"method\":\"test-method\",\"id\":\"test-id\",\"params\":{\"key\":\"value\"}}"));
|
||||
assertThat(uriArgumentCaptor.getValue().toString()).startsWith(host + "/message?sessionId=");
|
||||
|
||||
// Clean up
|
||||
customizedTransport.closeGracefully().block();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAsyncRequestCustomizer() {
|
||||
var mockCustomizer = mock(AsyncHttpRequestCustomizer.class);
|
||||
when(mockCustomizer.customize(any(), any(), any(), any()))
|
||||
.thenAnswer(invocation -> Mono.just(invocation.getArguments()[0]));
|
||||
|
||||
// Create a transport with the customizer
|
||||
var customizedTransport = HttpClientSseClientTransport.builder(host)
|
||||
.asyncHttpRequestCustomizer(mockCustomizer)
|
||||
.build();
|
||||
|
||||
// Connect
|
||||
StepVerifier.create(customizedTransport.connect(Function.identity())).verifyComplete();
|
||||
|
||||
// Verify the customizer was called
|
||||
verify(mockCustomizer).customize(any(), eq("GET"),
|
||||
eq(UriComponentsBuilder.fromUriString(host).path("/sse").build().toUri()), isNull());
|
||||
clearInvocations(mockCustomizer);
|
||||
|
||||
// Send test message
|
||||
var testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
|
||||
Map.of("key", "value"));
|
||||
|
||||
// Subscribe to messages and verify
|
||||
StepVerifier.create(customizedTransport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
// Verify the customizer was called
|
||||
var uriArgumentCaptor = ArgumentCaptor.forClass(URI.class);
|
||||
verify(mockCustomizer).customize(any(), eq("POST"), uriArgumentCaptor.capture(), eq(
|
||||
"{\"jsonrpc\":\"2.0\",\"method\":\"test-method\",\"id\":\"test-id\",\"params\":{\"key\":\"value\"}}"));
|
||||
assertThat(uriArgumentCaptor.getValue().toString()).startsWith(host + "/message?sessionId=");
|
||||
|
||||
// Clean up
|
||||
customizedTransport.closeGracefully().block();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,115 @@
|
|||
/*
|
||||
* Copyright 2024-2025 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import java.net.URI;
|
||||
import java.net.URISyntaxException;
|
||||
import java.util.function.Consumer;
|
||||
import org.junit.jupiter.api.AfterAll;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.atLeastOnce;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
/**
|
||||
* Tests for the {@link HttpClientStreamableHttpTransport} class.
|
||||
*
|
||||
* @author Daniel Garnier-Moiroux
|
||||
*/
|
||||
class HttpClientStreamableHttpTransportTest {
|
||||
|
||||
static String host = "http://localhost:3001";
|
||||
|
||||
@SuppressWarnings("resource")
|
||||
static GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v2")
|
||||
.withCommand("node dist/index.js streamableHttp")
|
||||
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
|
||||
.withExposedPorts(3001)
|
||||
.waitingFor(Wait.forHttp("/").forStatusCode(404));
|
||||
|
||||
@BeforeAll
|
||||
static void startContainer() {
|
||||
container.start();
|
||||
int port = container.getMappedPort(3001);
|
||||
host = "http://" + container.getHost() + ":" + port;
|
||||
}
|
||||
|
||||
@AfterAll
|
||||
static void stopContainer() {
|
||||
container.stop();
|
||||
}
|
||||
|
||||
void withTransport(HttpClientStreamableHttpTransport transport, Consumer<HttpClientStreamableHttpTransport> c) {
|
||||
try {
|
||||
c.accept(transport);
|
||||
}
|
||||
finally {
|
||||
StepVerifier.create(transport.closeGracefully()).verifyComplete();
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRequestCustomizer() throws URISyntaxException {
|
||||
var uri = new URI(host + "/mcp");
|
||||
var mockRequestCustomizer = mock(SyncHttpRequestCustomizer.class);
|
||||
|
||||
var transport = HttpClientStreamableHttpTransport.builder(host)
|
||||
.httpRequestCustomizer(mockRequestCustomizer)
|
||||
.build();
|
||||
|
||||
withTransport(transport, (t) -> {
|
||||
// Send test message
|
||||
var initializeRequest = new McpSchema.InitializeRequest(McpSchema.LATEST_PROTOCOL_VERSION,
|
||||
McpSchema.ClientCapabilities.builder().roots(true).build(),
|
||||
new McpSchema.Implementation("Spring AI MCP Client", "0.3.1"));
|
||||
var testMessage = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, McpSchema.METHOD_INITIALIZE,
|
||||
"test-id", initializeRequest);
|
||||
|
||||
StepVerifier.create(t.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
// Verify the customizer was called
|
||||
verify(mockRequestCustomizer, atLeastOnce()).customize(any(), eq("GET"), eq(uri), eq(
|
||||
"{\"jsonrpc\":\"2.0\",\"method\":\"initialize\",\"id\":\"test-id\",\"params\":{\"protocolVersion\":\"2025-03-26\",\"capabilities\":{\"roots\":{\"listChanged\":true}},\"clientInfo\":{\"name\":\"Spring AI MCP Client\",\"version\":\"0.3.1\"}}}"));
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAsyncRequestCustomizer() throws URISyntaxException {
|
||||
var uri = new URI(host + "/mcp");
|
||||
var mockRequestCustomizer = mock(AsyncHttpRequestCustomizer.class);
|
||||
when(mockRequestCustomizer.customize(any(), any(), any(), any()))
|
||||
.thenAnswer(invocation -> Mono.just(invocation.getArguments()[0]));
|
||||
|
||||
var transport = HttpClientStreamableHttpTransport.builder(host)
|
||||
.asyncHttpRequestCustomizer(mockRequestCustomizer)
|
||||
.build();
|
||||
|
||||
withTransport(transport, (t) -> {
|
||||
// Send test message
|
||||
var initializeRequest = new McpSchema.InitializeRequest(McpSchema.LATEST_PROTOCOL_VERSION,
|
||||
McpSchema.ClientCapabilities.builder().roots(true).build(),
|
||||
new McpSchema.Implementation("Spring AI MCP Client", "0.3.1"));
|
||||
var testMessage = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, McpSchema.METHOD_INITIALIZE,
|
||||
"test-id", initializeRequest);
|
||||
|
||||
StepVerifier.create(t.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
// Verify the customizer was called
|
||||
verify(mockRequestCustomizer, atLeastOnce()).customize(any(), eq("GET"), eq(uri), eq(
|
||||
"{\"jsonrpc\":\"2.0\",\"method\":\"initialize\",\"id\":\"test-id\",\"params\":{\"protocolVersion\":\"2025-03-26\",\"capabilities\":{\"roots\":{\"listChanged\":true}},\"clientInfo\":{\"name\":\"Spring AI MCP Client\",\"version\":\"0.3.1\"}}}"));
|
||||
});
|
||||
}
|
||||
|
||||
}
|
||||
Loading…
Reference in a new issue