fix: refresh stale credentials on auth failure without restart (#245)

* fix: refresh stale credentials on auth failure without restart

When AWS credentials expire, the proxy now automatically picks up
refreshed credentials on the next request without requiring a restart.

Problem:
- boto3.Session was created once at startup and cached forever
- session.get_credentials() returned the same stale Credentials object
- Even after refreshing creds on disk (ada, aws sso login), the proxy
  kept using the old frozen session until restarted

Fix:
- Add SessionHolder that wraps the boto3 session with lazy refresh
- On 401/403, mark the session for refresh (don't refresh immediately,
  since creds on disk may not be updated yet)
- On the next request's signing, create a fresh boto3.Session that
  reads the current credentials from disk
- Improve error messages: credential errors now clearly say 'expired
  or invalid AWS credentials' instead of 'Unknown tool'

The lazy refresh pattern ensures the new session is created at signing
time (after the user has refreshed creds), not at error time (when
creds may still be stale).

* fix: address review comments on credential refresh

* chore: fix ruff lint and format issues

* test: add unit tests for credential refresh flow

* fix: address PR review comments

---------

Co-authored-by: Ian de Villiers <iddv@amazon.com>

Author: anasstahr

README.md

@@ -353,8 +353,10 @@ For long-running sessions, consider using long-lived credentials:
 - Use an AWS profile via `--profile`
 - Use IAM Identity Center and run `aws sso login` before starting the proxy
 
+If your credentials do expire during a session, the proxy will automatically detect the auth failure and pick up refreshed credentials on the next request — no restart required. Simply refresh your credentials (e.g., `aws sso login`) and retry.
+
 ### Client hangs on tool calls
-If your MCP client hangs waiting for a tool call response (e.g., due to expired credentials or an unresponsive endpoint), use `--tool-timeout` to set a maximum duration in seconds for each tool call. When the timeout is exceeded, the proxy returns a graceful error to the agent instead of hanging indefinitely.
+If your MCP client hangs waiting for a tool call response (e.g., due to an unresponsive endpoint), use `--tool-timeout` to set a maximum duration in seconds for each tool call. When the timeout is exceeded, the proxy returns a graceful error to the agent instead of hanging indefinitely.
 
 ## Development & Contributing
 

mcp_proxy_for_aws/middleware/tool_error_middleware.py

@@ -62,17 +62,27 @@ async def on_call_tool(
             logger.error('Tool call %r failed: %s.', tool_name, e)
             message = f'Tool call {tool_name!r} failed: {e}. Please retry.'
             if self._is_credential_error(e):
-                message += (
-                    ' This may be caused by expired or invalid AWS credentials.'
-                    ' Consider using long-lived credentials such as an AWS profile'
-                    ' (--profile) or IAM Identity Center (aws sso login).'
+                message = (
+                    f'Tool call {tool_name!r} failed due to expired or invalid AWS credentials.'
+                    ' Please refresh your credentials and retry.'
+                    ' The proxy will automatically use the new credentials on the next request.'
                 )
             raise ToolError(message) from e
 
     @staticmethod
     def _is_credential_error(error: Exception) -> bool:
-        """Check if the error is likely caused by expired or invalid credentials."""
-        return isinstance(error, httpx.HTTPStatusError) and error.response.status_code in (
-            401,
-            403,
-        )
+        """Check if the error is likely caused by expired or invalid credentials.
+
+        Walks the exception chain (__cause__/__context__) because the
+        HTTPStatusError may be wrapped.  isinstance already respects the
+        MRO, so subclasses are caught too.
+        """
+        current: BaseException | None = error
+        while current is not None:
+            if isinstance(current, httpx.HTTPStatusError) and current.response.status_code in (
+                401,
+                403,
+            ):
+                return True
+            current = current.__cause__ or current.__context__
+        return False

mcp_proxy_for_aws/sigv4_helper.py

@@ -120,12 +120,42 @@ def create_aws_session(profile: Optional[str] = None) -> boto3.Session:
     return session
 
 
+class SessionHolder:
+    """Holds a boto3 session that can be refreshed on credential errors.
+
+    Wraps a boto3.Session so the signing hook always uses the current session,
+    and can create a fresh session when the previous request got an auth error.
+    """
+
+    def __init__(self, session: boto3.Session, profile: Optional[str] = None) -> None:
+        """Initialize SessionHolder with the given session and optional profile."""
+        self.session = session
+        self._profile = profile
+        self._needs_refresh = False
+
+    def mark_needs_refresh(self) -> None:
+        """Mark that the next request should use a fresh session."""
+        self._needs_refresh = True
+
+    def refresh_if_needed(self) -> None:
+        """Create a fresh session if a previous request got an auth error."""
+        if not self._needs_refresh:
+            return
+        logger.info('Refreshing AWS session to pick up new credentials')
+        try:
+            self.session = create_aws_session(self._profile)
+            self._needs_refresh = False
+        except Exception:
+            logger.warning(
+                'Failed to create fresh AWS session, keeping current session', exc_info=True
+            )
+
+
 def create_sigv4_client(
     service: str,
     region: str,
+    session_holder: SessionHolder,
     timeout: Optional[httpx.Timeout] = None,
-    profile: Optional[str] = None,
-    session: Optional[boto3.Session] = None,
     headers: Optional[Dict[str, str]] = None,
     metadata: Optional[Dict[str, Any]] = None,
     disable_telemetry: bool = False,
@@ -135,9 +165,9 @@ def create_sigv4_client(
 
     Args:
         service: AWS service name for SigV4 signing
-        profile: AWS profile to use (optional, only used if session is not provided)
-        session: AWS boto3 session to use (optional, takes precedence over profile)
-        region: AWS region (optional, defaults to AWS_REGION env var or us-east-1)
+        region: AWS region for SigV4 signing
+        session_holder: SessionHolder that provides the current boto3 session
+            and can refresh it on credential errors
         timeout: Timeout configuration for the HTTP client
         headers: Headers to include in requests
         metadata: Metadata to inject into MCP _meta field
@@ -147,10 +177,6 @@ def create_sigv4_client(
     Returns:
         httpx.AsyncClient with SigV4 authentication
     """
-    # Create or use provided AWS session
-    if session is None:
-        session = create_aws_session(profile)
-
     # Create a copy of kwargs to avoid modifying the passed dict
     client_kwargs = {
         'follow_redirects': True,
@@ -184,28 +210,33 @@ def create_sigv4_client(
     return httpx.AsyncClient(
         **client_kwargs,
         event_hooks={
-            'response': [_handle_error_response],
+            'response': [partial(_handle_error_response, session_holder)],
             'request': [
                 partial(_inject_metadata_hook, metadata or {}),
-                partial(_sign_request_hook, region, service, session),
+                partial(_sign_request_hook, region, service, session_holder),
             ],
         },
     )
 
 
-async def _handle_error_response(response: httpx.Response) -> None:
+async def _handle_error_response(session_holder: SessionHolder, response: httpx.Response) -> None:
     """Event hook to handle HTTP error responses and extract details.
 
     This function is called for every HTTP response to check for errors
     and provide more detailed error information when requests fail.
 
     Args:
+        session_holder: SessionHolder to refresh on credential errors
         response: The HTTP response object
 
     Raises:
         No raises. let the mcp http client handle the errors.
     """
     if response.is_error:
+        # Mark session for refresh so the next request picks up new credentials
+        if response.status_code in (401, 403):
+            session_holder.mark_needs_refresh()
+
         # warning only because the SDK logs error
         log_level = logging.WARNING
         if (
@@ -246,7 +277,7 @@ async def _handle_error_response(response: httpx.Response) -> None:
 async def _sign_request_hook(
     region: str,
     service: str,
-    session: boto3.Session,
+    session_holder: SessionHolder,
     request: httpx.Request,
 ) -> None:
     """Request hook to sign HTTP requests with AWS SigV4.
@@ -258,14 +289,19 @@ async def _sign_request_hook(
     Args:
         region: AWS region for SigV4 signing
         service: AWS service name for SigV4 signing
-        session: AWS boto3 session to use for credentials
+        session_holder: Holder providing the current boto3 session (refreshed on auth errors)
         request: The HTTP request object to sign (modified in-place)
     """
     # Set Content-Length for signing
     request.headers['Content-Length'] = str(len(request.content))
 
+    # Refresh session if a previous request got an auth error.
+    # Done here (at signing time) so the new session reads credentials
+    # that the user may have refreshed since the error occurred.
+    session_holder.refresh_if_needed()
+
     # Get AWS credentials from the session
-    credentials = session.get_credentials()
+    credentials = session_holder.session.get_credentials()
 
     # Create SigV4 auth and use its signing logic
     auth = SigV4HTTPXAuth(credentials, service, region)

mcp_proxy_for_aws/utils.py

@@ -19,7 +19,7 @@
 import logging
 import os
 from fastmcp.client.transports import StreamableHttpTransport
-from mcp_proxy_for_aws.sigv4_helper import create_aws_session, create_sigv4_client
+from mcp_proxy_for_aws.sigv4_helper import SessionHolder, create_aws_session, create_sigv4_client
 from typing import Any, Dict, Optional, Tuple
 from urllib.parse import urlparse
 
@@ -90,9 +90,10 @@ def create_transport_with_sigv4(
     Returns:
         StreamableHttpTransport instance with SigV4 authentication
     """
-    # Create AWS session once and reuse it for all httpx clients
+    # Create AWS session with a holder that can refresh on credential errors
     logger.debug('Creating AWS session with profile: %s', profile)
     session = create_aws_session(profile)
+    session_holder = SessionHolder(session, profile)
 
     def client_factory(
         headers: Optional[Dict[str, str]] = None,
@@ -102,7 +103,7 @@ def client_factory(
     ) -> httpx.AsyncClient:
         return create_sigv4_client(
             service=service,
-            session=session,
+            session_holder=session_holder,
             region=region,
             headers=headers,
             timeout=custom_timeout,