vent_agent_2.0/tests/test_execution.py

"""中立执行入口、HTTP 适配和模块依赖边界的回归测试。"""
import ast
import asyncio
from contextlib import aclosing
import json
from pathlib import Path
import subprocess
import sys
from types import SimpleNamespace

import pytest
from langchain_core.messages import AIMessageChunk

import config
from db import chat_store, billing_store, run_store, schedule_store
from runtime import concurrency
from runtime.execution import admit, build_execution, execute


@pytest.fixture(autouse=True)
def isolated_runtime(monkeypatch, tmp_path):
    monkeypatch.setattr(config, "ROOT", tmp_path)
    concurrency.reset_for_tests()
    yield
    concurrency.reset_for_tests()


class Agent:
    def __init__(self, *, wait=None, failure=None, pending=()):
        self.wait, self.failure, self.pending = wait, failure, pending
        self.closed = False
        self.called = False

    async def aget_state(self, options):
        return SimpleNamespace(tasks=[SimpleNamespace(interrupts=self.pending)])

    async def astream(self, graph_input, **kwargs):
        self.called = True
        try:
            yield {"type": "messages", "data": (AIMessageChunk(content="回答", usage_metadata={
                "input_tokens": 8, "output_tokens": 2, "total_tokens": 10}), {})}
            if self.wait:
                await self.wait.wait()
            if self.failure:
                raise self.failure
        finally:
            self.closed = True


async def prepared(agent):
    session = chat_store.create_session("u")
    admission = await admit("u", session["id"])
    return await build_execution(admission, agent=agent, session_id=session["id"],
        thread_id=session["id"], user_name="u", message="问题", trace_id="execution-test",
        config={"configurable": {"thread_id": session["id"]}})


def assert_released():
    assert concurrency.stats()["global_active"] == 0
    assert not concurrency._sessions


@pytest.mark.asyncio
async def test_events_are_dicts_and_result_is_available_after_cleanup(pricing_ready):
    agent = Agent()
    task = await prepared(agent)
    with pytest.raises(RuntimeError, match="尚未完成"):
        task.result()
    async with aclosing(task.events()) as events:
        received = [event async for event in events]
    assert all(isinstance(event, dict) for event in received)
    assert all(event["session_id"] == task.req.session_id for event in received)
    assert all(event["trace_id"] == "execution-test" for event in received)
    assert next(e["content"] for e in received if e["type"] == "token") == "回答"
    result = task.result()
    assert result.status == "success" and result.output == "回答"
    assert result.usage["total_tokens"] == 10
    assert agent.closed
    assert_released()
    with pytest.raises(RuntimeError, match="只能执行一次"):
        await task.run()
    assert len(billing_store.recent_usage("u")) == 1


@pytest.mark.asyncio
async def test_sse_adapter_preserves_event_payload(pricing_ready):
    from api.sse_core import stream_sse
    from runtime.events import EventBuilder
    from api.sse_events import serialize_event
    builder = EventBuilder("s", "t", encode=lambda event: event)
    raw = builder.tool_call("calc", args={"value": "中文\n换行"}, tool_call_id="c")
    assert json.loads(serialize_event(raw)[6:]) == raw
    task = await prepared(Agent())
    frames = [frame async for frame in stream_sse(task)]
    events = [json.loads(frame[6:]) for frame in frames]
    assert events[-1]["type"] == "done" and events[-1]["run_id"] == task.req.run_id
    assert all(frame.startswith("data: ") and frame.endswith("\n\n") for frame in frames)
    assert_released()


@pytest.mark.asyncio
@pytest.mark.parametrize("cancel", [False, True])
async def test_http_adapter_close_and_cancellation_finalize_inner_execution(pricing_ready, cancel):
    from api.sse_core import stream_sse
    agent = Agent(wait=asyncio.Event())
    task = await prepared(agent)
    stream = stream_sse(task)
    while json.loads((await anext(stream))[6:])["type"] != "token":
        pass
    if cancel:
        pending = asyncio.create_task(anext(stream))
        await asyncio.sleep(0)
        pending.cancel()
        with pytest.raises(asyncio.CancelledError):
            await pending
    await stream.aclose()
    assert task.status == "cancelled"
    assert agent.closed
    assert_released()
    assert len(billing_store.recent_usage("u")) == 1
    assert run_store.list_runs(task.req.session_id)[0]["status"] == "cancelled"


