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.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 = 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"], ), ) fact_ref = await components.facts.verify( question="What does an agent framework supply?", root_trace_id=root.trace_id, uid=root.uid, ) 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, llm_calls) frozen_a = validation_a.result.model_dump(mode="json") restored_runner, _restored_config = await runtime.restore(root.trace_id) 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) 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 runtime.reconcile_events(root.trace_id) await 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 ], )