Agent/tests/test_recursive_context_policy.py

import json
import tempfile
import unittest

from pydantic import ValidationError

from cyber_agent.core.agent_mode import (
    RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY,
    policy_from_context,
)
from cyber_agent.core.runner import AgentRunner, RunConfig
from cyber_agent.core.context_policy import (
    MAX_CONTEXT_REFS,
    MAX_CONTEXT_SNAPSHOTS_CHARS,
    MAX_ROOT_TASK_ANCHOR_CHARS,
    MAX_TASK_BRIEF_CHARS,
    ContextPolicyError,
    add_context_snapshot,
    build_child_context_access,
    canonical_json,
    context_chars,
    create_context_snapshot,
    get_authorized_context_snapshot,
    normalize_root_task_anchor,
    normalize_task_brief,
    persist_root_task_anchor,
    prune_context_access,
    replace_context_access,
    require_matching_root_task_anchor,
    require_root_task_anchor,
    root_task_anchor_hash,
)
from cyber_agent.core.task_protocol import (
    ContextRef,
    TaskBrief,
    initialize_task_progress,
    new_task_protocol,
    replace_task_brief,
)
from cyber_agent.tools.builtin.context import read_context_ref
from cyber_agent.tools.builtin.subagent import _get_allowed_tools
from cyber_agent.trace.models import Trace
from cyber_agent.trace.run_api import CreateRequest
from cyber_agent.trace.store import FileSystemTraceStore


ROOT_ANCHOR = {
    "objective": "五层委托仍围绕同一个根任务",
    "completion_criteria": ["所有局部结论可追溯", "根结果通过独立验收"],
    "constraints": ["不得编造证据"],
}


def brief(level: int, **updates) -> TaskBrief:
    values = {
        "objective": f"完成第 {level} 层局部任务",
        "reason": f"第 {level - 1} 层需要该结果",
        "completion_criteria": [f"第 {level} 层结论可验收"],
        "expected_outputs": [f"第 {level} 层结论"],
        "parent_findings": [f"直接父级结论 {level - 1}"],
        "context": {"local_level": level},
        "constraints": [f"约束-{level}"],
    }
    values.update(updates)
    return TaskBrief.model_validate(values)


def anchored_context(task_brief: TaskBrief | None = None) -> dict:
    context = {
        "agent_mode": "recursive",
        "agent_mode_revision": 3,
        "root_trace_id": "root",
        "agent_depth": 0,
        "task_protocol": new_task_protocol(task_brief),
    }
    anchor = persist_root_task_anchor(context, ROOT_ANCHOR)
    initialize_task_progress(
        context["task_protocol"],
        root_task_anchor_hash=context.get("root_task_anchor_hash"),
    )
    replace_context_access(
        context,
        [],
        root_task_anchor=anchor,
        task_brief=task_brief,
    )
    return context


