64 lines
2.3 KiB
Python
64 lines
2.3 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
|
||
|
||
# 终态事件：新旧两套运行时共用本总线转发 SSE；Kernel 契约（run.finished/run.interrupted）
|
||
# 与旧运行时（run.completed/run.failed）的终态都在此登记
|
||
_TERMINAL_EVENT_TYPES = {
|
||
"run.completed",
|
||
"run.failed",
|
||
"run.cancelled",
|
||
"run.finished",
|
||
"run.interrupted",
|
||
}
|
||
|
||
|
||
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()
|