video-image-production-build/tests/run/test_persistence.py

from __future__ import annotations

import json
import tempfile
import unittest
from pathlib import Path
from typing import TypedDict, get_type_hints

from langchain_core.messages import AIMessage, ToolMessage
from langgraph.graph import END, START, StateGraph

from production_build_agents.agents.invocation import recovered_tool_progress
from production_build_agents.run.langgraph_checkpointer import (
    create_sqlite_checkpointer,
)
from production_build_agents.run.lock import RunLockError, acquire_run_lock
from production_build_agents.run.operation_journal import (
    OperationJournal,
    OperationOutcomeUnknownError,
    UncertainSideEffectError,
)
from production_build_agents.run.records import (
    VersionConflictError,
    save_executor_candidate,
    save_run_summary,
    write_once_json,
    write_once_text,
)


class RunPersistenceTest(unittest.TestCase):
    def test_executor_candidate_writer_has_resolvable_annotations(self) -> None:
        hints = get_type_hints(save_executor_candidate)
        self.assertEqual(hints["candidate"].__name__, "ExecutorCandidate")

    def test_run_summary_is_written_once_with_merged_event_log(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            state = {
                "run_id": "Run-summary",
                "status": "RUNNING",
                "event_log": ["started"],
            }
            update = {
                "status": "COMPLETED",
                "phase": "FINALIZE",
                "event_log": ["completed"],
            }

            path = save_run_summary(root, state, update)
            self.assertEqual(path, save_run_summary(root, state, update))
            saved = json.loads(path.read_text(encoding="utf-8"))
            self.assertEqual(saved["status"], "COMPLETED")
            self.assertEqual(saved["event_log"], ["started", "completed"])
            with self.assertRaises(VersionConflictError):
                save_run_summary(
                    root,
                    state,
                    {**update, "status": "FAILED"},
                )

    def test_recovered_progress_counts_known_failed_side_effect(self) -> None:
        messages = [
            AIMessage(
                content="",
                tool_calls=[
                    {
                        "name": "run_tool",
                        "args": {
                            "tool_id": "fake",
                            "params": {"value": 1},
                        },
                        "id": "call-1",
                        "type": "tool_call",
                    }
                ],
            ),
            ToolMessage(
                content=json.dumps(
                    {
                        "success": False,
                        "_operation_id": "executor:1",
                    }
                ),
                tool_call_id="call-1",
                name="run_tool",
            ),
        ]
        inspected, completed = recovered_tool_progress(
            messages,
            side_effect_tools={"run_tool"},
        )
        self.assertEqual(inspected, set())
        self.assertEqual(completed, 1)

    def test_write_once_reuses_identical_content_and_rejects_conflict(
        self,
    ) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            json_path = root / "record.json"
            text_path = root / "record.txt"

            self.assertEqual(
                write_once_json(json_path, {"value": 1}),
                write_once_json(json_path, {"value": 1}),
            )
            self.assertEqual(
                write_once_text(text_path, "stable\n"),
                write_once_text(text_path, "stable\n"),
            )
            with self.assertRaises(VersionConflictError):
                write_once_json(json_path, {"value": 2})
            with self.assertRaises(VersionConflictError):
                write_once_text(text_path, "changed\n")

    def test_operation_journal_replays_success(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            calls = 0

            def operation() -> dict[str, object]:
                nonlocal calls
                calls += 1
                return {"success": True, "artifact": "local://one"}

            first = OperationJournal(root, "executor-Task1-v1")
            result = first.execute(
                tool_name="run_tool",
                params={"prompt": "stable"},
                operation=operation,
            )
            restarted = OperationJournal(root, "executor-Task1-v1")
            replay = restarted.execute(
                tool_name="run_tool",
                params={"prompt": "stable"},
                operation=operation,
            )

            self.assertTrue(result["success"])
            self.assertTrue(replay["_operation_replayed"])
            self.assertEqual(
                replay["_operation_id"],
                result["_operation_id"],
            )
            self.assertEqual(calls, 1)

    def test_operation_journal_rejects_parameter_drift(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            OperationJournal(root, "executor-Task1-v1").execute(
                tool_name="run_tool",
                params={"prompt": "first"},
                operation=lambda: {"success": True},
            )
            with self.assertRaises(VersionConflictError):
                OperationJournal(root, "executor-Task1-v1").execute(
                    tool_name="run_tool",
                    params={"prompt": "changed"},
                    operation=lambda: {"success": True},
                )

    def test_operation_journal_replays_known_failure(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            first = OperationJournal(root, "executor-Task1-v1")
            result = first.execute(
                tool_name="crop_image",
                params={"image": "missing.png"},
                operation=lambda: (_ for _ in ()).throw(
                    RuntimeError("known failure")
                ),
            )
            self.assertFalse(result["success"])
            self.assertIn("known failure", result["error"])

            calls = 0

            def should_not_run() -> dict[str, object]:
                nonlocal calls
                calls += 1
                return {"success": True}

            replay = OperationJournal(
                root,
                "executor-Task1-v1",
            ).execute(
                tool_name="crop_image",
                params={"image": "missing.png"},
                operation=should_not_run,
            )
            self.assertFalse(replay["success"])
            self.assertTrue(replay["_operation_replayed"])
            self.assertEqual(
                replay["_operation_id"],
                result["_operation_id"],
            )
            self.assertEqual(calls, 0)

    def test_operation_journal_never_replays_unknown_outcome(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            first = OperationJournal(root, "executor-Task1-v1")
            with self.assertRaises(OperationOutcomeUnknownError):
                first.execute(
                    tool_name="run_tool",
                    params={"prompt": "test"},
                    operation=lambda: {
                        "success": False,
                        "outcome": "OUTCOME_UNKNOWN",
                    },
                )

            calls = 0

            def should_not_run() -> dict[str, object]:
                nonlocal calls
                calls += 1
                return {"success": True}

            with self.assertRaises(OperationOutcomeUnknownError):
                OperationJournal(root, "executor-Task1-v1").execute(
                    tool_name="run_tool",
                    params={"prompt": "test"},
                    operation=should_not_run,
                )
            self.assertEqual(calls, 0)
            entries = list((root / "tool_operations").glob("*.json"))
            self.assertEqual(len(entries), 1)
            self.assertEqual(
                json.loads(entries[0].read_text(encoding="utf-8"))["status"],
                "OUTCOME_UNKNOWN",
            )

    def test_uncertain_side_effect_exception_is_never_retried(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            with self.assertRaises(OperationOutcomeUnknownError):
                OperationJournal(root, "executor-Task1-v1").execute(
                    tool_name="video_concat",
                    params={"videos": ["a.mp4", "b.mp4"]},
                    operation=lambda: (_ for _ in ()).throw(
                        UncertainSideEffectError("upload response lost")
                    ),
                )
            with self.assertRaises(OperationOutcomeUnknownError):
                OperationJournal(root, "executor-Task1-v1").execute(
                    tool_name="video_concat",
                    params={"videos": ["a.mp4", "b.mp4"]},
                    operation=lambda: {"success": True},
                )

    def test_sqlite_checkpoint_survives_reopen(self) -> None:
        class CounterState(TypedDict):
            value: int

        builder = StateGraph(CounterState)
        builder.add_node(
            "increment",
            lambda state: {"value": state["value"] + 1},
        )
        builder.add_edge(START, "increment")
        builder.add_edge("increment", END)
        config = {"configurable": {"thread_id": "durable-thread"}}

        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            with create_sqlite_checkpointer(root) as checkpointer:
                graph = builder.compile(checkpointer=checkpointer)
                self.assertEqual(
                    graph.invoke({"value": 0}, config)["value"],
                    1,
                )
            with create_sqlite_checkpointer(root) as checkpointer:
                graph = builder.compile(checkpointer=checkpointer)
                self.assertEqual(graph.get_state(config).values["value"], 1)

    def test_run_lock_rejects_second_owner(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            root = Path(temp_dir)
            with acquire_run_lock(root):
                with self.assertRaises(RunLockError):
                    with acquire_run_lock(root):
                        self.fail("不应取得第二把锁")


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