chore: import upstream snapshot with attribution

application/streaming/async_broadcast_channel.py

@@ -0,0 +1,119 @@
+"""Async Redis pub/sub Topic for the native-async SSE reader.
+
+Event-loop twin of :class:`application.streaming.broadcast_channel.Topic`.
+Same contract — ``subscribe`` yields ``None`` on poll timeout (so the
+caller can emit keepalives / run the watchdog) and ``bytes`` per delivered
+message, fires ``on_subscribe`` once after Redis acks SUBSCRIBE, and tears
+the pubsub down cleanly on client disconnect — but awaitable so an idle
+stream costs a coroutine instead of a WSGI thread.
+
+Publishing stays on the sync side (the producer writes via
+``broadcast_channel.Topic.publish``); this is read-only fan-out.
+"""
+
+from __future__ import annotations
+
+import inspect
+import logging
+from typing import AsyncIterator, Awaitable, Callable, Optional, Union
+
+import anyio
+
+from application.streaming.async_redis import get_async_redis_instance
+
+logger = logging.getLogger(__name__)
+
+OnSubscribe = Callable[[], Union[None, Awaitable[None]]]
+
+
+class AsyncTopic:
+    """An async pub/sub channel identified by a string name."""
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+    async def subscribe(
+        self,
+        on_subscribe: Optional[OnSubscribe] = None,
+        poll_timeout: float = 1.0,
+    ) -> AsyncIterator[Optional[bytes]]:
+        """Subscribe to the topic; yield raw payloads or ``None`` on tick.
+
+        ``on_subscribe`` runs (and is awaited if it returns a coroutine)
+        after Redis acks SUBSCRIBE — use it to seed snapshot state that
+        must be ordered after the subscriber is live but before the first
+        live message is processed. If Redis is unavailable, returns
+        immediately without yielding so the caller can fall back to a
+        direct snapshot read. Cleanly unsubscribes on close / disconnect.
+        """
+        redis = await get_async_redis_instance()
+        if redis is None:
+            logger.debug(
+                "Async Redis unavailable; subscribe to %s yielded nothing",
+                self.name,
+            )
+            return
+        pubsub = redis.pubsub()
+        on_subscribe_fired = False
+        try:
+            try:
+                await pubsub.subscribe(self.name)
+            except Exception:
+                # Transient subscribe failure is treated like "Redis
+                # unavailable": yield nothing, let the caller fall back to
+                # its own snapshot read. The finally block still tears the
+                # pubsub down cleanly.
+                logger.exception("async pubsub.subscribe failed for %s", self.name)
+                return
+            while True:
+                try:
+                    msg = await pubsub.get_message(timeout=poll_timeout)
+                except Exception:
+                    logger.exception(
+                        "async pubsub.get_message failed for %s", self.name
+                    )
+                    return
+                if msg is None:
+                    yield None
+                    continue
+                msg_type = msg.get("type")
+                if msg_type == "subscribe":
+                    if not on_subscribe_fired and on_subscribe is not None:
+                        try:
+                            result = on_subscribe()
+                            if inspect.isawaitable(result):
+                                await result
+                        except Exception:
+                            logger.exception(
+                                "on_subscribe callback failed for %s", self.name
+                            )
+                    on_subscribe_fired = True
+                    continue
+                if msg_type != "message":
+                    continue
+                data = msg.get("data")
+                if data is None:
+                    continue
+                yield data if isinstance(data, bytes) else str(data).encode("utf-8")
+        finally:
+            # Client disconnect cancels this generator at the ``await
+            # get_message`` above; without shielding, the cancellation could
+            # re-fire mid-teardown and skip ``aclose()``, leaking the pooled
+            # connection back to nothing. Shield so unsubscribe + aclose
+            # always complete and the connection returns to the pool.
+            with anyio.CancelScope(shield=True):
+                if on_subscribe_fired:
+                    try:
+                        await pubsub.unsubscribe(self.name)
+                    except Exception:
+                        logger.debug(
+                            "async pubsub unsubscribe error for %s",
+                            self.name,
+                            exc_info=True,
+                        )
+                try:
+                    await pubsub.aclose()
+                except Exception:
+                    logger.debug(
+                        "async pubsub close error for %s", self.name, exc_info=True
+                    )

application/streaming/async_event_replay.py

