zhifu-platform/tests/provider_timeout_model_acceptance.py

import asyncio
import json
from pathlib import Path
from time import perf_counter
from uuid import UUID, uuid4

import pytest
from openai import APITimeoutError
from sqlalchemy import select
from test_after_sales import scenario  # noqa: F401
from test_money_budget import PRICE, approve, budget, setup

from zhifu.modules.agent_runtime.budget_callback import CallAllowance
from zhifu.modules.agent_runtime.config import AgentSettings
from zhifu.modules.agent_runtime.models import AgentRun, ModelAttempt, TokenEntry
from zhifu.modules.agent_runtime.worker import execute_one
from zhifu.modules.billing.models import MoneyEntry


def reserve_sample():
    with Path(".work/provider-timeout-recheck-allowance-20260922.json").open(
        "x", encoding="utf-8"
    ) as file:
        json.dump({"request_limit": 1, "output_limit": 128, "sdk_retries": 0}, file)


async def test_actual_timeout_keeps_unknown_token_and_money(scenario):  # noqa: F811
    settings = AgentSettings(max_output_tokens=128, model_timeout=30)
    sessions, state = scenario["app"].state.sessions, scenario["app"].state.state
    await setup(scenario)
    await budget(scenario, 1000000)
    await approve(scenario, {**PRICE, "model": settings.model_name, "provider": settings.model_url})
    key = str(uuid4())
    body = {"store_id": str(scenario["store"]), "message": "你好，请简短介绍售后服务。"}
    queued = await scenario["call"]("POST", "/openapi/v1/agent/runs", payload=body, key=key)
    assert queued.status_code == 202, queued.text
    run_id = UUID(queued.json()["run_id"]).bytes
    async with sessions.begin() as session:
        run = await session.get(AgentRun, run_id)
        run.release_snapshot = {**run.release_snapshot, "max_output_tokens": 128,
                                "model_timeout": 0.25}
    await asyncio.to_thread(reserve_sample)
    allowance, started = CallAllowance(1), perf_counter()
    report = {"status": "INCOMPLETE", "request_limit": 1, "max_output_tokens": 128,
              "sdk_retries": 0, "request_timeout_seconds": 0.25, "total_timeout_seconds": 90}
    try:
        async with asyncio.timeout(90):
            with pytest.raises((TimeoutError, APITimeoutError)) as failure:
                await execute_one(sessions, state, settings, allowance=allowance, run_ids=[run_id])
            repeated = await scenario["call"](
                "POST", "/openapi/v1/agent/runs", payload=body, key=key)
            assert repeated.json()["run_id"] == str(UUID(bytes=run_id))
            assert not await execute_one(
                sessions, state, settings, allowance=allowance, run_ids=[run_id])
            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()
                entries = (await session.scalars(select(TokenEntry).where(
                    TokenEntry.tenant_id == scenario["tenant"].bytes))).all()
                money = (await session.scalars(select(MoneyEntry).where(
                    MoneyEntry.tenant_id == scenario["tenant"].bytes))).all()
                report.update({
                    "run_id": str(UUID(bytes=run_id)), "run_status": run.status,
                    "error": type(failure.value).__name__, "supplier_delivery": "UNKNOWN",
                    "attempts": [{"status": row.status, "usage": row.usage} for row in attempts],
                    "token_entries": [{"status": row.status, "reserved": row.reserved,
                                       "usage": row.usage} for row in entries],
                    "money_entries": [{"status": row.status, "reserved": row.reserved}
                                      for row in money],
                    "price_basis": "真实审批的事务测试价格，不作为供应商实际费用",
                })
                assert run.status == "FAILED"
                assert len(attempts) == len(entries) == len(money) == allowance.started == 1
                assert attempts[0].status == entries[0].status == money[0].status == "UNKNOWN"
                assert entries[0].usage is None and entries[0].reserved > 0
                assert money[0].reserved > 0
            report["status"] = "PASSED"
    finally:
        report["actual_attempts"] = allowance.started
        report["elapsed_seconds"] = round(perf_counter() - started, 3)
        await asyncio.to_thread(Path(".work/provider-timeout-recheck-20260922.json").write_text,
            json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")