Merge pull request #711 from cbcoutinho/fix/mcp-client-session-cancel-scope

fix(tests): convert create_mcp_client_session to asynccontextmanager

Author: cbcoutinho

tests/conftest.py

@@ -8,8 +8,9 @@
 import threading
 import time
 import uuid
+from contextlib import asynccontextmanager
 from http.server import BaseHTTPRequestHandler, HTTPServer
-from typing import Any, AsyncGenerator
+from typing import Any, AsyncGenerator, AsyncIterator
 from urllib.parse import parse_qs, quote, urlparse
 
 import anyio
@@ -118,14 +119,15 @@ async def wait_for_nextcloud(
     return False
 
 
+@asynccontextmanager
 async def create_mcp_client_session(
     url: str,
     token: str | None = None,
     client_name: str = "MCP",
     elicitation_callback: Any = None,
     sampling_callback: Any = None,
     headers: dict[str, str] | None = None,
-) -> AsyncGenerator[ClientSession, Any]:
+) -> AsyncIterator[ClientSession]:
     """
     Factory function to create an MCP client session with proper lifecycle management.
 
@@ -227,10 +229,10 @@ async def nc_mcp_client(anyio_backend) -> AsyncGenerator[ClientSession, Any]:
 
     Uses anyio pytest plugin for proper async fixture handling.
     """
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8000/mcp",
         client_name="Basic MCP (HTTP)",
-    ):
+    ) as session:
         yield session
 
 
@@ -246,11 +248,11 @@ async def nc_mcp_oauth_client(
     Uses headless browser automation suitable for CI/CD.
     Uses anyio pytest plugin for proper async fixture handling.
     """
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8001/mcp",
         token=playwright_oauth_token,
         client_name="OAuth MCP (Playwright)",
-    ):
+    ) as session:
         yield session
 
 
@@ -271,11 +273,11 @@ async def nc_mcp_basic_auth_client(
     credentials = base64.b64encode(b"admin:admin").decode("utf-8")
     auth_header = f"Basic {credentials}"
 
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8003/mcp",
         headers={"Authorization": auth_header},
         client_name="BasicAuth MCP (Multi-User)",
-    ):
+    ) as session:
         yield session
 
 
@@ -296,11 +298,11 @@ async def nc_mcp_oauth_jwt_client(
     Uses headless browser automation suitable for CI/CD.
     Uses anyio pytest plugin for proper async fixture handling.
     """
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8001/mcp",
         token=playwright_oauth_token_jwt,
         client_name="OAuth JWT MCP (Playwright)",
-    ):
+    ) as session:
         yield session
 
 
@@ -456,12 +458,12 @@ async def elicitation_callback(
             await page.close()
 
     # Create client session with elicitation callback
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8001/mcp",
         token=playwright_oauth_token,
         client_name="OAuth MCP with Elicitation",
         elicitation_callback=elicitation_callback,
-    ):
+    ) as session:
         # Attach elicitation metadata for test validation
         session.elicitation_triggered = elicitation_triggered
         yield session
@@ -482,11 +484,11 @@ async def nc_mcp_oauth_client_read_only(
     Uses JWT tokens because they embed scope information in claims,
     enabling proper scope-based tool filtering.
     """
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8001/mcp",
         token=playwright_oauth_token_read_only,
         client_name="OAuth JWT MCP Read-Only (Playwright)",
-    ):
+    ) as session:
         yield session
 
 
@@ -505,11 +507,11 @@ async def nc_mcp_oauth_client_write_only(
     Uses JWT tokens because they embed scope information in claims,
     enabling proper scope-based tool filtering.
     """
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8001/mcp",
         token=playwright_oauth_token_write_only,
         client_name="OAuth JWT MCP Write-Only (Playwright)",
-    ):
+    ) as session:
         yield session
 
 
@@ -527,11 +529,11 @@ async def nc_mcp_oauth_client_full_access(
     Uses JWT tokens because they embed scope information in claims,
     enabling proper scope-based tool filtering.
     """
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8001/mcp",
         token=playwright_oauth_token_full_access,
         client_name="OAuth JWT MCP Full Access (Playwright)",
-    ):
+    ) as session:
         yield session
 
 
@@ -552,11 +554,11 @@ async def nc_mcp_oauth_client_no_custom_scopes(
     Uses JWT tokens because they embed scope information in claims,
     enabling proper scope-based tool filtering.
     """
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8001/mcp",
         token=playwright_oauth_token_no_custom_scopes,
         client_name="OAuth JWT MCP No Custom Scopes (Playwright)",
-    ):
+    ) as session:
         yield session
 
 
@@ -2726,11 +2728,11 @@ async def alice_mcp_client(
     alice_oauth_token: str,
 ) -> AsyncGenerator[ClientSession, Any]:
     """MCP client authenticated as alice (owner role)."""
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8001/mcp",
         token=alice_oauth_token,
         client_name="Alice MCP",
-    ):
+    ) as session:
         yield session
 
 
@@ -2739,11 +2741,11 @@ async def bob_mcp_client(
     anyio_backend, bob_oauth_token: str
 ) -> AsyncGenerator[ClientSession, Any]:
     """MCP client authenticated as bob (viewer role)."""
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8001/mcp",
         token=bob_oauth_token,
         client_name="Bob MCP",
-    ):
+    ) as session:
         yield session
 
 
@@ -2753,11 +2755,11 @@ async def charlie_mcp_client(
     charlie_oauth_token: str,
 ) -> AsyncGenerator[ClientSession, Any]:
     """MCP client authenticated as charlie (editor role, in 'editors' group)."""
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8001/mcp",
         token=charlie_oauth_token,
         client_name="Charlie MCP",
