Skip to content

Commit e750a06

Browse files
fix: avoid uncessary retries in OAuth authenticated requests (#1206)
Co-authored-by: Felix Weinberger <fweinberger@anthropic.com>
1 parent f4b2957 commit e750a06

File tree

2 files changed

+66
-4
lines changed

2 files changed

+66
-4
lines changed

src/mcp/client/auth.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -546,6 +546,6 @@ async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.
546546
logger.exception("OAuth flow error")
547547
raise
548548

549-
# Retry with new tokens
550-
self._add_auth_header(request)
551-
yield request
549+
# Retry with new tokens
550+
self._add_auth_header(request)
551+
yield request

tests/client/test_auth.py

Lines changed: 63 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -361,7 +361,19 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl
361361
),
362362
request=token_request,
363363
)
364-
token_request = await auth_flow.asend(token_response)
364+
365+
# After OAuth flow completes, the original request is retried with auth header
366+
final_request = await auth_flow.asend(token_response)
367+
assert final_request.headers["Authorization"] == "Bearer new_access_token"
368+
assert final_request.method == "GET"
369+
assert str(final_request.url) == "https://api.example.com/v1/mcp"
370+
371+
# Send final success response to properly close the generator
372+
final_response = httpx.Response(200, request=final_request)
373+
try:
374+
await auth_flow.asend(final_response)
375+
except StopAsyncIteration:
376+
pass # Expected - generator should complete
365377

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

709+
# Send final success response to properly close the generator
710+
final_response = httpx.Response(200, request=final_request)
711+
try:
712+
await auth_flow.asend(final_response)
713+
except StopAsyncIteration:
714+
pass # Expected - generator should complete
715+
697716
# Verify tokens were stored
698717
assert oauth_provider.context.current_tokens is not None
699718
assert oauth_provider.context.current_tokens.access_token == "new_access_token"
700719
assert oauth_provider.context.token_expiry_time is not None
701720

721+
@pytest.mark.anyio
722+
async def test_auth_flow_no_unnecessary_retry_after_oauth(
723+
self, oauth_provider: OAuthClientProvider, mock_storage: MockTokenStorage, valid_tokens: OAuthToken
724+
):
725+
"""Test that requests are not retried unnecessarily - the core bug that caused 2x performance degradation."""
726+
# Pre-store valid tokens so no OAuth flow is needed
727+
await mock_storage.set_tokens(valid_tokens)
728+
oauth_provider.context.current_tokens = valid_tokens
729+
oauth_provider.context.token_expiry_time = time.time() + 1800
730+
oauth_provider._initialized = True
731+
732+
test_request = httpx.Request("GET", "https://api.example.com/mcp")
733+
auth_flow = oauth_provider.async_auth_flow(test_request)
734+
735+
# Count how many times the request is yielded
736+
request_yields = 0
737+
738+
# First request - should have auth header already
739+
request = await auth_flow.__anext__()
740+
request_yields += 1
741+
assert request.headers["Authorization"] == "Bearer test_access_token"
742+
743+
# Send a successful 200 response
744+
response = httpx.Response(200, request=request)
745+
746+
# In the buggy version, this would yield the request AGAIN unconditionally
747+
# In the fixed version, this should end the generator
748+
try:
749+
await auth_flow.asend(response) # extra request
750+
request_yields += 1
751+
# If we reach here, the bug is present
752+
pytest.fail(
753+
f"Unnecessary retry detected! Request was yielded {request_yields} times. "
754+
f"This indicates the retry logic bug that caused 2x performance degradation. "
755+
f"The request should only be yielded once for successful responses."
756+
)
757+
except StopAsyncIteration:
758+
# This is the expected behavior - no unnecessary retry
759+
pass
760+
761+
# Verify exactly one request was yielded (no double-sending)
762+
assert request_yields == 1, f"Expected 1 request yield, got {request_yields}"
763+
702764

703765
@pytest.mark.parametrize(
704766
(

0 commit comments

Comments
 (0)