zhifu-platform/tests/test_token_budget.py

import asyncio
import json
from datetime import timedelta
from pathlib import Path
from uuid import UUID, uuid4

import pytest
from langchain_core.messages import HumanMessage
from sqlalchemy import func, select
from test_after_sales import scenario  # noqa: F401

from zhifu.core.errors import BusinessError
from zhifu.modules.after_sales.models import BusinessEvent
from zhifu.modules.agent_runtime.budget import change_entry, reserve
from zhifu.modules.agent_runtime.budget_callback import BudgetCallback
from zhifu.modules.agent_runtime.models import TokenAccount, TokenEntry
from zhifu.modules.identity.auth import now
from zhifu.operations.token_budget import BudgetChange, entries, inspect

pytestmark = pytest.mark.integration


def sample():
    # 回放受限真实采样的 usage，只验证账本事务，不模拟模型响应。
    document = json.loads(Path("docs/verification/token-budget-sample.json").read_text("utf-8"))
    return document["samples"][0]


async def set_limit(context, limit, version=0):
    return await inspect(
        context["app"].state.sessions,
        context["tenant"].bytes,
        change=BudgetChange(
            expected_version=version,
            token_limit=limit,
            operator="integration",
            reason="预算事务验收",
        ),
    )


async def test_concurrent_reservation_and_duplicate_settlement(scenario):  # noqa: F811
    sessions, tenant = scenario["app"].state.sessions, scenario["tenant"].bytes
    observed = sample()
    amount = observed["usage"]["total_tokens"]
    await set_limit(scenario, amount * 2)
    ids = [uuid4().bytes for _ in range(8)]
    results = await asyncio.gather(
        *(
            reserve(sessions, tenant, identifier, "agent", uuid4().bytes, observed["model"], amount)
            for identifier in ids
        ),
        return_exceptions=True,
    )
    assert sum(result is None for result in results) == 2
    assert all(
        result is None
        or (isinstance(result, BusinessError) and result.code == "TOKEN_BUDGET_EXCEEDED")
        for result in results
    )
    accepted = [
        identifier for identifier, result in zip(ids, results, strict=True) if result is None
    ]
    await asyncio.gather(
        *(change_entry(sessions, tenant, accepted[0], usage=observed["usage"]) for _ in range(6))
    )
    report = await inspect(sessions, tenant)
    assert report["consumed"] == amount and report["reserved"] == amount
    with pytest.raises(BusinessError, match="标识已存在"):
        await reserve(
            sessions, tenant, accepted[0], "agent", uuid4().bytes, observed["model"], amount
        )
    async with sessions() as session:
        count = await session.scalar(
            select(func.count())
            .select_from(BusinessEvent)
            .where(
                BusinessEvent.aggregate_id == accepted[0],
                BusinessEvent.event_type == "token.settled",
            )
        )
        assert count == 1


async def test_unknown_release_late_usage_and_tenant_isolation(scenario):  # noqa: F811
    sessions, tenant = scenario["app"].state.sessions, scenario["tenant"].bytes
    observed = sample()
    identifier = uuid4().bytes
    await set_limit(scenario, 10000)
    await reserve(sessions, tenant, identifier, "memory", uuid4().bytes, observed["model"], 5000)
    assert await change_entry(sessions, tenant, identifier) == "UNKNOWN"
    assert (await inspect(sessions, tenant))["reserved"] == 5000
    other = (await scenario["call"]("GET", "/openapi/v1/auth/me", "foreign")).json()
    with pytest.raises(BusinessError) as error:
        await change_entry(
            sessions, UUID(other["tenant_id"]).bytes, identifier, usage=observed["usage"]
        )
    assert error.value.status == 404
    with pytest.raises(BusinessError, match="版本"):
        await change_entry(
            sessions,
            tenant,
            identifier,
            release_reason="确认未计费",
            operator="integration",
            expected_version=1,
        )
    await change_entry(
        sessions,
        tenant,
        identifier,
        release_reason="确认未计费",
        operator="integration",
        expected_version=2,
    )
    assert (await inspect(sessions, tenant))["reserved"] == 0
    await set_limit(scenario, 0, 1)
    await change_entry(sessions, tenant, identifier, usage=observed["usage"])
    assert (await inspect(sessions, tenant))["available"] == -observed["usage"]["total_tokens"]


async def test_configuration_conflict_and_historical_window(scenario):  # noqa: F811
    sessions, tenant = scenario["app"].state.sessions, scenario["tenant"].bytes
    old = now().date() - timedelta(days=1)
    empty = await inspect(sessions, tenant, old)
    assert empty["recorded"] is False and empty["token_limit"] is None
    changes = await asyncio.gather(
        set_limit(scenario, 10000), set_limit(scenario, 20000), return_exceptions=True
    )
    assert sum(isinstance(result, BusinessError) for result in changes) == 1
    observed = sample()
    identifier = uuid4().bytes
    await reserve(sessions, tenant, identifier, "agent", uuid4().bytes, observed["model"], 100)
    # 使用真实数据库创建前一日期账户，验证迟到 usage 仍结算在原预占日期。
    async with sessions.begin() as session:
        session.add(
            TokenAccount(tenant_id=tenant, period=old, token_limit=100, reserved=100, consumed=0)
        )
        await session.flush()
        current = await session.get(TokenAccount, (tenant, now().date()))
        current.reserved = 0
        row = await session.get(TokenEntry, identifier)
        row.period = old
    await change_entry(sessions, tenant, identifier, usage=observed["usage"])
    assert (await inspect(sessions, tenant, old))["consumed"] == observed["usage"]["total_tokens"]
    assert (await inspect(sessions, tenant))["consumed"] == 0
    page = await entries(sessions, tenant, limit=1)
    assert len(page["items"]) == 1 and page["next_cursor"] is None


async def test_budget_callback_rejects_before_external_request(scenario):  # noqa: F811
    sessions, tenant = scenario["app"].state.sessions, scenario["tenant"].bytes
    await set_limit(scenario, 0)
    callback = BudgetCallback(sessions, tenant, "agent", uuid4().bytes, sample()["model"], 128)
    assert callback.raise_error and callback.run_inline
    with pytest.raises(BusinessError) as error:
        await callback.on_chat_model_start({}, [[HumanMessage(content="你好")]], run_id=uuid4())
    assert error.value.code == "TOKEN_BUDGET_EXCEEDED"
    assert (await entries(sessions, tenant))["items"] == []


async def test_conflicting_usage_keeps_original_settlement(scenario):  # noqa: F811
    sessions, tenant = scenario["app"].state.sessions, scenario["tenant"].bytes
    observed = sample()
    identifier = uuid4().bytes
    await reserve(sessions, tenant, identifier, "agent", uuid4().bytes, observed["model"], 10000)
    await change_entry(sessions, tenant, identifier, usage=observed["usage"])
    conflict = dict(observed["usage"])
    conflict["output_tokens"] += 1
    conflict["total_tokens"] += 1
    with pytest.raises(BusinessError) as error:
        await change_entry(sessions, tenant, identifier, usage=conflict)
    assert error.value.code == "TOKEN_USAGE_CONFLICT"
    assert (await inspect(sessions, tenant))["consumed"] == observed["usage"]["total_tokens"]