class RootAnchorAndTaskBriefTest(unittest.TestCase):
    def test_anchor_is_canonical_and_hash_stays_identical_for_five_levels(self):
        anchor = normalize_root_task_anchor({
            "constraints": ["不得编造证据", "不得编造证据"],
            "completion_criteria": ["所有局部结论可追溯", "根结果通过独立验收"],
            "objective": "五层委托仍围绕同一个根任务",
        })
        expected_json = canonical_json(anchor.model_dump(mode="json"))
        expected_hash = root_task_anchor_hash(anchor)
        for _ in range(6):
            context = {}
            persisted = persist_root_task_anchor(context, anchor)
            self.assertEqual(expected_json, canonical_json(persisted.model_dump(mode="json")))
            self.assertEqual(expected_hash, context["root_task_anchor_hash"])
            self.assertEqual(expected_json, canonical_json(require_root_task_anchor(context).model_dump(mode="json")))

    def test_anchor_rejects_tampering_and_enforces_real_4000_char_limit(self):
        context = {}
        persist_root_task_anchor(context, ROOT_ANCHOR)
        context["root_task_anchor"]["objective"] = "被篡改"
        with self.assertRaisesRegex(ContextPolicyError, "hash"):
            require_root_task_anchor(context)

        low, high = 1, MAX_ROOT_TASK_ANCHOR_CHARS * 2
        while low < high:
            middle = (low + high + 1) // 2
            candidate = {**ROOT_ANCHOR, "objective": "根" * middle}
            try:
                normalize_root_task_anchor(candidate)
                low = middle
            except ContextPolicyError:
                high = middle - 1
        largest = normalize_root_task_anchor({**ROOT_ANCHOR, "objective": "根" * low})
        self.assertLessEqual(
            context_chars(largest.model_dump(mode="json")),
            MAX_ROOT_TASK_ANCHOR_CHARS,
        )
        with self.assertRaisesRegex(ContextPolicyError, "4000"):
            normalize_root_task_anchor({**ROOT_ANCHOR, "objective": "根" * (low + 1)})

    def test_self_consistent_child_anchor_must_still_match_authoritative_root(self):
        root_context = {}
        child_context = {}
        persist_root_task_anchor(root_context, ROOT_ANCHOR)
        persist_root_task_anchor(child_context, {
            **ROOT_ANCHOR,
            "objective": "另一份自洽但非权威的根目标",
        })
        with self.assertRaisesRegex(ContextPolicyError, "does not match"):
            require_matching_root_task_anchor(root_context, child_context)

    def test_five_level_constraints_accumulate_without_other_context_drift(self):
        parent = None
        inherited_constraints = ROOT_ANCHOR["constraints"][:]
        for level in range(1, 6):
            current = normalize_task_brief(
                brief(level),
                parent_task_brief=parent,
                root_task_anchor=ROOT_ANCHOR,
            )
            inherited_constraints.append(f"约束-{level}")
            self.assertEqual(inherited_constraints, current.constraints)
            self.assertEqual([f"直接父级结论 {level - 1}"], current.parent_findings)
            self.assertEqual({"local_level": level}, current.context)
            self.assertEqual([], current.context_refs)
            parent = current

    def test_task_brief_version_preserves_old_normalized_content(self):
        first = normalize_task_brief(brief(1), root_task_anchor=ROOT_ANCHOR)
        state = new_task_protocol(first)
        second = normalize_task_brief(
            brief(1, objective="修订后的局部任务"),
            root_task_anchor=ROOT_ANCHOR,
        )
        replace_task_brief(state, second, effective_at_sequence=9)
        self.assertEqual(2, state["task_brief_version"])
        self.assertEqual(9, state["task_brief_effective_at_sequence"])
        self.assertEqual(first.model_dump(), state["task_brief_history"][0]["task_brief"])
        self.assertEqual(second.model_dump(), state["task_brief"])

        second.objective = "外部再次修改"
        self.assertEqual("修订后的局部任务", state["task_brief"]["objective"])

    def test_task_brief_uses_stable_json_and_enforces_16000_chars(self):
        escaped = normalize_task_brief(
            brief(1, context={"quote": "\"中文\\n", "ordered": {"b": 2, "a": 1}}),
            root_task_anchor=ROOT_ANCHOR,
        )
        encoded = canonical_json(escaped.model_dump(mode="json"))
        self.assertEqual(encoded, canonical_json(json.loads(encoded)))

        with self.assertRaisesRegex(ContextPolicyError, str(MAX_TASK_BRIEF_CHARS)):
            normalize_task_brief(
                brief(1, context={"payload": "资料" * MAX_TASK_BRIEF_CHARS}),
                root_task_anchor=ROOT_ANCHOR,
            )

    def test_create_request_uses_new_anchor_and_rejects_removed_field(self):
        request = CreateRequest(
            messages=[{"role": "user", "content": "开始"}],
            root_task_anchor=ROOT_ANCHOR,
        )
        self.assertEqual(ROOT_ANCHOR["objective"], request.root_task_anchor.objective)
        with self.assertRaisesRegex(ValidationError, "root_task_anchor"):
            CreateRequest(
                messages=[{"role": "user", "content": "开始"}],
                root_completion_criteria=["完成"],
            )


