zhifu-platform/scripts/check_browser_model.py

import argparse
import asyncio
import json
import os
import secrets
import time
from pathlib import Path
from urllib.parse import urlparse
from uuid import UUID, uuid4

from anyio import Path as AsyncPath
from redis.asyncio import Redis
from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import async_sessionmaker

from zhifu.core.config import Settings
from zhifu.core.database import create_engine
from zhifu.core.errors import BusinessError
from zhifu.modules.agent_runtime.budget_callback import CallAllowance
from zhifu.modules.agent_runtime.capacity import release
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
from zhifu.modules.billing.ledger import ADMISSION_ERRORS
from zhifu.modules.billing.models import MoneyEntry
from zhifu.modules.conversations.models import Conversation
from zhifu.modules.conversations.service import finish_run
from zhifu.modules.identity.access_models import AccessSession, Principal, StoreGrant
from zhifu.modules.identity.auth import password_hasher
from zhifu.modules.identity.models import Store, Tenant
from zhifu.modules.knowledge.vector import retrieval_client

BUDGET_PATH = ".work/frontend-verification/customer-model-allowance.json"


async def read_allowance():
    path = AsyncPath(BUDGET_PATH)
    ledger = json.loads(await path.read_text(encoding="utf-8")) if await path.exists() else {
        "maximum_provider_attempts": 3, "runs": [],
    }
    if ledger["maximum_provider_attempts"] != 3:
        raise ValueError("专项累计供应商次数上限必须为 3")
    remaining = 3 - sum(item["charged_attempts"] for item in ledger["runs"])
    if remaining <= 0:
        raise RuntimeError("本专项累计供应商尝试已达到上限，禁止继续生成")
    return ledger, remaining


def write_report(report, path):
    destination = Path(path).resolve()
    workspace = Path(__file__).resolve().parents[1] / ".work"
    if not destination.is_relative_to(workspace):
        raise ValueError("报告必须保存在项目 .work 目录")
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8",
    )


async def prepare_account(sessions, configuration):
    if configuration.environment == "production":
        raise ValueError("浏览器采样账户只允许在 development 或 test 环境创建")
    source = json.loads(await AsyncPath(".work/frontend-accounts.json").read_text(encoding="utf-8"))
    actor = uuid4()
    tenant_id, store_id = UUID(source["tenant_id"]), UUID(source["store_id"])
    username, password = "browser-model-" + actor.hex, secrets.token_urlsafe(24)
    encoded = await asyncio.to_thread(password_hasher.hash, password)
    async with sessions.begin() as session:
        tenant = await session.get(Tenant, tenant_id.bytes)
        store = await session.get(Store, store_id.bytes)
        if tenant.status != "ACTIVE" or store.status != "ACTIVE":
            raise RuntimeError("开发商家或店铺未启用")
        session.add(Principal(
            id=actor.bytes, tenant_id=tenant_id.bytes, username=username,
            password_hash=encoded, kind="CUSTOMER",
        ))
        await session.flush()
        session.add(StoreGrant(
            tenant_id=tenant_id.bytes, principal_id=actor.bytes,
            store_id=store_id.bytes, permission="orders:read",
        ))
    await asyncio.to_thread(write_report, {
        "tenant_id": str(tenant_id), "store_id": str(store_id),
        "actor_id": str(actor), "username": username, "password": password,
    }, ".work/browser-model-account.json")


async def cleanup_account(sessions):
    credentials = json.loads(
        await AsyncPath(".work/browser-model-account.json").read_text(encoding="utf-8"),
    )
    actor_id = UUID(credentials["actor_id"]).bytes
    async with sessions.begin() as session:
        principal = await session.get(Principal, actor_id, with_for_update=True)
        if principal.username != credentials["username"] or not principal.username.startswith(
            "browser-model-",
        ):
            raise RuntimeError("清理目标与专项账户不同")
        # 关闭专项账户与全部会话，保留已发生的生成用量和业务事件供核验。
        principal.status = "SUSPENDED"
        principal.version += 1
        await session.execute(update(AccessSession).where(
            AccessSession.principal_id == actor_id,
        ).values(revoked=True))
    await AsyncPath(".work/browser-model-account.json").unlink()


async def check_queue(sessions, state, expected=None):
    # 专项期间只允许当前浏览器任务，持续生成工作进程会使次数限制失去完整范围。
    for role in ("agent", "memory"):
        async for _ in state.scan_iter(match=f"monitor:worker:{role}:*"):
            raise RuntimeError("生成工作进程仍有有效心跳，请停止后等待心跳过期")
    async with sessions() as session:
        pending = set((await session.scalars(select(AgentRun.id).where(
            AgentRun.status.in_(["QUEUED", "RUNNING"]),
        ))).all())
    if pending != ({expected} if expected else set()):
        raise RuntimeError("队列存在其他运行或目标任务状态不同，专项未调用模型")


