feat: support AsyncMcpTool

Author: wenhaozhao

core/src/main/java/com/google/adk/tools/mcp/McpAsyncTool.java

@@ -26,13 +26,16 @@
 import com.google.genai.types.FunctionDeclaration;
 import com.google.genai.types.Schema;
 import io.modelcontextprotocol.client.McpAsyncClient;
+import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
 import io.modelcontextprotocol.spec.McpSchema.JsonSchema;
 import io.modelcontextprotocol.spec.McpSchema.Tool;
 import io.reactivex.rxjava3.core.Maybe;
 import io.reactivex.rxjava3.core.Single;
 import java.util.Map;
 import java.util.Optional;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 // TODO(b/413489523): Add support for auth. This is a TODO for Python as well.
 
@@ -44,8 +47,11 @@
  */
 public final class McpAsyncTool extends BaseTool {
 
+  private static final Logger logger = LoggerFactory.getLogger(McpAsyncTool.class);
+
   Tool mcpTool;
-  Single<McpAsyncClient> mcpSession;
+  // Volatile ensures write visibility in the asynchronous chain.
+  volatile McpAsyncClient mcpSession;
   McpSessionManager mcpSessionManager;
   ObjectMapper objectMapper;
 
@@ -58,7 +64,7 @@ public final class McpAsyncTool extends BaseTool {
    * @throws IllegalArgumentException If mcpTool or mcpSession are null.
    */
   public McpAsyncTool(
-      Tool mcpTool, Single<McpAsyncClient> mcpSession, McpSessionManager mcpSessionManager) {
+      Tool mcpTool, McpAsyncClient mcpSession, McpSessionManager mcpSessionManager) {
     this(mcpTool, mcpSession, mcpSessionManager, JsonBaseModel.getMapper());
   }
 
@@ -73,7 +79,7 @@ public McpAsyncTool(
    */
   public McpAsyncTool(
       Tool mcpTool,
-      Single<McpAsyncClient> mcpSession,
+      McpAsyncClient mcpSession,
       McpSessionManager mcpSessionManager,
       ObjectMapper objectMapper) {
     super(
@@ -95,16 +101,32 @@ public McpAsyncTool(
     this.objectMapper = objectMapper;
   }
 
-  public Single<McpAsyncClient> getMcpSession() {
+  public McpAsyncClient getMcpSession() {
     return this.mcpSession;
   }
 
   public Schema toGeminiSchema(JsonSchema openApiSchema) {
     return Schema.fromJson(objectMapper.valueToTree(openApiSchema).toString());
   }
 
-  private void reintializeSession() {
-    this.mcpSession = this.mcpSessionManager.createAsyncSession();
+  private Single<McpSchema.InitializeResult> reintializeSession() {
+    McpAsyncClient client = this.mcpSessionManager.createAsyncSession();
+    return Single.fromCompletionStage(
+        client
+            .initialize()
+            .doOnSuccess(
+                initResult -> {
+                  logger.debug("Initialize McpAsyncClient Result: {}", initResult);
+                })
+            .doOnError(
+                e -> {
+                  logger.error("Initialize McpAsyncClient Failed: {}", e.getMessage(), e);
+                })
+            .doOnNext(
+                _initResult -> {
+                  this.mcpSession = client;
+                })
+            .toFuture());
   }
 
   @Override
@@ -121,14 +143,10 @@ public Optional<FunctionDeclaration> declaration() {
   public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContext toolContext) {
     return Single.defer(
             () ->
-                this.mcpSession
-                    .flatMapMaybe(
-                        client ->
-                            Maybe.fromCompletionStage(
-                                client
-                                    .callTool(
-                                        new CallToolRequest(this.name(), ImmutableMap.copyOf(args)))
-                                    .toFuture()))
+                Maybe.fromCompletionStage(
+                        this.mcpSession
+                            .callTool(new CallToolRequest(this.name(), ImmutableMap.copyOf(args)))
+                            .toFuture())
                     .map(
                         callResult ->
                             McpTool.wrapCallResult(this.objectMapper, this.name(), callResult))
@@ -141,9 +159,8 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
                     .delay(100, MILLISECONDS)
                     .take(3)
                     .doOnNext(
-                        error -> {
-                          System.err.println("Retrying callTool due to: " + error);
-                          reintializeSession();
-                        }));
+                        error ->
+                            logger.error("Retrying callTool due to: {}", error.getMessage(), error))
+                    .flatMapSingle(_ignore -> this.reintializeSession()));
   }
 }

core/src/main/java/com/google/adk/tools/mcp/McpSessionManager.java

@@ -22,7 +22,6 @@
 import io.modelcontextprotocol.spec.McpClientTransport;
 import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
 import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
-import io.reactivex.rxjava3.core.Single;
 import java.time.Duration;
 import java.util.Optional;
 import org.slf4j.Logger;
@@ -73,15 +72,15 @@ public static McpSyncClient initializeSession(
     return client;
   }
 
-  public Single<McpAsyncClient> createAsyncSession() {
+  public McpAsyncClient createAsyncSession() {
     return initializeAsyncSession(this.connectionParams);
   }
 
-  public static Single<McpAsyncClient> initializeAsyncSession(Object connectionParams) {
+  public static McpAsyncClient initializeAsyncSession(Object connectionParams) {
     return initializeAsyncSession(connectionParams, new DefaultMcpTransportBuilder());
   }
 
-  public static Single<McpAsyncClient> initializeAsyncSession(
+  public static McpAsyncClient initializeAsyncSession(
       Object connectionParams, McpTransportBuilder transportBuilder) {
     Duration initializationTimeout = null;
     Duration requestTimeout = null;
@@ -90,25 +89,11 @@ public static Single<McpAsyncClient> initializeAsyncSession(
       initializationTimeout = sseServerParams.timeout();
       requestTimeout = sseServerParams.sseReadTimeout();
     }
-    McpAsyncClient client =
-        McpClient.async(transport)
-            .initializationTimeout(
-                initializationTimeout == null ? Duration.ofSeconds(10) : initializationTimeout)
-            .requestTimeout(requestTimeout == null ? Duration.ofSeconds(10) : requestTimeout)
-            .capabilities(ClientCapabilities.builder().build())
-            .build();
-    return Single.fromCompletionStage(
-        client
-            .initialize()
-            .doOnSuccess(
-                initResult -> {
-                  logger.debug("Initialize McpAsyncClient Result: {}", initResult);
-                })
-            .doOnError(
-                e -> {
-                  logger.error("Initialize McpAsyncClient Failed: {}", e.getMessage(), e);
-                })
-            .map(_initResult -> client)
-            .toFuture());
+    return McpClient.async(transport)
+        .initializationTimeout(
+            initializationTimeout == null ? Duration.ofSeconds(10) : initializationTimeout)
+        .requestTimeout(requestTimeout == null ? Duration.ofSeconds(10) : requestTimeout)
+        .capabilities(ClientCapabilities.builder().build())
+        .build();
   }
 }