Agent/tests/test_recursive_replan_context.py

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,
    new_task_protocol,
    pending_review_entry,
    rebuild_pending_replans,
)
from cyber_agent.core.validation import 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)
    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:
    return ValidationResult(
        validator_trace_id=f"{child_trace_id}@validator",
        evaluated_trace_id=child_trace_id,
        outcome=outcome,
        scope="task",
        reason=(
            "Independent evidence is still missing"
            if outcome != "passed"
            else "All criteria passed"
        ),
        issues=[] if outcome == "passed" else ["Missing independent evidence"],
        retry_from=retry_from,
    )


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": 2,
                "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": 2,
                "agent_depth": 2,
                "root_trace_id": self.parent_id,
                "task_protocol": new_task_protocol(BRIEF),
            }),
        ))
        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=validation_result(
                child_id,
                outcome=validation_outcome,
                retry_from=retry_from,
            ).model_dump(),
            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_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": 2,
            },
        ))
        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()