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
1 changed files with 110 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

81
mcp/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java Normal file
View file

@ -0,0 +1,81 @@
package io.modelcontextprotocol.client;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.modelcontextprotocol.spec.McpClientTransport;
import io.modelcontextprotocol.spec.McpSchema;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Mono;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import static org.assertj.core.api.Assertions.assertThatCode;
class McpAsyncClientTests {
public static final McpSchema.Implementation MOCK_SERVER_INFO = new McpSchema.Implementation("test-server",
"1.0.0");
public static final McpSchema.ServerCapabilities MOCK_SERVER_CAPABILITIES = McpSchema.ServerCapabilities.builder()
.build();
public static final McpSchema.InitializeResult MOCK_INIT_RESULT = new McpSchema.InitializeResult(
McpSchema.LATEST_PROTOCOL_VERSION, MOCK_SERVER_CAPABILITIES, MOCK_SERVER_INFO, "Test instructions");
private static final String CONTEXT_KEY = "context.key";
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
@Test
void validateContextPassedToTransportConnect() {
McpClientTransport transport = new McpClientTransport() {
Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler;
final AtomicReference<String> contextValue = new AtomicReference<>();
@Override
public Mono<Void> connect(
Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
return Mono.deferContextual(ctx -> {
this.handler = handler;
if (ctx.hasKey(CONTEXT_KEY)) {
this.contextValue.set(ctx.get(CONTEXT_KEY));
}
return Mono.empty();
});
}
@Override
public Mono<Void> closeGracefully() {
return Mono.empty();
}
@Override
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
if (!"hello".equals(this.contextValue.get())) {
return Mono.error(new RuntimeException("Context value not propagated via #connect method"));
}
// We're only interested in handling the init request to provide an init
// response
if (!(message instanceof McpSchema.JSONRPCRequest)) {
return Mono.empty();
}
McpSchema.JSONRPCResponse initResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION,
((McpSchema.JSONRPCRequest) message).id(), MOCK_INIT_RESULT, null);
return handler.apply(Mono.just(initResponse)).then();
}
@Override
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
return OBJECT_MAPPER.convertValue(data, typeRef);
}
};
assertThatCode(() -> {
McpAsyncClient client = McpClient.async(transport).build();
client.initialize().contextWrite(ctx -> ctx.put(CONTEXT_KEY, "hello")).block();
}).doesNotThrowAnyException();
}
}