vent_agent_2.0/api/sse_events.py

"""SSE 事件契约（唯一事实源）。

**这个文件定义了 2.0 与 1.5 前端的兼容边界。** 字段名不是猜的，是从
1.5 的 ``static/index.html`` 的 ``handleSSEEvent`` 逐行反查出来的消费点：

===============  ==========================================================
事件            前端读取的字段
===============  ==========================================================
``agent_start``  ``agent`` / ``cn_agent``
``thinking``     ``message`` / ``cn_agent``
``thinking_token`` ``content``
``executing``    ``message`` / ``cn_agent`` / ``tools[]`` / ``cn_tools[]``
``tool_call``    ``tool`` / ``cn_tool`` / ``agent_id``
``tool_result``  ``tool`` / ``cn_tool`` / ``agent_id`` / ``download_url`` / ``filename``
``token``        ``content``
``updated_todo_list`` ``todos``
``progress``     ``message`` / ``download_url`` / ``filename``
``citations``    ``items[]`` / ``model_only``
``interrupt``    ``tool_approval{tool,args,reason,side_effect}`` 或
                 ``ask_user{questions[{question,type,choices,required}]}`` / ``message``
``agent_done``   ``agent`` / ``cn_agent`` / ``duration_ms`` / ``progress``
``agent_error``  ``agent`` / ``cn_agent`` / ``duration_ms`` / ``progress`` / ``message``
``done``         ``duration_ms`` / ``download_url`` / ``filename`` / ``session_id``
``error``        ``message`` / ``code`` / ``retry_after``
===============  ==========================================================

帧格式：``data: {json}\\n\\n``；每帧恒带 ``session_id``，绑定 trace 时带
``trace_id``。

**改动纪律**：新增事件或字段是安全的（前端忽略未知字段）；重命名或删除
字段会静默破坏前端渲染——所以所有帧都必须经由 :class:`EventWriter`
构造，不允许在别处手拼 ``data:`` 字符串。
"""

from __future__ import annotations

import json
from typing import Any, Iterable

# 事件类型全集（与 1.5 对齐）
AGENT_START = "agent_start"
THINKING = "thinking"
THINKING_TOKEN = "thinking_token"
EXECUTING = "executing"
TOOL_CALL = "tool_call"
TOOL_RESULT = "tool_result"
UPDATED_TODO_LIST = "updated_todo_list"
TOKEN = "token"
PROGRESS = "progress"
INTERRUPT = "interrupt"
CITATIONS = "citations"
AGENT_DONE = "agent_done"
AGENT_ERROR = "agent_error"
ERROR = "error"
DONE = "done"

ALL_EVENT_TYPES: tuple[str, ...] = (
    AGENT_START,
    THINKING,
    THINKING_TOKEN,
    EXECUTING,
    TOOL_CALL,
    TOOL_RESULT,
    UPDATED_TODO_LIST,
    TOKEN,
    PROGRESS,
    INTERRUPT,
    CITATIONS,
    AGENT_DONE,
    AGENT_ERROR,
    ERROR,
    DONE,
)

# 错误码 → 前端提示文案。前端会用 event.code 查这张表兜底。
ERROR_MESSAGES: dict[str, str] = {
    "quota_exceeded": "今日额度已用尽",
    "too_many_tasks": "并发任务已达上限，请稍后重试",
    "task_timeout": "任务执行超时",
    "recursion_limit": "任务步骤过多已中止，请拆分为更小的任务",
    "model_error": "模型调用失败",
    "internal_error": "服务内部错误",
}


