zhifu-platform/tests/test_interrupted_model_budget.py

import asyncio
from uuid import uuid4

import pytest
from test_after_sales import scenario  # noqa: F401
from test_money_budget import PRICE, PROVIDER, approve, budget, setup
from test_token_budget import sample

from zhifu.modules.agent_runtime.budget import change_entry, mark_pending_unknown, reserve
from zhifu.modules.agent_runtime.models import TokenAccount, TokenEntry
from zhifu.modules.billing.models import MoneyEntry
from zhifu.modules.identity.auth import now


@pytest.mark.parametrize("kind", ["agent", "memory"])
async def test_interrupted_budget_preserves_scope_and_late_settlement(scenario, kind):  # noqa: F811
    sessions, tenant = scenario["app"].state.sessions, scenario["tenant"].bytes
    await setup(scenario)
    await budget(scenario, 1000000)
    await approve(scenario, PRICE)
    observed = sample()
    reference, other_reference = uuid4().bytes, uuid4().bytes
    pending, completed, other = [uuid4().bytes for _ in range(3)]
    for identifier, target in (
        (pending, reference), (completed, reference), (other, other_reference)
    ):
        await reserve(sessions, tenant, identifier, kind, target,
                      PRICE["model"], 6000, provider=PROVIDER)
    await change_entry(sessions, tenant, completed, usage=observed["usage"])
    # 真实事务创建待计量记录；错误租户和调用类型不能改变其他业务的预算状态。
    await mark_pending_unknown(sessions, uuid4().bytes, kind, reference)
    await mark_pending_unknown(
        sessions, tenant, "memory" if kind == "agent" else "agent", reference)
    async with sessions() as session:
        assert (await session.get(TokenEntry, pending)).status == "PENDING"
    await asyncio.gather(*(
        mark_pending_unknown(sessions, tenant, kind, reference) for _ in range(3)))
    async with sessions() as session:
        assert (await session.get(TokenEntry, pending)).status == "UNKNOWN"
        assert (await session.get(MoneyEntry, pending)).status == "UNKNOWN"
        assert (await session.get(TokenEntry, pending)).version == 2
        assert (await session.get(TokenEntry, completed)).status == "SETTLED"
        assert (await session.get(TokenEntry, other)).status == "PENDING"
        account = await session.get(TokenAccount, (tenant, now().date()))
        assert account.reserved == 12000
        assert account.consumed == observed["usage"]["total_tokens"]
    # 使用已有供应商真实usage验证迟到结算与清理竞争，不生成模型响应。
    await asyncio.gather(
        change_entry(sessions, tenant, pending, usage=observed["usage"]),
        mark_pending_unknown(sessions, tenant, kind, reference),
    )
    async with sessions() as session:
        assert (await session.get(TokenEntry, pending)).status == "SETTLED"
        assert (await session.get(MoneyEntry, pending)).status == "SETTLED"
        account = await session.get(TokenAccount, (tenant, now().date()))
        assert account.reserved == 6000
        assert account.consumed == 2 * observed["usage"]["total_tokens"]