avoid collapsing exception groups from user code (#2830)

Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>

Author: graingert

starlette/_utils.py

@@ -2,10 +2,12 @@
 
 import functools
 import sys
-from collections.abc import Awaitable, Callable, Generator
-from contextlib import AbstractAsyncContextManager, contextmanager
+from collections.abc import AsyncGenerator, Awaitable, Callable, Generator
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
 from typing import Any, Generic, Protocol, TypeVar, overload
 
+import anyio.abc
+
 from starlette.types import Scope
 
 if sys.version_info >= (3, 13):  # pragma: no cover
@@ -16,12 +18,14 @@
 
     from typing_extensions import TypeIs
 
-has_exceptiongroups = True
 if sys.version_info < (3, 11):  # pragma: no cover
     try:
-        from exceptiongroup import BaseExceptionGroup  # type: ignore[unused-ignore,import-not-found]
+        from exceptiongroup import BaseExceptionGroup
     except ImportError:
-        has_exceptiongroups = False
+
+        class BaseExceptionGroup(BaseException):  # type: ignore[no-redef]
+            pass
+
 
 T = TypeVar("T")
 AwaitableCallable = Callable[..., Awaitable[T]]
@@ -75,16 +79,18 @@ async def __aexit__(self, *args: Any) -> None | bool:
         return None
 
 
-@contextmanager
-def collapse_excgroups() -> Generator[None, None, None]:
+@asynccontextmanager
+async def create_collapsing_task_group() -> AsyncGenerator[anyio.abc.TaskGroup, None]:
     try:
-        yield
-    except BaseException as exc:
-        if has_exceptiongroups:  # pragma: no cover
-            while isinstance(exc, BaseExceptionGroup) and len(exc.exceptions) == 1:
-                exc = exc.exceptions[0]
-
-        raise exc
+        async with anyio.create_task_group() as tg:
+            yield tg
+    except BaseExceptionGroup as excs:
+        if len(excs.exceptions) != 1:
+            raise
+
+        exc = excs.exceptions[0]
+        context = None if exc.__suppress_context__ else exc.__context__
+        raise exc from exc.__cause__ or context
 
 
 def get_route_path(scope: Scope) -> str:

starlette/middleware/base.py

@@ -5,7 +5,7 @@
 
 import anyio
 
-from starlette._utils import collapse_excgroups
+from starlette._utils import create_collapsing_task_group
 from starlette.requests import ClientDisconnect, Request
 from starlette.responses import Response
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
@@ -188,8 +188,8 @@ async def body_stream() -> BodyStreamGenerator:
 
         streams: anyio.create_memory_object_stream[Message] = anyio.create_memory_object_stream()
         send_stream, recv_stream = streams
-        with recv_stream, send_stream, collapse_excgroups():
-            async with anyio.create_task_group() as task_group:
+        with recv_stream, send_stream:
+            async with create_collapsing_task_group() as task_group:
                 response = await self.dispatch_func(request, call_next)
                 await response(scope, wrapped_receive, send)
                 response_sent.set()

starlette/middleware/wsgi.py

@@ -10,6 +10,7 @@
 import anyio
 from anyio.abc import ObjectReceiveStream, ObjectSendStream
 
+from starlette._utils import create_collapsing_task_group
 from starlette.types import Receive, Scope, Send
 
 warnings.warn(
@@ -104,7 +105,7 @@ async def __call__(self, receive: Receive, send: Send) -> None:
             more_body = message.get("more_body", False)
         environ = build_environ(self.scope, body)
 
-        async with anyio.create_task_group() as task_group:
+        async with create_collapsing_task_group() as task_group:
             task_group.start_soon(self.sender, send)
             async with self.stream_send:
                 await anyio.to_thread.run_sync(self.wsgi, environ, self.start_response)

starlette/responses.py

@@ -18,7 +18,7 @@
 import anyio
 import anyio.to_thread
 
-from starlette._utils import collapse_excgroups
+from starlette._utils import create_collapsing_task_group
 from starlette.background import BackgroundTask
 from starlette.concurrency import iterate_in_threadpool
 from starlette.datastructures import URL, Headers, MutableHeaders
@@ -270,15 +270,14 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
             except OSError:
                 raise ClientDisconnect()
         else:
-            with collapse_excgroups():
-                async with anyio.create_task_group() as task_group:
+            async with create_collapsing_task_group() as task_group:
 
-                    async def wrap(func: Callable[[], Awaitable[None]]) -> None:
-                        await func()
-                        task_group.cancel_scope.cancel()
+                async def wrap(func: Callable[[], Awaitable[None]]) -> None:
+                    await func()
+                    task_group.cancel_scope.cancel()
 
-                    task_group.start_soon(wrap, partial(self.stream_response, send))
-                    await wrap(partial(self.listen_for_disconnect, receive))
+                task_group.start_soon(wrap, partial(self.stream_response, send))
+                await wrap(partial(self.listen_for_disconnect, receive))
 
         if self.background is not None:
             await self.background()

tests/middleware/test_base.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextvars
+import sys
 from collections.abc import AsyncGenerator, AsyncIterator, Generator
 from contextlib import AsyncExitStack
 from pathlib import Path
@@ -22,6 +23,9 @@
 from starlette.websockets import WebSocket
 from tests.types import TestClientFactory
 
+if sys.version_info < (3, 11):  # pragma: no cover
+    from exceptiongroup import ExceptionGroup
+
 
 class CustomMiddleware(BaseHTTPMiddleware):
     async def dispatch(
@@ -42,6 +46,10 @@ def exc(request: Request) -> None:
     raise Exception("Exc")
 
 
+def exc_group(request: Request) -> None:
+    raise ExceptionGroup("my exception group", [ValueError("TEST")])
+
+
 def exc_stream(request: Request) -> StreamingResponse:
     return StreamingResponse(_generate_faulty_stream())
 
@@ -77,6 +85,7 @@ async def websocket_endpoint(session: WebSocket) -> None:
     routes=[
         Route("/", endpoint=homepage),
         Route("/exc", endpoint=exc),
+        Route("/exc-group", endpoint=exc_group),
         Route("/exc-stream", endpoint=exc_stream),
         Route("/no-response", endpoint=NoResponse),
         WebSocketRoute("/ws", endpoint=websocket_endpoint),
@@ -90,13 +99,18 @@ def test_custom_middleware(test_client_factory: TestClientFactory) -> None:
     response = client.get("/")
     assert response.headers["Custom-Header"] == "Example"
 
-    with pytest.raises(Exception) as ctx:
+    with pytest.raises(Exception) as ctx1:
         response = client.get("/exc")
-    assert str(ctx.value) == "Exc"
+    assert str(ctx1.value) == "Exc"
 
-    with pytest.raises(Exception) as ctx:
+    with pytest.raises(Exception) as ctx2:
         response = client.get("/exc-stream")
-    assert str(ctx.value) == "Faulty Stream"
+    assert str(ctx2.value) == "Faulty Stream"
+
+    with pytest.raises(ExceptionGroup, match="my exception group") as ctx3:
+        client.get("/exc-group")
+    assert len(ctx3.value.exceptions) == 1
+    assert isinstance(ctx3.value.exceptions[0], ValueError)
 
     with pytest.raises(RuntimeError):
         response = client.get("/no-response")

tests/middleware/test_wsgi.py

@@ -4,7 +4,6 @@
 
 import pytest
 
-from starlette._utils import collapse_excgroups
 from starlette.middleware.wsgi import WSGIMiddleware, build_environ
 from tests.types import TestClientFactory
 
@@ -86,7 +85,7 @@ def test_wsgi_exception(test_client_factory: TestClientFactory) -> None:
     # The HTTP protocol implementations would catch this error and return 500.
     app = WSGIMiddleware(raise_exception)
     client = test_client_factory(app)
-    with pytest.raises(RuntimeError), collapse_excgroups():
+    with pytest.raises(RuntimeError):
         client.get("/")
 
 

tests/test__utils.py

@@ -1,12 +1,16 @@
 import functools
+import sys
 from typing import Any
 from unittest.mock import create_autospec
 
 import pytest
 
-from starlette._utils import get_route_path, is_async_callable
+from starlette._utils import create_collapsing_task_group, get_route_path, is_async_callable
 from starlette.types import Scope
 
+if sys.version_info < (3, 11):  # pragma: no cover
+    from exceptiongroup import ExceptionGroup
+
 
 def test_async_func() -> None:
     async def async_func() -> None: ...  # pragma: no cover
@@ -102,3 +106,31 @@ async def async_func() -> None: ...  # pragma: no cover
 )
 def test_get_route_path(scope: Scope, expected_result: str) -> None:
     assert get_route_path(scope) == expected_result
+
+
+@pytest.mark.anyio
+async def test_collapsing_task_group_one_exc() -> None:
+    class MyException(Exception):
+        pass
+
+    with pytest.raises(MyException):
+        async with create_collapsing_task_group():
+            raise MyException
+
+
+@pytest.mark.anyio
+async def test_collapsing_task_group_two_exc() -> None:
+    class MyException(Exception):
+        pass
+
+    async def raise_exc() -> None:
+        raise MyException
+
+    with pytest.raises(ExceptionGroup) as exc:
+        async with create_collapsing_task_group() as task_group:
+            task_group.start_soon(raise_exc)
+            raise MyException
+
+    exc1, exc2 = exc.value.exceptions
+    assert isinstance(exc1, MyException)
+    assert isinstance(exc2, MyException)