Agent/tests/test_task_progress.py

import asyncio
from copy import deepcopy
import tempfile
import unittest

from pydantic import ValidationError

from cyber_agent.core.context_policy import (
    add_context_snapshot,
    create_context_snapshot,
    normalize_root_task_anchor,
    persist_root_task_anchor,
    replace_context_access,
)
from cyber_agent.core.resource_budget import (
    RESOURCE_BUDGET_CONTEXT_KEY,
    ResourceBudget,
    ResourceBudgetController,
)
from cyber_agent.core.run_snapshot import (
    RunConfigSnapshotV1,
    persist_run_config_snapshot,
)
from cyber_agent.core.runner import AgentRunner, RunConfig
from cyber_agent.core.task_protocol import (
    Blocker,
    Finding,
    Hypothesis,
    Question,
    TaskProgress,
    TaskBrief,
    WorkItem,
    append_task_progress_revision,
    current_task_progress,
    initialize_task_progress,
    new_task_protocol,
    replace_task_brief,
    rewind_task_progress,
    task_contract_ref,
    task_progress_readiness_error,
)
from cyber_agent.core.task_protocol_service import TaskProtocolService
from cyber_agent.tools.builtin.task_protocol import submit_task_report
from cyber_agent.tools.builtin.subagent import _load_or_create_task_report
from cyber_agent.core.validation import (
    ValidationPolicy,
    ValidatorSettings,
    persist_validation_policy,
)
from cyber_agent.trace.api import get_task_progress, set_trace_store
from cyber_agent.trace.models import Message, Trace
from cyber_agent.trace.store import FileSystemTraceStore


ANCHOR = {
    "objective": "produce a verified result",
    "completion_criteria": ["the result is verified"],
    "constraints": ["cite evidence"],
}


def ready_progress() -> TaskProgress:
    return TaskProgress(
        phase="ready_to_submit",
        questions=[Question(item_id="q1", text="What is true?", state="answered", answer="12%")],
        findings=[Finding(item_id="f1", statement="The value is 12%", basis="official evidence")],
        hypotheses=[Hypothesis(item_id="h1", statement="The value is 12%", state="supported", rationale="official evidence")],
        work_items=[WorkItem(item_id="w1", description="verify value", state="done", result_summary="verified")],
        decision_rationale="Use the supported value.",
    )


