Fixed uses of async generators and removed the pytest warning ignore

Author: agronholm

pyproject.toml

@@ -126,5 +126,4 @@ filterwarnings = [
     "ignore:Returning str or bytes.*:DeprecationWarning:mcp.server.lowlevel",
     # This is to avoid test failures on Trio due to httpx's failure to explicitly close
     # async generators
-    "ignore::pytest.PytestUnraisableExceptionWarning"
 ]

src/mcp/client/streamable_http.py

@@ -11,7 +11,6 @@
 from contextlib import aclosing, asynccontextmanager
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import cast
 
 import anyio
 import httpx
@@ -241,15 +240,16 @@ async def _handle_resumption_request(self, ctx: RequestContext) -> None:
             event_source.response.raise_for_status()
             logger.debug("Resumption GET SSE connection established")
 
-            async for sse in event_source.aiter_sse():
-                is_complete = await self._handle_sse_event(
-                    sse,
-                    ctx.read_stream_writer,
-                    original_request_id,
-                    ctx.metadata.on_resumption_token_update if ctx.metadata else None,
-                )
-                if is_complete:
-                    break
+            async with aclosing(event_source.aiter_sse()) as iterator:
+                async for sse in iterator:
+                    is_complete = await self._handle_sse_event(
+                        sse,
+                        ctx.read_stream_writer,
+                        original_request_id,
+                        ctx.metadata.on_resumption_token_update if ctx.metadata else None,
+                    )
+                    if is_complete:
+                        break
 
     async def _handle_post_request(self, ctx: RequestContext) -> None:
         """Handle a POST request with response processing."""
@@ -320,9 +320,7 @@ async def _handle_sse_response(
     ) -> None:
         """Handle SSE response from the server."""
         try:
-            event_source = EventSource(response)
-            sse_iter = cast(AsyncGenerator[ServerSentEvent], event_source.aiter_sse())
-            async with aclosing(sse_iter) as items:
+            async with aclosing(EventSource(response).aiter_sse()) as items:
                 async for sse in items:
                     is_complete = await self._handle_sse_event(
                         sse,

tests/client/test_auth.py

@@ -5,6 +5,7 @@
 import base64
 import hashlib
 import time
+from contextlib import aclosing
 from unittest.mock import AsyncMock, Mock, patch
 from urllib.parse import parse_qs, urlparse
 
@@ -654,17 +655,17 @@ async def test_async_auth_flow_401_response(self, oauth_provider, oauth_token):
         mock_response = Mock()
         mock_response.status_code = 401
 
-        auth_flow = oauth_provider.async_auth_flow(request)
-        await auth_flow.__anext__()
+        async with aclosing(oauth_provider.async_auth_flow(request)) as auth_flow:
+            await auth_flow.__anext__()
 
-        # Send 401 response
-        try:
-            await auth_flow.asend(mock_response)
-        except StopAsyncIteration:
-            pass
+            # Send 401 response
+            try:
+                await auth_flow.asend(mock_response)
+            except StopAsyncIteration:
+                pass
 
-        # Should clear current tokens
-        assert oauth_provider._current_tokens is None
+            # Should clear current tokens
+            assert oauth_provider._current_tokens is None
 
     @pytest.mark.anyio
     async def test_async_auth_flow_no_token(self, oauth_provider):
@@ -675,14 +676,14 @@ async def test_async_auth_flow_no_token(self, oauth_provider):
             patch.object(oauth_provider, "initialize") as mock_init,
             patch.object(oauth_provider, "ensure_token") as mock_ensure,
         ):
-            auth_flow = oauth_provider.async_auth_flow(request)
-            updated_request = await auth_flow.__anext__()
+            async with aclosing(oauth_provider.async_auth_flow(request)) as auth_flow:
+                updated_request = await auth_flow.__anext__()
 
-            mock_init.assert_called_once()
-            mock_ensure.assert_called_once()
+                mock_init.assert_called_once()
+                mock_ensure.assert_called_once()
 
-            # No Authorization header should be added if no token
-            assert "Authorization" not in updated_request.headers
+                # No Authorization header should be added if no token
+                assert "Authorization" not in updated_request.headers
 
     @pytest.mark.anyio
     async def test_scope_priority_client_metadata_first(self, oauth_provider, oauth_client_info):

tests/shared/test_sse.py

@@ -3,6 +3,7 @@
 import socket
 import time
 from collections.abc import AsyncGenerator, Generator
+from contextlib import aclosing
 
 import anyio
 import httpx
@@ -160,14 +161,15 @@ async def connection_test() -> None:
                 assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
 
                 line_number = 0
-                async for line in response.aiter_lines():
-                    if line_number == 0:
-                        assert line == "event: endpoint"
-                    elif line_number == 1:
-                        assert line.startswith("data: /messages/?session_id=")
-                    else:
-                        return
-                    line_number += 1
+                async with aclosing(response.aiter_lines()) as lines:
+                    async for line in lines:
+                        if line_number == 0:
+                            assert line == "event: endpoint"
+                        elif line_number == 1:
+                            assert line.startswith("data: /messages/?session_id=")
+                        else:
+                            return
+                        line_number += 1
 
         # Add timeout to prevent test from hanging if it fails
         with anyio.fail_after(3):