feat: add subscriber and broker exception handlers

faststream/_internal/broker/broker.py

@@ -11,6 +11,7 @@
     ConnectionType,
     MsgType,
 )
+from faststream._internal.utils import apply_types, to_async
 
 from .pub_base import BrokerPublishMixin
 from .registrator import Registrator
@@ -21,6 +22,7 @@
     from faststream._internal.context.repository import ContextRepo
     from faststream._internal.di import FastDependsConfig
     from faststream._internal.producer import ProducerProto
+    from faststream._internal.types import AsyncExceptionHandler
     from faststream.specification.schema import BrokerSpec
 
 
@@ -48,6 +50,15 @@ def __init__(
             routers=routers,
             config=config,
         )
+        if config.broker_exception_handler is not None:
+            async_handler: AsyncExceptionHandler = to_async(
+                config.broker_exception_handler,
+            )
+            config._broker_exception_handler = apply_types(
+                async_handler,
+                serializer_cls=config.fd_config._serializer,
+                context__=config.context,
+            )
         self.specification = specification
 
         self.running = False

faststream/_internal/configs/broker.py

@@ -16,7 +16,12 @@
 
     from faststream._internal.context import ContextRepo
     from faststream._internal.parser import CodecProto
-    from faststream._internal.types import BrokerMiddleware, CustomCallable
+    from faststream._internal.types import (
+        AsyncExceptionHandler,
+        BrokerMiddleware,
+        CustomCallable,
+        ExceptionHandler,
+    )
     from faststream.middlewares import AckPolicy
 
 
@@ -29,6 +34,12 @@ class BrokerConfig:
     broker_parser: Optional["CustomCallable"] = None
     broker_decoder: Optional["CustomCallable"] = None
     broker_codec: Optional["CodecProto"] = None
+    broker_exception_handler: Optional["ExceptionHandler"] = None
+    _broker_exception_handler: Optional["AsyncExceptionHandler"] = field(
+        default=None,
+        init=False,
+        repr=False,
+    )
 
     producer: "ProducerProto[Any]" = field(default_factory=ProducerUnset)
     logger: "LoggerState" = field(default_factory=LoggerState)

faststream/_internal/configs/endpoint.py

@@ -5,7 +5,7 @@
 from faststream.middlewares import AckPolicy
 
 if TYPE_CHECKING:
-    from faststream._internal.types import AsyncCallable
+    from faststream._internal.types import AsyncCallable, ExceptionHandler
 
     from .broker import BrokerConfig
 
@@ -23,6 +23,7 @@ class PublisherUsecaseConfig(EndpointConfig):
 @dataclass(kw_only=True)
 class SubscriberUsecaseConfig(EndpointConfig):
     no_reply: bool = False
+    exception_handler: "ExceptionHandler | None" = None
 
     _ack_policy: AckPolicy = field(default_factory=lambda: EMPTY, repr=False)
 

faststream/_internal/endpoint/subscriber/usecase.py

@@ -24,6 +24,7 @@
     P_HandlerParams,
     T_HandlerReturn,
 )
+from faststream._internal.utils import apply_types
 from faststream._internal.utils.functions import FakeContext, to_async
 from faststream.exceptions import StopConsume, SubscriberNotFound
 from faststream.middlewares import AcknowledgementMiddleware
