Skip to content

Fix: avoid uncessary retries in OAuth authenticated requests #1206

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Aug 22, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions src/mcp/client/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -546,6 +546,6 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
logger.exception("OAuth flow error")
raise

# Retry with new tokens
self._add_auth_header(request)
yield request
# Retry with new tokens
self._add_auth_header(request)
yield request
64 changes: 63 additions & 1 deletion tests/client/test_auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -361,7 +361,19 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl
),
request=token_request,
)
token_request = await auth_flow.asend(token_response)

# After OAuth flow completes, the original request is retried with auth header
final_request = await auth_flow.asend(token_response)
assert final_request.headers["Authorization"] == "Bearer new_access_token"
assert final_request.method == "GET"
assert str(final_request.url) == "https://api.example.com/v1/mcp"

# Send final success response to properly close the generator
final_response = httpx.Response(200, request=final_request)
try:
await auth_flow.asend(final_response)
except StopAsyncIteration:
pass # Expected - generator should complete

@pytest.mark.anyio
async def test_handle_metadata_response_success(self, oauth_provider: OAuthClientProvider):
Expand Down Expand Up @@ -694,11 +706,61 @@ async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvide
assert final_request.method == "GET"
assert str(final_request.url) == "https://api.example.com/mcp"

# Send final success response to properly close the generator
final_response = httpx.Response(200, request=final_request)
try:
await auth_flow.asend(final_response)
except StopAsyncIteration:
pass # Expected - generator should complete

Copy link
Contributor

@felixweinberger felixweinberger Aug 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These tests were inadvertently tuned to the buggy behavior where requests were being unconditionally retried after every response. The indentation fix ensures retries only happen after OAuth completion (401 responses), not after successful responses. This exposed that the tests weren't properly closing the auth flow generator after the legitimate OAuth retry, causing lock release errors.

# Verify tokens were stored
assert oauth_provider.context.current_tokens is not None
assert oauth_provider.context.current_tokens.access_token == "new_access_token"
assert oauth_provider.context.token_expiry_time is not None

@pytest.mark.anyio
async def test_auth_flow_no_unnecessary_retry_after_oauth(
self, oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken
):
"""Test that requests are not retried unnecessarily - the core bug that caused 2x performance degradation."""
# Pre-store valid tokens so no OAuth flow is needed
await mock_storage.set_tokens(valid_tokens)
oauth_provider.context.current_tokens = valid_tokens
oauth_provider.context.token_expiry_time = time.time() + 1800
oauth_provider._initialized = True

test_request = httpx.Request("GET", "https://api.example.com/mcp")
auth_flow = oauth_provider.async_auth_flow(test_request)

# Count how many times the request is yielded
request_yields = 0

# First request - should have auth header already
request = await auth_flow.__anext__()
request_yields += 1
assert request.headers["Authorization"] == "Bearer test_access_token"

# Send a successful 200 response
response = httpx.Response(200, request=request)

# In the buggy version, this would yield the request AGAIN unconditionally
# In the fixed version, this should end the generator
try:
await auth_flow.asend(response) # extra request
request_yields += 1
# If we reach here, the bug is present
pytest.fail(
f"Unnecessary retry detected! Request was yielded {request_yields} times. "
f"This indicates the retry logic bug that caused 2x performance degradation. "
f"The request should only be yielded once for successful responses."
)
except StopAsyncIteration:
# This is the expected behavior - no unnecessary retry
pass

# Verify exactly one request was yielded (no double-sending)
assert request_yields == 1, f"Expected 1 request yield, got {request_yields}"


@pytest.mark.parametrize(
(
Expand Down
Loading