zhifu-platform/tests/billing_retry_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 APIStatusError
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 import change_entry
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
from zhifu.operations.token_budget import inspect


def reserve_sample():
    with Path(".work/long-billing-model-allowance.json").open("x", encoding="utf-8") as stream:
        json.dump({"request_limit": 3, "sdk_retries": 0, "max_output_tokens": 128}, stream)


async def test_provider_error_and_explicit_retry_accounting(scenario):  # noqa: F811
    sessions, state = scenario["app"].state.sessions, scenario["app"].state.state
    settings = AgentSettings(max_output_tokens=128, model_timeout=30)
    await setup(scenario)
    await budget(scenario, 1000000)
    invalid_model = "zhifu-unavailable-" + uuid4().hex
    for name in (invalid_model, settings.model_name):
        await approve(scenario, {**PRICE, "model": name, "provider": settings.model_url})
    for role in ("agent", "memory"):
        async for _ in state.scan_iter(match=f"monitor:worker:{role}:*"):
            raise RuntimeError("采样要求其他生成工作进程停止")
    body = {"store_id": str(scenario["store"]), "message": "你好，请仅回复一句问候。"}
    key = str(uuid4())
    response = await scenario["call"]("POST", "/openapi/v1/agent/runs", payload=body, key=key)
    assert response.status_code == 202, response.text
    failed_id = UUID(response.json()["run_id"]).bytes
    # 在独立测试租户设置无效模型配置，真实供应商返回配置错误，全部计入请求上限。
    async with sessions.begin() as session:
        run = await session.get(AgentRun, failed_id)
        run.model_name = invalid_model
        run.release_snapshot = {**run.release_snapshot, "max_output_tokens": 128,
                                "model_timeout": 30, "model_name": invalid_model}
    await asyncio.to_thread(reserve_sample)
    allowance, started = CallAllowance(3), perf_counter()
    http_status, provider_request_id, success_id = None, None, None
    try:
        async with asyncio.timeout(100):
            with pytest.raises(APIStatusError) as failure:
                await execute_one(sessions, state, settings, allowance=allowance,
                                  run_ids=[failed_id])
            http_status = failure.value.status_code
            provider_request_id = failure.value.request_id
            assert 400 <= http_status < 500
            repeated = await scenario["call"](
                "POST", "/openapi/v1/agent/runs", payload=body, key=key)
            assert repeated.json()["run_id"] == str(UUID(bytes=failed_id))
            assert not await execute_one(sessions, state, settings, allowance=allowance,
                                         run_ids=[failed_id])
            # 配置恢复后由客户使用新请求标识明确重试，旧失败记录和未知预占继续保留。
            retry = await scenario["call"]("POST", "/openapi/v1/agent/runs", payload=body)
            assert retry.status_code == 202, retry.text
            success_id = UUID(retry.json()["run_id"]).bytes
            async with sessions.begin() as session:
                run = await session.get(AgentRun, success_id)
                run.release_snapshot = {**run.release_snapshot, "max_output_tokens": 128,
                                        "model_timeout": 30}
            assert await execute_one(sessions, state, settings, allowance=allowance,
                                     run_ids=[success_id])
            async with sessions() as session:
                failed = await session.get(AgentRun, failed_id)
                success = await session.get(AgentRun, success_id)
                assert failed.status == "FAILED" and success.status == "COMPLETED"
                rows = (await session.scalars(select(TokenEntry).where(
                    TokenEntry.tenant_id == scenario["tenant"].bytes))).all()
                assert len(rows) == allowance.started <= 3
                unknown = [row for row in rows if row.status == "UNKNOWN"]
                settled = [row for row in rows if row.status == "SETTLED"]
                assert len(unknown) == 1 and settled
                for row in rows:
                    money = await session.get(MoneyEntry, row.id)
                    assert money.status == row.status
            before = await inspect(sessions, scenario["tenant"].bytes)
            assert before["reserved"] == unknown[0].reserved
            assert before["consumed"] == sum(row.usage["total_tokens"] for row in settled)
            for row in settled:
                await change_entry(sessions, row.tenant_id, row.id, usage=row.usage)
            assert await inspect(sessions, scenario["tenant"].bytes) == before
    finally:
        async with sessions() as session:
            attempts = (await session.scalars(select(ModelAttempt).where(
                ModelAttempt.run_id.in_([value for value in (failed_id, success_id) if value])
            ))).all()
            entries = (await session.scalars(select(TokenEntry).where(
                TokenEntry.tenant_id == scenario["tenant"].bytes))).all()
            report = {
                "request_limit": 3, "actual_calls": allowance.started,
                "max_output_tokens": 128, "sdk_retries": 0,
                "timeout_seconds": 100, "seconds": perf_counter() - started,
                "provider_error_status": http_status, "provider_request_id": provider_request_id,
                "failed_run": str(UUID(bytes=failed_id)),
                "retry_run": str(UUID(bytes=success_id)) if success_id else None,
                "attempts": [{"run_id": str(UUID(bytes=row.run_id)),
                              "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],
                "price_basis": "事务测试价格；未核对供应商金额",
            }
        await asyncio.to_thread(Path(".work/billing-retry-model.json").write_text,
            json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")