281 lines
10 KiB
Python
281 lines
10 KiB
Python
"""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"]
|