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"]