Merge pull request #166 from scouturier/feat/cache-oidc-id-token-silent-refresh

scouturier · web-flow · commit fc576f13d7c9 · 2026-03-30T11:28:22.000-04:00
feat: cache OIDC id_token for silent credential refresh
diff --git a/source/credential_provider/__main__.py b/source/credential_provider/__main__.py
@@ -584,8 +584,8 @@ def get_monitoring_token(self):
             exp_time = token_data.get("expires", 0)
             now = int(datetime.now(timezone.utc).timestamp())
 
-            # Return token if it expires in more than 10 minutes
-            if exp_time - now > 600:
+            # Return token if it expires in more than 60 seconds
+            if exp_time - now > 60:
                 token = token_data["token"]
                 # Set in environment for this session
                 os.environ["CLAUDE_CODE_MONITORING_TOKEN"] = token
@@ -1759,6 +1759,30 @@ def _handle_quota_warning(self, quota_result: dict):
     # End Quota Check Methods
     # ===========================================
 
+    def _try_silent_refresh(self):
+        """Attempt to refresh AWS credentials using a cached, still-valid OIDC id_token.
+
+        Returns:
+            Tuple of (credentials, id_token, token_claims) if successful, (None, None, None) otherwise.
+        """
+        try:
+            id_token = self.get_monitoring_token()
+            if not id_token:
+                self._debug_print("No valid cached id_token for silent refresh")
+                return None, None, None
+
+            self._debug_print("Found valid cached id_token, attempting silent credential refresh...")
+            token_claims = jwt.decode(id_token, options={"verify_signature": False})
+
+            credentials = self.get_aws_credentials(id_token, token_claims)
+            self.save_credentials(credentials)
+            self.save_monitoring_token(id_token, token_claims)
+            self._debug_print("Silent credential refresh succeeded")
+            return credentials, id_token, token_claims
+        except Exception as e:
+            self._debug_print(f"Silent refresh failed, will require browser auth: {e}")
+            return None, None, None
+
     def run(self):
         """Main execution flow"""
         try:
@@ -1817,7 +1841,23 @@ def run(self):
                 print(json.dumps(cached))  # noqa: S105
                 return 0
 
-            # Authenticate with OIDC provider
+            # Try silent refresh using cached id_token before opening browser
+            silent_creds, id_token, token_claims = self._try_silent_refresh()
+            if silent_creds:
+                # Check quota if configured (reuse token/claims already fetched above)
+                if self._should_check_quota():
+                    if id_token and token_claims:
+                        quota_result = self._check_quota(token_claims, id_token)
+                        self._save_quota_check_timestamp()
+                        if not quota_result.get("allowed", True):
+                            return self._handle_quota_blocked(quota_result)
+                        else:
+                            self._handle_quota_warning(quota_result)
+
+                print(json.dumps(silent_creds))
+                return 0
+
+            # Authenticate with OIDC provider (browser popup - only when id_token is also expired)
             self._debug_print(f"Authenticating with {self.provider_config['name']} for profile '{self.profile}'...")
             id_token, token_claims = self.authenticate_oidc()
 
diff --git a/source/otel_helper/__main__.py b/source/otel_helper/__main__.py
@@ -202,20 +202,23 @@ def get_cache_path():
 
 
 def read_cached_headers():
-    """Read cached OTEL headers if they exist and the token hasn't expired."""
+    """Read cached OTEL headers if they exist.
+
+    User attributes (email, team, etc.) don't change between sessions,
+    so cached headers are served regardless of token expiry. Headers are
+    refreshed opportunistically when a valid token is available.
+    """
     try:
         cache_path = get_cache_path()
         if not cache_path.exists():
             return None
         with open(cache_path) as f:
             cached = json.load(f)