@@ -45,6 +46,7 @@
     from faststream._internal.endpoint.publisher import PublisherProto
     from faststream._internal.parser import CodecProto
     from faststream._internal.types import (
+        AsyncExceptionHandler,
         AsyncFilter,
         BrokerMiddleware,
         CustomCallable,
@@ -87,6 +89,16 @@ def __init__(
         self._no_reply = config.no_reply
         self._parser = config.parser
         self._decoder = config.decoder
+        self._exception_handler_call: AsyncExceptionHandler | None = None
+        if config.exception_handler is not None:
+            async_handler: AsyncExceptionHandler = to_async(
+                config.exception_handler,
+            )
+            self._exception_handler_call = apply_types(
+                async_handler,
+                serializer_cls=self._outer_config.fd_config._serializer,
+                context__=self._outer_config.context,
+            )
 
         self.ack_policy = config.ack_policy
         self.__auto_ack_disabled = config.auto_ack_disabled
@@ -313,6 +325,19 @@ def real_wrapper(
 
         return real_wrapper(func)
 
+    async def _handle_exception(self, exc: BaseException) -> bool:
+        local_handler = self._exception_handler_call
+        if local_handler is not None and await local_handler(exc):
+            return True
+
+        broker_handler: AsyncExceptionHandler | None = (
+            self._outer_config._broker_exception_handler
+        )
+        if broker_handler is not None:
+            return await broker_handler(exc)
+
+        return False
+
     async def consume(self, msg: MsgType) -> Any:
         """Consume a message asynchronously."""
         if not self.running:
@@ -332,9 +357,14 @@ async def consume(self, msg: MsgType) -> Any:
             if app := self._outer_config.context.get("app"):
                 app.exit()
 
-        except Exception:  # nosec B110  # noqa: S110
+        except Exception as exc:
             # All other exceptions were logged by CriticalLogMiddleware
-            pass
+            has_handler = (
+                self._exception_handler_call is not None
+                or self._outer_config._broker_exception_handler is not None
+            )
+            if has_handler and not await self._handle_exception(exc):
+                raise
 
     async def process_message(self, msg: MsgType) -> "Response":
         """Execute all message processing stages."""

faststream/_internal/types.py

@@ -42,6 +42,10 @@
 AsyncFilter: TypeAlias = Callable[[StreamMsg], Awaitable[bool]]
 Filter: TypeAlias = SyncFilter[StreamMsg] | AsyncFilter[StreamMsg]
 
+SyncExceptionHandler: TypeAlias = Callable[[BaseException], bool]
+AsyncExceptionHandler: TypeAlias = Callable[[BaseException], Awaitable[bool]]
+ExceptionHandler: TypeAlias = SyncExceptionHandler | AsyncExceptionHandler
+
 SyncCallable: TypeAlias = Callable[
     [Any],
     Any,

faststream/kafka/broker/broker.py

@@ -59,6 +59,7 @@
     from faststream._internal.types import (
         BrokerMiddleware,
         CustomCallable,
+        ExceptionHandler,
     )
     from faststream.kafka.message import KafkaMessage
     from faststream.kafka.types import KafkaSendableMessage
@@ -234,6 +235,7 @@ def __init__(
         decoder: Optional["CustomCallable"] = None,
         codec: Optional["CodecProto"] = None,
         parser: Optional["CustomCallable"] = None,
+        exception_handler: "ExceptionHandler | None" = None,
         dependencies: Sequence["Dependant"] = (),
         middlewares: Sequence["BrokerMiddleware[Any, Any]"] = (),
         routers: Iterable[KafkaRegistrator] = (),
@@ -348,6 +350,10 @@ def __init__(
                 Custom codec object.
             parser (Optional[CustomCallable]):
                 Custom parser object.
+            exception_handler (ExceptionHandler | None):
+                Sync or async exception handler called when a subscriber's
+                exception handler is absent or returns False.
+                Return True if handled, otherwise False for default handling.
             dependencies (Sequence[Dependant]):
                 Dependencies to apply to all broker subscribers.
             middlewares (Sequence[BrokerMiddlewarep[Any, Any]]):
@@ -455,6 +461,7 @@ def __init__(
                 broker_decoder=decoder,
                 broker_codec=codec,
                 broker_parser=parser,
+                broker_exception_handler=exception_handler,
                 broker_middlewares=middlewares,
                 logger=make_kafka_logger_state(
                     logger=logger,

faststream/kafka/broker/registrator.py

@@ -30,6 +30,7 @@
     from faststream._internal.types import (
         BrokerMiddleware,
         CustomCallable,
+        ExceptionHandler,
     )
     from faststream.kafka.publisher.usecase import (
         BatchPublisher,
@@ -96,6 +97,7 @@ def subscriber(
         max_workers: None = None,
         ack_policy: AckPolicy = EMPTY,
         no_reply: bool = False,
+        exception_handler: "ExceptionHandler | None" = None,
         # Specification args
         title: str | None = None,
         description: str | None = None,
@@ -146,6 +148,7 @@ def subscriber(
         max_workers: None = None,
         ack_policy: AckPolicy = EMPTY,
         no_reply: bool = False,
+        exception_handler: "ExceptionHandler | None" = None,
         # Specification args
         title: str | None = None,
         description: str | None = None,
@@ -196,6 +199,7 @@ def subscriber(
         max_workers: int = ...,
         ack_policy: AckPolicy = EMPTY,
         no_reply: bool = False,
+        exception_handler: "ExceptionHandler | None" = None,
         # Specification args
         title: str | None = None,
         description: str | None = None,
@@ -246,6 +250,7 @@ def subscriber(
         max_workers: int = ...,
         ack_policy: AckPolicy = EMPTY,
         no_reply: bool = False,
+        exception_handler: "ExceptionHandler | None" = None,
         # Specification args
         title: str | None = None,
         description: str | None = None,
@@ -296,6 +301,7 @@ def subscriber(
         max_workers: int | None = None,
         ack_policy: AckPolicy = EMPTY,
         no_reply: bool = False,
+        exception_handler: "ExceptionHandler | None" = None,
         # Specification args
         title: str | None = None,
         description: str | None = None,
@@ -351,6 +357,7 @@ def subscriber(
         max_workers: int | None = None,
         ack_policy: AckPolicy = EMPTY,
         no_reply: bool = False,
+        exception_handler: "ExceptionHandler | None" = None,
         # Specification args
         title: str | None = None,
         description: str | None = None,
@@ -552,6 +559,9 @@ def subscriber(
             no_ack: Whether to disable **FastStream** auto acknowledgement logic or not.
             ack_policy: Acknowledgement policy for the subscriber.
             no_reply: Whether to disable **FastStream** RPC and Reply To auto responses or not.
+            exception_handler: Sync or async exception handler called before
+                the broker's handler. Return True if handled, otherwise False
+                to delegate to the broker.
             title: Specification subscriber object title.
             description: Specification subscriber object description. " "Uses decorated docstring as default.
             include_in_schema: Whetever to include operation in Specification schema or not.
@@ -595,6 +605,7 @@ def subscriber(
             partitions=partitions,
             ack_policy=ack_policy,
             no_reply=no_reply,
+            exception_handler=exception_handler,
             config=cast("KafkaBrokerConfig", self.config),
             # Specification
             title_=title,

faststream/kafka/subscriber/factory.py

@@ -19,6 +19,7 @@
 if TYPE_CHECKING:
     from aiokafka.abc import ConsumerRebalanceListener
 
+    from faststream._internal.types import ExceptionHandler
     from faststream.kafka.configs import KafkaBrokerConfig
     from faststream.kafka.schemas import TopicPartition
 
@@ -38,6 +39,7 @@ def create_subscriber(
     ack_policy: "AckPolicy",
     max_workers: int,
     no_reply: bool,
+    exception_handler: "ExceptionHandler | None",
     config: "KafkaBrokerConfig",
     # Specification args
     title_: str | None,
@@ -66,6 +68,7 @@ def create_subscriber(
         listener=listener,
         pattern=pattern,
         no_reply=no_reply,
+        exception_handler=exception_handler,
         _outer_config=config,
         _ack_policy=ack_policy,
     )