mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:21:58 +00:00
ThearchyHelios
135 lines
5.1 KiB
Python
135 lines
5.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
from collections import defaultdict
|
|
from collections.abc import Callable
|
|
from typing import AsyncIterator
|
|
|
|
from app.core.events import RunEventModel
|
|
|
|
# 单个订阅者的积压上限。慢消费者（比如卡住的浏览器标签）不应该拖垮整个 run，
|
|
# 所以队列满了就丢最旧的事件并打标记，而不是阻塞发布方。
|
|
_QUEUE_MAXSIZE = 2048
|
|
|
|
|
|
class _Subscriber:
|
|
def __init__(self) -> None:
|
|
self.queue: asyncio.Queue[RunEventModel | None] = asyncio.Queue(maxsize=_QUEUE_MAXSIZE)
|
|
self.dropped = 0
|
|
|
|
|
|
class EventBus:
|
|
"""按 run_id 分频道的内存 pub/sub。
|
|
|
|
只负责实时分发；历史事件由 DB 承担，订阅者接入时先读库补齐再接实时流，
|
|
这样刷新页面或中途接入都能看到完整时间线。
|
|
|
|
多进程（app/cluster/relay.py）：本进程发布的每一条也交给 tap（转发器据此给别的进程发通知、给在看的进程转流式增量）；
|
|
别的进程执行的运行，转发器从库里按 seq 读出来再在本进程发布一遍（relayed=True，不再交给 tap，免得来回转）。
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self._subs: dict[str, set[_Subscriber]] = defaultdict(set)
|
|
self._seq: dict[str, int] = defaultdict(int)
|
|
self._lock = asyncio.Lock()
|
|
self.tap: Callable[[RunEventModel], None] | None = None
|
|
|
|
def next_seq(self, run_id: str) -> int:
|
|
self._seq[run_id] += 1
|
|
return self._seq[run_id]
|
|
|
|
def set_seq(self, run_id: str, value: int) -> None:
|
|
"""恢复一个已存在的 run 时，把序号接上，保证前端排序不乱。"""
|
|
self._seq[run_id] = max(self._seq[run_id], value)
|
|
|
|
async def publish(self, event: RunEventModel, *, relayed: bool = False) -> None:
|
|
if not relayed and self.tap is not None:
|
|
try:
|
|
self.tap(event)
|
|
except Exception: # noqa: BLE001 - 转发出错不能拦住本进程的推送
|
|
pass
|
|
for sub in list(self._subs.get(event.run_id, ())):
|
|
try:
|
|
sub.queue.put_nowait(event)
|
|
except asyncio.QueueFull:
|
|
with contextlib.suppress(asyncio.QueueEmpty):
|
|
sub.queue.get_nowait()
|
|
sub.dropped += 1
|
|
with contextlib.suppress(asyncio.QueueFull):
|
|
sub.queue.put_nowait(event)
|
|
|
|
async def close(self, run_id: str) -> None:
|
|
"""给所有订阅者发结束哨兵，让 WebSocket 干净地收尾。"""
|
|
for sub in list(self._subs.get(run_id, ())):
|
|
with contextlib.suppress(asyncio.QueueFull):
|
|
sub.queue.put_nowait(None)
|
|
|
|
def attach(self, run_id: str) -> "Subscription":
|
|
"""立刻占一个位，之后再慢慢消费。
|
|
|
|
必须能在读历史**之前**占位。`subscribe()` 做不到这件事：它是异步
|
|
生成器，订阅者要等到第一次迭代才真正注册——而调用方通常先读完历史
|
|
才开始迭代。这中间发出的事件既不在历史里、也不在实时流里，就这么
|
|
没了。一次跑图从发起到订阅之间正好有几毫秒，丢掉的往往是开头那
|
|
两三条 node.started/node.finished，表现为"某个节点永远在转圈、
|
|
它的子步骤跑到了顶层"。
|
|
"""
|
|
sub = _Subscriber()
|
|
self._subs[run_id].add(sub)
|
|
return Subscription(self, run_id, sub)
|
|
|
|
def _detach(self, run_id: str, sub: _Subscriber) -> None:
|
|
self._subs.get(run_id, set()).discard(sub)
|
|
if run_id in self._subs and not self._subs[run_id]:
|
|
self._subs.pop(run_id, None)
|
|
|
|
async def subscribe(self, run_id: str) -> AsyncIterator[RunEventModel]:
|
|
"""便捷入口：占位即消费。只适合不需要补历史的场景。"""
|
|
with self.attach(run_id) as sub:
|
|
async for event in sub.stream():
|
|
yield event
|
|
|
|
def subscriber_count(self, run_id: str) -> int:
|
|
return len(self._subs.get(run_id, ()))
|
|
|
|
|
|
class Subscription:
|
|
"""一个已经在收事件的订阅位。"""
|
|
|
|
def __init__(self, bus: EventBus, run_id: str, sub: _Subscriber) -> None:
|
|
self._bus = bus
|
|
self._run_id = run_id
|
|
self._sub = sub
|
|
|
|
@property
|
|
def dropped(self) -> int:
|
|
"""因为消费太慢被丢掉的条数。不是 0 就说明这条时间线有洞。"""
|
|
return self._sub.dropped
|
|
|
|
async def stream(self, after: int = 0) -> AsyncIterator[RunEventModel]:
|
|
"""吐出实时事件。
|
|
|
|
after 用来去重：占位到读完历史之间发出的事件，两边都有一份。
|
|
按 seq 过滤掉已经补过的，前端就不会看到同一步出现两次。
|
|
"""
|
|
while True:
|
|
item = await self._sub.queue.get()
|
|
if item is None:
|
|
break
|
|
if after and (item.seq or 0) <= after:
|
|
continue
|
|
yield item
|
|
|
|
def close(self) -> None:
|
|
self._bus._detach(self._run_id, self._sub)
|
|
|
|
def __enter__(self) -> "Subscription":
|
|
return self
|
|
|
|
def __exit__(self, *exc: object) -> None:
|
|
self.close()
|
|
|
|
|
|
bus = EventBus()
|