-        # Check if token expires in more than 10 minutes
-        now = int(time.time())
-        if cached.get("token_exp", 0) - now > 600:
-            logger.debug("Using cached OTEL headers (token still valid)")
-            return cached["headers"]
-        logger.debug("Cached OTEL headers expired or expiring soon")
-        return None
+        headers = cached.get("headers")
+        if not headers:
+            return None
+        logger.debug("Using cached OTEL headers")
+        return headers
     except Exception as e:
         logger.debug(f"Failed to read cached headers: {e}")
         return None
@@ -332,10 +335,6 @@ def main():
 
         # Generate headers dictionary
         headers_dict = format_as_headers_dict(user_info)
-        # Only include Bearer token when ALB JWT validation is enabled (set by installer)
-        if os.environ.get("OTEL_JWT_AUTH", "").lower() in ("true", "1", "yes"):
-            headers_dict["authorization"] = f"Bearer {token}"
-
         # In test mode, print detailed output
         if TEST_MODE:
             print("===== TEST MODE OUTPUT =====\n")
diff --git a/source/otel_helper/otel-helper.sh b/source/otel_helper/otel-helper.sh
@@ -6,15 +6,10 @@ CACHE_DIR="$HOME/.claude-code-session"
 CACHE_FILE="$CACHE_DIR/${PROFILE}-otel-headers.json"
 RAW_FILE="$CACHE_DIR/${PROFILE}-otel-headers.raw"
 
-if [ -f "$CACHE_FILE" ] && [ -f "$RAW_FILE" ]; then
-    # Extract token_exp from metadata cache (date +%s is GNU/BSD, works on macOS and Linux)
-    TOKEN_EXP=$(grep -o '"token_exp": *[0-9]*' "$CACHE_FILE" | grep -o '[0-9]*')
-    NOW=$(date +%s)
-    if [ -n "$TOKEN_EXP" ] && [ "$((TOKEN_EXP - NOW))" -gt 600 ]; then
-        # Serve raw headers directly — no JSON parsing needed
-        cat "$RAW_FILE"
-        exit 0
-    fi
+if [ -f "$RAW_FILE" ]; then
+    # Serve raw headers directly — no JSON parsing needed
+    cat "$RAW_FILE"
+    exit 0
 fi
 
 # Cache miss - fall back to full PyInstaller binary (which writes the cache)
