"""Conversation-scoped in-process event hub for MAF roundtable SSE streaming. This hub mediates between the writers (store/service/workflow executor that persist conversation events to the DB) and the SSE subscribers (the HTTP stream endpoint). It is deliberately **notification-only**: the database is the authoritative event store, and this hub merely signals "new events may be available" so subscribers can re-query. This avoids duplicating the event log in memory and guarantees that a reconnecting client can always replay from the DB cursor (``after_seq``). Writer contract (§4.2 "先持久化、后通知"): 1. Persist events to ``maf_roundtable_conversation_events`` (via ``MafRoundtableStore.append_conversation_item``) inside a transaction. 2. **After commit**, call ``await hub.publish(conversation_id)``. Subscriber contract (P0-1 race-free): 1. ``sub = hub.open_subscription(conversation_id)`` — register **before** any DB query so publishes during replay are captured. 2. Replay from DB: ``list_events(after_seq=last_seq)``. 3. **Stabilise**: if ``sub.current_generation > sub.seen_generation``, re-query DB and repeat until stable (events arrived during replay). 4. Send snapshot. 5. ``await sub.wait_for_change(timeout=heartbeat)`` → on True, query DB for new events; on False, send heartbeat. This eliminates the ``signal.clear()`` race: generation is monotonic and never cleared, so a publish that happens between the DB read and the ``wait()`` is always detected on the next stabilisation check. Multi-worker limitation: Single-process, in-memory. Multi-worker requires Redis Pub/Sub. """ from __future__ import annotations import asyncio import logging from typing import TYPE_CHECKING if TYPE_CHECKING: pass logger = logging.getLogger(__name__) _DEFAULT_HEARTBEAT_INTERVAL = 15.0 _MAX_SUBSCRIBERS_PER_CONVERSATION = 64 class _ConvState: """Per-conversation notification state.""" __slots__ = ("generation", "condition", "subscriber_count") def __init__(self) -> None: self.generation: int = 0 self.condition: asyncio.Condition = asyncio.Condition() self.subscriber_count: int = 0 class Subscription: """Explicit subscription handle — race-free notification reader. Created via :meth:`ConversationEventHub.open_subscription`. The caller owns its lifecycle: call :meth:`close` when the SSE connection ends. Key invariant: ``seen_generation`` is monotonic — it is never reset. A publish that races ahead of ``wait_for_change`` simply increments ``current_generation``, and the next ``wait_for_change`` returns immediately (no lost wakeup). """ __slots__ = ("_hub", "_conv_id", "_seen_gen", "_closed") def __init__(self, hub: ConversationEventHub, conversation_id: str, gen_at_open: int) -> None: self._hub = hub self._conv_id = conversation_id self._seen_gen = gen_at_open self._closed = False @property def conversation_id(self) -> str: return self._conv_id @property def current_generation(self) -> int: """Latest generation known to the hub for this conversation.""" return self._hub._get_generation(self._conv_id) @property def seen_generation(self) -> int: """Generation as of the last successful stabilisation / wait.""" return self._seen_gen @property def has_pending(self) -> bool: """True when the hub generation is ahead of what we've seen.""" return self.current_generation > self._seen_gen def mark_seen(self) -> None: """Advance ``seen_generation`` to the hub's current generation.""" self._seen_gen = self.current_generation def mark_seen_if_higher(self, gen: int) -> None: if gen > self._seen_gen: self._seen_gen = gen async def wait_for_change(self, *, timeout: float) -> bool: """Block until ``current_generation > seen_generation`` or timeout. Returns ``True`` if the generation advanced (new events *might* be available — caller must query DB to confirm). Returns ``False`` on timeout (caller should send a heartbeat). """ state = self._hub._convs.get(self._conv_id) if state is None: return False async with state.condition: if state.generation > self._seen_gen: self._seen_gen = state.generation return True try: await asyncio.wait_for(state.condition.wait(), timeout=timeout) except TimeoutError: return False if state.generation > self._seen_gen: self._seen_gen = state.generation return True return False def close(self) -> None: """Release the subscription. Safe to call multiple times.""" if not self._closed: self._closed = True self._hub._release_subscription(self._conv_id) def __enter__(self) -> Subscription: return self def __exit__(self, *exc: object) -> None: self.close() class ConversationEventHub: """In-process notification hub keyed by ``conversation_id``. Not thread-safe; intended for a single asyncio event loop (one worker). """ def __init__(self, *, heartbeat_interval: float = _DEFAULT_HEARTBEAT_INTERVAL) -> None: self._heartbeat_interval = heartbeat_interval self._convs: dict[str, _ConvState] = {} self._closed = False @property def heartbeat_interval(self) -> float: return self._heartbeat_interval @property def total_subscribers(self) -> int: return sum(s.subscriber_count for s in self._convs.values()) @property def is_closed(self) -> bool: return self._closed def conversation_subscriber_count(self, conversation_id: str) -> int: state = self._convs.get(conversation_id) return state.subscriber_count if state else 0 def _get_generation(self, conversation_id: str) -> int: state = self._convs.get(conversation_id) return state.generation if state else 0 def open_subscription(self, conversation_id: str) -> Subscription: """Register a new subscriber **before** any DB replay. Returns a :class:`Subscription` whose ``seen_generation`` is the hub's current generation at registration time. Any ``publish()`` after this call will be detectable via ``sub.has_pending``. """ if self._closed: raise RuntimeError("ConversationEventHub is shut down") state = self._convs.get(conversation_id) if state is None: state = _ConvState() self._convs[conversation_id] = state if state.subscriber_count >= _MAX_SUBSCRIBERS_PER_CONVERSATION: raise RuntimeError( f"Too many SSE subscribers for conversation {conversation_id} " f"(max {_MAX_SUBSCRIBERS_PER_CONVERSATION})" ) state.subscriber_count += 1 logger.debug( "SSE subscribe conv=%s gen=%d (subs=%d)", conversation_id, state.generation, state.subscriber_count, ) return Subscription(self, conversation_id, state.generation) def _release_subscription(self, conversation_id: str) -> None: state = self._convs.get(conversation_id) if state is None: return state.subscriber_count -= 1 logger.debug( "SSE unsubscribe conv=%s (subs=%d)", conversation_id, max(0, state.subscriber_count), ) if state.subscriber_count <= 0: self._convs.pop(conversation_id, None) async def publish(self, conversation_id: str) -> None: """Signal that new events have been persisted for a conversation. Must be called **after** the events are committed to the DB, never before. Atomically increments the per-conversation generation counter and wakes all waiting subscribers. Safe to call when no subscribers are active (the next subscriber's stabilisation check will see the incremented generation). """ if self._closed: return state = self._convs.get(conversation_id) if state is None: # No active subscribers — nothing to wake. The generation is not # tracked for conversations with zero subscribers; a future # open_subscription starts fresh at 0, which is correct because # there was nobody to notify. return async with state.condition: state.generation += 1 state.condition.notify_all() async def wait_for_change( self, conversation_id: str, seen_gen: int, *, timeout: float ) -> int: """Convenience: wait until generation > ``seen_gen``. Returns the new generation (> ``seen_gen``) on change, or ``seen_gen`` unchanged on timeout. Prefer using a :class:`Subscription` for the full lifecycle. """ state = self._convs.get(conversation_id) if state is None: return seen_gen async with state.condition: if state.generation > seen_gen: return state.generation try: await asyncio.wait_for(state.condition.wait(), timeout=timeout) except TimeoutError: return seen_gen return state.generation async def aclose(self) -> None: """Async shutdown that immediately wakes all subscribers. Acquires each Condition lock, bumps generation, and calls ``notify_all()`` so every waiter returns within milliseconds — not after the next heartbeat timeout. """ self._closed = True for state in list(self._convs.values()): async with state.condition: state.generation += 1 state.condition.notify_all() def shutdown(self) -> None: """Deprecated sync stub — bumps generation but cannot wake waiters immediately. Callers should use ``await hub.aclose()`` instead.""" self._closed = True for state in self._convs.values(): state.generation += 1 async def hub_aclose(hub: ConversationEventHub) -> None: """Async shutdown that immediately wakes all subscribers.""" await hub.aclose() __all__ = ["ConversationEventHub", "Subscription", "hub_aclose"]