refactor: consolidate MCP client initialization logic (#334)
- Move session setup, completion, and error callbacks into doInitialize method - Rename variables in McpClientSession for better clarity (sink -> pendingResponseSink/deliveredResponseSink) Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
This commit is contained in:
parent
ba5bc94963
commit
be719f53d8
2 changed files with 18 additions and 16 deletions
|
|
@ -389,7 +389,12 @@ public class McpAsyncClient {
|
|||
return withSession("by explicit API call", init -> Mono.just(init.get()));
|
||||
}
|
||||
|
||||
private Mono<McpSchema.InitializeResult> doInitialize(McpClientSession mcpClientSession) {
|
||||
private Mono<McpSchema.InitializeResult> doInitialize(Initialization initialization) {
|
||||
|
||||
initialization.setMcpClientSession(this.sessionSupplier.get());
|
||||
|
||||
McpClientSession mcpClientSession = initialization.mcpSession();
|
||||
|
||||
String latestVersion = this.protocolVersions.get(this.protocolVersions.size() - 1);
|
||||
|
||||
McpSchema.InitializeRequest initializeRequest = new McpSchema.InitializeRequest(// @formatter:off
|
||||
|
|
@ -412,6 +417,9 @@ public class McpAsyncClient {
|
|||
|
||||
return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null)
|
||||
.thenReturn(initializeResult);
|
||||
}).doOnNext(initialization::complete).onErrorResume(ex -> {
|
||||
initialization.error(ex);
|
||||
return Mono.error(ex);
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -479,15 +487,9 @@ public class McpAsyncClient {
|
|||
|
||||
boolean needsToInitialize = previous == null;
|
||||
logger.debug(needsToInitialize ? "Initialization process started" : "Joining previous initialization");
|
||||
if (needsToInitialize) {
|
||||
newInit.setMcpClientSession(this.sessionSupplier.get());
|
||||
}
|
||||
|
||||
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize
|
||||
? doInitialize(newInit.mcpSession()).doOnNext(newInit::complete).onErrorResume(ex -> {
|
||||
newInit.error(ex);
|
||||
return Mono.error(ex);
|
||||
}) : previous.await();
|
||||
Mono<McpSchema.InitializeResult> initializationJob = needsToInitialize ? doInitialize(newInit)
|
||||
: previous.await();
|
||||
|
||||
return initializationJob.map(initializeResult -> this.initializationRef.get())
|
||||
.timeout(this.initializationTimeout)
|
||||
|
|
|
|||
|
|
@ -229,27 +229,27 @@ public class McpClientSession implements McpSession {
|
|||
public <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
|
||||
String requestId = this.generateRequestId();
|
||||
|
||||
return Mono.deferContextual(ctx -> Mono.<McpSchema.JSONRPCResponse>create(sink -> {
|
||||
return Mono.deferContextual(ctx -> Mono.<McpSchema.JSONRPCResponse>create(pendingResponseSink -> {
|
||||
logger.debug("Sending message for method {}", method);
|
||||
this.pendingResponses.put(requestId, sink);
|
||||
this.pendingResponses.put(requestId, pendingResponseSink);
|
||||
McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method,
|
||||
requestId, requestParams);
|
||||
this.transport.sendMessage(jsonrpcRequest).contextWrite(ctx).subscribe(v -> {
|
||||
}, error -> {
|
||||
this.pendingResponses.remove(requestId);
|
||||
sink.error(error);
|
||||
pendingResponseSink.error(error);
|
||||
});
|
||||
})).timeout(this.requestTimeout).handle((jsonRpcResponse, sink) -> {
|
||||
})).timeout(this.requestTimeout).handle((jsonRpcResponse, deliveredResponseSink) -> {
|
||||
if (jsonRpcResponse.error() != null) {
|
||||
logger.error("Error handling request: {}", jsonRpcResponse.error());
|
||||
sink.error(new McpError(jsonRpcResponse.error()));
|
||||
deliveredResponseSink.error(new McpError(jsonRpcResponse.error()));
|
||||
}
|
||||
else {
|
||||
if (typeRef.getType().equals(Void.class)) {
|
||||
sink.complete();
|
||||
deliveredResponseSink.complete();
|
||||
}
|
||||
else {
|
||||
sink.next(this.transport.unmarshalFrom(jsonRpcResponse.result(), typeRef));
|
||||
deliveredResponseSink.next(this.transport.unmarshalFrom(jsonRpcResponse.result(), typeRef));
|
||||
}
|
||||
}
|
||||
});
|
||||
|
|
|
|||
Loading…
Reference in a new issue