zhifu-platform/tests/provider_limits_model_acceptance.py

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

import pytest
from langchain_core.messages import message_to_dict
from langchain_openai import ChatOpenAI
from openai import APITimeoutError, DefaultAsyncHttpxClient, DefaultHttpxClient
from sqlalchemy import select
from test_after_sales import scenario  # noqa: F401

from zhifu.core.model_parameters import chat_parameters
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


def reserve_sample():
    # 本次授权独立保存，已有额度文件阻止重复付费执行；失败也保留实际请求证据。
    with Path(".work/provider-limits-allowance-20260922.json").open("x", encoding="utf-8") as file:
        json.dump({"request_limit": 2, "output_limit": 128, "sdk_retries": 0}, file)


async def test_provider_output_limit_and_actual_request_timeout(scenario):  # noqa: F811
    settings = AgentSettings(max_output_tokens=128, model_timeout=30)
    sessions, state = scenario["app"].state.sessions, scenario["app"].state.state
    report_path = Path(".work/provider-limits-20260922.json")
    await asyncio.to_thread(reserve_sample)
    allowance, started = CallAllowance(2), perf_counter()
    report = {"request_limit": 2, "max_output_tokens": 128, "sdk_retries": 0,
              "total_timeout_seconds": 90, "status": "INCOMPLETE"}
    run_id = None
    try:
        async with asyncio.timeout(90):
            # 直接使用与Agent和记忆相同的SDK参数，验证供应商在输出上限结束生成。
            client = ChatOpenAI(
                model=settings.model_name, base_url=settings.model_url,
                api_key=os.environ[settings.credential_env], temperature=0,
                timeout=30, max_retries=0, cache=False, use_responses_api=False,
                **chat_parameters(settings.model_url, 128, settings.model_extra_body),
                http_client=DefaultHttpxClient(timeout=30),
                http_async_client=DefaultAsyncHttpxClient(timeout=30),
            )
            try:
                allowance.take()
                response = await client.ainvoke(
                    "请依次列出从1到1000的每个整数，每个整数单独一行，不要省略，也不要其他说明。")
                report["output_sample"] = message_to_dict(response)
                assert response.usage_metadata
                assert response.usage_metadata["output_tokens"] <= 128
                assert response.response_metadata["finish_reason"] == "length"
            finally:
                await client.root_async_client.close()
                client.root_client.close()

            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
            # 正常供应商请求使用很短的真实期限，不生成错误响应或替换网络客户端。
            # 期限包含预算登记；供应商是否已收到请求未知时保留预占和UNKNOWN。
            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}
            with pytest.raises((TimeoutError, APITimeoutError)) as failure:
                await execute_one(sessions, state, settings, allowance=allowance, run_ids=[run_id])
            report["timeout_error"] = type(failure.value).__name__
            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()
                assert run.status == "FAILED"
                assert len(attempts) == 1 and attempts[0].status == "UNKNOWN"
                assert len(entries) == 1 and entries[0].status == "UNKNOWN"
                assert entries[0].reserved > 0 and entries[0].usage is None
                report["timeout_run"] = {
                    "run_id": str(UUID(bytes=run_id)), "status": run.status,
                    "error_code": run.error_code, "attempt_status": attempts[0].status,
                    "token_status": entries[0].status, "reserved_tokens": entries[0].reserved,
                    "usage": entries[0].usage, "request_timeout_seconds": 0.25,
                    "supplier_delivery": "UNKNOWN",
                }
            assert allowance.started == 2
            report["status"] = "PASSED"
    finally:
        report["actual_attempts"] = allowance.started
        report["elapsed_seconds"] = round(perf_counter() - started, 3)
        report["scope"] = "SDK真实截断1次、Agent真实超时尝试1次；无供应商账单与限流结论"
        await asyncio.to_thread(report_path.write_text,
            json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")