Add token_expiry_threshold_seconds for proactive token refresh (#4142)

Co-authored-by: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>

fastmcp_slim/fastmcp/server/auth/oauth_proxy/proxy.py

@@ -271,6 +271,8 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
         # Token expiry fallback
         fallback_access_token_expiry_seconds: int | None = None,
         fallback_refresh_token_expiry_seconds: int | None = None,
+        # Token refresh threshold
+        token_expiry_threshold_seconds: int = 0,
         # CIMD (Client ID Metadata Document) support
         enable_cimd: bool = True,
     ):
@@ -346,6 +348,10 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
                 lifetime — the actual upstream refresh remains the source of truth. If the
                 upstream rejects the refresh, the client gets `invalid_grant` and re-auths,
                 regardless of how much life is left on the FastMCP refresh token.
+            token_expiry_threshold_seconds: Number of seconds before actual expiry to consider
+                a token as expired (default 0). This prevents race conditions where a token
+                passes the expiry check but expires before the next operation completes.
+                For example, set to 30 to refresh tokens that will expire within 30 seconds.
             enable_cimd: Enable CIMD (Client ID Metadata Document) support for URL-based
                 client IDs. When True, clients can authenticate using HTTPS URLs as client
                 IDs, with metadata fetched from the URL. Supports private_key_jwt auth.
@@ -450,6 +456,7 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
             if fallback_refresh_token_expiry_seconds is not None
             else DEFAULT_REFRESH_TOKEN_EXPIRY_SECONDS
         )
+        self._token_expiry_threshold_seconds: int = token_expiry_threshold_seconds
 
         if jwt_signing_key is None:
             if upstream_client_secret is None:
@@ -1704,16 +1711,19 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
                 return None
             validated = await self._token_validator.verify_token(verification_token)
 
-            # 4. If upstream validation failed due to token expiry and we
-            # have a refresh token, attempt transparent refresh to avoid
-            # forcing the client into a full re-auth flow. Only refresh on
-            # expiry — other failures (scope mismatch, revocation) won't be
-            # helped by a refresh and would just burn tokens.
-            if (
+            # 4. Determine if refresh is needed. Two cases:
+            #    a) Validation failed and token is expired/within threshold
+            #    b) Validation passed but token is within threshold (proactive)
+            needs_refresh = upstream_token_set.refresh_token and (
+                upstream_token_set.expires_at
+                <= time.time() + self._token_expiry_threshold_seconds
+            )
+            should_refresh = needs_refresh and (
                 not validated
-                and upstream_token_set.refresh_token
-                and upstream_token_set.expires_at <= time.time()
-            ):
+                or (validated and self._token_expiry_threshold_seconds > 0)
+            )
+
+            if should_refresh:
                 try:
                     token_id = upstream_token_set.upstream_token_id
 
@@ -1738,12 +1748,11 @@ class OAuthProxy(OAuthProvider, ConsentMixin):
                             )
 
                         # Only refresh if the (possibly reloaded) token is
-                        # still expired — a non-expiry failure on a fresh
-                        # token (scope mismatch, revocation) won't be
-                        # helped by refreshing.
+                        # still within threshold — a freshly-refreshed token
+                        # doesn't need another refresh.
                         if (
-                            not validated
-                            and upstream_token_set.expires_at <= time.time()
+                            upstream_token_set.expires_at
+                            <= time.time() + self._token_expiry_threshold_seconds
                         ):
                             upstream_token_set = await self._try_transparent_refresh(
                                 upstream_token_set

fastmcp_slim/fastmcp/server/auth/oidc_proxy.py

@@ -234,6 +234,8 @@ class OIDCProxy(OAuthProxy):
         # Token expiry fallback
         fallback_access_token_expiry_seconds: int | None = None,
         fallback_refresh_token_expiry_seconds: int | None = None,
+        # Token refresh threshold
+        token_expiry_threshold_seconds: int = 0,
         # CIMD configuration
         enable_cimd: bool = True,
     ) -> None:
@@ -300,6 +302,9 @@ class OIDCProxy(OAuthProxy):
                 Defaults to 1 year. The actual upstream refresh remains the source of
                 truth — if upstream rejects the refresh, the client gets `invalid_grant`
                 and re-auths.
+            token_expiry_threshold_seconds: Number of seconds before actual expiry to consider
+                a token as expired (default 0). Prevents race conditions where a token
+                passes the expiry check but expires before the next operation completes.
             enable_cimd: Whether to enable CIMD (Client ID Metadata Document) client support.
                 When True, clients can use their metadata document URL as client_id instead of
                 Dynamic Client Registration. Default is True.
