fastmcp/tests/server/auth/oauth_proxy/test_tokens.py

"""Tests for OAuth proxy token endpoint and handling."""

import time
from unittest.mock import AsyncMock, Mock, patch

import pytest
from key_value.aio.stores.memory import MemoryStore
from mcp.server.auth.handlers.token import TokenErrorResponse
from mcp.server.auth.handlers.token import TokenHandler as SDKTokenHandler
from mcp.server.auth.provider import AccessToken, AuthorizationCode
from mcp.shared.auth import OAuthClientInformationFull
from pydantic import AnyUrl

from fastmcp.server.auth.auth import RefreshToken, TokenHandler, TokenVerifier
from fastmcp.server.auth.oauth_proxy import OAuthProxy
from fastmcp.server.auth.oauth_proxy.models import (
    DEFAULT_ACCESS_TOKEN_EXPIRY_NO_REFRESH_SECONDS,
    DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS,
    ClientCode,
    JTIMapping,
    UpstreamTokenSet,
    _hash_token,
)
from fastmcp.server.auth.providers.jwt import JWTVerifier


class TestOAuthProxyTokenEndpointAuth:
    """Tests for token endpoint authentication methods."""

    def test_token_auth_method_initialization(self, jwt_verifier):
        """Test different token endpoint auth methods."""
        # client_secret_post
        proxy_post = OAuthProxy(
            upstream_authorization_endpoint="https://oauth.example.com/authorize",
            upstream_token_endpoint="https://oauth.example.com/token",
            upstream_client_id="client",
            upstream_client_secret="secret",
            token_verifier=jwt_verifier,
            base_url="https://proxy.example.com",
            token_endpoint_auth_method="client_secret_post",
            jwt_signing_key="test-secret",
            client_storage=MemoryStore(),
        )
        assert proxy_post._token_endpoint_auth_method == "client_secret_post"

        # client_secret_basic (default)
        proxy_basic = OAuthProxy(
            upstream_authorization_endpoint="https://oauth.example.com/authorize",
            upstream_token_endpoint="https://oauth.example.com/token",
            upstream_client_id="client",
            upstream_client_secret="secret",
            token_verifier=jwt_verifier,
            base_url="https://proxy.example.com",
            token_endpoint_auth_method="client_secret_basic",
            jwt_signing_key="test-secret",
            client_storage=MemoryStore(),
        )
        assert proxy_basic._token_endpoint_auth_method == "client_secret_basic"

        # None (use authlib default)
        proxy_default = OAuthProxy(
            upstream_authorization_endpoint="https://oauth.example.com/authorize",
            upstream_token_endpoint="https://oauth.example.com/token",
            upstream_client_id="client",
            upstream_client_secret="secret",
            token_verifier=jwt_verifier,
            base_url="https://proxy.example.com",
            jwt_signing_key="test-secret",
            client_storage=MemoryStore(),
        )
        assert proxy_default._token_endpoint_auth_method is None

    async def test_token_auth_method_passed_to_client(self, jwt_verifier):
        """Test that auth method is passed to AsyncOAuth2Client."""
        proxy = OAuthProxy(
            upstream_authorization_endpoint="https://oauth.example.com/authorize",
            upstream_token_endpoint="https://oauth.example.com/token",
            upstream_client_id="client-id",
            upstream_client_secret="client-secret",
            token_verifier=jwt_verifier,
            base_url="https://proxy.example.com",
            token_endpoint_auth_method="client_secret_post",
            jwt_signing_key="test-secret",
            client_storage=MemoryStore(),
        )

        # Initialize JWT issuer before token operations
        proxy.set_mcp_path("/mcp")

        # First, create a valid FastMCP token via full OAuth flow
        client = OAuthClientInformationFull(
            client_id="test-client",
            client_secret="test-secret",
            redirect_uris=[AnyUrl("http://localhost:12345/callback")],
        )

        # Mock the upstream OAuth provider response
        with patch(
            "fastmcp.server.auth.oauth_proxy.proxy.AsyncOAuth2Client"
        ) as MockClient:
            mock_client = AsyncMock()

            # Mock initial token exchange (authorization code flow)
            mock_client.fetch_token = AsyncMock(
                return_value={
                    "access_token": "upstream-access-token",
                    "refresh_token": "upstream-refresh-token",
                    "expires_in": 3600,
                    "token_type": "Bearer",
                }
            )

            # Mock token refresh
            mock_client.refresh_token = AsyncMock(
                return_value={
                    "access_token": "new-upstream-token",
                    "refresh_token": "new-upstream-refresh",
                    "expires_in": 3600,
                    "token_type": "Bearer",
                }
            )
            MockClient.return_value = mock_client

            # Register client and do initial OAuth flow to get valid FastMCP tokens
            await proxy.register_client(client)

            # Store client code that would be created during OAuth callback
            client_code = ClientCode(
                code="test-auth-code",
                client_id="test-client",
                redirect_uri="http://localhost:12345/callback",
                code_challenge="",
                code_challenge_method="S256",
                scopes=["read"],
                idp_tokens={
                    "access_token": "upstream-access-token",
                    "refresh_token": "upstream-refresh-token",
                    "expires_in": 3600,
                    "token_type": "Bearer",
                },
                expires_at=time.time() + 300,
                created_at=time.time(),
            )
            await proxy._code_store.put(key=client_code.code, value=client_code)

            # Exchange authorization code to get FastMCP tokens
            auth_code = AuthorizationCode(
                code="test-auth-code",
                scopes=["read"],
                expires_at=time.time() + 300,
                client_id="test-client",
                code_challenge="",
                redirect_uri=AnyUrl("http://localhost:12345/callback"),
                redirect_uri_provided_explicitly=True,
            )
            result = await proxy.exchange_authorization_code(
                client=client,
                authorization_code=auth_code,
            )

            # Now test refresh with the valid FastMCP refresh token
            assert result.refresh_token is not None
            fastmcp_refresh = RefreshToken(
                token=result.refresh_token,
                client_id="test-client",
                scopes=["read"],
                expires_at=None,
            )

            # Reset mock to check refresh call
            MockClient.reset_mock()
            mock_client.refresh_token = AsyncMock(
                return_value={
                    "access_token": "new-upstream-token-2",
                    "refresh_token": "new-upstream-refresh-2",
                    "expires_in": 3600,
                    "token_type": "Bearer",
                }
            )
            MockClient.return_value = mock_client

            await proxy.exchange_refresh_token(client, fastmcp_refresh, ["read"])

            # Verify auth method was passed to OAuth client
            MockClient.assert_called_with(
                client_id="client-id",
                client_secret="client-secret",
                token_endpoint_auth_method="client_secret_post",
                timeout=30.0,
            )