@pytest.mark.asyncio
async def test_encoder_failure_still_closes_execution(monkeypatch, pricing_ready):
    from api import sse_core
    original = sse_core.serialize_event
    def broken(event):
        if event["type"] == "token":
            raise ValueError("encoding failed")
        return original(event)
    monkeypatch.setattr(sse_core, "serialize_event", broken)
    agent = Agent(wait=asyncio.Event())
    task = await prepared(agent)
    with pytest.raises(ValueError, match="encoding failed"):
        async for _ in sse_core.stream_sse(task):
            pass
    assert agent.closed
    assert_released()
    assert len(billing_store.recent_usage("u")) == 1


@pytest.mark.asyncio
async def test_batch_entry_reuses_admission_and_persistence(pricing_ready):
    sessions = [chat_store.create_session("u") for _ in range(2)]
    results = await asyncio.gather(*(execute(user_name="u", session_id=s["id"], message="批处理", agent=Agent()) for s in sessions))
    assert all(r.status == "success" and r.output == "回答" for r in results)
    assert len({r.run_id for r in results}) == 2
    assert len(billing_store.recent_usage("u")) == 2
    for result in results:
        assert run_store.list_runs(result.session_id)[0]["kind"] == "batch"
        assert len(chat_store.get_messages(result.session_id)) == 2
    assert_released()


@pytest.mark.asyncio
async def test_batch_rejects_wrong_owner_and_pending_approval(pricing_ready):
    from runtime.interrupts import ResumeError
    session = chat_store.create_session("u")
    agent = Agent()
    with pytest.raises(concurrency.SessionMissingError):
        await execute(user_name="other", session_id=session["id"], message="q", agent=agent)
    assert not agent.called
    agent.pending = [SimpleNamespace(id="pending", value={"action_requests": [{"name": "write"}]})]
    with pytest.raises(ResumeError) as error:
        await execute(user_name="u", session_id=session["id"], message="q", agent=agent)
    assert error.value.code == "pending_interrupt"
    assert not agent.called
    assert_released()


@pytest.mark.asyncio
async def test_scheduler_does_not_encode_sse(monkeypatch, pricing_ready):
    from scheduler.runner import run_schedule
    from api import sse_core
    def forbidden(event):
        pytest.fail("定时任务不能经过 SSE 编码")
    monkeypatch.setattr(sse_core, "serialize_event", forbidden)
    monkeypatch.setattr("agents.agent_registry.get_agent", lambda: Agent())
    schedule = schedule_store.create_schedule("u", "report", "0 9 * * *", "生成报告")
    result = await run_schedule(schedule["id"])
    assert result["ok"] and result["output"] == "回答"
    assert result["input_tokens"] == 8
    assert_released()


@pytest.mark.asyncio
async def test_failed_run_returns_error_and_settles_usage(pricing_ready):
    task = await prepared(Agent(failure=RuntimeError("model error")))
    result = await task.run()
    assert result.status == "failed" and result.error == "RuntimeError"
    assert result.output == "回答"
    assert len(billing_store.recent_usage("u")) == 1
    assert_released()


def test_resume_validation_is_independent_of_http():
    from runtime.interrupts import ResumeDecision, ResumeError, pack, resume_value, token_for
    item = SimpleNamespace(id="approval", value={"action_requests": [{"name": "a"}, {"name": "b"}]})
    state = SimpleNamespace(tasks=[SimpleNamespace(interrupts=[item])])
    decision = ResumeDecision(approved=True, interrupt_token=token_for(pack(item)))
    assert len(resume_value(decision, state)["approval"]["decisions"]) == 2
    with pytest.raises(ResumeError) as error:
        resume_value(ResumeDecision(approved=True, interrupt_token="old"), state)
    assert error.value.code == "stale_interrupt"


def test_runtime_and_scheduler_do_not_import_http_layer():
    root = Path(__file__).resolve().parents[1]
    for directory in ("runtime", "scheduler", "agents"):
        for path in (root / directory).rglob("*.py"):
            for node in ast.walk(ast.parse(path.read_text())):
                modules = ([node.module or ""] if isinstance(node, ast.ImportFrom) else
                           [alias.name for alias in node.names] if isinstance(node, ast.Import) else [])
                assert not any(m.split(".")[0] in {"api", "fastapi", "starlette"} for m in modules), path
    # 新进程防止测试会话里提前导入 API 掩盖反向依赖。
    script = """
import runtime.engine
import runtime.execution
import runtime.interrupts
import scheduler.runner
import sys
assert not any(name.split('.')[0] in {'api', 'fastapi', 'starlette'} for name in sys.modules)
"""
    subprocess.run([sys.executable, "-B", "-c", script], cwd=root, check=True)