class TaskProgressModelTest(unittest.TestCase):
    def test_strict_models_reject_duplicate_ids_and_forbidden_fields(self):
        with self.assertRaises(ValidationError):
            TaskProgress.model_validate({"phase": "completed"})
        with self.assertRaises(ValidationError):
            TaskProgress.model_validate({"phase": "planning", "target": "copy"})
        with self.assertRaisesRegex(ValidationError, "globally unique"):
            TaskProgress(
                phase="planning",
                questions=[Question(item_id="same", text="q")],
                blockers=[Blocker(item_id="same", text="b")],
            )

    def test_item_state_requires_its_structured_result(self):
        with self.assertRaisesRegex(ValidationError, "requires answer"):
            Question(item_id="q", text="q", state="answered")
        with self.assertRaisesRegex(ValidationError, "requires resolution"):
            Blocker(item_id="b", text="b", state="resolved")
        with self.assertRaisesRegex(ValidationError, "requires rationale"):
            Hypothesis(item_id="h", statement="h", state="supported")
        with self.assertRaisesRegex(ValidationError, "requires result_summary"):
            WorkItem(item_id="w", description="w", state="done")

    def test_readiness_is_derived_without_copying_terminal_status(self):
        state = new_task_protocol()
        initialize_task_progress(state, root_task_anchor_hash="a" * 64)
        self.assertIn("ready_to_submit", task_progress_readiness_error(state))
        append_task_progress_revision(
            state,
            ready_progress(),
            expected_revision=1,
            effective_at_sequence=3,
            contract_ref=task_contract_ref(
                state,
                root_task_anchor_hash="a" * 64,
            ),
        )
        self.assertIsNone(task_progress_readiness_error(state))

    def test_rewind_follows_ancestry_not_largest_sequence(self):
        state = new_task_protocol()
        initialize_task_progress(state, root_task_anchor_hash="a" * 64)
        contract = task_contract_ref(state, root_task_anchor_hash="a" * 64)
        r2 = append_task_progress_revision(
            state,
            TaskProgress(phase="investigating"),
            expected_revision=1,
            effective_at_sequence=6,
            contract_ref=contract,
        )
        self.assertEqual(2, r2.revision)
        state["task_report_progress_revision"] = r2.revision
        rewind_task_progress(state, 4)
        self.assertIsNone(state["task_report_progress_revision"])
        r3 = append_task_progress_revision(
            state,
            TaskProgress(phase="executing"),
            expected_revision=1,
            effective_at_sequence=10,
            contract_ref=contract,
        )
        self.assertEqual(3, r3.revision)
        self.assertEqual(1, r3.parent_revision)
        rewind_task_progress(state, 7)
        self.assertEqual(1, current_task_progress(state).revision)
        self.assertEqual(3, len(state["task_progress_revisions"]))

    def test_new_brief_starts_empty_planning_revision_without_erasing_history(self):
        first_brief = TaskBrief(
            objective="check the first claim",
            reason="the parent needs evidence",
            completion_criteria=["reach a checked conclusion"],
            expected_outputs=["an evidence note"],
        )
        state = new_task_protocol(first_brief)
        initialize_task_progress(state, root_task_anchor_hash="a" * 64)
        append_task_progress_revision(
            state,
            TaskProgress(
                phase="investigating",
                findings=[Finding(
                    item_id="old-finding",
                    statement="old claim",
                    basis="old evidence",
                )],
            ),
            expected_revision=1,
            effective_at_sequence=4,
            contract_ref=task_contract_ref(state),
        )

        replace_task_brief(
            state,
            first_brief.model_copy(update={
                "objective": "check the corrected claim",
            }),
            effective_at_sequence=8,
        )

        current = current_task_progress(state)
        self.assertEqual(3, current.revision)
        self.assertEqual(2, current.parent_revision)
        self.assertEqual("planning", current.phase)
        self.assertEqual([], current.findings)
        self.assertEqual(2, current.task_contract_ref.task_brief_version)
        self.assertEqual(3, len(state["task_progress_revisions"]))

    def test_validation_plan_hash_binds_exact_progress_revision(self):
        state = new_task_protocol()
        first = initialize_task_progress(
            state,
            root_task_anchor_hash="a" * 64,
        )
        second = append_task_progress_revision(
            state,
            TaskProgress(phase="investigating"),
            expected_revision=1,
            effective_at_sequence=2,
            contract_ref=task_contract_ref(
                state,
                root_task_anchor_hash="a" * 64,
            ),
        )
        policy = ValidationPolicy()
        common = {
            "task_brief": None,
            "task_brief_version": 0,
            "root_task_anchor": ANCHOR,
            "task_report": None,
            "candidate_output": "candidate",
            "evaluated_head_sequence": 3,
            "materials": [],
            "material_issues": [],
            "model_by_scope": {"root": "fake"},
            "root": True,
        }

        first_plan = policy.compile_plan(task_progress=first, **common)
        second_plan = policy.compile_plan(task_progress=second, **common)

        self.assertNotEqual(first_plan.plan_hash, second_plan.plan_hash)
        self.assertEqual(1, first_plan.task_progress_revision)
        self.assertEqual(2, second_plan.task_progress_revision)


class TaskProtocolServiceTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.store = FileSystemTraceStore(self.temp_dir.name)
        context = {
            "agent_mode": "recursive",
            "agent_mode_revision": 3,
            "root_trace_id": "root",
            "task_protocol": new_task_protocol(),
        }
        anchor = normalize_root_task_anchor(ANCHOR)
        persist_root_task_anchor(context, anchor)
        initialize_task_progress(
            context["task_protocol"],
            root_task_anchor_hash=context["root_task_anchor_hash"],
        )
        replace_context_access(
            context,
            [],
            root_task_anchor=anchor,
            task_brief=None,
        )
        await self.store.create_trace(Trace(
            trace_id="root",
            mode="agent",
            task="test",
            uid="user-1",
            context=context,
        ))
        self.service = TaskProtocolService(self.store)

    async def prepare_resume_contract(self) -> tuple[AgentRunner, RunConfig]:
        trace = await self.store.get_trace("root")
        budget = ResourceBudget()
        trace.context[RESOURCE_BUDGET_CONTEXT_KEY] = budget.to_dict()
        persist_validation_policy(
            trace.context,
            ValidationPolicy(),
            ValidatorSettings(),
        )
        config = RunConfig(
            trace_id="root",
            model="fake",
            uid="user-1",
            tools=[],
            tool_groups=[],
            enable_research_flow=False,
        )
        persist_run_config_snapshot(
            trace.context,
            RunConfigSnapshotV1.from_run_config(
                config,
                memory_identity=None,
            ),
        )
        await self.store.update_trace(
            "root",
            context=trace.context,
            status="stopped",
        )
        await ResourceBudgetController(self.store).initialize("root", budget)
        return (
            AgentRunner(trace_store=self.store, llm_call=lambda **_: None),
            config,
        )

    async def asyncTearDown(self):
        self.temp_dir.cleanup()

    async def test_concurrent_same_revision_has_one_winner(self):
        results = await asyncio.gather(
            self.service.update_progress(
                "root",
                expected_revision=1,
                progress=TaskProgress(phase="investigating"),
                effective_at_sequence=2,
            ),
            self.service.update_progress(
                "root",
                expected_revision=1,
                progress=TaskProgress(phase="executing"),
                effective_at_sequence=2,
            ),
            return_exceptions=True,
        )
        self.assertEqual(1, sum(not isinstance(item, Exception) for item in results))
        self.assertEqual(1, sum("stale" in str(item) for item in results if isinstance(item, Exception)))
        trace = await self.store.get_trace("root")
        state = trace.context["task_protocol"]
        self.assertEqual(2, len(state["task_progress_revisions"]))

    async def test_progress_and_report_share_one_atomic_command_boundary(self):
        root = await self.store.get_trace("root")
        brief = TaskBrief(
            objective="prepare one partial result",
            reason="exercise concurrent protocol commands",
            completion_criteria=["record current state"],
            expected_outputs=["partial report"],
        )
        child_context = deepcopy(root.context)
        child_context["task_protocol"] = new_task_protocol(brief)
        initialize_task_progress(
            child_context["task_protocol"],
            effective_at_sequence=1,
        )
        await self.store.create_trace(Trace(
            trace_id="child",
            mode="agent",
            task="child",
            uid=root.uid,
            parent_trace_id="root",
            context=child_context,
        ))

        progress_result, report_result = await asyncio.gather(
            self.service.update_progress(
                "child",
                expected_revision=1,
                progress=TaskProgress(
                    phase="investigating",
                    decision_rationale="recorded before partial submission",
                ),
                effective_at_sequence=2,
            ),
            submit_task_report(
                task_report={
                    "summary": "partial result",
                    "outcome": "partial",
                    "validation": {
                        "hard_passed": False,
                        "open_issues": ["work remains"],
                    },
                    "next_step_suggestion": {
                        "direction": "NONE",
                        "reason": "parent should decide after partial submission",
                    },
                    "remaining_issues": ["work remains"],
                },
                context={
                    "store": self.store,
                    "task_protocol_service": self.service,
                    "trace_id": "child",
                    "sequence": 3,
                },
            ),
            return_exceptions=True,
        )
        persisted = await self.store.get_trace("child")
        state = persisted.context["task_protocol"]
        if isinstance(progress_result, Exception):
            self.assertEqual(1, state["task_progress_head_revision"])
        else:
            self.assertEqual(2, state["task_progress_head_revision"])
            self.assertEqual(2, len(state["task_progress_revisions"]))
        if not isinstance(report_result, Exception):
            self.assertEqual("completed", report_result["status"], report_result)
            self.assertEqual(
                state["task_progress_head_revision"],
                state["task_report_progress_revision"],
            )

    async def test_progress_and_framework_report_cannot_overwrite_each_other(self):
        root = await self.store.get_trace("root")
        child_context = deepcopy(root.context)
        child_context["task_protocol"] = new_task_protocol(TaskBrief(
            objective="record a recoverable result",
            reason="exercise framework report convergence",
            completion_criteria=["preserve the latest progress revision"],
            expected_outputs=["one report"],
        ))
        initialize_task_progress(
            child_context["task_protocol"],
            effective_at_sequence=1,
        )
        await self.store.create_trace(Trace(
            trace_id="framework-child",
            mode="agent",
            uid=root.uid,
            parent_trace_id="root",
            context=child_context,
        ))

        progress_result, report = await asyncio.gather(
            self.service.update_progress(
                "framework-child",
                expected_revision=1,
                progress=TaskProgress(
                    phase="investigating",
                    decision_rationale="must not be lost by framework recovery",
                ),
                effective_at_sequence=2,
            ),
            _load_or_create_task_report(
                self.store,
                "framework-child",
                "simulated child interruption",
                "failed",
                3,
                task_protocol_service=self.service,
            ),
            return_exceptions=True,
        )
        persisted = await self.store.get_trace("framework-child")
        state = persisted.context["task_protocol"]
        self.assertEqual("protocol_error", report.outcome)
        if isinstance(progress_result, Exception):
            self.assertEqual(1, state["task_progress_head_revision"])
        else:
            self.assertEqual(2, state["task_progress_head_revision"])
            self.assertEqual(2, len(state["task_progress_revisions"]))
        self.assertEqual(
            state["task_progress_head_revision"],
            state["task_report_progress_revision"],
        )

    async def test_read_projection_hides_other_protocol_state(self):
        set_trace_store(self.store)
        response = await get_task_progress("root", include_history=False)
        self.assertEqual(3, response.protocol_revision)
        self.assertEqual(1, response.revision_count)
        self.assertIsNone(response.revisions)
        self.assertNotIn("pending_reviews", response.model_dump())

    async def test_progress_rejects_context_ref_from_another_owner(self):
        trace = await self.store.get_trace("root")
        anchor = normalize_root_task_anchor(ANCHOR)
        snapshot = create_context_snapshot(
            kind="task_brief",
            summary="authorized evidence",
            source_trace_id="source",
            root_trace_id="root",
            uid="user-1",
            content={"claim": "12%"},
            granted_at_sequence=1,
        )
        ref = add_context_snapshot(
            trace.context,
            snapshot,
            root_task_anchor=anchor,
            task_brief=None,
        )
        await self.store.update_trace(
            "root",
            context=trace.context,
            uid="another-user",
        )

        with self.assertRaisesRegex(Exception, "another tree or owner"):
            await self.service.update_progress(
                "root",
                expected_revision=1,
                progress=TaskProgress(
                    phase="investigating",
                    findings=[Finding(
                        item_id="finding",
                        statement="The value is 12%",
                        basis="authorized evidence",
                        context_refs=[ref],
                    )],
                ),
                effective_at_sequence=2,
            )
        persisted = await self.store.get_trace("root")
        self.assertEqual(
            1,
            persisted.context["task_protocol"]["task_progress_head_revision"],
        )

    async def test_resume_rolls_back_progress_without_persisted_tool_result(self):
        runner, config = await self.prepare_resume_contract()
        await self.store.add_message(Message.create(
            trace_id="root",
            role="assistant",
            sequence=1,
            content={"tool_calls": [{"function": {"name": "update_task_progress"}}]},
        ))
        await self.service.update_progress(
            "root",
            expected_revision=1,
            progress=TaskProgress(phase="investigating"),
            effective_at_sequence=2,
        )

        await runner._prepare_existing_trace(config)

        persisted = await self.store.get_trace("root")
        state = persisted.context["task_protocol"]
        self.assertEqual(1, state["task_progress_head_revision"])
        self.assertEqual(2, len(state["task_progress_revisions"]))

    async def test_revision_three_resume_rejects_missing_or_partial_protocol(self):
        runner, config = await self.prepare_resume_contract()
        trace = await self.store.get_trace("root")
        trace.context.pop("task_protocol")
        await self.store.update_trace("root", context=trace.context)

        with self.assertRaisesRegex(ValueError, "task_protocol is missing"):
            await runner._prepare_existing_trace(config)
        persisted = await self.store.get_trace("root")
        self.assertNotIn("task_protocol", persisted.context)
        self.assertEqual("stopped", persisted.status)

        trace.context["task_protocol"] = new_task_protocol()
        trace.context["task_protocol"].pop("task_progress_revisions")
        await self.store.update_trace("root", context=trace.context)
        with self.assertRaisesRegex(ValueError, "missing required fields"):
            await runner._prepare_existing_trace(config)

    async def test_resume_rejects_report_bound_to_abandoned_progress_branch(self):
        runner, config = await self.prepare_resume_contract()
        trace = await self.store.get_trace("root")
        state = trace.context["task_protocol"]
        contract = task_contract_ref(
            state,
            root_task_anchor_hash=trace.context["root_task_anchor_hash"],
        )
        abandoned = append_task_progress_revision(
            state,
            TaskProgress(phase="investigating"),
            expected_revision=1,
            effective_at_sequence=1,
            contract_ref=contract,
        )
        rewind_task_progress(state, 0)
        active = append_task_progress_revision(
            state,
            TaskProgress(phase="executing"),
            expected_revision=1,
            effective_at_sequence=1,
            contract_ref=contract,
        )
        state["task_report"] = {"summary": "corrupt binding"}
        state["task_report_progress_revision"] = abandoned.revision
        await self.store.update_trace("root", context=trace.context)
        await self.store.add_message(Message.create(
            trace_id="root",
            role="assistant",
            sequence=1,
            content="persisted active branch",
        ))

        with self.assertRaisesRegex(ValueError, "outside the current revision ancestry"):
            await runner._prepare_existing_trace(config)
        self.assertEqual(3, active.revision)


if __name__ == "__main__":
    unittest.main()