|
|
@@ -0,0 +1,420 @@
|
|
|
+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()
|
SamLee