fix(client/auth): finalize the discovery sub-generator when the flow is closed mid-discovery

httpx2's _send_handling_auth acloses the auth flow on any transport error or
cancellation, throwing GeneratorExit at the relay's yield — caught by neither
except StopAsyncIteration nor except Exception — so both hand-relay sites (the
401 branch and the 403 step-up) abandoned the suspended
_discover_authorization_server_metadata sub-generator to GC (a ResourceWarning
on trio, which filterwarnings=["error"] turns into a test failure for any
future abort-mid-discovery test).

Both relays now drive the sub-generator under contextlib.aclosing, matching the
eager-refresh relay in #3263, so closing the outer flow finalizes it in the
same unwind. Covered by two regression tests that close the flow mid-discovery
and assert the captured sub-generator reports exhaustion instead of a live
suspended frame.

Author: claude

src/mcp/client/auth/oauth2.py

@@ -10,6 +10,7 @@
 import string
 import time
 from collections.abc import AsyncGenerator, Awaitable, Callable
+from contextlib import aclosing
 from dataclasses import dataclass, field
 from typing import Any, Protocol, get_args
 from urllib.parse import quote, urlencode, urljoin, urlparse
@@ -812,15 +813,17 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
                     # metadata, applying the SEP-2352 issuer checks along the way. The
                     # sequence lives in a sub-generator shared with the 403 step-up;
                     # its requests are relayed by hand (`yield from` cannot cross an
-                    # async generator).
-                    discovery = self._discover_authorization_server_metadata(response)
-                    try:
-                        discovery_request = await anext(discovery)
-                        while True:
-                            discovery_response = yield discovery_request
-                            discovery_request = await discovery.asend(discovery_response)
-                    except StopAsyncIteration:
-                        pass
+                    # async generator). `aclosing` finalizes the sub-generator when
+                    # httpx2 closes this flow mid-discovery (transport error or
+                    # cancellation throws `GeneratorExit` at the relay's `yield`).
+                    async with aclosing(self._discover_authorization_server_metadata(response)) as discovery:
+                        try:
+                            discovery_request = await anext(discovery)
+                            while True:
+                                discovery_response = yield discovery_request
+                                discovery_request = await discovery.asend(discovery_response)
+                        except StopAsyncIteration:
+                            pass
 
                     # Step 3: Apply scope selection strategy
                     self.context.client_metadata.scope = get_client_metadata_scopes(
@@ -884,14 +887,16 @@ async def async_auth_flow(self, request: httpx2.Request) -> AsyncGenerator[httpx
                         if self.context.oauth_metadata is None and (
                             self.context.registration_secret_expired() or self.context.client_info is None
                         ):
-                            discovery = self._discover_authorization_server_metadata(response)
-                            try:
-                                discovery_request = await anext(discovery)
-                                while True:
-                                    discovery_response = yield discovery_request
-                                    discovery_request = await discovery.asend(discovery_response)
-                            except StopAsyncIteration:
-                                pass
+                            # `aclosing` mirrors the 401 relay above: it finalizes the
+                            # sub-generator when httpx2 closes this flow mid-discovery.
+                            async with aclosing(self._discover_authorization_server_metadata(response)) as discovery:
+                                try:
+                                    discovery_request = await anext(discovery)
+                                    while True:
+                                        discovery_response = yield discovery_request
+                                        discovery_request = await discovery.asend(discovery_response)
+                                except StopAsyncIteration:
+                                    pass
 
                         # Step 2a: Union previously requested scopes with the newly challenged
                         # scopes (SEP-2350) so escalating one operation keeps the others' grants.

tests/client/test_auth.py

@@ -3,6 +3,7 @@
 import base64
 import json
 import time
+from collections.abc import AsyncGenerator
 from unittest import mock
 from urllib.parse import parse_qs, quote, unquote, urlparse
 
@@ -4031,3 +4032,85 @@ async def mock_callback() -> AuthorizationCodeResult:
         await auth_flow.asend(httpx2.Response(200, request=final_request))
     except StopAsyncIteration:
         pass
+
+
+@pytest.mark.anyio
+async def test_closing_the_flow_mid_discovery_finalizes_the_401_discovery_sub_generator(
+    oauth_provider: OAuthClientProvider,
+):
+    """httpx2's `_send_handling_auth` acloses the auth flow when a discovery request
+    fails at the transport level (or the request is cancelled), throwing `GeneratorExit`
+    at the 401 relay's `yield`; the relay must finalize the discovery sub-generator with
+    the flow rather than abandon it suspended to GC (a `ResourceWarning` under trio).
+    The sub-generator is captured via a wrapper because its finalization is not
+    observable through the auth-flow protocol itself.
+    """
+    captured: list[AsyncGenerator[httpx2.Request, httpx2.Response]] = []
+    original = oauth_provider._discover_authorization_server_metadata
+
+    def capturing(response: httpx2.Response) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
+        discovery = original(response)
+        captured.append(discovery)
+        return discovery
+
+    oauth_provider._discover_authorization_server_metadata = capturing
+    oauth_provider._initialized = True
+
+    auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
+    request = await auth_flow.__anext__()
+
+    # 401 → the flow yields the first discovery request, suspending both generators mid-relay.
+    prm_req = await auth_flow.asend(httpx2.Response(401, request=request))
+    assert "oauth-protected-resource" in str(prm_req.url)
+    assert len(captured) == 1
+
+    # Closing the flow here mirrors httpx2 aborting it on a transport error.
+    await auth_flow.aclose()
+
+    # The sub-generator was closed with the flow: probing it reports exhaustion instead
+    # of resuming a suspended frame.
+    with pytest.raises(StopAsyncIteration):
+        await captured[0].__anext__()
+
+
+@pytest.mark.anyio
+async def test_closing_the_flow_mid_discovery_finalizes_the_403_step_up_discovery_sub_generator(
+    oauth_provider: OAuthClientProvider,
+):
+    """The 403 step-up's discovery relay must finalize its sub-generator when httpx2
+    acloses the flow mid-discovery, exactly like the 401 relay (same abandonment
+    hazard, same fix); the sub-generator is captured via a wrapper because its
+    finalization is not observable through the auth-flow protocol itself.
+    """
+    captured: list[AsyncGenerator[httpx2.Request, httpx2.Response]] = []
+    original = oauth_provider._discover_authorization_server_metadata
+
+    def capturing(response: httpx2.Response) -> AsyncGenerator[httpx2.Request, httpx2.Response]:
+        discovery = original(response)
+        captured.append(discovery)
+        return discovery
+
+    oauth_provider._discover_authorization_server_metadata = capturing
+    # Restart shape: a live token and no stored registration, so the 403 step-up must
+    # discover before registering (`oauth_metadata` is never restored by `_initialize`).
+    oauth_provider.context.current_tokens = OAuthToken(access_token="live-token")
+    oauth_provider._initialized = True
+
+    auth_flow = oauth_provider.async_auth_flow(httpx2.Request("GET", "https://api.example.com/v1/mcp"))
+    request = await auth_flow.__anext__()
+
+    response_403 = httpx2.Response(
+        403,
+        headers={"WWW-Authenticate": 'Bearer error="insufficient_scope", scope="write"'},
+        request=request,
+    )
+
+    # 403 → the step-up yields the first discovery request, suspending both generators mid-relay.
+    prm_req = await auth_flow.asend(response_403)
+    assert "oauth-protected-resource" in str(prm_req.url)
+    assert len(captured) == 1
+
+    await auth_flow.aclose()
+
+    with pytest.raises(StopAsyncIteration):
+        await captured[0].__anext__()