Introduce support for "token publisher" in MCP Client, for client-side OAuth2

This commit is contained in:
Daniel Garnier-Moiroux 2025-06-13 14:59:46 +02:00
commit 5e590151bc
6 changed files with 207 additions and 98 deletions

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;
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.McpSchema;
import io.modelcontextprotocol.spec.McpTransport;
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest;
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.Implementation;
import io.modelcontextprotocol.spec.McpSchema.Root;
import io.modelcontextprotocol.spec.McpTransport;
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;
/**
@ -179,6 +179,8 @@ public interface McpClient {
private Function<ElicitRequest, ElicitResult> elicitationHandler;
private McpSyncTokenSupplier tokenSupplier = Optional::empty;
private SyncSpec(McpClientTransport transport) {
Assert.notNull(transport, "Transport must not be null");
this.transport = transport;
@ -375,6 +377,12 @@ public interface McpClient {
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
* sensible defaults.
@ -388,7 +396,8 @@ public interface McpClient {
McpClientFeatures.Async asyncFeatures = McpClientFeatures.Async.fromSync(syncFeatures);
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;
import java.time.Duration;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
import io.modelcontextprotocol.spec.McpSchema.GetPromptRequest;
import io.modelcontextprotocol.spec.McpSchema.GetPromptResult;
import io.modelcontextprotocol.spec.McpSchema.ListPromptsResult;
import io.modelcontextprotocol.util.Assert;
import java.time.Duration;
import java.util.Optional;
import org.slf4j.Logger;
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
@ -29,6 +32,7 @@ import org.slf4j.LoggerFactory;
* <li>Prompt template handling for standardized AI interactions
* <li>Real-time notifications for tools, resources, and prompts changes
* <li>Structured logging with configurable severity levels
* <li>OAuth2 authorization support
* </ul>
*
* <p>
@ -62,14 +66,13 @@ public class McpSyncClient implements AutoCloseable {
private final McpAsyncClient delegate;
/**
* Create a new McpSyncClient with the given delegate.
* @param delegate the asynchronous kernel on top of which this synchronous client
* provides a blocking API.
*/
McpSyncClient(McpAsyncClient delegate) {
private final McpSyncTokenSupplier tokenSupplier;
McpSyncClient(McpAsyncClient delegate, McpSyncTokenSupplier tokenSupplier) {
Assert.notNull(delegate, "The delegate can not be null");
Assert.notNull(tokenSupplier, "The tokenSupplier can not be null");
this.delegate = delegate;
this.tokenSupplier = tokenSupplier;
}
/**
@ -168,28 +171,38 @@ public class McpSyncClient implements AutoCloseable {
public McpSchema.InitializeResult initialize() {
// TODO: block takes no argument here as we assume the async client is
// 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.
*/
public void rootsListChangedNotification() {
this.delegate.rootsListChangedNotification().block();
withOAuth(this.delegate.rootsListChangedNotification()).block();
}
/**
* Add a roots dynamically.
*/
public void addRoot(McpSchema.Root root) {
this.delegate.addRoot(root).block();
withOAuth(this.delegate.addRoot(root)).block();
}
/**
* Remove a root dynamically.
*/
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
*/
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)
*/
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
*/
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
*/
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.
*/
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.
*/
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.
*/
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.
*/
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.
*/
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.
*/
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.
*/
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.
*/
public void unsubscribeResource(McpSchema.UnsubscribeRequest unsubscribeRequest) {
this.delegate.unsubscribeResource(unsubscribeRequest).block();
withOAuth(this.delegate.unsubscribeResource(unsubscribeRequest)).block();
}
// --------------------------
// Prompts
// --------------------------
public ListPromptsResult listPrompts(String cursor) {
return this.delegate.listPrompts(cursor).block();
return withOAuth(this.delegate.listPrompts(cursor)).block();
}
public ListPromptsResult listPrompts() {
return this.delegate.listPrompts().block();
return withOAuth(this.delegate.listPrompts()).block();
}
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
*/
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.
*/
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
*/
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("Cache-Control", "no-cache")
.GET()
.build();
.GET();
if (bearerToken != null) {
builder.header("Authorization", "Bearer " + bearerToken);
}
HttpRequest request = builder.build();
StringBuilder eventBuilder = new StringBuilder();
AtomicReference<String> currentEventId = new AtomicReference<>();

