fix: refuse OutboxResponse on a closed session, write the recorded outcome on config errors, take correlation_id from headers (#225)

Author: lesnik512

docs/operations/troubleshooting.md

@@ -17,6 +17,7 @@ owns the underlying design.
 | [`AckPolicy.ACK_FIRST` raises `ValueError` at registration](#ackpolicyack_first-raises-valueerror-at-registration) | By design (would defeat outbox reliability) |
 | [`OutboxResponse(...)` + foreign-publisher decorator logs a configuration error](#outboxresponse-foreign-publisher-decorator-config-error) | By design (dual-fire footgun) |
 | [Chained `OutboxResponse` retries after the handler "succeeded"](#outboxresponse-relay-publish-failure) | Follow-on publish fails post-handler; nacks the inbound row |
+| [`OutboxResponse` raises "no open transaction"](#outboxresponse-no-open-transaction) | The response's session closed before the publish |
 | [`validate_schema()` raises `ImportError`](#validate_schema-raises-importerror) | `[validate]` extra not installed |
 
 ## `event=lease_lost` recurring in logs { #event-lease_lost-recurring-in-logs }
@@ -245,14 +246,14 @@ See [Subscriber § Ack policy](../usage/subscriber.md#ack-policy).
 
 A handler with both `@kafka_pub` and an `OutboxResponse(...)` return
 value logs an ERROR on every dispatch:
-`Outbox configuration error (fix required; row left to lease-expiry retry)`.
+`Outbox configuration error (fix required; row retried per its strategy)`.
 
 This is by design. The combination would both insert a row into the
 outbox *and* publish to Kafka, a dual-fire that doubles delivery. The
 subscriber refuses the chain composition after the handler returns. The
-worker logs the error and moves on without nacking, so the retry
-strategy is not consulted. The row's lease expires, a later fetch
-reclaims it, and the cycle repeats until the configuration is fixed.
+worker logs the error and writes the nack the ack middleware recorded, so
+the row follows its retry strategy until the configuration is fixed (or,
+with `max_deliveries` set, until it goes terminal).
 
 To confirm, inspect the handler's decorator stack and return type.
 
@@ -291,6 +292,24 @@ a redelivery's insert is a no-op.
 
 See [Publisher § Chained publishing](../usage/publisher.md#chained-publishing).
 
+## `OutboxResponse` raises "no open transaction" { #outboxresponse-no-open-transaction }
+
+A handler that returns `OutboxResponse(...)` logs an ERROR on every
+dispatch saying the response's session has no open transaction, and the
+inbound row keeps retrying.
+
+The response is published after the handler returns. A session from
+`Depends(...)`, or from an `async with` block inside the handler, is
+already closed by then, so the follow-on row would sit in a transaction
+nobody commits. The subscriber refuses to publish it and the inbound row
+follows its retry strategy.
+
+To fix it, open the session in a broker middleware so it outlives the
+handler, or call `broker.publish(..., session=session)` inside the
+handler instead of returning `OutboxResponse`.
+
+See [Publisher § Chained publishing](../usage/publisher.md#chained-publishing).
+
 ## `validate_schema()` raises `ImportError` { #validate_schema-raises-importerror }
 
 Calling `await broker.validate_schema()` raises:

docs/usage/fastapi.md

@@ -96,9 +96,12 @@ delivery, opened in a `session.begin()` block, committed on handler return,
 rolled back on exception.
 
 A handler's `AsyncSession` is therefore resolved exactly as in an HTTP
-route (a fresh session per delivery, not a shared instance), and
-`OutboxResponse(session=...)` commits the follow-on row with the handler's
-domain writes. See [Chained
+route (a fresh session per delivery, not a shared instance). Publish
+follow-on rows with `router.broker.publish(..., session=session)` inside
+the handler so they commit with its domain writes. Returning
+`OutboxResponse(session=session)` with this session does not work: the
+dependency is torn down before the response is published, so the dispatch
+raises and the inbound row is retried. See [Chained
 publishing](./publisher.md#chained-publishing).
 
 ## Annotated context shortcuts

docs/usage/publisher.md

@@ -149,44 +149,56 @@ with your domain writes):
 
 !!! note "The `session` must outlive the handler return"
     The returned `OutboxResponse` is published after the handler
-    returns, so its `session` must still be open at that point. The
-    requirement is about session *lifetime*, not any particular framework:
-    provide the session through a dependency that the framework tears down
-    *after* the response flow. FastAPI's `Depends(get_session)` or
-    FastStream's own `Depends` / `Context` session both do this. Opening
-    your own `async with session_factory() as session:` inside the handler
-    does not work here: that session closes on `return`, before the row
-    is inserted. In that case, call `broker.publish(..., session=session)`
-    directly inside the `async with` instead (see
-    [§ Not a relay decorator](#not-a-relay-decorator)).
+    returns, so its `session` must still hold an open transaction at that
+    point. A session from `Depends(...)` does not: FastStream and FastAPI
+    both tear dependencies down inside the handler call, before the
+    response is published. The same goes for an
+    `async with session_factory() as session:` block inside the handler.
+    In those cases the dispatch raises `RuntimeError` naming the missing
+    transaction, logs it at ERROR, and retries the inbound row. Either
+    open the session in a broker middleware, as below, or call
+    `broker.publish(..., session=session)` inside the handler instead
+    (see [§ Not a relay decorator](#not-a-relay-decorator)).
+
+A broker middleware's `on_receive` / `after_processed` hooks wrap the
+whole dispatch, response publish included, so a session opened there
+covers both the handler's writes and the follow-on row. The commit runs
+before the ack decision, so a failed commit retries the inbound row:
 
 ```python
-from fastapi import Depends
+from faststream import BaseMiddleware, Context
 from sqlalchemy.ext.asyncio import AsyncSession
 
-from faststream_outbox import OutboxResponse
-from faststream_outbox.annotations import OutboxMessage
+from faststream_outbox import OutboxBroker, OutboxResponse
 
-# `get_session` is your app's session dependency — the same one your HTTP
-# routes use. `Depends(...)` resolves inside a handler only under the FastAPI
-# integration; see the engine / sessionmaker setup in the FastAPI guide.
+
+class SessionPerMessage(BaseMiddleware):
+    async def on_receive(self) -> None:
+        self._begin = session_factory.begin()
+        self.context.set_local("session", await self._begin.__aenter__())
+        await super().on_receive()
+
+    async def after_processed(self, exc_type=None, exc_val=None, exc_tb=None) -> bool | None:
+        await self._begin.__aexit__(exc_type, exc_val, exc_tb)
+        return await super().after_processed(exc_type, exc_val, exc_tb)
+
+
+broker = OutboxBroker(engine, outbox_table=outbox_table, middlewares=[SessionPerMessage])
 
 
 @broker.subscriber("orders")
-async def handle(
-    msg: OutboxMessage,
-    session: AsyncSession = Depends(get_session),
-) -> OutboxResponse:
-    ...  # process inbound
+async def handle(body: dict, session: AsyncSession = Context()) -> OutboxResponse:
+    ...  # domain writes on `session`
     return OutboxResponse(
         body={"chained": True},
         queue="downstream",
         session=session,
     )
 ```
 
-`correlation_id` propagates from the inbound message if you don't set one
-explicitly, which is useful for trace stitching. Plain returns (`None`, `dict`,
+`correlation_id` comes from `headers["correlation_id"]` when you set the
+header but not the argument, and otherwise from the inbound message, which
+is useful for trace stitching. Plain returns (`None`, `dict`,
 etc.) are silently skipped, so handlers that don't want to chain just
 return normally.
 

faststream_outbox/publisher/fake.py

@@ -19,7 +19,7 @@
 
 from faststream.response.publish_type import PublishType
 
-from faststream_outbox.response import OutboxPublishCommand
+from faststream_outbox.response import OutboxPublishCommand, _OutboxConfigError
 
 
 if typing.TYPE_CHECKING:
@@ -43,6 +43,14 @@ async def _publish(
     ) -> typing.Any:
         if not isinstance(cmd, OutboxPublishCommand):
             return None
+        if not cmd.session.in_transaction():
+            msg = (
+                "OutboxResponse is published after the handler returns, but its session has no open "
+                "transaction by then, so the follow-on row would never commit. Open the session and its "
+                "transaction in a broker middleware so they outlive the handler, or call "
+                "broker.publish(..., session=session) inside the handler instead."
+            )
+            raise _OutboxConfigError(msg)
         cmd.publish_type = PublishType.REPLY
         # ty diagnostic: ``producer.publish`` narrows to ``OutboxPublishCommand``
         # while ``PublisherMiddleware`` is generic over the base ``PublishCommand``.

faststream_outbox/response.py

@@ -19,6 +19,8 @@
 from faststream.response.response import BatchPublishCommand, PublishCommand, Response
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from faststream_outbox.message import CORRELATION_ID_HEADER
+
 
 # Matches the ``queue`` column width in ``make_outbox_table`` (``String(255)``).
 _MAX_QUEUE_LENGTH = 255
@@ -68,6 +70,14 @@ def _validate_publish_args(
         raise ValueError(msg)
 
 
+# Marker exception raised by programming guards on the dispatch path (e.g. the
+# OutboxResponse + foreign-publisher dual-fire guard). Inherits from
+# ``RuntimeError`` so ``pytest.raises(RuntimeError, ...)`` in tests catches it, but
+# is a distinct subclass so ``consume()`` and ``dispatch_one`` can re-raise it while
+# still swallowing plain handler-raised ``RuntimeError``s.
+class _OutboxConfigError(RuntimeError): ...
+
+
 class OutboxPublishCommand(BatchPublishCommand):
     """Outbox-specific publish command: carries session + scheduling fields end-to-end."""
 
@@ -166,9 +176,13 @@ class OutboxResponse(Response):
     re-runs the same validator on ``as_publish_command()``, so it stays the authoritative
     single source of truth and the eager checks can never drift from it.
 
-    ``correlation_id`` defaults to the inbound message's correlation_id when not set
-    — FastStream's ``SubscriberUsecase.process_message`` does the inheritance before
-    publishing the response.
+    ``correlation_id`` defaults to ``headers["correlation_id"]``, then to the inbound
+    message's correlation_id, which FastStream's ``SubscriberUsecase.process_message``
+    fills in before publishing the response.
+
+    The response is published after the handler returns, so ``session`` must still
+    hold an open transaction then; otherwise the dispatch raises ``RuntimeError`` and
+    the inbound row is retried.
     """
 
     def __init__(
@@ -190,6 +204,8 @@ def __init__(
             activate_in=activate_in,
             activate_at=activate_at,
         )
+        if correlation_id is None and headers:
+            correlation_id = headers.get(CORRELATION_ID_HEADER)
         super().__init__(body=body, headers=headers, correlation_id=correlation_id)
         self.queue = queue
         self.session = session

faststream_outbox/subscriber/usecase.py

@@ -43,7 +43,7 @@
 from faststream_outbox.message import ENVELOPE_MANAGED_HEADERS, OutboxInnerMessage, Retry, Terminal
 from faststream_outbox.parser.parser import OutboxParser
 from faststream_outbox.publisher.fake import OutboxFakePublisher
-from faststream_outbox.response import OutboxResponse
+from faststream_outbox.response import OutboxResponse, _OutboxConfigError
 from faststream_outbox.subscriber.config import OutboxSubscriberConfig, OutboxSubscriberSpecificationConfig
 
 
@@ -79,14 +79,6 @@ class _PendingFlush:
     tags: dict[str, typing.Any]
 
 
-# Marker exception raised by programming guards inside ``process_message`` (e.g.
-# the OutboxResponse + foreign-publisher dual-fire guard). Inherits from
-# ``RuntimeError`` so ``pytest.raises(RuntimeError, ...)`` in tests catches it, but
-# is a distinct subclass so ``consume()`` and ``dispatch_one`` can re-raise it while
-# still swallowing plain handler-raised ``RuntimeError``s.
-class _OutboxConfigError(RuntimeError): ...
-
-
 # Cap for the ``last_exception`` string written to the DLQ. Some exceptions carry
 # huge payloads (validation errors with the full request body, asyncpg ``DataError``
 # with the rejected row, etc.). An unbounded ``repr`` would extend the writer
@@ -643,11 +635,11 @@ async def _worker_inner(self, *, writer_conn: "AsyncConnection | None") -> None:
                     # P18: a config error (e.g. OutboxResponse + foreign publisher) is not a
                     # connection failure. Letting it propagate to _run_with_reconnect would tear
                     # down the writer connection and back off (up to 30s), throttling unrelated
-                    # rows. Log it and continue; the row's lease expires and it is reclaimed.
-                    # Fix the configuration to stop the error.
+                    # rows. dispatch_one already wrote the recorded outcome, so the retry
+                    # strategy applies; log it and continue until the configuration is fixed.
                     self._log(
                         log_level=logging.ERROR,
-                        message=f"Outbox configuration error (fix required; row left to lease-expiry retry): {e!r}",
+                        message=f"Outbox configuration error (fix required; row retried per its strategy): {e!r}",
                         exc_info=e,
                     )
                 finally:
@@ -722,14 +714,12 @@ async def dispatch_one(  # linear pipeline: guard, consume, branch on outcome, f
         try:
             await self.consume(row)
         except _OutboxConfigError:
-            # _OutboxConfigError escaping consume() is a programming guard (e.g.
-            # the OutboxResponse + foreign-publisher dual-fire guard). Re-raise so
-            # callers see the error immediately — sync test dispatch surfaces it
-            # from outbox.publish; the worker loop's outer except re-raises it via
-            # _run_with_reconnect so it's logged at ERROR level with a reconnect
-            # backoff rather than silently ignored.
+            # A programming guard tripped after the handler ran. Write the outcome the ack
+            # middleware recorded, then re-raise: sync test dispatch surfaces it from
+            # outbox.publish, and _worker_inner logs it at ERROR. Inline (buffer=None),
+            # because a buffered row would be task_done'd twice once the error unwinds.
             row.handler_duration_seconds = time.perf_counter() - start_perf
-            self._emit_dispatch_abandoned(row)
+            await self._write_outcome(row, base, buffer=None, writer_conn=writer_conn)
             raise
         except Exception as e:  # noqa: BLE001
             # No ack/nack metric here: the row was never marked terminal/retry, so
@@ -745,6 +735,18 @@ async def dispatch_one(  # linear pipeline: guard, consume, branch on outcome, f
             self._emit_dispatch_abandoned(row)
             return False
         row.handler_duration_seconds = time.perf_counter() - start_perf
+        return await self._write_outcome(row, base, buffer=buffer, writer_conn=writer_conn)
+
+    async def _write_outcome(
+        self,
+        row: OutboxInnerMessage,
+        base: dict[str, typing.Any],
+        *,
+        buffer: "list[_PendingFlush] | None",
+        writer_conn: "AsyncConnection | None",
+    ) -> bool:
+        """Write the outcome recorded during consume and emit its metric. Returns True iff buffered."""
+        logger = self._outer_config.logger.logger.logger if self._outer_config.logger else None
         # Shutdown race: SubscriberUsecase.consume() returns None without invoking
         # process_message when self.running has been flipped to False by stop().
         # Detecting that here lets us preserve the row instead of falling through
@@ -1006,10 +1008,9 @@ async def consume(self, msg: OutboxInnerMessage) -> typing.Any:
         ``OutboxResponse + foreign-publisher`` dual-fire check in
         ``process_message``) raise ``_OutboxConfigError`` (a ``RuntimeError``
         subclass) so the guard is distinguishable from handler-raised
-        ``RuntimeError``s. Re-raising here lets ``dispatch_one`` surface it to
-        the caller (sync test dispatch exposes it from ``outbox.publish``; the
-        worker loop re-raises so the error is logged and the lease expires rather
-        than silently swallowing a configuration mistake).
+        ``RuntimeError``s. Re-raising here lets ``dispatch_one`` write the recorded
+        outcome and surface the error: sync test dispatch exposes it from
+        ``outbox.publish``, and the worker loop logs it at ERROR.
 
         # Upstream equivalent (extended):
         #   SubscriberUsecase.consume

tests/test_fake.py

@@ -6,7 +6,7 @@
 import warnings as _warnings
 from collections.abc import Awaitable, Callable, Mapping
 from types import TracebackType
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from faststream import BaseMiddleware, FastStream, StreamMessage
@@ -2375,3 +2375,63 @@ async def handle(body: str) -> None:
         await _wait_until(lambda: not test_broker.fake_client.rows, timeout=5.0)
 
     assert sorted(handled) == sorted(f"m{i}" for i in range(30))
+
+
+async def test_handler_returning_outbox_response_takes_correlation_id_from_headers() -> None:
+    broker = _make_broker()
+    downstream_cors: list[str] = []
+
+    @broker.subscriber("orders")
+    async def handle(body: dict) -> OutboxResponse:
+        del body
+        return OutboxResponse(
+            body={"v": 1}, queue="downstream", session=_fake_session(), headers={"correlation_id": "from-header"}
+        )
+
+    @broker.subscriber("downstream")
+    async def handle_next(body: dict, correlation_id: str = _Context("message.correlation_id")) -> None:
+        del body
+        downstream_cors.append(correlation_id)
+
+    async with TestOutboxBroker(broker):
+        await broker.publish({"x": 1}, queue="orders", correlation_id="inbound")  # ty: ignore[missing-argument]
+
+    assert downstream_cors == ["from-header"]
+
+
+async def test_outbox_response_on_a_session_without_a_transaction_retries_the_inbound_row() -> None:
+    broker = _make_broker()
+    closed_session = _fake_session()
+    closed_session.in_transaction = MagicMock(return_value=False)
+
+    @broker.subscriber("orders")
+    async def handle(body: dict) -> OutboxResponse:
+        return OutboxResponse(body=body, queue="downstream", session=closed_session)
+
+    test_broker = TestOutboxBroker(broker)
+    async with test_broker:
+        with pytest.raises(RuntimeError, match="transaction"):
+            await broker.publish({"x": 1}, queue="orders")  # ty: ignore[missing-argument]
+
+    assert [(r.queue, r.attempts_count, r.acquired_token) for r in test_broker.fake_client.rows] == [
+        ("orders", 1, None)
+    ]
+
+
+async def test_config_error_under_manual_ack_retries_the_row() -> None:
+    broker = _make_broker()
+    closed_session = _fake_session()
+    closed_session.in_transaction = MagicMock(return_value=False)
+
+    @broker.subscriber("orders", ack_policy=AckPolicy.MANUAL)
+    async def handle(body: dict) -> OutboxResponse:
+        return OutboxResponse(body=body, queue="downstream", session=closed_session)
+
+    test_broker = TestOutboxBroker(broker)
+    async with test_broker:
+        with pytest.raises(RuntimeError, match="transaction"):
+            await broker.publish({"x": 1}, queue="orders")  # ty: ignore[missing-argument]
+
+    assert [(r.queue, r.attempts_count, r.acquired_token) for r in test_broker.fake_client.rows] == [
+        ("orders", 1, None)
+    ]