vent_agent_2.0/tests/test_sse_contract.py

"""SSE 契约测试。

两层：
1. **字段级**：逐个事件断言 ``type`` 与必需字段名——这是与 1.5 前端兼容的
   硬约束，字段被改名或删除必须在这里失败。
2. **内核级**：用假 agent（手工构造 astream chunk）驱动 :class:`ChatStream`，
   验证「哪个通道产出哪种事件」的分工，不需要模型网络。
"""

from __future__ import annotations

import json

import pytest
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage

from api.sse_events import ALL_EVENT_TYPES, EventWriter

SESSION = "sess-contract-001"
TRACE = "trace-abc"


def writer() -> EventWriter:
    return EventWriter(SESSION, TRACE)


def parse(frame: str) -> dict:
    assert frame.startswith("data: "), f"帧必须以 'data: ' 开头：{frame[:40]!r}"
    assert frame.endswith("\n\n"), "帧必须以空行结束"
    return json.loads(frame[6:])


# ============================================================
# 1. 字段级契约
# ============================================================

class TestEventContract:
    """每个事件类型的最小字段集。字段名改动会在这里失败。"""

    def test_event_types_frozen(self):
        """事件类型全集固定为 15 种——前端 switch 分支依赖它。"""
        assert len(ALL_EVENT_TYPES) == 15
        for name in (
            "agent_start", "thinking", "thinking_token", "executing", "tool_call",
            "tool_result", "updated_todo_list", "token", "progress", "interrupt",
            "citations", "agent_done", "agent_error", "error", "done",
        ):
            assert name in ALL_EVENT_TYPES, f"缺少事件类型 {name}"

    def test_every_frame_carries_session_and_trace(self):
        w = writer()
        frames = [
            w.agent_start("a", "A"),
            w.thinking("m"),
            w.thinking_token("t"),
            w.executing("m"),
            w.tool_call("t"),
            w.tool_result("t"),
            w.token("x"),
            w.todo_list([]),
            w.progress("p"),
            w.interrupt_raw("i"),
            w.citations([]),
            w.agent_done("a"),
            w.agent_error("a", "e"),
            w.error("e"),
            w.done(),
        ]
        for frame in frames:
            payload = parse(frame)
            assert payload["session_id"] == SESSION
            assert payload["trace_id"] == TRACE
            assert payload["type"] in ALL_EVENT_TYPES

    def test_trace_absent_when_empty(self):
        """未绑定 trace 时不发 trace_id 字段（而不是发空串）。"""
        payload = parse(EventWriter(SESSION).token("x"))
        assert "trace_id" not in payload

    def test_agent_start_fields(self):
        p = parse(writer().agent_start("deep-agent", "深度智能体"))
        assert p["type"] == "agent_start"
        # 前端 `agentCardKey` 依次读 agent_id → cn_agent → agent
        assert p["agent"] == "deep-agent"
        assert p["cn_agent"] == "深度智能体"

    def test_agent_start_cn_falls_back_to_agent(self):
        assert parse(writer().agent_start("x"))["cn_agent"] == "x"

    def test_thinking_fields(self):
        p = parse(writer().thinking("正在分析", "深度智能体"))
        assert p["type"] == "thinking"
        assert p["message"] == "正在分析"
        assert p["cn_agent"] == "深度智能体"

    def test_thinking_token_content(self):
        p = parse(writer().thinking_token("推理中"))
        assert p["type"] == "thinking_token" and p["content"] == "推理中"

    def test_executing_tools_and_cn_tools(self):
        p = parse(writer().executing("调用工具：计算", tools=["calculate"], cn_tools=["安全算术计算"]))
        assert p["type"] == "executing"
        assert p["message"] == "调用工具：计算"
        assert p["tools"] == ["calculate"]
        assert p["cn_tools"] == ["安全算术计算"]

    def test_executing_without_tools_omits_keys(self):
        p = parse(writer().executing("执行中"))
        assert "tools" not in p and "cn_tools" not in p

    def test_tool_call_fields(self):
        p = parse(writer().tool_call("calculate", {"expression": "1+1"}, "安全算术计算", "sub-1"))
        assert p["type"] == "tool_call"
        assert p["tool"] == "calculate"
        assert p["cn_tool"] == "安全算术计算"
        assert p["args"] == {"expression": "1+1"}
        assert p["agent_id"] == "sub-1"

    def test_tool_call_id_is_additive_2_0_field(self):
        """tool_call_id 是 2.0 新增字段（1.5 没有）：同一工具多次调用时按它
        索引才不会把结果挂错卡。附加字段不破坏兼容。"""
        call = parse(writer().tool_call("calculate", {}, "计算", tool_call_id="call-abc"))
        result = parse(writer().tool_result("calculate", 2, "计算", tool_call_id="call-abc"))
        assert call["tool_call_id"] == "call-abc"
        assert result["tool_call_id"] == call["tool_call_id"]

    def test_tool_call_id_omitted_when_absent(self):
        """没有 id 时不发该字段，而不是发空串。"""
        assert "tool_call_id" not in parse(writer().tool_call("t", {}))

    def test_tool_result_fields(self):
        p = parse(writer().tool_result("calculate", {"result": 2}, "安全算术计算", error=False))
        assert p["type"] == "tool_result"
        assert p["tool"] == "calculate" and p["error"] is False
        assert p["content"] == {"result": 2}

    def test_token_content(self):
        assert parse(writer().token("你好"))["content"] == "你好"

    def test_todo_list_field_is_todos(self):
        todos = [{"content": "第一步", "status": "in_progress"}]
        p = parse(writer().todo_list(todos))
        assert p["type"] == "updated_todo_list"
        assert p["todos"] == todos

    def test_progress_message(self):
        p = parse(writer().progress("进行中"))
        assert p["type"] == "progress" and p["message"] == "进行中"

    def test_citations_items_and_model_only(self):
        items = [{"title": "T", "url": "u"}]
        p = parse(writer().citations(items))
        assert p["type"] == "citations"
        assert p["items"] == items
        assert p["model_only"] is False
        assert parse(writer().citations([], model_only=True))["model_only"] is True

    def test_interrupt_tool_approval_shape(self):
        """前端 `showToolApprovalCard` 从 event.tool_approval 读四个字段。"""
        p = parse(writer().interrupt_tool_approval(
            "request_plan_approval", {"plan": "..."}, "需确认计划", "需人工确认"
        ))
        assert p["type"] == "interrupt"
        ta = p["tool_approval"]
        assert ta["tool"] == "request_plan_approval"
        assert ta["args"] == {"plan": "..."}
        assert ta["reason"] == "需确认计划"
        assert ta["side_effect"] == "需人工确认"

    def test_interrupt_ask_user_shape(self):
        """前端 `showAskUserCard` 从 event.ask_user.questions 读 q.question 等。"""
        questions = [{"question": "选哪个？", "type": "multiple_choice",
                      "choices": ["A", "B"], "required": True}]
        p = parse(writer().interrupt_ask_user(questions, "请选择"))
        assert p["type"] == "interrupt"
        assert p["ask_user"]["questions"] == questions
        assert p["message"] == "请选择"

    def test_agent_done_fields(self):
        p = parse(writer().agent_done("deep-agent", "深度智能体", duration_ms=1234, progress="已完成"))
        assert p["type"] == "agent_done"
        assert p["agent"] == "deep-agent" and p["cn_agent"] == "深度智能体"
        assert p["duration_ms"] == 1234 and p["progress"] == "已完成"

    def test_agent_error_fields(self):
        p = parse(writer().agent_error("deep-agent", "炸了", "深度智能体", 50))
        assert p["type"] == "agent_error"
        assert p["message"] == "炸了" and p["duration_ms"] == 50

    def test_error_code_and_retry_after(self):
        p = parse(writer().error("太挤了", code="too_many_tasks", retry_after=15))
        assert p["type"] == "error"
        assert p["code"] == "too_many_tasks"
        assert p["retry_after"] == 15

    def test_error_without_retry_after_omits_key(self):
        assert "retry_after" not in parse(writer().error("x", code="internal_error"))

    def test_done_duration(self):
        p = parse(writer().done(duration_ms=999))
        assert p["type"] == "done" and p["duration_ms"] == 999

    def test_comment_frame_is_ignored_form(self):
        """心跳帧是 SSE 注释，客户端会忽略。"""
        assert writer().comment("ping").startswith(": ")

    def test_unknown_types_serializable(self):
        """args 里出现不可 JSON 化的对象时不能让整帧崩掉。"""
        from datetime import datetime

        p = parse(writer().tool_call("t", {"when": datetime(2026, 1, 1)}))
        assert isinstance(p["args"]["when"], str)


