video-image-production-build/tests/graph/test_recovery.py

from __future__ import annotations

import hashlib
import json
import tempfile
import unittest
from pathlib import Path
from typing import Any
from unittest.mock import patch

from langchain_core.messages import AIMessage

from run_global_data import run_global_data

from production_build_agents.contracts.models import (
    ExecutorDelivery,
    TaskPackage,
    ValidationReport,
)
from production_build_agents.agents.executor.agent import run_executor_agent as real_run_executor_agent
from production_build_agents.run.records import VersionConflictError
from production_build_agents.tools.registry import (
    create_default_tool_registry as real_create_default_tool_registry,
)
from production_build_agents.contracts.identifiers import validator_run_id_for
from tests.support.executor_fixtures import (
    build_executor_candidate,
    build_executor_model,
)
from tests.support.fake_models import ToolAwareFakeChatModel
from tests.support.planner_fixtures import (
    build_global_data_plan,
    build_planned_task,
    build_planner_model,
    build_task_package,
)
from tests.support.validator_fixtures import (
    build_stage_read_call,
    build_validator_candidate,
    build_validator_model,
    build_stage_validator_candidate,
)


INPUT_PATH = (
    Path(__file__).parents[1] / "fixtures" / "minimal_input.json"
)


def _tree_hashes(root: Path) -> dict[str, str]:
    return {
        str(path.relative_to(root)): hashlib.sha256(path.read_bytes()).hexdigest()
        for path in sorted(root.rglob("*"))
        if path.is_file()
    }


def _completed_run(output_root: Path, run_id: str) -> dict[str, Any]:
    return run_global_data(
        INPUT_PATH,
        output_root=output_root,
        thread_id=run_id,
        registry_dir=output_root / ".registry",
        planner_model=build_planner_model(build_global_data_plan()),
        executor_model=build_executor_model(
            build_executor_candidate(run_id=run_id)
        ),
        validator_model=build_validator_model(
            build_validator_candidate(run_id=run_id)
        ),
    )


