refactor: Modify the Mcp retry logic to be completely async

PiperOrigin-RevId: 805881432

Author: google-genai-bot

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

@@ -16,7 +16,7 @@
 
 package com.google.adk.tools.mcp;
 
-import static com.google.common.collect.ImmutableList.toImmutableList;
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.adk.JsonBaseModel;
@@ -198,73 +198,72 @@ public McpToolset(StreamableHttpServerParameters connectionParams) {
 
   @Override
   public Flowable<BaseTool> getTools(ReadonlyContext readonlyContext) {
-    return Flowable.fromCallable(
+    return Flowable.defer(
             () -> {
-              for (int i = 0; i < MAX_RETRIES; i++) {
-                try {
-                  if (this.mcpSession == null) {
-                    logger.info("MCP session is null or closed, initializing (attempt {}).", i + 1);
-                    this.mcpSession = this.mcpSessionManager.createSession();
-                  }
-
-                  ListToolsResult toolsResponse = this.mcpSession.listTools();
-                  return toolsResponse.tools().stream()
+              if (this.mcpSession == null) {
+                logger.info("MCP session is null, initializing.");
+                this.mcpSession = this.mcpSessionManager.createSession();
+              }
+
+              // Retrieve tools from the MCP session, wrap them in McpTool, filter them, and return
+              // as a Flowable.
+              ListToolsResult toolsResponse = this.mcpSession.listTools();
+              return Flowable.fromStream(
+                  toolsResponse.tools().stream()
                       .map(
                           tool ->
                               new McpTool(
                                   tool, this.mcpSession, this.mcpSessionManager, this.objectMapper))
                       .filter(
                           tool ->
                               isToolSelected(
-                                  tool, toolFilter, Optional.ofNullable(readonlyContext)))
-                      .collect(toImmutableList());
-                } catch (IllegalArgumentException e) {
-                  // This could happen if parameters for tool loading are somehow invalid.
-                  // This is likely a fatal error and should not be retried.
-                  logger.error("Invalid argument encountered during tool loading.", e);
-                  throw new McpToolsetException.McpToolLoadingException(
-                      "Invalid argument encountered during tool loading.", e);
-                } catch (RuntimeException e) { // Catch any other unexpected runtime exceptions
-                  logger.error("Unexpected error during tool loading, retry attempt " + (i + 1), e);
-                  if (i < MAX_RETRIES - 1) {
-                    // For other general exceptions, we might still want to retry if they are
-                    // potentially transient, or if we don't have more specific handling. But it's
-                    // better to be specific. For now, we'll treat them as potentially retryable but
-                    // log
-                    // them at a higher level.
-                    try {
-                      logger.info(
-                          "Reinitializing MCP session before next retry for unexpected error.");
-                      this.mcpSession = this.mcpSessionManager.createSession();
-                      Thread.sleep(RETRY_DELAY_MILLIS);
-                    } catch (InterruptedException ie) {
-                      Thread.currentThread().interrupt();
-                      logger.error(
-                          "Interrupted during retry delay for loadTools (unexpected error).", ie);
-                      throw new McpToolsetException.McpToolLoadingException(
-                          "Interrupted during retry delay (unexpected error)", ie);
-                    } catch (RuntimeException reinitE) {
-                      logger.error(
-                          "Failed to reinitialize session during retry (unexpected error).",
-                          reinitE);
-                      throw new McpToolsetException.McpInitializationException(
-                          "Failed to reinitialize session during tool loading retry (unexpected"
-                              + " error).",
-                          reinitE);
-                    }
-                  } else {
-                    logger.error(
-                        "Failed to load tools after multiple retries due to unexpected error.", e);
-                    throw new McpToolsetException.McpToolLoadingException(
-                        "Failed to load tools after multiple retries due to unexpected error.", e);
-                  }
-                }
-              }
-              // This line should ideally not be reached if retries are handled correctly or an
-              // exception is always thrown.
-              throw new IllegalStateException("Unexpected state in getTools retry loop");
+                                  tool, toolFilter, Optional.ofNullable(readonlyContext))));
             })
-        .flatMapIterable(tools -> tools);
+        .retryWhen(
+            errorObservable ->
+                errorObservable.zipWith(
+                    Flowable.range(1, MAX_RETRIES),
+                    (error, retryCount) -> {
+                      if (error instanceof IllegalArgumentException) {
+                        // This could happen if parameters for tool loading are somehow invalid.
+                        // This is likely a fatal error and should not be retried.
+                        logger.error("Invalid argument encountered during tool loading.", error);
+                        throw new McpToolsetException.McpToolLoadingException(
+                            "Invalid argument encountered during tool loading.", error);
+                      } else if (error instanceof RuntimeException) {
+                        // Catch any other unexpected runtime exceptions
+                        logger.error(
+                            "Unexpected error during tool loading, retry attempt " + retryCount,
+                            error);
+                        logger.info(
+                            "Reinitializing MCP session before next retry for unexpected error.");
+                        this.mcpSession = null;
+
+                        if (retryCount < MAX_RETRIES) {
+                          // For other general exceptions, we might still want to retry if they are
+                          // potentially transient, or if we don't have more specific handling. But
+                          // it's better to be specific. For now, we'll treat them as potentially
+                          // retryable but log them at a higher level.
+
+                          // Delay before retrying
+                          return Flowable.timer(RETRY_DELAY_MILLIS, MILLISECONDS);
+                        } else {
+                          logger.error(
+                              "Failed to load tools after multiple retries due to unexpected"
+                                  + " error.",
+                              error);
+                          throw new McpToolsetException.McpToolLoadingException(
+                              "Failed to load tools after multiple retries due to unexpected"
+                                  + " error.",
+                              error);
+                        }
+                      }
+                      // This line should ideally not be reached if retries are handled correctly or
+                      // an exception is always thrown.
+                      // If an unhandled error type occurs, propagate it.
+                      return Flowable.error(error);
+                    }))
+        .map(tools -> tools);
   }
 
   @Override

