Implement Robust Nonce Error Handling and Recovery Mechanism

agent-framework/prometheus_swarm/utils/errors.py

@@ -1,4 +1,9 @@
-"""Error type for API errors."""
+"""Error types for API and nonce-related errors."""
+
+import logging
+import time
+from functools import wraps
+from typing import Optional, Callable, Any, Union
 
 
 class ClientAPIError(Exception):
@@ -14,3 +19,69 @@ def __init__(self, original_error: Exception):
             super().__init__(original_error.message)
         else:
             super().__init__(str(original_error))
+
+
+class NonceError(Exception):
+    """Raised when a nonce (number used once) is invalid or has been reused."""
+
+    def __init__(self, message: str, current_nonce: Optional[str] = None):
+        """
+        Initialize NonceError with a descriptive message and optional current nonce.
+
+        Args:
+            message (str): Description of the nonce error
+            current_nonce (Optional[str]): The problematic nonce value
+        """
+        self.current_nonce = current_nonce
+        super().__init__(message)
+
+
+def nonce_error_handler(
+    func: Optional[Callable] = None,
+    *,
+    max_retries: int = 3,
+    retry_delay: float = 1.0,
+    logger: Optional[logging.Logger] = None
+) -> Union[Callable, Any]:
+    """
+    Decorator to handle nonce-related errors with automatic recovery and retries.
+
+    Args:
+        func (Optional[Callable]): The function to wrap
+        max_retries (int): Maximum number of retry attempts
+        retry_delay (float): Delay between retry attempts in seconds
+        logger (Optional[logging.Logger]): Logger for tracking retry attempts
+
+    Returns:
+        Wrapped function with nonce error handling
+    """
+    def decorator(func_to_wrap: Callable) -> Callable:
+        @wraps(func_to_wrap)
+        def wrapper(*args, **kwargs):
+            retries = 0
+            last_error = None
+            while retries < max_retries:
+                try:
+                    return func_to_wrap(*args, **kwargs)
+                except NonceError as e:
+                    retries += 1
+                    last_error = e
+                    if logger:
+                        logger.warning(
+                            f"Nonce error encountered: {e}. "
+                            f"Retry attempt {retries}/{max_retries}"
+                        )
+
+                    if retries >= max_retries:
+                        break
+
+                    time.sleep(retry_delay * (2 ** retries))  # Exponential backoff
+
+            raise RuntimeError("Max nonce error retries exceeded") from last_error
+
+        return wrapper
+
+    # Support using decorator with or without arguments
+    if func is None:
+        return decorator
+    return decorator(func)

agent-framework/tests/unit/utils/test_nonce_errors.py

@@ -0,0 +1,59 @@
+"""Tests for nonce error handling and recovery."""
+
+import logging
+import pytest
+from prometheus_swarm.utils.errors import NonceError, nonce_error_handler
+
+
+def test_nonce_error_initialization():
+    """Test NonceError can be created with message and optional nonce."""
+    error = NonceError("Invalid nonce", "abc123")
+    assert str(error) == "Invalid nonce"
+    assert error.current_nonce == "abc123"
+
+
+def test_nonce_error_without_nonce():
+    """Test NonceError can be created without specifying a nonce."""
+    error = NonceError("Nonce reused")
+    assert str(error) == "Nonce reused"
+    assert error.current_nonce is None
+
+
+def test_nonce_error_handler_success():
+    """Test nonce error handler successfully recovers from nonce errors."""
+    mock_attempts = [0]
+
+    @nonce_error_handler
+    def mock_function():
+        mock_attempts[0] += 1
+        if mock_attempts[0] < 3:
+            raise NonceError("Simulated nonce error")
+        return "Success"
+
+    result = mock_function()
+    assert result == "Success"
+    assert mock_attempts[0] == 3
+
+
+def test_nonce_error_handler_max_retries():
+    """Test nonce error handler raises error after max retries."""
+    @nonce_error_handler(max_retries=2)
+    def mock_function():
+        raise NonceError("Persistent nonce error")
+
+    with pytest.raises(RuntimeError, match="Max nonce error retries exceeded"):
+        mock_function()
+
+
+def test_nonce_error_handler_with_logger(caplog):
+    """Test nonce error handler works with a logger."""
+    caplog.set_level(logging.WARNING)
+
+    @nonce_error_handler(logger=logging.getLogger())
+    def mock_function():
+        raise NonceError("Logged nonce error")
+
+    with pytest.raises(RuntimeError, match="Max nonce error retries exceeded"):
+        mock_function()
+
+    assert "Nonce error encountered: Logged nonce error" in caplog.text