@@ -391,6 +396,7 @@ class OIDCProxy(OAuthProxy):
             "forward_resource": forward_resource,
             "fallback_access_token_expiry_seconds": fallback_access_token_expiry_seconds,
             "fallback_refresh_token_expiry_seconds": fallback_refresh_token_expiry_seconds,
+            "token_expiry_threshold_seconds": token_expiry_threshold_seconds,
             "enable_cimd": enable_cimd,
         }
 

fastmcp_slim/fastmcp/server/auth/providers/auth0.py

@@ -76,6 +76,7 @@ class Auth0Provider(OIDCProxy):
         consent_csp_policy: str | None = None,
         forward_resource: bool = True,
         fallback_refresh_token_expiry_seconds: int | None = None,
+        token_expiry_threshold_seconds: int = 0,
     ) -> None:
         """Initialize Auth0 OAuth provider.
 
@@ -128,6 +129,7 @@ class Auth0Provider(OIDCProxy):
             consent_csp_policy=consent_csp_policy,
             forward_resource=forward_resource,
             fallback_refresh_token_expiry_seconds=fallback_refresh_token_expiry_seconds,
+            token_expiry_threshold_seconds=token_expiry_threshold_seconds,
         )
 
         logger.debug(

fastmcp_slim/fastmcp/server/auth/providers/aws.py

@@ -137,6 +137,7 @@ class AWSCognitoProvider(OIDCProxy):
         consent_csp_policy: str | None = None,
         forward_resource: bool = True,
         fallback_refresh_token_expiry_seconds: int | None = None,
+        token_expiry_threshold_seconds: int = 0,
     ):
         """Initialize AWS Cognito OAuth provider.
 
@@ -198,6 +199,7 @@ class AWSCognitoProvider(OIDCProxy):
             consent_csp_policy=consent_csp_policy,
             forward_resource=forward_resource,
             fallback_refresh_token_expiry_seconds=fallback_refresh_token_expiry_seconds,
+            token_expiry_threshold_seconds=token_expiry_threshold_seconds,
         )
 
         logger.debug(

fastmcp_slim/fastmcp/server/auth/providers/azure.py

@@ -116,6 +116,7 @@ class AzureProvider(OAuthProxy):
         consent_csp_policy: str | None = None,
         forward_resource: bool = True,
         fallback_refresh_token_expiry_seconds: int | None = None,
+        token_expiry_threshold_seconds: int = 0,
         base_authority: str = "login.microsoftonline.com",
         token_issuer: str | None = None,
         http_client: httpx.AsyncClient | None = None,
@@ -262,6 +263,7 @@ class AzureProvider(OAuthProxy):
             consent_csp_policy=consent_csp_policy,
             forward_resource=forward_resource,
             fallback_refresh_token_expiry_seconds=fallback_refresh_token_expiry_seconds,
+            token_expiry_threshold_seconds=token_expiry_threshold_seconds,
             valid_scopes=parsed_required_scopes,
             enable_cimd=enable_cimd,
         )

fastmcp_slim/fastmcp/server/auth/providers/clerk.py

@@ -290,6 +290,7 @@ class ClerkProvider(OAuthProxy):
         consent_csp_policy: str | None = None,
         forward_resource: bool = True,
         fallback_refresh_token_expiry_seconds: int | None = None,
+        token_expiry_threshold_seconds: int = 0,
         extra_authorize_params: dict[str, str] | None = None,
         http_client: httpx.AsyncClient | None = None,
         enable_cimd: bool = True,
@@ -378,6 +379,7 @@ class ClerkProvider(OAuthProxy):
             consent_csp_policy=consent_csp_policy,
             forward_resource=forward_resource,
             fallback_refresh_token_expiry_seconds=fallback_refresh_token_expiry_seconds,
+            token_expiry_threshold_seconds=token_expiry_threshold_seconds,
             extra_authorize_params=extra_authorize_params_final or None,
             valid_scopes=parsed_valid_scopes,
             enable_cimd=enable_cimd,

fastmcp_slim/fastmcp/server/auth/providers/discord.py

@@ -208,6 +208,7 @@ class DiscordProvider(OAuthProxy):
         consent_csp_policy: str | None = None,
         forward_resource: bool = True,
         fallback_refresh_token_expiry_seconds: int | None = None,