class DurableRunTest(unittest.TestCase):

    def test_executor_agent_resumes_after_inspect_without_losing_gate(self) -> None:
        class CountingRemoteClient:
            def __init__(self) -> None:
                self.invoke_count = 0

            def search(
                self,
                query: str,
                *,
                limit: int = 5,
            ) -> dict[str, Any]:
                return {
                    "success": True,
                    "results": [
                        {"tool_id": "fake-research", "summary": query}
                    ],
                    "total": 1,
                }

            def inspect(
                self,
                tool_ids: list[str],
            ) -> dict[str, Any]:
                return {
                    "success": True,
                    "tools": {
                        tool_id: {
                            "success": True,
                            "input_schema": {"type": "object"},
                        }
                        for tool_id in tool_ids
                    },
                }

            def invoke(
                self,
                tool_id: str,
                payload: dict[str, Any],
            ) -> dict[str, Any]:
                self.invoke_count += 1
                return {
                    "success": True,
                    "tool_id": tool_id,
                    "payload": payload,
                    "url": "https://example.test/recovered-source",
                }

        run_id = "Run-resume-agent-loop"
        planned_task = build_planned_task(
            skill_id="external-research",
            deliverable_type="research_result",
        )
        client = CountingRemoteClient()

        def registry_factory(**kwargs):
            return real_create_default_tool_registry(
                **kwargs,
                remote_client=client,
            )

        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            task = build_task_package(
                Path(__file__).parents[1]
                / "fixtures"
                / "minimal_brief.json",
                run_id=run_id,
                planned_task=planned_task,
            )
            interrupted_model = build_executor_model(
                build_executor_candidate(
                    run_id=run_id,
                    deliverable_type="research_result",
                    artifact_uri="https://example.test/recovered-source",
                )
            )
            interrupted_model.responses.insert(
                0,
                AIMessage(
                    content="",
                    tool_calls=[
                        {
                            "name": "inspect_tool",
                            "args": {"tool_ids": ["fake-research"]},
                            "id": "inspect-before-interrupt",
                            "type": "tool_call",
                        }
                    ],
                ),
            )
            original_generate = interrupted_model._generate
            model_calls = 0

            def interrupt_second_model_call(*args, **kwargs):
                nonlocal model_calls
                model_calls += 1
                if model_calls == 2:
                    raise RuntimeError("模拟 inspect 后进程中断")
                return original_generate(*args, **kwargs)

            with patch(
                "production_build_agents.agents.executor.agent.create_default_tool_registry",
                side_effect=registry_factory,
            ), patch.object(
                interrupted_model,
                "_generate",
                side_effect=interrupt_second_model_call,
            ):
                with self.assertRaisesRegex(RuntimeError, "进程中断"):
                    real_run_executor_agent(
                        task,
                        run_dir=root,
                        model=interrupted_model,
                    )

            resumed_model = build_executor_model(
                build_executor_candidate(
                    run_id=run_id,
                    deliverable_type="research_result",
                    artifact_uri="https://example.test/recovered-source",
                    evidence_tool_call_ids=["run-after-resume"],
                )
            )
            resumed_model.responses.insert(
                0,
                AIMessage(
                    content="",
                    tool_calls=[
                        {
                            "name": "run_tool",
                            "args": {
                                "tool_id": "fake-research",
                                "params": {"query": "恢复测试"},
                            },
                            "id": "run-after-resume",
                            "type": "tool_call",
                        }
                    ],
                ),
            )
            with patch(
                "production_build_agents.agents.executor.agent.create_default_tool_registry",
                side_effect=registry_factory,
            ):
                delivery = real_run_executor_agent(
                    task,
                    run_dir=root,
                    model=resumed_model,
                )

        self.assertEqual(client.invoke_count, 1)
        self.assertEqual(delivery.task_id, "Task1")
        self.assertEqual(
            [record.tool_name for record in delivery.tool_calls],
            ["inspect_tool", "run_tool"],
        )

    def test_restart_after_executor_resumes_at_validator(self) -> None:
        run_id = "Run-resume-validator"
        with tempfile.TemporaryDirectory() as temp_dir:
            output_root = Path(temp_dir)
            with patch(
                "production_build_agents.global_data.nodes.validate_task."
                "run_validator_agent",
                side_effect=RuntimeError("模拟 Validator 前进程中断"),
            ):
                with self.assertRaisesRegex(RuntimeError, "进程中断"):
                    run_global_data(
                        INPUT_PATH,
                        output_root=output_root,
                        thread_id=run_id,
                        registry_dir=output_root / ".registry",
                        planner_model=build_planner_model(
                            build_global_data_plan()
                        ),
                        executor_model=build_executor_model(
                            build_executor_candidate(run_id=run_id)
                        ),
                    )

            run_dir = output_root / run_id
            delivery_path = (
                run_dir / "executor_results" / "Task1.v1.json"
            )
            original_delivery = delivery_path.read_bytes()
            result = run_global_data(
                INPUT_PATH,
                output_root=output_root,
                thread_id=run_id,
                registry_dir=output_root / ".registry",
                validator_model=build_validator_model(
                    build_validator_candidate(run_id=run_id)
                ),
            )

            self.assertEqual(result["status"], "COMPLETED")
            self.assertEqual(result["executor_calls"], 1)
            self.assertEqual(result["validator_calls"], 1)
            self.assertEqual(delivery_path.read_bytes(), original_delivery)

    def test_restart_at_stage_validator_does_not_repeat_task_work(self) -> None:
        run_id = "Run-resume-stage-validator"
        with tempfile.TemporaryDirectory() as temp_dir:
            output_root = Path(temp_dir)
            with patch(
                "production_build_agents.global_data.nodes.validate_stage."
                "run_global_data_stage_validator",
                side_effect=RuntimeError("模拟 Stage Validator 前进程中断"),
            ):
                with self.assertRaisesRegex(RuntimeError, "进程中断"):
                    run_global_data(
                        INPUT_PATH,
                        output_root=output_root,
                        thread_id=run_id,
                        registry_dir=output_root / ".registry",
                        planner_model=build_planner_model(
                            build_global_data_plan()
                        ),
                        executor_model=build_executor_model(
                            build_executor_candidate(run_id=run_id)
                        ),
                        validator_model=build_validator_model(
                            build_validator_candidate(run_id=run_id),
                            include_stage=False,
                        ),
                    )
            task_delivery = (
                output_root
                / run_id
                / "executor_results"
                / "Task1.v1.json"
            )
            original = task_delivery.read_bytes()
            stage_model = ToolAwareFakeChatModel(
                responses=[
                    build_stage_read_call("stage-resume-read"),
                    AIMessage(
                        content=build_stage_validator_candidate(
                            run_id=run_id,
                        ).model_dump_json()
                    )
                ]
            )
            result = run_global_data(
                INPUT_PATH,
                output_root=output_root,
                thread_id=run_id,
                registry_dir=output_root / ".registry",
                validator_model=stage_model,
            )
            resumed = task_delivery.read_bytes()

        self.assertEqual(result["status"], "COMPLETED")
        self.assertEqual(result["executor_calls"], 1)
        self.assertEqual(result["validator_calls"], 1)
        self.assertEqual(result["stage_validator_calls"], 1)
        self.assertEqual(resumed, original)

    def test_restart_before_task2_does_not_repeat_passed_task1(self) -> None:
        run_id = "Run-resume-task2"
        plan = build_global_data_plan(
            tasks=[
                build_planned_task(task_id="Task1"),
                build_planned_task(
                    task_id="Task2",
                    objective="基于第一项资料完成第二项",
                    depends_on=["Task1"],
                ),
            ]
        )

        def interrupt_task2(task, **kwargs):
            if task.task_id == "Task2":
                raise RuntimeError("模拟 Task2 前进程中断")
            return real_run_executor_agent(task, **kwargs)

        with tempfile.TemporaryDirectory() as temp_dir:
            output_root = Path(temp_dir)
            with patch(
                "production_build_agents.global_data.nodes.execute_task."
                "run_executor_agent",
                side_effect=interrupt_task2,
            ):
                with self.assertRaisesRegex(RuntimeError, "Task2"):
                    run_global_data(
                        INPUT_PATH,
                        output_root=output_root,
                        thread_id=run_id,
                        registry_dir=output_root / ".registry",
                        planner_model=build_planner_model(plan),
                        executor_model=build_executor_model(
                            build_executor_candidate(
                                run_id=run_id,
                                task_id="Task1",
                            )
                        ),
                        validator_model=build_validator_model(
                            build_validator_candidate(
                                run_id=run_id,
                                task_id="Task1",
                            )
                        ),
                    )

            task1_delivery = (
                output_root
                / run_id
                / "executor_results"
                / "Task1.v1.json"
            )
            original = task1_delivery.read_bytes()
            result = run_global_data(
                INPUT_PATH,
                output_root=output_root,
                thread_id=run_id,
                registry_dir=output_root / ".registry",
                executor_model=build_executor_model(
                    build_executor_candidate(
                        run_id=run_id,
                        task_id="Task2",
                    )
                ),
                validator_model=build_validator_model(
                    build_validator_candidate(
                        run_id=run_id,
                        task_id="Task2",
                    )
                ),
            )

            self.assertEqual(result["status"], "COMPLETED")
            self.assertEqual(result["executor_calls"], 2)
            self.assertEqual(result["validator_calls"], 2)
            self.assertEqual(
                result["task_records"]["Task1"]["accepted_plan_version"],
                1,
            )
            self.assertEqual(task1_delivery.read_bytes(), original)

    def test_completed_run_is_noop_and_delivery_is_byte_stable(self) -> None:
        run_id = "Run-completed-noop"
        with tempfile.TemporaryDirectory() as temp_dir:
            output_root = Path(temp_dir)
            first = _completed_run(output_root, run_id)
            delivery_path = Path(first["global_data_delivery_path"])
            original = delivery_path.read_bytes()
            before = _tree_hashes(output_root)

            second = run_global_data(
                INPUT_PATH,
                output_root=output_root,
                thread_id=run_id,
                registry_dir=output_root / ".registry",
            )

            self.assertEqual(second["status"], "COMPLETED")
            self.assertEqual(second["phase"], "FINALIZE")
            self.assertIsNone(second["replan_scope"])
            for field in (
                "planner_calls",
                "executor_calls",
                "validator_calls",
                "stage_validator_calls",
                "replan_count",
                "event_log",
            ):
                self.assertEqual(second[field], first[field])
            self.assertEqual(delivery_path.read_bytes(), original)
            self.assertEqual(_tree_hashes(output_root), before)

    def test_completed_noop_rejects_each_tampered_business_record(
        self,
    ) -> None:
        cases = (
            (
                "FinalDelivery",
                lambda result, run_dir: Path(
                    result["global_data_delivery_path"]
                ),
                "plan_id",
                "TamperedPlan",
            ),
            (
                "Plan",
                lambda result, run_dir: Path(
                    result["global_data_plan_path"]
                ),
                "plan_id",
                "TamperedPlan",
            ),
            (
                "Task",
                lambda result, run_dir: run_dir / "tasks" / "Task1.v1.json",
                "plan_id",
                "TamperedPlan",
            ),
            (
                "ExecutorDelivery",
                lambda result, run_dir: Path(
                    result["task_records"]["Task1"][
                        "executor_delivery_path"
                    ]
                ),
                "plan_id",
                "TamperedPlan",
            ),
            (
                "TaskReport",
                lambda result, run_dir: Path(
                    result["task_records"]["Task1"][
                        "validation_report_path"
                    ]
                ),
                "verdict",
                "FAIL",
            ),
        )
        for index, (name, select_path, field, value) in enumerate(cases):
            with (
                self.subTest(record=name),
                tempfile.TemporaryDirectory() as temp_dir,
            ):
                output_root = Path(temp_dir)
                run_id = f"Run-noop-tamper-{index}"
                result = _completed_run(output_root, run_id)
                path = select_path(result, output_root / run_id)
                payload = json.loads(path.read_text(encoding="utf-8"))
                payload[field] = value
                path.write_text(
                    json.dumps(payload, ensure_ascii=False),
                    encoding="utf-8",
                )
                with self.assertRaises(ValueError):
                    run_global_data(
                        INPUT_PATH,
                        output_root=output_root,
                        thread_id=run_id,
                        registry_dir=output_root / ".registry",
                    )

    def test_completed_noop_rejects_tampered_artifact_content(self) -> None:
        run_id = "Run-noop-artifact-tamper"
        with tempfile.TemporaryDirectory() as temp_dir:
            output_root = Path(temp_dir)
            result = _completed_run(output_root, run_id)
            delivery = ExecutorDelivery.model_validate_json(
                Path(
                    result["task_records"]["Task1"][
                        "executor_delivery_path"
                    ]
                ).read_text(encoding="utf-8")
            )
            artifact_path = Path(delivery.artifacts[0].uri)
            artifact_path.write_text(
                '{"tampered": true}\n',
                encoding="utf-8",
            )
            with self.assertRaisesRegex(
                ValueError,
                "Artifact 内容已变化",
            ):
                run_global_data(
                    INPUT_PATH,
                    output_root=output_root,
                    thread_id=run_id,
                    registry_dir=output_root / ".registry",
                )

    def test_resume_rejects_input_drift(self) -> None:
        run_id = "Run-input-conflict"
        with tempfile.TemporaryDirectory() as temp_dir:
            output_root = Path(temp_dir)
            run_global_data(
                INPUT_PATH,
                output_root=output_root,
                thread_id=run_id,
                registry_dir=output_root / ".registry",
                planner_model=build_planner_model(
                    build_global_data_plan()
                ),
                executor_model=build_executor_model(
                    build_executor_candidate(run_id=run_id)
                ),
                validator_model=build_validator_model(
                    build_validator_candidate(run_id=run_id)
                ),
            )
            copied_input = output_root / "same-content-new-path.json"
            copied_input.write_bytes(INPUT_PATH.read_bytes())

            with self.assertRaisesRegex(
                VersionConflictError,
                "input_path",
            ):
                run_global_data(
                    copied_input,
                    output_root=output_root,
                    thread_id=run_id,
                    registry_dir=output_root / ".registry",
                )

    def test_thread_registry_rejects_output_directory_drift(self) -> None:
        run_id = "Run-output-conflict"
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            registry = root / ".registry"
            run_global_data(
                INPUT_PATH,
                output_root=root / "output-a",
                thread_id=run_id,
                registry_dir=registry,
                planner_model=build_planner_model(
                    build_global_data_plan()
                ),
                executor_model=build_executor_model(
                    build_executor_candidate(run_id=run_id)
                ),
                validator_model=build_validator_model(
                    build_validator_candidate(run_id=run_id)
                ),
            )
            with self.assertRaisesRegex(
                VersionConflictError,
                "run_dir",
            ):
                run_global_data(
                    INPUT_PATH,
                    output_root=root / "output-b",
                    thread_id=run_id,
                    registry_dir=registry,
                )

    def test_thread_registry_rejects_same_path_content_drift(self) -> None:
        run_id = "Run-content-conflict"
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            source = root / "input.json"
            source.write_bytes(INPUT_PATH.read_bytes())
            run_global_data(
                source,
                output_root=root / "output",
                thread_id=run_id,
                registry_dir=root / ".registry",
                planner_model=build_planner_model(
                    build_global_data_plan()
                ),
                executor_model=build_executor_model(
                    build_executor_candidate(run_id=run_id)
                ),
                validator_model=build_validator_model(
                    build_validator_candidate(run_id=run_id)
                ),
            )
            source.write_text(
                source.read_text(encoding="utf-8") + "\n",
                encoding="utf-8",
            )
            with self.assertRaisesRegex(
                VersionConflictError,
                "input_sha256",
            ):
                run_global_data(
                    source,
                    output_root=root / "output",
                    thread_id=run_id,
                    registry_dir=root / ".registry",
                )

    def test_inconsistent_saved_report_is_rejected_on_resume(self) -> None:
        run_id = "Run-invalid-saved-report"
        with tempfile.TemporaryDirectory() as temp_dir:
            output_root = Path(temp_dir)
            registry = output_root / ".registry"
            with patch(
                "production_build_agents.global_data.nodes.validate_task."
                "run_validator_agent",
                side_effect=RuntimeError("模拟 Validator 前中断"),
            ):
                with self.assertRaises(RuntimeError):
                    run_global_data(
                        INPUT_PATH,
                        output_root=output_root,
                        thread_id=run_id,
                        registry_dir=registry,
                        planner_model=build_planner_model(
                            build_global_data_plan()
                        ),
                        executor_model=build_executor_model(
                            build_executor_candidate(run_id=run_id)
                        ),
                    )

            run_dir = output_root / run_id
            task = TaskPackage.model_validate_json(
                (run_dir / "tasks" / "Task1.v1.json").read_text(
                    encoding="utf-8"
                )
            )
            delivery = ExecutorDelivery.model_validate_json(
                (
                    run_dir / "executor_results" / "Task1.v1.json"
                ).read_text(encoding="utf-8")
            )
            failed = build_validator_candidate(
                run_id=run_id,
                verdict="FAIL",
            )
            inconsistent = ValidationReport(
                run_id=run_id,
                plan_id=task.plan_id,
                plan_version=1,
                task_id="Task1",
                executor_run_id=delivery.executor_run_id,
                validator_run_id=validator_run_id_for(delivery),
                criterion_results=failed.criterion_results,
                verdict="PASS",
                summary="故意构造的自相矛盾报告",
            )
            report_path = (
                run_dir / "validation_results" / "Task1.v1.json"
            )
            report_path.parent.mkdir(parents=True, exist_ok=True)
            report_path.write_text(
                inconsistent.model_dump_json(),
                encoding="utf-8",
            )

            result = run_global_data(
                INPUT_PATH,
                output_root=output_root,
                thread_id=run_id,
                registry_dir=registry,
            )
            self.assertEqual(result["status"], "FAILED")
            self.assertEqual(result["phase"], "VALIDATE_TASK")
            self.assertEqual(
                result["failure_code"],
                "VALIDATOR_OUTPUT_INVALID",
            )
            self.assertNotEqual(
                result["task_records"]["Task1"]["status"],
                "passed",
            )


if __name__ == "__main__":
    unittest.main()