Agent/tests/test_application_reference.py

from copy import deepcopy
import asyncio
import json
import tempfile
import unittest
from unittest.mock import patch

from cyber_agent.application import (
    ApplicationRegistry,
    ApplicationRuntime,
    CandidateReviewAction,
)
from cyber_agent.application.candidate import (
    CandidateLedger,
    CandidatePointer,
    CandidateRef,
)
from cyber_agent.core.artifacts import ArtifactRef
from cyber_agent.core.run_snapshot import (
    RunConfigSnapshotV2,
    persist_run_config_snapshot,
)
from cyber_agent.core.runner import RunConfig
from cyber_agent.tools.builtin.knowledge import KnowledgeConfig
from cyber_agent.core.task_protocol import (
    Finding,
    Hypothesis,
    Question,
    TaskBrief,
    TaskProgress,
    WorkItem,
    initialize_task_progress,
    new_task_protocol,
)
from cyber_agent.trace.models import Message, Trace
from cyber_agent.trace.store import FileSystemTraceStore
from cyber_agent.tools.builtin.subagent import _record_pending_task_reports
from examples.application_reference import build_reference_components


def _tool_call(call_id, name, arguments):
    return {
        "content": "",
        "tool_calls": [{
            "id": call_id,
            "type": "function",
            "function": {
                "name": name,
                "arguments": json.dumps(arguments),
            },
        }],
        "finish_reason": "tool_calls",
    }


def _has_tool_result(messages, name):
    return any(
        item.get("role") == "tool" and item.get("name") == name
        for item in messages
    )


def _tool_names(schemas):
    return {item["function"]["name"] for item in schemas or []}