class EventWriter:
    """把事件序列化为 SSE 帧。

    一次请求一个实例，绑定 ``session_id`` 与 ``trace_id``——这两个字段
    每帧都要带，放在实例里避免每个方法都传一遍。
    """

    def __init__(self, session_id: str, trace_id: str = "") -> None:
        self.session_id = session_id
        self.trace_id = trace_id

    # -------- 底层 --------

    def frame(self, event_type: str, **fields: Any) -> str:
        """构造一帧。``type`` 与 ``session_id`` 由本方法统一注入。"""
        payload: dict[str, Any] = {"type": event_type, "session_id": self.session_id}
        if self.trace_id:
            payload["trace_id"] = self.trace_id
        payload.update(fields)
        return f"data: {json.dumps(payload, ensure_ascii=False, default=str)}\n\n"

    def comment(self, text: str) -> str:
        """SSE 注释帧，用作心跳。客户端会忽略，但能防止中间代理掐空闲连接。"""
        return f": {text}\n\n"

    # -------- 生命周期 --------

    def agent_start(self, agent: str, cn_agent: str = "") -> str:
        return self.frame(AGENT_START, agent=agent, cn_agent=cn_agent or agent)

    def agent_done(
        self, agent: str, cn_agent: str = "", duration_ms: int | None = None, **extra: Any
    ) -> str:
        fields: dict[str, Any] = {"agent": agent, "cn_agent": cn_agent or agent}
        if duration_ms is not None:
            fields["duration_ms"] = duration_ms
        fields.update(extra)
        return self.frame(AGENT_DONE, **fields)

    def agent_error(
        self, agent: str, message: str, cn_agent: str = "", duration_ms: int | None = None
    ) -> str:
        fields: dict[str, Any] = {
            "agent": agent,
            "cn_agent": cn_agent or agent,
            "message": message,
        }
        if duration_ms is not None:
            fields["duration_ms"] = duration_ms
        return self.frame(AGENT_ERROR, **fields)

    def done(self, duration_ms: int | None = None, **extra: Any) -> str:
        fields: dict[str, Any] = {}
        if duration_ms is not None:
            fields["duration_ms"] = duration_ms
        fields.update(extra)
        return self.frame(DONE, **fields)

    def error(
        self, message: str, code: str = "internal_error", retry_after: int | None = None
    ) -> str:
        fields: dict[str, Any] = {"message": message, "code": code}
        if retry_after is not None:
            fields["retry_after"] = retry_after
        return self.frame(ERROR, **fields)

    # -------- 过程 --------

    def thinking(self, message: str, cn_agent: str = "") -> str:
        fields: dict[str, Any] = {"message": message}
        if cn_agent:
            fields["cn_agent"] = cn_agent
        return self.frame(THINKING, **fields)

    def thinking_token(self, content: str) -> str:
        return self.frame(THINKING_TOKEN, content=content)

    def executing(
        self,
        message: str,
        cn_agent: str = "",
        tools: Iterable[str] | None = None,
        cn_tools: Iterable[str] | None = None,
    ) -> str:
        fields: dict[str, Any] = {"message": message}
        if cn_agent:
            fields["cn_agent"] = cn_agent
        tool_list = list(tools or [])
        if tool_list:
            fields["tools"] = tool_list
            fields["cn_tools"] = list(cn_tools or []) or tool_list
        return self.frame(EXECUTING, **fields)

    def tool_call(
        self,
        tool: str,
        args: Any = None,
        cn_tool: str = "",
        agent_id: str = "",
        tool_call_id: str = "",
    ) -> str:
        """工具调用卡。

        ``tool_call_id`` 是 **2.0 的附加字段**（1.5 没有）：同一轮里同一个工具
        可能被调多次，只按 ``tool`` 名索引卡片会导致结果挂到错误的卡上。
        附加字段不破坏兼容——旧前端忽略未知字段即可。
        """
        fields: dict[str, Any] = {"tool": tool, "cn_tool": cn_tool or tool, "args": args}
        if tool_call_id:
            fields["tool_call_id"] = tool_call_id
        if agent_id:
            fields["agent_id"] = agent_id
        return self.frame(TOOL_CALL, **fields)

    def tool_result(
        self,
        tool: str,
        content: Any = None,
        cn_tool: str = "",
        agent_id: str = "",
        error: bool = False,
        tool_call_id: str = "",
    ) -> str:
        """工具结果卡。``tool_call_id`` 与对应的 tool_call 帧一致。"""
        fields: dict[str, Any] = {
            "tool": tool,
            "cn_tool": cn_tool or tool,
            "content": content,
            "error": error,
        }
        if tool_call_id:
            fields["tool_call_id"] = tool_call_id
        if agent_id:
            fields["agent_id"] = agent_id
        return self.frame(TOOL_RESULT, **fields)

    def token(self, content: str) -> str:
        return self.frame(TOKEN, content=content)

    def todo_list(self, todos: list[dict[str, Any]]) -> str:
        return self.frame(UPDATED_TODO_LIST, todos=todos)

    def progress(self, message: str, **extra: Any) -> str:
        fields: dict[str, Any] = {"message": message}
        fields.update(extra)
        return self.frame(PROGRESS, **fields)

    def citations(self, items: list[dict[str, Any]], model_only: bool = False) -> str:
        return self.frame(CITATIONS, items=items, model_only=model_only)

    # -------- 中断 --------

    def interrupt_tool_approval(
        self,
        tool: str,
        args: Any = None,
        reason: str = "",
        side_effect: str = "",
        **extra: Any,
    ) -> str:
        """工具审批卡（对应前端 ``showToolApprovalCard``）。"""
        payload: dict[str, Any] = {
            "tool": tool,
            "args": args,
            "reason": reason,
            "side_effect": side_effect,
        }
        payload.update(extra)
        return self.frame(INTERRUPT, tool_approval=payload)

    def interrupt_ask_user(
        self, questions: list[dict[str, Any]], message: str = ""
    ) -> str:
        """提问卡（对应前端 ``showAskUserCard``）。

        每个问题形如 ``{"question": str, "type": "text"|"multiple_choice",
        "choices": [...], "required": bool}``。
        """
        return self.frame(
            INTERRUPT,
            ask_user={"questions": questions},
            message=message or "智能体暂停，等待输入",
        )

    def interrupt_raw(self, message: str = "", **extra: Any) -> str:
        """无法归类的中断：只给一条提示，前端显示 toast。"""
        return self.frame(INTERRUPT, message=message, **extra)