class ContextReferencePolicyTest(unittest.TestCase):
    def test_root_child_only_gets_anchor_and_deeper_child_gets_parent_brief(self):
        root_context = anchored_context()
        level1 = normalize_task_brief(brief(1), root_task_anchor=ROOT_ANCHOR)
        child_access = build_child_context_access(
            parent_context=root_context,
            parent_trace_id="root",
            root_trace_id="root",
            uid="user-1",
            parent_task_state=root_context["task_protocol"],
            child_task_brief=level1,
            root_task_anchor=normalize_root_task_anchor(ROOT_ANCHOR),
        )
        self.assertEqual({}, child_access["snapshots"])

        level1_context = anchored_context(level1)
        level2 = normalize_task_brief(
            brief(2),
            parent_task_brief=level1,
            root_task_anchor=ROOT_ANCHOR,
        )
        grandchild_access = build_child_context_access(
            parent_context=level1_context,
            parent_trace_id="depth-1",
            root_trace_id="root",
            uid="user-1",
            parent_task_state=level1_context["task_protocol"],
            child_task_brief=level2,
            root_task_anchor=normalize_root_task_anchor(ROOT_ANCHOR),
        )
        snapshots = list(grandchild_access["snapshots"].values())
        self.assertEqual(1, len(snapshots))
        self.assertEqual("task_brief", snapshots[0]["kind"])
        self.assertEqual(level1.model_dump(mode="json"), snapshots[0]["content"])

    def test_reference_requires_exact_local_grant_version_tree_and_owner(self):
        context = anchored_context(brief(1))
        snapshot = create_context_snapshot(
            kind="reviewed_task_result",
            summary="已审核结果",
            source_trace_id="child",
            root_trace_id="root",
            uid="user-1",
            content={"task_report": {"summary": "真实结论"}},
            granted_at_sequence=7,
        )
        ref = add_context_snapshot(
            context,
            snapshot,
            root_task_anchor=ROOT_ANCHOR,
            task_brief=brief(1),
        )
        loaded = get_authorized_context_snapshot(
            context,
            ref_id=ref.ref_id,
            version=ref.version,
            root_trace_id="root",
            uid="user-1",
        )
        self.assertEqual("真实结论", loaded.content["task_report"]["summary"])

        for changes, message in (
            ({"version": "0" * 64}, "version"),
            ({"root_trace_id": "other-root"}, "another tree"),
            ({"uid": "user-2"}, "another tree"),
        ):
            values = {
                "ref_id": ref.ref_id,
                "version": ref.version,
                "root_trace_id": "root",
                "uid": "user-1",
                **changes,
            }
            with self.assertRaisesRegex(ContextPolicyError, message):
                get_authorized_context_snapshot(context, **values)
        with self.assertRaisesRegex(ContextPolicyError, "not authorized"):
            get_authorized_context_snapshot(
                context,
                ref_id="ctx_unknown",
                version="0" * 64,
                root_trace_id="root",
                uid="user-1",
            )

        context["context_access"]["snapshots"][ref.ref_id]["content"] = {
            "task_report": {"summary": "伪造结论"}
        }
        with self.assertRaisesRegex(ContextPolicyError, "version"):
            get_authorized_context_snapshot(
                context,
                ref_id=ref.ref_id,
                version=ref.version,
                root_trace_id="root",
                uid="user-1",
            )

        context = anchored_context(brief(1))
        ref = add_context_snapshot(
            context,
            snapshot,
            root_task_anchor=ROOT_ANCHOR,
            task_brief=brief(1),
        )
        context["context_access"]["snapshots"][ref.ref_id]["source_trace_id"] = "sibling"
        with self.assertRaisesRegex(ContextPolicyError, "ref_id"):
            get_authorized_context_snapshot(
                context,
                ref_id=ref.ref_id,
                version=ref.version,
                root_trace_id="root",
                uid="user-1",
            )

        context = anchored_context(brief(1))
        ref = add_context_snapshot(
            context,
            snapshot,
            root_task_anchor=ROOT_ANCHOR,
            task_brief=brief(1),
        )
        raw = context["context_access"]["snapshots"].pop(ref.ref_id)
        mismatched_key = "ctx_000000000000000000000000"
        context["context_access"]["snapshots"][mismatched_key] = raw
        with self.assertRaisesRegex(ContextPolicyError, "does not match its key"):
            get_authorized_context_snapshot(
                context,
                ref_id=mismatched_key,
                version=ref.version,
                root_trace_id="root",
                uid="user-1",
            )
        forwarded = brief(
            2,
            context_refs=[ContextRef(ref_id=mismatched_key, version=ref.version)],
        )
        with self.assertRaisesRegex(ContextPolicyError, "does not match its key"):
            build_child_context_access(
                parent_context=context,
                parent_trace_id="depth-1",
                root_trace_id="root",
                uid="user-1",
                parent_task_state=context["task_protocol"],
                child_task_brief=forwarded,
                root_task_anchor=normalize_root_task_anchor(ROOT_ANCHOR),
            )

    def test_reference_must_be_forwarded_at_every_edge(self):
        parent_context = anchored_context(brief(1))
        ancestor_snapshot = create_context_snapshot(
            kind="reviewed_task_result",
            summary="祖先已审核结果",
            source_trace_id="ancestor-child",
            root_trace_id="root",
            uid="user-1",
            content={"task_review": {"reason": "已确认的完整结论"}},
            granted_at_sequence=4,
        )
        ancestor_ref = add_context_snapshot(
            parent_context,
            ancestor_snapshot,
            root_task_anchor=ROOT_ANCHOR,
            task_brief=brief(1),
        )

        forwarded_brief = normalize_task_brief(
            brief(2, context_refs=[ancestor_ref]),
            parent_task_brief=brief(1),
            root_task_anchor=ROOT_ANCHOR,
        )
        forwarded_access = build_child_context_access(
            parent_context=parent_context,
            parent_trace_id="depth-1",
            root_trace_id="root",
            uid="user-1",
            parent_task_state=parent_context["task_protocol"],
            child_task_brief=forwarded_brief,
            root_task_anchor=normalize_root_task_anchor(ROOT_ANCHOR),
        )
        self.assertIn(ancestor_ref.ref_id, forwarded_access["snapshots"])

        child_context = anchored_context(forwarded_brief)
        child_context["context_access"] = forwarded_access
        not_forwarded = normalize_task_brief(
            brief(3),
            parent_task_brief=forwarded_brief,
            root_task_anchor=ROOT_ANCHOR,
        )
        deeper_access = build_child_context_access(
            parent_context=child_context,
            parent_trace_id="depth-2",
            root_trace_id="root",
            uid="user-1",
            parent_task_state=child_context["task_protocol"],
            child_task_brief=not_forwarded,
            root_task_anchor=normalize_root_task_anchor(ROOT_ANCHOR),
        )
        self.assertNotIn(ancestor_ref.ref_id, deeper_access["snapshots"])

    def test_reference_limits_and_rewind_cleanup_are_enforced(self):
        context = anchored_context(brief(1))
        snapshots = [
            create_context_snapshot(
                kind="reviewed_task_result",
                summary=f"结果 {index}",
                source_trace_id=f"child-{index}",
                root_trace_id="root",
                uid="user-1",
                content={"index": index},
                granted_at_sequence=index,
            )
            for index in range(MAX_CONTEXT_REFS + 1)
        ]
        replace_context_access(
            context,
            snapshots[:MAX_CONTEXT_REFS],
            root_task_anchor=ROOT_ANCHOR,
            task_brief=brief(1),
        )
        self.assertEqual(
            MAX_CONTEXT_REFS,
            context["context_access"]["metrics"]["authorized_ref_count"],
        )
        with self.assertRaisesRegex(ContextPolicyError, str(MAX_CONTEXT_REFS)):
            replace_context_access(
                context,
                snapshots,
                root_task_anchor=ROOT_ANCHOR,
                task_brief=brief(1),
            )

        with self.assertRaisesRegex(ContextPolicyError, "16000"):
            create_context_snapshot(
                kind="reviewed_task_result",
                summary="超大结果",
                source_trace_id="oversized",
                root_trace_id="root",
                uid="user-1",
                content={"payload": "x" * 16_000},
                granted_at_sequence=1,
            )

        prune_context_access(context, 3)
        kept = context["context_access"]["snapshots"].values()
        self.assertTrue(all(item["granted_at_sequence"] <= 3 for item in kept))

        oversized = [
            create_context_snapshot(
                kind="reviewed_task_result",
                summary=f"大结果 {index}",
                source_trace_id=f"large-{index}",
                root_trace_id="root",
                uid="user-1",
                content={"payload": "x" * 13_000, "index": index},
                granted_at_sequence=index,
            )
            for index in range(5)
        ]
        self.assertGreater(
            sum(context_chars(item.content) for item in oversized),
            MAX_CONTEXT_SNAPSHOTS_CHARS,
        )
        with self.assertRaisesRegex(ContextPolicyError, str(MAX_CONTEXT_SNAPSHOTS_CHARS)):
            replace_context_access(
                context,
                oversized,
                root_task_anchor=ROOT_ANCHOR,
                task_brief=brief(1),
            )


