import asyncio import json import tempfile import unittest from unittest.mock import patch from pydantic import ValidationError from cyber_agent.application import ( AgentApplication, AgentRole, ApplicationRegistry, ApplicationServices, ProviderRef, ) from cyber_agent.application.candidate import ( CandidateAdoptionReceipt, CandidateLedger, CandidatePointer, CandidateReviewAction, ) from cyber_agent.application.quality import CandidateValidationRecord from cyber_agent.application.candidate_service import ( CandidateService, CandidateStateError, ) from cyber_agent.core.artifacts import ( ArtifactRef, ValidationMaterial, material_content_hash, ) from cyber_agent.core.task_protocol import TaskBrief, new_task_protocol from cyber_agent.core.validation import ( ScopeValidationResult, ValidationCheck, ValidationResult, ) from cyber_agent.tools.registry import ToolRegistry from cyber_agent.trace.models import Trace from cyber_agent.trace.store import ( FileSystemTraceStore, TraceStoreCorruptionError, ) class MemoryCandidateRepository: def __init__(self): self.contents = {} self.put_calls = 0 self.merge_calls = 0 self.bad_artifact = False self.adoptions = {} self.adoption_calls = 0 async def put_version(self, request): self.put_calls += 1 return self._save(request) async def merge(self, request): self.merge_calls += 1 return self._save(request) def _save(self, request): key = request.operation_id self.contents.setdefault(key, request) persisted = self.contents[key] digest = material_content_hash(persisted.content) if self.bad_artifact: digest = "0" * 64 return ArtifactRef( artifact_id=( f"candidate:{persisted.candidate_id}:{persisted.revision}" ), version=str(persisted.revision), content_hash=digest, kind="candidate.output", mime_type="application/json", ) async def resolve(self, ref, root_trace_id, uid): request = next( item for item in self.contents.values() if ( ref.artifact_id == f"candidate:{item.candidate_id}:{item.revision}" ) ) return ValidationMaterial( **ref.model_dump(), root_trace_id=root_trace_id, uid=uid, content=request.content, ) async def load_version(self, candidate_ref): return next( item.content for item in self.contents.values() if item.candidate_id == candidate_ref.candidate_id and item.revision == candidate_ref.revision ) async def commit_adoption(self, request): self.adoption_calls += 1 self.adoptions.setdefault(request.operation_id, request.candidate_ref) return CandidateAdoptionReceipt( operation_id=request.operation_id, external_id=f"published:{request.candidate_ref.candidate_id}", ) def application(): return AgentApplication( application_id="candidate.test", application_version="1", root_role="writer", roles=(AgentRole( role_id="writer", model="fake", system_prompt="write", ),), artifact_resolver_ref=ProviderRef( provider_id="artifacts", provider_version="1", ), candidate_repository_ref=ProviderRef( provider_id="candidates", provider_version="1", ), ) class CandidateServiceTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.temp = tempfile.TemporaryDirectory() self.store = FileSystemTraceStore(self.temp.name) self.repository = MemoryCandidateRepository() self.binding = ApplicationRegistry().register( application(), ApplicationServices( tool_registry=ToolRegistry(), artifact_resolver=self.repository, candidate_repository=self.repository, ), ) self.service = CandidateService( store=self.store, application_binding=self.binding, repository=self.repository, artifact_resolver=self.repository, ) await self._create_root("root-a") async def asyncTearDown(self): self.temp.cleanup() def _context(self, root_trace_id, *, task_brief=None): return { "agent_mode": "recursive", "agent_mode_revision": 3, "root_trace_id": root_trace_id, "root_task_anchor_hash": "a" * 64, "application_ref": self.binding.application_ref.model_dump(mode="json"), "application_role_id": "writer", "application_role_hash": self.binding.role("writer").role_hash, "task_protocol": new_task_protocol(task_brief), } async def _create_root(self, trace_id, *, uid="user-1"): await self.store.create_trace(Trace( trace_id=trace_id, mode="agent", agent_type="writer", uid=uid, context=self._context(trace_id), )) async def _create_child(self, trace_id, root_trace_id="root-a"): await self.store.create_trace(Trace( trace_id=trace_id, mode="agent", agent_type="writer", uid="user-1", parent_trace_id=root_trace_id, context=self._context( root_trace_id, task_brief=TaskBrief( objective="write candidate", reason="the root needs an option", completion_criteria=["candidate is complete"], expected_outputs=["one candidate"], ), ), )) async def _record_validation(self, trace_id, candidate, outcome="passed"): plan_hash = material_content_hash({ "candidate": candidate.model_dump(mode="json"), "outcome": outcome, }) status = "passed" if outcome == "passed" else "failed" scope = ScopeValidationResult( validator_trace_id=f"{trace_id}@validator", scope="output", outcome=outcome, checks=[ValidationCheck( check_id="quality.rule", status=status, issue=None if status == "passed" else "needs revision", )], reason="checked", retry_from=None if outcome == "passed" else "output", plan_hash=plan_hash, ) result = ValidationResult( evaluated_trace_id=trace_id, outcome=outcome, scope_results=[scope], issues=[] if outcome == "passed" else ["needs revision"], retry_from=None if outcome == "passed" else "output", plan_hash=plan_hash, ) await self.service.record_validation( trace_id, CandidateValidationRecord( candidate_ref=candidate, plan_hash=plan_hash, validation_result=result.model_dump(mode="json"), validated_at_sequence=candidate.created_at_sequence, ), ) async def test_create_fork_merge_and_idempotent_recovery(self): first, duplicate = await asyncio.gather( self.service.manage( "root-a", operation="create", content={"text": "A"}, parent_refs=[], effective_at_sequence=1, ), self.service.manage( "root-a", operation="create", content={"text": "A"}, parent_refs=[], effective_at_sequence=1, ), ) self.assertEqual(first, duplicate) self.assertEqual(1, self.repository.put_calls) second = await self.service.manage( "root-a", operation="create", content={"text": "B"}, parent_refs=[], effective_at_sequence=2, ) fork = await self.service.manage( "root-a", operation="fork", content={"text": "A2"}, parent_refs=[CandidatePointer( candidate_id=first.candidate_id, revision=first.revision, )], effective_at_sequence=3, ) self.assertEqual(first.candidate_id, fork.candidate_id) self.assertEqual(2, fork.revision) merged = await self.service.manage( "root-a", operation="merge", content={"text": "AB"}, parent_refs=[ CandidatePointer( candidate_id=fork.candidate_id, revision=fork.revision, ), CandidatePointer( candidate_id=second.candidate_id, revision=second.revision, ), ], effective_at_sequence=4, ) self.assertNotIn( merged.candidate_id, {first.candidate_id, second.candidate_id}, ) self.assertEqual(2, len(merged.parent_refs)) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(4, len(ledger.candidates)) self.assertTrue(all( ledger.current_state(item) == "proposed" for item in ledger.candidates )) async def test_cross_root_and_cross_task_parent_access_is_rejected(self): first = await self.service.manage( "root-a", operation="create", content={"text": "A"}, parent_refs=[], effective_at_sequence=1, ) await self._create_root("root-b") with self.assertRaisesRegex(CandidateStateError, "not found"): await self.service.manage( "root-b", operation="fork", content={"text": "stolen"}, parent_refs=[CandidatePointer( candidate_id=first.candidate_id, revision=first.revision, )], effective_at_sequence=2, ) await self._create_child("child-a") with self.assertRaisesRegex(CandidateStateError, "another application, root, or Task"): await self.service.manage( "child-a", operation="fork", content={"text": "stolen"}, parent_refs=[CandidatePointer( candidate_id=first.candidate_id, revision=first.revision, )], effective_at_sequence=3, ) async def test_bad_artifact_fails_without_registering_candidate(self): self.repository.bad_artifact = True with self.assertRaisesRegex(CandidateStateError, "invalid ArtifactRef"): await self.service.manage( "root-a", operation="create", content={"text": "bad"}, parent_refs=[], effective_at_sequence=1, ) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(0, len(ledger.candidates)) self.assertEqual("failed", ledger.operations[0].status) async def test_report_validation_rejects_tamper_and_wrong_owner(self): candidate = await self.service.manage( "root-a", operation="create", content={"text": "A"}, parent_refs=[], effective_at_sequence=1, ) await self.service.validate_report_refs("root-a", [candidate]) tampered = candidate.model_copy(update={ "artifact_ref": candidate.artifact_ref.model_copy(update={ "version": "tampered", }), }) with self.assertRaisesRegex(CandidateStateError, "does not match"): await self.service.validate_report_refs("root-a", [tampered]) await self._create_child("child-a") with self.assertRaisesRegex(CandidateStateError, "another application, root, or Task"): await self.service.validate_report_refs("child-a", [candidate]) async def test_ledger_rejects_future_lineage_and_store_reports_corruption(self): candidate = await self.service.manage( "root-a", operation="create", content={"text": "A"}, parent_refs=[], effective_at_sequence=1, ) raw = (await self.store.get_candidate_ledger("root-a")) future = candidate.model_copy(update={ "revision": 2, "parent_refs": (CandidatePointer( candidate_id=candidate.candidate_id, revision=3, ),), }) raw["candidates"].append(future.model_dump(mode="json")) with self.assertRaisesRegex(ValidationError, "unknown parent"): CandidateLedger.model_validate(raw) ledger_path = self.store._get_candidate_ledger_file("root-a") ledger_path.write_text("{broken", encoding="utf-8") with self.assertRaisesRegex(TraceStoreCorruptionError, "candidate ledger"): await self.store.get_candidate_ledger("root-a") async def test_cyclic_lineage_is_rejected_before_repository_side_effect(self): first = await self.service.manage( "root-a", operation="create", content={"text": "A"}, parent_refs=[], effective_at_sequence=1, ) raw = await self.store.get_candidate_ledger("root-a") second = first.model_copy(update={ "candidate_id": "candidate-b", "parent_refs": (CandidatePointer( candidate_id=first.candidate_id, revision=first.revision, ),), }) raw["candidates"][0]["parent_refs"] = [{ "candidate_id": second.candidate_id, "revision": second.revision, }] raw["candidates"].append(second.model_dump(mode="json")) ledger_path = self.store._get_candidate_ledger_file("root-a") ledger_path.write_text(json.dumps(raw), encoding="utf-8") put_calls = self.repository.put_calls with self.assertRaisesRegex(ValidationError, "lineage cycle"): await self.service.manage( "root-a", operation="create", content={"text": "must not be persisted"}, parent_refs=[], effective_at_sequence=2, ) self.assertEqual(put_calls, self.repository.put_calls) async def test_review_adoption_is_idempotent_and_creates_rewind_barrier(self): await self._create_child("child-a") candidate = await self.service.manage( "child-a", operation="create", content={"text": "A"}, parent_refs=[], effective_at_sequence=4, ) await self._record_validation("child-a", candidate) action = CandidateReviewAction( action="adopt", candidate_ref=candidate, reason="selected exact revision", ) first = await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[action], effective_at_sequence=9, ) second = await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[action], effective_at_sequence=9, ) self.assertEqual(first, second) self.assertEqual(1, self.repository.adoption_calls) self.assertEqual(1, len(self.repository.adoptions)) with self.assertRaisesRegex(CandidateStateError, "committed"): await self.service.assert_rewind_allowed("root-a", 8) await self.service.assert_rewind_allowed("root-a", 9) with self.assertRaisesRegex(CandidateStateError, "terminal"): await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[CandidateReviewAction( action="discard", candidate_ref=candidate, reason="too late", )], effective_at_sequence=10, ) async def test_failed_candidate_can_only_be_revised_or_discarded(self): await self._create_child("child-a") candidate = await self.service.manage( "child-a", operation="create", content={"text": "placeholder"}, parent_refs=[], effective_at_sequence=4, ) await self._record_validation("child-a", candidate, outcome="failed") with self.assertRaisesRegex(CandidateStateError, "passed"): await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[CandidateReviewAction( action="adopt", candidate_ref=candidate, reason="invalid", )], effective_at_sequence=9, ) records = await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[CandidateReviewAction( action="revise", candidate_ref=candidate, reason="remove placeholder", )], effective_at_sequence=10, ) self.assertEqual("discarded", records[0].state) self.assertIn("retry_from=output", records[0].reason) async def test_adoption_replays_same_operation_after_finalize_crash(self): await self._create_child("child-a") candidate = await self.service.manage( "child-a", operation="create", content={"text": "A"}, parent_refs=[], effective_at_sequence=4, ) await self._record_validation("child-a", candidate) action = CandidateReviewAction( action="adopt", candidate_ref=candidate, reason="publish once", ) original_replace = self.store.replace_candidate_ledger calls = 0 async def fail_finalize(root_trace_id, ledger): nonlocal calls calls += 1 if calls == 2: raise OSError("crash before framework finalize") await original_replace(root_trace_id, ledger) with patch.object( self.store, "replace_candidate_ledger", side_effect=fail_finalize, ): with self.assertRaisesRegex(OSError, "finalize"): await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[action], effective_at_sequence=9, ) pending = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual("adoption_pending", pending.current_state(candidate)) await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[action], effective_at_sequence=9, ) self.assertEqual(2, self.repository.adoption_calls) self.assertEqual(1, len(self.repository.adoptions)) completed = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual("adopted", completed.current_state(candidate)) if __name__ == "__main__": unittest.main()