Agent/tests/test_application_reference.py

from copy import deepcopy
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
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.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 examples.application_reference import build_reference_components


class ApplicationReferenceEndToEndTest(unittest.IsolatedAsyncioTestCase):
    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 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 = await runner.candidate_service.manage(
                writer.trace_id,
                operation="create",
                content={"text": "A finished explanation of reusable contracts."},
                parent_refs=[],
                effective_at_sequence=4,
            )
            candidate_b = await runner.candidate_service.manage(
                writer.trace_id,
                operation="create",
                content={"text": "B says {{placeholder}}."},
                parent_refs=[],
                effective_at_sequence=5,
            )
            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")

            await runner._mark_trace_stopped(root.trace_id, root.head_sequence)
            await runner._mark_trace_stopped(writer.trace_id, 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,
            )
            candidate_b2 = await restored_runner.candidate_service.manage(
                writer.trace_id,
                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,
                )],
                effective_at_sequence=8,
            )
            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
                ],
            )