vent_agent_2.0/tests/test_business_agents.py

"""多业务路由、真实图结构化结果、启停及持久化恢复。"""
import asyncio
from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
import json
from types import SimpleNamespace

import httpx
import pytest
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessage, AIMessageChunk
from langchain_core.outputs import ChatGenerationChunk

from agents import agent_factory, agent_registry
from db import business_store as store, chat_store, db_conn
from extensions import registry as extensions
from extensions.registry import ExtensionRegistry, GatewayBinding
from extensions.main_fan.schemas import MainFanInput, MainFanResult
from extensions.main_fan.tools import inspect_main_fan_materials
from runtime import business
from runtime.business_models import JudgmentResult
from runtime.business_tasks import BusinessTaskController, submit
from runtime.engine import ExecutionResult
from runtime.environment import RuntimeEnvironment
from runtime.execution import execute
from sandbox_agent_helpers import sandbox_messages, tool_call
from runtime.python_sandbox import capability


INPUT = {"mine_code": "test-mine", "mine_name": "测试煤矿", "materials": []}


class StructuredAgent:
    def __init__(self, value, waiting=None):
        self.value, self.waiting, self.called = value, waiting, False

    async def aget_state(self, config):
        return SimpleNamespace(tasks=[])

    async def astream(self, *args, **kwargs):
        self.called = True
        yield {"type": "messages", "data": (AIMessageChunk(content="", usage_metadata={
            "input_tokens": 8, "output_tokens": 2, "total_tokens": 10}), {})}
        if self.waiting:
            await self.waiting.wait()
        yield {"type": "updates", "data": {"model": {"structured_response": self.value}}}


class ToolModel(GenericFakeChatModel):
    def bind_tools(self, tools, *, tool_choice=None, **kwargs):
        return self

    async def _astream(self, *args, **kwargs):
        message = next(self.messages)
        yield ChatGenerationChunk(message=AIMessageChunk(content=message.content,
            tool_calls=message.tool_calls, usage_metadata=message.usage_metadata))


@pytest.fixture
def two_businesses(monkeypatch):
    base = extensions.load_extensions()
    catalog = ExtensionRegistry()
    catalog.tools = dict(base.tools)
    catalog.agents = dict(base.agents)
    catalog.business_agents = dict(base.business_agents)
    catalog.register_business_agent(name="test_business", display_name="测试业务",
        description="测试独立业务", prompt="仅用于离线测试", tools=[], input_schema=MainFanInput,
        result_schema=JudgmentResult, gateway=GatewayBinding("test-business-code", "test-check"))
    monkeypatch.setattr(extensions, "_loaded", catalog)
    agent_registry.invalidate()
    yield catalog
    agent_registry.invalidate()


def test_business_registration_checks_platform_identity_and_schema():
    catalog = ExtensionRegistry()
    args = dict(display_name="业务", description="业务", prompt="处理业务", tools=[],
                input_schema=MainFanInput, result_schema=JudgmentResult)
    catalog.register_business_agent(name="one", gateway=GatewayBinding("code-one", "check-one"), **args)
    with pytest.raises(ValueError, match="重复"):
        catalog.register_business_agent(name="two", gateway=GatewayBinding("code-one", "check-two"), **args)
    assert list(catalog.business_agents) == ["one"]
    with pytest.raises(ValueError, match="BaseModel"):
        catalog.register_business_agent(name="bad", **{**args, "input_schema": dict})
    with pytest.raises(ValueError, match="并发"):
        catalog.register_business_agent(name="bad", max_concurrency=0, **args)


def test_named_cache_is_single_flight_and_invalidated_independently(monkeypatch, two_businesses):
    calls = []
    def build(name="main"):
        calls.append(name)
        return object()
    monkeypatch.setattr(agent_factory, "build_agent", build)
    with ThreadPoolExecutor(max_workers=8) as pool:
        fans = list(pool.map(lambda _: agent_registry.get_agent("main_fan"), range(16)))
    assert all(a is fans[0] for a in fans)
    other = agent_registry.get_agent("test_business")
    main = agent_registry.get_agent()
    assert len({id(fans[0]), id(other), id(main)}) == 3
    assert calls.count("main_fan") == 1
    agent_registry.invalidate("main_fan")
    assert agent_registry.get_agent("main_fan") is not fans[0]
    assert agent_registry.get_agent("test_business") is other
    with pytest.raises(ValueError, match="未注册"):
        agent_registry.get_agent("does_not_exist")