class TestTokenHandlerErrorTransformation:
    """Tests for TokenHandler's OAuth 2.1 compliant error transformation."""

    async def test_transforms_client_auth_failure_to_invalid_client_401(self):
        """Test that client authentication failures return invalid_client with 401."""
        handler = TokenHandler(provider=Mock(), client_authenticator=Mock())

        # Create a mock 401 response like the SDK returns for auth failures
        mock_response = Mock()
        mock_response.status_code = 401
        mock_response.body = (
            b'{"error":"unauthorized_client","error_description":"Invalid client_id"}'
        )

        # Patch the parent class's handle() to return our mock response
        with patch.object(
            SDKTokenHandler,
            "handle",
            new_callable=AsyncMock,
            return_value=mock_response,
        ):
            response = await handler.handle(Mock())

        # Should transform to OAuth 2.1 compliant response
        assert response.status_code == 401
        assert b'"error":"invalid_client"' in response.body
        assert b'"error_description":"Invalid client_id"' in response.body

    def test_does_not_transform_grant_type_unauthorized_to_invalid_client(self):
        """Test that grant type authorization errors stay as unauthorized_client with 400."""
        handler = TokenHandler(provider=Mock(), client_authenticator=Mock())

        # Simulate error from grant_type not in client_info.grant_types
        error_response = TokenErrorResponse(
            error="unauthorized_client",
            error_description="Client not authorized for this grant type",
        )

        response = handler.response(error_response)

        # Should NOT transform - keep as 400 unauthorized_client
        assert response.status_code == 400
        assert b'"error":"unauthorized_client"' in response.body

    async def test_transforms_invalid_grant_to_401(self):
        """Test that invalid_grant errors return 401 per MCP spec.

        Per MCP spec: "Invalid or expired tokens MUST receive a HTTP 401 response."
        The SDK incorrectly returns 400 for all TokenErrorResponse including invalid_grant.
        """
        handler = TokenHandler(provider=Mock(), client_authenticator=Mock())

        # Create a mock 400 response like the SDK returns for invalid_grant
        mock_response = Mock()
        mock_response.status_code = 400
        mock_response.body = (
            b'{"error":"invalid_grant","error_description":"refresh token has expired"}'
        )

        # Patch the parent class's handle() to return our mock response
        with patch.object(
            SDKTokenHandler,
            "handle",
            new_callable=AsyncMock,
            return_value=mock_response,
        ):
            response = await handler.handle(Mock())

        # Should transform to MCP-compliant 401 response
        assert response.status_code == 401
        assert b'"error":"invalid_grant"' in response.body
        assert b'"error_description":"refresh token has expired"' in response.body

    def test_does_not_transform_other_400_errors(self):
        """Test that non-invalid_grant 400 errors pass through unchanged."""
        handler = TokenHandler(provider=Mock(), client_authenticator=Mock())

        # Test with invalid_request error (should stay 400)
        error_response = TokenErrorResponse(
            error="invalid_request",
            error_description="Missing required parameter",
        )

        response = handler.response(error_response)

        # Should pass through unchanged as 400
        assert response.status_code == 400
        assert b'"error":"invalid_request"' in response.body


