fix: collapse BaseExceptionGroup to surface real errors from task groups

g97iulio1609 · Copilot · g97iulio1609 · commit 2e8a963140b6 · 2026-02-28T13:15:32.000+01:00
When a task in an anyio task group fails, sibling tasks are cancelled. The resulting BaseExceptionGroup contains the real error alongside Cancelled exceptions from those siblings. This makes error classification extremely difficult for callers. Add open_task_group() context manager and collapse_exception_group() utility that detect this pattern and re-raise just the original error, keeping the full group as __cause__ for debugging. Applied to all 16 create_task_group() sites across: - Client transports (sse, stdio, websocket, streamable_http) - Server transports (sse, stdio, websocket, streamable_http) - Session __aexit__ - Server lowlevel run loop - StreamableHTTP session manager - SessionGroup, InMemoryTransport - Experimental task support Fixes #2114 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
diff --git a/src/mcp/client/_memory.py b/src/mcp/client/_memory.py
@@ -12,6 +12,7 @@
 from mcp.client._transport import TransportStreams
 from mcp.server import Server
 from mcp.server.mcpserver import MCPServer
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared.memory import create_client_server_memory_streams
 
 
@@ -48,7 +49,7 @@ async def _connect(self) -> AsyncIterator[TransportStreams]:
             client_read, client_write = client_streams
             server_read, server_write = server_streams
 