-    ):
+    ) as session:
         yield session
 
 
@@ -2767,11 +2769,11 @@ async def diana_mcp_client(
     diana_oauth_token: str,
 ) -> AsyncGenerator[ClientSession, Any]:
     """MCP client authenticated as diana (no-access role)."""
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8001/mcp",
         token=diana_oauth_token,
         client_name="Diana MCP",
-    ):
+    ) as session:
         yield session
 
 

tests/integration/sampling_support.py

@@ -39,10 +39,10 @@ def create_sampling_callback(provider: Provider):
         if provider.supports_generation:
             callback = create_sampling_callback(provider)
 
-            async for session in create_mcp_client_session(
+            async with create_mcp_client_session(
                 url="http://localhost:8000/mcp",
                 sampling_callback=callback,
-            ):
+            ) as session:
                 # Session now supports sampling
                 pass
         ```

tests/integration/test_astrolabe_plotly_visualization.py

@@ -145,11 +145,11 @@ async def test_astrolabe_plotly_visualization_with_basic_auth(
         logger.info(f"Authorization result: {auth_result}")
 
         # Create MCP client session as alice - all MCP operations inside this block
-        async for alice_mcp_client in create_mcp_client_session(
+        async with create_mcp_client_session(
             url="http://localhost:8003/mcp",
             headers={"Authorization": auth_header},
             client_name="Alice BasicAuth MCP",
-        ):
+        ) as alice_mcp_client:
             # Phase 3: Get initial indexed count
             initial_sync = await alice_mcp_client.call_tool(
                 "nc_get_vector_sync_status", {}
@@ -355,11 +355,11 @@ async def test_astrolabe_plotly_visualization_with_basic_auth(
         # Cleanup note if not already cleaned (create new client for cleanup)
         if note_id:
             try:
-                async for cleanup_client in create_mcp_client_session(
+                async with create_mcp_client_session(
                     url="http://localhost:8003/mcp",
                     headers={"Authorization": auth_header},
                     client_name="Cleanup MCP",
-                ):
+                ) as cleanup_client:
                     delete_response = await cleanup_client.call_tool(
                         "nc_notes_delete_note", {"note_id": note_id}
                     )

tests/integration/test_rag.py

@@ -227,11 +227,11 @@ async def nc_mcp_client_with_sampling(
     """
     sampling_callback = create_sampling_callback(generation_provider)
 
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url="http://localhost:8000/mcp",
         client_name=f"Sampling MCP ({provider_name})",
         sampling_callback=sampling_callback,
-    ):
+    ) as session:
         yield session
 
 

tests/server/login_flow/conftest.py

@@ -367,12 +367,12 @@ async def elicitation_callback(
             content={"acknowledged": True},
         )
 
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url=LOGIN_FLOW_MCP_URL,
         token=login_flow_oauth_token,
         client_name="Login Flow MCP",
         elicitation_callback=elicitation_callback,
-    ):
+    ) as session:
         # Step 1: Provision access via Login Flow v2
         logger.info("Starting Login Flow v2 provisioning...")
         provision_result = await session.call_tool(
@@ -499,11 +499,11 @@ async def nc_mcp_login_flow_client_read_only(
     anyio_backend, login_flow_read_only_token: str
 ) -> AsyncGenerator[ClientSession, Any]:
     """MCP client with read-only scopes on the login-flow server."""
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url=LOGIN_FLOW_MCP_URL,
         token=login_flow_read_only_token,
         client_name="Login Flow MCP Read-Only",
-    ):
+    ) as session:
         yield session
 
 
@@ -512,11 +512,11 @@ async def nc_mcp_login_flow_client_write_only(
     anyio_backend, login_flow_write_only_token: str
 ) -> AsyncGenerator[ClientSession, Any]:
     """MCP client with write-only scopes on the login-flow server."""
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url=LOGIN_FLOW_MCP_URL,
         token=login_flow_write_only_token,
         client_name="Login Flow MCP Write-Only",
-    ):
+    ) as session:
         yield session
 
 
@@ -525,11 +525,11 @@ async def nc_mcp_login_flow_client_full_access(
     anyio_backend, login_flow_full_access_token: str
 ) -> AsyncGenerator[ClientSession, Any]:
     """MCP client with full access scopes on the login-flow server."""
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url=LOGIN_FLOW_MCP_URL,
         token=login_flow_full_access_token,
         client_name="Login Flow MCP Full Access",
-    ):
+    ) as session:
         yield session
 
 
@@ -538,11 +538,11 @@ async def nc_mcp_login_flow_client_no_custom_scopes(
     anyio_backend, login_flow_no_custom_scopes_token: str
 ) -> AsyncGenerator[ClientSession, Any]:
     """MCP client with no custom scopes on the login-flow server."""
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url=LOGIN_FLOW_MCP_URL,
         token=login_flow_no_custom_scopes_token,
         client_name="Login Flow MCP No Custom Scopes",
-    ):
+    ) as session:
         yield session
 
 
@@ -724,12 +724,12 @@ async def elicitation_callback(
 
         return ElicitResult(action="accept", content={"acknowledged": True})
 
-    async for session in create_mcp_client_session(
+    async with create_mcp_client_session(
         url=LOGIN_FLOW_MCP_URL,
         token=token,
         client_name=f"Login Flow MCP ({username})",
         elicitation_callback=elicitation_callback,
-    ):
+    ) as session:
         # Provision access
         provision_result = await session.call_tool(
             "nc_auth_provision_access", {"scopes": None}