class TestFallbackAccessTokenExpiry:
    """Test fallback access token expiry constants and configuration."""

    def test_default_constants(self):
        """Verify the default expiry constants are set correctly."""
        assert DEFAULT_ACCESS_TOKEN_EXPIRY_SECONDS == 60 * 60  # 1 hour
        assert (
            DEFAULT_ACCESS_TOKEN_EXPIRY_NO_REFRESH_SECONDS == 60 * 60 * 24 * 365
        )  # 1 year

    def test_fallback_parameter_stored(self):
        """Verify fallback_access_token_expiry_seconds is stored on provider."""
        provider = OAuthProxy(
            upstream_authorization_endpoint="https://idp.example.com/authorize",
            upstream_token_endpoint="https://idp.example.com/token",
            upstream_client_id="test-client",
            upstream_client_secret="test-secret",
            token_verifier=JWTVerifier(
                jwks_uri="https://idp.example.com/.well-known/jwks.json",
                issuer="https://idp.example.com",
            ),
            base_url="http://localhost:8000",
            jwt_signing_key="test-signing-key",
            fallback_access_token_expiry_seconds=86400,
            client_storage=MemoryStore(),
        )

        assert provider._fallback_access_token_expiry_seconds == 86400

    def test_fallback_parameter_defaults_to_none(self):
        """Verify fallback defaults to None (enabling smart defaults)."""
        provider = OAuthProxy(
            upstream_authorization_endpoint="https://idp.example.com/authorize",
            upstream_token_endpoint="https://idp.example.com/token",
            upstream_client_id="test-client",
            upstream_client_secret="test-secret",
            token_verifier=JWTVerifier(
                jwks_uri="https://idp.example.com/.well-known/jwks.json",
                issuer="https://idp.example.com",
            ),
            base_url="http://localhost:8000",
            jwt_signing_key="test-signing-key",
            client_storage=MemoryStore(),
        )

        assert provider._fallback_access_token_expiry_seconds is None