@pytest.mark.asyncio
async def test_real_business_graph_uses_own_tools_and_structured_result(monkeypatch):
    if not capability()['available']:
        pytest.skip('当前平台未安装执行沙箱')
    model = ToolModel(messages=sandbox_messages('named-user', INPUT, encoded_lists=True))
    monkeypatch.setattr(agent_factory.model_factory, "get_model", lambda: model)
    async with RuntimeEnvironment(prewarm=False, billing_retry=False):
        session = chat_store.create_session("named-user")
        result = await execute(agent_name="main_fan", user_name="named-user", session_id=session["id"],
                               message=json.dumps(INPUT))
        assert result.status == "success", result.error
        assert result.structured_output["status"] == "needs_materials"
        assert store.session_agent(session["id"]) == "main_fan"
        assert chat_store.get_messages(session["id"])[-1]["content"]
        captured = {}
        monkeypatch.setattr("deepagents.create_deep_agent", lambda **kwargs: captured.update(kwargs) or kwargs)
        agent_factory.build_agent("main_fan")
        names = {tool.name for tool in captured["tools"]}
        assert {"prepare_main_fan_task", "fetch_main_fan_data", "query_fan_timeseries",
                "query_fan_sensor_realtime", "query_fan_equipment_realtime", "query_fan_equipment_timeseries",
                "run_python", "submit_main_fan_report"} <= names
        assert "compare_main_fan_data" not in names
        assert "search_demo_knowledge" not in names
        assert captured["response_format"] is None


@pytest.mark.asyncio
async def test_real_coordinator_can_delegate_business_without_overwriting_its_answer(monkeypatch):
    if not capability()['available']:
        pytest.skip('当前平台未安装执行沙箱')
    def script():
        yield tool_call('task', {'subagent_type': 'main_fan', 'description': json.dumps(INPUT)})
        yield from sandbox_messages('coordinator', INPUT)
        yield AIMessage(content="主通风机业务尚缺少材料，当前无法形成正式结论。")
    model = ToolModel(messages=script())
    monkeypatch.setattr(agent_factory.model_factory, "get_model", lambda: model)
    async with RuntimeEnvironment(prewarm=False, billing_retry=False):
        session = chat_store.create_session("coordinator")
        result = await execute(user_name="coordinator", session_id=session["id"], message="请让主通风机智能体检查材料")
        assert result.status == "success", result.error
        assert result.output == "主通风机业务尚缺少材料，当前无法形成正式结论。"
        assert result.structured_output is None


@pytest.mark.asyncio
async def test_named_chat_history_keeps_agent_and_rejects_switch(monkeypatch, pricing_ready):
    from main import app
    fake = StructuredAgent(inspect_main_fan_materials([], "m", "矿"))
    selected = []
    def get_agent(name="main"):
        selected.append(name)
        return fake
    monkeypatch.setattr(agent_registry, "get_agent", get_agent)
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
        first = await client.post("/api/chat", data={"message": "检查材料", "agent_name": "main_fan", "user_name": "chat-user"})
        assert first.status_code == 200, first.text
        session_id = first.headers["X-Session-Id"]
        assert first.headers["X-Agent-Name"] == "main_fan"
        events = [json.loads(frame[6:]) for frame in first.text.split("\n\n") if frame.startswith("data: ")]
        assert any(e["type"] == "token" and "无法形成正式研判" in e["content"] for e in events)
        again = await client.post("/api/chat", data={"message": "继续", "session_id": session_id, "user_name": "chat-user"})
        assert again.status_code == 200
        assert selected == ["main_fan", "main_fan"]
        history = await client.get(f"/api/chat/history/{session_id}", params={"user_name": "chat-user"})
        assert history.json()["agent_name"] == "main_fan"
        switched = await client.post("/api/chat", data={"message": "切换", "session_id": session_id, "agent_name": "main", "user_name": "chat-user"})
        assert switched.status_code == 409


@pytest.mark.asyncio
async def test_structured_child_result_does_not_replace_root(monkeypatch, pricing_ready):
    parent = StructuredAgent({"root": 1})
    original = parent.astream
    async def stream(*args, **kwargs):
        async for item in original(*args, **kwargs):
            yield item
        yield {"type": "updates", "ns": ("child:123",), "data": {"model": {"structured_response": {"child": 2}}}}
    parent.astream = stream
    session = chat_store.create_session("structured")
    result = await execute(user_name="structured", session_id=session["id"], message="任务", agent=parent)
    assert result.structured_output == {"root": 1}


@pytest.mark.asyncio
async def test_session_cannot_change_agent_and_stop_is_scoped(pricing_ready, two_businesses):
    fan = StructuredAgent(inspect_main_fan_materials([], **{k: INPUT[k] for k in ("mine_code", "mine_name")}))
    session = chat_store.create_session("scope")
    result = await execute(agent_name="main_fan", user_name="scope", session_id=session["id"], message="任务", agent=fan)
    assert result.status == "success"
    other = StructuredAgent({})
    with pytest.raises(ValueError, match="绑定"):
        await execute(agent_name="test_business", user_name="scope", session_id=session["id"], message="任务", agent=other)
    assert not other.called
    await business.stop("main_fan")
    fresh = chat_store.create_session("scope")
    with pytest.raises(business.BusinessStopped):
        await execute(agent_name="main_fan", user_name="scope", session_id=fresh["id"], message="任务", agent=fan)
    business.ensure_enabled("test_business")
    await business.start("main_fan")
    business.ensure_enabled("main_fan")


