feat: add subscriber and broker exception handlers

faststream/_internal/broker/broker.py

@@ -54,6 +54,8 @@ def __init__(
             routers=routers,
             config=config,
         )
+        self._exception_handler = config.broker_exception_handler
+
         self.specification = specification
 
         self.running = False

faststream/_internal/configs/broker.py

@@ -9,14 +9,19 @@
 from faststream._internal.logger import LoggerState
 from faststream._internal.producer import ProducerProto, ProducerUnset
 from faststream._internal.types import IdGenerator
+from faststream._internal.utils import apply_types
 from faststream.message import gen_cor_id
 
 if TYPE_CHECKING:
     from fast_depends.dependencies import Dependant
 
     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,
+    )
     from faststream.middlewares import AckPolicy
 
 
@@ -29,6 +34,7 @@ class BrokerConfig:
     broker_parser: Optional["CustomCallable"] = None
     broker_decoder: Optional["CustomCallable"] = None
     broker_codec: Optional["CodecProto"] = None
+    broker_exception_handler: Optional["AsyncExceptionHandler"] = None
 
     producer: "ProducerProto[Any]" = field(default_factory=ProducerUnset)
     logger: "LoggerState" = field(default_factory=LoggerState)
@@ -45,6 +51,13 @@ def __post_init__(self) -> None:
         # untyped callers still pass a generator: the first subscriber would spend it
         self.broker_dependencies = tuple(self.broker_dependencies)
 
+        if self.broker_exception_handler is not None:
+            self.broker_exception_handler = apply_types(
+                self.broker_exception_handler,
+                serializer_cls=self.fd_config._serializer,
+                context__=self.context,
+            )
+
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(id: {id(self)})"
 

faststream/_internal/configs/endpoint.py

@@ -2,10 +2,11 @@
 from typing import TYPE_CHECKING
 
 from faststream._internal.constants import EMPTY
+from faststream._internal.utils import apply_types
 from faststream.middlewares import AckPolicy
 
 if TYPE_CHECKING:
-    from faststream._internal.types import AsyncCallable
+    from faststream._internal.types import AsyncCallable, AsyncExceptionHandler
 
     from .broker import BrokerConfig
 
@@ -23,12 +24,21 @@ class PublisherUsecaseConfig(EndpointConfig):
 @dataclass(kw_only=True)
 class SubscriberUsecaseConfig(EndpointConfig):
     no_reply: bool = False
+    exception_handler: "AsyncExceptionHandler | None" = None
 
     _ack_policy: AckPolicy = field(default_factory=lambda: EMPTY, repr=False)
 
     parser: "AsyncCallable" = field(init=False)
     decoder: "AsyncCallable" = field(init=False)
 
+    def __post_init__(self) -> None:
+        if self.exception_handler is not None:
+            self.exception_handler = apply_types(
+                self.exception_handler,
+                serializer_cls=self._outer_config.fd_config._serializer,
+                context__=self._outer_config.context,
+            )
+
     @property
     def auto_ack_disabled(self) -> bool:
         return self.ack_policy is AckPolicy.MANUAL

faststream/_internal/endpoint/subscriber/usecase.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from collections.abc import AsyncIterator, Callable, Iterable, Sequence
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
 from contextlib import AbstractContextManager, AsyncExitStack, ExitStack, nullcontext
 from itertools import chain
 from types import TracebackType
