diff --git a/agent-framework/prometheus_swarm/utils/errors.py b/agent-framework/prometheus_swarm/utils/errors.py index 8e2959c4..4d63344e 100644 --- a/agent-framework/prometheus_swarm/utils/errors.py +++ b/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) \ No newline at end of file diff --git a/agent-framework/tests/unit/utils/test_nonce_errors.py b/agent-framework/tests/unit/utils/test_nonce_errors.py new file mode 100644 index 00000000..fff5ea9b --- /dev/null +++ b/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 \ No newline at end of file