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()))