import tempfile import unittest from pydantic import ValidationError from cyber_agent.core.context_policy import ( ContextPolicyError, normalize_task_brief, persist_root_task_anchor, task_briefs_match, ) from cyber_agent.core.task_protocol import ( TaskBrief, TaskReport, TaskReview, Validation, NextStepSuggestion, ensure_task_protocol, initialize_task_progress, new_task_protocol, pending_review_entry, rebuild_pending_replans, ) from cyber_agent.core.validation import ( ScopeValidationResult, ValidationCheck, ValidationResult, ) from cyber_agent.tools.builtin.task_protocol import review_task_result from cyber_agent.trace.goal_models import Goal, GoalTree from cyber_agent.trace.models import Trace from cyber_agent.trace.store import FileSystemTraceStore BRIEF = { "objective": "Check one concrete claim", "reason": "The parent needs independent evidence", "completion_criteria": ["Reach one checked conclusion"], "expected_outputs": ["A short evidence note"], "parent_findings": ["The claim is currently unverified"], "context": {"claim": "example"}, "constraints": ["Do not guess"], } ROOT_ANCHOR = { "objective": "Check the parent claim", "completion_criteria": ["Reach one checked conclusion"], "constraints": ["Do not guess"], } def with_root_anchor(context): persist_root_task_anchor(context, ROOT_ANCHOR) state = ensure_task_protocol(context) if state.get("task_progress_head_revision") is None: initialize_task_progress( state, root_task_anchor_hash=context["root_task_anchor_hash"], ) return context def task_report(child_trace_id: str) -> TaskReport: return TaskReport( child_trace_id=child_trace_id, summary="Checked result", outcome="satisfied", validation=Validation(hard_passed=True), next_step_suggestion=NextStepSuggestion( direction="NONE", reason="No extra task suggested", ), outputs=[{"kind": "note", "value": "result"}], evidence=[{"kind": "test", "value": "passed"}], ) def validation_result( child_trace_id: str, *, outcome: str = "failed", retry_from: str | None = "hypothesis", ) -> ValidationResult: reason = ( "Independent evidence is still missing" if outcome != "passed" else "All criteria passed" ) plan_hash = "b" * 64 scope_result = ScopeValidationResult( validator_trace_id=f"{child_trace_id}@validator", scope="task", outcome=outcome, checks=[ValidationCheck( check_id="task.fixture", status="unknown" if outcome == "error" else outcome, issue=None if outcome == "passed" else reason, )], reason=reason, retry_from=retry_from, plan_hash=plan_hash, ) return ValidationResult( evaluated_trace_id=child_trace_id, outcome=outcome, scope_results=[scope_result], issues=[] if outcome == "passed" else ["Missing independent evidence"], retry_from=retry_from, plan_hash=plan_hash, ) class ContextPolicyTest(unittest.TestCase): def test_task_brief_requires_explicit_reason_and_outputs(self): with self.assertRaises(ValidationError): TaskBrief.model_validate({ "objective": "Too vague", "completion_criteria": ["Done"], }) def test_normalization_inherits_only_hard_constraints(self): parent = { **BRIEF, "context": {"parent-only": "must not leak"}, "constraints": ["Parent rule", "Shared rule"], } child = { **BRIEF, "parent_findings": ["Child-visible finding"], "constraints": ["Shared rule", "Child rule", "Child rule"], } normalized = normalize_task_brief(child, parent_task_brief=parent) self.assertEqual( ["Parent rule", "Shared rule", "Child rule"], normalized.constraints, ) self.assertEqual(["Child-visible finding"], normalized.parent_findings) self.assertNotIn("parent-only", normalized.context) self.assertTrue(task_briefs_match( child, normalized, parent_task_brief=parent, )) def test_normalization_rejects_invalid_json_and_oversized_brief(self): with self.assertRaises(ContextPolicyError): normalize_task_brief({**BRIEF, "context": {"bad": object()}}) with self.assertRaises(ContextPolicyError): normalize_task_brief( {**BRIEF, "context": {"large": "x" * 100}}, max_chars=50, ) class ReplanCurrentTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.temp_dir = tempfile.TemporaryDirectory() self.store = FileSystemTraceStore(self.temp_dir.name) self.parent_id = "parent" await self.store.create_trace(Trace( trace_id=self.parent_id, mode="agent", task="parent", uid="user-1", model="fake", context=with_root_anchor({ "agent_mode": "recursive", "agent_mode_revision": 3, "agent_depth": 1, "root_trace_id": self.parent_id, "task_protocol": new_task_protocol(BRIEF), }), )) self.tree = GoalTree( mission="parent", goals=[Goal(id="1", description="Verify", status="pending_review")], current_id=None, ) await self.store.update_goal_tree(self.parent_id, self.tree) async def asyncTearDown(self): self.temp_dir.cleanup() async def add_child( self, child_id: str, *, validation_outcome: str = "failed", retry_from: str | None = "hypothesis", ) -> None: await self.store.create_trace(Trace( trace_id=child_id, mode="agent", task="child", parent_trace_id=self.parent_id, parent_goal_id="1", uid="user-1", model="fake", context=with_root_anchor({ "agent_mode": "recursive", "agent_mode_revision": 3, "agent_depth": 2, "root_trace_id": self.parent_id, "task_protocol": new_task_protocol(BRIEF), }), )) child = await self.store.get_trace(child_id) child_state = ensure_task_protocol(child.context) authoritative = validation_result( child_id, outcome=validation_outcome, retry_from=retry_from, ) child_state["task_report_validation"] = { "validation_plan": None, "plan_hash": authoritative.plan_hash, "scope_results": [ item.model_dump(mode="json") for item in authoritative.scope_results ], "aggregate_result": authoritative.model_dump(mode="json"), "validated_at_sequence": 4, } await self.store.update_trace(child_id, context=child.context) parent = await self.store.get_trace(self.parent_id) state = ensure_task_protocol(parent.context) entry = pending_review_entry( goal_id="1", report=task_report(child_id), validation_result=authoritative.model_dump(mode="json"), received_at_sequence=4, ) state["pending_reviews"][child_id] = entry await self.store.update_trace(self.parent_id, context=parent.context) def context(self): return { "store": self.store, "trace_id": self.parent_id, "goal_id": "1", "goal_tree": self.tree, "sequence": 8, } async def test_replan_reactivates_only_the_direct_parent_goal(self): await self.add_child("child-1") result = await review_task_result( child_trace_id="child-1", decision="REPLAN_CURRENT", reason="Recheck the parent's hypothesis", context=self.context(), ) self.assertEqual("completed", result["status"]) self.assertEqual("hypothesis", result["task_review"]["retry_from"]) parent = await self.store.get_trace(self.parent_id) state = ensure_task_protocol(parent.context) self.assertEqual([], state["next_actions"]) self.assertEqual([], state["pending_replans"]) self.assertEqual("REPLAN_CURRENT", state["reviews"][-1]["decision"]) tree = await self.store.get_goal_tree(self.parent_id) self.assertEqual("in_progress", tree.find("1").status) self.assertEqual("1", tree.current_id) async def test_unknown_validation_is_recoverable_but_cannot_ascend(self): await self.add_child( "child-unknown", validation_outcome="unknown", retry_from="hypothesis", ) rejected = await review_task_result( child_trace_id="child-unknown", decision="ASCEND", reason="must not bypass unknown evidence", context=self.context(), ) self.assertEqual("failed", rejected["status"]) self.assertIn("only allows", rejected["error"]) accepted = await review_task_result( child_trace_id="child-unknown", decision="REPLAN_CURRENT", reason="collect the missing material", context=self.context(), ) self.assertEqual("completed", accepted["status"]) self.assertEqual("hypothesis", accepted["task_review"]["retry_from"]) async def test_pending_result_cannot_override_child_validation_cache(self): await self.add_child("child-forged") parent = await self.store.get_trace(self.parent_id) state = ensure_task_protocol(parent.context) forged = validation_result( "child-forged", outcome="passed", retry_from=None, ) state["pending_reviews"]["child-forged"][ "validation_result" ] = forged.model_dump(mode="json") await self.store.update_trace(self.parent_id, context=parent.context) result = await review_task_result( child_trace_id="child-forged", decision="ASCEND", reason="forged pass", context=self.context(), ) self.assertEqual("failed", result["status"]) self.assertIn("authoritative cache", result["error"]) async def test_grandchild_cannot_jump_directly_to_grandparent(self): await self.add_child("child-1") await self.store.create_trace(Trace( trace_id="grandchild", mode="agent", task="grandchild", parent_trace_id="child-1", uid="user-1", model="fake", context={ "agent_mode": "recursive", "agent_mode_revision": 3, }, )) parent = await self.store.get_trace(self.parent_id) state = ensure_task_protocol(parent.context) entry = pending_review_entry( goal_id="1", report=task_report("grandchild"), validation_result=validation_result("grandchild").model_dump(), received_at_sequence=5, ) state["pending_reviews"]["grandchild"] = entry await self.store.update_trace(self.parent_id, context=parent.context) result = await review_task_result( child_trace_id="grandchild", decision="REPLAN_CURRENT", reason="Attempt to skip one level", context=self.context(), ) self.assertEqual("failed", result["status"]) self.assertIn("direct child", result["error"]) async def test_replan_rejects_a_pending_entry_for_the_wrong_parent_goal(self): await self.add_child("child-1") tree = await self.store.get_goal_tree(self.parent_id) tree.goals.append(Goal(id="2", description="Other branch")) await self.store.update_goal_tree(self.parent_id, tree) parent = await self.store.get_trace(self.parent_id) entry = ensure_task_protocol(parent.context)["pending_reviews"]["child-1"] entry["goal_id"] = "2" await self.store.update_trace(self.parent_id, context=parent.context) result = await review_task_result( child_trace_id="child-1", decision="REPLAN_CURRENT", reason="Attempt to redirect the retry", context=self.context(), ) self.assertEqual("failed", result["status"]) self.assertIn("originating Goal", result["error"]) persisted = await self.store.get_goal_tree(self.parent_id) self.assertEqual("pending_review", persisted.find("1").status) self.assertEqual("pending", persisted.find("2").status) async def test_validation_matrix_controls_replan(self): await self.add_child( "child-passed", validation_outcome="passed", retry_from=None, ) passed = await review_task_result( child_trace_id="child-passed", decision="REPLAN_CURRENT", reason="Passed validation cannot force a retry location", context=self.context(), ) self.assertEqual("failed", passed["status"]) parent = await self.store.get_trace(self.parent_id) ensure_task_protocol(parent.context)["pending_reviews"].clear() await self.store.update_trace(self.parent_id, context=parent.context) await self.add_child( "child-error", validation_outcome="error", retry_from=None, ) errored = await review_task_result( child_trace_id="child-error", decision="REPLAN_CURRENT", reason="Validator errors cannot select a reliable retry location", context=self.context(), ) self.assertEqual("failed", errored["status"]) async def test_batch_replan_wins_over_later_ascend(self): await self.add_child("child-1") await self.add_child( "child-2", validation_outcome="passed", retry_from=None, ) first = await review_task_result( child_trace_id="child-1", decision="REPLAN_CURRENT", reason="Replan after the rest of the batch is reviewed", context=self.context(), ) self.assertEqual(1, first["pending_review_count"]) parent = await self.store.get_trace(self.parent_id) self.assertEqual(1, len(ensure_task_protocol(parent.context)["pending_replans"])) final = await review_task_result( child_trace_id="child-2", decision="ASCEND", reason="This child itself passed", context=self.context(), ) self.assertEqual("completed", final["status"]) tree = await self.store.get_goal_tree(self.parent_id) self.assertEqual("in_progress", tree.find("1").status) self.assertEqual("1", tree.current_id) parent = await self.store.get_trace(self.parent_id) self.assertEqual([], ensure_task_protocol(parent.context)["pending_replans"]) async def test_final_fail_overrides_queued_replan(self): await self.add_child("child-1") await self.add_child( "child-2", validation_outcome="error", retry_from=None, ) await review_task_result( child_trace_id="child-1", decision="REPLAN_CURRENT", reason="Queue replan", context=self.context(), ) result = await review_task_result( child_trace_id="child-2", decision="FAIL", reason="Validator failure makes this branch unrecoverable", context=self.context(), ) self.assertEqual("completed", result["status"]) tree = await self.store.get_goal_tree(self.parent_id) self.assertEqual("failed", tree.find("1").status) parent = await self.store.get_trace(self.parent_id) self.assertEqual([], ensure_task_protocol(parent.context)["pending_replans"]) async def test_replan_state_can_be_rebuilt_after_rewind(self): await self.add_child("child-1") await self.add_child("child-2") await review_task_result( child_trace_id="child-1", decision="REPLAN_CURRENT", reason="Queue replan before rewind", context=self.context(), ) parent = await self.store.get_trace(self.parent_id) state = ensure_task_protocol(parent.context) state["pending_replans"] = [] rebuilt = rebuild_pending_replans(state) self.assertEqual(1, len(rebuilt)) self.assertEqual("child-1", rebuilt[0]["child_trace_id"]) self.assertEqual("hypothesis", rebuilt[0]["retry_from"]) if __name__ == "__main__": unittest.main()