+        token_expiry_threshold_seconds: int = 0,
         http_client: httpx.AsyncClient | None = None,
         enable_cimd: bool = True,
     ):
@@ -280,6 +281,7 @@ class DiscordProvider(OAuthProxy):
             consent_csp_policy=consent_csp_policy,
             forward_resource=forward_resource,
             fallback_refresh_token_expiry_seconds=fallback_refresh_token_expiry_seconds,
+            token_expiry_threshold_seconds=token_expiry_threshold_seconds,
             enable_cimd=enable_cimd,
         )
 

fastmcp_slim/fastmcp/server/auth/providers/github.py

@@ -223,6 +223,7 @@ class GitHubProvider(OAuthProxy):
         consent_csp_policy: str | None = None,
         forward_resource: bool = True,
         fallback_refresh_token_expiry_seconds: int | None = None,
+        token_expiry_threshold_seconds: int = 0,
         http_client: httpx.AsyncClient | None = None,
         enable_cimd: bool = True,
     ):
@@ -295,6 +296,7 @@ class GitHubProvider(OAuthProxy):
             consent_csp_policy=consent_csp_policy,
             forward_resource=forward_resource,
             fallback_refresh_token_expiry_seconds=fallback_refresh_token_expiry_seconds,
+            token_expiry_threshold_seconds=token_expiry_threshold_seconds,
             enable_cimd=enable_cimd,
         )
 

fastmcp_slim/fastmcp/server/auth/providers/google.py

@@ -248,6 +248,7 @@ class GoogleProvider(OAuthProxy):
         consent_csp_policy: str | None = None,
         forward_resource: bool = True,
         fallback_refresh_token_expiry_seconds: int | None = None,
+        token_expiry_threshold_seconds: int = 0,
         extra_authorize_params: dict[str, str] | None = None,
         http_client: httpx.AsyncClient | None = None,
         enable_cimd: bool = True,
@@ -355,6 +356,7 @@ class GoogleProvider(OAuthProxy):
             consent_csp_policy=consent_csp_policy,
             forward_resource=forward_resource,
             fallback_refresh_token_expiry_seconds=fallback_refresh_token_expiry_seconds,
+            token_expiry_threshold_seconds=token_expiry_threshold_seconds,
             extra_authorize_params=extra_authorize_params_final,
             valid_scopes=valid_scopes_final,
             enable_cimd=enable_cimd,

fastmcp_slim/fastmcp/server/auth/providers/oci.py

@@ -135,6 +135,7 @@ class OCIProvider(OIDCProxy):
         consent_csp_policy: str | None = None,
         forward_resource: bool = True,
         fallback_refresh_token_expiry_seconds: int | None = None,
