Try refresh token first

Author: yurikunash

src/mcp/client/auth.py

@@ -520,59 +520,56 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
             # Capture protocol version from request headers
             self.context.protocol_version = request.headers.get(MCP_PROTOCOL_VERSION)
 
+            if not self.context.is_token_valid() and self.context.can_refresh_token():
+                # Try to refresh token
+                refresh_request = await self._refresh_token()
+                refresh_response = yield refresh_request
+
+                if not await self._handle_refresh_response(refresh_response):
+                    # Refresh failed, need full re-authentication
+                    self._initialized = False
+
             if self.context.is_token_valid():
                 self._add_auth_header(request)
 
             response = yield request
 
             if response.status_code == 401:
-                if self.context.can_refresh_token():
-                    # Try to refresh token
-                    refresh_request = await self._refresh_token()
-                    refresh_response = yield refresh_request
-
-                    if not await self._handle_refresh_response(refresh_response):
-                        # Refresh failed, need full re-authentication
-                        self._initialized = False
-                else:
-                    self.context.clear_tokens()
-
-                # If we don't have valid tokens after refresh, perform OAuth flow
-                if not self.context.is_token_valid():
-                    try:
-                        # OAuth flow must be inline due to generator constraints
-                        # Step 1: Discover protected resource metadata (RFC9728 with WWW-Authenticate support)
-                        discovery_request = await self._discover_protected_resource(response)
-                        discovery_response = yield discovery_request
-                        await self._handle_protected_resource_response(discovery_response)
-
-                        # Step 2: Discover OAuth metadata (with fallback for legacy servers)
-                        oauth_request = await self._discover_oauth_metadata()
-                        oauth_response = yield oauth_request
-                        handled = await self._handle_oauth_metadata_response(oauth_response, is_fallback=False)
-
-                        # If path-aware discovery failed with 404, try fallback to root
-                        if not handled:
-                            fallback_request = await self._discover_oauth_metadata_fallback()
-                            fallback_response = yield fallback_request
-                            await self._handle_oauth_metadata_response(fallback_response, is_fallback=True)
-
-                        # Step 3: Register client if needed
-                        registration_request = await self._register_client()
-                        if registration_request:
-                            registration_response = yield registration_request
-                            await self._handle_registration_response(registration_response)
-
-                        # Step 4: Perform authorization
-                        auth_code, code_verifier = await self._perform_authorization()
-
-                        # Step 5: Exchange authorization code for tokens
-                        token_request = await self._exchange_token(auth_code, code_verifier)
-                        token_response = yield token_request
-                        await self._handle_token_response(token_response)
-                    except Exception:
-                        logger.exception("OAuth flow error")
-                        raise
+                # Perform full OAuth flow
+                try:
+                    # OAuth flow must be inline due to generator constraints
+                    # Step 1: Discover protected resource metadata (RFC9728 with WWW-Authenticate support)
+                    discovery_request = await self._discover_protected_resource(response)
+                    discovery_response = yield discovery_request
+                    await self._handle_protected_resource_response(discovery_response)
+
+                    # Step 2: Discover OAuth metadata (with fallback for legacy servers)
+                    oauth_request = await self._discover_oauth_metadata()
+                    oauth_response = yield oauth_request
+                    handled = await self._handle_oauth_metadata_response(oauth_response, is_fallback=False)
+
+                    # If path-aware discovery failed with 404, try fallback to root
+                    if not handled:
+                        fallback_request = await self._discover_oauth_metadata_fallback()
+                        fallback_response = yield fallback_request
+                        await self._handle_oauth_metadata_response(fallback_response, is_fallback=True)
+
+                    # Step 3: Register client if needed
+                    registration_request = await self._register_client()
+                    if registration_request:
+                        registration_response = yield registration_request
+                        await self._handle_registration_response(registration_response)
+
+                    # Step 4: Perform authorization
+                    auth_code, code_verifier = await self._perform_authorization()
+
+                    # Step 5: Exchange authorization code for tokens
+                    token_request = await self._exchange_token(auth_code, code_verifier)
+                    token_response = yield token_request
+                    await self._handle_token_response(token_response)
+                except Exception:
+                    logger.exception("OAuth flow error")
+                    raise
 
                 # Retry with new tokens
                 self._add_auth_header(request)