async def snapshot(sessions, identifier):
    async with sessions() as session:
        run = await session.get(AgentRun, identifier)
        attempts = (await session.scalars(select(ModelAttempt).where(
            ModelAttempt.run_id == identifier,
        ).order_by(ModelAttempt.created_at, ModelAttempt.id))).all()
        entries = (await session.scalars(select(TokenEntry).where(
            TokenEntry.reference_id.in_([item.id for item in attempts]),
            TokenEntry.kind == "agent",
        ))).all() if attempts else []
        money = (await session.scalars(select(MoneyEntry).where(
            MoneyEntry.id.in_([item.id for item in entries]),
        ))).all() if entries else []
        return {
            "run_status": run.status if run else None,
            "run_error_code": run.error_code if run else None,
            "model_attempts": [{
                "id": str(UUID(bytes=item.id)), "status": item.status, "usage": item.usage,
            } for item in attempts],
            "token_entries": [{
                "id": str(UUID(bytes=item.id)),
                "model_attempt_id": str(UUID(bytes=item.reference_id)),
                "model": item.model_name, "status": item.status,
                "reserved": item.reserved, "usage": item.usage,
                "cache_read_tokens": (
                    item.usage.get("input_token_details", {}).get("cache_read")
                    if item.usage else None
                ),
            } for item in entries],
            "money_entries": [{
                "id": str(UUID(bytes=item.id)), "status": item.status,
                "charged_micro_cny": item.charged, "reserved_micro_cny": item.reserved,
                "basis": item.basis, "price_id": item.price["price_id"], "price": item.price,
            } for item in money],
        }


