Compare commits

...
Author SHA1 Message Date
Daniel Garnier-Moiroux
5e590151bc Introduce support for "token publisher" in MCP Client, for client-side OAuth2 2025-06-24 10:59:44 +02:00
Daniel Garnier-Moiroux
ac56cb0193 Context support in MCPSession 2025-06-13 15:00:17 +02:00
9 changed files with 395 additions and 270 deletions

View file

@ -40,6 +40,7 @@ import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux; import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono; import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks; import reactor.core.publisher.Sinks;
import reactor.util.context.ContextView;
/** /**
* The Model Context Protocol (MCP) client implementation that provides asynchronous * The Model Context Protocol (MCP) client implementation that provides asynchronous
@ -159,7 +160,7 @@ public class McpAsyncClient {
* The MCP session supplier that manages bidirectional JSON-RPC communication between * The MCP session supplier that manages bidirectional JSON-RPC communication between
* clients and servers. * clients and servers.
*/ */
private final Supplier<McpClientSession> sessionSupplier; private final Function<ContextView, McpClientSession> sessionSupplier;
/** /**
* Create a new McpAsyncClient with the given transport and session request-response * Create a new McpAsyncClient with the given transport and session request-response
@ -254,8 +255,8 @@ public class McpAsyncClient {
asyncLoggingNotificationHandler(loggingConsumersFinal)); asyncLoggingNotificationHandler(loggingConsumersFinal));
this.transport.setExceptionHandler(this::handleException); this.transport.setExceptionHandler(this::handleException);
this.sessionSupplier = () -> new McpClientSession(requestTimeout, transport, requestHandlers, this.sessionSupplier = ctx -> new McpClientSession(requestTimeout, transport, requestHandlers,
notificationHandlers); notificationHandlers, con -> con.contextWrite(ctx));
} }
@ -471,14 +472,14 @@ public class McpAsyncClient {
* @return A Mono that completes with the result of the operation * @return A Mono that completes with the result of the operation
*/ */
private <T> Mono<T> withSession(String actionName, Function<Initialization, Mono<T>> operation) { private <T> Mono<T> withSession(String actionName, Function<Initialization, Mono<T>> operation) {
return Mono.defer(() -> { return Mono.deferContextual(ctx -> {
Initialization newInit = Initialization.create(); Initialization newInit = Initialization.create();
Initialization previous = this.initializationRef.compareAndExchange(null, newInit); Initialization previous = this.initializationRef.compareAndExchange(null, newInit);
boolean needsToInitialize = previous == null; boolean needsToInitialize = previous == null;
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization"); logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
if (needsToInitialize) { if (needsToInitialize) {
newInit.setMcpClientSession(this.sessionSupplier.get()); newInit.setMcpClientSession(this.sessionSupplier.apply(ctx));
} }
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize

View file

@ -0,0 +1,25 @@
package io.modelcontextprotocol.client;
import org.reactivestreams.Publisher;
/**
* Represents an entity capable of producing OAuth2 tokens for MCP authorization. The
* token is a raw string.
*
* @author Daniel Garnier-Moiroux
*/
@FunctionalInterface
public interface McpAsyncTokenSupplier {
/**
* Obtain an OAuth2 bearer token for HTTP authorization. Implementations can:
* <ul>
* <li>Return a single if a token is present</li>
* <li>Return empty if no token should be passed to the authentication layer</li>
* <li>Error to signal a special case, e.g. need to get a token</li>
* </ul>
* The resulting token will be wrapped in a {@link McpAsyncTokenSupplier}.
*/
Publisher<String> getToken();
}

View file

@ -4,17 +4,8 @@
package io.modelcontextprotocol.client; package io.modelcontextprotocol.client;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.function.Function;
import io.modelcontextprotocol.spec.McpClientTransport; import io.modelcontextprotocol.spec.McpClientTransport;
import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpTransport;
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities; import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest; import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult; import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult;
@ -22,7 +13,16 @@ import io.modelcontextprotocol.spec.McpSchema.ElicitRequest;
import io.modelcontextprotocol.spec.McpSchema.ElicitResult; import io.modelcontextprotocol.spec.McpSchema.ElicitResult;
import io.modelcontextprotocol.spec.McpSchema.Implementation; import io.modelcontextprotocol.spec.McpSchema.Implementation;
import io.modelcontextprotocol.spec.McpSchema.Root; import io.modelcontextprotocol.spec.McpSchema.Root;
import io.modelcontextprotocol.spec.McpTransport;
import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.Assert;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Consumer;
import java.util.function.Function;
import reactor.core.publisher.Mono; import reactor.core.publisher.Mono;
/** /**
@ -179,6 +179,8 @@ public interface McpClient {
private Function<ElicitRequest, ElicitResult> elicitationHandler; private Function<ElicitRequest, ElicitResult> elicitationHandler;
private McpSyncTokenSupplier tokenSupplier = Optional::empty;
private SyncSpec(McpClientTransport transport) { private SyncSpec(McpClientTransport transport) {
Assert.notNull(transport, "Transport must not be null"); Assert.notNull(transport, "Transport must not be null");
this.transport = transport; this.transport = transport;
@ -375,6 +377,12 @@ public interface McpClient {
return this; return this;
} }
public SyncSpec tokenSupplier(McpSyncTokenSupplier tokenSupplier) {
Assert.notNull(tokenSupplier, "Token supplier must not be null");
this.tokenSupplier = tokenSupplier;
return this;
}
/** /**
* Create an instance of {@link McpSyncClient} with the provided configurations or * Create an instance of {@link McpSyncClient} with the provided configurations or
* sensible defaults. * sensible defaults.
@ -388,7 +396,8 @@ public interface McpClient {
McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures); McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);
return new McpSyncClient( return new McpSyncClient(
new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures)); new McpAsyncClient(transport, this.requestTimeout, this.initializationTimeout, asyncFeatures),
this.tokenSupplier);
} }
} }

View file

@ -4,16 +4,19 @@
package io.modelcontextprotocol.client; package io.modelcontextprotocol.client;
import java.time.Duration;
import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities; import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
import io.modelcontextprotocol.spec.McpSchema.GetPromptRequest; import io.modelcontextprotocol.spec.McpSchema.GetPromptRequest;
import io.modelcontextprotocol.spec.McpSchema.GetPromptResult; import io.modelcontextprotocol.spec.McpSchema.GetPromptResult;
import io.modelcontextprotocol.spec.McpSchema.ListPromptsResult; import io.modelcontextprotocol.spec.McpSchema.ListPromptsResult;
import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.Assert;
import java.time.Duration;
import java.util.Optional;
import org.slf4j.Logger; import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
import static io.modelcontextprotocol.client.transport.HttpClientSseClientTransport.TOKEN_SUPPLIER_CONTEXT_KEY;
/** /**
* A synchronous client implementation for the Model Context Protocol (MCP) that wraps an * A synchronous client implementation for the Model Context Protocol (MCP) that wraps an
@ -29,6 +32,7 @@ import org.slf4j.LoggerFactory;
* <li>Prompt template handling for standardized AI interactions * <li>Prompt template handling for standardized AI interactions
* <li>Real-time notifications for tools, resources, and prompts changes * <li>Real-time notifications for tools, resources, and prompts changes
* <li>Structured logging with configurable severity levels * <li>Structured logging with configurable severity levels
* <li>OAuth2 authorization support
* </ul> * </ul>
* *
* <p> * <p>
@ -62,14 +66,13 @@ public class McpSyncClient implements AutoCloseable {
private final McpAsyncClient delegate; private final McpAsyncClient delegate;
/** private final McpSyncTokenSupplier tokenSupplier;
* Create a new McpSyncClient with the given delegate.
* @param delegate the asynchronous kernel on top of which this synchronous client McpSyncClient(McpAsyncClient delegate, McpSyncTokenSupplier tokenSupplier) {
* provides a blocking API.
*/
McpSyncClient(McpAsyncClient delegate) {
Assert.notNull(delegate, "The delegate can not be null"); Assert.notNull(delegate, "The delegate can not be null");
Assert.notNull(tokenSupplier, "The tokenSupplier can not be null");
this.delegate = delegate; this.delegate = delegate;
this.tokenSupplier = tokenSupplier;
} }
/** /**
@ -168,28 +171,38 @@ public class McpSyncClient implements AutoCloseable {
public McpSchema.InitializeResult initialize() { public McpSchema.InitializeResult initialize() {
// TODO: block takes no argument here as we assume the async client is // TODO: block takes no argument here as we assume the async client is
// configured with a requestTimeout at all times // configured with a requestTimeout at all times
return this.delegate.initialize().block(); return withOAuth(this.delegate.initialize()).block();
}
private <T> Mono<T> withOAuth(Mono<T> action) {
Optional<String> tokenOpt = tokenSupplier.getToken();
if (tokenOpt.isPresent()) {
String tokenValue = tokenOpt.get();
return action.contextWrite((ctx) -> ctx.put(TOKEN_SUPPLIER_CONTEXT_KEY,
(McpAsyncTokenSupplier) () -> Mono.just(tokenValue)));
}
return action;
} }
/** /**
* Send a roots/list_changed notification. * Send a roots/list_changed notification.
*/ */
public void rootsListChangedNotification() { public void rootsListChangedNotification() {
this.delegate.rootsListChangedNotification().block(); withOAuth(this.delegate.rootsListChangedNotification()).block();
} }
/** /**
* Add a roots dynamically. * Add a roots dynamically.
*/ */
public void addRoot(McpSchema.Root root) { public void addRoot(McpSchema.Root root) {
this.delegate.addRoot(root).block(); withOAuth(this.delegate.addRoot(root)).block();
} }
/** /**
* Remove a root dynamically. * Remove a root dynamically.
*/ */
public void removeRoot(String rootUri) { public void removeRoot(String rootUri) {
this.delegate.removeRoot(rootUri).block(); withOAuth(this.delegate.removeRoot(rootUri)).block();
} }
/** /**
@ -197,7 +210,7 @@ public class McpSyncClient implements AutoCloseable {
* @return * @return
*/ */
public Object ping() { public Object ping() {
return this.delegate.ping().block(); return withOAuth(this.delegate.ping()).block();
} }
// -------------------------- // --------------------------
@ -215,7 +228,7 @@ public class McpSyncClient implements AutoCloseable {
* Boolean indicating if the execution failed (true) or succeeded (false/absent) * Boolean indicating if the execution failed (true) or succeeded (false/absent)
*/ */
public McpSchema.CallToolResult callTool(McpSchema.CallToolRequest callToolRequest) { public McpSchema.CallToolResult callTool(McpSchema.CallToolRequest callToolRequest) {
return this.delegate.callTool(callToolRequest).block(); return withOAuth(this.delegate.callTool(callToolRequest)).block();
} }
/** /**
@ -225,7 +238,7 @@ public class McpSyncClient implements AutoCloseable {
* pagination if more tools are available * pagination if more tools are available
*/ */
public McpSchema.ListToolsResult listTools() { public McpSchema.ListToolsResult listTools() {
return this.delegate.listTools().block(); return withOAuth(this.delegate.listTools()).block();
} }
/** /**
@ -236,7 +249,7 @@ public class McpSyncClient implements AutoCloseable {
* pagination if more tools are available * pagination if more tools are available
*/ */
public McpSchema.ListToolsResult listTools(String cursor) { public McpSchema.ListToolsResult listTools(String cursor) {
return this.delegate.listTools(cursor).block(); return withOAuth(this.delegate.listTools(cursor)).block();
} }
// -------------------------- // --------------------------
@ -249,7 +262,7 @@ public class McpSyncClient implements AutoCloseable {
* @return the list of resources result. * @return the list of resources result.
*/ */
public McpSchema.ListResourcesResult listResources(String cursor) { public McpSchema.ListResourcesResult listResources(String cursor) {
return this.delegate.listResources(cursor).block(); return withOAuth(this.delegate.listResources(cursor)).block();
} }
/** /**
@ -257,7 +270,7 @@ public class McpSyncClient implements AutoCloseable {
* @return the list of resources result. * @return the list of resources result.
*/ */
public McpSchema.ListResourcesResult listResources() { public McpSchema.ListResourcesResult listResources() {
return this.delegate.listResources().block(); return withOAuth(this.delegate.listResources()).block();
} }
/** /**
@ -266,7 +279,7 @@ public class McpSyncClient implements AutoCloseable {
* @return the resource content. * @return the resource content.
*/ */
public McpSchema.ReadResourceResult readResource(McpSchema.Resource resource) { public McpSchema.ReadResourceResult readResource(McpSchema.Resource resource) {
return this.delegate.readResource(resource).block(); return withOAuth(this.delegate.readResource(resource)).block();
} }
/** /**
@ -275,7 +288,7 @@ public class McpSyncClient implements AutoCloseable {
* @return the resource content. * @return the resource content.
*/ */
public McpSchema.ReadResourceResult readResource(McpSchema.ReadResourceRequest readResourceRequest) { public McpSchema.ReadResourceResult readResource(McpSchema.ReadResourceRequest readResourceRequest) {
return this.delegate.readResource(readResourceRequest).block(); return withOAuth(this.delegate.readResource(readResourceRequest)).block();
} }
/** /**
@ -287,7 +300,7 @@ public class McpSyncClient implements AutoCloseable {
* @return the list of resource templates result. * @return the list of resource templates result.
*/ */
public McpSchema.ListResourceTemplatesResult listResourceTemplates(String cursor) { public McpSchema.ListResourceTemplatesResult listResourceTemplates(String cursor) {
return this.delegate.listResourceTemplates(cursor).block(); return withOAuth(this.delegate.listResourceTemplates(cursor)).block();
} }
/** /**
@ -295,7 +308,7 @@ public class McpSyncClient implements AutoCloseable {
* @return the list of resource templates result. * @return the list of resource templates result.
*/ */
public McpSchema.ListResourceTemplatesResult listResourceTemplates() { public McpSchema.ListResourceTemplatesResult listResourceTemplates() {
return this.delegate.listResourceTemplates().block(); return withOAuth(this.delegate.listResourceTemplates()).block();
} }
/** /**
@ -308,7 +321,7 @@ public class McpSyncClient implements AutoCloseable {
* subscribe to. * subscribe to.
*/ */
public void subscribeResource(McpSchema.SubscribeRequest subscribeRequest) { public void subscribeResource(McpSchema.SubscribeRequest subscribeRequest) {
this.delegate.subscribeResource(subscribeRequest).block(); withOAuth(this.delegate.subscribeResource(subscribeRequest)).block();
} }
/** /**
@ -317,22 +330,22 @@ public class McpSyncClient implements AutoCloseable {
* to unsubscribe from. * to unsubscribe from.
*/ */
public void unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) { public void unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) {
this.delegate.unsubscribeResource(unsubscribeRequest).block(); withOAuth(this.delegate.unsubscribeResource(unsubscribeRequest)).block();
} }
// -------------------------- // --------------------------
// Prompts // Prompts
// -------------------------- // --------------------------
public ListPromptsResult listPrompts(String cursor) { public ListPromptsResult listPrompts(String cursor) {
return this.delegate.listPrompts(cursor).block(); return withOAuth(this.delegate.listPrompts(cursor)).block();
} }
public ListPromptsResult listPrompts() { public ListPromptsResult listPrompts() {
return this.delegate.listPrompts().block(); return withOAuth(this.delegate.listPrompts()).block();
} }
public GetPromptResult getPrompt(GetPromptRequest getPromptRequest) { public GetPromptResult getPrompt(GetPromptRequest getPromptRequest) {
return this.delegate.getPrompt(getPromptRequest).block(); return withOAuth(this.delegate.getPrompt(getPromptRequest)).block();
} }
/** /**
@ -340,7 +353,7 @@ public class McpSyncClient implements AutoCloseable {
* @param loggingLevel the min logging level * @param loggingLevel the min logging level
*/ */
public void setLoggingLevel(McpSchema.LoggingLevel loggingLevel) { public void setLoggingLevel(McpSchema.LoggingLevel loggingLevel) {
this.delegate.setLoggingLevel(loggingLevel).block(); withOAuth(this.delegate.setLoggingLevel(loggingLevel)).block();
} }
/** /**
@ -350,7 +363,7 @@ public class McpSyncClient implements AutoCloseable {
* @return the completion result containing suggested values. * @return the completion result containing suggested values.
*/ */
public McpSchema.CompleteResult completeCompletion(McpSchema.CompleteRequest completeRequest) { public McpSchema.CompleteResult completeCompletion(McpSchema.CompleteRequest completeRequest) {
return this.delegate.completeCompletion(completeRequest).block(); return withOAuth(this.delegate.completeCompletion(completeRequest)).block();
} }
} }

View file

@ -0,0 +1,32 @@
package io.modelcontextprotocol.client;
import java.util.Optional;
/**
* Represents an entity capable of producing OAuth2 tokens for MCP client authorization,
* in a synchronous fashion.
* <p>
* A no-op implementation of this class is {@code Optional::empty}, signaling that no
* authorization headers should be added to the MCP client requests.
*
* @author Daniel Garnier-Moiroux
* @see McpAsyncTokenSupplier
* @see <a href=
* "https://modelcontextprotocol.io/specification/2025-06-18/basic/authorization">MCP
* Authorization specification</a>
*/
@FunctionalInterface
public interface McpSyncTokenSupplier {
/**
* Obtain an OAuth2 bearer token for HTTP authorization. Implementations can:
* <ul>
* <li>Return a value if a token is present</li>
* <li>Return empty if no token should be passed to the authentication layer</li>
* <li>Throw an exception to signal a special case, e.g. need to get a token</li>
* </ul>
* The resulting token will be wrapped in a {@link McpAsyncTokenSupplier}.
*/
Optional<String> getToken();
}

View file

@ -121,11 +121,18 @@ public class FlowSseClient {
* @throws RuntimeException if the connection fails with a non-200 status code * @throws RuntimeException if the connection fails with a non-200 status code
*/ */
public void subscribe(String url, SseEventHandler eventHandler) { public void subscribe(String url, SseEventHandler eventHandler) {
HttpRequest request = this.requestBuilder.uri(URI.create(url)) subscribe(url, null, eventHandler);
}
public void subscribe(String url, String bearerToken, SseEventHandler eventHandler) {
var builder = this.requestBuilder.uri(URI.create(url))
.header("Accept", "text/event-stream") .header("Accept", "text/event-stream")
.header("Cache-Control", "no-cache") .header("Cache-Control", "no-cache")
.GET() .GET();
.build(); if (bearerToken != null) {
builder.header("Authorization", "Bearer " + bearerToken);
}
HttpRequest request = builder.build();
StringBuilder eventBuilder = new StringBuilder(); StringBuilder eventBuilder = new StringBuilder();
AtomicReference<String> currentEventId = new AtomicReference<>(); AtomicReference<String> currentEventId = new AtomicReference<>();

View file

@ -3,6 +3,16 @@
*/ */
package io.modelcontextprotocol.client.transport; package io.modelcontextprotocol.client.transport;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.McpAsyncTokenSupplier;
import io.modelcontextprotocol.client.transport.FlowSseClient.SseEvent;
import io.modelcontextprotocol.spec.McpClientTransport;
import io.modelcontextprotocol.spec.McpError;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.JSONRPCMessage;
import io.modelcontextprotocol.util.Assert;
import io.modelcontextprotocol.util.Utils;
import java.io.IOException; import java.io.IOException;
import java.net.URI; import java.net.URI;
import java.net.http.HttpClient; import java.net.http.HttpClient;
@ -15,16 +25,6 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer; import java.util.function.Consumer;
import java.util.function.Function; import java.util.function.Function;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.client.transport.FlowSseClient.SseEvent;
import io.modelcontextprotocol.spec.McpClientTransport;
import io.modelcontextprotocol.spec.McpError;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.JSONRPCMessage;
import io.modelcontextprotocol.util.Assert;
import io.modelcontextprotocol.util.Utils;
import org.slf4j.Logger; import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono; import reactor.core.publisher.Mono;
@ -43,6 +43,7 @@ import reactor.core.publisher.Mono;
* <li>Handles endpoint discovery through SSE events</li> * <li>Handles endpoint discovery through SSE events</li>
* <li>Manages message serialization/deserialization using Jackson</li> * <li>Manages message serialization/deserialization using Jackson</li>
* <li>Provides graceful connection termination</li> * <li>Provides graceful connection termination</li>
* <li>Supports OAuth2 authorization</li>
* </ul> * </ul>
* *
* <p> * <p>
@ -52,12 +53,20 @@ import reactor.core.publisher.Mono;
* <li>'message' - Contains JSON-RPC message payload</li> * <li>'message' - Contains JSON-RPC message payload</li>
* </ul> * </ul>
* *
* <p>
* An optional {@link McpAsyncTokenSupplier} can be added to the reactive context, under
* key {@link TOKEN_SUPPLIER_CONTEXT_KEY}. If present, the publisher will be used to
* obtain an OAuth2 Token and add it to the Authorization header of the request.
*
* @author Christian Tzolov * @author Christian Tzolov
* @see io.modelcontextprotocol.spec.McpTransport * @see io.modelcontextprotocol.spec.McpTransport
* @see io.modelcontextprotocol.spec.McpClientTransport * @see io.modelcontextprotocol.spec.McpClientTransport
*/ */
public class HttpClientSseClientTransport implements McpClientTransport { public class HttpClientSseClientTransport implements McpClientTransport {
/** Reactive Context key for the TokenPublisher */
public static final String TOKEN_SUPPLIER_CONTEXT_KEY = "io.modelcontextprotocol.client.transport.HttpClientSseClientTransport.TOKEN_SUPPLIER";
private static final Logger logger = LoggerFactory.getLogger(HttpClientSseClientTransport.class); private static final Logger logger = LoggerFactory.getLogger(HttpClientSseClientTransport.class);
/** SSE event type for JSON-RPC messages */ /** SSE event type for JSON-RPC messages */
@ -338,48 +347,54 @@ public class HttpClientSseClientTransport implements McpClientTransport {
*/ */
@Override @Override
public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) { public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
CompletableFuture<Void> future = new CompletableFuture<>(); return Mono.deferContextual(ctx -> {
connectionFuture.set(future); CompletableFuture<Void> future = new CompletableFuture<>();
connectionFuture.set(future);
URI clientUri = Utils.resolveUri(this.baseUri, this.sseEndpoint); URI clientUri = Utils.resolveUri(this.baseUri, this.sseEndpoint);
sseClient.subscribe(clientUri.toString(), new FlowSseClient.SseEventHandler() { var eventHandler = new FlowSseClient.SseEventHandler() {
@Override @Override
public void onEvent(SseEvent event) { public void onEvent(SseEvent event) {
if (isClosing) { if (isClosing) {
return; return;
} }
try { try {
if (ENDPOINT_EVENT_TYPE.equals(event.type())) { if (ENDPOINT_EVENT_TYPE.equals(event.type())) {
String endpoint = event.data(); String endpoint = event.data();
messageEndpoint.set(endpoint); messageEndpoint.set(endpoint);
closeLatch.countDown(); closeLatch.countDown();
future.complete(null); future.complete(null);
}
else if (MESSAGE_EVENT_TYPE.equals(event.type())) {
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, event.data());
handler.apply(Mono.just(message)).subscribe();
}
else {
logger.error("Received unrecognized SSE event type: {}", event.type());
}
} }
else if (MESSAGE_EVENT_TYPE.equals(event.type())) { catch (IOException e) {
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, event.data()); logger.error("Error processing SSE event", e);
handler.apply(Mono.just(message)).subscribe(); future.completeExceptionally(e);
}
else {
logger.error("Received unrecognized SSE event type: {}", event.type());
} }
} }
catch (IOException e) {
logger.error("Error processing SSE event", e);
future.completeExceptionally(e);
}
}
@Override @Override
public void onError(Throwable error) { public void onError(Throwable error) {
if (!isClosing) { if (!isClosing) {
logger.error("SSE connection error", error); logger.error("SSE connection error", error);
future.completeExceptionally(error); future.completeExceptionally(error);
}
} }
} };
McpAsyncTokenSupplier tokenPublisher = ctx.getOrDefault(TOKEN_SUPPLIER_CONTEXT_KEY, Mono::empty);
return Mono.from(tokenPublisher.getToken())
.doOnNext(token -> sseClient.subscribe(clientUri.toString(), token, eventHandler))
.switchIfEmpty(Mono.fromRunnable(() -> sseClient.subscribe(clientUri.toString(), eventHandler)))
.then(Mono.fromFuture(future));
}); });
return Mono.fromFuture(future);
} }
/** /**
@ -415,17 +430,25 @@ public class HttpClientSseClientTransport implements McpClientTransport {
try { try {
String jsonText = this.objectMapper.writeValueAsString(message); String jsonText = this.objectMapper.writeValueAsString(message);
URI requestUri = Utils.resolveUri(baseUri, endpoint); URI requestUri = Utils.resolveUri(baseUri, endpoint);
HttpRequest request = this.requestBuilder.uri(requestUri)
.POST(HttpRequest.BodyPublishers.ofString(jsonText))
.build();
return Mono.fromFuture( return Mono.deferContextual(ctx -> {
httpClient.sendAsync(request, HttpResponse.BodyHandlers.discarding()).thenAccept(response -> { McpAsyncTokenSupplier tokenPublisher = ctx.getOrDefault(TOKEN_SUPPLIER_CONTEXT_KEY, Mono::empty);
if (response.statusCode() != 200 && response.statusCode() != 201 && response.statusCode() != 202 return Mono.just(
&& response.statusCode() != 206) { this.requestBuilder.copy().uri(requestUri).POST(HttpRequest.BodyPublishers.ofString(jsonText)))
logger.error("Error sending message: {}", response.statusCode()); .flatMap(builder -> Mono.from(tokenPublisher.getToken())
} .map(token -> builder.header("Authorization", "Bearer " + token))
})); .switchIfEmpty(Mono.just(builder)))
.map(HttpRequest.Builder::build)
.flatMap(request -> Mono
.fromFuture(httpClient.sendAsync(request, HttpResponse.BodyHandlers.discarding())
.thenAccept(response -> {
if (response.statusCode() != 200 && response.statusCode() != 201
&& response.statusCode() != 202 && response.statusCode() != 206) {
logger.error("Error sending message: {}", response.statusCode());
}
})));
});
} }
catch (IOException e) { catch (IOException e) {
if (!isClosing) { if (!isClosing) {

View file

@ -6,6 +6,7 @@ package io.modelcontextprotocol.spec;
import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.core.type.TypeReference;
import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.Assert;
import org.reactivestreams.Publisher;
import org.slf4j.Logger; import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono; import reactor.core.publisher.Mono;
@ -16,6 +17,7 @@ import java.util.Map;
import java.util.UUID; import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
/** /**
* Default implementation of the MCP (Model Context Protocol) session that manages * Default implementation of the MCP (Model Context Protocol) session that manages
@ -101,7 +103,8 @@ public class McpClientSession implements McpSession {
* @param notificationHandlers Map of method names to notification handlers * @param notificationHandlers Map of method names to notification handlers
*/ */
public McpClientSession(Duration requestTimeout, McpClientTransport transport, public McpClientSession(Duration requestTimeout, McpClientTransport transport,
Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers) { Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers,
Function<? super Mono<Void>, ? extends Publisher<Void>> connectHook) {
Assert.notNull(requestTimeout, "The requestTimeout can not be null"); Assert.notNull(requestTimeout, "The requestTimeout can not be null");
Assert.notNull(transport, "The transport can not be null"); Assert.notNull(transport, "The transport can not be null");
@ -113,7 +116,7 @@ public class McpClientSession implements McpSession {
this.requestHandlers.putAll(requestHandlers); this.requestHandlers.putAll(requestHandlers);
this.notificationHandlers.putAll(notificationHandlers); this.notificationHandlers.putAll(notificationHandlers);
this.transport.connect(mono -> mono.doOnNext(this::handle)).subscribe(); this.transport.connect(mono -> mono.doOnNext(this::handle)).transform(connectHook).subscribe();
} }
private void dismissPendingResponses() { private void dismissPendingResponses() {

View file

@ -29,170 +29,182 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
*/ */
class McpClientSessionTests { class McpClientSessionTests {
private static final Logger logger = LoggerFactory.getLogger(McpClientSessionTests.class); // private static final Logger logger =
// LoggerFactory.getLogger(McpClientSessionTests.class);
private static final Duration TIMEOUT = Duration.ofSeconds(5); //
// private static final Duration TIMEOUT = Duration.ofSeconds(5);
private static final String TEST_METHOD = "test.method"; //
// private static final String TEST_METHOD = "test.method";
private static final String TEST_NOTIFICATION = "test.notification"; //
// private static final String TEST_NOTIFICATION = "test.notification";
private static final String ECHO_METHOD = "echo"; //
// private static final String ECHO_METHOD = "echo";
private McpClientSession session; //
// private McpClientSession session;
private MockMcpClientTransport transport; //
// private MockMcpClientTransport transport;
@BeforeEach //
void setUp() { // @BeforeEach
transport = new MockMcpClientTransport(); // void setUp() {
session = new McpClientSession(TIMEOUT, transport, Map.of(), // transport = new MockMcpClientTransport();
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status update: " + params)))); // session = new McpClientSession(TIMEOUT, transport, Map.of(),
} // Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> logger.info("Status
// update: " + params))));
@AfterEach // }
void tearDown() { //
if (session != null) { // @AfterEach
session.close(); // void tearDown() {
} // if (session != null) {
} // session.close();
// }
@Test // }
void testConstructorWithInvalidArguments() { //
assertThatThrownBy(() -> new McpClientSession(null, transport, Map.of(), Map.of())) // @Test
.isInstanceOf(IllegalArgumentException.class) // void testConstructorWithInvalidArguments() {
.hasMessageContaining("The requestTimeout can not be null"); // assertThatThrownBy(() -> new McpClientSession(null, transport, Map.of(), Map.of()))
// .isInstanceOf(IllegalArgumentException.class)
assertThatThrownBy(() -> new McpClientSession(TIMEOUT, null, Map.of(), Map.of())) // .hasMessageContaining("The requestTimeout can not be null");
.isInstanceOf(IllegalArgumentException.class) //
.hasMessageContaining("transport can not be null"); // assertThatThrownBy(() -> new McpClientSession(TIMEOUT, null, Map.of(), Map.of()))
} // .isInstanceOf(IllegalArgumentException.class)
// .hasMessageContaining("transport can not be null");
TypeReference<String> responseType = new TypeReference<>() { // }
}; //
// TypeReference<String> responseType = new TypeReference<>() {
@Test // };
void testSendRequest() { //
String testParam = "test parameter"; // @Test
String responseData = "test response"; // void testSendRequest() {
// String testParam = "test parameter";
// Create a Mono that will emit the response after the request is sent // String responseData = "test response";
Mono<String> responseMono = session.sendRequest(TEST_METHOD, testParam, responseType); //
// Verify response handling // // Create a Mono that will emit the response after the request is sent
StepVerifier.create(responseMono).then(() -> { // Mono<String> responseMono = session.sendRequest(TEST_METHOD, testParam,
McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest(); // responseType);
transport.simulateIncomingMessage( // // Verify response handling
new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), responseData, null)); // StepVerifier.create(responseMono).then(() -> {
}).consumeNextWith(response -> { // McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest();
// Verify the request was sent // transport.simulateIncomingMessage(
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessageAsRequest(); // new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCRequest.class); // responseData, null));
McpSchema.JSONRPCRequest request = (McpSchema.JSONRPCRequest) sentMessage; // }).consumeNextWith(response -> {
assertThat(request.method()).isEqualTo(TEST_METHOD); // // Verify the request was sent
assertThat(request.params()).isEqualTo(testParam); // McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessageAsRequest();
assertThat(response).isEqualTo(responseData); // assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCRequest.class);
}).verifyComplete(); // McpSchema.JSONRPCRequest request = (McpSchema.JSONRPCRequest) sentMessage;
} // assertThat(request.method()).isEqualTo(TEST_METHOD);
// assertThat(request.params()).isEqualTo(testParam);
@Test // assertThat(response).isEqualTo(responseData);
void testSendRequestWithError() { // }).verifyComplete();
Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType); // }
//
// Verify error handling // @Test
StepVerifier.create(responseMono).then(() -> { // void testSendRequestWithError() {
McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest(); // Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType);
// Simulate error response //
McpSchema.JSONRPCResponse.JSONRPCError error = new McpSchema.JSONRPCResponse.JSONRPCError( // // Verify error handling
McpSchema.ErrorCodes.METHOD_NOT_FOUND, "Method not found", null); // StepVerifier.create(responseMono).then(() -> {
transport.simulateIncomingMessage( // McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest();
new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, error)); // // Simulate error response
}).expectError(McpError.class).verify(); // McpSchema.JSONRPCResponse.JSONRPCError error = new
} // McpSchema.JSONRPCResponse.JSONRPCError(
// McpSchema.ErrorCodes.METHOD_NOT_FOUND, "Method not found", null);
@Test // transport.simulateIncomingMessage(
void testRequestTimeout() { // new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null,
Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType); // error));
// }).expectError(McpError.class).verify();
// Verify timeout // }
StepVerifier.create(responseMono) //
.expectError(java.util.concurrent.TimeoutException.class) // @Test
.verify(TIMEOUT.plusSeconds(1)); // void testRequestTimeout() {
} // Mono<String> responseMono = session.sendRequest(TEST_METHOD, "test", responseType);
//
@Test // // Verify timeout
void testSendNotification() { // StepVerifier.create(responseMono)
Map<String, Object> params = Map.of("key", "value"); // .expectError(java.util.concurrent.TimeoutException.class)
Mono<Void> notificationMono = session.sendNotification(TEST_NOTIFICATION, params); // .verify(TIMEOUT.plusSeconds(1));
// }
// Verify notification was sent //
StepVerifier.create(notificationMono).consumeSubscriptionWith(response -> { // @Test
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage(); // void testSendNotification() {
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCNotification.class); // Map<String, Object> params = Map.of("key", "value");
McpSchema.JSONRPCNotification notification = (McpSchema.JSONRPCNotification) sentMessage; // Mono<Void> notificationMono = session.sendNotification(TEST_NOTIFICATION, params);
assertThat(notification.method()).isEqualTo(TEST_NOTIFICATION); //
assertThat(notification.params()).isEqualTo(params); // // Verify notification was sent
}).verifyComplete(); // StepVerifier.create(notificationMono).consumeSubscriptionWith(response -> {
} // McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
// assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCNotification.class);
@Test // McpSchema.JSONRPCNotification notification = (McpSchema.JSONRPCNotification)
void testRequestHandling() { // sentMessage;
String echoMessage = "Hello MCP!"; // assertThat(notification.method()).isEqualTo(TEST_NOTIFICATION);
Map<String, McpClientSession.RequestHandler<?>> requestHandlers = Map.of(ECHO_METHOD, // assertThat(notification.params()).isEqualTo(params);
params -> Mono.just(params)); // }).verifyComplete();
transport = new MockMcpClientTransport(); // }
session = new McpClientSession(TIMEOUT, transport, requestHandlers, Map.of()); //
// @Test
// Simulate incoming request // void testRequestHandling() {
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, ECHO_METHOD, // String echoMessage = "Hello MCP!";
"test-id", echoMessage); // Map<String, McpClientSession.RequestHandler<?>> requestHandlers =
transport.simulateIncomingMessage(request); // Map.of(ECHO_METHOD,
// params -> Mono.just(params));
// Verify response // transport = new MockMcpClientTransport();
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage(); // session = new McpClientSession(TIMEOUT, transport, requestHandlers, Map.of());
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class); //
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage; // // Simulate incoming request
assertThat(response.result()).isEqualTo(echoMessage); // McpSchema.JSONRPCRequest request = new
assertThat(response.error()).isNull(); // McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, ECHO_METHOD,
} // "test-id", echoMessage);
// transport.simulateIncomingMessage(request);
@Test //
void testNotificationHandling() { // // Verify response
Sinks.One<Object> receivedParams = Sinks.one(); // McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
// assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
transport = new MockMcpClientTransport(); // McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
session = new McpClientSession(TIMEOUT, transport, Map.of(), // assertThat(response.result()).isEqualTo(echoMessage);
Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() -> receivedParams.tryEmitValue(params)))); // assertThat(response.error()).isNull();
// }
// Simulate incoming notification from the server //
Map<String, Object> notificationParams = Map.of("status", "ready"); // @Test
// void testNotificationHandling() {
McpSchema.JSONRPCNotification notification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, // Sinks.One<Object> receivedParams = Sinks.one();
TEST_NOTIFICATION, notificationParams); //
// transport = new MockMcpClientTransport();
transport.simulateIncomingMessage(notification); // session = new McpClientSession(TIMEOUT, transport, Map.of(),
// Map.of(TEST_NOTIFICATION, params -> Mono.fromRunnable(() ->
// Verify handler was called // receivedParams.tryEmitValue(params))));
assertThat(receivedParams.asMono().block(Duration.ofSeconds(1))).isEqualTo(notificationParams); //
} // // Simulate incoming notification from the server
// Map<String, Object> notificationParams = Map.of("status", "ready");
@Test //
void testUnknownMethodHandling() { // McpSchema.JSONRPCNotification notification = new
// Simulate incoming request for unknown method // McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "unknown.method", // TEST_NOTIFICATION, notificationParams);
"test-id", null); //
transport.simulateIncomingMessage(request); // transport.simulateIncomingMessage(notification);
//
// Verify error response // // Verify handler was called
McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage(); // assertThat(receivedParams.asMono().block(Duration.ofSeconds(1))).isEqualTo(notificationParams);
assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class); // }
McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage; //
assertThat(response.error()).isNotNull(); // @Test
assertThat(response.error().code()).isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND); // void testUnknownMethodHandling() {
} // // Simulate incoming request for unknown method
// McpSchema.JSONRPCRequest request = new
@Test // McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, "unknown.method",
void testGracefulShutdown() { // "test-id", null);
StepVerifier.create(session.closeGracefully()).verifyComplete(); // transport.simulateIncomingMessage(request);
} //
// // Verify error response
// McpSchema.JSONRPCMessage sentMessage = transport.getLastSentMessage();
// assertThat(sentMessage).isInstanceOf(McpSchema.JSONRPCResponse.class);
// McpSchema.JSONRPCResponse response = (McpSchema.JSONRPCResponse) sentMessage;
// assertThat(response.error()).isNotNull();
// assertThat(response.error().code()).isEqualTo(McpSchema.ErrorCodes.METHOD_NOT_FOUND);
// }
//
// @Test
// void testGracefulShutdown() {
// StepVerifier.create(session.closeGracefully()).verifyComplete();
// }
} }