-            async with anyio.create_task_group() as tg:
+            async with open_task_group() as tg:
                 # Start server in background
                 tg.start_soon(
                     lambda: actual_server.run(
diff --git a/src/mcp/client/session_group.py b/src/mcp/client/session_group.py
@@ -24,6 +24,7 @@
 from mcp.client.sse import sse_client
 from mcp.client.stdio import StdioServerParameters
 from mcp.client.streamable_http import streamable_http_client
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared._httpx_utils import create_mcp_http_client
 from mcp.shared.exceptions import MCPError
 from mcp.shared.session import ProgressFnT
@@ -166,7 +167,7 @@ async def __aexit__(
             await self._exit_stack.aclose()
 
         # Concurrently close session stacks.
-        async with anyio.create_task_group() as tg:
+        async with open_task_group() as tg:
             for exit_stack in self._session_exit_stacks.values():
                 tg.start_soon(exit_stack.aclose)
 
diff --git a/src/mcp/client/sse.py b/src/mcp/client/sse.py
@@ -12,6 +12,7 @@
 from httpx_sse._exceptions import SSEError
 
 from mcp import types
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared._httpx_utils import McpHttpClientFactory, create_mcp_http_client
 from mcp.shared.message import SessionMessage
 
@@ -60,7 +61,7 @@ async def sse_client(
     read_stream_writer, read_stream = anyio.create_memory_object_stream(0)
     write_stream, write_stream_reader = anyio.create_memory_object_stream(0)
 
-    async with anyio.create_task_group() as tg:
+    async with open_task_group() as tg:
         try:
             logger.debug(f"Connecting to SSE endpoint: {remove_request_params(url)}")
             async with httpx_client_factory(
diff --git a/src/mcp/client/stdio.py b/src/mcp/client/stdio.py
@@ -20,6 +20,7 @@
     get_windows_executable_command,
     terminate_windows_process_tree,
 )
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared.message import SessionMessage
 
 logger = logging.getLogger(__name__)
@@ -177,7 +178,7 @@ async def stdin_writer():
         except anyio.ClosedResourceError:  # pragma: no cover
             await anyio.lowlevel.checkpoint()
 
-    async with anyio.create_task_group() as tg, process:
+    async with open_task_group() as tg, process:
         tg.start_soon(stdout_reader)
         tg.start_soon(stdin_writer)
         try:
diff --git a/src/mcp/client/streamable_http.py b/src/mcp/client/streamable_http.py
@@ -16,6 +16,7 @@
 from pydantic import ValidationError
 
 from mcp.client._transport import TransportStreams
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared._httpx_utils import create_mcp_http_client
 from mcp.shared.message import ClientMessageMetadata, SessionMessage
 from mcp.types import (
@@ -546,7 +547,7 @@ async def streamable_http_client(
 
     transport = StreamableHTTPTransport(url)
 
-    async with anyio.create_task_group() as tg:
+    async with open_task_group() as tg:
         try:
             logger.debug(f"Connecting to StreamableHTTP endpoint: {url}")
 
diff --git a/src/mcp/client/websocket.py b/src/mcp/client/websocket.py
@@ -9,6 +9,7 @@
 from websockets.typing import Subprotocol
 
 from mcp import types
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared.message import SessionMessage
 
 
@@ -68,7 +69,7 @@ async def ws_writer():
                     msg_dict = session_message.message.model_dump(by_alias=True, mode="json", exclude_unset=True)
                     await ws.send(json.dumps(msg_dict))
 
-        async with anyio.create_task_group() as tg:
+        async with open_task_group() as tg:
             # Start reader and writer tasks
             tg.start_soon(ws_reader)
             tg.start_soon(ws_writer)
diff --git a/src/mcp/server/experimental/task_result_handler.py b/src/mcp/server/experimental/task_result_handler.py
@@ -15,6 +15,7 @@
 import anyio
 
 from mcp.server.session import ServerSession
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared.exceptions import MCPError
 from mcp.shared.experimental.tasks.helpers import RELATED_TASK_METADATA_KEY, is_terminal
 from mcp.shared.experimental.tasks.message_queue import TaskMessageQueue
@@ -162,7 +163,7 @@ async def _wait_for_task_update(self, task_id: str) -> None:
 
         Races between store update and queue message - first one wins.
         """
-        async with anyio.create_task_group() as tg:
+        async with open_task_group() as tg:
 
             async def wait_for_store() -> None:
                 try:
diff --git a/src/mcp/server/experimental/task_support.py b/src/mcp/server/experimental/task_support.py
@@ -13,6 +13,7 @@
 
 from mcp.server.experimental.task_result_handler import TaskResultHandler
 from mcp.server.session import ServerSession
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
 from mcp.shared.experimental.tasks.message_queue import InMemoryTaskMessageQueue, TaskMessageQueue
 from mcp.shared.experimental.tasks.store import TaskStore
@@ -79,7 +80,7 @@ async def run(self) -> AsyncIterator[None]:
                 # Task group is now available
                 ...
         """
-        async with anyio.create_task_group() as tg:
+        async with open_task_group() as tg:
             self._task_group = tg
             try:
                 yield
diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py
@@ -66,6 +66,7 @@ async def main():
 from mcp.server.streamable_http import EventStore
 from mcp.server.streamable_http_manager import StreamableHTTPASGIApp, StreamableHTTPSessionManager
 from mcp.server.transport_security import TransportSecuritySettings
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared.exceptions import MCPError
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 from mcp.shared.session import RequestResponder
@@ -389,7 +390,7 @@ async def run(
                 task_support.configure_session(session)
                 await stack.enter_async_context(task_support.run())
 
-            async with anyio.create_task_group() as tg:
+            async with open_task_group() as tg:
                 async for message in session.incoming_messages:
                     logger.debug("Received message: %s", message)
 
diff --git a/src/mcp/server/sse.py b/src/mcp/server/sse.py
@@ -55,6 +55,7 @@ async def handle_sse(request):
     TransportSecurityMiddleware,
     TransportSecuritySettings,
 )
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 
 logger = logging.getLogger(__name__)
@@ -174,7 +175,7 @@ async def sse_writer():
                         }
                     )
 
-        async with anyio.create_task_group() as tg:
+        async with open_task_group() as tg:
 
             async def response_wrapper(scope: Scope, receive: Receive, send: Send):
                 """The EventSourceResponse returning signals a client close / disconnect.
diff --git a/src/mcp/server/stdio.py b/src/mcp/server/stdio.py
@@ -26,6 +26,7 @@ async def run_server():
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 
 from mcp import types
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared.message import SessionMessage
 
 
@@ -77,7 +78,7 @@ async def stdout_writer():
         except anyio.ClosedResourceError:  # pragma: no cover
             await anyio.lowlevel.checkpoint()
 
-    async with anyio.create_task_group() as tg:
+    async with open_task_group() as tg:
         tg.start_soon(stdin_reader)
         tg.start_soon(stdout_writer)
         yield read_stream, write_stream
diff --git a/src/mcp/server/streamable_http.py b/src/mcp/server/streamable_http.py
@@ -25,6 +25,7 @@
 from starlette.types import Receive, Scope, Send
 
 from mcp.server.transport_security import TransportSecurityMiddleware, TransportSecuritySettings
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared.message import ServerMessageMetadata, SessionMessage
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 from mcp.types import (
@@ -614,7 +615,7 @@ async def sse_writer():  # pragma: lax no cover
                 # Start the SSE response (this will send headers immediately)
                 try:
                     # First send the response to establish the SSE connection
-                    async with anyio.create_task_group() as tg:
+                    async with open_task_group() as tg:
                         tg.start_soon(response, scope, receive, send)
                         # Then send the message to be processed by the server
                         session_message = self._create_session_message(message, request, request_id, protocol_version)
@@ -970,7 +971,7 @@ async def connect(
         self._write_stream = write_stream
 
         # Start a task group for message routing
-        async with anyio.create_task_group() as tg:
+        async with open_task_group() as tg:
             # Create a message router that distributes messages to request streams
             async def message_router():
                 try:
diff --git a/src/mcp/server/streamable_http_manager.py b/src/mcp/server/streamable_http_manager.py
@@ -21,6 +21,7 @@
     StreamableHTTPServerTransport,
 )
 from mcp.server.transport_security import TransportSecuritySettings
+from mcp.shared._exception_utils import open_task_group
 from mcp.types import INVALID_REQUEST, ErrorData, JSONRPCError
 
 if TYPE_CHECKING:
@@ -122,7 +123,7 @@ async def lifespan(app: Starlette) -> AsyncIterator[None]:
                 )
             self._has_started = True
 
-        async with anyio.create_task_group() as tg:
+        async with open_task_group() as tg:
             # Store the task group for later use
             self._task_group = tg
             logger.info("StreamableHTTP session manager started")
diff --git a/src/mcp/server/websocket.py b/src/mcp/server/websocket.py
@@ -7,6 +7,7 @@
 from starlette.websockets import WebSocket
 
 from mcp import types
+from mcp.shared._exception_utils import open_task_group
 from mcp.shared.message import SessionMessage
 
 
@@ -52,7 +53,7 @@ async def ws_writer():
         except anyio.ClosedResourceError:
             await websocket.close()
 
-    async with anyio.create_task_group() as tg:
+    async with open_task_group() as tg:
         tg.start_soon(ws_reader)
         tg.start_soon(ws_writer)
         yield (read_stream, write_stream)
diff --git a/src/mcp/shared/_exception_utils.py b/src/mcp/shared/_exception_utils.py
@@ -0,0 +1,82 @@
+"""Utilities for collapsing BaseExceptionGroup noise from anyio task groups.
+
+When a task in an anyio task group fails, sibling tasks are cancelled.  The
+resulting ``BaseExceptionGroup`` contains the real error alongside
+``Cancelled`` exceptions from those siblings.  This module provides helpers
+that detect that pattern and re-raise just the original error, preserving the
+full group as ``__cause__`` for debugging.
+
+If multiple tasks fail with non-cancellation errors concurrently, the full
+``BaseExceptionGroup`` is preserved unchanged.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from collections.abc import AsyncIterator
+
+import anyio
+import anyio.abc
+
+
+def collapse_exception_group(
+    eg: BaseExceptionGroup,
+    cancelled_type: type[BaseException] = asyncio.CancelledError,
+) -> BaseException:
+    """Extract the single real error from a *BaseExceptionGroup* if possible.
+
+    Args:
+        eg: The exception group to collapse.
+        cancelled_type: The cancellation exception class to filter out.
+            Defaults to ``asyncio.CancelledError``.  The ``open_task_group``
+            context manager passes ``anyio.get_cancelled_exc_class()`` so the
+            correct type is used for any backend.
+
+    Returns:
+        * The single non-cancelled exception if exactly one exists.
+        * A filtered group (cancelled noise stripped) if multiple real errors.
+        * A single ``Cancelled`` if all exceptions are cancellations.
+    """
+
+    # split(type) uses isinstance on leaf exceptions, NOT on the group.
+    # Using split(lambda) is incorrect because the lambda would first be
+    # called on the group object itself.
+    cancelled, non_cancelled = eg.split(cancelled_type)
+
+    if non_cancelled is None:
+        # Every exception is a cancellation – surface just one.
+        return eg.exceptions[0]
+
+    if len(non_cancelled.exceptions) == 1:
+        return non_cancelled.exceptions[0]
+
+    # Multiple real errors – return the filtered group (without Cancelled).
+    return non_cancelled if non_cancelled is not eg else eg
+
+
+@contextlib.asynccontextmanager
+async def open_task_group() -> AsyncIterator[anyio.abc.TaskGroup]:
+    """Drop-in replacement for ``anyio.create_task_group()`` that collapses
+    exception groups containing a single real error plus cancellation noise.
+
+    Usage::
+
+        async with open_task_group() as tg:
+            tg.start_soon(some_task)
+            ...
+
+    If *some_task* raises ``ConnectionError`` and all siblings are cancelled,
+    this context manager will raise ``ConnectionError`` directly (with the
+    original ``BaseExceptionGroup`` attached as ``__cause__``).
+    """
+
+    try:
+        async with anyio.create_task_group() as tg:
+            yield tg
+    except BaseExceptionGroup as eg:
+        cancelled_cls = anyio.get_cancelled_exc_class()
+        collapsed = collapse_exception_group(eg, cancelled_type=cancelled_cls)
+        if collapsed is not eg:
+            raise collapsed from eg
+        raise
diff --git a/src/mcp/shared/session.py b/src/mcp/shared/session.py
@@ -11,6 +11,7 @@
 from pydantic import BaseModel, TypeAdapter
 from typing_extensions import Self
 
+from mcp.shared._exception_utils import collapse_exception_group
 from mcp.shared.exceptions import MCPError
 from mcp.shared.message import MessageMetadata, ServerMessageMetadata, SessionMessage
 from mcp.shared.response_router import ResponseRouter
@@ -228,7 +229,14 @@ 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)
+        try:
+            return await self._task_group.__aexit__(exc_type, exc_val, exc_tb)
+        except BaseExceptionGroup as eg:
+            cancelled_cls = anyio.get_cancelled_exc_class()
+            collapsed = collapse_exception_group(eg, cancelled_type=cancelled_cls)
+            if collapsed is not eg:
+                raise collapsed from eg
+            raise
 
     async def send_request(
         self,
diff --git a/tests/shared/test_exception_utils.py b/tests/shared/test_exception_utils.py

Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,7 @@ async def handle_sse(request):`
`55`	`55`	`TransportSecurityMiddleware,`
`56`	`56`	`TransportSecuritySettings,`
`57`	`57`	`)`
	`58`	`+from mcp.shared._exception_utils import open_task_group`
`58`	`59`	`from mcp.shared.message import ServerMessageMetadata, SessionMessage`
`59`	`60`
`60`	`61`	`logger = logging.getLogger(__name__)`
`@@ -174,7 +175,7 @@ async def sse_writer():`
`174`	`175`	`}`
`175`	`176`	`)`
`176`	`177`
`177`		`- async with anyio.create_task_group() as tg:`
	`178`	`+ async with open_task_group() as tg:`
`178`	`179`
`179`	`180`	`async def response_wrapper(scope: Scope, receive: Receive, send: Send):`
`180`	`181`	`"""The EventSourceResponse returning signals a client close / disconnect.`