import asyncio 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.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_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_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()