@pytest.mark.asyncio
async def test_stop_cancels_named_execution_and_releases_limits(pricing_ready):
    gate = asyncio.Event()
    agent = StructuredAgent({}, gate)
    session = chat_store.create_session("stop-user")
    task = asyncio.create_task(execute(agent_name="main_fan", user_name="stop-user", session_id=session["id"], message="任务", agent=agent))
    for _ in range(100):
        if agent.called:
            break
        await asyncio.sleep(.01)
    assert agent.called
    await business.stop("main_fan")
    with pytest.raises(asyncio.CancelledError):
        await task
    assert not business._counts.get("main_fan")
    assert not business._active.get("main_fan")
    await business.start("main_fan")


@pytest.mark.asyncio
async def test_data_gateway_status_does_not_stop_business_execution(monkeypatch, pricing_ready, two_businesses):
    spec = two_businesses.business_agents["main_fan"]
    monkeypatch.setitem(two_businesses.business_agents, "main_fan", replace(spec, gateway=GatewayBinding("fan-code", "fan-check")))
    monkeypatch.setenv("SHANXI_GATEWAY_ENABLED", "1")
    store.set_state("main_fan", gateway_online=False)
    business.ensure_enabled("main_fan")
    agent = StructuredAgent(inspect_main_fan_materials([], "m", "矿"))
    session = chat_store.create_session("remote-stop")
    result = await execute(agent_name="main_fan", user_name="remote-stop", session_id=session["id"], message="任务", agent=agent)
    assert result.status == "success" and agent.called
    assert store.state("main_fan")["local_enabled"] == 1
    business.ensure_enabled("main_fan")
    business.ensure_enabled("test_business")


def test_job_idempotency_rejects_different_input_and_skeleton_is_not_completed():
    first = submit("main_fan", INPUT, request_id="one")
    assert submit("main_fan", INPUT, request_id="one")["id"] == first["id"]
    with pytest.raises(ValueError, match="不同"):
        submit("main_fan", {**INPUT, "fan_id": "other"}, request_id="one")
    with pytest.raises(ValueError):
        MainFanResult.model_validate({**inspect_main_fan_materials([], "m", "n"), "status": "completed"})


@pytest.mark.asyncio
async def test_controller_runs_skeleton_without_publishing_and_cleans_up():
    calls = []
    async def executor(**kwargs):
        calls.append(kwargs)
        value = inspect_main_fan_materials([], "test-mine", "测试煤矿")
        return ExecutionResult("fake-run", kwargs["session_id"], "success", "", "", {}, [], [], value)
    controller = BusinessTaskController(executor=executor)
    await controller.start()
    try:
        job = submit("main_fan", INPUT)
        for _ in range(100):
            if store.get(job["id"])["status"] != "queued":
                await asyncio.sleep(.03)
            if store.get(job["id"])["status"] == "needs_materials":
                break
            await asyncio.sleep(.01)
        saved = store.get(job["id"])
        assert saved["status"] == "needs_materials", saved
        assert saved["push_status"] == "none"
        assert calls[0]["agent_name"] == "main_fan"
        assert store.metrics("main_fan")["completed"] == 0
    finally:
        await controller.close()
    assert not [task for task in asyncio.all_tasks() if task.get_name().startswith("business-") and not task.done()]


def test_sqlite_reopen_retains_jobs_and_pending_pushes(monkeypatch, tmp_path):
    adapter = db_conn._SqliteAdapter(str(tmp_path / "business.sqlite"))
    monkeypatch.setattr(db_conn, "_adapter", adapter)
    monkeypatch.setattr(db_conn, "_schema_ready", False)
    db_conn.init_schema()
    job, _ = store.enqueue("main_fan", "u", INPUT, [], request_id="persist")
    store.update(job["id"], status="completed", push_status="pending", result_json='{"result":"saved"}')
    running, _ = store.enqueue("main_fan", "u", INPUT, [], request_id="running")
    store.update(running["id"], status="running")
    adapter.close()
    store.recover()
    assert store.get(running["id"])["status"] == "interrupted"
    assert store.pending_pushes(db_conn.fmt_dt())[0]["id"] == job["id"]
    adapter.close()


@pytest.mark.asyncio
async def test_business_api_requires_own_token_and_keeps_job_ownership(monkeypatch):
    from main import app
    from runtime import business_tasks
    monkeypatch.setenv("BUSINESS_API_TOKEN", "local-test-token")
    monkeypatch.setattr(business_tasks, "_controller", object())
    headers = {"Authorization": "Bearer local-test-token"}
    base = "/api/business-agents/main_fan"
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
        assert (await client.get("/api/business-agents")).json()["items"][0]["name"] == "main_fan"
        assert (await client.post(base + "/jobs", json={"input": INPUT})).status_code == 401
        created = await client.post(base + "/jobs", headers=headers, json={"input": INPUT, "request_id": "api-1"})
        assert created.status_code == 202, created.text
        job_id = created.json()["id"]
        assert (await client.get(base + "/jobs/" + job_id, headers=headers, params={"user_name": "other"})).status_code == 404
        assert (await client.post(base + "/manage?action=stop", headers=headers)).status_code == 200
        assert (await client.post(base + "/jobs", headers=headers, json={"input": INPUT})).status_code == 409
        assert (await client.post(base + "/manage?action=start", headers=headers)).status_code == 200