zhifu-platform/tests/business_wait_model_acceptance.py

import asyncio
import json
from pathlib import Path
from uuid import UUID

from sqlalchemy import func, select
from test_after_sales import scenario  # noqa: F401
from test_business_wait import confirm, decide, draft_wait

from zhifu.modules.after_sales.models import Approval, Ticket
from zhifu.modules.agent_runtime.budget_callback import CallAllowance
from zhifu.modules.agent_runtime.business_wait import reconcile
from zhifu.modules.agent_runtime.config import AgentSettings
from zhifu.modules.agent_runtime.models import AgentRun, ModelAttempt
from zhifu.modules.agent_runtime.worker import execute_one
from zhifu.modules.integrations.models import ExternalOperation


def reserve():
    with Path(".work/business-wait-model-allowance.json").open("x", encoding="utf-8") as stream:
        json.dump({"request_limit": 3, "automatic_retries": 0}, stream)


async def test_real_model_resumes_from_business_decision(scenario):  # noqa: F811
    sessions, state = scenario["app"].state.sessions, scenario["app"].state.state
    settings = AgentSettings(max_output_tokens=512, model_timeout=30)
    async with sessions() as session:
        assert await session.scalar(select(func.count()).select_from(AgentRun).where(
            AgentRun.status.in_(["QUEUED", "RUNNING"]))) == 0
    for role in ("agent", "memory"):
        async for _ in state.scan_iter(match=f"monitor:worker:{role}:*"):
            raise RuntimeError("当前模型专项要求停止其他生成工作进程")
    run_id, path, _, ticket, _ = await draft_wait(scenario)
    await decide(scenario, await confirm(scenario, ticket), "REJECTED")
    await reconcile(sessions)
    async with sessions.begin() as session:
        run = await session.get(AgentRun, run_id)
        run.model_name, run.model_url = settings.model_name, settings.model_url
        run.model_credential_env = settings.credential_env
    allowance = CallAllowance(3)
    await asyncio.to_thread(reserve)
    try:
        async with asyncio.timeout(100):
            assert await execute_one(sessions, state, settings,
                scenario["app"].state.knowledge.search, allowance=allowance)
            view = (await scenario["call"]("GET", path)).json()
            assert view["status"] == "COMPLETED", view
            assert view["answer"]
            async with sessions() as session:
                for model, expected in ((Ticket, 1), (Approval, 1), (ExternalOperation, 0)):
                    assert await session.scalar(select(func.count()).select_from(model).where(
                        model.tenant_id == scenario["tenant"].bytes)) == expected
                run = await session.get(AgentRun, run_id)
                assert any(tool["name"] == "get_ticket" for message in run.messages
                           for tool in message["data"].get("tool_calls", []))
                assert any(
                    message["type"] == "tool"
                    and json.loads(message["data"]["content"]).get("status") == "REJECTED"
                    for message in run.messages
                )
    finally:
        async with sessions() as session:
            run = await session.get(AgentRun, run_id)
            attempts = (await session.scalars(select(ModelAttempt).where(
                ModelAttempt.run_id == run_id))).all()
            report = {
                "run_id": str(UUID(bytes=run_id)), "status": run.status,
                "request_limit": 3, "actual_calls": allowance.started,
                "max_output_tokens": 512, "model_timeout_seconds": 30,
                "total_timeout_seconds": 100, "automatic_retries": 0,
                "attempts": [{"status": item.status, "usage": item.usage} for item in attempts],
                "messages": run.messages,
            }
        await asyncio.to_thread(Path(".work/business-wait-model.json").write_text,
            json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")