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