zhifu-platform/tests/token_budget_model_acceptance.py

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

from sqlalchemy import select
from test_after_sales import scenario  # noqa: F401

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.runtime import build_graph
from zhifu.modules.agent_runtime.worker import claim, execute_one
from zhifu.modules.memory.worker import claim as claim_memory
from zhifu.modules.memory.worker import process
from zhifu.operations.token_budget import BudgetChange, entries, inspect


async def test_bounded_actual_usage_enters_budget(scenario):  # noqa: F811
    """最多三次供应商请求，Agent 与记忆提取共用次数控制，不进行完整付费压测。"""
    sessions, state = scenario["app"].state.sessions, scenario["app"].state.state
    for role in ("agent", "memory"):
        async for _ in state.scan_iter(match=f"monitor:worker:{role}:*"):
            raise RuntimeError("请停止生成模型工作进程并等待心跳过期")
    settings = AgentSettings(max_output_tokens=128)
    allowance = CallAllowance(3)
    call = scenario["call"]
    async with asyncio.timeout(100):
        response = await call(
            "POST",
            "/openapi/v1/agent/runs",
            payload={"store_id": str(scenario["store"]), "message": "你好，请仅回复一句问候。"},
        )
        assert response.status_code == 202, response.text
        run_id, token = await claim(sessions, settings.tenant_concurrency)
        await build_graph(sessions, settings, allowance=allowance).ainvoke(
            {"run_id": run_id, "token": token, "next_node": "model"}
        )
        response = await call(
            "PUT",
            "/openapi/v1/me/memory/consent",
            payload={"expected_revision": 0, "enabled": True},
        )
        assert response.status_code == 200
        response = await call(
            "POST", "/openapi/v1/conversations", payload={"store_id": str(scenario["store"])}
        )
        conversation_id = response.json()["conversation_id"]
        response = await call(
            "POST",
            f"/openapi/v1/conversations/{conversation_id}/messages",
            payload={"text": "以后请一直用中文简短回复。", "expected_control_version": 1},
        )
        assert response.status_code == 201
        pending = response.json()["payload"]["run_id"]
        assert (await call("POST", f"/openapi/v1/agent/runs/{pending}/cancel")).status_code == 200
        job, previous = await claim_memory(sessions)
        await process(sessions, job, previous, settings, allowance=allowance)
    async with sessions() as session:
        rows = (
            await session.scalars(
                select(TokenEntry).where(TokenEntry.tenant_id == scenario["tenant"].bytes)
            )
        ).all()
        assert {row.kind for row in rows} == {"agent", "memory"}
        assert len(rows) == allowance.started <= 3
        assert all(row.status == "SETTLED" and row.usage["total_tokens"] > 0 for row in rows)
        samples = [
            {
                "attempt": str(UUID(bytes=row.id)),
                "kind": row.kind,
                "model": row.model_name,
                "reserved": row.reserved,
                "usage": row.usage,
            }
            for row in rows
        ]
    budget = await inspect(sessions, scenario["tenant"].bytes)
    assert budget["reserved"] == 0
    assert budget["consumed"] == sum(item["usage"]["total_tokens"] for item in samples)
    await asyncio.to_thread(
        Path("docs/verification/token-budget-sample.json").write_text,
        json.dumps(
            {"model_requests": allowance.started, "samples": samples, "budget": budget},
            ensure_ascii=False,
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )


async def test_zero_budget_worker_returns_without_sending(scenario):  # noqa: F811
    sessions = scenario["app"].state.sessions
    tenant = scenario["tenant"].bytes
    await inspect(
        sessions,
        tenant,
        change=BudgetChange(
            expected_version=0, token_limit=0, operator="acceptance", reason="验证调用前预算拒绝"
        ),
    )
    response = await scenario["call"](
        "POST",
        "/openapi/v1/agent/runs",
        payload={"store_id": str(scenario["store"]), "message": "你好，请仅回复一句问候。"},
    )
    assert response.status_code == 202
    run_id = UUID(response.json()["run_id"]).bytes
    assert await execute_one(
        sessions,
        scenario["app"].state.state,
        AgentSettings(max_output_tokens=128),
        allowance=CallAllowance(1),
    )
    async with sessions() as session:
        run = await session.get(AgentRun, run_id)
        assert run.status == "FAILED" and run.error_code == "TOKEN_BUDGET_EXCEEDED"
        attempt = await session.scalar(select(ModelAttempt).where(ModelAttempt.run_id == run_id))
        assert attempt.status == "NOT_SENT"
    assert (await entries(sessions, tenant))["items"] == []