# ============================================================
# 2. 内核级：假 agent 驱动 ChatStream
# ============================================================

class FakeAgent:
    """按脚本产出 astream chunk 的假 agent。

    只实现 :class:`ChatStream` 用到的接口：``astream`` 与 ``aget_state``。
    """

    def __init__(self, chunks: list, state=None) -> None:
        self._chunks = chunks
        self._state = state
        self.inputs: list = []

    async def astream(self, graph_input, **kwargs):
        self.inputs.append(graph_input)
        for chunk in self._chunks:
            yield chunk

    async def aget_state(self, config):
        return self._state


def updates(data: dict, ns: tuple = ()) -> dict:
    return {"type": "updates", "ns": ns, "data": data}


def messages(chunk, metadata: dict | None = None, ns: tuple = ()) -> dict:
    return {"type": "messages", "ns": ns, "data": (chunk, metadata or {})}


class _State:
    def __init__(self, interrupts=()) -> None:
        self.tasks = [type("T", (), {"interrupts": list(interrupts)})()] if interrupts else []


class _Interrupt:
    def __init__(self, value) -> None:
        self.value = value


async def drive(agent: FakeAgent, session_id: str, **overrides) -> list[dict]:
    """跑一遍 ChatStream，返回解析后的事件列表。"""
    from runtime import concurrency
    from api.sse_core import ChatStream, StreamRequest

    kwargs = {
        "user_name": "contract_user",
        "kind": "test",
        "model": "deepseek-flash",
    }
    kwargs.update(overrides)

    slot = concurrency.acquire_task_slot(kwargs["user_name"])
    req = StreamRequest(
        agent=agent,
        session_id=session_id,
        thread_id=session_id,
        config={"configurable": {"thread_id": session_id}},
        message="测试消息",
        slot=slot,
        **kwargs,
    )
    events = []
    async for frame in ChatStream(req).generate():
        if frame.startswith("data: "):
            events.append(json.loads(frame[6:]))
    return events


