zhifu-platform/tests/process_recovery_acceptance.py

import asyncio
import json
import sys
from datetime import datetime
from pathlib import Path
from time import perf_counter
from uuid import UUID

import pytest
from sqlalchemy import func, select
from test_after_sales import scenario  # noqa: F401
from test_agent_runtime import tool_run

from zhifu.core.errors import BusinessError
from zhifu.modules.after_sales.models import Ticket
from zhifu.modules.agent_runtime.models import AgentRun, ModelAttempt
from zhifu.modules.agent_runtime.runtime import build_tool_graph
from zhifu.modules.agent_runtime.worker import claim
from zhifu.modules.identity.auth import now


@pytest.mark.parametrize("stage", ["claimed", "committed"])
async def test_process_termination_and_natural_lease_recovery(scenario, stage):  # noqa: F811
    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("恢复验收要求停止生成工作进程")
    async with sessions() as session:
        count = await session.scalar(
            select(func.count())
            .select_from(AgentRun)
            .where(AgentRun.status.in_(["QUEUED", "RUNNING"]))
        )
        assert count == 0, "需要没有其他待执行任务的独立验收环境"
    run = await tool_run(scenario, "create_exchange_proposal", scenario["proposal"])
    started = perf_counter()
    process = await asyncio.create_subprocess_exec(
        sys.executable,
        "scripts/recovery_probe.py",
        str(UUID(bytes=run.id)),
        stage,
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )
    try:
        async with asyncio.timeout(180):
            async with asyncio.timeout(35):
                raw = await process.stdout.readline()
                assert raw, (await process.stderr.read()).decode(errors="replace")
            marker = json.loads(raw)
            assert marker["run_id"] == str(UUID(bytes=run.id))
            process.kill()
            await process.wait()
            assert process.returncode != 0
            # 等待数据库中实际125秒租约自然到期，期间不能重新领取。
            assert await claim(sessions, 2) is None
            expiry = datetime.fromisoformat(marker["lease_until"])
            await asyncio.sleep(max(0, (expiry - now()).total_seconds()) + 0.05)
            recovered = await claim(sessions, 2)
            assert recovered and recovered[0] == run.id and recovered[1] > marker["token"]
            graph = build_tool_graph(sessions)
            with pytest.raises(BusinessError, match="执行权已失效"):
                await graph.ainvoke({"run_id": run.id, "token": marker["token"]})
            await graph.ainvoke({"run_id": run.id, "token": recovered[1]})
            await graph.ainvoke({"run_id": run.id, "token": recovered[1]})
            async with sessions() as session:
                tickets = await session.scalar(
                    select(func.count())
                    .select_from(Ticket)
                    .where(
                        Ticket.tenant_id == scenario["tenant"].bytes,
                    )
                )
                attempts = await session.scalar(
                    select(func.count())
                    .select_from(ModelAttempt)
                    .where(
                        ModelAttempt.tenant_id == scenario["tenant"].bytes,
                    )
                )
                assert tickets == 1 and attempts == 0
            output = Path(".work/process-recovery")
            await asyncio.to_thread(output.mkdir, parents=True, exist_ok=True)
            report = {
                "stage": stage,
                "child_exit": process.returncode,
                "seconds": perf_counter() - started,
                "tickets": tickets,
                "generation_attempts": attempts,
                "natural_lease_expiry": True,
                "old_token": marker["token"],
                "recovered_token": recovered[1],
            }
            await asyncio.to_thread(
                (output / f"{stage}.json").write_text,
                json.dumps(report, ensure_ascii=False, indent=2) + "\n",
                encoding="utf-8",
            )
    finally:
        if process.returncode is None:
            process.kill()
        await process.wait()