Isolate Sentry context per FastStream message

lite_bootstrap/bootstrappers/faststream_bootstrapper.py

@@ -4,6 +4,7 @@
 import logging
 import typing
 from collections.abc import Mapping, Sequence
+from types import TracebackType
 
 from lite_bootstrap import import_checker
 from lite_bootstrap.bootstrappers.base import BaseBootstrapper
@@ -20,6 +21,24 @@
     from faststream._internal.logger.params_storage import ManualLoggerStorage
     from faststream.asgi import AsgiFastStream, AsgiResponse
     from faststream.asgi import get as handle_get
+    from faststream.middlewares import BaseMiddleware
+
+    class _SentryIsolationMiddleware(BaseMiddleware[typing.Any, typing.Any]):
+        async def on_receive(self) -> None:
+            self._scope_context = sentry_sdk.isolation_scope()
+            self._scope_context.__enter__()
+
+        async def after_processed(
+            self,
+            exc_type: type[BaseException] | None = None,
+            exc_val: BaseException | None = None,
+            exc_tb: TracebackType | None = None,
+        ) -> None:
+            self._scope_context.__exit__(exc_type, exc_val, exc_tb)
+
+
+if import_checker.is_sentry_installed:
+    import sentry_sdk
 
 if import_checker.is_structlog_installed:
     import structlog
@@ -227,13 +246,20 @@ def bootstrap(self) -> None:
                 broker.add_middleware(middleware_cls(registry=self.collector_registry))
 
 
+class FastStreamSentryInstrument(SentryInstrument["FastStreamConfig"]):
+    def bootstrap(self) -> None:
+        super().bootstrap()
+        for broker in self.bootstrap_config.application.brokers:
+            broker.insert_middleware(_SentryIsolationMiddleware)
+
+
 class FastStreamBootstrapper(BaseBootstrapper["AsgiFastStream", FastStreamConfig]):
     __slots__ = "bootstrap_config", "instruments"
 
     instruments_types: typing.ClassVar = [
         FastStreamOpenTelemetryInstrument,
         PyroscopeInstrument,
-        SentryInstrument,
+        FastStreamSentryInstrument,
         FastStreamHealthChecksInstrument,
         FastStreamLoggingInstrument,
         FastStreamPrometheusInstrument,

tests/test_faststream_bootstrap.py

@@ -5,14 +5,17 @@
 import typing
 import uuid
 import warnings
+from types import TracebackType
 from unittest.mock import AsyncMock, patch
 
 import faststream.asgi
 import prometheus_client
 import pytest
+import sentry_sdk
 import structlog
 from faststream._internal.broker import BrokerUsecase
 from faststream._internal.logger.params_storage import ManualLoggerStorage
+from faststream.middlewares import BaseMiddleware
 from faststream.redis import RedisBroker, TestRedisBroker
 from faststream.redis.opentelemetry import RedisTelemetryMiddleware
 from faststream.redis.prometheus import RedisPrometheusMiddleware
@@ -124,6 +127,58 @@ async def test_faststream_bootstrap(broker: RedisBroker) -> None:
     assert not bootstrapper.is_bootstrapped
 
 
+async def test_faststream_sentry_context_is_isolated_per_message(broker: RedisBroker) -> None:
+    transport = SentryTestTransport()
+    config = FastStreamConfig(
+        application=faststream.asgi.AsgiFastStream(broker),
+        sentry_dsn="https://testdsn@localhost/1",
+        sentry_additional_params={"transport": transport},
+    )
+    bootstrapper = FastStreamBootstrapper(bootstrap_config=config)
+
+    class CaptureFailure(BaseMiddleware[typing.Any, typing.Any]):
+        async def after_processed(
+            self,
+            exc_type: type[BaseException] | None = None,
+            exc_val: BaseException | None = None,
+            exc_tb: TracebackType | None = None,
+        ) -> None:
+            if exc_type is not None and exc_val is not None:
+                sentry_sdk.capture_exception((exc_type, exc_val, exc_tb))
+
+    broker.add_middleware(CaptureFailure)
+
+    @broker.subscriber("messages")
+    async def consume(message: str) -> None:
+        if message == "first":
+            sentry_sdk.set_tag("customer", "alice")
+            sentry_sdk.add_breadcrumb(message="first message breadcrumb")
+            return
+        sentry_sdk.set_tag("message", "second")
+        sentry_sdk.add_breadcrumb(message="second message breadcrumb")
+        msg = "second message failed"
+        raise ValueError(msg)
+
+    with sentry_sdk.isolation_scope():
+        bootstrapper.bootstrap()
+        try:
+            async with TestRedisBroker(broker):
+                await broker.publish("first", "messages")
+                with pytest.raises(ValueError, match="second message failed"):
+                    await broker.publish("second", "messages")
+        finally:
+            bootstrapper.teardown()
+
+    (event,) = [event for envelope in transport.mock_envelopes if (event := envelope.get_event())]
+    assert event["exception"]["values"][-1]["value"] == "second message failed"
+    assert event["tags"] == {"message": "second"}
+    breadcrumbs = event.get("breadcrumbs")
+    assert isinstance(breadcrumbs, dict)
+    messages = [crumb.get("message") for crumb in breadcrumbs["values"]]
+    assert "second message breadcrumb" in messages
+    assert "first message breadcrumb" not in messages
+
+
 async def test_faststream_bootstrap_health_check_wo_broker() -> None:
     bootstrap_config = build_faststream_config()
     bootstrapper = FastStreamBootstrapper(bootstrap_config=bootstrap_config)
@@ -624,7 +679,7 @@ def _middleware_types(broker: BrokerUsecase[typing.Any, typing.Any, typing.Any])
 def _bootstrap_with_brokers(
     brokers: list[BrokerUsecase[typing.Any, typing.Any, typing.Any]], **overrides: object
 ) -> FastStreamBootstrapper:
-    config = dataclasses.replace(build_faststream_config(broker=brokers[0]), **overrides)
+    config = dataclasses.replace(build_faststream_config(broker=brokers[0]), sentry_dsn=None, **overrides)
     for extra in brokers[1:]:
         config.application.add_broker(extra)
     bootstrapper = FastStreamBootstrapper(bootstrap_config=config)