@pytest.fixture
def _reset_concurrency():
    from runtime import concurrency

    concurrency.reset_for_tests()
    yield
    concurrency.reset_for_tests()


pytestmark = pytest.mark.usefixtures("_reset_concurrency")


class TestKernel:
    """通道分工：updates 出工具事件，messages 出正文 token。"""

    @pytest.mark.asyncio
    async def test_full_sequence(self, session_id, pricing_ready):
        tool_call = {"name": "calculate", "args": {"expression": "1+1"}, "id": "call-1"}
        agent = FakeAgent([
            updates({"model": {"messages": [AIMessage(content="", tool_calls=[tool_call])]}}),
            messages(AIMessageChunk(content="答案")),
            updates({"tools": {"messages": [
                ToolMessage(content='{"result": 2}', name="calculate", tool_call_id="call-1")
            ]}}),
        ])
        events = await drive(agent, session_id)
        types = [e["type"] for e in events]

        assert types[0] == "agent_start"
        assert "tool_call" in types
        assert "tool_result" in types
        assert "token" in types
        assert types[-1] == "done"
        assert "error" not in types

        # 正文只应包含模型说的话
        text = "".join(e["content"] for e in events if e["type"] == "token")
        assert text == "答案"
        assert '{"result"' not in text

    @pytest.mark.asyncio
    async def test_tool_message_not_streamed_as_token(self, session_id, pricing_ready):
        """ToolMessage 也走 messages 通道，但它的内容不能当正文推出去。"""
        agent = FakeAgent([
            messages(ToolMessage(content="工具返回的原始数据", name="t", tool_call_id="c")),
        ])
        events = await drive(agent, session_id)
        assert not [e for e in events if e["type"] == "token"]

    @pytest.mark.asyncio
    async def test_summarization_tokens_filtered(self, session_id, pricing_ready):
        """摘要中间件的模型调用也在 messages 通道上，它的 token 不能给用户看。"""
        agent = FakeAgent([
            messages(AIMessageChunk(content="## SESSION INTENT 摘要正文"),
                     {"lc_source": "summarization"}),
            messages(AIMessageChunk(content="真正的回答")),
        ])
        events = await drive(agent, session_id)
        text = "".join(e["content"] for e in events if e["type"] == "token")
        assert text == "真正的回答"
        assert "SESSION INTENT" not in text

    @pytest.mark.asyncio
    async def test_todo_list_event(self, session_id, pricing_ready):
        todos = [{"content": "步骤一", "status": "in_progress"}]
        agent = FakeAgent([updates({"model": {"todos": todos}})])
        events = await drive(agent, session_id)
        todo_events = [e for e in events if e["type"] == "updated_todo_list"]
        assert todo_events and todo_events[0]["todos"] == todos

    @pytest.mark.asyncio
    async def test_reasoning_content_becomes_thinking_token(self, session_id, pricing_ready):
        chunk = AIMessageChunk(content="", additional_kwargs={"reasoning_content": "我在想"})
        agent = FakeAgent([messages(chunk)])
        events = await drive(agent, session_id)
        thinking = [e for e in events if e["type"] == "thinking_token"]
        assert thinking and thinking[0]["content"] == "我在想"

    @pytest.mark.asyncio
    async def test_unknown_chunk_ignored(self, session_id, pricing_ready):
        """无法识别的 chunk 形态应被忽略，而不是让流崩掉。"""
        agent = FakeAgent([{"unexpected": "shape"}, messages(AIMessageChunk(content="ok"))])
        events = await drive(agent, session_id)
        assert "error" not in [e["type"] for e in events]
        assert any(e["type"] == "token" for e in events)

    @pytest.mark.asyncio
    async def test_legacy_tuple_chunk_supported(self, session_id, pricing_ready):
        """旧版 astream 的 (mode, data) 元组形态也要能解析。"""
        agent = FakeAgent([("messages", (AIMessageChunk(content="旧格式"), {}))])
        events = await drive(agent, session_id)
        assert any(e["type"] == "token" and e["content"] == "旧格式" for e in events)

    @pytest.mark.asyncio
    async def test_graph_input_is_dict(self, session_id, pricing_ready):
        """图输入必须是 {"messages": [...]}：传裸 list 会触发 InvalidUpdateError。"""
        agent = FakeAgent([messages(AIMessageChunk(content="x"))])
        await drive(agent, session_id)
        assert isinstance(agent.inputs[0], dict)
        assert "messages" in agent.inputs[0]

    @pytest.mark.asyncio
    async def test_user_message_persisted_before_streaming(self, session_id, pricing_ready):
        """用户消息要在流开始前落库，否则流挂了历史里就没有用户问了什么。"""
        from db import chat_store

        agent = FakeAgent([messages(AIMessageChunk(content="回答"))])
        await drive(agent, session_id)
        rows = chat_store.get_messages(session_id)
        assert rows[0]["msg_role"] == "user"
        assert rows[1]["msg_role"] == "assistant"
        assert rows[1]["content"] == "回答"

    @pytest.mark.asyncio
    async def test_interrupt_from_state_is_adapted(self, session_id, pricing_ready):
        """官方 HITL 载荷要转成 1.5 的 tool_approval 形状。"""
        payload = {
            "action_requests": [
                {"name": "request_plan_approval", "args": {"plan": "P"}, "description": "审批计划"}
            ],
            "review_configs": [{"action_name": "request_plan_approval",
                                "allowed_decisions": ["approve", "reject"]}],
        }
        agent = FakeAgent(
            [messages(AIMessageChunk(content=""))],
            state=_State(interrupts=[_Interrupt(payload)]),
        )
        events = await drive(agent, session_id)
        interrupts = [e for e in events if e["type"] == "interrupt"]
        assert interrupts, [e["type"] for e in events]
        ta = interrupts[0]["tool_approval"]
        assert ta["tool"] == "request_plan_approval"
        assert ta["reason"] == "审批计划"

    @pytest.mark.asyncio
    async def test_interrupt_from_stream_channel(self, session_id, pricing_ready):
        """流里捕获到的 __interrupt__ 优先于图状态。"""
        payload = {"action_requests": [{"name": "x", "args": {}, "description": "d"}]}
        agent = FakeAgent([updates({"__interrupt__": [_Interrupt(payload)]})])
        events = await drive(agent, session_id)
        assert any(e["type"] == "interrupt" for e in events)

    @pytest.mark.asyncio
    async def test_citations_from_tool_sources(self, session_id, pricing_ready):
        """来源由工具侧结构化产出，不依赖模型自觉。"""
        from tools.registry import source, sourced

        payload = json.dumps(
            sourced(data=[{"x": 1}], sources=[source("标题A", "https://e.invalid/a", "摘要", 0.9)]),
            ensure_ascii=False,
        )
        agent = FakeAgent([
            updates({"tools": {"messages": [
                ToolMessage(content=payload, name="search_demo_knowledge", tool_call_id="c1")
            ]}}),
        ])
        events = await drive(agent, session_id)
        cites = [e for e in events if e["type"] == "citations"]
        assert cites and len(cites[0]["items"]) == 1
        assert cites[0]["items"][0]["title"] == "标题A"
        assert cites[0]["model_only"] is False

    @pytest.mark.asyncio
    async def test_model_only_when_no_sources(self, session_id, pricing_ready):
        agent = FakeAgent([messages(AIMessageChunk(content="纯模型回答"))])
        events = await drive(agent, session_id)
        cites = [e for e in events if e["type"] == "citations"]
        assert cites and cites[0]["model_only"] is True

    @pytest.mark.asyncio
    async def test_usage_recorded_with_real_numbers(self, session_id, pricing_ready, unique_user):
        """有真实 usage 时按 provider 口径入账。"""
        chunk = AIMessageChunk(
            content="答",
            usage_metadata={"input_tokens": 1000, "output_tokens": 50, "total_tokens": 1050},
        )
        agent = FakeAgent([messages(chunk)])
        await drive(agent, session_id, user_name=unique_user)

        from billing import guard

        snapshot = guard.snapshot(unique_user)
        assert snapshot["requests"] == 1
        assert snapshot["input_tokens"] == 1000
        assert snapshot["output_tokens"] == 50

    @pytest.mark.asyncio
    async def test_billing_settles_once(self, session_id, pricing_ready, unique_user):
        """SSE 有正常/异常/断连三条退出路径，全汇到同一个 finally：不能重复扣费。"""
        from billing import guard

        agent = FakeAgent([
            messages(AIMessageChunk(
                content="答",
                usage_metadata={"input_tokens": 500, "output_tokens": 10, "total_tokens": 510},
            ))
        ])
        await drive(agent, session_id, user_name=unique_user)
        snapshot = guard.snapshot(unique_user)
        assert snapshot["requests"] == 1, f"重复入账：{snapshot}"

    @pytest.mark.asyncio
    async def test_concurrency_slot_released(self, session_id, pricing_ready):
        """流结束后并发位必须归还，否则服务会被自己锁死。"""
        from runtime import concurrency

        agent = FakeAgent([messages(AIMessageChunk(content="x"))])
        await drive(agent, session_id)
        assert concurrency.stats()["global_active"] == 0

    @pytest.mark.asyncio
    async def test_slot_released_on_error(self, session_id, pricing_ready):
        """异常路径同样要归还并发位。"""
        from runtime import concurrency

        class Boom(FakeAgent):
            async def astream(self, graph_input, **kwargs):
                raise RuntimeError("模拟图崩溃")
                yield  # pragma: no cover

        await drive(Boom([]), session_id)
        assert concurrency.stats()["global_active"] == 0

    @pytest.mark.asyncio
    async def test_error_frame_on_exception(self, session_id, pricing_ready):
        class Boom(FakeAgent):
            async def astream(self, graph_input, **kwargs):
                raise RuntimeError("模拟图崩溃")
                yield  # pragma: no cover

        events = await drive(Boom([]), session_id)
        errors = [e for e in events if e["type"] == "error"]
        assert errors and errors[0]["code"] == "internal_error"
        assert events[-1]["type"] == "done"

    @pytest.mark.asyncio
    async def test_resume_sends_command(self, session_id, pricing_ready):
        """恢复时必须喂 Command(resume=...)，而不是组装消息。"""
        from langgraph.types import Command

        agent = FakeAgent([messages(AIMessageChunk(content="继续"))])
        await drive(agent, session_id, resume_payload={"decisions": [{"type": "approve"}]})
        assert isinstance(agent.inputs[0], Command)