Skip to content
Merged
Prev Previous commit
Try refresh token first
  • Loading branch information
yurikunash committed Jul 15, 2025
commit 2c1c35d06fc7ae7a05bf74c9fcf41c36d3a947bd
91 changes: 44 additions & 47 deletions src/mcp/client/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down