zhifu-platform/tests/mixed_capacity_acceptance.py

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

from sqlalchemy import func, select
from test_after_sales import scenario  # noqa: F401
from test_knowledge_memory import cleanup_vectors, document

from zhifu.modules.agent_runtime.capacity import acquire, release
from zhifu.modules.agent_runtime.config import AgentSettings
from zhifu.modules.agent_runtime.models import ModelAttempt
from zhifu.modules.agent_runtime.worker import claim
from zhifu.operations.tenant_quota import QuotaChange, configure


async def mixed_admission(
    scenario,  # noqa: F811
    *,
    rounds=8,
    b_requests=12,
    knowledge_queries=4,
    pause=0.01,
    report_path="docs/verification/mixed-capacity.json",
):
    """有界真实准入负载：领取后取消，不执行生成图；知识请求调用本地向量服务。"""
    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()
    identity = (await scenario["call"]("GET", "/openapi/v1/auth/me", "foreign")).json()
    other = UUID(identity["tenant_id"])
    change = QuotaChange(
        expected_version=0,
        interactive_limit=3,
        background_limit=1,
        interactive_reserved=1,
        operator="acceptance",
        reason="商家执行份额验收",
    )
    await configure(sessions, other.bytes, change)
    path = await document(scenario)
    held = []
    background = uuid4().hex
    results = {"a_statuses": {}, "b_admission_ms": [], "b_queue_ms": [], "knowledge_ms": []}

    async def enqueue(actor, store):
        return await scenario["call"](
            "POST",
            "/openapi/v1/agent/runs",
            actor,
            {"store_id": store, "message": "查询订单"},
        )

    async def consume_one():
        identifier = uuid4().hex
        selected = await claim(
            sessions,
            settings.tenant_concurrency,
            state=state,
            settings=settings,
            capacity_id=identifier,
        )
        if selected:
            held.append(identifier)
        else:
            await release(state, identifier)
        return selected, identifier

    async def flood_a():
        for _ in range(rounds):
            responses = await asyncio.gather(
                *(enqueue("customer", str(scenario["store"])) for _ in range(8))
            )
            for response in responses:
                assert response.status_code == 429, response.text
                status = str(response.status_code)
                results["a_statuses"][status] = results["a_statuses"].get(status, 0) + 1
            await asyncio.sleep(pause)

    async def requests_b():
        for _ in range(b_requests):
            started = time.perf_counter()
            response = await enqueue("foreign", identity["grants"][0]["store_id"])
            assert response.status_code == 202, response.text
            accepted = time.perf_counter()
            results["b_admission_ms"].append((accepted - started) * 1000)
            selected, identifier = await consume_one()
            assert selected and str(UUID(bytes=selected[0])) == response.json()["run_id"]
            results["b_queue_ms"].append((time.perf_counter() - accepted) * 1000)
            response = await scenario["call"](
                "POST", f"/openapi/v1/agent/runs/{UUID(bytes=selected[0])}/cancel", "foreign"
            )
            assert response.status_code == 200
            await release(state, identifier)
            await asyncio.sleep(pause)

    async def knowledge():
        started = time.perf_counter()
        response = await scenario["call"](
            "POST", path + "/publish", "agent", {"expected_version": 1}
        )
        assert response.status_code == 200, response.text
        results["knowledge_publish_ms"] = (time.perf_counter() - started) * 1000
        for _ in range(knowledge_queries):
            started = time.perf_counter()
            response = await scenario["call"](
                "POST",
                f"/openapi/v1/stores/{scenario['store']}/knowledge/search",
                payload={"query": "键盘如何清洁"},
            )
            assert response.status_code == 200 and response.json()["items"], response.text
            results["knowledge_ms"].append((time.perf_counter() - started) * 1000)
            await asyncio.sleep(pause)

    try:
        assert await acquire(state, settings, "background", background, reserved={other.hex: 1})
        for _ in range(3):
            assert (await enqueue("customer", str(scenario["store"]))).status_code == 202
        for _ in range(2):
            selected, _ = await consume_one()
            assert selected
        # A 与后台持续占用三个额度，B 的一个预留机会仍可用。
        load_started = time.perf_counter()
        async with asyncio.timeout(60):
            await asyncio.gather(flood_a(), requests_b(), knowledge())
        async with sessions() as session:
            attempts = await session.scalar(
                select(func.count())
                .select_from(ModelAttempt)
                .where(ModelAttempt.tenant_id.in_([scenario["tenant"].bytes, other.bytes]))
            )
        assert attempts == 0
        results["generation_attempts"] = attempts
        results["b_accepted"] = b_requests
        results["load_seconds"] = time.perf_counter() - load_started
        results["duration_limit_seconds"] = 60
        results["a_request_limit"] = rounds * 8
        results["b_request_limit"] = b_requests
        results["knowledge_query_limit"] = knowledge_queries
        results["generation_token_checks"] = "人工验收，不阻塞开发"
        await asyncio.to_thread(
            Path(report_path).write_text,
            json.dumps(results, ensure_ascii=False, indent=2) + "\n",
            encoding="utf-8",
        )
    finally:
        for identifier in [background, *held]:
            await release(state, identifier)
        await cleanup_vectors(scenario["tenant"])


async def test_mixed_admission_without_generation(scenario):  # noqa: F811
    await mixed_admission(scenario)