@@ -0,0 +1,232 @@
+"""Native-async snapshot+tail iterator for chat-stream reconnect.
+
+The sole reconnect reader: a ``: connected`` prelude, a snapshot flush
+inside the SUBSCRIBE-ack callback, a dedup'd live tail, keepalive +
+producer-liveness watchdog, and close-on-terminal — all as an async
+generator driven off the event loop instead of a WSGI thread.
+
+Wire format, dedup floor, and terminal detection come from
+``event_replay`` (``read_snapshot_lines``, ``format_sse_event``,
+``_decode_pubsub_message``, ``_payload_is_terminal``,
+``_check_producer_liveness``), the same primitives the producer's journal
+writes through — so the reader and writer cannot drift on wire shape. The
+only sync I/O (snapshot read, watchdog DB probe) is pushed to a worker
+thread via ``anyio.to_thread`` so it never blocks the loop.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import AsyncIterator, Optional
+
+import anyio
+
+from application.streaming.async_broadcast_channel import AsyncTopic
+from application.streaming.event_replay import (
+    DEFAULT_KEEPALIVE_SECONDS,
+    DEFAULT_POLL_TIMEOUT_SECONDS,
+    DEFAULT_PRODUCER_IDLE_SECONDS,
+    DEFAULT_WATCHDOG_INTERVAL_SECONDS,
+    _check_producer_liveness,
+    _decode_pubsub_message,
+    _payload_is_terminal,
+    format_sse_event,
+    read_snapshot_lines,
+)
+from application.streaming.keys import message_topic_name
+
+logger = logging.getLogger(__name__)
+
+# The snapshot read and watchdog probe are the reader's only DB I/O; they run
+# in worker threads and each borrows a connection from the app-wide SQLAlchemy
+# pool (pool_size=10 + max_overflow=20 = 30). The reader scales to many
+# concurrent event-loop streams, so without a bound a burst of aligned
+# watchdog/snapshot ticks could exhaust the pool and starve every other route.
+# Cap the reader's concurrent DB-thread usage well below the pool so it can
+# never monopolise it; excess ticks queue briefly (watchdog cadence is 5s, so
+# a short queue delay is harmless).
+_MAX_CONCURRENT_DB_READS = 8
+_db_read_limiter: Optional[anyio.CapacityLimiter] = None
+
+
+def _get_db_read_limiter() -> anyio.CapacityLimiter:
+    """Lazily build the shared limiter on the (single) event loop.
+
+    Created on first use rather than at import so it binds to the running
+    loop; creation is synchronous, so the single-worker loop has no race.
+    """
+    global _db_read_limiter
+    if _db_read_limiter is None:
+        _db_read_limiter = anyio.CapacityLimiter(_MAX_CONCURRENT_DB_READS)
+    return _db_read_limiter
+
+
+async def build_message_event_stream_async(
+    message_id: str,
+    last_event_id: Optional[int] = None,
+    *,
+    user_id: Optional[str] = None,
+    keepalive_seconds: float = DEFAULT_KEEPALIVE_SECONDS,
+    poll_timeout_seconds: float = DEFAULT_POLL_TIMEOUT_SECONDS,
+    watchdog_interval_seconds: float = DEFAULT_WATCHDOG_INTERVAL_SECONDS,
+    producer_idle_seconds: float = DEFAULT_PRODUCER_IDLE_SECONDS,
+) -> AsyncIterator[str]:
+    """Yield SSE-formatted lines for one ``message_id`` reconnect stream.
+
+    First frame is ``: connected``; subsequent frames are snapshot rows,
+    live-tail events, or ``: keepalive`` comments. Runs until the client
+    disconnects or a terminal event is delivered.
+    """
+    yield ": connected\n\n"
+
+    replay_buffer: list[str] = []
+    max_replayed_seq: Optional[int] = last_event_id
+    replay_done = False
+    replay_failed = False
+    terminal_in_snapshot = False
+
+    async def _load_snapshot() -> None:
+        nonlocal max_replayed_seq, replay_failed, terminal_in_snapshot
+        try:
+            lines, max_seq, terminal = await anyio.to_thread.run_sync(
+                read_snapshot_lines,
+                message_id,
+                last_event_id,
+                user_id,
+                limiter=_get_db_read_limiter(),
+            )
+        except Exception:
+            logger.exception(
+                "Snapshot read failed for message_id=%s last_event_id=%s",
+                message_id,
+                last_event_id,
+            )
+            replay_failed = True
+            return
+        replay_buffer.extend(lines)
+        max_replayed_seq = max_seq
+        terminal_in_snapshot = terminal
+
+    async def _on_subscribe() -> None:
+        # SUBSCRIBE acked — Postgres reads from this point capture every
+        # committed row; pub/sub messages published after this are queued
+        # at the connection level until the loop polls again.
+        nonlocal replay_done
+        try:
+            await _load_snapshot()
+        finally:
+            replay_done = True
+
+    topic = AsyncTopic(message_topic_name(message_id))
+    last_keepalive = time.monotonic()
+    last_watchdog_check = float("-inf")
+    watchdog_synthetic_seq = -1
+
+    try:
+        async for payload in topic.subscribe(
+            on_subscribe=_on_subscribe,
+            poll_timeout=poll_timeout_seconds,
+        ):
+            # Flush snapshot exactly once after the SUBSCRIBE callback ran.
+            if replay_done and replay_buffer:
+                for line in replay_buffer:
+                    yield line
+                replay_buffer.clear()
+                if terminal_in_snapshot:
+                    # Original stream already finished; tailing would just
+                    # emit keepalives forever.
+                    return
+
+            if replay_failed:
+                yield format_sse_event(
+                    {
+                        "type": "error",
+                        "error": "Stream replay failed; please refresh to load the latest state.",
+                        "code": "snapshot_failed",
+                        "message_id": message_id,
+                    },
+                    sequence_no=-1,
+                )
+                return
+
+            now = time.monotonic()
+            if payload is None:
+                # Idle tick — gate the watchdog on ``replay_done`` so we
+                # don't race the snapshot read on the first iteration.
+                if (
+                    replay_done
+                    and watchdog_interval_seconds >= 0
+                    and now - last_watchdog_check >= watchdog_interval_seconds
+                ):
+                    last_watchdog_check = now
+                    terminal_payload = await anyio.to_thread.run_sync(
+                        _check_producer_liveness,
+                        message_id,
+                        user_id,
+                        producer_idle_seconds,
+                        limiter=_get_db_read_limiter(),
+                    )
+                    if terminal_payload is not None:
+                        yield format_sse_event(
+                            terminal_payload,
+                            sequence_no=watchdog_synthetic_seq,
+                        )
+                        return
+                if now - last_keepalive >= keepalive_seconds:
+                    yield ": keepalive\n\n"
+                    last_keepalive = now
+                continue
+
+            envelope = _decode_pubsub_message(payload)
+            if envelope is None:
+                continue
+            seq = envelope.get("sequence_no")
+            inner = envelope.get("payload")
+            if (
+                not isinstance(seq, int)
+                or isinstance(seq, bool)
+                or not isinstance(inner, dict)
+            ):
+                continue
+            if max_replayed_seq is not None and seq <= max_replayed_seq:
+                # Snapshot already covered this id — drop the duplicate.
+                continue
+            yield format_sse_event(inner, seq)
+            max_replayed_seq = seq
+            last_keepalive = now
+            if _payload_is_terminal(inner, envelope.get("event_type")):
+                return
+
+        # Subscribe exited without yielding (Redis unavailable / subscribe
+        # raised). The snapshot half is still in Postgres — read it
+        # directly so a Redis-only outage doesn't cost the client their
+        # backlog. Gate on ``replay_done`` so we don't double-read.
+        if not replay_done:
+            await _load_snapshot()
+            replay_done = True
+        for line in replay_buffer:
+            yield line
+        replay_buffer.clear()
+        if replay_failed:
+            yield format_sse_event(
+                {
+                    "type": "error",
+                    "error": "Stream replay failed; please refresh to load the latest state.",
+                    "code": "snapshot_failed",
+                    "message_id": message_id,
+                },
+                sequence_no=-1,
+            )
+            return
+        if terminal_in_snapshot:
+            return
+    except Exception:
+        # GeneratorExit / CancelledError are BaseException subclasses, so a
+        # client disconnect bypasses this handler and propagates to close
+        # the inner AsyncTopic generator (tearing its pubsub down in that
+        # generator's finally). Only genuine bugs land here.
+        logger.exception(
+            "Async reconnect stream crashed for message_id=%s", message_id
+        )
+        return

application/streaming/async_redis.py

@@ -0,0 +1,47 @@
+"""Lazy async Redis client for the native-async SSE reader.
+
+Async twin of :func:`application.cache.get_redis_instance`. The
+Starlette-mounted reader (``application.api.async_sse``) tails pub/sub on
+the event loop, so it needs a ``redis.asyncio`` client rather than the
+sync one used by the producer side. The app runs a single ASGI worker /
+event loop, so a module-level singleton is sufficient and avoids
+reconnecting per request.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import redis.asyncio as aioredis
+
+from application.core.settings import settings
+
+logger = logging.getLogger(__name__)
+
+_async_redis: Optional[aioredis.Redis] = None
+_creation_failed = False
+
+
+async def get_async_redis_instance() -> Optional[aioredis.Redis]:
+    """Return a process-wide async Redis client, or ``None`` if unavailable.
+
+    ``from_url`` builds the client without opening a socket (connection is
+    lazy), so a transient broker outage surfaces later on the first command
+    rather than here. Mirrors the sync client's ``socket_connect_timeout``
+    and ``health_check_interval`` so a half-open TCP can't wedge the tail
+    loop past its keepalive cadence.
+    """
+    global _async_redis, _creation_failed
+    if _async_redis is None and not _creation_failed:
+        try:
+            _async_redis = aioredis.Redis.from_url(
+                settings.CACHE_REDIS_URL,
+                socket_connect_timeout=2,
+                health_check_interval=10,
+            )
+        except ValueError as e:
+            logger.error("Invalid Redis URL for async client: %s", e)
+            _creation_failed = True
+            _async_redis = None
+    return _async_redis