class ContextReferenceToolTest(unittest.IsolatedAsyncioTestCase):
    def test_legacy_and_recursive_revision_one_never_expose_context_tool(self):
        class Tools:
            names = [
                "agent",
                "read_context_ref",
                "submit_task_report",
                "review_task_result",
            ]

            def get_tool_names(self, current_url=None, groups=None):
                return list(self.names)

            def get_schemas(self, tool_names=None):
                selected = self.names if tool_names is None else tool_names
                return [{"type": "function", "function": {"name": name}}
                        for name in selected]

        runner = AgentRunner(tool_registry=Tools(), llm_call=lambda **_kwargs: None)
        config = RunConfig()
        for context in (
            {"agent_mode": "legacy", "agent_mode_revision": 1},
            {"agent_mode": "recursive", "agent_mode_revision": 1},
        ):
            trace = Trace(trace_id="trace", mode="agent", task="task", context=context)
            names = {
                item["function"]["name"]
                for item in runner._get_runtime_tool_schemas(config, trace)
            }
            self.assertNotIn("read_context_ref", names)
            self.assertNotIn("submit_task_report", names)
            self.assertNotIn("review_task_result", names)

    async def test_read_tool_returns_only_local_authorized_snapshot(self):
        with tempfile.TemporaryDirectory() as directory:
            store = FileSystemTraceStore(directory)
            context = anchored_context(brief(1))
            snapshot = create_context_snapshot(
                kind="reviewed_task_result",
                summary="已审核证据",
                source_trace_id="child",
                root_trace_id="root",
                uid="user-1",
                content={"evidence": [{"fact": "授权原文"}]},
                granted_at_sequence=5,
            )
            ref = add_context_snapshot(
                context,
                snapshot,
                root_task_anchor=ROOT_ANCHOR,
                task_brief=brief(1),
            )
            await store.create_trace(Trace(
                trace_id="root",
                mode="agent",
                task="root",
                uid="user-1",
                context=context,
            ))
            result = await read_context_ref(
                ref_id=ref.ref_id,
                version=ref.version,
                context={"store": store, "trace_id": "root"},
            )
            payload = json.loads(result.output)
            self.assertEqual("授权原文", payload["content"]["evidence"][0]["fact"])

            denied = await read_context_ref(
                ref_id="ctx_unknown",
                version="0" * 64,
                context={"store": store, "trace_id": "root"},
            )
            self.assertIn("not authorized", denied.error)

    async def test_read_and_validator_reject_child_anchor_that_differs_from_root(self):
        with tempfile.TemporaryDirectory() as directory:
            store = FileSystemTraceStore(directory)
            root_context = anchored_context()
            child_context = anchored_context(brief(1))
            persist_root_task_anchor(child_context, {
                **ROOT_ANCHOR,
                "objective": "非权威但自洽的根目标",
            })
            child_context["root_trace_id"] = "root"
            child_context["agent_depth"] = 1
            await store.create_trace(Trace(
                trace_id="root",
                mode="agent",
                task="root",
                uid="user-1",
                context=root_context,
            ))
            await store.create_trace(Trace(
                trace_id="child",
                mode="agent",
                task="child",
                parent_trace_id="root",
                uid="user-1",
                context=child_context,
            ))

            denied = await read_context_ref(
                ref_id="ctx_000000000000000000000000",
                version="0" * 64,
                context={"store": store, "trace_id": "child"},
            )
            self.assertIn("does not match", denied.error)

            async def must_not_call(**_kwargs):
                raise AssertionError("mismatched child must fail before Validator LLM")

            runner = AgentRunner(trace_store=store, llm_call=must_not_call)
            with self.assertRaisesRegex(ContextPolicyError, "does not match"):
                await runner.validate_recursive_trace(
                    "child",
                    scope="task",
                )

    def test_depth_five_keeps_read_permission_but_loses_agent(self):
        class Tools:
            @staticmethod
            def get_tool_names():
                return [
                    "agent",
                    "read_context_ref",
                    "submit_task_report",
                    "evaluate",
                    "bash_command",
                ]

        class Runner:
            tools = Tools()

        trace_context = {
            "agent_mode": "recursive",
            "agent_mode_revision": 2,
        }
        policy = policy_from_context(trace_context)
        context = {
            "runner": Runner(),
            RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY: [
                "agent",
                "read_context_ref",
                "submit_task_report",
            ],
        }
        depth_four = _get_allowed_tools(context, 4, policy)
        depth_five = _get_allowed_tools(context, 5, policy)
        self.assertIn("agent", depth_four)
        self.assertIn("read_context_ref", depth_four)
        self.assertNotIn("agent", depth_five)
        self.assertIn("read_context_ref", depth_five)
        self.assertNotIn("evaluate", depth_five)
        self.assertNotIn("bash_command", depth_five)

        context[RECURSIVE_CAPABILITY_TOOLS_CONTEXT_KEY].remove("read_context_ref")
        self.assertNotIn("read_context_ref", _get_allowed_tools(context, 4, policy))
        self.assertNotIn("read_context_ref", _get_allowed_tools(context, 5, policy))


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