async def main(args):
    configuration = Settings()
    engine = create_engine(configuration)
    sessions = async_sessionmaker(engine, expire_on_commit=False)
    state = Redis.from_url(configuration.state_url.get_secret_value(), socket_timeout=3)
    allowance = CallAllowance(0)
    capacity_id = uuid4().hex
    started = time.monotonic()
    identifier = UUID(args.run_id).bytes if args.run_id else None
    token = None
    budget_lock = None
    ledger = None
    ledger_run = None
    report = {"maximum_provider_attempts": 3, "maximum_duration_seconds": args.timeout_seconds}
    try:
        async with asyncio.timeout(args.timeout_seconds):
            if args.cleanup_account:
                await cleanup_account(sessions)
                print(json.dumps({"cleanup": "passed", "model_requests": 0}))
                return
            await check_queue(sessions, state, identifier)
            if args.preflight:
                if not os.environ.get("DEEPSEEK_API_KEY"):
                    raise RuntimeError("进程环境缺少 DEEPSEEK_API_KEY")
                await read_allowance()
                if args.prepare_account:
                    await prepare_account(sessions, configuration)
                print(json.dumps({"preflight": "passed", "model_requests": 0}))
                return
            if identifier is None or not args.conversation_id or not args.report:
                raise ValueError("执行需要 run-id、conversation-id 和 report")
            settings = AgentSettings(max_output_tokens=1024)
            async with sessions() as session:
                run = await session.get(AgentRun, identifier)
                if run.conversation_id != UUID(args.conversation_id).bytes or run.model_calls:
                    raise RuntimeError("运行所属会话不同或已有模型调用，专项终止")
                endpoint = urlparse(run.model_url)
                if endpoint.scheme != "https" or endpoint.hostname != "api.deepseek.com":
                    raise RuntimeError("本专项只允许已配置的 DeepSeek HTTPS 模型地址")
                effective = run.release_snapshot or {
                    "max_output_tokens": settings.max_output_tokens,
                    "model_timeout": settings.model_timeout,
                    "model_extra_body": settings.model_extra_body,
                }
                output_limit = effective["max_output_tokens"]
                if not 1 <= output_limit <= 1024 or not 0 < effective["model_timeout"] <= 30:
                    raise RuntimeError("发布快照的输出或超时参数超过本次采样上限")
                extra = effective["model_extra_body"]
                for name in ("max_tokens", "max_completion_tokens", "max_output_tokens"):
                    if name in extra and extra[name] > output_limit:
                        raise RuntimeError("模型附加参数超过发布快照的输出上限")
                if extra.get("n", 1) != 1:
                    raise RuntimeError("本专项仅允许单候选响应")
                os.environ[run.model_credential_env] = os.environ["DEEPSEEK_API_KEY"]
                report.update({
                    "run_id": args.run_id, "conversation_id": args.conversation_id,
                    "model": run.model_name, "provider": run.model_url,
                    "effective_max_output_tokens": output_limit,
                    "effective_model_timeout_seconds": effective["model_timeout"],
                    "automatic_provider_retries": 0,
                })
            # 跨次执行共享三次总额度；进程被终止时保留已预留次数，禁止自动增加额度。
            budget_lock = state.lock(
                "verification:browser-model:allowance", timeout=100, blocking_timeout=2,
            )
            if not await budget_lock.acquire():
                raise RuntimeError("另一次浏览器模型专项正在执行")
            ledger, remaining = await read_allowance()
            ledger_run = {
                "run_id": args.run_id, "status": "RESERVED", "charged_attempts": remaining,
            }
            ledger["runs"].append(ledger_run)
            await asyncio.to_thread(write_report, ledger, BUDGET_PATH)
            allowance = CallAllowance(remaining)
            claimed = await claim(
                sessions, settings.tenant_concurrency, state=state,
                settings=settings, capacity_id=capacity_id,
            )
            if claimed is None or claimed[0] != identifier:
                raise RuntimeError("领取结果与浏览器运行不同，专项未调用模型")
            _, token = claimed
            async with retrieval_client() as knowledge:
                graph = build_graph(sessions, settings, knowledge.search, allowance=allowance)
                await graph.ainvoke(
                    {"run_id": identifier, "token": token, "next_node": "model"},
                    {"recursion_limit": 24},
                )
    except BaseException as failure:
        report["failure_type"] = type(failure).__name__
        failure_code = (
            failure.code if isinstance(failure, BusinessError) else type(failure).__name__
        )
        report["failure_code"] = failure_code
        not_sent = isinstance(failure, BusinessError) and failure.code in (
            ADMISSION_ERRORS | {"MODEL_SAMPLE_LIMIT"}
        )
        if identifier and token is not None:
            # 终止仅更新本人领取的运行，真实已发生用量由模型回调保留。
            async with sessions.begin() as session:
                run = await session.get(AgentRun, identifier)
                await session.get(Conversation, run.conversation_id, with_for_update=True)
                await session.refresh(run, with_for_update=True)
                if run.status == "RUNNING" and run.lease_token == token:
                    run.status = "FAILED"
                    run.error_code = failure_code[:64]
                    run.lease_until = None
                    await finish_run(session, run, "run.failed", {"code": run.error_code})
                    await session.execute(update(ModelAttempt).where(
                        ModelAttempt.run_id == identifier, ModelAttempt.status == "STARTED",
                    ).values(status="NOT_SENT" if not_sent else "UNKNOWN"))
        raise
    finally:
        try:
            await release(state, capacity_id)
            if identifier and args.report:
                report.update(await snapshot(sessions, identifier))
                # 回调先检查次数和预算；只有预算事务通过后才进入供应商请求。
                report["callback_admission_attempts"] = allowance.started
                report["provider_attempts_started"] = len(report["token_entries"])
                report["provider_attempts_received"] = sum(
                    item["status"] == "RECEIVED" for item in report["model_attempts"]
                )
                report["provider_attempts_result_unknown"] = sum(
                    item["status"] in {"UNKNOWN", "PENDING"} for item in report["token_entries"]
                )
                report["admission_rejections"] = sum(
                    item["status"] == "NOT_SENT" for item in report["model_attempts"]
                )
                if ledger_run is not None:
                    ledger_run.update({
                        "status": "RECORDED",
                        "charged_attempts": report["provider_attempts_started"],
                        "callback_admission_attempts": allowance.started,
                        "failure_code": report.get("failure_code"),
                    })
                    await asyncio.to_thread(write_report, ledger, BUDGET_PATH)
                    report["cumulative_provider_attempts"] = sum(
                        item["charged_attempts"] for item in ledger["runs"]
                    )
                report["elapsed_seconds"] = round(time.monotonic() - started, 3)
                await asyncio.to_thread(write_report, report, args.report)
                evidence = str(Path(args.report).with_name(
                    f"customer-model-worker-{args.run_id}.json",
                ))
                await asyncio.to_thread(write_report, report, evidence)
        finally:
            if budget_lock is not None and await budget_lock.owned():
                await budget_lock.release()
            await state.aclose()
            await engine.dispose()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="浏览器 AI 有界验证的单次工作进程")
    parser.add_argument("--preflight", action="store_true")
    parser.add_argument("--prepare-account", action="store_true")
    parser.add_argument("--cleanup-account", action="store_true")
    parser.add_argument("--run-id")
    parser.add_argument("--conversation-id")
    parser.add_argument("--report")
    parser.add_argument("--timeout-seconds", type=int, default=80, choices=range(1, 86))
    asyncio.run(main(parser.parse_args()))