+        token_expiry_threshold_seconds: int = 0,
     ) -> None:
         """Initialize OCI OIDC provider.
 
@@ -173,6 +174,7 @@ class OCIProvider(OIDCProxy):
             consent_csp_policy=consent_csp_policy,
             forward_resource=forward_resource,
             fallback_refresh_token_expiry_seconds=fallback_refresh_token_expiry_seconds,
+            token_expiry_threshold_seconds=token_expiry_threshold_seconds,
         )
 
         logger.debug(

fastmcp_slim/fastmcp/server/auth/providers/workos.py

@@ -177,6 +177,7 @@ class WorkOSProvider(OAuthProxy):
         consent_csp_policy: str | None = None,
         forward_resource: bool = True,
         fallback_refresh_token_expiry_seconds: int | None = None,
+        token_expiry_threshold_seconds: int = 0,
         extra_authorize_params: dict[str, str] | None = None,
         http_client: httpx.AsyncClient | None = None,
         enable_cimd: bool = True,
@@ -216,6 +217,9 @@ class WorkOSProvider(OAuthProxy):
             extra_authorize_params: Additional parameters to forward to WorkOS's authorization endpoint.
                 Useful for forcing scopes like `offline_access` so WorkOS issues a refresh token,
                 e.g. ``{"scope": "openid profile email offline_access"}``.
+            token_expiry_threshold_seconds: Number of seconds before actual expiry to consider
+                a token as expired (default 0). Prevents race conditions where a token
+                passes the expiry check but expires before the next operation completes.
             http_client: Optional httpx.AsyncClient for connection pooling in token verification.
                 When provided, the client is reused across verify_token calls and the caller
                 is responsible for its lifecycle. When None (default), a fresh client is created per call.
@@ -260,6 +264,7 @@ class WorkOSProvider(OAuthProxy):
             consent_csp_policy=consent_csp_policy,
             forward_resource=forward_resource,
             fallback_refresh_token_expiry_seconds=fallback_refresh_token_expiry_seconds,
+            token_expiry_threshold_seconds=token_expiry_threshold_seconds,
             extra_authorize_params=extra_authorize_params,
             valid_scopes=valid_scopes_final,
             enable_cimd=enable_cimd,

tests/server/auth/oauth_proxy/test_tokens.py

@@ -1269,3 +1269,262 @@ class TestTransparentUpstreamRefresh:
             returned = await mock_verifier.verify_token(call.args[0])
             if returned:
                 assert "upstream_claims" not in returned.claims
+
+    async def test_transparent_refresh_triggered_by_threshold(self, mock_verifier):
+        """Token within expiry threshold is treated as expired and refreshed."""
+        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(),
+            token_expiry_threshold_seconds=60,
+        )
+        proxy.set_mcp_path("/mcp")
+
+        upstream_token_id = "upstream-tok-threshold"
+        access_jti = "test-threshold-jti"
+
+        upstream_token_set = UpstreamTokenSet(
+            upstream_token_id=upstream_token_id,
+            access_token="almost-expired-upstream-access",
+            refresh_token="upstream-refresh-tok",
+            refresh_token_expires_at=time.time() + 86400,
+            expires_at=time.time() + 30,  # expires in 30s, within 60s threshold
+            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,
+        )
+
+        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_no_refresh_when_within_threshold_but_no_refresh_token(
+        self, mock_verifier
+    ):
+        """Token within threshold but without refresh token is not refreshed."""
+        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(),
+            token_expiry_threshold_seconds=60,
+        )
+        proxy.set_mcp_path("/mcp")
+
+        upstream_token_id = "upstream-tok-no-refresh"
+        access_jti = "test-no-refresh-jti"
+
+        upstream_token_set = UpstreamTokenSet(
+            upstream_token_id=upstream_token_id,
+            access_token="valid-upstream-access",
+            refresh_token=None,
+            refresh_token_expires_at=None,
+            expires_at=time.time() + 30,  # expires in 30s, within 60s threshold
+            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,
+        )
+
+        result = await proxy.load_access_token(fastmcp_jwt)
+
+        # Token validates via the verifier (starts with "valid-"), so even though
+        # it's within threshold, no refresh is attempted without a refresh token.
+        assert result is not None
+        assert result.token == "valid-upstream-access"
+
+    async def test_zero_threshold_only_refreshes_on_actual_expiry(self, mock_verifier):
+        """With threshold=0, token that's not yet expired is not refreshed."""
+        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(),
+            token_expiry_threshold_seconds=0,
+        )
+        proxy.set_mcp_path("/mcp")
+
+        upstream_token_id = "upstream-tok-zero-threshold"
+        access_jti = "test-zero-threshold-jti"
+
+        upstream_token_set = UpstreamTokenSet(
+            upstream_token_id=upstream_token_id,
+            access_token="valid-upstream-access",
+            refresh_token="upstream-refresh-tok",
+            refresh_token_expires_at=time.time() + 86400,
+            expires_at=time.time() + 30,  # expires in 30s, but threshold is 0
+            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,
+        )
+
+        result = await proxy.load_access_token(fastmcp_jwt)
+
+        # With threshold=0, token with 30s remaining is still valid
+        assert result is not None
+        assert result.token == "valid-upstream-access"
+
+    async def test_proactive_refresh_when_validated_but_within_threshold(
+        self, mock_verifier
+    ):
+        """Token that passes validation but is within threshold gets proactively refreshed."""
+        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(),
+            token_expiry_threshold_seconds=60,
+        )
+        proxy.set_mcp_path("/mcp")
+
+        upstream_token_id = "upstream-tok-proactive"
+        access_jti = "test-proactive-jti"
+
+        # Token starts with "valid-" so the mock verifier accepts it,
+        # but it expires in 30s which is within the 60s threshold.
+        upstream_token_set = UpstreamTokenSet(
+            upstream_token_id=upstream_token_id,
+            access_token="valid-but-near-expiry",
+            refresh_token="upstream-refresh-tok",
+            refresh_token_expires_at=time.time() + 86400,
+            expires_at=time.time() + 30,
+            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,
+        )
+
+        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()