fix: BaseSession.__aexit__ healing checkpoint for gh-106749 (CI 3.11/3.14)

The second gh-106749 site (the first being _memory.py): BaseSession's
task-group cancel is also delivered via coro.throw() into the host task,
desyncing CTracer past the caller's frame on 3.11. A
cancel_shielded_checkpoint() after the join resumes via .send() and
re-stamps the missing 'call' events. Shielded so a pending outer cancel
isn't re-delivered here.

The added tick shifts whether tests reach streamable_http.py's outer
POST-error handler (it was timing-dependent before too); marked lax no
cover.

Also: collapse a nested with in test_jsonrpc_dispatcher.py and pragma: no
branch it (3.14 misreports the 254->255 arc on the nested-with bytecode
shape).

Author: maxisbey

src/mcp/server/streamable_http.py

@@ -635,15 +635,18 @@ async def sse_writer():
                 finally:
                     await sse_stream_reader.aclose()
 
-        except Exception as err:
+        except Exception as err:  # pragma: lax no cover
+            # Reached only when something raises during POST handling outside
+            # the per-SSE-stream guard above; whether tests reach this depends
+            # on client teardown timing.
             logger.exception("Error handling POST request")
             response = self._create_error_response(
                 f"Error handling POST request: {err}",
                 HTTPStatus.INTERNAL_SERVER_ERROR,
                 INTERNAL_ERROR,
             )
             await response(scope, receive, send)
-            if writer:  # pragma: no cover
+            if writer:
                 await writer.send(Exception(err))
             return  # pragma: no cover
 

src/mcp/shared/session.py

@@ -7,6 +7,7 @@
 from typing import Any, Generic, Protocol, TypeVar
 
 import anyio
+import anyio.lowlevel
 from anyio.streams.memory import MemoryObjectSendStream
 from opentelemetry.trace import SpanKind
 from pydantic import BaseModel, TypeAdapter
@@ -173,7 +174,15 @@ async def __aexit__(
         # would be very surprising behavior), so make sure to cancel the tasks
         # in the task group.
         self._task_group.cancel_scope.cancel()
-        return await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
+        result = await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
+        # The cancel above is delivered via `coro.throw()` into this task; on
+        # CPython 3.11 (gh-106749) that drops `'call'` trace events for the
+        # outer await chain and desyncs coverage's CTracer past the caller's
+        # frame. Yielding once here resumes via `.send()`, which re-stamps the
+        # missing `'call'` events and resyncs the tracer. Shielded so a pending
+        # outer cancel is not re-delivered at this point.
+        await anyio.lowlevel.cancel_shielded_checkpoint()
+        return result
 
     async def send_request(
         self,

tests/shared/test_jsonrpc_dispatcher.py

@@ -249,10 +249,9 @@ async def test_run_closes_write_stream_on_exit():
     async with anyio.create_task_group() as tg:
         await tg.start(server.run, on_request, on_notify)
         c2s_send.close()  # EOF the read side; run() exits
-        with anyio.fail_after(5):
+        with anyio.fail_after(5), pytest.raises(anyio.EndOfStream):  # pragma: no branch
             # Write end was entered and released by run(); peer's receive sees EOF.
-            with pytest.raises(anyio.EndOfStream):
-                await s2c_recv.receive()
+            await s2c_recv.receive()
     s2c_recv.close()