test(hardening): 覆盖运行恢复与文件账本故障窗口

新增 RunConfig snapshot 全字段 round-trip、哈希篡改、旧 Recursive 拒绝、Legacy 推导和历史项目 Runner 恢复测试。

覆盖工具精确白名单、URL 门禁、混合审批批次、参数编辑与篡改、拒绝、重复决策、停止取消，以及 pending、decided、executing 和副作用后持久化失败等状态。

增加 TraceStore 并发消息与事件压力、原子替换故障、路径穿越、JSON 损坏、计数对账、附件后缀碰撞和重复 Message 测试。

对 compression、reflection、knowledge_eval 三类 SideBranch 验证跨续跑轮次重建、邻近分支隔离和崩溃后不超预算。

tests/test_runtime_hardening.py

@@ -0,0 +1,999 @@
+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_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, "predates RunConfig snapshots"):
+                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_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()