application/streaming/broadcast_channel.py

@@ -0,0 +1,144 @@
+"""Redis pub/sub Topic abstraction for SSE fan-out.
+
+A Topic is a named channel for one-shot live event delivery. Canonical uses:
+
+- ``user:{user_id}`` for per-user notifications
+- ``channel:{message_id}`` for per-chat-message streams
+
+Subscription is race-free via ``on_subscribe``: the callback fires only
+after Redis acknowledges ``SUBSCRIBE``, so a publisher dispatched inside
+the callback cannot lose its first event to a not-yet-registered
+subscriber.
+
+The subscribe iterator yields ``None`` on poll timeout so the caller can
+emit SSE keepalive comments without spawning a separate timer thread.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Callable, Iterator, Optional
+
+import redis as redis_lib
+
+from application.cache import get_pubsub_redis_instance, get_redis_instance
+
+logger = logging.getLogger(__name__)
+
+
+class Topic:
+    """A pub/sub channel identified by a string name."""
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+    def publish(self, payload: str | bytes) -> int:
+        """Fan out a payload to currently subscribed clients.
+
+        Returns the number Redis reports as receiving the message (limited
+        to subscribers connected to *this* Redis instance), or 0 if Redis
+        is unavailable. Never raises.
+        """
+        redis = get_redis_instance()
+        if redis is None:
+            logger.debug("Redis unavailable; dropping publish to %s", self.name)
+            return 0
+        try:
+            return int(redis.publish(self.name, payload))
+        except Exception:
+            logger.exception("Topic.publish failed for %s", self.name)
+            return 0
+
+    def subscribe(
+        self,
+        on_subscribe: Optional[Callable[[], None]] = None,
+        poll_timeout: float = 1.0,
+    ) -> Iterator[Optional[bytes]]:
+        """Subscribe to the topic; yield raw payloads or ``None`` on tick.
+
+        Yields ``None`` every ``poll_timeout`` seconds while idle so the
+        caller can emit keepalive frames or check cancellation. Yields
+        ``bytes`` for each delivered message.
+
+        ``on_subscribe`` runs synchronously after Redis acknowledges the
+        SUBSCRIBE — use it to seed any state (e.g. read backlog) that
+        must be ordered after the subscriber is live but before the
+        first pub/sub message is processed.
+
+        If Redis is unavailable, returns immediately without yielding.
+        Cleanly unsubscribes on ``GeneratorExit`` (client disconnect).
+
+        Uses the pub/sub-dedicated client (bounded ``socket_timeout``): a
+        subscriber whose connection went half-open must fail within seconds
+        and release its WSGI thread, not block in ``get_message`` until the
+        worker restarts.
+        """
+        redis = get_pubsub_redis_instance()
+        if redis is None:
+            logger.debug("Redis unavailable; subscribe to %s yielded nothing", self.name)
+            return
+        pubsub = None
+        on_subscribe_fired = False
+        try:
+            pubsub = redis.pubsub()
+            try:
+                pubsub.subscribe(self.name)
+            except Exception:
+                # Subscribe failure (transient Redis hiccup, conn reset, etc.)
+                # is treated like "Redis unavailable": yield nothing, let the
+                # caller fall back to its own resilience strategy. The finally
+                # block will still tear down the pubsub object cleanly.
+                logger.exception("pubsub.subscribe failed for %s", self.name)
+                return
+            while True:
+                try:
+                    msg = pubsub.get_message(timeout=poll_timeout)
+                except redis_lib.exceptions.TimeoutError:
+                    # A bounded read timed out mid-message/health-check: the
+                    # connection is half-open (NAT/IPVS dropped the flow).
+                    # End the subscription so the SSE client reconnects on a
+                    # fresh connection; expected during network churn, so no
+                    # stack trace.
+                    logger.info(
+                        "pubsub read timed out for %s; closing subscriber",
+                        self.name,
+                    )
+                    return
+                except Exception:
+                    logger.exception("pubsub.get_message failed for %s", self.name)
+                    return
+                if msg is None:
+                    yield None
+                    continue
+                msg_type = msg.get("type")
+                if msg_type == "subscribe":
+                    if not on_subscribe_fired and on_subscribe is not None:
+                        try:
+                            on_subscribe()
+                        except Exception:
+                            logger.exception(
+                                "on_subscribe callback failed for %s", self.name
+                            )
+                    on_subscribe_fired = True
+                    continue
+                if msg_type != "message":
+                    continue
+                data = msg.get("data")
+                if data is None:
+                    continue
+                yield data if isinstance(data, bytes) else str(data).encode("utf-8")
+        finally:
+            if pubsub is not None:
+                if on_subscribe_fired:
+                    try:
+                        pubsub.unsubscribe(self.name)
+                    except Exception:
+                        logger.debug(
+                            "pubsub unsubscribe error for %s",
+                            self.name,
+                            exc_info=True,
+                        )
+                try:
+                    pubsub.close()
+                except Exception:
+                    logger.debug("pubsub close error for %s", self.name, exc_info=True)

application/streaming/event_replay.py

@@ -0,0 +1,260 @@
+"""Shared snapshot/replay primitives for chat-stream reconnect.
+
+The reconnect reader itself is the native-async generator in
+``async_event_replay.build_message_event_stream_async``; this module holds
+the pieces both it and the producer's journal depend on: the SSE wire
+format (``format_sse_event``), the ``message_events`` snapshot read
+(``read_snapshot_lines``), the producer-liveness watchdog probe
+(``_check_producer_liveness``), and the pub/sub envelope encode/decode.
+Keeping them here lets the async reader and the sync journal agree on the
+exact wire shape and dedup/terminal rules. See
+``docs/runbooks/sse-notifications.md``.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from typing import Optional
+
+from sqlalchemy import text as sql_text
+
+from application.storage.db.repositories.message_events import (
+    MessageEventsRepository,
+)
+from application.storage.db.session import db_readonly
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_KEEPALIVE_SECONDS = 15.0
+DEFAULT_POLL_TIMEOUT_SECONDS = 1.0
+# When the live tail has no events and no terminal in snapshot, fall
+# back to checking ``conversation_messages`` directly. If the row has
+# already gone terminal (worker journaled ``end``/``error`` to the DB
+# but the matching pub/sub publish was lost, or the row was finalized
+# without a journal write at all) we surface a terminal event so the
+# client doesn't hang on keepalives. If the row is still non-terminal
+# but the producer heartbeat is older than ``PRODUCER_IDLE_SECONDS``
+# the producer is presumed dead (worker crash / recycle between chunks
+# and finalize) and we emit a terminal ``error`` so the UI can recover.
+DEFAULT_WATCHDOG_INTERVAL_SECONDS = 5.0
+# 1.5× the route's 60s heartbeat interval — long enough that a normal
+# heartbeat skew doesn't false-positive, short enough that a stuck
+# stream surfaces before the 5-minute reconciler sweep escalates.
+DEFAULT_PRODUCER_IDLE_SECONDS = 90.0
+
+# WHATWG SSE accepts CRLF, CR, LF — split on any of them so a stray CR
+# can't smuggle a record boundary into the wire format.
+_SSE_LINE_SPLIT_PATTERN = re.compile(r"\r\n|\r|\n")
+
+# Event types that mark the end of a chat answer. After delivering one
+# we close the reconnect stream — keeping the connection open past a
+# terminal event would leak both the client's reconnect promise and
+# the server's WSGI thread waiting on keepalives that the user no
+# longer cares about. The agent loop emits ``end`` for normal /
+# tool-paused completion and ``error`` for the catch-all failure path
+# (which doesn't get a trailing ``end``).
+_TERMINAL_EVENT_TYPES = frozenset({"end", "error"})
+
+
+def _payload_is_terminal(
+    payload: object, event_type: Optional[str] = None
+) -> bool:
+    """True if ``payload['type']`` or ``event_type`` is a terminal sentinel."""
+    if isinstance(payload, dict) and payload.get("type") in _TERMINAL_EVENT_TYPES:
+        return True
+    return event_type in _TERMINAL_EVENT_TYPES
+
+
+def format_sse_event(payload: dict, sequence_no: int) -> str:
+    """Encode a journal event as one ``id:``/``data:`` SSE record.
+
+    The body is the payload's JSON serialisation. ``complete_stream``
+    payloads are flat JSON dicts with no embedded newlines, so a
+    single ``data:`` line is sufficient — but we still split on any
+    line terminator in case a future caller passes a multi-line string.
+    """
+    body = json.dumps(payload)
+    lines = [f"id: {sequence_no}"]
+    for line in _SSE_LINE_SPLIT_PATTERN.split(body):
+        lines.append(f"data: {line}")
+    return "\n".join(lines) + "\n\n"
+
+
+def _check_producer_liveness(
+    message_id: str, user_id: Optional[str], idle_seconds: float
+) -> Optional[dict]:
+    """Inspect ``conversation_messages`` and return a terminal SSE
+    payload when the producer is no longer alive, else ``None``.
+
+    When ``user_id`` is given the lookup is scoped to ``AND user_id = :u``
+    (defence in depth: this long-lived re-read re-asserts the ownership the
+    route gated on, so a stream cannot keep tailing a row it no longer
+    owns). A non-matching row reads as missing → a terminal ``error``.
+
+    Three terminal cases collapse into a single DB round-trip:
+
+    - ``status='complete'`` — the live finalize ran but its journal
+      terminal write didn't reach us (or never happened). Synthesise
+      ``end`` so the client closes cleanly on the row's user-visible
+      state.
+    - ``status='failed'`` — same, but for the failure path. Carry the
+      stashed ``error`` from ``message_metadata`` so the UI shows the
+      real reason.
+    - non-terminal status and ``last_heartbeat_at`` (or ``timestamp``)
+      older than ``idle_seconds`` — the producing worker is gone.
+      Synthesise ``error`` so the client doesn't hang on keepalives
+      until the proxy idle-timeout kicks in.
+    """
+    owner_clause = " AND user_id = :u" if user_id is not None else ""
+    params = {"id": message_id, "idle_secs": float(idle_seconds)}
+    if user_id is not None:
+        params["u"] = user_id
+    try:
+        with db_readonly() as conn:
+            row = conn.execute(
+                sql_text(
+                    # ``owner_clause`` is a fixed literal (no user input in the
+                    # SQL string); ``user_id`` is bound via ``:u``.
+                    f"""
+                    SELECT
+                        status,
+                        message_metadata->>'error' AS err,
+                        GREATEST(
+                            timestamp,
+                            COALESCE(
+                                (message_metadata->>'last_heartbeat_at')
+                                    ::timestamptz,
+                                timestamp
+                            )
+                        ) < now() - make_interval(secs => :idle_secs)
+                            AS is_stale
+                    FROM conversation_messages
+                    WHERE id = CAST(:id AS uuid){owner_clause}
+                    """
+                ),
+                params,
+            ).first()
+    except Exception:
+        logger.exception(
+            "Watchdog liveness check failed for message_id=%s", message_id
+        )
+        return None
+
+    if row is None:
+        # Row deleted out from under us — treat as terminal so the
+        # client doesn't keep tailing a message that no longer exists.
+        return {
+            "type": "error",
+            "error": "Message no longer exists; please refresh.",
+            "code": "message_missing",
+            "message_id": message_id,
+        }
+
+    status, err, is_stale = row[0], row[1], bool(row[2])
+    if status == "complete":
+        return {"type": "end"}
+    if status == "failed":
+        return {
+            "type": "error",
+            "error": err or "Stream failed; please try again.",
+            "code": "producer_failed",
+            "message_id": message_id,
+        }
+    if is_stale:
+        return {
+            "type": "error",
+            "error": (
+                "Stream producer is no longer responding; please try again."
+            ),
+            "code": "producer_stale",
+            "message_id": message_id,
+        }
+    return None
+
+
+def read_snapshot_lines(
+    message_id: str, last_event_id: Optional[int], user_id: Optional[str] = None
+) -> tuple[list[str], Optional[int], bool]:
+    """Read journal rows after ``last_event_id`` as SSE-formatted lines.
+
+    Returns ``(lines, max_sequence_no, terminal)``: ``max_sequence_no`` is
+    seeded with ``last_event_id`` and advanced past every row read,
+    ``terminal`` is True if any row carried a terminal ``end``/``error``.
+    Raises on DB error so the caller can drive its replay-failed path.
+
+    Used by ``async_event_replay.build_message_event_stream_async`` (the
+    reconnect reader); it shares ``format_sse_event`` / ``_payload_is_terminal``
+    with the producer's journal writer so reader and writer never drift on
+    wire shape or terminal semantics.
+    """
+    lines: list[str] = []
+    max_seq = last_event_id
+    terminal = False
+    with db_readonly() as conn:
+        rows = MessageEventsRepository(conn).read_after(
+            message_id, last_sequence_no=last_event_id, user_id=user_id
+        )
+    for row in rows:
+        seq = int(row["sequence_no"])
+        payload = row.get("payload")
+        if not isinstance(payload, dict):
+            # ``record_event`` rejects non-dict payloads at the write gate,
+            # so this is a legacy/direct-SQL row — drop it rather than ship
+            # a malformed envelope that would poison a reconnect.
+            logger.warning(
+                "Skipping non-dict payload from message_events: "
+                "message_id=%s seq=%s type=%s",
+                message_id,
+                seq,
+                row.get("event_type"),
+            )
+            continue
+        lines.append(format_sse_event(payload, seq))
+        if max_seq is None or seq > max_seq:
+            max_seq = seq
+        if _payload_is_terminal(payload, row.get("event_type")):
+            terminal = True
+    return lines, max_seq, terminal
+
+
+def _decode_pubsub_message(raw) -> Optional[dict]:
+    """Parse a ``Topic.publish`` payload to ``{sequence_no, payload, ...}``.
+
+    Returns ``None`` for malformed messages (drop silently — the
+    journal is still authoritative on reconnect).
+    """
+    try:
+        if isinstance(raw, (bytes, bytearray)):
+            text_value = raw.decode("utf-8")
+        else:
+            text_value = str(raw)
+        envelope = json.loads(text_value)
+    except Exception:
+        return None
+    if not isinstance(envelope, dict):
+        return None
+    return envelope
+
+
+def encode_pubsub_message(
+    message_id: str,
+    sequence_no: int,
+    event_type: str,
+    payload: dict,
+) -> str:
+    """Build the JSON envelope used for ``channel:{message_id}`` publishes.
+
+    Kept here (not in ``message_journal.py``) so the encode/decode pair
+    stays in one file — replay's ``_decode_pubsub_message`` and the
+    journal's publish must agree on the shape exactly.
+    """
+    return json.dumps(
+        {
+            "message_id": str(message_id),
+            "sequence_no": int(sequence_no),
+            "event_type": event_type,
+            "payload": payload,
+        }
+    )

