Explorar o código

补齐 recursive 上下文规范化与授权安全测试

覆盖五层根锚点稳定性、硬约束累积、字符边界、快照哈希和 TaskBrief 版本不可变性。

系统验证未授权、跨 root、跨 UID、兄弟引用、伪造版本和停止转授权等越权路径均被拒绝。
SamLee hai 1 día
pai
achega
9c51f0dc9a
Modificáronse 1 ficheiros con 626 adicións e 0 borrados
  1. 626 0
      tests/test_recursive_context_policy.py

+ 626 - 0
tests/test_recursive_context_policy.py

@@ -0,0 +1,626 @@
+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,
+    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": 2,
+        "root_trace_id": "root",
+        "agent_depth": 0,
+        "task_protocol": new_task_protocol(task_brief),
+    }
+    anchor = persist_root_task_anchor(context, ROOT_ANCHOR)
+    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()