56 lines
2.1 KiB
Python
56 lines
2.1 KiB
Python
"""进程内 SSE 事件总线。
|
|
|
|
后台任务 ``append_event`` 写入 DB（持久化）后，同时 publish 到本总线；SSE generator
|
|
从内存队列实时读取并转发给前端，避免轮询 DB 带来的延迟与开销。
|
|
|
|
约束：仅进程内有效，进程重启即失。多进程 / 多副本部署需替换为 Redis Pub/Sub 等共享
|
|
通道；本地开发（单进程 uvicorn）适用。事件历史仍以 DB 为持久事实源，总线仅承担实时转发。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections import defaultdict
|
|
|
|
from app.schemas.runtime import EventResponse
|
|
|
|
_TERMINAL_EVENT_TYPES = {"run.completed", "run.failed", "run.cancelled"}
|
|
|
|
|
|
class RunEventBus:
|
|
"""按 run_id 缓存事件并广播给订阅者。"""
|
|
|
|
def __init__(self) -> None:
|
|
self._events: dict[str, list[EventResponse]] = defaultdict(list)
|
|
self._subscribers: dict[str, list[asyncio.Queue[EventResponse]]] = defaultdict(list)
|
|
|
|
def publish(self, run_id: str, event: EventResponse) -> None:
|
|
self._events[run_id].append(event)
|
|
for queue in self._subscribers.get(run_id, []):
|
|
queue.put_nowait(event)
|
|
|
|
def list_after(self, run_id: str, *, after_sequence: int) -> list[EventResponse]:
|
|
"""返回 run 内 sequence 大于 after_sequence 的事件（用于断线重连补发）。"""
|
|
events = self._events.get(run_id, [])
|
|
if after_sequence <= 0:
|
|
return list(events)
|
|
return [event for event in events if event.sequence > after_sequence]
|
|
|
|
def subscribe(self, run_id: str) -> asyncio.Queue[EventResponse]:
|
|
queue: asyncio.Queue[EventResponse] = asyncio.Queue()
|
|
self._subscribers[run_id].append(queue)
|
|
return queue
|
|
|
|
def unsubscribe(self, run_id: str, queue: asyncio.Queue[EventResponse]) -> None:
|
|
subscribers = self._subscribers.get(run_id)
|
|
if subscribers is None:
|
|
return
|
|
self._subscribers[run_id] = [q for q in subscribers if q is not queue]
|
|
|
|
@staticmethod
|
|
def is_terminal(event: EventResponse) -> bool:
|
|
return event.type in _TERMINAL_EVENT_TYPES
|
|
|
|
|
|
bus = RunEventBus()
|