import asyncio import importlib import json import os import tempfile import unittest from datetime import datetime from pathlib import Path from unittest.mock import patch from fastapi import HTTPException from cyber_agent.core.memory import MemoryConfig, compute_memory_identity from cyber_agent.core.run_snapshot import ( RunConfigSnapshotError, RunConfigSnapshotV1, load_run_config_snapshot, persist_run_config_snapshot, ) from cyber_agent.core.runner import AgentRunner, RunConfig from cyber_agent.tools.approval import ( ToolApprovalBatchV1, ToolApprovalCallV1, tool_argument_hash, ) from cyber_agent.tools.builtin.knowledge import KnowledgeConfig from cyber_agent.tools.registry import ToolRegistry from cyber_agent.trace import run_api from cyber_agent.trace.goal_models import GoalTree from cyber_agent.trace.models import Message, Trace from cyber_agent.trace.run_api import ( DecideToolApprovalRequest, TraceRunRequest, ToolApprovalDecision, decide_tool_approval, _restore_runner_and_config, reconcile_traces, run_trace, set_runner, stop_trace, ) from cyber_agent.trace.store import ( FileSystemTraceStore, TraceStoreCorruptionError, ) def quiet_knowledge() -> KnowledgeConfig: return KnowledgeConfig( enable_extraction=False, enable_completion_extraction=False, enable_injection=False, ) class RunSnapshotHardeningTest(unittest.IsolatedAsyncioTestCase): async def test_v1_snapshot_shape_and_hash_remain_byte_compatible(self): config = RunConfig( model="v1-fixture", temperature=0.25, max_iterations=3, tools=[], tool_groups=[], exclude_tools=[], auto_execute_tools=True, agent_type="default", uid="fixture", skills=[], enable_memory=False, knowledge=quiet_knowledge(), parallel_tool_execution=False, child_execution_mode="sequential", max_parallel_children=1, side_branch_max_turns=2, goal_compression="none", enable_prompt_caching=True, enable_research_flow=False, extra_llm_params={}, context={}, ) snapshot = RunConfigSnapshotV1.from_run_config( config, memory_identity=None, ) self.assertEqual(1, snapshot.schema_version) self.assertNotIn("application_ref", snapshot.model_dump()) self.assertEqual( "1fdd11b90c121e13fa30aa5d246c20a50b602baa6b3a1aeed186b2c1ea72745b", snapshot.snapshot_hash, ) async def test_full_snapshot_round_trip_and_hash_tamper_detection(self): with tempfile.TemporaryDirectory() as temp_dir: memory = MemoryConfig( base_path=str(Path(temp_dir) / "memory"), files={"identity.md": "identity", "journals/**": "journal"}, dream_prompt="dream", reflect_prompt="reflect", ) config = RunConfig( model="model-x", temperature=0.17, max_iterations=37, tools=["alpha", "beta"], tool_groups=["ignored-when-tools-explicit"], exclude_tools=["beta"], auto_execute_tools=False, agent_type="writer", uid="user-7", skills=["planning"], enable_memory=True, memory=memory, knowledge=KnowledgeConfig( enable_extraction=False, default_scopes=["project:test"], ), parallel_tool_execution=True, child_execution_mode="parallel", max_parallel_children=4, side_branch_max_turns=7, goal_compression="on_complete", enable_prompt_caching=False, enable_research_flow=False, extra_llm_params={"top_p": 0.8}, context={"project_name": "demo", "business": {"lane": 2}}, ) identity = compute_memory_identity(memory) snapshot = RunConfigSnapshotV1.from_run_config( config, memory_identity=identity, system_prompt_hash="a" * 64, ) context = {} persist_run_config_snapshot(context, snapshot) loaded = load_run_config_snapshot(context) restored = RunConfig() restored.apply_snapshot(loaded) self.assertEqual("model-x", restored.model) self.assertEqual(["alpha", "beta"], restored.tools) self.assertEqual(["ignored-when-tools-explicit"], restored.tool_groups) self.assertEqual("parallel", restored.child_execution_mode) self.assertEqual(4, restored.max_parallel_children) self.assertEqual(7, restored.side_branch_max_turns) self.assertEqual(["project:test"], restored.knowledge.default_scopes) self.assertEqual(memory.files, restored.memory.files) self.assertEqual(identity, loaded.memory_identity) self.assertEqual("demo", loaded.project_name) context["run_config_snapshot"]["model"] = "tampered" with self.assertRaisesRegex(RunConfigSnapshotError, "hash mismatch"): load_run_config_snapshot(context) async def test_precreated_trace_binds_actual_system_prompt_hash_on_first_run(self): with tempfile.TemporaryDirectory() as temp_dir: store = FileSystemTraceStore(temp_dir) config = RunConfig( trace_id="precreated-child", model="fake", tools=[], tool_groups=[], system_prompt="fixed child prompt", knowledge=quiet_knowledge(), ) context = {"agent_mode": "legacy", "agent_mode_revision": 1} persist_run_config_snapshot( context, RunConfigSnapshotV1.from_run_config( config, memory_identity=None, ), ) await store.create_trace(Trace( trace_id=config.trace_id, mode="agent", task="child", model="fake", context=context, )) runner = AgentRunner(trace_store=store) await runner._build_history( config.trace_id, [], GoalTree(mission="child"), config, 1, ) persisted = await store.get_trace(config.trace_id) snapshot = load_run_config_snapshot(persisted.context) self.assertIsNotNone(snapshot.system_prompt_hash) self.assertEqual(64, len(snapshot.system_prompt_hash)) async def test_old_recursive_refuses_resume_but_legacy_is_inferred_once(self): with tempfile.TemporaryDirectory() as temp_dir: store = FileSystemTraceStore(temp_dir) recursive = Trace( trace_id="old-recursive", mode="agent", task="old", context={ "agent_mode": "recursive", "agent_mode_revision": 2, }, ) legacy = Trace( trace_id="old-legacy", mode="agent", task="old", model="legacy-model", context={"agent_mode": "legacy", "agent_mode_revision": 1}, ) await store.create_trace(recursive) await store.create_trace(legacy) runner = AgentRunner(trace_store=store, llm_call=lambda **_: None) with self.assertRaisesRegex(ValueError, "recursive_revision_2_read_only"): await runner._prepare_existing_trace( RunConfig(trace_id=recursive.trace_id) ) _trace, _goal, _sequence = await runner._prepare_existing_trace( RunConfig(trace_id=legacy.trace_id) ) persisted = await store.get_trace(legacy.trace_id) self.assertTrue(load_run_config_snapshot(persisted.context).legacy_inferred) async def test_old_legacy_project_snapshot_is_inferred_before_runner_selection(self): with tempfile.TemporaryDirectory() as temp_dir: store = FileSystemTraceStore(temp_dir) trace = Trace( trace_id="old-project-legacy", mode="agent", task="old project", model="legacy-project-model", context={ "agent_mode": "legacy", "agent_mode_revision": 1, "project_name": "legacy_demo", }, ) await store.create_trace(trace) base_runner = AgentRunner(trace_store=store) project_runner = AgentRunner(trace_store=store) async def init_project_env(): return project_runner, [], RunConfig() module = type("ProjectModule", (), { "init_project_env": staticmethod(init_project_env), }) real_import_module = importlib.import_module def import_project(name, *args, **kwargs): if name == "examples.legacy_demo.run": return module return real_import_module(name, *args, **kwargs) with patch("importlib.import_module", side_effect=import_project): restored_runner, config = await _restore_runner_and_config( trace=trace, base_runner=base_runner, ) self.assertIs(project_runner, restored_runner) self.assertEqual("legacy-project-model", config.model) persisted = await store.get_trace(trace.trace_id) snapshot = load_run_config_snapshot(persisted.context) self.assertTrue(snapshot.legacy_inferred) self.assertEqual("legacy_demo", snapshot.project_name) class ToolPolicyHardeningTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.registry = ToolRegistry() self.executed = [] async def exact_one(value: int = 0): self.executed.append(("exact_one", value)) return "one" async def grouped_two(value: int = 0): self.executed.append(("grouped_two", value)) return "two" self.registry.register(exact_one, groups=["shared"]) self.registry.register(grouped_two, groups=["shared"]) async def test_explicit_tools_are_exact_for_schema_and_dispatch(self): runner = AgentRunner(tool_registry=self.registry) names = runner._get_configured_tool_names( tools=["exact_one"], tool_groups=["shared"], exclude_tools=[], ) schemas = runner._get_tool_schemas( tools=["exact_one"], tool_groups=["shared"], exclude_tools=[], ) self.assertEqual({"exact_one"}, names) self.assertEqual( {"exact_one"}, {schema["function"]["name"] for schema in schemas}, ) denied = await self.registry.execute( "grouped_two", {}, allowed_tool_names=names, ) self.assertIn("not allowed", denied) self.assertEqual([], self.executed) async def test_url_and_confirmation_are_hard_execution_gates(self): async def guarded(value: int = 0): self.executed.append(("guarded", value)) return "ok" self.registry.register( guarded, requires_confirmation=True, editable_params=["value"], url_patterns=["https://allowed.example.com/**"], ) call_id = "call-1" arguments = {"value": 4} digest = tool_argument_hash( tool_call_id=call_id, tool_name="guarded", arguments=arguments, ) grant = { "batch_id": "batch-1", "trace_id": "trace-1", "tool_call_id": call_id, "tool_name": "guarded", "argument_hash": digest, "decision": "approved", } missing_url = await self.registry.execute( "guarded", arguments, tool_call_id=call_id, approval_grant=grant, context={"trace_id": "trace-1"}, ) wrong_url = await self.registry.execute( "guarded", arguments, tool_call_id=call_id, approval_grant=grant, context={"trace_id": "trace-1", "page_url": "https://evil.example/x"}, ) missing_grant = await self.registry.execute( "guarded", arguments, tool_call_id=call_id, context={ "trace_id": "trace-1", "page_url": "https://allowed.example.com/path", }, ) self.assertIn("URL policy denied", missing_url) self.assertIn("URL policy denied", wrong_url) self.assertIn("confirmation required", missing_grant) self.assertEqual([], self.executed) accepted = await self.registry.execute( "guarded", arguments, tool_call_id=call_id, approval_grant=grant, context={ "trace_id": "trace-1", "page_url": "https://allowed.example.com/path", }, ) self.assertEqual("ok", accepted) self.assertEqual([("guarded", 4)], self.executed) class ToolApprovalStateMachineTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.temp = tempfile.TemporaryDirectory() self.store = FileSystemTraceStore(self.temp.name) self.registry = ToolRegistry() self.side_effects = [] run_api._running_tasks.clear() run_api._running_runners.clear() run_api._approval_locks.clear() async def dangerous_tool(value: int, fixed: str = "fixed"): self.side_effects.append(("dangerous", value, fixed)) return {"value": value} async def safe_tool(label: str): self.side_effects.append(("safe", label)) return "safe" self.registry.register( dangerous_tool, requires_confirmation=True, editable_params=["value"], groups=["tests"], ) self.registry.register(safe_tool, groups=["tests"]) self.model_calls = 0 async def llm_call(**kwargs): self.model_calls += 1 if self.model_calls == 1: return { "content": "", "tool_calls": [ { "id": "danger-1", "type": "function", "function": { "name": "dangerous_tool", "arguments": json.dumps({"value": 1, "fixed": "fixed"}), }, }, { "id": "safe-1", "type": "function", "function": { "name": "safe_tool", "arguments": json.dumps({"label": "after-approval"}), }, }, ], "finish_reason": "tool_calls", } return { "content": "done", "tool_calls": [], "finish_reason": "stop", } self.runner = AgentRunner( trace_store=self.store, tool_registry=self.registry, llm_call=llm_call, ) set_runner(self.runner) async def asyncTearDown(self): tasks = list(run_api._running_tasks.values()) for task in tasks: if not task.done(): task.cancel() if tasks: await asyncio.gather(*tasks, return_exceptions=True) run_api._running_tasks.clear() run_api._running_runners.clear() run_api._approval_locks.clear() self.temp.cleanup() async def start_waiting_batch(self): with patch.dict(os.environ, {"AGENT_MODE": "legacy"}): result = await self.runner.run_result( [{"role": "user", "content": "run tools"}], RunConfig( model="fake", tools=["dangerous_tool", "safe_tool"], tool_groups=[], enable_research_flow=False, knowledge=quiet_knowledge(), ), ) self.assertEqual("waiting_confirmation", result["status"]) self.assertEqual([], self.side_effects) batch = await self.store.get_tool_approval_batch(result["trace_id"]) self.assertEqual("pending", batch.status) self.assertEqual(1, len(batch.pending_calls)) return result["trace_id"], batch async def test_mixed_batch_waits_then_executes_approved_edit_and_auto_call(self): trace_id, batch = await self.start_waiting_batch() response = await decide_tool_approval( trace_id, batch.batch_id, DecideToolApprovalRequest(decisions=[ToolApprovalDecision( tool_call_id="danger-1", decision="approve", edited_arguments={"value": 9, "fixed": "fixed"}, )]), ) self.assertEqual("started", response.status) await run_api._running_tasks[trace_id] self.assertEqual( [("dangerous", 9, "fixed"), ("safe", "after-approval")], self.side_effects, ) persisted = await self.store.get_tool_approval_batch(trace_id) self.assertEqual("completed", persisted.status) self.assertTrue(all( call.execution_status == "executed" for call in persisted.calls )) self.assertEqual("completed", (await self.store.get_trace(trace_id)).status) async def test_auto_execute_false_requires_confirmation_for_safe_tools_too(self): registry = ToolRegistry() effects = [] async def ordinary_tool(value: int): effects.append(value) return "ok" registry.register(ordinary_tool) async def llm_call(**_kwargs): return { "content": "", "tool_calls": [{ "id": "ordinary-1", "type": "function", "function": { "name": "ordinary_tool", "arguments": json.dumps({"value": 3}), }, }], "finish_reason": "tool_calls", } runner = AgentRunner( trace_store=self.store, tool_registry=registry, llm_call=llm_call, ) result = await runner.run_result( [{"role": "user", "content": "confirm everything"}], RunConfig( model="fake", tools=["ordinary_tool"], tool_groups=[], auto_execute_tools=False, enable_research_flow=False, knowledge=quiet_knowledge(), ), ) self.assertEqual("waiting_confirmation", result["status"]) self.assertEqual([], effects) batch = await self.store.get_tool_approval_batch(result["trace_id"]) self.assertEqual(["pending"], [call.decision for call in batch.calls]) async def test_reject_is_persisted_and_repeat_decision_is_conflict(self): trace_id, batch = await self.start_waiting_batch() await decide_tool_approval( trace_id, batch.batch_id, DecideToolApprovalRequest(decisions=[ToolApprovalDecision( tool_call_id="danger-1", decision="reject", )]), ) await run_api._running_tasks[trace_id] self.assertEqual([("safe", "after-approval")], self.side_effects) persisted = await self.store.get_tool_approval_batch(trace_id) dangerous = next(c for c in persisted.calls if c.tool_call_id == "danger-1") self.assertEqual("rejected", dangerous.execution_status) with self.assertRaises(HTTPException) as raised: await decide_tool_approval( trace_id, batch.batch_id, DecideToolApprovalRequest(decisions=[ToolApprovalDecision( tool_call_id="danger-1", decision="approve", )]), ) self.assertEqual(409, raised.exception.status_code) async def test_editing_non_editable_argument_is_rejected_without_mutation(self): trace_id, batch = await self.start_waiting_batch() with self.assertRaises(HTTPException) as raised: await decide_tool_approval( trace_id, batch.batch_id, DecideToolApprovalRequest(decisions=[ToolApprovalDecision( tool_call_id="danger-1", decision="approve", edited_arguments={"value": 1, "fixed": "tampered"}, )]), ) self.assertEqual(422, raised.exception.status_code) unchanged = await self.store.get_tool_approval_batch(trace_id) self.assertEqual("pending", unchanged.status) self.assertEqual("pending", unchanged.pending_calls[0].decision) self.assertEqual([], self.side_effects) async def test_waiting_trace_rejects_plain_run_and_stop_cancels_batch(self): trace_id, batch = await self.start_waiting_batch() with self.assertRaises(HTTPException) as raised: await run_trace(trace_id, TraceRunRequest(messages=[])) self.assertEqual(409, raised.exception.status_code) stopped = await stop_trace(trace_id) self.assertEqual("stopping", stopped.status) self.assertEqual("stopped", (await self.store.get_trace(trace_id)).status) self.assertEqual( "cancelled", (await self.store.get_tool_approval_batch(trace_id)).status, ) async def test_persisted_argument_hash_tampering_is_detected(self): trace_id, _batch = await self.start_waiting_batch() approval_path = Path(self.temp.name) / trace_id / "tool_approvals.json" raw = json.loads(approval_path.read_text(encoding="utf-8")) raw["calls"][0]["effective_arguments"]["value"] = 999 approval_path.write_text(json.dumps(raw), encoding="utf-8") with self.assertRaises(TraceStoreCorruptionError): await self.store.get_tool_approval_batch(trace_id) async def test_restart_marks_executing_side_effect_unknown_and_fails_closed(self): trace_id, batch = await self.start_waiting_batch() pending = batch.pending_calls[0] pending.decision = "approved" pending.execution_status = "executing" batch.status = "executing" await self.store.replace_tool_approval_batch(trace_id, batch) await self.store.update_trace(trace_id, status="running") await reconcile_traces() persisted = await self.store.get_tool_approval_batch(trace_id) trace = await self.store.get_trace(trace_id) self.assertEqual("execution_unknown", persisted.status) self.assertEqual( "execution_unknown", next( call for call in persisted.calls if call.tool_call_id == pending.tool_call_id ).execution_status, ) self.assertEqual("failed", trace.status) self.assertIn("unknown", trace.error_message) self.assertEqual([], self.side_effects) async def test_restart_restores_pending_batch_to_waiting_confirmation(self): trace_id, _batch = await self.start_waiting_batch() # Crash window: approval state is durable, but the following Trace # status update has not happened yet. await self.store.update_trace(trace_id, status="running") await reconcile_traces() trace = await self.store.get_trace(trace_id) persisted = await self.store.get_tool_approval_batch(trace_id) self.assertEqual("waiting_confirmation", trace.status) self.assertEqual("pending", persisted.status) self.assertEqual([], self.side_effects) async def test_orphaned_review_restores_confirmation_and_never_replays_unknown(self): registry = ToolRegistry() review_side_effects = [] async def review_task_result(child_trace_id: str): review_side_effects.append(child_trace_id) return {"status": "completed"} registry.register(review_task_result) runner = AgentRunner( trace_store=self.store, tool_registry=registry, llm_call=self.runner.llm_call, ) trace = Trace( trace_id="orphaned-review-approval", mode="agent", uid="user-1", context={}, ) await self.store.create_trace(trace) assistant = Message.create( trace_id=trace.trace_id, role="assistant", sequence=1, content={ "text": "", "tool_calls": [{ "id": "orphaned-review-call", "type": "function", "function": { "name": "review_task_result", "arguments": json.dumps({"child_trace_id": "child-1"}), }, }], }, ) await self.store.add_message(assistant) await self.store.update_trace( trace.trace_id, head_sequence=1, last_sequence=1, ) config = RunConfig(auto_execute_tools=False) with self.assertRaisesRegex(RuntimeError, "restored to persistent approval"): await runner._heal_orphaned_tool_calls( [assistant], trace, None, config, 2, ) batch = await self.store.get_tool_approval_batch(trace.trace_id) self.assertEqual("pending", batch.status) self.assertEqual("waiting_confirmation", ( await self.store.get_trace(trace.trace_id) ).status) self.assertEqual([], review_side_effects) batch.status = "executing" batch.calls[0].decision = "approved" batch.calls[0].execution_status = "executing" await self.store.replace_tool_approval_batch(trace.trace_id, batch) await self.store.update_trace(trace.trace_id, status="running") with self.assertRaisesRegex(RuntimeError, "outcome is unknown"): await runner._heal_orphaned_tool_calls( [assistant], await self.store.get_trace(trace.trace_id), None, config, 2, ) unknown = await self.store.get_tool_approval_batch(trace.trace_id) self.assertEqual("execution_unknown", unknown.status) self.assertEqual([], review_side_effects) async def test_restart_resumes_decided_batch_before_any_side_effect(self): trace_id, batch = await self.start_waiting_batch() # Crash window: decisions are durable, but no call has yet been marked # executing and no side effect has begun. pending = batch.pending_calls[0] pending.decision = "approved" pending.decided_at = datetime.now().isoformat() batch.status = "decided" batch.updated_at = datetime.now().isoformat() await self.store.replace_tool_approval_batch(trace_id, batch) await self.store.update_trace(trace_id, status="running") await reconcile_traces() self.assertIn(trace_id, run_api._running_tasks) await run_api._running_tasks[trace_id] persisted = await self.store.get_tool_approval_batch(trace_id) trace = await self.store.get_trace(trace_id) self.assertEqual("completed", persisted.status) self.assertEqual("completed", trace.status) self.assertEqual( [("dangerous", 1, "fixed"), ("safe", "after-approval")], self.side_effects, ) async def test_post_side_effect_persistence_failure_is_immediately_unknown(self): trace_id, batch = await self.start_waiting_batch() original_add_message = self.store.add_message async def fail_tool_result_persistence(message): if message.role == "tool": raise OSError("injected result persistence failure") return await original_add_message(message) with patch.object( self.store, "add_message", side_effect=fail_tool_result_persistence, ): await decide_tool_approval( trace_id, batch.batch_id, DecideToolApprovalRequest(decisions=[ToolApprovalDecision( tool_call_id="danger-1", decision="approve", )]), ) await run_api._running_tasks[trace_id] persisted = await self.store.get_tool_approval_batch(trace_id) trace = await self.store.get_trace(trace_id) self.assertEqual("execution_unknown", persisted.status) self.assertEqual("failed", trace.status) self.assertEqual([("dangerous", 1, "fixed")], self.side_effects) class TraceStoreHardeningTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.temp = tempfile.TemporaryDirectory() self.store = FileSystemTraceStore(self.temp.name) self.trace = Trace(trace_id="trace@safe-1", mode="agent", task="test") await self.store.create_trace(self.trace) async def asyncTearDown(self): self.temp.cleanup() async def test_concurrent_messages_and_events_keep_unique_monotonic_counters(self): await asyncio.gather(*[ self.store.add_message(Message.create( trace_id=self.trace.trace_id, role="user", sequence=index, content=f"message-{index}", )) for index in range(1, 41) ]) trace = await self.store.get_trace(self.trace.trace_id) events = await self.store.get_events(self.trace.trace_id) self.assertEqual(40, trace.total_messages) self.assertEqual(40, trace.last_sequence) self.assertEqual(list(range(1, 41)), [e["event_id"] for e in events]) extra_ids = await asyncio.gather(*[ self.store.append_event(self.trace.trace_id, "extra", {"index": index}) for index in range(20) ]) self.assertEqual(list(range(41, 61)), sorted(extra_ids)) self.assertEqual(60, (await self.store.get_trace(self.trace.trace_id)).last_event_id) async def test_duplicate_message_id_fails_without_double_counting(self): message = Message.create( trace_id=self.trace.trace_id, role="user", sequence=1, content="once", ) await self.store.add_message(message) with self.assertRaisesRegex(ValueError, "message_id already exists"): await self.store.add_message(message) trace = await self.store.get_trace(self.trace.trace_id) self.assertEqual(1, trace.total_messages) self.assertEqual(1, len(await self.store.get_trace_messages(trace.trace_id))) self.assertEqual(1, len(await self.store.get_events(trace.trace_id))) async def test_attachment_suffix_cannot_overwrite_message_json(self): message = Message.create( trace_id=self.trace.trace_id, role="user", sequence=1, content="durable", ) await self.store.add_message(message) attachment = await self.store.write_message_attachment( self.trace.trace_id, message.message_id, suffix=".json", data=b'{"attachment":true}', ) self.assertEqual("attachments", attachment.parent.name) self.assertEqual(b'{"attachment":true}', attachment.read_bytes()) messages = await self.store.get_trace_messages(self.trace.trace_id) self.assertEqual(["durable"], [item.content for item in messages]) async def test_atomic_replace_failure_keeps_previous_meta_valid(self): meta_path = Path(self.temp.name) / self.trace.trace_id / "meta.json" before = json.loads(meta_path.read_text(encoding="utf-8")) with patch("cyber_agent.trace.store.os.replace", side_effect=OSError("fault")): with self.assertRaises(OSError): await self.store.update_trace(self.trace.trace_id, status="failed") after = json.loads(meta_path.read_text(encoding="utf-8")) self.assertEqual(before, after) async def test_corruption_and_path_traversal_fail_closed(self): with self.assertRaises(ValueError): await self.store.get_trace("../outside") with self.assertRaises(ValueError): await self.store.get_message("bad/name") goal_path = Path(self.temp.name) / self.trace.trace_id / "goal.json" goal_path.write_text("{broken", encoding="utf-8") with self.assertRaises(TraceStoreCorruptionError): await self.store.get_goal_tree(self.trace.trace_id) corrupt_dir = Path(self.temp.name) / "corrupt-trace" corrupt_dir.mkdir() (corrupt_dir / "meta.json").write_text("[]", encoding="utf-8") with self.assertRaises(TraceStoreCorruptionError): await self.store.list_traces() async def test_startup_reconciliation_restores_file_derived_counters(self): for index in range(1, 4): await self.store.add_message(Message.create( trace_id=self.trace.trace_id, role="user", sequence=index, content=f"message-{index}", prompt_tokens=index, )) await self.store.update_trace( self.trace.trace_id, total_messages=0, last_sequence=0, last_event_id=0, ) report = await self.store.reconcile_trace_integrity(self.trace.trace_id) trace = await self.store.get_trace(self.trace.trace_id) self.assertEqual({"messages": 3, "last_sequence": 3, "last_event_id": 3}, report) self.assertEqual(3, trace.total_messages) self.assertEqual(6, trace.total_prompt_tokens) class SideBranchPersistenceTest(unittest.IsolatedAsyncioTestCase): async def test_all_side_branch_types_reconstruct_and_respect_total_turn_budget(self): for branch_type in ("compression", "reflection", "knowledge_eval"): with self.subTest(branch_type=branch_type), tempfile.TemporaryDirectory() as temp_dir: store = FileSystemTraceStore(temp_dir) branch_id = f"{branch_type}_persisted" context = { "agent_mode": "legacy", "agent_mode_revision": 1, "active_side_branch": { "type": branch_type, "branch_id": branch_id, "start_head_seq": 1, "start_sequence": 2, "start_iteration": 99, "max_turns": 2, # Simulate a crash after the assistant Message was # durable but before this counter update. "turns_used": 0, }, } config = RunConfig( trace_id=f"trace-{branch_type}", model="fake", max_iterations=5, side_branch_max_turns=9, tools=[], tool_groups=[], enable_research_flow=False, knowledge=quiet_knowledge(), ) persist_run_config_snapshot( context, RunConfigSnapshotV1.from_run_config( config, memory_identity=None, ), ) trace = Trace( trace_id=config.trace_id, mode="agent", task="resume branch", model="fake", context=context, head_sequence=1, ) await store.create_trace(trace) await store.update_goal_tree(trace.trace_id, GoalTree(mission="resume")) await store.add_message(Message.create( trace_id=trace.trace_id, role="user", sequence=1, content="main", )) await store.add_message(Message.create( trace_id=trace.trace_id, role="user", sequence=2, parent_sequence=1, branch_type=branch_type, branch_id=branch_id, content="branch instruction", )) await store.add_message(Message.create( trace_id=trace.trace_id, role="assistant", sequence=3, parent_sequence=2, branch_type=branch_type, branch_id=branch_id, content={"text": "first branch turn", "tool_calls": []}, )) await store.add_message(Message.create( trace_id=trace.trace_id, role="assistant", sequence=4, parent_sequence=1, branch_type=branch_type, branch_id="neighbor-branch", content={"text": "must not enter recovered history", "tool_calls": []}, )) await store.update_trace(trace.trace_id, head_sequence=1) first_prompt = [] calls = 0 async def llm_call(**kwargs): nonlocal calls calls += 1 if calls == 1: first_prompt.extend(kwargs["messages"]) content = ( "[[SUMMARY]] compacted" if branch_type == "compression" else ( json.dumps({"evaluations": []}) if branch_type == "knowledge_eval" else "reflection complete" ) ) return {"content": content, "tool_calls": [], "finish_reason": "stop"} return {"content": "main complete", "tool_calls": [], "finish_reason": "stop"} runner = AgentRunner(trace_store=store, llm_call=llm_call) await runner.run_result([], config) messages = await store.get_trace_messages(trace.trace_id) branch_assistants = [ message for message in messages if message.role == "assistant" and message.branch_id == branch_id and message.branch_type == branch_type ] self.assertEqual(2, len(branch_assistants)) self.assertNotIn( "must not enter recovered history", json.dumps(first_prompt, ensure_ascii=False), ) persisted = await store.get_trace(trace.trace_id) self.assertNotIn("active_side_branch", persisted.context) async def test_recovered_exhausted_branch_never_receives_an_extra_turn(self): for branch_type in ("compression", "reflection", "knowledge_eval"): with self.subTest(branch_type=branch_type), tempfile.TemporaryDirectory() as temp_dir: store = FileSystemTraceStore(temp_dir) trace_id = f"exhausted-{branch_type}" branch_id = f"{branch_type}-crash" context = { "agent_mode": "legacy", "agent_mode_revision": 1, "active_side_branch": { "type": branch_type, "branch_id": branch_id, "start_head_seq": 1, "start_sequence": 2, "start_iteration": 0, "max_turns": 1, "turns_used": 0, }, } config = RunConfig( trace_id=trace_id, model="fake", max_iterations=2, tools=[], tool_groups=[], enable_research_flow=False, knowledge=quiet_knowledge(), ) persist_run_config_snapshot( context, RunConfigSnapshotV1.from_run_config( config, memory_identity=None, ), ) await store.create_trace(Trace( trace_id=trace_id, mode="agent", task="resume exhausted branch", model="fake", context=context, head_sequence=1, )) await store.update_goal_tree(trace_id, GoalTree(mission="resume")) await store.add_message(Message.create( trace_id=trace_id, role="user", sequence=1, content="main", )) await store.add_message(Message.create( trace_id=trace_id, role="user", sequence=2, parent_sequence=1, branch_type=branch_type, branch_id=branch_id, content="branch instruction", )) await store.add_message(Message.create( trace_id=trace_id, role="assistant", sequence=3, parent_sequence=2, branch_type=branch_type, branch_id=branch_id, content={"text": "already used turn", "tool_calls": []}, )) await store.update_trace(trace_id, head_sequence=1) prompts = [] async def llm_call(**kwargs): prompts.append(kwargs["messages"]) return { "content": "main complete", "tool_calls": [], "finish_reason": "stop", } runner = AgentRunner(trace_store=store, llm_call=llm_call) await runner.run_result([], config) messages = await store.get_trace_messages(trace_id) branch_assistants = [ message for message in messages if message.role == "assistant" and message.branch_id == branch_id and message.branch_type == branch_type ] self.assertEqual(1, len(branch_assistants)) self.assertNotIn( "already used turn", json.dumps(prompts, ensure_ascii=False), ) self.assertNotIn( "active_side_branch", (await store.get_trace(trace_id)).context, ) if __name__ == "__main__": unittest.main()