Agent/tests/test_runtime_hardening.py

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