zhifu-platform/tests/test_agent_runtime.py

import asyncio
import hashlib
import json
from datetime import timedelta
from uuid import UUID, uuid4

import pytest
from langchain_core.messages import AIMessage, message_to_dict
from sqlalchemy import func, select
from test_after_sales import scenario  # noqa: F401

from zhifu.core.errors import BusinessError
from zhifu.modules.after_sales.models import Ticket
from zhifu.modules.agent_runtime.capacity import acquire, keys
from zhifu.modules.agent_runtime.config import AgentSettings
from zhifu.modules.agent_runtime.models import AgentRun
from zhifu.modules.agent_runtime.runtime import build_tool_graph
from zhifu.modules.agent_runtime.worker import claim
from zhifu.modules.identity.access_models import AccessSession
from zhifu.modules.identity.auth import now

pytestmark = pytest.mark.integration


async def tool_run(context, name, arguments):
    token = context["headers"]["customer"]["Authorization"].split()[1]
    async with context["app"].state.sessions.begin() as session:
        run = AgentRun(
            id=uuid4().bytes,
            tenant_id=context["tenant"].bytes,
            actor_id=context["accounts"]["customer"].bytes,
            store_id=context["store"].bytes,
            token_hash=hashlib.sha256(token.encode()).digest(),
            request_key=uuid4().hex,
            request_hash=hashlib.sha256(json.dumps(arguments).encode()).digest(),
            messages=[
                message_to_dict(
                    AIMessage(
                        content="",
                        tool_calls=[
                            {
                                "name": name,
                                "args": arguments,
                                "id": uuid4().hex,
                                "type": "tool_call",
                            }
                        ],
                    )
                )
            ],
            next_node="tools",
            model_name="",
            model_url="",
            model_credential_env="",
            created_at=now(),
            deadline_at=now() + timedelta(minutes=10),
        )
        session.add(run)
        await session.flush()
        return run


async def test_tool_graph_creates_one_proposal_on_replay(scenario):  # noqa: F811
    run = await tool_run(scenario, "create_exchange_proposal", scenario["proposal"])
    sessions = scenario["app"].state.sessions
    run_id, token = await claim(sessions, 2)
    graph = build_tool_graph(sessions)
    state = {"run_id": run_id, "token": token, "next_node": "tools"}
    await graph.ainvoke(state)
    await graph.ainvoke(state)
    async with sessions() as session:
        count = await session.scalar(
            select(func.count())
            .select_from(Ticket)
            .where(Ticket.tenant_id == scenario["tenant"].bytes)
        )
        assert count == 1
        current = await session.get(AgentRun, run.id)
        result = json.loads(current.messages[-1]["data"]["content"])
        assert result["status"] == "DRAFT"
        assert current.step == 1


@pytest.mark.parametrize("name", ["confirm", "decide", "submit_exchange", "cancel"])
async def test_unauthorized_tools_are_rejected(scenario, name):  # noqa: F811
    await tool_run(scenario, name, {})
    sessions = scenario["app"].state.sessions
    run_id, token = await claim(sessions, 2)
    with pytest.raises(BusinessError, match="工具不在授权清单"):
        await build_tool_graph(sessions).ainvoke({"run_id": run_id, "token": token})


async def test_tool_cannot_read_another_customer_order(scenario):  # noqa: F811
    run = await tool_run(scenario, "get_order", {"order_id": str(scenario["order_id"])})
    sessions = scenario["app"].state.sessions
    other_token = scenario["headers"]["other"]["Authorization"].split()[1]
    async with sessions.begin() as session:
        current = await session.get(AgentRun, run.id)
        current.actor_id = scenario["accounts"]["other"].bytes
        current.token_hash = hashlib.sha256(other_token.encode()).digest()
    run_id, token = await claim(sessions, 2)
    with pytest.raises(BusinessError, match="不可访问"):
        await build_tool_graph(sessions).ainvoke({"run_id": run_id, "token": token})


async def test_revoked_session_stops_tool_execution(scenario):  # noqa: F811
    run = await tool_run(scenario, "search_orders", {})
    sessions = scenario["app"].state.sessions
    run_id, token = await claim(sessions, 2)
    async with sessions.begin() as session:
        access = await session.get(AccessSession, run.token_hash)
        access.revoked = True
    with pytest.raises(BusinessError, match="访问令牌无效"):
        await build_tool_graph(sessions).ainvoke({"run_id": run_id, "token": token})


async def test_cancelled_run_rejects_old_execution(scenario):  # noqa: F811
    run = await tool_run(scenario, "search_orders", {})
    sessions = scenario["app"].state.sessions
    run_id, token = await claim(sessions, 2)
    response = await scenario["call"]("POST", f"/openapi/v1/agent/runs/{UUID(bytes=run.id)}/cancel")
    assert response.status_code == 200
    with pytest.raises(BusinessError, match="执行权已失效"):
        await build_tool_graph(sessions).ainvoke({"run_id": run_id, "token": token})
    response = await scenario["call"](
        "GET", f"/openapi/v1/agent/runs/{UUID(bytes=run.id)}", "foreign"
    )
    assert response.status_code == 404


async def test_tenant_and_global_concurrency_limits(scenario):  # noqa: F811
    for _ in range(3):
        await tool_run(scenario, "search_orders", {})
    sessions = scenario["app"].state.sessions
    claims = await asyncio.gather(*(claim(sessions, 2) for _ in range(3)))
    assert sum(value is not None for value in claims) == 2
    state = scenario["app"].state.state
    key = "verify:agent:capacity:" + uuid4().hex
    try:
        results = await asyncio.gather(
            *(
                acquire(state, AgentSettings(), "background", uuid4().hex, prefix=key)
                for _ in range(8)
            )
        )
        assert sum(results) == 2
    finally:
        await state.delete(*keys(key))


async def test_expired_lease_resumes_committed_tool_without_duplicate(scenario):  # noqa: F811
    run = await tool_run(scenario, "create_exchange_proposal", scenario["proposal"])
    sessions = scenario["app"].state.sessions
    run_id, token = await claim(sessions, 2)
    await build_tool_graph(sessions).ainvoke({"run_id": run_id, "token": token})
    async with sessions.begin() as session:
        current = await session.get(AgentRun, run.id)
        current.lease_until = now() - timedelta(seconds=1)
    recovered_id, recovered_token = await claim(sessions, 2)
    assert recovered_id == run_id and recovered_token > token
    with pytest.raises(BusinessError, match="执行权已失效"):
        await build_tool_graph(sessions).ainvoke({"run_id": run_id, "token": token})
    result = await build_tool_graph(sessions).ainvoke({"run_id": run_id, "token": recovered_token})
    assert result["next_node"] == "model"
    async with sessions() as session:
        count = await session.scalar(
            select(func.count())
            .select_from(Ticket)
            .where(Ticket.tenant_id == scenario["tenant"].bytes)
        )
        assert count == 1