Create retry utility and mark certain errors as non-retriable

jtcorbett · jtcorbett · commit 5b8c79025fd0 · 2025-09-15T22:00:09.000Z
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)`