@@ -45,6 +45,7 @@
     from faststream._internal.endpoint.publisher import PublisherProto
     from faststream._internal.parser import CodecProto
     from faststream._internal.types import (
+        AsyncExceptionHandler,
         AsyncFilter,
         BrokerMiddleware,
         CustomCallable,
@@ -102,6 +103,7 @@ def __init__(
         self._no_reply = config.no_reply
         self._parser = config.parser
         self._decoder = config.decoder
+        self._exception_handler = config.exception_handler
 
         self.ack_policy = config.ack_policy
         self.__auto_ack_disabled = config.auto_ack_disabled
@@ -328,6 +330,38 @@ def real_wrapper(
 
         return real_wrapper(func)
 
+    @property
+    def _app_exception_handler(self) -> "AsyncExceptionHandler | None":
+        app_handler: AsyncExceptionHandler | None = getattr(
+            self._outer_config.context.get("app"),
+            "_exception_handler",
+            None,
+        )
+        return app_handler
+
+    async def _handle_exception(self, exc: BaseException) -> bool:
+        context = self._outer_config.context
+
+        local_handler = cast(
+            "Callable[..., Awaitable[bool]] | None",
+            self._exception_handler,
+        )
+        if local_handler is not None and await local_handler(exc, context__=context):
+            return True
+
+        broker_handler = cast(
+            "Callable[..., Awaitable[bool]] | None",
+            self._outer_config.broker_exception_handler,
+        )
+        if broker_handler is not None and await broker_handler(exc, context__=context):
+            return True
+
+        app_handler = self._app_exception_handler
+        if app_handler is not None:
+            return await app_handler(exc)
+
+        return False
+
     async def consume(self, msg: MsgType) -> Any:
         """Consume a message asynchronously."""
         if not self.running:
@@ -347,9 +381,15 @@ 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 is not None
+                or self._outer_config.broker_exception_handler is not None
+                or self._app_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

@@ -43,6 +43,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/app.py

@@ -19,6 +19,7 @@
 from faststream._internal.context import ContextRepo
 from faststream._internal.di import FastDependsConfig
 from faststream._internal.logger import logger
+from faststream._internal.utils import apply_types, to_async
 from faststream.asgi.app import AsgiFastStream
 
 if TYPE_CHECKING:
@@ -31,6 +32,7 @@
         SettingField,
     )
     from faststream._internal.broker import BrokerUsecase
+    from faststream._internal.types import AsyncExceptionHandler, ExceptionHandler
     from faststream.asgi import AsyncAPIRoute
     from faststream.asgi.types import ASGIApp
     from faststream.specification.base import SpecificationFactory
@@ -49,6 +51,7 @@ def __init__(
         provider: Optional["Provider"] = None,
         serializer: Optional["SerializerProto"] = EMPTY,
         context: ContextRepo | None = None,
+        exception_handler: "ExceptionHandler | None" = None,
         lifespan: Optional["Lifespan"] = None,
         on_startup: Sequence["AnyCallable"] = (),
         after_startup: Sequence["AnyCallable"] = (),
@@ -72,6 +75,14 @@ def __init__(
             specification=specification,
         )
 
+        self._exception_handler: AsyncExceptionHandler | None = None
+        if exception_handler is not None:
+            self._exception_handler = apply_types(
+                to_async(exception_handler),
+                serializer_cls=self.config._serializer,
+                context__=self.context,
+            )
+
         self._should_exit = False
 
     async def run(

faststream/confluent/subscriber/config.py

@@ -31,6 +31,7 @@ class KafkaSubscriberConfig(SubscriberUsecaseConfig):
     connection_data: dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
+        super().__post_init__()
         self.connection_data["enable_auto_commit"] = self.ack_first
 
     @property

faststream/kafka/broker/broker.py

@@ -26,6 +26,7 @@
 from faststream._internal.context.repository import ContextRepo
 from faststream._internal.di import FastDependsConfig
 from faststream._internal.types import IdGenerator
+from faststream._internal.utils import to_async
 from faststream._internal.utils.data import filter_by_dict
 from faststream.exceptions import IncorrectState
 from faststream.kafka._compat import AIOKAFKA_V013, validate_client_rack
@@ -60,6 +61,7 @@
     from faststream._internal.types import (
         BrokerMiddleware,
         CustomCallable,
+        ExceptionHandler,
     )
     from faststream.kafka.message import KafkaMessage
     from faststream.kafka.types import KafkaSendableMessage
@@ -237,6 +239,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] = (),
@@ -357,6 +360,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]]):
@@ -468,6 +475,9 @@ def __init__(
                 broker_decoder=decoder,
                 broker_codec=codec,
                 broker_parser=parser,
+                broker_exception_handler=(
+                    to_async(exception_handler) if exception_handler is not None else None
+                ),
                 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,
@@ -553,6 +560,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.
@@ -596,6 +606,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/config.py

@@ -35,6 +35,7 @@ class KafkaSubscriberConfig(SubscriberUsecaseConfig):
     partitions: Iterable["TopicPartition"] = field(default_factory=list)
 
     def __post_init__(self) -> None:
+        super().__post_init__()
         self.connection_args["enable_auto_commit"] = self.ack_first
 
     @property

faststream/kafka/subscriber/factory.py

@@ -4,6 +4,7 @@
 
 from faststream._internal.constants import EMPTY
 from faststream._internal.endpoint.subscriber.call_item import CallsCollection
+from faststream._internal.utils import to_async
 from faststream.exceptions import SetupError
 from faststream.kafka._compat import validate_client_rack
 from faststream.kafka.schemas import Topic
@@ -21,6 +22,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
 
@@ -40,6 +42,7 @@ def create_subscriber(
     ack_policy: "AckPolicy",
     max_workers: int,
     no_reply: bool,
+    exception_handler: "ExceptionHandler | None",
     config: "KafkaBrokerConfig",
     # Specification args
     title_: str | None,
@@ -71,6 +74,9 @@ def create_subscriber(
         listener=listener,
         pattern=pattern,
         no_reply=no_reply,
+        exception_handler=(
+            to_async(exception_handler) if exception_handler is not None else None
+        ),
         _outer_config=config,
         _ack_policy=ack_policy,
     )