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 ], )