|
|
@@ -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 = {
|