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

feat: Propagate Context to eager connect via McpClientSession (#339)

Browse source 
* feat: Propagate Context to eager connect via McpClientSession

In order to allow the initial connection to have contextual information
in the reactive chain, the McpClientSession should be able to transform
the McpClientTransport#connect result, e.g. to attach Context items.
This change introduces a new constructor for sessions that makes it
possible.

Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>

---------

Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
This commit is contained in:
Dariusz Jędrzejczyk 2025-06-26 13:10:17 +02:00 • committed by GitHub
commit 4f93ce8eea
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 110 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

18
mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java
View file

@ -11,7 +11,6 @@ import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.function.Supplier;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@ -42,6 +41,7 @@ import io.modelcontextprotocol.util.Utils;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.util.context.ContextView;
/**
* The Model Context Protocol (MCP) client implementation that provides asynchronous
@ -161,7 +161,7 @@ public class McpAsyncClient {
* The MCP session supplier that manages bidirectional JSON-RPC communication between
* 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
@ -268,9 +268,8 @@ public class McpAsyncClient {
asyncLoggingNotificationHandler(loggingConsumersFinal));
this.transport.setExceptionHandler(this::handleException);
this.sessionSupplier = () -> new McpClientSession(requestTimeout, transport, requestHandlers,
notificationHandlers);
this.sessionSupplier = ctx -> new McpClientSession(requestTimeout, transport, requestHandlers,
notificationHandlers, con -> con.contextWrite(ctx));
}
private void handleException(Throwable t) {
@ -401,9 +400,8 @@ public class McpAsyncClient {
return withSession("by explicit API call", init -> Mono.just(init.get()));
}
private Mono<McpSchema.InitializeResult> doInitialize(Initialization initialization) {
initialization.setMcpClientSession(this.sessionSupplier.get());
private Mono<McpSchema.InitializeResult> doInitialize(Initialization initialization, ContextView ctx) {
initialization.setMcpClientSession(this.sessionSupplier.apply(ctx));
McpClientSession mcpClientSession = initialization.mcpSession();
@ -493,14 +491,14 @@ public class McpAsyncClient {
* @return A Mono that completes with the result of the 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 previous = this.initializationRef.compareAndExchange(null, newInit);
boolean needsToInitialize = previous == null;
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize ? doInitialize(newInit)
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize ? doInitialize(newInit, ctx)
: previous.await();
return initializationJob.map(initializeResult -> this.initializationRef.get())

22
mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java
View file

@ -6,6 +6,7 @@ package io.modelcontextprotocol.spec;
import com.fasterxml.jackson.core.type.TypeReference;
import io.modelcontextprotocol.util.Assert;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
@ -16,6 +17,7 @@ import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Function;
/**
* Default implementation of the MCP (Model Context Protocol) session that manages
@ -99,9 +101,27 @@ public class McpClientSession implements McpSession {
* @param transport Transport implementation for message exchange
* @param requestHandlers Map of method names to request handlers
* @param notificationHandlers Map of method names to notification handlers
* @deprecated Use
* {@link #McpClientSession(Duration, McpClientTransport, Map, Map, Function)}
*/
@Deprecated
public McpClientSession(Duration requestTimeout, McpClientTransport transport,
Map<String, RequestHandler<?>> requestHandlers, Map<String, NotificationHandler> notificationHandlers) {
this(requestTimeout, transport, requestHandlers, notificationHandlers, Function.identity());
}
/**
* Creates a new McpClientSession with the specified configuration and handlers.
* @param requestTimeout Duration to wait for responses
* @param transport Transport implementation for message exchange
* @param requestHandlers Map of method names to request handlers
* @param notificationHandlers Map of method names to notification handlers
* @param connectHook Hook that allows transforming the connection Publisher prior to
* subscribing
*/
public McpClientSession(Duration requestTimeout, McpClientTransport transport,
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(transport, "The transport can not be null");
@ -113,7 +133,7 @@ public class McpClientSession implements McpSession {
this.requestHandlers.putAll(requestHandlers);
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() {