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