Create retry utility and mark certain errors as non-retriable (#476)

jtcorbett · web-flow · commit fe2567c4bc3a · 2025-09-15T18:43:43.000-04:00
### TL;DR

Added retry mechanism for CLI commands with proper error handling for non-retriable errors.

### What changed?

- Added a new `retriable` flag to `CLIError` class to indicate whether an error can be retried
- Set authentication and validation errors as non-retriable by adding `retriable=False` to relevant error messages
- Created a new retry utility module (`src/mcp_agent/cli/utils/retry.py`) with:
  - Functions to handle exponential backoff retries for both synchronous and asynchronous operations
  - Logic to determine if errors are retriable based on error type and the new `retriable` flag
  - A `RetryError` class to provide context about retry failures

### How to test?

1. Test CLI commands that previously failed with authentication errors to ensure they still fail appropriately
2. Verify that validation errors (like invalid API keys or formats) fail immediately without retries
3. Test network-related operations to confirm they properly retry on transient failures
4. Check that the exponential backoff works as expected with increasing delays between retry attempts

&lt;!-- This is an auto-generated comment: release notes by coderabbit.ai --&gt;
## Summary by CodeRabbit

- New Features
  - Adds automatic retry utilities with exponential backoff for CLI operations to improve reliability.
- Bug Fixes
  - Marks authentication- and parsing-related failures as non-retriable to avoid futile retries (login, whoami, workflows, secrets, and related commands).
  - Ensures unauthenticated and expired-token errors are handled consistently across commands.
- Chores
  - Standardizes error metadata to include an explicit retriable flag without changing messages or prompts.
&lt;!-- end of auto-generated comment: release notes by coderabbit.ai --&gt;
diff --git a/src/mcp_agent/cli/cloud/commands/app/status/main.py b/src/mcp_agent/cli/cloud/commands/app/status/main.py
@@ -57,7 +57,8 @@ def get_app_status(
 
     if not effective_api_key:
         raise CLIError(
-            "Must be logged in to get app status. Run 'mcp-agent login', set MCP_API_KEY environment variable or specify --api-key option."
+            "Must be logged in to get app status. Run 'mcp-agent login', set MCP_API_KEY environment variable or specify --api-key option.",
+            retriable=False
         )
 
     client = MCPAppClient(
@@ -92,7 +93,8 @@ def get_app_status(
 
     except UnauthenticatedError as e:
         raise CLIError(
-            "Invalid API key. Run 'mcp-agent login' or set MCP_API_KEY environment variable with new API key."
+            "Invalid API key. Run 'mcp-agent login' or set MCP_API_KEY environment variable with new API key.",
+            retriable=False
         ) from e
     except Exception as e:
         # Re-raise with more context - top-level CLI handler will show clean message
diff --git a/src/mcp_agent/cli/cloud/commands/auth/login/main.py b/src/mcp_agent/cli/cloud/commands/auth/login/main.py
@@ -88,7 +88,7 @@ def login(
     if api_key:
         print_info("Using provided API key for authentication (MCP_API_KEY).")
         if not _is_valid_api_key(api_key):
-            raise CLIError("Invalid API key provided.")
+            raise CLIError("Invalid API key provided.", retriable=False)
 
         credentials = _load_user_credentials(api_key)
 
@@ -139,11 +139,11 @@ def _handle_manual_key_input() -> str:
 
     if not input_api_key:
         print_warning("No API key provided.")
-        raise CLIError("Failed to set valid API key")
+        raise CLIError("Failed to set valid API key", retriable=False)
 
     if not _is_valid_api_key(input_api_key):
         print_warning("Invalid API key provided.")
-        raise CLIError("Failed to set valid API key")
+        raise CLIError("Failed to set valid API key", retriable=False)
 
     credentials = _load_user_credentials(input_api_key)
 
diff --git a/src/mcp_agent/cli/cloud/commands/auth/whoami/main.py b/src/mcp_agent/cli/cloud/commands/auth/whoami/main.py
@@ -29,13 +29,14 @@ def whoami() -> None:
         )
     if not credentials:
         raise CLIError(
-            "Not authenticated. Set MCP_API_KEY or run 'mcp-agent login'.", exit_code=4
+            "Not authenticated. Set MCP_API_KEY or run 'mcp-agent login'.", exit_code=4, retriable=False
         )
 
     if credentials.is_token_expired:
         raise CLIError(
             "Authentication token has expired. Use 'mcp-agent login' to re-authenticate.",
             exit_code=4,
+            retriable=False,
         )
 
     user_table = Table(show_header=False, box=None)
diff --git a/src/mcp_agent/cli/cloud/commands/utils.py b/src/mcp_agent/cli/cloud/commands/utils.py
@@ -29,9 +29,7 @@ def setup_authenticated_client() -> MCPAppClient:
     effective_api_key = settings.API_KEY or load_api_key_credentials()
 
     if not effective_api_key:
-        raise CLIError(
-            "Must be authenticated. Set MCP_API_KEY or run 'mcp-agent login'."
-        )
+        raise CLIError("Must be authenticated. Set MCP_API_KEY or run 'mcp-agent login'.", retriable=False)
 
     return MCPAppClient(api_url=DEFAULT_API_BASE_URL, api_key=effective_api_key)
 
@@ -48,7 +46,8 @@ def validate_output_format(format: str) -> None:
     valid_formats = ["text", "json", "yaml"]
     if format not in valid_formats:
         raise CLIError(
-            f"Invalid format '{format}'. Valid options are: {', '.join(valid_formats)}"
+            f"Invalid format '{format}'. Valid options are: {', '.join(valid_formats)}",
+            retriable=False
         )
 
 
@@ -100,7 +99,8 @@ def wrapper(*args, **kwargs):
             return func(*args, **kwargs)
         except UnauthenticatedError as e:
             raise CLIError(
-                "Invalid API key. Run 'mcp-agent login' or set MCP_API_KEY environment variable with new API key."
+                "Invalid API key. Run 'mcp-agent login' or set MCP_API_KEY environment variable with new API key.",
+                retriable=False
             ) from e
         except CLIError:
             # Re-raise CLIErrors as-is
diff --git a/src/mcp_agent/cli/cloud/commands/workflows/cancel/main.py b/src/mcp_agent/cli/cloud/commands/workflows/cancel/main.py
@@ -41,7 +41,7 @@ async def _cancel_workflow_async(
     effective_api_key = _settings.API_KEY or load_api_key_credentials()
 
     if not effective_api_key:
-        raise CLIError("Must be logged in to access server. Run 'mcp-agent login'.")
+        raise CLIError("Must be logged in to access server. Run 'mcp-agent login'.", retriable=False)
 
     try:
         async with mcp_connection_session(
diff --git a/src/mcp_agent/cli/cloud/commands/workflows/describe/main.py b/src/mcp_agent/cli/cloud/commands/workflows/describe/main.py
@@ -46,7 +46,7 @@ async def _describe_workflow_async(
     effective_api_key = _settings.API_KEY or load_api_key_credentials()
 
     if not effective_api_key:
-        raise CLIError("Must be logged in to access server. Run 'mcp-agent login'.")
+        raise CLIError("Must be logged in to access server. Run 'mcp-agent login'.", retriable=False)
 
     try:
         async with mcp_connection_session(
diff --git a/src/mcp_agent/cli/cloud/commands/workflows/list/main.py b/src/mcp_agent/cli/cloud/commands/workflows/list/main.py
@@ -43,7 +43,7 @@ async def _list_workflows_async(server_id_or_url: str, format: str = "text") ->
     effective_api_key = _settings.API_KEY or load_api_key_credentials()
 
     if not effective_api_key:
-        raise CLIError("Must be logged in to access server. Run 'mcp-agent login'.")
+        raise CLIError("Must be logged in to access server. Run 'mcp-agent login'.", retriable=False)
 
     try:
         async with mcp_connection_session(
diff --git a/src/mcp_agent/cli/cloud/commands/workflows/resume/main.py b/src/mcp_agent/cli/cloud/commands/workflows/resume/main.py
@@ -45,7 +45,7 @@ async def _signal_workflow_async(
     effective_api_key = _settings.API_KEY or load_api_key_credentials()
 
     if not effective_api_key:
-        raise CLIError("Must be logged in to access server. Run 'mcp-agent login'.")
+        raise CLIError("Must be logged in to access server. Run 'mcp-agent login'.", retriable=False)
 
     try:
         async with mcp_connection_session(
diff --git a/src/mcp_agent/cli/cloud/commands/workflows/runs/main.py b/src/mcp_agent/cli/cloud/commands/workflows/runs/main.py
@@ -47,7 +47,7 @@ async def _list_workflow_runs_async(
     effective_api_key = _settings.API_KEY or load_api_key_credentials()
 
     if not effective_api_key:
-        raise CLIError("Must be logged in to access server. Run 'mcp-agent login'.")
+        raise CLIError("Must be logged in to access server. Run 'mcp-agent login'.", retriable=False)
 
     try:
         async with mcp_connection_session(
diff --git a/src/mcp_agent/cli/exceptions.py b/src/mcp_agent/cli/exceptions.py
@@ -4,6 +4,7 @@
 class CLIError(Exception):
     """Exception for expected CLI errors that should show clean user-facing messages."""
 
-    def __init__(self, message: str, exit_code: int = 1):
+    def __init__(self, message: str, exit_code: int = 1, retriable: bool = True):
         super().__init__(message)
         self.exit_code = exit_code
+        self.retriable = retriable
diff --git a/src/mcp_agent/cli/secrets/processor.py b/src/mcp_agent/cli/secrets/processor.py
@@ -93,7 +93,8 @@ async def process_config_secrets(
 
         if not effective_api_key:
             raise CLIError(
-                "Must have API key to process secrets. Login via 'mcp-agent login'."
+                "Must have API key to process secrets. Login via 'mcp-agent login'.",
+                retriable=False
             )
 
         # Create a new client
@@ -191,7 +192,7 @@ async def process_secrets_in_config_str(
     try:
         input_config = yaml.safe_load(input_secrets_content)
     except Exception as e:
-        raise CLIError(f"Failed to parse input YAML: {str(e)}") from e
+        raise CLIError(f"Failed to parse input YAML: {str(e)}", retriable=False) from e
 
     # Parse the existing secrets YAML if provided
     existing_config = None
@@ -200,7 +201,7 @@ async def process_secrets_in_config_str(
             existing_config = load_yaml_with_secrets(existing_secrets_content)
             print_info("Loaded existing secrets configuration for reuse")
         except Exception as e:
-            raise CLIError(f"Failed to parse existing secrets YAML: {str(e)}") from e
+            raise CLIError(f"Failed to parse existing secrets YAML: {str(e)}", retriable=False) from e
 
     # Make sure the existing config secrets are actually valid for the user
     if existing_config:
diff --git a/src/mcp_agent/cli/utils/retry.py b/src/mcp_agent/cli/utils/retry.py
@@ -0,0 +1,152 @@
+"""Retry utilities for CLI operations."""
+
+import asyncio
+import time
+from typing import Any, Callable, Optional
+
+from mcp_agent.cli.core.api_client import UnauthenticatedError
+from mcp_agent.cli.exceptions import CLIError
+from mcp_agent.cli.utils.ux import print_warning
+
+
+class RetryError(Exception):
+    """Exception raised when all retry attempts are exhausted."""
+
+    def __init__(self, original_error: Exception, attempts: int):
+        self.original_error = original_error
+        self.attempts = attempts
+        super().__init__(f"Failed after {attempts} attempts. Last error: {original_error}")
+
+
+def is_retryable_error(error: Exception) -> bool:
+    """Determine if an error should trigger a retry.
+
+    Args:
+        error: The exception to evaluate
+
+    Returns:
+        True if the error is retryable, False otherwise
+    """
+    if isinstance(error, UnauthenticatedError):
+        return False
+
+    if isinstance(error, CLIError):
+        return error.retriable
+
+    return True
+
+
+def retry_with_exponential_backoff(
+    func: Callable,
+    max_attempts: int = 3,
+    initial_delay: float = 1.0,
+    backoff_multiplier: float = 2.0,
+    max_delay: float = 60.0,
+    retryable_check: Optional[Callable[[Exception], bool]] = None,
+    *args,
+    **kwargs
+) -> Any:
+    """Retry a function with exponential backoff.
+
+    Args:
+        func: The function to retry
+        max_attempts: Maximum number of attempts (including the first one)
+        initial_delay: Initial delay in seconds before first retry
+        backoff_multiplier: Multiplier for delay between attempts
+        max_delay: Maximum delay between attempts
+        retryable_check: Function to determine if an error is retryable
+        *args: Arguments to pass to func
+        **kwargs: Keyword arguments to pass to func
+
+    Returns:
+        Result of the successful function call
+
+    Raises:
+        RetryError: If all attempts fail with a retryable error
+        Exception: The original exception if it's not retryable
+    """
+    if retryable_check is None:
+        retryable_check = is_retryable_error
+
+    last_exception = None
+    delay = initial_delay
+
+    for attempt in range(1, max_attempts + 1):
+        try:
+            return func(*args, **kwargs)
+        except Exception as e:
+            last_exception = e
+
+            if attempt == max_attempts or not retryable_check(e):
+                break
+
+            print_warning(f"Attempt {attempt}/{max_attempts} failed: {e}. Retrying in {delay:.1f}s...")
+
+            time.sleep(delay)
+            delay = min(delay * backoff_multiplier, max_delay)
+
+    if last_exception:
+        if max_attempts > 1 and retryable_check(last_exception):
+            raise RetryError(last_exception, max_attempts) from last_exception
+        else:
+            raise last_exception
+
+    raise RuntimeError("Unexpected error in retry logic")
+
+
+async def retry_async_with_exponential_backoff(
+    func: Callable,
+    max_attempts: int = 3,
+    initial_delay: float = 1.0,
+    backoff_multiplier: float = 2.0,
+    max_delay: float = 60.0,
+    retryable_check: Optional[Callable[[Exception], bool]] = None,
+    *args,
+    **kwargs
+) -> Any:
+    """Async version of retry with exponential backoff.
+
+    Args:
+        func: Async function to retry
+        max_attempts: Maximum number of attempts (including the first one)
+        initial_delay: Initial delay in seconds before first retry
+        backoff_multiplier: Multiplier for delay between attempts
+        max_delay: Maximum delay between attempts
+        retryable_check: Function to determine if an error is retryable
+        *args: Arguments to pass to func
+        **kwargs: Keyword arguments to pass to func
+
+    Returns:
+        Result of the successful function call
+
+    Raises:
+        RetryError: If all attempts fail with a retryable error
+        Exception: The original exception if it's not retryable
+    """
+    if retryable_check is None:
+        retryable_check = is_retryable_error
+
+    last_exception = None
+    delay = initial_delay
+
+    for attempt in range(1, max_attempts + 1):
+        try:
+            return await func(*args, **kwargs)
+        except Exception as e:
+            last_exception = e
+
+            if attempt == max_attempts or not retryable_check(e):
+                break
+
+            print_warning(f"Attempt {attempt}/{max_attempts} failed: {e}. Retrying in {delay:.1f}s...")
+
+            await asyncio.sleep(delay)
+            delay = min(delay * backoff_multiplier, max_delay)
+
+    if last_exception:
+        if max_attempts > 1 and retryable_check(last_exception):
+            raise RetryError(last_exception, max_attempts) from last_exception
+        else:
+            raise last_exception
+
+    raise RuntimeError("Unexpected error in retry logic")

Original file line number	Diff line number	Diff line change
`@@ -29,13 +29,14 @@ def whoami() -> None:`
`29`	`29`	`)`
`30`	`30`	`if not credentials:`
`31`	`31`	`raise CLIError(`
`32`		`- "Not authenticated. Set MCP_API_KEY or run 'mcp-agent login'.", exit_code=4`
	`32`	`+ "Not authenticated. Set MCP_API_KEY or run 'mcp-agent login'.", exit_code=4, retriable=False`
`33`	`33`	`)`
`34`	`34`
`35`	`35`	`if credentials.is_token_expired:`
`36`	`36`	`raise CLIError(`
`37`	`37`	`"Authentication token has expired. Use 'mcp-agent login' to re-authenticate.",`
`38`	`38`	`exit_code=4,`
	`39`	`+ retriable=False,`
`39`	`40`	`)`
`40`	`41`
`41`	`42`	`user_table = Table(show_header=False, box=None)`