class ApplicationReferenceEndToEndTest(unittest.IsolatedAsyncioTestCase):
    async def test_real_agent_recovers_orphaned_adoption_review_after_restart(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            store = FileSystemTraceStore(temp_dir)
            components = build_reference_components()
            registry = ApplicationRegistry()
            registry.register(components.application, components.services)
            delegated = False
            reviewed = False

            async def llm_call(**kwargs):
                nonlocal delegated, reviewed
                messages = kwargs["messages"]
                if any(
                    "recursive_validation_protocol" in str(item.get("content", ""))
                    for item in messages
                    if item.get("role") == "system"
                ):
                    packet = next(
                        json.loads(item["content"])
                        for item in messages
                        if item.get("role") == "user"
                        and "recursive_validation_protocol"
                        in str(item.get("content", ""))
                    )
                    scope = packet["validation_scope"]
                    return {
                        "content": json.dumps({
                            "scope": scope,
                            "outcome": "passed",
                            "checks": [{
                                "check_id": item["check_id"],
                                "status": "passed",
                                "evidence_refs": [],
                                "issue": None,
                            } for item in packet["validation_plan"]["checks"]
                            if item["scope"] == scope],
                            "reason": "reference output is complete",
                            "retry_from": None,
                        }),
                        "tool_calls": [],
                        "finish_reason": "stop",
                    }
                system = "\n".join(
                    str(item.get("content", ""))
                    for item in messages
                    if item.get("role") == "system"
                )
                names = _tool_names(kwargs.get("tools"))
                if "Create replaceable candidate versions" in system:
                    if not _has_tool_result(messages, "manage_candidate"):
                        self.assertIn("manage_candidate", names)
                        self.assertIn("submit_task_report", names)
                        return _tool_call("candidate-create", "manage_candidate", {
                            "request": {
                                "operation": "create",
                                "content": {"text": "A finished candidate."},
                                "parent_refs": [],
                            },
                        })
                    candidate_message = next(
                        item for item in reversed(messages)
                        if item.get("role") == "tool"
                        and item.get("name") == "manage_candidate"
                    )
                    candidate_ref = json.loads(
                        candidate_message["content"]
                    )["candidate_ref"]
                    if not _has_tool_result(messages, "update_task_progress"):
                        return _tool_call("writer-progress", "update_task_progress", {
                            "expected_revision": 1,
                            "progress": {
                                "phase": "ready_to_submit",
                                "questions": [],
                                "blockers": [],
                                "findings": [],
                                "hypotheses": [],
                                "work_items": [],
                                "decision_rationale": "The exact candidate is registered.",
                            },
                        })
                    if not _has_tool_result(messages, "submit_task_report"):
                        return _tool_call("writer-report", "submit_task_report", {
                            "task_report": {
                                "summary": "One complete candidate is ready.",
                                "outcome": "satisfied",
                                "validation": {
                                    "hard_passed": True,
                                    "open_issues": [],
                                },
                                "next_step_suggestion": {
                                    "direction": "NONE",
                                    "reason": "Parent may adopt the exact revision.",
                                },
                                "outputs": [],
                                "evidence": [],
                                "remaining_issues": [],
                                "candidate_refs": [candidate_ref],
                            },
                        })
                    return {"content": "writer submitted", "tool_calls": []}
                if names == {"review_task_result"}:
                    root = next(
                        item for item in await store.list_traces(limit=20)
                        if item.parent_trace_id is None
                    )
                    child_id, pending = next(iter(
                        root.context["task_protocol"]["pending_reviews"].items()
                    ))
                    candidate_ref = pending["task_report"]["candidate_refs"][0]
                    reviewed = True
                    return _tool_call("adopt-review", "review_task_result", {
                        "child_trace_id": child_id,
                        "decision": "ASCEND",
                        "reason": "Adopt the validated exact candidate revision.",
                        "candidate_actions": [{
                            "action": "adopt",
                            "candidate_ref": candidate_ref,
                            "reason": "The candidate passed independent validation.",
                        }],
                    })
                if not delegated:
                    delegated = True
                    return _tool_call("delegate-writer", "agent", {
                        "agent_type": "writer",
                        "task_brief": {
                            "objective": "Produce one finished candidate",
                            "reason": "The editor needs an adoptable output",
                            "completion_criteria": ["No placeholder remains"],
                            "expected_outputs": ["One candidate revision"],
                            "validation_scopes": ["output"],
                        },
                    })
                if reviewed and not _has_tool_result(
                    messages,
                    "update_task_progress",
                ):
                    return _tool_call("root-progress", "update_task_progress", {
                        "expected_revision": 1,
                        "progress": {
                            "phase": "ready_to_submit",
                            "questions": [],
                            "blockers": [],
                            "findings": [],
                            "hypotheses": [],
                            "work_items": [],
                            "decision_rationale": "The validated candidate was adopted.",
                        },
                    })
                return {"content": "Final adopted content", "tool_calls": []}

            runtime = ApplicationRuntime(
                registry=registry,
                trace_store=store,
                llm_call=llm_call,
            )
            runner, config = runtime.new_run(
                "application_reference",
                "1",
                uid="reference-user",
                root_task_anchor={
                    "objective": "Create one complete output",
                    "completion_criteria": ["Adopt one validated candidate"],
                    "constraints": ["No placeholders"],
                },
            )
            config.knowledge = KnowledgeConfig(
                enable_extraction=False,
                enable_completion_extraction=False,
                enable_injection=False,
            )
            original_update_trace = store.update_trace
            interrupted = False

            async def interrupt_after_adoption(trace_id, **updates):
                nonlocal interrupted
                context = updates.get("context")
                reviews = (
                    context.get("task_protocol", {}).get("reviews", [])
                    if isinstance(context, dict)
                    else []
                )
                if (
                    not interrupted
                    and any(
                        item.get("review_command_id") == "adopt-review"
                        for item in reviews
                    )
                ):
                    interrupted = True
                    raise asyncio.CancelledError()
                return await original_update_trace(trace_id, **updates)

            with patch.object(
                store,
                "update_trace",
                side_effect=interrupt_after_adoption,
            ), patch.dict(
                "os.environ",
                {"AGENT_MODE": "recursive"},
                clear=False,
            ):
                with self.assertRaises(asyncio.CancelledError):
                    await runner.run_result(
                        [{"role": "user", "content": "Create the output"}],
                        config,
                    )
            self.assertTrue(interrupted)
            roots = [
                item for item in await store.list_traces(limit=20)
                if item.parent_trace_id is None
            ]
            self.assertEqual(1, len(roots))
            root = roots[0]
            self.assertEqual(1, components.candidates.adoption_calls)
            self.assertTrue(
                root.context["task_protocol"]["pending_reviews"]
            )
            orphaned = [
                item for item in await store.get_trace_messages(root.trace_id)
                if item.role == "assistant"
                and isinstance(item.content, dict)
                and any(
                    call.get("id") == "adopt-review"
                    for call in item.content.get("tool_calls", [])
                )
            ]
            self.assertEqual(1, len(orphaned))
            self.assertFalse(any(
                item.role == "tool" and item.tool_call_id == "adopt-review"
                for item in await store.get_trace_messages(root.trace_id)
            ))

            restarted_registry = ApplicationRegistry()
            restarted_registry.register(
                components.application,
                components.services,
            )
            restarted_runtime = ApplicationRuntime(
                registry=restarted_registry,
                trace_store=store,
                llm_call=llm_call,
            )
            restored_runner, restored_config = await restarted_runtime.restore(
                root.trace_id
            )
            with patch.dict("os.environ", {"AGENT_MODE": "recursive"}, clear=False):
                result = await restored_runner.run_result([], restored_config)
            self.assertEqual("completed", result["status"])
            self.assertTrue(reviewed)
            self.assertEqual(1, len(components.candidates.versions))
            self.assertEqual(1, components.candidates.adoption_calls)
            root = await store.get_trace(result["trace_id"])
            ledger = CandidateLedger.model_validate(
                await store.get_candidate_ledger(root.trace_id)
            )
            self.assertEqual("adopted", ledger.current_state(ledger.candidates[0]))
            self.assertIn(
                "candidate-create",
                {item.command_id for item in ledger.operations},
            )
            self.assertIn(
                "adopt-review",
                {item.command_id for item in ledger.operations},
            )
            self.assertEqual(1, len([
                item for item in await store.get_trace_messages(root.trace_id)
                if item.role == "tool" and item.tool_call_id == "adopt-review"
            ]))
            self.assertEqual(
                1,
                len(root.context["task_protocol"]["reviews"]),
            )

    async def test_local_candidate_retry_preserves_fact_and_peer_validation(self):
        with tempfile.TemporaryDirectory() as temp_dir:
            store = FileSystemTraceStore(temp_dir)
            components = build_reference_components()
            registry = ApplicationRegistry()
            binding = registry.register(
                components.application,
                components.services,
            )
            llm_calls = 0

            async def llm_call(**kwargs):
                nonlocal llm_calls
                llm_calls += 1
                packet = json.loads(kwargs["messages"][-1]["content"])
                scope = packet["validation_scope"]
                return {
                    "content": json.dumps({
                        "scope": scope,
                        "outcome": "passed",
                        "checks": [
                            {
                                "check_id": item["check_id"],
                                "status": "passed",
                                "evidence_refs": [],
                                "issue": None,
                            }
                            for item in packet["validation_plan"]["checks"]
                        ],
                        "reason": "reference candidate is complete",
                        "retry_from": None,
                    }),
                    "tool_calls": [],
                }

            runtime = ApplicationRuntime(
                registry=registry,
                trace_store=store,
                llm_call=llm_call,
            )
            runner, config = runtime.new_run(
                "application_reference",
                "1",
                uid="reference-user",
                root_task_anchor={
                    "objective": "Create an evidence-grounded explanation",
                    "completion_criteria": ["Use verified facts and finished copy"],
                    "constraints": ["Do not publish placeholder text"],
                },
            )
            with patch.dict("os.environ", {"AGENT_MODE": "recursive"}, clear=False):
                root, _goal, _sequence = await runner._prepare_new_trace(
                    [{"role": "user", "content": "write the explanation"}],
                    config,
                )
            await store.add_message(Message.create(
                trace_id=root.trace_id,
                role="user",
                sequence=1,
                content="write the explanation",
            ))
            await store.add_message(Message.create(
                trace_id=root.trace_id,
                role="assistant",
                sequence=2,
                parent_sequence=1,
                content="drafting",
            ))
            await store.update_trace(root.trace_id, head_sequence=2)

            async def execute_tool(trace_id, tool_name, arguments, sequence):
                tool_call_id = f"reference-{tool_name}-{sequence}"
                result = await binding.tool_registry.execute(
                    tool_name,
                    arguments,
                    uid=root.uid,
                    context={
                        "store": store,
                        "trace_id": trace_id,
                        "sequence": sequence,
                        "tool_call_id": tool_call_id,
                        "task_protocol_service": runner.task_protocol_service,
                        "candidate_service": runner.candidate_service,
                        "event_service": runner.event_service,
                    },
                    allowed_tool_names={tool_name},
                )
                if tool_name == "manage_candidate":
                    trace = await store.get_trace(trace_id)
                    await store.add_message(Message.create(
                        trace_id=trace_id,
                        role="tool",
                        sequence=sequence,
                        parent_sequence=(trace.head_sequence or None),
                        tool_call_id=tool_call_id,
                        content={"tool_name": tool_name, "result": "registered"},
                    ))
                    await store.update_trace(trace_id, head_sequence=sequence)
                if isinstance(result, str):
                    return json.loads(result)
                text = result.get("text") if isinstance(result, dict) else None
                return json.loads(text) if isinstance(text, str) else result

            async def create_child(trace_id, role_id, brief):
                context = deepcopy(root.context)
                context.pop("run_config_snapshot", None)
                context["application_role_id"] = role_id
                context["application_role_hash"] = binding.role(role_id).role_hash
                context["agent_depth"] = 1
                state = new_task_protocol(brief)
                initialize_task_progress(state, effective_at_sequence=1)
                context["task_protocol"] = state
                child_config = RunConfig(
                    trace_id=trace_id,
                    uid=root.uid,
                    parent_trace_id=root.trace_id,
                    context=context,
                    enable_research_flow=False,
                )
                binding.configure_run_config(child_config, role_id)
                context["effective_run_limits"] = dict(
                    child_config.effective_run_limits
                )
                child_config.context = context
                persist_run_config_snapshot(
                    context,
                    RunConfigSnapshotV2.from_run_config(
                        child_config,
                        memory_identity=None,
                    ),
                )
                child = Trace(
                    trace_id=trace_id,
                    mode="agent",
                    agent_type=role_id,
                    uid=root.uid,
                    model=binding.role(role_id).role.model,
                    parent_trace_id=root.trace_id,
                    context=context,
                )
                await store.create_trace(child)
                return child

            fact_child = await create_child(
                "reference-fact",
                "fact_checker",
                TaskBrief(
                    objective="Verify what an agent framework supplies",
                    reason="The content needs one traceable fact",
                    completion_criteria=["Return one sourced finding"],
                    expected_outputs=["fact artifact"],
                    validation_scopes=["evidence"],
                ),
            )
            raw_fact_result = await binding.tool_registry.execute(
                "verify_fact",
                {"question": "What does an agent framework supply?"},
                uid=root.uid,
                context={
                    "trace_id": fact_child.trace_id,
                    "store": store,
                },
                allowed_tool_names={"verify_fact"},
            )
            fact_ref = ArtifactRef.model_validate(
                json.loads(raw_fact_result)["artifact_ref"]
            )
            await runner.task_protocol_service.update_progress(
                fact_child.trace_id,
                expected_revision=1,
                progress=TaskProgress(
                    phase="ready_to_submit",
                    questions=[Question(
                        item_id="fact-question",
                        text="What does the framework supply?",
                        state="answered",
                        answer="Reusable execution contracts",
                        artifact_refs=[fact_ref],
                    )],
                    findings=[Finding(
                        item_id="fact-finding",
                        statement="The framework supplies reusable execution contracts.",
                        basis="reference fixture",
                        artifact_refs=[fact_ref],
                    )],
                    hypotheses=[Hypothesis(
                        item_id="fact-hypothesis",
                        statement="The explanation can distinguish framework from business code.",
                        state="supported",
                        rationale="The verified definition supports the distinction.",
                        artifact_refs=[fact_ref],
                    )],
                    work_items=[WorkItem(
                        item_id="fact-work",
                        description="verify the definition",
                        state="done",
                        result_summary="one fact artifact recorded",
                        artifact_refs=[fact_ref],
                    )],
                    decision_rationale="The question is answered by one stable fixture.",
                ),
                effective_at_sequence=2,
            )

            writer = await create_child(
                "reference-writer",
                "writer",
                TaskBrief(
                    objective="Produce two candidate explanations",
                    reason="The editor needs alternatives",
                    completion_criteria=["Both candidates use the verified fact"],
                    expected_outputs=["candidate A", "candidate B"],
                    validation_scopes=["output"],
                ),
            )
            candidate_a = CandidateRef.model_validate((await execute_tool(
                writer.trace_id,
                "manage_candidate",
                {"request": {
                    "operation": "create",
                    "content": {
                        "text": "A finished explanation of reusable contracts."
                    },
                    "parent_refs": [],
                }},
                4,
            ))["candidate_ref"])
            candidate_b = CandidateRef.model_validate((await execute_tool(
                writer.trace_id,
                "manage_candidate",
                {"request": {
                    "operation": "create",
                    "content": {"text": "B says {{placeholder}}."},
                    "parent_refs": [],
                }},
                5,
            ))["candidate_ref"])
            await runner.task_protocol_service.update_progress(
                writer.trace_id,
                expected_revision=1,
                progress=TaskProgress(
                    phase="ready_to_submit",
                    findings=[Finding(
                        item_id="writer-finding",
                        statement="Both candidates derive from the verified definition.",
                        basis="fact artifact",
                        artifact_refs=[fact_ref],
                    )],
                    work_items=[WorkItem(
                        item_id="writer-work",
                        description="draft two alternatives",
                        state="done",
                        result_summary="A and B registered",
                    )],
                    decision_rationale="Both candidate revisions are ready for validation.",
                ),
                effective_at_sequence=6,
            )

            validation_a = await runner.validate_recursive_trace(
                writer.trace_id,
                candidate_ref=candidate_a,
            )
            validation_b = await runner.validate_recursive_trace(
                writer.trace_id,
                candidate_ref=candidate_b,
            )
            self.assertEqual("passed", validation_a.result.outcome)
            self.assertEqual("failed", validation_b.result.outcome)
            self.assertEqual(1, components.facts.calls)
            self.assertEqual(
                1,
                binding.tool_registry.get_stats("verify_fact")[
                    "verify_fact"
                ]["call_count"],
            )
            self.assertEqual(1, llm_calls)
            frozen_a = validation_a.result.model_dump(mode="json")

            fresh_root = await store.get_trace(root.trace_id)
            fresh_writer = await store.get_trace(writer.trace_id)
            await runner._mark_trace_stopped(
                root.trace_id,
                fresh_root.head_sequence,
            )
            await runner._mark_trace_stopped(
                writer.trace_id,
                fresh_writer.head_sequence,
            )
            restarted_components = build_reference_components(
                artifacts=components.artifacts,
                candidates=components.candidates,
                projector=components.projector,
            )
            restarted_registry = ApplicationRegistry()
            restarted_registry.register(
                restarted_components.application,
                restarted_components.services,
            )
            restarted_runtime = ApplicationRuntime(
                registry=restarted_registry,
                trace_store=store,
                llm_call=llm_call,
            )
            restored_runner, restored_config = await restarted_runtime.restore(
                writer.trace_id
            )
            self.assertEqual("writer", restored_config.role_id)
            self.assertEqual(0, restarted_components.context.list_calls)
            self.assertEqual(0, restarted_components.context.resolve_calls)
            await restored_runner.candidate_service.apply_review_actions(
                root.trace_id,
                writer.trace_id,
                report_refs=[candidate_a, candidate_b],
                actions=[CandidateReviewAction(
                    action="revise",
                    candidate_ref=candidate_b,
                    reason="remove the placeholder without repeating fact research",
                )],
                effective_at_sequence=7,
            )
            runner = restored_runner
            binding = restarted_registry.resolve("application_reference", "1")
            candidate_b2 = CandidateRef.model_validate((await execute_tool(
                writer.trace_id,
                "manage_candidate",
                {"request": {
                    "operation": "fork",
                    "content": {
                        "text": (
                            "B2 is a finished explanation of reusable contracts."
                        )
                    },
                    "parent_refs": [CandidatePointer(
                        candidate_id=candidate_b.candidate_id,
                        revision=candidate_b.revision,
                    ).model_dump(mode="json")],
                }},
                8,
            ))["candidate_ref"])
            validation_b2 = await restored_runner.validate_recursive_trace(
                writer.trace_id,
                candidate_ref=candidate_b2,
            )
            self.assertEqual("passed", validation_b2.result.outcome)
            self.assertEqual(2, llm_calls)
            self.assertEqual(
                1,
                components.facts.calls + restarted_components.facts.calls,
            )
            self.assertEqual(
                frozen_a,
                (await restored_runner.validate_recursive_trace(
                    writer.trace_id,
                    candidate_ref=candidate_a,
                )).result.model_dump(mode="json"),
            )
            self.assertEqual(2, llm_calls)

            messages = await store.get_trace_messages(root.trace_id)
            cutoff = min(item.sequence for item in messages)
            await restored_runner._rewind(root.trace_id, cutoff, None)
            await restored_runner.candidate_service.apply_review_actions(
                root.trace_id,
                writer.trace_id,
                report_refs=[candidate_a, candidate_b2],
                actions=[CandidateReviewAction(
                    action="adopt",
                    candidate_ref=candidate_b2,
                    reason="publish the corrected exact revision",
                )],
                effective_at_sequence=20,
            )
            self.assertEqual(1, components.candidates.adoption_calls)
            self.assertEqual(1, len(components.candidates.adoptions))

            await restored_runner.event_service.try_pump(root.trace_id)
            before_replay = dict(components.projector.rows)
            await restarted_runtime.reconcile_events(root.trace_id)
            await restarted_runtime.reconcile_events(root.trace_id)
            self.assertEqual(before_replay, components.projector.rows)
            adopted_keys = [
                key for key in components.projector.rows
                if key.endswith(":adopted")
            ]
            self.assertEqual(1, len(adopted_keys))
            with self.assertRaisesRegex(ValueError, "committed candidate adoption"):
                await restored_runner._rewind(root.trace_id, cutoff, None)

            ledger = CandidateLedger.model_validate(
                await store.get_candidate_ledger(root.trace_id)
            )
            self.assertEqual(3, len(ledger.candidates))
            self.assertEqual(3, len(ledger.validations))
            self.assertEqual(
                [(candidate_b.candidate_id, candidate_b.revision)],
                [
                    (item.candidate_id, item.revision)
                    for item in candidate_b2.parent_refs
                ],
            )