class TestUpstreamTokenStorageTTL:
    """Tests for upstream token storage TTL calculation (issue #2670).

    The TTL should use max(refresh_expires_in, expires_in) to handle cases where
    the refresh token has a shorter lifetime than the access token (e.g., Keycloak
    with sliding session windows).
    """

    @pytest.fixture
    def jwt_verifier(self):
        """Create a mock JWT verifier."""
        verifier = Mock(spec=TokenVerifier)
        verifier.required_scopes = ["read", "write"]
        verifier.verify_token = AsyncMock(return_value=None)
        return verifier

    @pytest.fixture
    def proxy(self, jwt_verifier):
        """Create an OAuth proxy for testing."""
        proxy = OAuthProxy(
            upstream_authorization_endpoint="https://idp.example.com/authorize",
            upstream_token_endpoint="https://idp.example.com/token",
            upstream_client_id="test-client",
            upstream_client_secret="test-secret",
            token_verifier=jwt_verifier,
            base_url="https://proxy.example.com",
            jwt_signing_key="test-secret-key",
            client_storage=MemoryStore(),
        )
        proxy.set_mcp_path("/mcp")
        return proxy

    async def test_ttl_uses_max_when_refresh_shorter_than_access(self, proxy):
        """TTL should use access token expiry when refresh is shorter.

        This is the xsreality case: Keycloak returns refresh_expires_in=120 (2 min)
        but expires_in=28800 (8 hours). The upstream tokens should persist for
        8 hours (the access token lifetime), not 2 minutes.
        """
        # Register client
        client = OAuthClientInformationFull(
            client_id="test-client",
            client_secret="test-secret",
            redirect_uris=[AnyUrl("http://localhost:12345/callback")],
        )
        await proxy.register_client(client)

        # Simulate xsreality's Keycloak setup: short refresh, long access
        client_code = ClientCode(
            code="test-auth-code",
            client_id="test-client",
            redirect_uri="http://localhost:12345/callback",
            code_challenge="test-challenge",
            code_challenge_method="S256",
            scopes=["read", "write"],
            idp_tokens={
                "access_token": "upstream-access-token",
                "refresh_token": "upstream-refresh-token",
                "expires_in": 28800,  # 8 hours (access token)
                "refresh_expires_in": 120,  # 2 minutes (refresh token) - SHORTER!
                "token_type": "Bearer",
            },
            expires_at=time.time() + 300,
            created_at=time.time(),
        )
        await proxy._code_store.put(key=client_code.code, value=client_code)

        # Exchange the code
        auth_code = AuthorizationCode(
            code="test-auth-code",
            scopes=["read", "write"],
            expires_at=time.time() + 300,
            client_id="test-client",
            code_challenge="test-challenge",
            redirect_uri=AnyUrl("http://localhost:12345/callback"),
            redirect_uri_provided_explicitly=True,
        )

        result = await proxy.exchange_authorization_code(
            client=client,
            authorization_code=auth_code,
        )

        # Verify tokens were issued
        assert result.access_token is not None
        assert result.refresh_token is not None

        # The key test: verify upstream tokens are stored with TTL=max(120, 28800)=28800
        # We can verify this by checking the tokens are still accessible after 2 minutes
        # would have passed (if TTL was incorrectly set to 120)
        #
        # Since we can't easily time-travel in tests, we verify the storage directly
        # by checking that we can still look up the tokens for refresh purposes.
        #
        # Extract the JTI from the refresh token to look up the mapping
        refresh_payload = proxy.jwt_issuer.verify_token(
            result.refresh_token, expected_token_use="refresh"
        )
        refresh_jti = refresh_payload["jti"]

        # The JTI mapping should exist
        jti_mapping = await proxy._jti_mapping_store.get(key=refresh_jti)
        assert jti_mapping is not None

        # The upstream tokens should exist
        upstream_tokens = await proxy._upstream_token_store.get(
            key=jti_mapping.upstream_token_id
        )
        assert upstream_tokens is not None
        assert upstream_tokens.access_token == "upstream-access-token"
        assert upstream_tokens.refresh_token == "upstream-refresh-token"

    async def test_ttl_uses_refresh_when_refresh_longer_than_access(self, proxy):
        """TTL should use refresh token expiry when refresh is longer.

        This is the ianw case: IdP returns expires_in=300 (5 min) but
        refresh_expires_in=32318 (9 hours). The upstream tokens should persist
        for 9 hours (the refresh token lifetime).
        """
        # Register client
        client = OAuthClientInformationFull(
            client_id="test-client",
            client_secret="test-secret",
            redirect_uris=[AnyUrl("http://localhost:12345/callback")],
        )
        await proxy.register_client(client)

        # Simulate ianw's setup: short access, long refresh (typical)
        client_code = ClientCode(
            code="test-auth-code-2",
            client_id="test-client",
            redirect_uri="http://localhost:12345/callback",
            code_challenge="test-challenge",
            code_challenge_method="S256",
            scopes=["read", "write"],
            idp_tokens={
                "access_token": "upstream-access-token-2",
                "refresh_token": "upstream-refresh-token-2",
                "expires_in": 300,  # 5 minutes (access token)
                "refresh_expires_in": 32318,  # 9 hours (refresh token) - LONGER
                "token_type": "Bearer",
            },
            expires_at=time.time() + 300,
            created_at=time.time(),
        )
        await proxy._code_store.put(key=client_code.code, value=client_code)

        # Exchange the code
        auth_code = AuthorizationCode(
            code="test-auth-code-2",
            scopes=["read", "write"],
            expires_at=time.time() + 300,
            client_id="test-client",
            code_challenge="test-challenge",
            redirect_uri=AnyUrl("http://localhost:12345/callback"),
            redirect_uri_provided_explicitly=True,
        )

        result = await proxy.exchange_authorization_code(
            client=client,
            authorization_code=auth_code,
        )

        # Verify tokens were issued
        assert result.access_token is not None
        assert result.refresh_token is not None

        # Verify upstream tokens are accessible
        refresh_payload = proxy.jwt_issuer.verify_token(
            result.refresh_token, expected_token_use="refresh"
        )
        refresh_jti = refresh_payload["jti"]

        jti_mapping = await proxy._jti_mapping_store.get(key=refresh_jti)
        assert jti_mapping is not None

        upstream_tokens = await proxy._upstream_token_store.get(
            key=jti_mapping.upstream_token_id
        )
        assert upstream_tokens is not None

    async def test_refresh_expires_in_zero_issues_refresh_token(self, proxy):
        """refresh_expires_in=0 should fall back to 30-day default.

        Keycloak returns refresh_expires_in=0 for offline tokens (offline_access scope),
        meaning "no fixed time-based expiry". The proxy should still issue a PROXY_RT.
        """
        client = OAuthClientInformationFull(
            client_id="test-client",
            client_secret="test-secret",
            redirect_uris=[AnyUrl("http://localhost:12345/callback")],
        )
        await proxy.register_client(client)

        client_code = ClientCode(
            code="test-auth-code-keycloak-offline",
            client_id="test-client",
            redirect_uri="http://localhost:12345/callback",
            code_challenge="test-challenge",
            code_challenge_method="S256",
            scopes=["read", "write"],
            idp_tokens={
                "access_token": "upstream-access-token-kc",
                "refresh_token": "upstream-refresh-token-kc",
                "expires_in": 3600,
                "refresh_expires_in": 0,  # Keycloak offline token convention
                "token_type": "Bearer",
            },
            expires_at=time.time() + 300,
            created_at=time.time(),
        )
        await proxy._code_store.put(key=client_code.code, value=client_code)

        auth_code = AuthorizationCode(
            code="test-auth-code-keycloak-offline",
            scopes=["read", "write"],
            expires_at=time.time() + 300,
            client_id="test-client",
            code_challenge="test-challenge",
            redirect_uri=AnyUrl("http://localhost:12345/callback"),
            redirect_uri_provided_explicitly=True,
        )

        result = await proxy.exchange_authorization_code(
            client=client,
            authorization_code=auth_code,
        )

        # refresh_expires_in=0 must NOT prevent refresh token issuance
        assert result.access_token is not None
        assert result.refresh_token is not None

        # Verify refresh token metadata was stored
        refresh_meta = await proxy._refresh_token_store.get(
            key=_hash_token(result.refresh_token)
        )
        assert refresh_meta is not None


