Jelajahi Sumber

补齐 TaskBrief 版本与上下文协议测试

覆盖 RootTaskAnchor、ContextRef、TaskBrief 稳定格式、版本历史和替换语义。

验证审核结果引用生成、修订时授权替换以及协议状态的序列化往返,防止任务合同在多轮运行中漂移。
SamLee 21 jam lalu
induk
melakukan
8e2b45fe49
1 mengubah file dengan 220 tambahan dan 20 penghapusan
  1. 220 20
      tests/test_task_protocol.py

+ 220 - 20
tests/test_task_protocol.py

@@ -12,12 +12,18 @@ from cyber_agent.core.resource_budget import (
     ResourceBudgetController,
 )
 from cyber_agent.core.runner import AgentRunner, RunConfig
+from cyber_agent.core.context_policy import (
+    add_context_snapshot,
+    create_context_snapshot,
+    persist_root_task_anchor,
+)
 from cyber_agent.core.task_protocol import (
     TaskBrief,
     TaskReport,
     TaskReview,
     ensure_task_protocol,
     new_task_protocol,
+    replace_task_brief,
 )
 from cyber_agent.tools import get_tool_registry
 from cyber_agent.tools.builtin.knowledge import KnowledgeConfig
@@ -40,6 +46,16 @@ BRIEF = {
     "expected_outputs": ["one checked conclusion"],
     "constraints": ["do not guess"],
 }
+ROOT_ANCHOR = {
+    "objective": "finish the root task",
+    "completion_criteria": ["finish the root task"],
+    "constraints": ["do not guess"],
+}
+
+
+def with_root_anchor(context):
+    persist_root_task_anchor(context, ROOT_ANCHOR)
+    return context
 
 
 def report_payload(*, outcome="satisfied"):
@@ -88,10 +104,10 @@ async def create_recursive_root(store, trace_id, *, state=None, uid="user-1"):
         "agent_mode_revision": 2,
         "agent_depth": 0,
         "root_trace_id": trace_id,
-        "root_completion_criteria": ["finish the root task"],
         RESOURCE_BUDGET_CONTEXT_KEY: budget.to_dict(),
         "task_protocol": state or new_task_protocol(),
     }
+    persist_root_task_anchor(context, ROOT_ANCHOR)
     await store.create_trace(Trace(
         trace_id=trace_id,
         mode="agent",
@@ -118,6 +134,7 @@ class FakeToolRegistry:
             "read_file",
             "submit_task_report",
             "review_task_result",
+            "read_context_ref",
         ]
 
 
@@ -180,11 +197,11 @@ class ReportingRunner:
             uid=evaluated.uid,
             status="completed",
             tools=[],
-            context={
+            context=with_root_anchor({
                 "created_by_tool": "validator",
                 "root_trace_id": evaluated.context["root_trace_id"],
                 "agent_depth": evaluated.context["agent_depth"],
-            },
+            }),
         ))
         return ValidationRun(result=result, trace_id=validator_trace_id)
 
@@ -283,7 +300,12 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
     async def test_models_validate_and_round_trip(self):
         brief = TaskBrief.model_validate(BRIEF)
         self.assertEqual(
-            {**BRIEF, "parent_findings": [], "context": {}},
+            {
+                **BRIEF,
+                "parent_findings": [],
+                "context": {},
+                "context_refs": [],
+            },
             brief.model_dump(exclude_none=True),
         )
         with self.assertRaises(ValidationError):
@@ -330,7 +352,10 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
                 RunConfig(), restored
             )
         }
-        self.assertEqual({"review_task_result"}, restored_names)
+        self.assertEqual(
+            {"review_task_result", "read_context_ref"},
+            restored_names,
+        )
 
         blocked = await agent(task_brief=BRIEF, context=self.context())
         self.assertEqual("failed", blocked["status"])
@@ -437,6 +462,11 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
             context=self.context(),
         )
         self.assertEqual("completed", reviewed["status"])
+        self.assertIsNotNone(reviewed["context_ref"])
+        self.assertIn(
+            reviewed["context_ref"],
+            reviewed["next_action_task_brief"]["context_refs"],
+        )
         parent = await self.store.get_trace(self.root_id)
         self.assertEqual(
             0,
@@ -446,7 +476,7 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
             tool_registry=get_tool_registry()
         )._get_runtime_tool_schemas(RunConfig(), parent)
         self.assertEqual(
-            {"agent"},
+            {"agent", "read_context_ref"},
             {schema["function"]["name"] for schema in schemas},
         )
 
