PlantAI-rag-new/bioflow_service/app/services/bioflow/events.py

"""app/services/bioflow/events.py —— 进程内事件总线（pub/sub 供 SSE）+ 持久事件日志。

- 发布：runner/路由在任务状态变化时 publish(type, data)（事件类型清单 = 实施计划 §六
  12 类 + 内部 job_recovered）。
- 订阅：async queue（供 SSE 端点）；内存环形缓冲（断线重连回放）。
- 持久：events/{job_id}.jsonl 追加落盘（服务重启恢复回放；FR-BIO-031）。
- 快照：snapshot(job_id) 构造断线重连首帧（TC-BIO-SSE-022/024）。

服务层纪律：纯 dataclass/dict/stdlib，无 pydantic。
"""
from __future__ import annotations

import asyncio
import json
import logging
import os
import threading
from collections import deque
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

EVENT_TYPES = (
    "job_created", "plan_ready", "stage_started", "stage_log", "stage_progress",
    "stage_done", "stage_failed", "stage_skipped", "stage_await_confirm",
    "stage_await_files",
    "job_done", "job_paused", "job_archived", "job_recovered",
    # 用户自定义扩展分析事件帧（契约 §4；并行链路，不触碰现有 job_* 语义）
    "extension_created", "extension_build", "extension_trial", "extension_repair",
    "extension_validate", "extension_published", "extension_candidates",
    "extension_confirm", "extension_run_started", "extension_run_log",
    "extension_run_done", "extension_run_failed",
)

# 断线重连时需要回放的关键状态事件（含历史事件全量）
REPLAYABLE_TYPES = set(EVENT_TYPES)

MEMORY_BUFFER_LIMIT = 500
SNAPSHOT_TYPE = "snapshot"


@dataclass(eq=False)
class Subscriber:
    """订阅者（eq=False：身份哈希，可入 set）。"""

    queue: asyncio.Queue = field(default_factory=asyncio.Queue)
    last_seq: int = -1
    job_filter: str | None = None


def _stage_key(s: Any) -> str:
    """job_created stages 元素 → 阶段名键（str 或 dict 统一；BUG-R1-04）。"""
    if isinstance(s, dict):
        return str(s.get("stage") or "")
    return str(s)


def _stage_entry(s: Any) -> dict[str, Any]:
    """job_created stages 元素 → 干净 {stage, status}（dict 保真透传）。

    runner 发布的 job_created.stages 元素为 {stage, status} dict；此前直接
    str(s) 会把 dict repr 串当阶段名（前端进度卡乱串，BUG-R1-04）。
    """
    if isinstance(s, dict):
        return {"stage": str(s.get("stage") or ""), "status": str(s.get("status") or "waiting")}
    return {"stage": str(s), "status": "waiting"}


