zhifu-platform/scripts/recovery_probe.py

import argparse
import asyncio
import json
import sys
from uuid import UUID

from sqlalchemy.ext.asyncio import async_sessionmaker

from zhifu.app import create_app
from zhifu.core.config import Settings
from zhifu.core.database import create_engine
from zhifu.modules.agent_runtime.models import AgentRun
from zhifu.modules.agent_runtime.runtime import build_tool_graph
from zhifu.modules.agent_runtime.worker import claim
from zhifu.modules.integrations.fulfillment import execute_due


async def main(identifier, stage):
    # 仅供独立验收进程使用；生产租约和工具事务保持完整，不修改租约时间。
    settings = Settings()
    if settings.environment == "production":
        raise RuntimeError("恢复探针仅允许开发和测试环境")
    create_app()
    engine = create_engine(settings)
    sessions = async_sessionmaker(engine, expire_on_commit=False)
    try:
        if stage == "external":
            if not await execute_due(sessions, 30, identifier.bytes):
                raise RuntimeError("指定外部操作未被领取")
            return
        selected = await claim(sessions, 2)
        if selected is None or selected[0] != identifier.bytes:
            raise RuntimeError("领取结果与指定验收任务不一致")
        run_id, token = selected
        if stage == "committed":
            await build_tool_graph(sessions).ainvoke({"run_id": run_id, "token": token})
        async with sessions() as session:
            row = await session.get(AgentRun, run_id)
            marker = {
                "run_id": str(identifier),
                "token": token,
                "lease_until": row.lease_until.isoformat(),
                "stage": stage,
            }
            print(json.dumps(marker), flush=True)
        # 父进程收到真实事务完成标记后终止本进程；不执行生成模型。
        await asyncio.to_thread(sys.stdin.read)
    finally:
        await engine.dispose()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="验证工作进程终止后的租约恢复")
    parser.add_argument("run_id", type=UUID)
    parser.add_argument("stage", choices=("claimed", "committed", "external"))
    arguments = parser.parse_args()
    asyncio.run(main(arguments.run_id, arguments.stage))