View file

@ -3,6 +3,16 @@
*/
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.net.URI;
import java.net.http.HttpClient;
@ -15,16 +25,6 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
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.LoggerFactory;
import reactor.core.publisher.Mono;
@ -43,6 +43,7 @@ import reactor.core.publisher.Mono;
* <li>Handles endpoint discovery through SSE events</li>
* <li>Manages message serialization/deserialization using Jackson</li>
* <li>Provides graceful connection termination</li>
* <li>Supports OAuth2 authorization</li>
* </ul>
*
* <p>
@ -52,12 +53,20 @@ import reactor.core.publisher.Mono;
* <li>'message' - Contains JSON-RPC message payload</li>
* </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
* @see io.modelcontextprotocol.spec.McpTransport
* @see io.modelcontextprotocol.spec.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);
/** SSE event type for JSON-RPC messages */
@ -338,48 +347,54 @@ public class HttpClientSseClientTransport implements McpClientTransport {
*/
@Override
public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
CompletableFuture<Void> future = new CompletableFuture<>();
connectionFuture.set(future);
return Mono.deferContextual(ctx -> {
CompletableFuture<Void> future = new CompletableFuture<>();
connectionFuture.set(future);
URI clientUri = Utils.resolveUri(this.baseUri, this.sseEndpoint);
sseClient.subscribe(clientUri.toString(), new FlowSseClient.SseEventHandler() {
@Override
public void onEvent(SseEvent event) {
if (isClosing) {
return;
}
URI clientUri = Utils.resolveUri(this.baseUri, this.sseEndpoint);
var eventHandler = new FlowSseClient.SseEventHandler() {
@Override
public void onEvent(SseEvent event) {
if (isClosing) {
return;
}
try {
if (ENDPOINT_EVENT_TYPE.equals(event.type())) {
String endpoint = event.data();
messageEndpoint.set(endpoint);
closeLatch.countDown();
future.complete(null);
try {
if (ENDPOINT_EVENT_TYPE.equals(event.type())) {
String endpoint = event.data();
messageEndpoint.set(endpoint);
closeLatch.countDown();
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())) {
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, event.data());
handler.apply(Mono.just(message)).subscribe();
}
else {
logger.error("Received unrecognized SSE event type: {}", event.type());
catch (IOException e) {
logger.error("Error processing SSE event", e);
future.completeExceptionally(e);
}
}
catch (IOException e) {
logger.error("Error processing SSE event", e);
future.completeExceptionally(e);
}
}
@Override
public void onError(Throwable error) {
if (!isClosing) {
logger.error("SSE connection error", error);
future.completeExceptionally(error);
@Override
public void onError(Throwable error) {
if (!isClosing) {
logger.error("SSE connection error", 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 {
String jsonText = this.objectMapper.writeValueAsString(message);
URI requestUri = Utils.resolveUri(baseUri, endpoint);
HttpRequest request = this.requestBuilder.uri(requestUri)
.POST(HttpRequest.BodyPublishers.ofString(jsonText))
.build();
return 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());
}
}));
return Mono.deferContextual(ctx -> {
McpAsyncTokenSupplier tokenPublisher = ctx.getOrDefault(TOKEN_SUPPLIER_CONTEXT_KEY, Mono::empty);
return Mono.just(
this.requestBuilder.copy().uri(requestUri).POST(HttpRequest.BodyPublishers.ofString(jsonText)))
.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) {
if (!isClosing) {