393 lines
17 KiB
Python
393 lines
17 KiB
Python
"""RunBroker:pub/sub,把 agent run 产生的 event fan-out 给所有 SSE 订阅者。
|
||
|
||
DESIGN §7 简化(0004 一并)—— 单活 run 形态下 run_id 是冗余的(同 task 同时
|
||
最多 1 个活 run);broker 内部全用 task_id 索引,客户端只需要 task_id 即可
|
||
订阅 / cancel,不再需要先拿 run_id。
|
||
|
||
两个实现,env `ZCBOT_REDIS_URL` 决定(import 时定,进程生命周期不变):
|
||
- **LocalRunBroker**(默认,env 不设):进程内 asyncio queue,零依赖 —— 本地 dev /
|
||
单实例部署用。多进程不共享:蓝绿并存期跨实例订阅/cancel 送达不到(RUN.md B 档边缘)。
|
||
- **RedisRunBroker**(env 设了):event 走 Redis PUB/SUB(`zcbot:ev:<task_id>`),
|
||
done/cancel 走带 TTL 的 key —— 蓝绿双实例跨进程可达,也是稳态双实例分流的前提
|
||
(DESIGN §7.0)。启动 ping 不通 fail-fast(同 sandbox init 纪律);运行中 publish
|
||
失败 log + 丢帧(直播降级,run 不受影响)。
|
||
|
||
公共语义(两实现一致):
|
||
- emit() 从工作线程调(agent.run 在 to_thread 跑);SSE generator await queue.get()。
|
||
- 同一 task 多个订阅者(刷新页面 / 多 tab)— 每个订阅 1 个独立 queue。
|
||
- run 结束 → close(task_id) 给所有订阅者派 done;晚到的订阅者立即收 done 断流
|
||
(local 用 _done set;redis 用 `zcbot:done:<tid>` 60s TTL key —— 主保护是 SSE
|
||
端点先查 DB run_status,key 只兜"查完 DB 到订阅完成之间 run 恰好收尾"的窗口)。
|
||
- 同 task 起新 run → start(task_id) 清 done 标记。
|
||
- cancel:request_cancel 置位,AgentLoop 在 chunk 间 poll is_cancelled() 即退;
|
||
redis 版是 `zcbot:cancel:<tid>` key(GET ~0.1ms/次,localhost 可忽略),跨实例可达。
|
||
- 不持久化 event — messages 已落 PG,刷新走 GET /messages 看历史;订阅建立前
|
||
emit 的帧两个实现都会丢(设计选择)。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import threading
|
||
import time
|
||
from collections import defaultdict
|
||
from typing import Any, Optional
|
||
from uuid import UUID
|
||
|
||
|
||
class LocalRunBroker:
|
||
"""进程内实现(原 RunBroker)。
|
||
|
||
线程模型:
|
||
- bind_loop(loop) 在 FastAPI startup 调一次,记录 asyncio loop 引用。
|
||
- emit() 调用方可能在任意线程;put_nowait 是 thread-unsafe(asyncio.Queue 设计
|
||
前提是单 loop),所以走 call_soon_threadsafe 跨回 loop 线程再 put。
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self._subs: dict[UUID, set[asyncio.Queue]] = defaultdict(set)
|
||
# 已经发完 done 的 task — 后来订阅者直接收到 done,避免无限等
|
||
self._done: set[UUID] = set()
|
||
# cancel signal per-task。AgentLoop 在 BG 线程里 poll is_cancelled() 决定是否退;
|
||
# request_cancel 可在 BG 还没 register 时调用(setdefault),BG 启动后第一次
|
||
# check 即看到。run 完成在 finally 里 clear_cancel 回收。
|
||
self._cancel_flags: dict[UUID, threading.Event] = {}
|
||
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
||
|
||
def bind_loop(self, loop: asyncio.AbstractEventLoop) -> None:
|
||
"""FastAPI startup 调一次。"""
|
||
self._loop = loop
|
||
|
||
async def shutdown(self) -> None:
|
||
"""lifespan finally 调;local 无外部资源,no-op。"""
|
||
|
||
def start(self, task_id: UUID) -> None:
|
||
"""同 task 起新 run 时调:清 _done 标记,让新订阅者能看到流式。
|
||
cancel flag 在 finally 里 clear_cancel 清,这里不动(避免擦掉刚刚 request_cancel 的请求)。
|
||
"""
|
||
self._done.discard(task_id)
|
||
|
||
def subscribe(self, task_id: UUID) -> asyncio.Queue:
|
||
"""订阅 task 当前 run 的 event 流。已 done 的 task 立刻在 queue 放一条 done。
|
||
|
||
调用方:SSE handler(在 asyncio loop 线程内)。
|
||
"""
|
||
q: asyncio.Queue = asyncio.Queue()
|
||
if task_id in self._done:
|
||
q.put_nowait({"type": "done"})
|
||
else:
|
||
self._subs[task_id].add(q)
|
||
return q
|
||
|
||
def unsubscribe(self, task_id: UUID, q: asyncio.Queue) -> None:
|
||
"""SSE generator finally 清理。"""
|
||
self._subs.get(task_id, set()).discard(q)
|
||
if task_id in self._subs and not self._subs[task_id]:
|
||
del self._subs[task_id]
|
||
|
||
def emit(self, task_id: UUID, event: dict[str, Any]) -> None:
|
||
"""从工作线程调:把 event 推给所有订阅者。
|
||
|
||
如果没人订阅(run 在跑但没浏览器连上),event 丢弃 — 这是设计选择
|
||
(event 不持久化,messages 走 PG)。
|
||
"""
|
||
loop = self._loop
|
||
if loop is None:
|
||
return # 还没 bind,丢弃(测试 / 启动竞态)
|
||
for q in list(self._subs.get(task_id, [])):
|
||
loop.call_soon_threadsafe(q.put_nowait, event)
|
||
|
||
def close(self, task_id: UUID) -> None:
|
||
"""run 结束:派 done 给所有订阅者,标记 task 为已完成。
|
||
|
||
从工作线程调(agent.run 完成 / 抛异常 finally 清理)。
|
||
"""
|
||
self.emit(task_id, {"type": "done"})
|
||
self._done.add(task_id)
|
||
# subs 不在这里立即删 — SSE generator 会先收到 done、yield 它、走到
|
||
# finally unsubscribe;此处 emit 后立即删会让那次 emit 之后的清理无的放矢。
|
||
|
||
def n_subscribers(self, task_id: UUID) -> int:
|
||
"""供测试 / 监控用。"""
|
||
return len(self._subs.get(task_id, set()))
|
||
|
||
def total_subscribers(self) -> int:
|
||
"""本进程 SSE 订阅者总数(所有 task 累加),供 §8.4 并发监控用。"""
|
||
return sum(len(s) for s in self._subs.values())
|
||
|
||
def is_done(self, task_id: UUID) -> bool:
|
||
return task_id in self._done
|
||
|
||
# ─────────────── cancel signaling ───────────────
|
||
|
||
def request_cancel(self, task_id: UUID) -> None:
|
||
"""主线程(HTTP handler)发的 cancel 信号 — BG 线程 poll is_cancelled() 看见即退。
|
||
setdefault:即便 BG 还没注册 flag 也能 set,BG 启动后第一次 check 立刻看见。
|
||
"""
|
||
self._cancel_flags.setdefault(task_id, threading.Event()).set()
|
||
|
||
def is_cancelled(self, task_id: UUID) -> bool:
|
||
ev = self._cancel_flags.get(task_id)
|
||
return bool(ev and ev.is_set())
|
||
|
||
def clear_cancel(self, task_id: UUID) -> None:
|
||
"""run 真正退出(BG finally)清掉 flag,避免 dict 无限增长。"""
|
||
self._cancel_flags.pop(task_id, None)
|
||
|
||
|
||
class RedisRunBroker:
|
||
"""Redis pub/sub 实现(蓝绿双实例跨进程,DESIGN §7.0)。
|
||
|
||
- event:`PUBLISH zcbot:ev:<task_id>`(JSON);单条 async pubsub reader 任务
|
||
把消息路由到本进程各订阅 queue(订阅表 _subs 仍在本地 —— fan-out 的"最后
|
||
一公里"还是进程内,Redis 只负责跨进程那一跳)。
|
||
- done:close 先 `SETEX zcbot:done:<tid> 60`再 publish done;subscribe 先挂
|
||
channel 再查 key,晚到订阅者不漏。
|
||
- cancel:`SETEX zcbot:cancel:<tid> 3600` / GET / DEL,跨实例可达。
|
||
- 容错:bind_loop ping 不通直接 raise(fail-fast);运行中 redis 操作失败
|
||
log(10s 节流)+ 降级(emit 丢帧、is_cancelled 返 False),run 不受影响;
|
||
reader 断线 1s 退避重连并重挂全部活跃 channel。
|
||
"""
|
||
|
||
_EV_PREFIX = "zcbot:ev:"
|
||
_DONE_PREFIX = "zcbot:done:"
|
||
_CANCEL_PREFIX = "zcbot:cancel:"
|
||
_CTL_CHANNEL = "zcbot:ev:__ctl__" # 常挂空 channel,保 pubsub 始终有订阅可 get_message
|
||
_DONE_TTL = 60
|
||
_CANCEL_TTL = 3600
|
||
|
||
def __init__(self, url: str, *, sync_client=None, async_client=None) -> None:
|
||
import redis as _redis
|
||
import redis.asyncio as _aredis
|
||
self._url = url
|
||
# sync client:emit/close/cancel 从工作线程调(redis-py 自带连接池,线程安全);
|
||
# start/request_cancel 也会在 loop 线程上调 —— localhost 单次 ~0.1ms,超时收紧
|
||
# 到 2s 防 redis 挂掉时把 event loop 卡死太久。
|
||
self._redis = sync_client if sync_client is not None else _redis.Redis.from_url(
|
||
url, decode_responses=True, socket_timeout=2.0, socket_connect_timeout=2.0,
|
||
)
|
||
self._aredis = async_client if async_client is not None else _aredis.from_url(
|
||
url, decode_responses=True,
|
||
)
|
||
self._pubsub: Any = None
|
||
self._reader_task: Optional[asyncio.Task] = None
|
||
self._subs: dict[UUID, set[asyncio.Queue]] = defaultdict(set)
|
||
# Redis 写入期间的同实例瞬时旁路：HTTP 线程先置 Event，run 线程下一次
|
||
# 轮询立即命中；写入结束即清，Redis key 负责跨实例传递和持久兜底。
|
||
self._cancel_flags: dict[UUID, threading.Event] = {}
|
||
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
||
self._last_err_log = 0.0
|
||
|
||
# ─────────────── helpers ───────────────
|
||
|
||
def _chan(self, task_id: UUID) -> str:
|
||
return f"{self._EV_PREFIX}{task_id}"
|
||
|
||
def _log_err(self, what: str, e: Exception) -> None:
|
||
"""redis 故障时 emit/is_cancelled 每 chunk 都会撞,10s 节流防刷屏。"""
|
||
now = time.monotonic()
|
||
if now - self._last_err_log >= 10:
|
||
self._last_err_log = now
|
||
print(f"[broker] redis {what} failed: {type(e).__name__}: {e} (降级继续,10s 内同类不再报)")
|
||
|
||
# ─────────────── lifecycle ───────────────
|
||
|
||
def bind_loop(self, loop: asyncio.AbstractEventLoop) -> None:
|
||
"""FastAPI startup 调一次:ping 验活(不通 fail-fast)+ 起 pubsub reader。"""
|
||
self._loop = loop
|
||
try:
|
||
self._redis.ping()
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
f"ZCBOT_REDIS_URL 已设但 Redis 连不上({type(e).__name__}: {e})——"
|
||
f"确认 redis-server 在跑 / URL 对,或删掉该 env 回退进程内 broker"
|
||
)
|
||
self._reader_task = loop.create_task(self._reader(), name="redis-broker-reader")
|
||
host = self._url.rsplit("@", 1)[-1] # 不把密码打进日志
|
||
print(f"[broker] redis pub/sub enabled ({host})")
|
||
|
||
async def shutdown(self) -> None:
|
||
"""lifespan finally 调:停 reader、断两个客户端。"""
|
||
if self._reader_task is not None:
|
||
self._reader_task.cancel()
|
||
try:
|
||
await self._reader_task
|
||
except (asyncio.CancelledError, Exception):
|
||
pass
|
||
if self._pubsub is not None:
|
||
try:
|
||
await self._pubsub.aclose()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
await self._aredis.aclose()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
self._redis.close()
|
||
except Exception:
|
||
pass
|
||
|
||
async def _reader(self) -> None:
|
||
"""常驻:单条 pubsub 连接收所有活跃 channel 的消息,路由到本地订阅 queue。
|
||
断线 1s 退避重连,并把 _subs 里所有活跃 channel 重挂上(订阅状态不丢)。"""
|
||
while True:
|
||
try:
|
||
pubsub = self._aredis.pubsub()
|
||
self._pubsub = pubsub
|
||
channels = [self._CTL_CHANNEL] + [self._chan(t) for t in self._subs]
|
||
await pubsub.subscribe(*channels)
|
||
while True:
|
||
msg = await pubsub.get_message(
|
||
ignore_subscribe_messages=True, timeout=5.0
|
||
)
|
||
if msg is None or msg.get("type") != "message":
|
||
continue
|
||
ch = msg.get("channel") or ""
|
||
if ch == self._CTL_CHANNEL:
|
||
continue
|
||
try:
|
||
tid = UUID(ch[len(self._EV_PREFIX):])
|
||
event = json.loads(msg["data"])
|
||
except (ValueError, TypeError):
|
||
continue # 不认识的 channel / 坏 payload,丢
|
||
for q in list(self._subs.get(tid, ())):
|
||
q.put_nowait(event)
|
||
except asyncio.CancelledError:
|
||
raise
|
||
except Exception as e:
|
||
self._log_err("reader", e)
|
||
try:
|
||
if self._pubsub is not None:
|
||
await self._pubsub.aclose()
|
||
except Exception:
|
||
pass
|
||
await asyncio.sleep(1)
|
||
|
||
# ─────────────── pub/sub ───────────────
|
||
|
||
def start(self, task_id: UUID) -> None:
|
||
try:
|
||
self._redis.delete(f"{self._DONE_PREFIX}{task_id}")
|
||
except Exception as e:
|
||
self._log_err("start/del-done", e)
|
||
|
||
def subscribe(self, task_id: UUID) -> asyncio.Queue:
|
||
"""调用方:SSE handler(loop 线程)。queue 立即返回;channel 挂载与 done-key
|
||
补查异步完成 —— 挂载完成前 publish 的帧会丢(与 local 版"订阅前的帧丢"同类,
|
||
SSE 端点已先查过 DB run_status 兜大头)。"""
|
||
q: asyncio.Queue = asyncio.Queue()
|
||
first = not self._subs[task_id]
|
||
self._subs[task_id].add(q)
|
||
|
||
async def _setup() -> None:
|
||
try:
|
||
if first and self._pubsub is not None:
|
||
await self._pubsub.subscribe(self._chan(task_id))
|
||
# 先挂 channel 再查 done key:close() 是先 SETEX 再 publish,
|
||
# 两边一交叉,晚到订阅者要么收到 publish、要么查到 key,不漏
|
||
if await self._aredis.exists(f"{self._DONE_PREFIX}{task_id}"):
|
||
q.put_nowait({"type": "done"})
|
||
except Exception as e:
|
||
self._log_err("subscribe-setup", e)
|
||
q.put_nowait({"type": "done"}) # 保守断流,客户端自动重连
|
||
|
||
if self._loop is not None:
|
||
self._loop.create_task(_setup())
|
||
return q
|
||
|
||
def unsubscribe(self, task_id: UUID, q: asyncio.Queue) -> None:
|
||
subs = self._subs.get(task_id)
|
||
if subs is None:
|
||
return
|
||
subs.discard(q)
|
||
if not subs:
|
||
del self._subs[task_id]
|
||
ch = self._chan(task_id)
|
||
|
||
async def _unsub() -> None:
|
||
try:
|
||
if self._pubsub is not None:
|
||
await self._pubsub.unsubscribe(ch)
|
||
except Exception:
|
||
pass # reader 重连时按 _subs 重挂,残留订阅无害
|
||
|
||
if self._loop is not None:
|
||
self._loop.create_task(_unsub())
|
||
|
||
def emit(self, task_id: UUID, event: dict[str, Any]) -> None:
|
||
try:
|
||
self._redis.publish(
|
||
self._chan(task_id), json.dumps(event, ensure_ascii=False, default=str)
|
||
)
|
||
except Exception as e:
|
||
self._log_err("publish", e) # 丢帧降级,run 不受影响
|
||
|
||
def close(self, task_id: UUID) -> None:
|
||
try:
|
||
self._redis.set(f"{self._DONE_PREFIX}{task_id}", "1", ex=self._DONE_TTL)
|
||
except Exception as e:
|
||
self._log_err("close/set-done", e)
|
||
self.emit(task_id, {"type": "done"})
|
||
|
||
def n_subscribers(self, task_id: UUID) -> int:
|
||
return len(self._subs.get(task_id, set()))
|
||
|
||
def total_subscribers(self) -> int:
|
||
"""本进程的订阅数(蓝绿双实例各报各的,admin 面板看单实例)。"""
|
||
return sum(len(s) for s in self._subs.values())
|
||
|
||
def is_done(self, task_id: UUID) -> bool:
|
||
try:
|
||
return bool(self._redis.exists(f"{self._DONE_PREFIX}{task_id}"))
|
||
except Exception as e:
|
||
self._log_err("is_done", e)
|
||
return False
|
||
|
||
# ─────────────── cancel signaling ───────────────
|
||
|
||
def request_cancel(self, task_id: UUID) -> None:
|
||
self._cancel_flags.setdefault(task_id, threading.Event()).set()
|
||
try:
|
||
self._redis.set(f"{self._CANCEL_PREFIX}{task_id}", "1", ex=self._CANCEL_TTL)
|
||
except Exception as e:
|
||
self._log_err("request_cancel", e)
|
||
finally:
|
||
self._cancel_flags.pop(task_id, None)
|
||
|
||
def is_cancelled(self, task_id: UUID) -> bool:
|
||
"""AgentLoop 在 stream chunk 间 poll；同实例先读 Event，跨实例再读 Redis。
|
||
|
||
Redis 运行中故障时跨实例快信号不可达，但同实例 Event 与共享 DB fallback
|
||
仍可取消，不影响正常 run。
|
||
"""
|
||
local = self._cancel_flags.get(task_id)
|
||
if local is not None and local.is_set():
|
||
return True
|
||
try:
|
||
return self._redis.get(f"{self._CANCEL_PREFIX}{task_id}") is not None
|
||
except Exception as e:
|
||
self._log_err("is_cancelled", e)
|
||
return False
|
||
|
||
def clear_cancel(self, task_id: UUID) -> None:
|
||
self._cancel_flags.pop(task_id, None)
|
||
try:
|
||
self._redis.delete(f"{self._CANCEL_PREFIX}{task_id}")
|
||
except Exception as e:
|
||
self._log_err("clear_cancel", e)
|
||
|
||
|
||
def _make_broker():
|
||
"""env `ZCBOT_REDIS_URL` 在 → Redis 版(蓝绿跨进程);不在 → 进程内版(dev/单实例)。
|
||
import 时定一次,与 ZCBOT_INSTANCE 同款读法。"""
|
||
url = os.getenv("ZCBOT_REDIS_URL", "").strip()
|
||
if url:
|
||
return RedisRunBroker(url)
|
||
return LocalRunBroker()
|
||
|
||
|
||
# 进程内单例 — FastAPI lifespan 里 bind_loop;agent / sink / SSE handler 共享。
|
||
broker = _make_broker()
|