@@ -454,12 +484,84 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
         self.assertEqual("failed", wrong["status"])
         self.assertIn("approved next task", wrong["error"])
 
-        second = await agent(task_brief=next_brief, context=self.context())
+        second = await agent(
+            task_brief=reviewed["next_action_task_brief"],
+            context=self.context(),
+        )
         self.assertEqual("completed", second["status"])
         self.assertEqual(2, len(await self.children()))
+        second_child = await self.store.get_trace(second["sub_trace_id"])
+        self.assertIn(
+            reviewed["context_ref"]["ref_id"],
+            second_child.context["context_access"]["snapshots"],
+        )
         parent = await self.store.get_trace(self.root_id)
         self.assertEqual([], ensure_task_protocol(parent.context)["next_actions"])
 
+    async def test_review_rejects_an_unusable_context_ref_without_consuming_report(self):
+        first = await agent(task_brief=BRIEF, context=self.context())
+        invalid_next = {
+            "objective": "use an unauthorized context ref",
+            "reason": "exercise approval validation",
+            "completion_criteria": ["must not be approved"],
+            "expected_outputs": ["nothing"],
+            "context_refs": [{
+                "ref_id": "ctx_000000000000000000000000",
+                "version": "0" * 64,
+            }],
+        }
+        reviewed = await review_task_result(
+            child_trace_id=first["sub_trace_id"],
+            decision="ACCEPT_REFINE",
+            reason="try an invalid transfer",
+            approved_next_task=invalid_next,
+            context=self.context(),
+        )
+        self.assertEqual("failed", reviewed["status"])
+        self.assertIn("not authorized", reviewed["error"])
+        parent = await self.store.get_trace(self.root_id)
+        state = ensure_task_protocol(parent.context)
+        self.assertIn(first["sub_trace_id"], state["pending_reviews"])
+        self.assertEqual([], state["reviews"])
+
+    async def test_oversized_review_snapshot_warns_without_blocking_review(self):
+        first = await agent(task_brief=BRIEF, context=self.context())
+        parent = await self.store.get_trace(self.root_id)
+        state = ensure_task_protocol(parent.context)
+        state["pending_reviews"][first["sub_trace_id"]]["task_report"]["outputs"] = [
+            {"payload": "x" * 16_000}
+        ]
+        await self.store.update_trace(self.root_id, context=parent.context)
+
+        reviewed = await review_task_result(
+            child_trace_id=first["sub_trace_id"],
+            decision="ASCEND",
+            reason="The task passed but its transferable snapshot is too large",
+            context=self.context(),
+        )
+        self.assertEqual("completed", reviewed["status"])
+        self.assertIsNone(reviewed["context_ref"])
+        self.assertIn("16000", reviewed["context_ref_warning"])
+        parent = await self.store.get_trace(self.root_id)
+        self.assertEqual({}, ensure_task_protocol(parent.context)["pending_reviews"])
+
+    async def test_batch_context_refs_are_validated_before_any_trace_is_created(self):
+        invalid_ref_brief = {
+            **BRIEF,
+            "objective": "second child with an unauthorized reference",
+            "context_refs": [{
+                "ref_id": "ctx_000000000000000000000000",
+                "version": "0" * 64,
+            }],
+        }
+        result = await agent(
+            task_brief=[BRIEF, invalid_ref_brief],
+            context=self.context(),
+        )
+        self.assertEqual("failed", result["status"])
+        self.assertIn("not authorized", result["error"])
+        self.assertEqual([], await self.children())
+
     async def test_batch_reports_are_reviewed_one_by_one(self):
         briefs = [
             BRIEF,
@@ -582,6 +684,62 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
         self.assertEqual("failed", invalid_ascend["status"])
         self.assertIn("only allow", invalid_ascend["error"])
 
+    async def test_submit_and_review_reject_child_anchor_that_differs_from_root(self):
+        child_context = with_root_anchor({
+            "agent_mode": "recursive",
+            "agent_mode_revision": 2,
+            "agent_depth": 1,
+            "root_trace_id": self.root_id,
+            "created_by_tool": "agent",
+            "task_protocol": new_task_protocol(BRIEF),
+        })
+        persist_root_task_anchor(child_context, {
+            **ROOT_ANCHOR,
+            "objective": "self-consistent but unauthorized root objective",
+        })
+        child_id = "mismatched-anchor-child"
+        await self.store.create_trace(Trace(
+            trace_id=child_id,
+            mode="agent",
+            task="child",
+            parent_trace_id=self.root_id,
+            parent_goal_id="1",
+            uid="user-1",
+            model="fake-model",
+            context=child_context,
+        ))
+        submitted = await submit_task_report(
+            task_report=report_payload(),
+            context={"store": self.store, "trace_id": child_id},
+        )
+        self.assertEqual("failed", submitted["status"])
+        self.assertIn("does not match", submitted["error"])
+
+        parent = await self.store.get_trace(self.root_id)
+        parent_state = ensure_task_protocol(parent.context)
+        parent_state["pending_reviews"][child_id] = {
+            "goal_id": "1",
+            "task_report": TaskReport(
+                child_trace_id=child_id,
+                **report_payload(),
+            ).model_dump(),
+            "validation_result": validation_payload(child_id),
+            "received_at_sequence": 1,
+        }
+        await self.store.update_trace(self.root_id, context=parent.context)
+        reviewed = await review_task_result(
+            child_trace_id=child_id,
+            decision="ASCEND",
+            reason="must not accept a different root anchor",
+            context=self.context(),
+        )
+        self.assertEqual("failed", reviewed["status"])
+        self.assertIn("does not match", reviewed["error"])
+        parent = await self.store.get_trace(self.root_id)
+        state = ensure_task_protocol(parent.context)
+        self.assertIn(child_id, state["pending_reviews"])
+        self.assertEqual([], state["reviews"])
+
     async def test_review_rejects_wrong_owner_even_if_report_is_injected(self):
         first = await agent(task_brief=BRIEF, context=self.context())
         child_id = first["sub_trace_id"]
@@ -645,6 +803,8 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
             context=self.context(),
         )
         self.assertEqual("completed", reviewed["status"])