class EventBus:
    """进程内事件总线：publish → 订阅队列广播 + 内存环形缓冲 + jsonl 持久日志。"""

    def __init__(self, events_dir: str | Path) -> None:
        self.events_dir = Path(events_dir)
        self.events_dir.mkdir(parents=True, exist_ok=True)
        self._subscribers: set[Subscriber] = set()
        self._sub_lock = asyncio.Lock()
        self._buffer: deque[dict[str, Any]] = deque(maxlen=MEMORY_BUFFER_LIMIT)
        # 跨重启 seq 单调递增（BUG-FIX 2026-09-17）：旧进程事件已持久化在
        # events/*.jsonl，重启后 _seq 归零会让新事件 seq 小于浏览器 last_seq
        # 水位，前端回放被过滤、收不到 stage_await_files 等关键帧。
        # 启动时扫描历史日志取最大 seq 作为起始值。
        self._seq = self._scan_max_seq()
        self._file_lock = threading.Lock()

    def _scan_max_seq(self) -> int:
        """扫描 events_dir 下全部 jsonl，取历史最大 seq（无文件/坏行 → 0）。"""
        max_seq = 0
        try:
            for path in self.events_dir.glob("*.jsonl"):
                try:
                    for line in path.read_text(encoding="utf-8").splitlines():
                        if not line.strip():
                            continue
                        try:
                            seq = int(json.loads(line).get("seq") or 0)
                        except (ValueError, TypeError):
                            continue
                        if seq > max_seq:
                            max_seq = seq
                except OSError:
                    continue
        except OSError:  # pragma: no cover - 目录扫描失败(权限等)不影响启动
            return 0
        return max_seq

    # ---- 内部 ----

    def _events_file(self, job_id: str) -> Path:
        return self.events_dir / f"{job_id}.jsonl"

    def _append_file(self, event: dict[str, Any]) -> None:
        try:
            job_id = str((event.get("data") or {}).get("job_id") or "")
            with self._file_lock:
                with self._events_file(job_id).open("a", encoding="utf-8") as fh:
                    fh.write(json.dumps(event, ensure_ascii=False) + "\n")
        except OSError:
            logger.warning("bioflow events: 持久化失败 %s", event.get("type"))

    # ---- 发布 ----

    def publish(self, event_type: str, data: dict[str, Any]) -> dict[str, Any]:
        """发布事件（同步）：返回完整事件（含 seq），并唤醒所有订阅者。"""
        if event_type not in EVENT_TYPES:
            raise ValueError(f"未知事件类型: {event_type}")
        self._seq += 1
        event: dict[str, Any] = {
            "type": event_type,
            "seq": self._seq,
            "data": dict(data),
        }
        self._buffer.append(event)
        self._append_file(event)
        # 广播给订阅者（跨线程安全：queue.put_nowait 本身线程安全）
        for sub in list(self._subscribers):
            if sub.job_filter is not None and sub.job_filter != str(data.get("job_id") or ""):
                continue
            try:
                sub.queue.put_nowait(event)
            except asyncio.QueueFull:  # pragma: no cover - 防御
                logger.warning("bioflow events: 订阅队列满，丢弃事件 %s", event_type)
        return event

    def publish_async(self, event_type: str, data: dict[str, Any]) -> None:
        """异步发布（供路由协程调用，publish 本身为同步安全实现）。"""
        self.publish(event_type, data)

    # ---- 订阅（SSE 端点）----

    async def subscribe(self, job_filter: str | None = None, last_seq: int = -1) -> Subscriber:
        sub = Subscriber(last_seq=last_seq, job_filter=job_filter)
        async with self._sub_lock:
            self._subscribers.add(sub)
        return sub

    async def unsubscribe(self, sub: Subscriber) -> None:
        async with self._sub_lock:
            self._subscribers.discard(sub)

    # ---- 断线重连回放 ----

    def snapshot(self, job_id: str) -> dict[str, Any] | None:
        """构造任务快照（首帧 data.type=snapshot；TC-BIO-SSE-022）。

        从持久日志读该任务全部事件，拼出 job + stages 当前状态。
        事件不存在（未知任务）返回 None。
        """
        path = self._events_file(job_id)
        if not path.exists():
            return None
        events: list[dict[str, Any]] = []
        try:
            for line in path.read_text(encoding="utf-8").splitlines():
                if not line.strip():
                    continue
                try:
                    ev = json.loads(line)
                    if ev.get("type") in REPLAYABLE_TYPES:
                        events.append(ev)
                except ValueError:
                    continue
        except OSError:
            return None
        if not events:
            return None
        # 由事件序列推导 job/stages 状态
        job_status = "created"
        stage_states: dict[str, dict[str, Any]] = {}
        for ev in events:
            t, d = ev["type"], ev.get("data") or {}
            if t == "job_created":
                job_status = "created"
                for s in d.get("stages") or []:
                    key = _stage_key(s)
                    if key:
                        stage_states.setdefault(key, _stage_entry(s))
            elif t == "stage_started":
                st = stage_states.setdefault(str(d.get("stage")), {"stage": str(d.get("stage"))})
                st["status"] = "running"
            elif t == "stage_done":
                st = stage_states.setdefault(str(d.get("stage")), {"stage": str(d.get("stage"))})
                st["status"] = "done"
            elif t == "stage_failed":
                st = stage_states.setdefault(str(d.get("stage")), {"stage": str(d.get("stage"))})
                st["status"] = "failed"
            elif t == "stage_skipped":
                st = stage_states.setdefault(str(d.get("stage")), {"stage": str(d.get("stage"))})
                st["status"] = "skipped"
            elif t == "stage_progress":
                st = stage_states.setdefault(str(d.get("stage")), {"stage": str(d.get("stage"))})
                pct = d.get("percent")
                if isinstance(pct, (int, float)):
                    st["progress"] = max(float(st.get("progress") or 0.0), float(pct))
                msg = d.get("message")
                if msg:
                    st["progressMessage"] = str(msg)
            elif t == "stage_await_confirm":
                st = stage_states.setdefault(str(d.get("stage")), {"stage": str(d.get("stage"))})
                st["status"] = "waiting_confirm"
            elif t == "stage_await_files":
                st = stage_states.setdefault(str(d.get("stage")), {"stage": str(d.get("stage"))})
                st["status"] = "await_files"
            elif t == "job_done":
                job_status = "done"
            elif t == "job_paused":
                job_status = "paused"
            elif t == "job_archived":
                job_status = "archived"
        return {
            "type": SNAPSHOT_TYPE,
            "job_id": job_id,
            "job": {"status": job_status},
            "stages": [stage_states[k] for k in stage_states],
            "last_seq": events[-1]["seq"] if events else 0,
        }

    def replay(self, job_id: str, after_seq: int = -1) -> list[dict[str, Any]]:
        """回放任务历史事件（seq > after_seq；断线重连补齐；TC-BIO-SSE-023/024）。"""
        path = self._events_file(job_id)
        if not path.exists():
            return []
        out: list[dict[str, Any]] = []
        try:
            for line in path.read_text(encoding="utf-8").splitlines():
                if not line.strip():
                    continue
                try:
                    ev = json.loads(line)
                    if ev.get("type") in REPLAYABLE_TYPES and int(ev.get("seq") or 0) > after_seq:
                        out.append(ev)
                except (ValueError, TypeError):
                    continue
        except OSError:
            return []
        return out

    def memory_replay(self, after_seq: int = -1) -> list[dict[str, Any]]:
        """内存环形缓冲回放（进程重启后内存为空，此时走持久回放）。"""
        return [ev for ev in self._buffer if int(ev.get("seq") or 0) > after_seq]