class TestTransparentUpstreamRefresh:
    """Tests for transparent upstream token refresh in load_access_token.

    When the upstream token expires but a refresh token is available, the proxy
    should transparently refresh the upstream token rather than returning a 401
    that forces the client into a full re-authentication flow.
    """

    @pytest.fixture
    def mock_verifier(self):
        """Token verifier that rejects expired tokens, accepts refreshed ones."""
        verifier = Mock(spec=TokenVerifier)
        verifier.required_scopes = ["read"]

        async def verify(token: str) -> AccessToken | None:
            if token.startswith("refreshed-"):
                return AccessToken(
                    token=token,
                    client_id="test-client",
                    scopes=["read"],
                    expires_at=int(time.time() + 3600),
                )
            return None

        verifier.verify_token = AsyncMock(side_effect=verify)
        return verifier

    @pytest.fixture
    def proxy(self, mock_verifier):
        proxy = OAuthProxy(
            upstream_authorization_endpoint="https://idp.example.com/authorize",
            upstream_token_endpoint="https://idp.example.com/token",
            upstream_client_id="test-client",
            upstream_client_secret="test-secret",
            token_verifier=mock_verifier,
            base_url="https://proxy.example.com",
            jwt_signing_key="test-secret-key",
            client_storage=MemoryStore(),
        )
        proxy.set_mcp_path("/mcp")
        return proxy

    async def _setup_expired_session(
        self,
        proxy: OAuthProxy,
        *,
        upstream_refresh_token: str | None = "upstream-refresh-tok",
    ) -> str:
        """Set up a proxy JWT pointing at an expired upstream token.

        Returns the proxy JWT (access token) that can be passed to
        load_access_token.
        """
        upstream_token_id = "upstream-tok-id"
        access_jti = "test-access-jti"

        upstream_token_set = UpstreamTokenSet(
            upstream_token_id=upstream_token_id,
            access_token="expired-upstream-access",
            refresh_token=upstream_refresh_token,
            refresh_token_expires_at=time.time() + 86400
            if upstream_refresh_token
            else None,
            expires_at=time.time() - 60,  # expired 1 minute ago
            token_type="Bearer",
            scope="read",
            client_id="test-client",
            created_at=time.time() - 3600,
        )
        await proxy._upstream_token_store.put(
            key=upstream_token_id,
            value=upstream_token_set,
            ttl=86400,
        )

        await proxy._jti_mapping_store.put(
            key=access_jti,
            value=JTIMapping(
                jti=access_jti,
                upstream_token_id=upstream_token_id,
                created_at=time.time(),
            ),
            ttl=3600,
        )

        fastmcp_jwt = proxy.jwt_issuer.issue_access_token(
            client_id="test-client",
            scopes=["read"],
            jti=access_jti,
            expires_in=3600,
        )
        return fastmcp_jwt

    async def test_transparent_refresh_on_expired_upstream(self, proxy):
        """load_access_token refreshes upstream token when validation fails."""
        fastmcp_jwt = await self._setup_expired_session(proxy)

        mock_oauth_client = AsyncMock()
        mock_oauth_client.refresh_token = AsyncMock(
            return_value={
                "access_token": "refreshed-upstream-access",
                "token_type": "Bearer",
                "expires_in": 3600,
                "refresh_token": "upstream-refresh-tok",
                "scope": "read",
            }
        )

        with patch.object(
            proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
        ):
            result = await proxy.load_access_token(fastmcp_jwt)

        assert result is not None
        assert result.token == "refreshed-upstream-access"
        mock_oauth_client.refresh_token.assert_called_once()

    async def test_transparent_refresh_updates_stored_token(self, proxy):
        """After transparent refresh, the stored upstream token is updated."""
        fastmcp_jwt = await self._setup_expired_session(proxy)

        mock_oauth_client = AsyncMock()
        mock_oauth_client.refresh_token = AsyncMock(
            return_value={
                "access_token": "refreshed-upstream-access",
                "token_type": "Bearer",
                "expires_in": 7200,
                "refresh_token": "upstream-refresh-tok",
                "scope": "read",
            }
        )

        with patch.object(
            proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
        ):
            await proxy.load_access_token(fastmcp_jwt)

        stored = await proxy._upstream_token_store.get(key="upstream-tok-id")
        assert stored is not None
        assert stored.access_token == "refreshed-upstream-access"
        assert stored.expires_at > time.time()

    async def test_no_refresh_without_refresh_token(self, proxy):
        """Without a refresh token, load_access_token returns None immediately."""
        fastmcp_jwt = await self._setup_expired_session(
            proxy, upstream_refresh_token=None
        )

        result = await proxy.load_access_token(fastmcp_jwt)

        assert result is None

    async def test_returns_none_when_refresh_fails(self, proxy):
        """If the upstream refresh call fails, load_access_token returns None."""
        fastmcp_jwt = await self._setup_expired_session(proxy)

        mock_oauth_client = AsyncMock()
        mock_oauth_client.refresh_token = AsyncMock(
            side_effect=Exception("upstream refused refresh")
        )

        with patch.object(
            proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
        ):
            result = await proxy.load_access_token(fastmcp_jwt)

        assert result is None

    async def test_returns_none_when_refreshed_token_still_invalid(self, proxy):
        """If the refreshed token also fails validation, returns None."""
        fastmcp_jwt = await self._setup_expired_session(proxy)

        mock_oauth_client = AsyncMock()
        mock_oauth_client.refresh_token = AsyncMock(
            return_value={
                "access_token": "still-bad-token",
                "token_type": "Bearer",
                "expires_in": 3600,
                "scope": "read",
            }
        )

        with patch.object(
            proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
        ):
            result = await proxy.load_access_token(fastmcp_jwt)

        # "still-bad-token" doesn't start with "refreshed-" so verifier rejects it
        assert result is None

    async def test_no_refresh_when_token_not_expired(self, proxy):
        """Non-expiry validation failures (e.g. revocation) should not trigger refresh."""
        upstream_token_id = "upstream-tok-id"
        access_jti = "test-access-jti"

        # Token is NOT expired — verification failure is for another reason
        upstream_token_set = UpstreamTokenSet(
            upstream_token_id=upstream_token_id,
            access_token="revoked-upstream-access",
            refresh_token="upstream-refresh-tok",
            refresh_token_expires_at=time.time() + 86400,
            expires_at=time.time() + 3600,  # still valid for 1 hour
            token_type="Bearer",
            scope="read",
            client_id="test-client",
            created_at=time.time() - 60,
        )
        await proxy._upstream_token_store.put(
            key=upstream_token_id,
            value=upstream_token_set,
            ttl=86400,
        )
        await proxy._jti_mapping_store.put(
            key=access_jti,
            value=JTIMapping(
                jti=access_jti,
                upstream_token_id=upstream_token_id,
                created_at=time.time(),
            ),
            ttl=3600,
        )
        fastmcp_jwt = proxy.jwt_issuer.issue_access_token(
            client_id="test-client",
            scopes=["read"],
            jti=access_jti,
            expires_in=3600,
        )

        mock_oauth_client = AsyncMock()
        mock_oauth_client.refresh_token = AsyncMock()

        with patch.object(
            proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
        ):
            result = await proxy.load_access_token(fastmcp_jwt)

        assert result is None
        # Refresh should NOT have been attempted
        mock_oauth_client.refresh_token.assert_not_called()

    async def test_reload_from_storage_after_refresh_failure(self, proxy):
        """If refresh fails, re-read from storage in case another worker refreshed."""
        fastmcp_jwt = await self._setup_expired_session(proxy)

        # Simulate: refresh fails (stale refresh token), but another worker
        # already wrote a fresh upstream token to storage.
        refreshed_token_set = UpstreamTokenSet(
            upstream_token_id="upstream-tok-id",
            access_token="refreshed-upstream-access",
            refresh_token="new-refresh-tok",
            refresh_token_expires_at=time.time() + 86400,
            expires_at=time.time() + 3600,
            token_type="Bearer",
            scope="read",
            client_id="test-client",
            created_at=time.time(),
        )

        # The store is read three times:
        # 1. Initial lookup in load_access_token
        # 2. Re-read inside the advisory lock
        # 3. Re-read after refresh failure (recovery path)
        original_get = proxy._upstream_token_store.get
        call_count = 0

        async def mock_get(key: str) -> UpstreamTokenSet | None:
            nonlocal call_count
            call_count += 1
            if call_count <= 2:
                return await original_get(key)
            # After refresh failure, return the "other worker's" refreshed token
            return refreshed_token_set

        mock_oauth_client = AsyncMock()
        mock_oauth_client.refresh_token = AsyncMock(
            side_effect=Exception("refresh token rotated by another worker")
        )

        with (
            patch.object(
                proxy, "_create_upstream_oauth_client", return_value=mock_oauth_client
            ),
            patch.object(proxy._upstream_token_store, "get", side_effect=mock_get),
        ):
            result = await proxy.load_access_token(fastmcp_jwt)

        assert result is not None
        assert result.token == "refreshed-upstream-access"