+        cleared_child = await self.store.get_trace(child_id)
+        self.assertEqual({}, cleared_child.context["context_access"]["snapshots"])
 
         wrong = await agent(
             task_brief={
@@ -659,8 +819,9 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
         self.assertEqual("failed", wrong["status"])
         self.assertIn("approved revision", wrong["error"])
 
+        revised_brief = reviewed["next_action_task_brief"]
         revised = await agent(
-            task_brief=BRIEF,
+            task_brief=revised_brief,
             continue_from=child_id,
             context=self.context(),
         )
@@ -670,9 +831,14 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
         child = await self.store.get_trace(child_id)
         self.assertEqual(1, len(ensure_task_protocol(child.context)["report_history"]))
         self.assertEqual(
-            TaskBrief.model_validate(BRIEF).model_dump(),
+            TaskBrief.model_validate(revised_brief).model_dump(),
             ensure_task_protocol(child.context)["task_brief"],
         )
+        self.assertEqual(2, ensure_task_protocol(child.context)["task_brief_version"])
+        self.assertEqual(
+            1,
+            len(ensure_task_protocol(child.context)["task_brief_history"]),
+        )
 
     async def test_review_cannot_reserve_more_children_than_quota(self):
         briefs = [
@@ -771,20 +937,20 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
         self.assertEqual("pending_review", restored_tree.find("1").status)
         schemas = runner._get_runtime_tool_schemas(RunConfig(), parent)
         self.assertEqual(
-            {"review_task_result"},
+            {"review_task_result", "read_context_ref"},
             {schema["function"]["name"] for schema in schemas},
         )
 
     async def test_grandchild_report_is_reviewed_only_by_direct_child(self):
         child_id = "manual-child"
-        child_context = {
+        child_context = with_root_anchor({
             "agent_mode": "recursive",
             "agent_mode_revision": 2,
             "agent_depth": 1,
             "root_trace_id": self.root_id,
             "created_by_tool": "agent",
             "task_protocol": new_task_protocol(BRIEF),
-        }
+        })
         await self.store.create_trace(Trace(
             trace_id=child_id,
             mode="agent",
@@ -832,14 +998,14 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
 
     async def test_child_must_execute_approved_action_before_own_report(self):
         child_id = "child-with-approved-action"
-        child_context = {
+        child_context = with_root_anchor({
             "agent_mode": "recursive",
             "agent_mode_revision": 2,
             "agent_depth": 1,
             "root_trace_id": self.root_id,
             "created_by_tool": "agent",
             "task_protocol": new_task_protocol(BRIEF),
-        }
+        })
         await self.store.create_trace(Trace(
             trace_id=child_id,
             mode="agent",
@@ -880,7 +1046,10 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
         self.assertEqual("failed", blocked_report["status"])
         self.assertIn("approved next action", blocked_report["error"])
 
-        second = await agent(task_brief=next_brief, context=child_tool_context)
+        second = await agent(
+            task_brief=reviewed["next_action_task_brief"],
+            context=child_tool_context,
+        )
         self.assertEqual("completed", second["status"])
         child = await self.store.get_trace(child_id)
         state = ensure_task_protocol(child.context)
@@ -915,14 +1084,14 @@ class RecursiveTaskProtocolTest(unittest.IsolatedAsyncioTestCase):
             child_trace_id=child_id,
             **report_payload(outcome="failed"),
         )
-        child_context = {
+        child_context = with_root_anchor({
             "agent_mode": "recursive",
             "agent_mode_revision": 2,
             "agent_depth": 1,
             "root_trace_id": self.root_id,
             "created_by_tool": "agent",
             "task_protocol": new_task_protocol(BRIEF),
-        }
+        })
         child_context["task_protocol"]["task_report"] = failed_report.model_dump()
         await self.store.create_trace(Trace(
             trace_id=child_id,
@@ -968,14 +1137,14 @@ class RunnerCompletionGateTest(unittest.IsolatedAsyncioTestCase):
             task="must report",
             parent_trace_id="root",
             uid="user-1",
-            context={
+            context=with_root_anchor({
                 "agent_mode": "recursive",
                 "agent_mode_revision": 2,
                 "agent_depth": 1,
                 "root_trace_id": "root",
                 "created_by_tool": "agent",
                 "task_protocol": new_task_protocol(BRIEF),
-            },
+            }),
         ))
         await self.store.update_goal_tree(self.trace_id, GoalTree(mission="must report"))
 
@@ -1158,6 +1327,30 @@ class RunnerCompletionGateTest(unittest.IsolatedAsyncioTestCase):
     async def test_rewind_clears_protocol_state_after_cutoff(self):
         trace = await self.store.get_trace(self.trace_id)
         state = ensure_task_protocol(trace.context)
+        original_brief = TaskBrief.model_validate(state["task_brief"])
+        revised_brief = original_brief.model_copy(
+            update={"objective": "revised after the rewind cutoff"}
+        )
+        replace_task_brief(
+            state,
+            revised_brief,
+            effective_at_sequence=8,
+        )
+        late_snapshot = create_context_snapshot(
+            kind="reviewed_task_result",
+            summary="granted after cutoff",
+            source_trace_id="late-child",
+            root_trace_id="root",
+            uid="user-1",
+            content={"result": "must be removed by rewind"},
+            granted_at_sequence=8,
+        )
+        add_context_snapshot(
+            trace.context,
+            late_snapshot,
+            root_task_anchor=ROOT_ANCHOR,
+            task_brief=revised_brief,
+        )
         report = TaskReport(
             child_trace_id=self.trace_id,
             **report_payload(),
@@ -1200,6 +1393,10 @@ class RunnerCompletionGateTest(unittest.IsolatedAsyncioTestCase):
         self.assertEqual({}, state["pending_reviews"])
         self.assertEqual([], state["reviews"])
         self.assertEqual([], state["next_actions"])
+        self.assertEqual(1, state["task_brief_version"])
+        self.assertEqual(original_brief.model_dump(), state["task_brief"])
+        self.assertEqual([], state["task_brief_history"])
+        self.assertEqual({}, rewound.context["context_access"]["snapshots"])
         rewound_tree = await self.store.get_goal_tree(self.trace_id)
         self.assertEqual("pending", rewound_tree.find("1").status)
 
@@ -1343,7 +1540,10 @@ class RuntimeProtocolSchemaTest(unittest.TestCase):
         state = ensure_task_protocol(child.context)
         state["pending_reviews"]["grandchild"] = {"task_report": {}}
         review_schemas = self.runner._get_runtime_tool_schemas(RunConfig(), child)
-        self.assertEqual({"review_task_result"}, self.names(review_schemas))
+        self.assertEqual(
+            {"review_task_result", "read_context_ref"},
+            self.names(review_schemas),
+        )
 
     def test_protocol_tools_publish_nested_object_schemas(self):
         schemas = {