From 8445068313d18146e00ea1ad67d76e1bd66d1640 Mon Sep 17 00:00:00 2001 From: keurcien Date: Sun, 27 Jul 2025 17:21:15 +0200 Subject: [PATCH 1/6] Avoid retry everytime --- src/mcp/client/auth.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/mcp/client/auth.py b/src/mcp/client/auth.py index 775fb0f6c..17d2558f7 100644 --- a/src/mcp/client/auth.py +++ b/src/mcp/client/auth.py @@ -131,7 +131,7 @@ def is_token_valid(self) -> bool: and self.current_tokens.access_token and (not self.token_expiry_time or time.time() <= self.token_expiry_time) ) - + def can_refresh_token(self) -> bool: """Check if token can be refreshed.""" return bool(self.current_tokens and self.current_tokens.refresh_token and self.client_info) @@ -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 From b2ec4594f11e636521ab18a01c9049f4c2eba555 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Thu, 21 Aug 2025 16:12:40 +0100 Subject: [PATCH 2/6] Add proper regression test for OAuth retry performance bug This test validates the fix for PR #1206 which resolved a critical performance issue where every request was being sent twice due to incorrect indentation of retry logic. The test: - Simulates a successful authenticated request (200 response) - Counts how many times the request is yielded - FAILS with buggy version: 2 yields (double sending) - PASSES with fix: 1 yield (correct behavior) This directly tests the root cause that made OAuth requests 2x slower and ensures this regression cannot occur again. --- tests/client/test_auth.py | 41 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index 61d74df1e..df1c2549b 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -699,6 +699,47 @@ async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvide 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, mock_storage, valid_tokens): + """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: + extra_request = await auth_flow.asend(response) + 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( ( From 2c6c9091fcbfc60e1af5c2fdf8af977df337d1c9 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Thu, 21 Aug 2025 16:35:59 +0100 Subject: [PATCH 3/6] lint: fix formatting --- src/mcp/client/auth.py | 2 +- tests/client/test_auth.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mcp/client/auth.py b/src/mcp/client/auth.py index 17d2558f7..376036e8c 100644 --- a/src/mcp/client/auth.py +++ b/src/mcp/client/auth.py @@ -131,7 +131,7 @@ def is_token_valid(self) -> bool: and self.current_tokens.access_token and (not self.token_expiry_time or time.time() <= self.token_expiry_time) ) - + def can_refresh_token(self) -> bool: """Check if token can be refreshed.""" return bool(self.current_tokens and self.current_tokens.refresh_token and self.client_info) diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index df1c2549b..d64885abc 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -721,11 +721,11 @@ async def test_auth_flow_no_unnecessary_retry_after_oauth(self, oauth_provider, # 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: - extra_request = await auth_flow.asend(response) + await auth_flow.asend(response) # extra request request_yields += 1 # If we reach here, the bug is present pytest.fail( From 5f371588c840c999d596e510f001c73a0459f7c4 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Thu, 21 Aug 2025 16:43:54 +0100 Subject: [PATCH 4/6] Fix type annotations for regression test Add missing type hints to test_auth_flow_no_unnecessary_retry_after_oauth to satisfy pyright type checking requirements. --- tests/client/test_auth.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index d64885abc..4b70a6395 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -700,7 +700,9 @@ async def test_auth_flow_with_no_tokens(self, oauth_provider: OAuthClientProvide 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, mock_storage, valid_tokens): + 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) From 3b8e7b66e8ad7830ec41564216503b2a72e25ea3 Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Thu, 21 Aug 2025 16:50:44 +0100 Subject: [PATCH 5/6] Fix OAuth auth flow tests to properly close async generators The retry logic fix in PR-1206 moved the final request yield inside the OAuth flow block, which now correctly retries requests after OAuth completion. However, the existing tests didn't expect this final yield and would exit the generator early, causing a GeneratorExit exception while the async lock was still held. This resulted in RuntimeError: The current task is not holding this lock. Fix by properly closing the generators in both failing tests: - test_oauth_discovery_fallback_conditions - test_auth_flow_with_no_tokens Both tests now send a final success response to properly complete the auth flow generator and prevent the lock release error. --- tests/client/test_auth.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index 4b70a6395..31c5f8c36 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -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): @@ -694,6 +706,13 @@ 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 + # Verify tokens were stored assert oauth_provider.context.current_tokens is not None assert oauth_provider.context.current_tokens.access_token == "new_access_token" From b81bf9aac706cfe5b078d32da4b6d4e66a0443ab Mon Sep 17 00:00:00 2001 From: Felix Weinberger Date: Thu, 21 Aug 2025 16:57:35 +0100 Subject: [PATCH 6/6] lint: fix lints --- tests/client/test_auth.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/client/test_auth.py b/tests/client/test_auth.py index 31c5f8c36..6e58e496d 100644 --- a/tests/client/test_auth.py +++ b/tests/client/test_auth.py @@ -361,13 +361,13 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl ), request=token_request, ) - + # 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: