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
4 changed files with 539 additions and 150 deletions
|
|
@ -0,0 +1,52 @@
|
|||
/*
|
||||
* Copyright 2024-2025 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import java.net.URI;
|
||||
import java.net.http.HttpRequest;
|
||||
import org.reactivestreams.Publisher;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
import reactor.util.annotation.Nullable;
|
||||
|
||||
/**
|
||||
* Customize {@link HttpRequest.Builder} before executing the request, in either SSE or
|
||||
* Streamable HTTP transport.
|
||||
* <p>
|
||||
* When used in a non-blocking context, implementations MUST be non-blocking.
|
||||
*
|
||||
* @author Daniel Garnier-Moiroux
|
||||
*/
|
||||
public interface AsyncHttpRequestCustomizer {
|
||||
|
||||
Publisher<HttpRequest.Builder> customize(HttpRequest.Builder builder, String method, URI endpoint,
|
||||
@Nullable String body);
|
||||
|
||||
AsyncHttpRequestCustomizer NOOP = new Noop();
|
||||
|
||||
/**
|
||||
* Wrap a sync implementation in an async wrapper.
|
||||
* <p>
|
||||
* Do NOT wrap a blocking implementation for use in a non-blocking context. For a
|
||||
* blocking implementation, consider using {@link Schedulers#boundedElastic()}.
|
||||
*/
|
||||
static AsyncHttpRequestCustomizer fromSync(SyncHttpRequestCustomizer customizer) {
|
||||
return (builder, method, uri, body) -> Mono.fromSupplier(() -> {
|
||||
customizer.customize(builder, method, uri, body);
|
||||
return builder;
|
||||
});
|
||||
}
|
||||
|
||||
class Noop implements AsyncHttpRequestCustomizer {
|
||||
|
||||
@Override
|
||||
public Publisher<HttpRequest.Builder> customize(HttpRequest.Builder builder, String method, URI endpoint,
|
||||
String body) {
|
||||
return Mono.just(builder);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
/*
|
||||
* Copyright 2024 - 2024 the original author or authors.
|
||||
* Copyright 2024 - 2025 the original author or authors.
|
||||
*/
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
|
|
@ -102,6 +102,11 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
*/
|
||||
protected final Sinks.One<String> messageEndpointSink = Sinks.one();
|
||||
|
||||
/**
|
||||
* Customizer to modify requests before they are executed.
|
||||
*/
|
||||
private final AsyncHttpRequestCustomizer httpRequestCustomizer;
|
||||
|
||||
/**
|
||||
* Creates a new transport instance with default HTTP client and object mapper.
|
||||
* @param baseUri the base URI of the MCP server
|
||||
|
|
@ -172,18 +177,38 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
* @param objectMapper the object mapper for JSON serialization/deserialization
|
||||
* @throws IllegalArgumentException if objectMapper, clientBuilder, or headers is null
|
||||
*/
|
||||
@Deprecated(forRemoval = true)
|
||||
HttpClientSseClientTransport(HttpClient httpClient, HttpRequest.Builder requestBuilder, String baseUri,
|
||||
String sseEndpoint, ObjectMapper objectMapper) {
|
||||
this(httpClient, requestBuilder, baseUri, sseEndpoint, objectMapper, AsyncHttpRequestCustomizer.NOOP);
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a new transport instance with custom HTTP client builder, object mapper,
|
||||
* and headers.
|
||||
* @param httpClient the HTTP client to use
|
||||
* @param requestBuilder the HTTP request builder to use
|
||||
* @param baseUri the base URI of the MCP server
|
||||
* @param sseEndpoint the SSE endpoint path
|
||||
* @param objectMapper the object mapper for JSON serialization/deserialization
|
||||
* @param httpRequestCustomizer customizer for the requestBuilder before executing
|
||||
* requests
|
||||
* @throws IllegalArgumentException if objectMapper, clientBuilder, or headers is null
|
||||
*/
|
||||
HttpClientSseClientTransport(HttpClient httpClient, HttpRequest.Builder requestBuilder, String baseUri,
|
||||
String sseEndpoint, ObjectMapper objectMapper, AsyncHttpRequestCustomizer httpRequestCustomizer) {
|
||||
Assert.notNull(objectMapper, "ObjectMapper must not be null");
|
||||
Assert.hasText(baseUri, "baseUri must not be empty");
|
||||
Assert.hasText(sseEndpoint, "sseEndpoint must not be empty");
|
||||
Assert.notNull(httpClient, "httpClient must not be null");
|
||||
Assert.notNull(requestBuilder, "requestBuilder must not be null");
|
||||
Assert.notNull(httpRequestCustomizer, "httpRequestCustomizer must not be null");
|
||||
this.baseUri = URI.create(baseUri);
|
||||
this.sseEndpoint = sseEndpoint;
|
||||
this.objectMapper = objectMapper;
|
||||
this.httpClient = httpClient;
|
||||
this.requestBuilder = requestBuilder;
|
||||
this.httpRequestCustomizer = httpRequestCustomizer;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -213,6 +238,8 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
private HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
|
||||
.header("Content-Type", "application/json");
|
||||
|
||||
private AsyncHttpRequestCustomizer httpRequestCustomizer = AsyncHttpRequestCustomizer.NOOP;
|
||||
|
||||
/**
|
||||
* Creates a new builder instance.
|
||||
*/
|
||||
|
|
@ -310,31 +337,66 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the customizer for {@link HttpRequest.Builder}, to modify requests before
|
||||
* executing them.
|
||||
* <p>
|
||||
* This overrides the customizer from
|
||||
* {@link #asyncHttpRequestCustomizer(AsyncHttpRequestCustomizer)}.
|
||||
* <p>
|
||||
* Do NOT use a blocking {@link SyncHttpRequestCustomizer} in a non-blocking
|
||||
* context. Use {@link #asyncHttpRequestCustomizer(AsyncHttpRequestCustomizer)}
|
||||
* instead.
|
||||
* @param syncHttpRequestCustomizer the request customizer
|
||||
* @return this builder
|
||||
*/
|
||||
public Builder httpRequestCustomizer(SyncHttpRequestCustomizer syncHttpRequestCustomizer) {
|
||||
this.httpRequestCustomizer = AsyncHttpRequestCustomizer.fromSync(syncHttpRequestCustomizer);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the customizer for {@link HttpRequest.Builder}, to modify requests before
|
||||
* executing them.
|
||||
* <p>
|
||||
* This overrides the customizer from
|
||||
* {@link #httpRequestCustomizer(SyncHttpRequestCustomizer)}.
|
||||
* <p>
|
||||
* Do NOT use a blocking implementation in a non-blocking context.
|
||||
* @param asyncHttpRequestCustomizer the request customizer
|
||||
* @return this builder
|
||||
*/
|
||||
public Builder asyncHttpRequestCustomizer(AsyncHttpRequestCustomizer asyncHttpRequestCustomizer) {
|
||||
this.httpRequestCustomizer = asyncHttpRequestCustomizer;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Builds a new {@link HttpClientSseClientTransport} instance.
|
||||
* @return a new transport instance
|
||||
*/
|
||||
public HttpClientSseClientTransport build() {
|
||||
return new HttpClientSseClientTransport(clientBuilder.build(), requestBuilder, baseUri, sseEndpoint,
|
||||
objectMapper);
|
||||
objectMapper, httpRequestCustomizer);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
|
||||
var uri = Utils.resolveUri(this.baseUri, this.sseEndpoint);
|
||||
|
||||
return Mono.create(sink -> {
|
||||
|
||||
HttpRequest request = requestBuilder.copy()
|
||||
.uri(Utils.resolveUri(this.baseUri, this.sseEndpoint))
|
||||
return Mono.defer(() -> {
|
||||
var builder = requestBuilder.copy()
|
||||
.uri(uri)
|
||||
.header("Accept", "text/event-stream")
|
||||
.header("Cache-Control", "no-cache")
|
||||
.GET()
|
||||
.build();
|
||||
|
||||
.GET();
|
||||
return Mono.from(this.httpRequestCustomizer.customize(builder, "GET", uri, null));
|
||||
}).flatMap(requestBuilder -> Mono.create(sink -> {
|
||||
Disposable connection = Flux.<ResponseEvent>create(sseSink -> this.httpClient
|
||||
.sendAsync(request, responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo, sseSink))
|
||||
.sendAsync(requestBuilder.build(),
|
||||
responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo, sseSink))
|
||||
.exceptionallyCompose(e -> {
|
||||
sseSink.error(e);
|
||||
return CompletableFuture.failedFuture(e);
|
||||
|
|
@ -397,7 +459,7 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
.subscribe();
|
||||
|
||||
this.sseSubscription.set(connection);
|
||||
});
|
||||
}));
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -453,13 +515,13 @@ public class HttpClientSseClientTransport implements McpClientTransport {
|
|||
|
||||
private Mono<HttpResponse<String>> sendHttpPost(final String endpoint, final String body) {
|
||||
final URI requestUri = Utils.resolveUri(baseUri, endpoint);
|
||||
final HttpRequest request = this.requestBuilder.copy()
|
||||
.uri(requestUri)
|
||||
.POST(HttpRequest.BodyPublishers.ofString(body))
|
||||
.build();
|
||||
|
||||
// TODO: why discard the body?
|
||||
return Mono.fromFuture(httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofString()));
|
||||
return Mono.defer(() -> {
|
||||
var builder = this.requestBuilder.copy().uri(requestUri).POST(HttpRequest.BodyPublishers.ofString(body));
|
||||
return Mono.from(this.httpRequestCustomizer.customize(builder, "POST", requestUri, body));
|
||||
}).flatMap(customizedBuilder -> {
|
||||
var request = customizedBuilder.build();
|
||||
return Mono.fromFuture(httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofString()));
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
* Copyright 2024-2025 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
|
@ -109,6 +109,8 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
|
||||
private final boolean resumableStreams;
|
||||
|
||||
private final AsyncHttpRequestCustomizer httpRequestCustomizer;
|
||||
|
||||
private final AtomicReference<DefaultMcpTransportSession> activeSession = new AtomicReference<>();
|
||||
|
||||
private final AtomicReference<Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>>> handler = new AtomicReference<>();
|
||||
|
|
@ -117,7 +119,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
|
||||
private HttpClientStreamableHttpTransport(ObjectMapper objectMapper, HttpClient httpClient,
|
||||
HttpRequest.Builder requestBuilder, String baseUri, String endpoint, boolean resumableStreams,
|
||||
boolean openConnectionOnStartup) {
|
||||
boolean openConnectionOnStartup, AsyncHttpRequestCustomizer httpRequestCustomizer) {
|
||||
this.objectMapper = objectMapper;
|
||||
this.httpClient = httpClient;
|
||||
this.requestBuilder = requestBuilder;
|
||||
|
|
@ -126,6 +128,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
this.resumableStreams = resumableStreams;
|
||||
this.openConnectionOnStartup = openConnectionOnStartup;
|
||||
this.activeSession.set(createTransportSession());
|
||||
this.httpRequestCustomizer = httpRequestCustomizer;
|
||||
}
|
||||
|
||||
public static Builder builder(String baseUri) {
|
||||
|
|
@ -154,14 +157,18 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
}
|
||||
|
||||
private Publisher<Void> createDelete(String sessionId) {
|
||||
HttpRequest request = this.requestBuilder.copy()
|
||||
.uri(Utils.resolveUri(this.baseUri, this.endpoint))
|
||||
.header("Cache-Control", "no-cache")
|
||||
.header("mcp-session-id", sessionId)
|
||||
.DELETE()
|
||||
.build();
|
||||
|
||||
return Mono.fromFuture(() -> this.httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofString())).then();
|
||||
var uri = Utils.resolveUri(this.baseUri, this.endpoint);
|
||||
return Mono.defer(() -> {
|
||||
var builder = this.requestBuilder.copy()
|
||||
.uri(uri)
|
||||
.header("Cache-Control", "no-cache")
|
||||
.header("mcp-session-id", sessionId)
|
||||
.DELETE();
|
||||
return Mono.from(this.httpRequestCustomizer.customize(builder, "DELETE", uri, null));
|
||||
}).flatMap(requestBuilder -> {
|
||||
var request = requestBuilder.build();
|
||||
return Mono.fromFuture(() -> this.httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofString()));
|
||||
}).then();
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -208,100 +215,110 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
|
||||
final AtomicReference<Disposable> disposableRef = new AtomicReference<>();
|
||||
final McpTransportSession<Disposable> transportSession = this.activeSession.get();
|
||||
var uri = Utils.resolveUri(this.baseUri, this.endpoint);
|
||||
|
||||
HttpRequest.Builder requestBuilder = this.requestBuilder.copy();
|
||||
Disposable connection = Mono.defer(() -> {
|
||||
HttpRequest.Builder requestBuilder = this.requestBuilder.copy();
|
||||
|
||||
if (transportSession != null && transportSession.sessionId().isPresent()) {
|
||||
requestBuilder = requestBuilder.header("mcp-session-id", transportSession.sessionId().get());
|
||||
}
|
||||
if (transportSession != null && transportSession.sessionId().isPresent()) {
|
||||
requestBuilder = requestBuilder.header("mcp-session-id", transportSession.sessionId().get());
|
||||
}
|
||||
|
||||
if (stream != null && stream.lastId().isPresent()) {
|
||||
requestBuilder = requestBuilder.header("last-event-id", stream.lastId().get());
|
||||
}
|
||||
if (stream != null && stream.lastId().isPresent()) {
|
||||
requestBuilder = requestBuilder.header("last-event-id", stream.lastId().get());
|
||||
}
|
||||
|
||||
HttpRequest request = requestBuilder.uri(Utils.resolveUri(this.baseUri, this.endpoint))
|
||||
.header("Accept", TEXT_EVENT_STREAM)
|
||||
.header("Cache-Control", "no-cache")
|
||||
.GET()
|
||||
.build();
|
||||
var builder = requestBuilder.uri(uri)
|
||||
.header("Accept", TEXT_EVENT_STREAM)
|
||||
.header("Cache-Control", "no-cache")
|
||||
.GET();
|
||||
return Mono.from(this.httpRequestCustomizer.customize(builder, "GET", uri, null));
|
||||
})
|
||||
.flatMapMany(
|
||||
requestBuilder -> Flux.<ResponseEvent>create(
|
||||
sseSink -> this.httpClient
|
||||
.sendAsync(requestBuilder.build(),
|
||||
responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo,
|
||||
sseSink))
|
||||
.whenComplete((response, throwable) -> {
|
||||
if (throwable != null) {
|
||||
sseSink.error(throwable);
|
||||
}
|
||||
else {
|
||||
logger.debug("SSE connection established successfully");
|
||||
}
|
||||
}))
|
||||
.map(responseEvent -> (ResponseSubscribers.SseResponseEvent) responseEvent)
|
||||
.flatMap(responseEvent -> {
|
||||
int statusCode = responseEvent.responseInfo().statusCode();
|
||||
|
||||
Disposable connection = Flux.<ResponseEvent>create(sseSink -> this.httpClient
|
||||
.sendAsync(request, responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo, sseSink))
|
||||
.whenComplete((response, throwable) -> {
|
||||
if (throwable != null) {
|
||||
sseSink.error(throwable);
|
||||
}
|
||||
else {
|
||||
logger.debug("SSE connection established successfully");
|
||||
}
|
||||
}))
|
||||
.map(responseEvent -> (ResponseSubscribers.SseResponseEvent) responseEvent)
|
||||
.flatMap(responseEvent -> {
|
||||
int statusCode = responseEvent.responseInfo().statusCode();
|
||||
if (statusCode >= 200 && statusCode < 300) {
|
||||
|
||||
if (statusCode >= 200 && statusCode < 300) {
|
||||
if (MESSAGE_EVENT_TYPE.equals(responseEvent.sseEvent().event())) {
|
||||
try {
|
||||
// We don't support batching ATM and probably
|
||||
// won't since the next version considers
|
||||
// removing it.
|
||||
McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(
|
||||
this.objectMapper, responseEvent.sseEvent().data());
|
||||
|
||||
if (MESSAGE_EVENT_TYPE.equals(responseEvent.sseEvent().event())) {
|
||||
try {
|
||||
// We don't support batching ATM and probably won't since
|
||||
// the
|
||||
// next version considers removing it.
|
||||
McpSchema.JSONRPCMessage message = McpSchema
|
||||
.deserializeJsonRpcMessage(this.objectMapper, responseEvent.sseEvent().data());
|
||||
Tuple2<Optional<String>, Iterable<McpSchema.JSONRPCMessage>> idWithMessages = Tuples
|
||||
.of(Optional.ofNullable(responseEvent.sseEvent().id()),
|
||||
List.of(message));
|
||||
|
||||
Tuple2<Optional<String>, Iterable<McpSchema.JSONRPCMessage>> idWithMessages = Tuples
|
||||
.of(Optional.ofNullable(responseEvent.sseEvent().id()), List.of(message));
|
||||
McpTransportStream<Disposable> sessionStream = stream != null ? stream
|
||||
: new DefaultMcpTransportStream<>(this.resumableStreams,
|
||||
this::reconnect);
|
||||
logger.debug("Connected stream {}", sessionStream.streamId());
|
||||
|
||||
McpTransportStream<Disposable> sessionStream = stream != null ? stream
|
||||
: new DefaultMcpTransportStream<>(this.resumableStreams, this::reconnect);
|
||||
logger.debug("Connected stream {}", sessionStream.streamId());
|
||||
return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages)));
|
||||
|
||||
return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages)));
|
||||
}
|
||||
catch (IOException ioException) {
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(
|
||||
new McpError("Error parsing JSON-RPC message: "
|
||||
+ responseEvent.sseEvent().data()));
|
||||
}
|
||||
}
|
||||
else {
|
||||
logger.debug("Received SSE event with type: {}", responseEvent.sseEvent());
|
||||
return Flux.empty();
|
||||
}
|
||||
}
|
||||
else if (statusCode == METHOD_NOT_ALLOWED) { // NotAllowed
|
||||
logger
|
||||
.debug("The server does not support SSE streams, using request-response mode.");
|
||||
return Flux.empty();
|
||||
}
|
||||
else if (statusCode == NOT_FOUND) {
|
||||
String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
|
||||
"Session not found for session ID: " + sessionIdRepresentation);
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(exception);
|
||||
}
|
||||
else if (statusCode == BAD_REQUEST) {
|
||||
String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
|
||||
"Session not found for session ID: " + sessionIdRepresentation);
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(exception);
|
||||
}
|
||||
|
||||
}
|
||||
catch (IOException ioException) {
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(new McpError(
|
||||
"Error parsing JSON-RPC message: " + responseEvent.sseEvent().data()));
|
||||
}
|
||||
}
|
||||
else {
|
||||
logger.debug("Received SSE event with type: {}", responseEvent.sseEvent());
|
||||
return Flux.empty();
|
||||
}
|
||||
}
|
||||
else if (statusCode == METHOD_NOT_ALLOWED) { // NotAllowed
|
||||
logger.debug("The server does not support SSE streams, using request-response mode.");
|
||||
return Flux.empty();
|
||||
}
|
||||
else if (statusCode == NOT_FOUND) {
|
||||
String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
|
||||
"Session not found for session ID: " + sessionIdRepresentation);
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(exception);
|
||||
}
|
||||
else if (statusCode == BAD_REQUEST) {
|
||||
String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException(
|
||||
"Session not found for session ID: " + sessionIdRepresentation);
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(exception);
|
||||
}
|
||||
|
||||
return Flux.<McpSchema.JSONRPCMessage>error(
|
||||
new McpError("Received unrecognized SSE event type: " + responseEvent.sseEvent().event()));
|
||||
|
||||
}).<McpSchema
|
||||
.JSONRPCMessage>flatMap(jsonrpcMessage -> this.handler.get().apply(Mono.just(jsonrpcMessage)))
|
||||
.onErrorMap(CompletionException.class, t -> t.getCause())
|
||||
.onErrorComplete(t -> {
|
||||
this.handleException(t);
|
||||
return true;
|
||||
})
|
||||
.doFinally(s -> {
|
||||
Disposable ref = disposableRef.getAndSet(null);
|
||||
if (ref != null) {
|
||||
transportSession.removeConnection(ref);
|
||||
}
|
||||
})
|
||||
"Received unrecognized SSE event type: " + responseEvent.sseEvent().event()));
|
||||
}).<McpSchema
|
||||
.JSONRPCMessage>flatMap(
|
||||
jsonrpcMessage -> this.handler.get().apply(Mono.just(jsonrpcMessage)))
|
||||
.onErrorMap(CompletionException.class, t -> t.getCause())
|
||||
.onErrorComplete(t -> {
|
||||
this.handleException(t);
|
||||
return true;
|
||||
})
|
||||
.doFinally(s -> {
|
||||
Disposable ref = disposableRef.getAndSet(null);
|
||||
if (ref != null) {
|
||||
transportSession.removeConnection(ref);
|
||||
}
|
||||
}))
|
||||
.contextWrite(ctx)
|
||||
.subscribe();
|
||||
|
||||
|
|
@ -347,31 +364,33 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
}
|
||||
|
||||
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage sentMessage) {
|
||||
return Mono.create(messageSink -> {
|
||||
return Mono.create(deliveredSink -> {
|
||||
logger.debug("Sending message {}", sentMessage);
|
||||
|
||||
final AtomicReference<Disposable> disposableRef = new AtomicReference<>();
|
||||
final McpTransportSession<Disposable> transportSession = this.activeSession.get();
|
||||
|
||||
HttpRequest.Builder requestBuilder = this.requestBuilder.copy();
|
||||
|
||||
if (transportSession != null && transportSession.sessionId().isPresent()) {
|
||||
requestBuilder = requestBuilder.header("mcp-session-id", transportSession.sessionId().get());
|
||||
}
|
||||
|
||||
var uri = Utils.resolveUri(this.baseUri, this.endpoint);
|
||||
String jsonBody = this.toString(sentMessage);
|
||||
|
||||
HttpRequest request = requestBuilder.uri(Utils.resolveUri(this.baseUri, this.endpoint))
|
||||
.header("Accept", APPLICATION_JSON + ", " + TEXT_EVENT_STREAM)
|
||||
.header("Content-Type", APPLICATION_JSON)
|
||||
.header("Cache-Control", "no-cache")
|
||||
.POST(HttpRequest.BodyPublishers.ofString(jsonBody))
|
||||
.build();
|
||||
Disposable connection = Mono.defer(() -> {
|
||||
HttpRequest.Builder requestBuilder = this.requestBuilder.copy();
|
||||
|
||||
Disposable connection = Flux.<ResponseEvent>create(responseEventSink -> {
|
||||
if (transportSession != null && transportSession.sessionId().isPresent()) {
|
||||
requestBuilder = requestBuilder.header("mcp-session-id", transportSession.sessionId().get());
|
||||
}
|
||||
|
||||
var builder = requestBuilder.uri(uri)
|
||||
.header("Accept", APPLICATION_JSON + ", " + TEXT_EVENT_STREAM)
|
||||
.header("Content-Type", APPLICATION_JSON)
|
||||
.header("Cache-Control", "no-cache")
|
||||
.POST(HttpRequest.BodyPublishers.ofString(jsonBody));
|
||||
return Mono.from(this.httpRequestCustomizer.customize(builder, "GET", uri, jsonBody));
|
||||
}).flatMapMany(requestBuilder -> Flux.<ResponseEvent>create(responseEventSink -> {
|
||||
|
||||
// Create the async request with proper body subscriber selection
|
||||
Mono.fromFuture(this.httpClient.sendAsync(request, this.toSendMessageBodySubscriber(responseEventSink))
|
||||
Mono.fromFuture(this.httpClient
|
||||
.sendAsync(requestBuilder.build(), this.toSendMessageBodySubscriber(responseEventSink))
|
||||
.whenComplete((response, throwable) -> {
|
||||
if (throwable != null) {
|
||||
responseEventSink.error(throwable);
|
||||
|
|
@ -381,13 +400,13 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
}
|
||||
})).onErrorMap(CompletionException.class, t -> t.getCause()).onErrorComplete().subscribe();
|
||||
|
||||
}).flatMap(responseEvent -> {
|
||||
})).flatMap(responseEvent -> {
|
||||
if (transportSession.markInitialized(
|
||||
responseEvent.responseInfo().headers().firstValue("mcp-session-id").orElseGet(() -> null))) {
|
||||
// Once we have a session, we try to open an async stream for
|
||||
// the server to send notifications and requests out-of-band.
|
||||
|
||||
reconnect(null).contextWrite(messageSink.contextView()).subscribe();
|
||||
reconnect(null).contextWrite(deliveredSink.contextView()).subscribe();
|
||||
}
|
||||
|
||||
String sessionRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
|
|
@ -404,16 +423,18 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
|
||||
if (contentType.isBlank()) {
|
||||
logger.debug("No content type returned for POST in session {}", sessionRepresentation);
|
||||
// No content type means no response body, so we can just return
|
||||
// No content type means no response body, so we can just
|
||||
// return
|
||||
// an empty stream
|
||||
messageSink.success();
|
||||
deliveredSink.success();
|
||||
return Flux.empty();
|
||||
}
|
||||
else if (contentType.contains(TEXT_EVENT_STREAM)) {
|
||||
return Flux.just(((ResponseSubscribers.SseResponseEvent) responseEvent).sseEvent())
|
||||
.flatMap(sseEvent -> {
|
||||
try {
|
||||
// We don't support batching ATM and probably won't
|
||||
// We don't support batching ATM and probably
|
||||
// won't
|
||||
// since the
|
||||
// next version considers removing it.
|
||||
McpSchema.JSONRPCMessage message = McpSchema
|
||||
|
|
@ -427,7 +448,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
|
||||
logger.debug("Connected stream {}", sessionStream.streamId());
|
||||
|
||||
messageSink.success();
|
||||
deliveredSink.success();
|
||||
|
||||
return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages)));
|
||||
}
|
||||
|
|
@ -438,7 +459,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
});
|
||||
}
|
||||
else if (contentType.contains(APPLICATION_JSON)) {
|
||||
messageSink.success();
|
||||
deliveredSink.success();
|
||||
String data = ((ResponseSubscribers.AggregateResponseEvent) responseEvent).data();
|
||||
if (sentMessage instanceof McpSchema.JSONRPCNotification && Utils.hasText(data)) {
|
||||
logger.warn("Notification: {} received non-compliant response: {}", sentMessage, data);
|
||||
|
|
@ -483,7 +504,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
// handle the error first
|
||||
this.handleException(t);
|
||||
// inform the caller of sendMessage
|
||||
messageSink.error(t);
|
||||
deliveredSink.error(t);
|
||||
return true;
|
||||
})
|
||||
.doFinally(s -> {
|
||||
|
|
@ -493,7 +514,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
transportSession.removeConnection(ref);
|
||||
}
|
||||
})
|
||||
.contextWrite(messageSink.contextView())
|
||||
.contextWrite(deliveredSink.contextView())
|
||||
.subscribe();
|
||||
|
||||
disposableRef.set(connection);
|
||||
|
|
@ -531,6 +552,8 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
|
||||
private HttpRequest.Builder requestBuilder = HttpRequest.newBuilder();
|
||||
|
||||
private AsyncHttpRequestCustomizer httpRequestCustomizer = AsyncHttpRequestCustomizer.NOOP;
|
||||
|
||||
/**
|
||||
* Creates a new builder with the specified base URI.
|
||||
* @param baseUri the base URI of the MCP server
|
||||
|
|
@ -633,6 +656,40 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the customizer for {@link HttpRequest.Builder}, to modify requests before
|
||||
* executing them.
|
||||
* <p>
|
||||
* This overrides the customizer from
|
||||
* {@link #asyncHttpRequestCustomizer(AsyncHttpRequestCustomizer)}.
|
||||
* <p>
|
||||
* Do NOT use a blocking {@link SyncHttpRequestCustomizer} in a non-blocking
|
||||
* context. Use {@link #asyncHttpRequestCustomizer(AsyncHttpRequestCustomizer)}
|
||||
* instead.
|
||||
* @param syncHttpRequestCustomizer the request customizer
|
||||
* @return this builder
|
||||
*/
|
||||
public Builder httpRequestCustomizer(SyncHttpRequestCustomizer syncHttpRequestCustomizer) {
|
||||
this.httpRequestCustomizer = AsyncHttpRequestCustomizer.fromSync(syncHttpRequestCustomizer);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets the customizer for {@link HttpRequest.Builder}, to modify requests before
|
||||
* executing them.
|
||||
* <p>
|
||||
* This overrides the customizer from
|
||||
* {@link #httpRequestCustomizer(SyncHttpRequestCustomizer)}.
|
||||
* <p>
|
||||
* Do NOT use a blocking implementation in a non-blocking context.
|
||||
* @param asyncHttpRequestCustomizer the request customizer
|
||||
* @return this builder
|
||||
*/
|
||||
public Builder asyncHttpRequestCustomizer(AsyncHttpRequestCustomizer asyncHttpRequestCustomizer) {
|
||||
this.httpRequestCustomizer = asyncHttpRequestCustomizer;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Construct a fresh instance of {@link HttpClientStreamableHttpTransport} using
|
||||
* the current builder configuration.
|
||||
|
|
@ -642,7 +699,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
|
|||
ObjectMapper objectMapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper();
|
||||
|
||||
return new HttpClientStreamableHttpTransport(objectMapper, clientBuilder.build(), requestBuilder, baseUri,
|
||||
endpoint, resumableStreams, openConnectionOnStartup);
|
||||
endpoint, resumableStreams, openConnectionOnStartup, httpRequestCustomizer);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,21 @@
|
|||
/*
|
||||
* Copyright 2024-2025 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import java.net.URI;
|
||||
import java.net.http.HttpRequest;
|
||||
import reactor.util.annotation.Nullable;
|
||||
|
||||
/**
|
||||
* Customize {@link HttpRequest.Builder} before executing the request, either in SSE or
|
||||
* Streamable HTTP transport.
|
||||
*
|
||||
* @author Daniel Garnier-Moiroux
|
||||
*/
|
||||
public interface SyncHttpRequestCustomizer {
|
||||
|
||||
void customize(HttpRequest.Builder builder, String method, URI endpoint, @Nullable String body);
|
||||
|
||||
}
|
||||
Loading…
Reference in a new issue