fix: avoid uncessary retries in OAuth authenticated requests (#1206)

Co-authored-by: Felix Weinberger <fweinberger@anthropic.com>
This commit is contained in:
keurcien 2025-08-22 12:45:54 +02:00 • committed by GitHub
commit e750a06a99
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 66 additions and 4 deletions

View file

@ -546,6 +546,6 @@ class OAuthClientProvider(httpx.Auth):
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

View file

@ -361,7 +361,19 @@ class TestOAuthFallback:
),
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):
@ -694,11 +706,61 @@ class TestAuthFlow:
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
# 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(
(