core/src/test/java/com/google/adk/tools/mcp/McpToolsetTest.java

@@ -17,19 +17,38 @@
 package com.google.adk.tools.mcp;
 
 import static com.google.common.truth.Truth.assertThat;
+import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.junit.Assert.assertThrows;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+import com.google.adk.JsonBaseModel;
 import com.google.adk.agents.ConfigAgentUtils.ConfigurationException;
+import com.google.adk.agents.ReadonlyContext;
 import com.google.adk.tools.BaseTool;
 import com.google.adk.tools.mcp.McpToolset.McpToolsetConfig;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import io.modelcontextprotocol.client.McpSyncClient;
+import io.modelcontextprotocol.spec.McpSchema;
+import java.util.List;
+import java.util.Optional;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
 
 @RunWith(JUnit4.class)
 public class McpToolsetTest {
+  @Rule public final MockitoRule mocks = MockitoJUnit.rule();
+
+  @Mock private McpSessionManager mockMcpSessionManager;
+  @Mock private McpSyncClient mockMcpSyncClient;
+  @Mock private ReadonlyContext mockReadonlyContext;
 
   @Test
   public void testMcpToolsetConfig_withStdioServerParams_parsesCorrectly() {
@@ -135,7 +154,6 @@ public void testFromConfig_validStdioParams_createsToolset() throws Configuratio
     McpToolset toolset = McpToolset.fromConfig(config, configPath);
 
     assertThat(toolset).isNotNull();
-    // The toolset should be created successfully with stdio parameters
   }
 
   @Test
@@ -243,4 +261,68 @@ public void testFromConfig_emptyToolFilter_createsToolset() throws Configuration
     assertThat(toolset).isNotNull();
     // The toolset should be created successfully with empty tool filter
   }
+
+  @Test
+  public void getTools_withToolFilter_returnsFilteredTools() {
+    ImmutableList<String> toolFilter = ImmutableList.of("tool1", "tool3");
+    McpSchema.Tool mockTool1 =
+        McpSchema.Tool.builder().name("tool1").description("desc1").inputSchema("{}").build();
+    McpSchema.Tool mockTool2 =
+        McpSchema.Tool.builder().name("tool2").description("desc2").inputSchema("{}").build();
+    McpSchema.Tool mockTool3 =
+        McpSchema.Tool.builder().name("tool3").description("desc3").inputSchema("{}").build();
+    McpSchema.ListToolsResult mockResult =
+        new McpSchema.ListToolsResult(ImmutableList.of(mockTool1, mockTool2, mockTool3), null);
+
+    when(mockMcpSessionManager.createSession()).thenReturn(mockMcpSyncClient);
+    when(mockMcpSyncClient.listTools()).thenReturn(mockResult);
+
+    McpToolset toolset =
+        new McpToolset(mockMcpSessionManager, JsonBaseModel.getMapper(), Optional.of(toolFilter));
+
+    List<BaseTool> tools = toolset.getTools(mockReadonlyContext).toList().blockingGet();
+
+    assertThat(tools.stream().map(BaseTool::name).collect(ImmutableList.toImmutableList()))
+        .containsExactly("tool1", "tool3")
+        .inOrder();
+    verify(mockMcpSessionManager).createSession();
+    verify(mockMcpSyncClient).listTools();
+  }
+
+  @Test
+  public void getTools_retriesAndFailsAfterMaxRetries() {
+    when(mockMcpSessionManager.createSession()).thenReturn(mockMcpSyncClient);
+    when(mockMcpSyncClient.listTools()).thenThrow(new RuntimeException("Test Exception"));
+
+    McpToolset toolset =
+        new McpToolset(mockMcpSessionManager, JsonBaseModel.getMapper(), Optional.empty());
+
+    toolset
+        .getTools(mockReadonlyContext)
+        .test()
+        .awaitDone(5, SECONDS)
+        .assertError(McpToolsetException.McpToolLoadingException.class);
+
+    verify(mockMcpSessionManager, times(3)).createSession();
+    verify(mockMcpSyncClient, times(3)).listTools();
+  }
+
+  @Test
+  public void getTools_succeedsOnLastRetryAttempt() {
+    McpSchema.ListToolsResult mockResult = new McpSchema.ListToolsResult(ImmutableList.of(), null);
+    when(mockMcpSessionManager.createSession()).thenReturn(mockMcpSyncClient);
+    when(mockMcpSyncClient.listTools())
+        .thenThrow(new RuntimeException("Attempt 1 failed"))
+        .thenThrow(new RuntimeException("Attempt 2 failed"))
+        .thenReturn(mockResult);
+
+    McpToolset toolset =
+        new McpToolset(mockMcpSessionManager, JsonBaseModel.getMapper(), Optional.empty());
+
+    List<BaseTool> tools = toolset.getTools(mockReadonlyContext).toList().blockingGet();
+
+    assertThat(tools).isEmpty();
+    verify(mockMcpSessionManager, times(3)).createSession();
+    verify(mockMcpSyncClient, times(3)).listTools();
+  }
 }