@@ -361,7 +361,19 @@ async def test_oauth_discovery_fallback_conditions(self, oauth_provider: OAuthCl
361
361
),
362
362
request = token_request ,
363
363
)
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
365
377
366
378
@pytest .mark .anyio
367
379
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
694
706
assert final_request .method == "GET"
695
707
assert str (final_request .url ) == "https://api.example.com/mcp"
696
708
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
+
697
716
# Verify tokens were stored
698
717
assert oauth_provider .context .current_tokens is not None
699
718
assert oauth_provider .context .current_tokens .access_token == "new_access_token"
700
719
assert oauth_provider .context .token_expiry_time is not None
701
720
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
+
702
764
703
765
@pytest .mark .parametrize (
704
766
(
0 commit comments