diff --git a/source/tests/test_silent_refresh.py b/source/tests/test_silent_refresh.py
@@ -0,0 +1,215 @@
+# ABOUTME: Tests for OIDC id_token caching and silent credential refresh (issue #153)
+# ABOUTME: Verifies that expired AWS credentials can be refreshed without browser popup
+"""Tests for silent credential refresh using cached OIDC id_token."""
+
+import json
+import time
+from unittest.mock import MagicMock, patch
+
+import jwt as pyjwt
+import pytest
+
+
+def _make_id_token(exp_offset=3600, email="test@example.com"):
+    """Create a minimal JWT id_token for testing.
+
+    Args:
+        exp_offset: Seconds from now until expiration (positive = future).
+        email: Email claim to embed.
+    """
+    claims = {
+        "sub": "user-123",
+        "email": email,
+        "iss": "https://test.okta.com",
+        "aud": "test-client-id",
+        "exp": int(time.time()) + exp_offset,
+        "iat": int(time.time()),
+        "nonce": "test-nonce",
+    }
+    # Encode without signing (matches how the provider decodes with verify_signature=False)
+    return pyjwt.encode(claims, "secret", algorithm="HS256"), claims
+
+
+def _make_config():
+    """Return a minimal config dict for MultiProviderAuth."""
+    return {
+        "profiles": {
+            "TestProfile": {
+                "provider_domain": "test.okta.com",
+                "client_id": "test-client-id",
+                "identity_pool_id": "us-east-1:test-pool",
+                "aws_region": "us-east-1",
+                "credential_storage": "session",
+            }
+        }
+    }
+
+
+def _make_aws_credentials(exp_offset=900):
+    """Return fake AWS credentials dict."""
+    from datetime import datetime, timezone, timedelta
+
+    exp = datetime.now(timezone.utc) + timedelta(seconds=exp_offset)
+    return {
+        "Version": 1,
+        "AccessKeyId": "AKIAIOSFODNN7EXAMPLE",
+        "SecretAccessKey": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
+        "SessionToken": "FwoGZXIvYXdzEBYaDH...",
+        "Expiration": exp.isoformat(),
+    }
+
+
+@pytest.fixture
+def auth_instance(tmp_path):
+    """Create a MultiProviderAuth instance with mocked config."""
+    config_file = tmp_path / "config.json"
+    config_file.write_text(json.dumps(_make_config()))
+
+    with patch("credential_provider.__main__.Path") as mock_path_cls:
+        # Make _load_config find our temp config
+        mock_home = MagicMock()
+        mock_path_cls.home.return_value = mock_home
+        mock_home.__truediv__ = lambda self, key: tmp_path / key if key == "claude-code-with-bedrock" else MagicMock()
+
+        # Also mock __file__ parent for binary dir config lookup
+        mock_file_parent = MagicMock()
+        mock_file_parent.__truediv__ = lambda self, key: MagicMock(exists=lambda: False)
+        mock_path_cls.return_value = mock_file_parent
+
+        # Simpler approach: just patch _load_config and _init_credential_storage
+        with patch("credential_provider.__main__.MultiProviderAuth._load_config") as mock_load, \
+             patch("credential_provider.__main__.MultiProviderAuth._init_credential_storage"):
+            mock_load.return_value = {
+                "provider_domain": "test.okta.com",
+                "client_id": "test-client-id",
+                "identity_pool_id": "us-east-1:test-pool",
+                "aws_region": "us-east-1",
+                "credential_storage": "session",
+                "provider_type": "okta",
+                "federation_type": "cognito",
+                "max_session_duration": 28800,
+            }
+
+            from credential_provider.__main__ import MultiProviderAuth
+            instance = MultiProviderAuth(profile="TestProfile")
+            instance.cache_dir = tmp_path / "cache"
+            instance.cache_dir.mkdir(parents=True, exist_ok=True)
+            return instance
+
+
+class TestSilentRefresh:
+    """Tests for _try_silent_refresh method."""
+
+    def test_silent_refresh_succeeds_with_valid_id_token(self, auth_instance):
+        """When a valid id_token is cached, silent refresh should return new AWS creds."""
+        id_token, claims = _make_id_token(exp_offset=3600)
+        aws_creds = _make_aws_credentials()
+
+        with patch.object(auth_instance, "get_monitoring_token", return_value=id_token), \
+             patch.object(auth_instance, "get_aws_credentials", return_value=aws_creds) as mock_get_creds, \
+             patch.object(auth_instance, "save_credentials") as mock_save, \
+             patch.object(auth_instance, "save_monitoring_token") as mock_save_token:
+
+            creds, returned_token, returned_claims = auth_instance._try_silent_refresh()
+
+            assert creds is not None
+            assert creds["AccessKeyId"] == aws_creds["AccessKeyId"]
+            assert returned_token == id_token
+            assert returned_claims["sub"] == claims["sub"]
+            mock_get_creds.assert_called_once()
+            mock_save.assert_called_once_with(aws_creds)
+            # Verify the id_token is re-persisted so the next refresh also works
+            mock_save_token.assert_called_once_with(id_token, claims)
+
+    def test_silent_refresh_returns_none_when_id_token_expired(self, auth_instance):
+        """When cached id_token is within the 60-second expiry buffer, get_monitoring_token
+        returns None and silent refresh must not attempt an STS exchange."""
+        with patch.object(auth_instance, "get_monitoring_token", return_value=None) as mock_get_token, \
+             patch.object(auth_instance, "get_aws_credentials") as mock_get_creds:
+
+            creds, id_token, token_claims = auth_instance._try_silent_refresh()
+
+            assert creds is None
+            assert id_token is None
+            assert token_claims is None
+            mock_get_token.assert_called_once()
+            # STS must never be called when the token is expired
+            mock_get_creds.assert_not_called()
+
+    def test_silent_refresh_returns_none_when_no_cached_token(self, auth_instance):
+        """When no id_token is cached, silent refresh should return None."""
+        with patch.object(auth_instance, "get_monitoring_token", return_value=None):
+            creds, id_token, token_claims = auth_instance._try_silent_refresh()
+            assert creds is None
+            assert id_token is None
+            assert token_claims is None
+
+    def test_silent_refresh_returns_none_when_sts_exchange_fails(self, auth_instance):
+        """When id_token is valid but STS exchange fails, should return None (fallback to browser)."""
+        id_token, _ = _make_id_token(exp_offset=3600)
+
+        with patch.object(auth_instance, "get_monitoring_token", return_value=id_token), \
+             patch.object(auth_instance, "get_aws_credentials", side_effect=Exception("STS error")):
+
+            creds, returned_token, returned_claims = auth_instance._try_silent_refresh()
+            assert creds is None
+            assert returned_token is None
+            assert returned_claims is None
+
+    def test_silent_refresh_not_called_when_aws_creds_valid(self, auth_instance):
+        """When AWS credentials are still valid, silent refresh should not be attempted."""
+        aws_creds = _make_aws_credentials(exp_offset=3600)
+
+        with patch.object(auth_instance, "get_cached_credentials", return_value=aws_creds), \
+             patch.object(auth_instance, "_try_silent_refresh") as mock_silent, \
+             patch.object(auth_instance, "_should_recheck_quota", return_value=False):
+
+            # Capture stdout
+            with patch("builtins.print"):
+                auth_instance.run()
+
+            mock_silent.assert_not_called()
+
+    def test_run_uses_silent_refresh_before_browser(self, auth_instance):
+        """When AWS creds expired but id_token valid, run() should use silent refresh."""
+        aws_creds = _make_aws_credentials(exp_offset=3600)
+
+        with patch.object(auth_instance, "get_cached_credentials", return_value=None), \
+             patch("socket.socket") as mock_socket_cls, \
+             patch.object(auth_instance, "_try_silent_refresh", return_value=(aws_creds, None, None)), \
+             patch.object(auth_instance, "_should_check_quota", return_value=False), \
+             patch.object(auth_instance, "authenticate_oidc") as mock_browser, \
+             patch("builtins.print"):
+
+            # Mock socket to simulate port available
+            mock_socket = MagicMock()
+            mock_socket_cls.return_value = mock_socket
+
+            result = auth_instance.run()
+
+            assert result == 0
+            mock_browser.assert_not_called()
+
+    def test_run_falls_back_to_browser_when_silent_refresh_fails(self, auth_instance):
+        """When silent refresh fails, run() should fall back to browser auth."""
+        id_token, claims = _make_id_token(exp_offset=3600)
+        aws_creds = _make_aws_credentials(exp_offset=3600)
+
+        with patch.object(auth_instance, "get_cached_credentials", return_value=None), \
+             patch("socket.socket") as mock_socket_cls, \
+             patch.object(auth_instance, "_try_silent_refresh", return_value=(None, None, None)), \
+             patch.object(auth_instance, "authenticate_oidc", return_value=(id_token, claims)) as mock_browser, \
+             patch.object(auth_instance, "_should_check_quota", return_value=False), \
+             patch.object(auth_instance, "get_aws_credentials", return_value=aws_creds), \
+             patch.object(auth_instance, "save_credentials"), \
+             patch.object(auth_instance, "save_monitoring_token"), \
+             patch("builtins.print"):
+
+            mock_socket = MagicMock()
+            mock_socket_cls.return_value = mock_socket
+
+            result = auth_instance.run()
+
+            assert result == 0
+            mock_browser.assert_called_once()
+