application/streaming/keys.py

@@ -0,0 +1,19 @@
+"""Per-chat-message stream key derivations.
+
+Single source of truth for the Redis pub/sub topic name and any
+auxiliary keys that the chat-stream snapshot+tail reconnect path
+shares between the writer (``complete_stream`` + journal) and the
+reader (``/api/messages/<id>/events`` reconnect endpoint).
+"""
+
+from __future__ import annotations
+
+
+def message_topic_name(message_id: str) -> str:
+    """Redis pub/sub channel for live fan-out of one chat message.
+
+    Subscribers tail this topic for every event that ``complete_stream``
+    yielded after the SUBSCRIBE-ack arrived; older events are recovered
+    from the ``message_events`` snapshot half of the pattern.
+    """
+    return f"channel:{message_id}"

application/streaming/message_journal.py

@@ -0,0 +1,451 @@
+"""Per-yield journal write for the chat-stream snapshot+tail pattern.
+
+``record_event`` inserts into ``message_events`` and publishes to
+``channel:{message_id}``. Both are best-effort; the INSERT commits
+before the publish so a fast reconnect sees the row. See
+``docs/runbooks/sse-notifications.md``.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any, Optional
+
+from sqlalchemy.exc import IntegrityError
+
+from application.storage.db.repositories.message_events import (
+    MessageEventsRepository,
+)
+from application.storage.db.session import db_readonly, db_session
+from application.streaming.broadcast_channel import Topic
+from application.streaming.event_replay import encode_pubsub_message
+from application.streaming.keys import message_topic_name
+
+logger = logging.getLogger(__name__)
+
+
+# Tunables for ``BatchedJournalWriter``. A streaming answer emits ~100s
+# of ``answer`` chunks per response; without batching, that's one PG
+# transaction per yield in the WSGI thread. With these defaults, ~10x
+# fewer commits at the cost of a ≤100ms reconnect-visibility lag for
+# any event still sitting in the buffer.
+DEFAULT_BATCH_SIZE = 16
+DEFAULT_BATCH_INTERVAL_MS = 100
+
+
+def _strip_null_bytes(value: Any) -> Any:
+    """Recursively strip ``\\x00`` from string keys/values in ``value``.
+
+    Postgres JSONB rejects the NUL escape; an LLM emitting a stray NUL
+    in a chunk would otherwise raise ``DataError`` at INSERT and the row
+    would be lost from the journal (live stream proceeds, reconnect
+    snapshot misses the chunk). Mirrors the strip already done in
+    ``parser/embedding_pipeline.py`` and
+    ``api/user/attachments/routes.py``.
+    """
+    if isinstance(value, str):
+        return value.replace("\x00", "") if "\x00" in value else value
+    if isinstance(value, dict):
+        return {
+            (k.replace("\x00", "") if isinstance(k, str) and "\x00" in k else k):
+            _strip_null_bytes(v)
+            for k, v in value.items()
+        }
+    if isinstance(value, list):
+        return [_strip_null_bytes(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_strip_null_bytes(item) for item in value)
+    return value
+
+
+# Postgres SQLSTATE for a foreign-key violation. ``message_events`` has
+# an FK to ``conversation_messages(id)``, so a 23503 means that parent
+# row was never committed — the event can't be journaled and a
+# ``sequence_no`` retry is pointless churn. A 23505 (PK collision), by
+# contrast, is recoverable by rewriting at a fresh seq.
+_FK_VIOLATION_SQLSTATE = "23503"
+
+
+def _is_foreign_key_violation(exc: IntegrityError) -> bool:
+    """Return ``True`` when ``exc`` wraps a Postgres FK violation (23503).
+
+    ``IntegrityError.orig`` is the underlying driver error, which carries
+    the SQLSTATE on ``.sqlstate`` (psycopg3) or ``.pgcode`` (psycopg2).
+    When neither is present (an error we can't classify) this returns
+    ``False`` so the caller keeps its existing seq-retry behavior.
+    """
+    orig = getattr(exc, "orig", None)
+    sqlstate = getattr(orig, "sqlstate", None) or getattr(orig, "pgcode", None)
+    return sqlstate == _FK_VIOLATION_SQLSTATE
+
+
+def record_event(
+    message_id: str,
+    sequence_no: int,
+    event_type: str,
+    payload: Optional[dict[str, Any]] = None,
+) -> bool:
+    """Journal one SSE event and publish it live. Best-effort.
+
+    ``payload`` must be a ``dict`` or ``None`` (non-dicts are dropped so
+    live and replay envelopes stay byte-identical). Returns ``True`` when
+    the journal INSERT committed. Never raises.
+    """
+    if not message_id or not event_type:
+        logger.warning(
+            "record_event called without message_id/event_type "
+            "(message_id=%r, event_type=%r)",
+            message_id,
+            event_type,
+        )
+        return False
+
+    if payload is None:
+        materialised_payload: dict[str, Any] = {}
+    elif isinstance(payload, dict):
+        materialised_payload = _strip_null_bytes(payload)
+    else:
+        logger.warning(
+            "record_event called with non-dict payload "
+            "(message_id=%s seq=%s type=%s payload_type=%s) — dropping",
+            message_id,
+            sequence_no,
+            event_type,
+            type(payload).__name__,
+        )
+        return False
+
+    journal_committed = False
+    # The seq we actually managed to write. Diverges from
+    # ``sequence_no`` only on the IntegrityError-retry path below.
+    materialised_seq = sequence_no
+    try:
+        # Short-lived per-event transaction. Critical for visibility:
+        # the reconnect endpoint reads the journal from a separate
+        # connection and only sees committed rows.
+        with db_session() as conn:
+            MessageEventsRepository(conn).record(
+                message_id, sequence_no, event_type, materialised_payload
+            )
+        journal_committed = True
+    except IntegrityError as exc:
+        if _is_foreign_key_violation(exc):
+            # No ``conversation_messages`` parent row — the event
+            # references a message that was never committed. A
+            # ``sequence_no`` retry can't fix that, so log the real
+            # cause and drop instead of churning the readonly probe +
+            # retry (which would just FK-fail again).
+            logger.warning(
+                "record_event: no conversation_messages row for "
+                "message_id=%s (foreign-key violation); dropping "
+                "(seq=%s type=%s). The parent message row must be "
+                "committed before its events are journaled.",
+                message_id,
+                sequence_no,
+                event_type,
+            )
+        else:
+            # Composite-PK collision on (message_id, sequence_no). Most
+            # likely cause is a stale ``latest_sequence_no`` seed on a
+            # continuation retry — the route read MAX(seq) from a
+            # separate connection before another writer committed past
+            # it. Look up the live latest and retry once with latest+1
+            # so the event is not silently lost. Bounded to a single
+            # retry — if two writers keep racing in lockstep the
+            # route-level retry will converge them across attempts.
+            try:
+                with db_readonly() as conn:
+                    latest = MessageEventsRepository(conn).latest_sequence_no(
+                        message_id
+                    )
+                materialised_seq = (latest if latest is not None else -1) + 1
+                with db_session() as conn:
+                    MessageEventsRepository(conn).record(
+                        message_id,
+                        materialised_seq,
+                        event_type,
+                        materialised_payload,
+                    )
+                journal_committed = True
+                logger.info(
+                    "record_event: collision at seq=%s recovered → wrote at "
+                    "seq=%s message_id=%s type=%s",
+                    sequence_no,
+                    materialised_seq,
+                    message_id,
+                    event_type,
+                )
+            except IntegrityError:
+                # Second collision under the same retry — give up and
+                # log. The route's nonlocal counter will continue at
+                # ``sequence_no+1`` on the next emit; the next call may
+                # land cleanly past the contended window.
+                logger.warning(
+                    "record_event: IntegrityError persists after seq+1 "
+                    "retry; dropping. message_id=%s original_seq=%s "
+                    "retry_seq=%s type=%s",
+                    message_id,
+                    sequence_no,
+                    materialised_seq,
+                    event_type,
+                )
+            except Exception:
+                logger.exception(
+                    "record_event: retry path failed unexpectedly "
+                    "(message_id=%s seq=%s type=%s)",
+                    message_id,
+                    sequence_no,
+                    event_type,
+                )
+    except Exception:
+        logger.exception(
+            "message_events INSERT failed: message_id=%s seq=%s type=%s",
+            message_id,
+            sequence_no,
+            event_type,
+        )
+
+    try:
+        # Publish using ``materialised_seq`` so the live pubsub frame
+        # matches the journal row that other clients will snapshot on
+        # reconnect. The original POST stream's SSE ``id:`` still
+        # carries the caller's ``sequence_no`` — a reconnect from that
+        # client will receive the same event at ``materialised_seq``
+        # on the snapshot, which is a benign duplicate (the slice's
+        # ``max_replayed_seq`` advances past it). No-collision case:
+        # ``materialised_seq == sequence_no`` and this is identical to
+        # the prior behaviour.
+        wire = encode_pubsub_message(
+            message_id, materialised_seq, event_type, materialised_payload
+        )
+        Topic(message_topic_name(message_id)).publish(wire)
+    except Exception:
+        logger.exception(
+            "channel:%s publish failed: seq=%s type=%s",
+            message_id,
+            materialised_seq,
+            event_type,
+        )
+
+    return journal_committed
+
+
+class BatchedJournalWriter:
+    """Per-stream journal writer that batches PG INSERTs.
+
+    One writer per ``message_id``; ``record()`` buffers events and flushes
+    on size/time/``close()`` triggers. Pubsub publishes fire only after the
+    INSERT commits. On ``IntegrityError`` falls back to per-row writes.
+    """
+
+    def __init__(
+        self,
+        message_id: str,
+        *,
+        batch_size: int = DEFAULT_BATCH_SIZE,
+        batch_interval_ms: int = DEFAULT_BATCH_INTERVAL_MS,
+    ) -> None:
+        self._message_id = message_id
+        self._batch_size = batch_size
+        self._batch_interval_ms = batch_interval_ms
+        self._buffer: list[tuple[int, str, dict[str, Any]]] = []
+        self._last_flush_mono_ms = time.monotonic() * 1000.0
+        self._closed = False
+
+    def record(
+        self,
+        sequence_no: int,
+        event_type: str,
+        payload: Optional[dict[str, Any]] = None,
+    ) -> bool:
+        """Buffer one event; maybe flush. Publish happens after journal commit."""
+        if self._closed:
+            logger.warning(
+                "BatchedJournalWriter.record after close: "
+                "message_id=%s seq=%s type=%s",
+                self._message_id,
+                sequence_no,
+                event_type,
+            )
+            return False
+        if not event_type:
+            logger.warning(
+                "BatchedJournalWriter.record without event_type: "
+                "message_id=%s seq=%s",
+                self._message_id,
+                sequence_no,
+            )
+            return False
+        if payload is None:
+            materialised: dict[str, Any] = {}
+        elif isinstance(payload, dict):
+            materialised = _strip_null_bytes(payload)
+        else:
+            # Same contract as ``record_event`` — non-dict payloads
+            # are rejected so the live and replay paths can't diverge
+            # on envelope reconstruction.
+            logger.warning(
+                "BatchedJournalWriter.record with non-dict payload: "
+                "message_id=%s seq=%s type=%s payload_type=%s — dropping",
+                self._message_id,
+                sequence_no,
+                event_type,
+                type(payload).__name__,
+            )
+            return False
+
+        self._buffer.append((sequence_no, event_type, materialised))
+
+        if self._should_flush():
+            self.flush()
+        return True
+
+    def _should_flush(self) -> bool:
+        if len(self._buffer) >= self._batch_size:
+            return True
+        elapsed_ms = (time.monotonic() * 1000.0) - self._last_flush_mono_ms
+        return elapsed_ms >= self._batch_interval_ms and len(self._buffer) > 0
+
+    def flush(self) -> None:
+        """Commit buffered events to PG. Best-effort.
+
+        Tries one bulk INSERT first; on ``IntegrityError`` (composite
+        PK collision — typically a stale continuation seed) falls back
+        to per-row ``record_event`` so one bad seq doesn't drop the
+        rest of the batch. Always clears the buffer to bound memory,
+        even on failure — a journaled event missing from a snapshot
+        is degraded UX, but a runaway buffer is corruption.
+        """
+        if not self._buffer:
+            self._last_flush_mono_ms = time.monotonic() * 1000.0
+            return
+
+        # Snapshot and clear before the I/O so a concurrent record()
+        # call would land in a fresh buffer rather than racing the
+        # flush. ``complete_stream`` is single-threaded per stream, so
+        # this is belt-and-suspenders for any future change.
+        pending = self._buffer
+        self._buffer = []
+        self._last_flush_mono_ms = time.monotonic() * 1000.0
+
+        try:
+            with db_session() as conn:
+                MessageEventsRepository(conn).bulk_record(
+                    self._message_id, pending
+                )
+        except IntegrityError as exc:
+            if _is_foreign_key_violation(exc):
+                # The whole batch references a conversation_messages
+                # parent that doesn't exist — every per-row retry would
+                # FK-fail too, so drop the batch once instead of N
+                # pointless re-attempts. The buffer is already cleared.
+                logger.warning(
+                    "BatchedJournalWriter: no conversation_messages row "
+                    "for message_id=%s (foreign-key violation); dropping "
+                    "%d event(s). The parent message row must be committed "
+                    "before its events are journaled.",
+                    self._message_id,
+                    len(pending),
+                )
+                return
+            logger.info(
+                "BatchedJournalWriter: bulk INSERT collided for "
+                "message_id=%s n=%d; falling back to per-row writes",
+                self._message_id,
+                len(pending),
+            )
+            self._flush_per_row(pending)
+            return
+        except Exception:
+            logger.exception(
+                "BatchedJournalWriter: bulk INSERT failed for "
+                "message_id=%s n=%d; events dropped from journal",
+                self._message_id,
+                len(pending),
+            )
+            return
+
+        # Bulk INSERT committed — publish each frame in order. Best-effort:
+        # one failed publish must not poison the rest of the batch.
+        for seq, event_type, payload in pending:
+            self._publish(seq, event_type, payload)
+
+    def _flush_per_row(
+        self, pending: list[tuple[int, str, dict[str, Any]]]
+    ) -> None:
+        """Per-row fallback after a bulk collision. Publishes after each commit."""
+        for seq, event_type, payload in pending:
+            committed_seq: Optional[int] = None
+            try:
+                with db_session() as conn:
+                    MessageEventsRepository(conn).record(
+                        self._message_id, seq, event_type, payload
+                    )
+                committed_seq = seq
+            except IntegrityError:
+                try:
+                    with db_readonly() as conn:
+                        latest = MessageEventsRepository(
+                            conn
+                        ).latest_sequence_no(self._message_id)
+                    retry_seq = (latest if latest is not None else -1) + 1
+                    with db_session() as conn:
+                        MessageEventsRepository(conn).record(
+                            self._message_id, retry_seq, event_type, payload
+                        )
+                    committed_seq = retry_seq
+                except IntegrityError:
+                    logger.warning(
+                        "BatchedJournalWriter: IntegrityError persists "
+                        "after seq+1 retry; dropping. message_id=%s "
+                        "original_seq=%s type=%s",
+                        self._message_id,
+                        seq,
+                        event_type,
+                    )
+                except Exception:
+                    logger.exception(
+                        "BatchedJournalWriter: per-row retry failed "
+                        "(message_id=%s seq=%s type=%s)",
+                        self._message_id,
+                        seq,
+                        event_type,
+                    )
+            except Exception:
+                logger.exception(
+                    "BatchedJournalWriter: per-row INSERT failed "
+                    "(message_id=%s seq=%s type=%s)",
+                    self._message_id,
+                    seq,
+                    event_type,
+                )
+
+            if committed_seq is not None:
+                self._publish(committed_seq, event_type, payload)
+
+    def _publish(
+        self, sequence_no: int, event_type: str, payload: dict[str, Any]
+    ) -> None:
+        """Publish one frame to the per-message pubsub channel. Best-effort."""
+        try:
+            wire = encode_pubsub_message(
+                self._message_id, sequence_no, event_type, payload
+            )
+            Topic(message_topic_name(self._message_id)).publish(wire)
+        except Exception:
+            logger.exception(
+                "channel:%s publish failed: seq=%s type=%s",
+                self._message_id,
+                sequence_no,
+                event_type,
+            )
+
+    def close(self) -> None:
+        """Final flush. Idempotent — safe to call from multiple
+        finally clauses.
+        """
+        if self._closed:
+            return
+        self.flush()
+        self._closed = True