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, CandidateOperationRecord, CandidatePointer, CandidateRepositoryRejected, CandidateReviewAction, MAX_CANDIDATES_PER_ROOT, MAX_CANDIDATE_OPERATIONS_PER_ROOT, ) from cyber_agent.application.quality import CandidateValidationRecord from cyber_agent.application.candidate_service import ( CandidateService, CandidateStateError, ) from cyber_agent.core.runner import AgentRunner, RunConfig 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 import get_tool_registry from cyber_agent.tools.errors import RecoverableToolExecutionError from cyber_agent.tools.registry import ToolRegistry from cyber_agent.trace.models import Message as TraceMessage, 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_finalize_crash_retries_same_semantic_command_identity(self): original_replace = self.store.replace_candidate_ledger writes = 0 async def fail_completed_publish(root_trace_id, ledger): nonlocal writes writes += 1 if writes == 2: raise OSError("simulated ledger finalize crash") await original_replace(root_trace_id, ledger) with patch.object( self.store, "replace_candidate_ledger", side_effect=fail_completed_publish, ): with self.assertRaisesRegex( RecoverableToolExecutionError, "original command ID", ): await self.service.manage( "root-a", operation="create", content={"text": "survives retry"}, parent_refs=[], effective_at_sequence=2, command_id="candidate-crash-call", ) pending = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual("pending", pending.operations[0].status) operation_id = pending.operations[0].operation_id recovered = await self.service.manage( "root-a", operation="create", content={"text": "survives retry"}, parent_refs=[], effective_at_sequence=99, command_id="candidate-crash-call", ) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(operation_id, ledger.operations[0].operation_id) self.assertEqual("completed", ledger.operations[0].status) self.assertEqual(1, len(ledger.candidates)) self.assertEqual(1, len(self.repository.contents)) self.assertEqual( recovered, ledger.candidate(ledger.operations[0].candidate), ) async def test_distinct_tool_calls_with_identical_payload_create_versions(self): first = await self.service.manage( "root-a", operation="create", content={"text": "same payload"}, parent_refs=[], effective_at_sequence=10, command_id="candidate-call-one", ) await self.store.add_message(TraceMessage.create( trace_id="root-a", role="tool", sequence=10, tool_call_id="candidate-call-one", content={"tool_name": "manage_candidate", "result": "registered"}, )) await self.store.update_trace("root-a", head_sequence=10) second = await self.service.manage( "root-a", operation="create", content={"text": "same payload"}, parent_refs=[], effective_at_sequence=20, command_id="candidate-call-two", ) await self.store.add_message(TraceMessage.create( trace_id="root-a", role="tool", sequence=20, parent_sequence=10, tool_call_id="candidate-call-two", content={"tool_name": "manage_candidate", "result": "registered"}, )) await self.store.update_trace("root-a", head_sequence=20) self.assertNotEqual(first, second) self.assertEqual(2, self.repository.put_calls) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(2, len(ledger.candidates)) self.assertEqual(2, len(ledger.operations)) async def test_rewind_makes_future_branch_candidate_unreportable(self): await self.store.add_message(TraceMessage.create( trace_id="root-a", role="user", sequence=1, content="before candidate", )) await self.store.update_trace("root-a", head_sequence=1) candidate = await self.service.manage( "root-a", operation="create", content={"text": "future branch"}, parent_refs=[], effective_at_sequence=10, command_id="future-candidate-call", ) await self.store.add_message(TraceMessage.create( trace_id="root-a", role="tool", sequence=10, parent_sequence=1, tool_call_id="future-candidate-call", content={"tool_name": "manage_candidate", "result": "registered"}, )) await self.store.update_trace("root-a", head_sequence=10) await self.service.validate_report_refs("root-a", [candidate]) await self.store.update_trace("root-a", head_sequence=1) with self.assertRaisesRegex( CandidateStateError, "current owner Trace branch", ): await self.service.validate_report_refs("root-a", [candidate]) async def test_checkpoint_update_cannot_mutate_identity_or_sequence(self): candidate = await self.service.manage( "root-a", operation="create", content={"text": "checkpoint"}, parent_refs=[], effective_at_sequence=3, ) plan_hash = "a" * 64 await self.service.begin_validation_checkpoint( "root-a", candidate, plan={"checks": []}, plan_hash=plan_hash, validated_at_sequence=3, material_usage_recorded=False, ) with self.assertRaisesRegex( CandidateStateError, "fields are immutable", ): await self.service.update_validation_checkpoint( "root-a", candidate, plan_hash, validated_at_sequence=-1, ) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual( 3, ledger.validation_checkpoints[0]["validated_at_sequence"], ) 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_adoption_retry_uses_same_command_after_parent_sequence_changes(self): await self._create_child("child-a") candidate = await self.service.manage( "child-a", operation="create", content={"text": "publish once"}, parent_refs=[], effective_at_sequence=4, ) await self._record_validation("child-a", candidate) peer = await self.service.manage( "child-a", operation="create", content={"text": "peer"}, parent_refs=[], effective_at_sequence=5, ) await self._record_validation("child-a", peer) action = CandidateReviewAction( action="adopt", candidate_ref=candidate, reason="stable parent decision", ) first = await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate, peer], actions=[action], effective_at_sequence=10, command_id="adoption-review-call", ) recovered = await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate, peer], actions=[action], effective_at_sequence=999, command_id="adoption-review-call", ) self.assertEqual(first, recovered) self.assertEqual(1, self.repository.adoption_calls) self.assertEqual(1, len(self.repository.adoptions)) with self.assertRaisesRegex( CandidateStateError, "review batch payload changed", ): await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate, peer], actions=[action.model_copy(update={ "reason": "tampered decision payload", })], effective_at_sequence=1_000, command_id="adoption-review-call", ) self.assertEqual(1, self.repository.adoption_calls) with self.assertRaisesRegex( CandidateStateError, "review batch payload changed", ): await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate, peer], actions=[ action, CandidateReviewAction( action="discard", candidate_ref=peer, reason="injected tail action", ), ], effective_at_sequence=1_001, command_id="adoption-review-call", ) self.assertEqual("proposed", ( CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ).current_state(peer) )) with self.assertRaisesRegex( CandidateStateError, "review batch payload changed", ): await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate, peer], actions=[], effective_at_sequence=1_002, command_id="adoption-review-call", ) async def test_review_batch_replay_rejects_deleted_and_reordered_actions(self): await self._create_child("child-a") candidates = [] for sequence, text in ((4, "A"), (5, "B")): candidate = await self.service.manage( "child-a", operation="create", content={"text": text}, parent_refs=[], effective_at_sequence=sequence, ) await self._record_validation("child-a", candidate) candidates.append(candidate) actions = [ CandidateReviewAction( action="discard", candidate_ref=candidate, reason=f"discard {candidate.candidate_id}", ) for candidate in candidates ] await self.service.apply_review_actions( "root-a", "child-a", report_refs=candidates, actions=actions, effective_at_sequence=10, command_id="two-action-review", ) for altered in (actions[:1], list(reversed(actions))): with self.assertRaisesRegex( CandidateStateError, "review batch payload changed", ): await self.service.apply_review_actions( "root-a", "child-a", report_refs=candidates, actions=altered, effective_at_sequence=999, command_id="two-action-review", ) async def test_empty_review_batch_cannot_gain_an_action_during_replay(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) self.assertEqual([], await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[], effective_at_sequence=10, command_id="empty-review-batch", )) with self.assertRaisesRegex( CandidateStateError, "review batch payload changed", ): await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[CandidateReviewAction( action="discard", candidate_ref=candidate, reason="injected after empty batch", )], effective_at_sequence=999, command_id="empty-review-batch", ) self.assertEqual( "proposed", CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ).current_state(candidate), ) async def test_review_capacity_rejects_before_envelope_or_repository_write(self): await self._create_child("child-a") candidate = await self.service.manage( "child-a", operation="create", content={"text": "capacity"}, parent_refs=[], effective_at_sequence=4, ) await self._record_validation("child-a", candidate) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) fillers = [ CandidateOperationRecord( operation_id=f"capacity-filler-{index}", operation="review_batch", input_hash="0" * 64, status="completed", ) for index in range( MAX_CANDIDATE_OPERATIONS_PER_ROOT - len(ledger.operations) ) ] saturated = CandidateLedger.model_validate({ **ledger.model_dump(mode="json"), "operations": [ *[item.model_dump(mode="json") for item in ledger.operations], *[item.model_dump(mode="json") for item in fillers], ], }) await self.store.replace_candidate_ledger( "root-a", saturated.model_dump(mode="json"), ) before = await self.store.get_candidate_ledger("root-a") with self.assertRaisesRegex( CandidateStateError, "operation limit exceeded before execution", ): await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[CandidateReviewAction( action="adopt", candidate_ref=candidate, reason="must not reach repository", )], effective_at_sequence=10, command_id="capacity-overflow-review", ) self.assertEqual(0, self.repository.adoption_calls) self.assertEqual(before, await self.store.get_candidate_ledger("root-a")) async def test_review_capacity_allows_exact_envelope_and_action_boundary(self): await self._create_child("child-a") candidate = await self.service.manage( "child-a", operation="create", content={"text": "exact capacity"}, parent_refs=[], effective_at_sequence=4, ) await self._record_validation("child-a", candidate) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) target = MAX_CANDIDATE_OPERATIONS_PER_ROOT - 2 fillers = [ CandidateOperationRecord( operation_id=f"exact-filler-{index}", operation="review_batch", input_hash="1" * 64, status="completed", ) for index in range(target - len(ledger.operations)) ] near_limit = CandidateLedger.model_validate({ **ledger.model_dump(mode="json"), "operations": [ *[item.model_dump(mode="json") for item in ledger.operations], *[item.model_dump(mode="json") for item in fillers], ], }) await self.store.replace_candidate_ledger( "root-a", near_limit.model_dump(mode="json"), ) await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[CandidateReviewAction( action="discard", candidate_ref=candidate, reason="fits exact operation boundary", )], effective_at_sequence=10, command_id="exact-capacity-review", ) completed = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(MAX_CANDIDATE_OPERATIONS_PER_ROOT, len(completed.operations)) self.assertEqual("discarded", completed.current_state(candidate)) async def test_review_batch_reserves_action_slot_across_crash_recovery(self): await self._create_child("child-a") candidate = await self.service.manage( "child-a", operation="create", content={"text": "reserved review"}, parent_refs=[], effective_at_sequence=4, ) await self._record_validation("child-a", candidate) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) target = MAX_CANDIDATE_OPERATIONS_PER_ROOT - 2 fillers = [ CandidateOperationRecord( operation_id=f"review-reservation-filler-{index}", operation="review_batch", input_hash="2" * 64, status="completed", ) for index in range(target - len(ledger.operations)) ] near_limit = CandidateLedger.model_validate({ **ledger.model_dump(mode="json"), "operations": [ *[item.model_dump(mode="json") for item in ledger.operations], *[item.model_dump(mode="json") for item in fillers], ], }) await self.store.replace_candidate_ledger( "root-a", near_limit.model_dump(mode="json"), ) action = CandidateReviewAction( action="adopt", candidate_ref=candidate, reason="reserve the complete review batch", ) original_replace = self.store.replace_candidate_ledger writes = 0 async def crash_after_batch_reservation(root_trace_id, raw_ledger): nonlocal writes writes += 1 if writes == 2: raise OSError("crash after review reservation") await original_replace(root_trace_id, raw_ledger) with patch.object( self.store, "replace_candidate_ledger", side_effect=crash_after_batch_reservation, ): with self.assertRaisesRegex(OSError, "review reservation"): await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[action], effective_at_sequence=10, command_id="capacity-reserved-review", ) reserved = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(MAX_CANDIDATE_OPERATIONS_PER_ROOT, len(reserved.operations)) self.assertEqual(2, sum( item.status == "pending" for item in reserved.operations )) repository_version_calls = self.repository.put_calls with self.assertRaisesRegex( CandidateStateError, "operation limit exceeded before execution", ): await self.service.manage( "child-a", operation="create", content={"text": "must not steal review capacity"}, parent_refs=[], effective_at_sequence=11, command_id="sibling-after-review-reservation", ) self.assertEqual(repository_version_calls, self.repository.put_calls) records = await self.service.apply_review_actions( "root-a", "child-a", report_refs=[candidate], actions=[action], effective_at_sequence=999, command_id="capacity-reserved-review", ) self.assertEqual("adopted", records[0].state) self.assertEqual(1, self.repository.adoption_calls) completed = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(MAX_CANDIDATE_OPERATIONS_PER_ROOT, len(completed.operations)) self.assertTrue(all( item.status == "completed" for item in completed.operations if item.command_id == "capacity-reserved-review" )) async def test_candidate_capacity_rejects_before_repository_or_intent_write(self): await self._create_child("child-a") candidate = await self.service.manage( "child-a", operation="create", content={"text": "seed"}, parent_refs=[], effective_at_sequence=4, ) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) saturated_candidates = [candidate] saturated_candidates.extend( candidate.model_copy(update={ "candidate_id": f"capacity-candidate-{index}", }) for index in range(1, MAX_CANDIDATES_PER_ROOT) ) saturated = CandidateLedger.model_validate({ **ledger.model_dump(mode="json"), "candidates": [ item.model_dump(mode="json") for item in saturated_candidates ], }) await self.store.replace_candidate_ledger( "root-a", saturated.model_dump(mode="json"), ) before = await self.store.get_candidate_ledger("root-a") repository_calls = self.repository.put_calls with self.assertRaisesRegex( CandidateStateError, "candidate limit exceeded before execution", ): await self.service.manage( "child-a", operation="create", content={"text": "must remain external-orphan free"}, parent_refs=[], effective_at_sequence=5, command_id="candidate-capacity-overflow", ) self.assertEqual(repository_calls, self.repository.put_calls) self.assertEqual(before, await self.store.get_candidate_ledger("root-a")) async def test_pending_candidate_reserves_capacity_until_recovery(self): await self._create_child("child-a") seed = await self.service.manage( "child-a", operation="create", content={"text": "seed"}, parent_refs=[], effective_at_sequence=4, ) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) near_limit_candidates = [seed] near_limit_candidates.extend( seed.model_copy(update={ "candidate_id": f"reserved-candidate-{index}", }) for index in range(1, MAX_CANDIDATES_PER_ROOT - 1) ) near_limit = CandidateLedger.model_validate({ **ledger.model_dump(mode="json"), "candidates": [ item.model_dump(mode="json") for item in near_limit_candidates ], }) await self.store.replace_candidate_ledger( "root-a", near_limit.model_dump(mode="json"), ) original_replace = self.store.replace_candidate_ledger writes = 0 async def crash_after_repository(root_trace_id, raw_ledger): nonlocal writes writes += 1 if writes == 2: raise OSError("crash before candidate publish") await original_replace(root_trace_id, raw_ledger) with patch.object( self.store, "replace_candidate_ledger", side_effect=crash_after_repository, ): with self.assertRaisesRegex( RecoverableToolExecutionError, "original command ID", ): await self.service.manage( "child-a", operation="create", content={"text": "reserved A"}, parent_refs=[], effective_at_sequence=5, command_id="reserved-candidate-A", ) pending = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(MAX_CANDIDATES_PER_ROOT - 1, len(pending.candidates)) self.assertEqual( 1, sum( item.status == "pending" and item.operation == "create" for item in pending.operations ), ) calls_after_a = self.repository.put_calls with self.assertRaisesRegex( CandidateStateError, "candidate limit exceeded before execution", ): await self.service.manage( "child-a", operation="create", content={"text": "must not steal A's slot"}, parent_refs=[], effective_at_sequence=6, command_id="reserved-candidate-B", ) self.assertEqual(calls_after_a, self.repository.put_calls) recovered = await self.service.manage( "child-a", operation="create", content={"text": "reserved A"}, parent_refs=[], effective_at_sequence=99, command_id="reserved-candidate-A", ) completed = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(MAX_CANDIDATES_PER_ROOT, len(completed.candidates)) self.assertEqual(recovered, completed.candidate(recovered)) async def test_runner_replays_pending_candidate_with_original_tool_call_id(self): await self._create_child("child-a") registry = get_tool_registry().clone(["manage_candidate"]) runner = AgentRunner( trace_store=self.store, tool_registry=registry, llm_call=lambda **_kwargs: None, application_binding=self.binding, candidate_service=self.service, ) config = RunConfig( tools=["manage_candidate"], tool_groups=[], auto_execute_tools=True, ) request = { "request": { "operation": "create", "content": {"text": "recover with the original command"}, "parent_refs": [], }, } assistant = TraceMessage.create( trace_id="child-a", role="assistant", sequence=4, content={ "text": "", "tool_calls": [{ "id": "recoverable-candidate-call", "type": "function", "function": { "name": "manage_candidate", "arguments": json.dumps(request), }, }], }, ) await self.store.add_message(assistant) await self.store.update_trace( "child-a", head_sequence=4, last_sequence=4, ) trace = await self.store.get_trace("child-a") original_replace = self.store.replace_candidate_ledger writes = 0 async def fail_first_publish(root_trace_id, raw_ledger): nonlocal writes writes += 1 if writes == 2: raise OSError("candidate publish interrupted") await original_replace(root_trace_id, raw_ledger) with patch.object( self.store, "replace_candidate_ledger", side_effect=fail_first_publish, ): with self.assertRaises(RecoverableToolExecutionError): await registry.execute( "manage_candidate", request, context=runner._build_tool_context( config=config, trace=trace, trace_id=trace.trace_id, goal_id=None, goal_tree=None, sequence=5, head_sequence=4, tool_call_id="recoverable-candidate-call", side_branch_ctx=None, trigger_event=None, ), allowed_tool_names={"manage_candidate"}, tool_call_id="recoverable-candidate-call", ) pending = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual("pending", pending.operations[-1].status) self.assertEqual(1, len(self.repository.contents)) self.assertFalse(any( item.role == "tool" for item in await self.store.get_trace_messages("child-a") )) healed, next_sequence = await runner._heal_orphaned_tool_calls( [assistant], await self.store.get_trace("child-a"), None, config, 5, ) self.assertEqual(6, next_sequence) self.assertEqual("tool", healed[-1].role) self.assertEqual( "recoverable-candidate-call", healed[-1].tool_call_id, ) self.assertEqual(1, len(self.repository.contents)) completed = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) matching = [ item for item in completed.operations if item.command_id == "recoverable-candidate-call" ] self.assertEqual(1, len(matching)) self.assertEqual("completed", matching[0].status) self.assertEqual(1, len(completed.candidates)) async def test_approved_candidate_publish_failure_keeps_original_grant(self): await self._create_child("child-a") registry = get_tool_registry().clone(["manage_candidate"]) runner = AgentRunner( trace_store=self.store, tool_registry=registry, llm_call=lambda **_kwargs: None, application_binding=self.binding, candidate_service=self.service, ) request = { "request": { "operation": "create", "content": {"text": "approved durable candidate"}, "parent_refs": [], }, } assistant = TraceMessage.create( trace_id="child-a", role="assistant", sequence=4, content={ "text": "", "tool_calls": [{ "id": "approved-recoverable-candidate", "type": "function", "function": { "name": "manage_candidate", "arguments": json.dumps(request), }, }], }, ) await self.store.add_message(assistant) await self.store.update_trace("child-a", head_sequence=4, last_sequence=4) batch = runner._create_tool_approval_batch( trace_id="child-a", assistant_msg=assistant, tool_calls=assistant.content["tool_calls"], auto_execute_tools=False, ) batch.calls[0].decision = "approved" batch.status = "decided" await self.store.replace_tool_approval_batch("child-a", batch) config = RunConfig( trace_id="child-a", uid="user-1", tools=["manage_candidate"], tool_groups=[], auto_execute_tools=False, approval_batch_id=batch.batch_id, ) original_replace = self.store.replace_candidate_ledger writes = 0 async def fail_first_terminal_publish(root_trace_id, raw_ledger): nonlocal writes writes += 1 if writes == 2: raise OSError("approved candidate terminal publish interrupted") await original_replace(root_trace_id, raw_ledger) trace = await self.store.get_trace("child-a") original_approval_replace = self.store.replace_tool_approval_batch approval_writes = 0 async def fail_recoverable_approval_publish(trace_id, raw_batch): nonlocal approval_writes approval_writes += 1 if approval_writes == 3: raise OSError("recoverable approval state publish interrupted") await original_approval_replace(trace_id, raw_batch) with patch.object( self.store, "replace_candidate_ledger", side_effect=fail_first_terminal_publish, ), patch.object( self.store, "replace_tool_approval_batch", side_effect=fail_recoverable_approval_publish, ): with self.assertRaisesRegex( OSError, "recoverable approval state publish interrupted", ): await runner._resume_approved_tool_batch( trace=trace, history=[assistant.to_llm_dict()], goal_tree=None, config=config, sequence=5, head_seq=4, side_branch_ctx=None, runtime_tool_names={"manage_candidate"}, ) interrupted = await self.store.get_tool_approval_batch("child-a") self.assertEqual("executing", interrupted.status) self.assertEqual("executing", interrupted.calls[0].execution_status) self.assertFalse( await runner._mark_approval_execution_unknown("child-a") ) recoverable = await self.store.get_tool_approval_batch("child-a") self.assertEqual("recoverable_candidate_command", recoverable.status) self.assertEqual("approved", recoverable.calls[0].decision) self.assertEqual("not_started", recoverable.calls[0].execution_status) self.assertEqual( "approved-recoverable-candidate", recoverable.calls[0].tool_call_id, ) with self.assertRaisesRegex( ValueError, "recoverable_candidate_command_must_resume_first", ): await runner._prepare_trace( [{"role": "user", "content": "change direction"}], RunConfig(trace_id="child-a"), ) with self.assertRaisesRegex( ValueError, "recoverable_candidate_command_must_resume_first", ): await runner._prepare_trace( [], RunConfig(trace_id="child-a", after_sequence=3), ) from cyber_agent.trace import run_api async def restore_for_api(**_kwargs): return runner, RunConfig( trace_id="child-a", approval_batch_id=recoverable.batch_id, ) with patch.object(run_api, "_get_runner", return_value=runner), patch.object( run_api, "_restore_runner_and_config", side_effect=restore_for_api, ): with self.assertRaisesRegex( run_api.HTTPException, "recoverable_candidate_command_must_resume_first", ) as blocked_api: await run_api.run_trace( "child-a", run_api.TraceRunRequest(messages=[{ "role": "user", "content": "must wait for candidate recovery", }]), ) self.assertEqual(409, blocked_api.exception.status_code) resumed_config = config resumed_config.approval_batch_id = recoverable.batch_id _history, next_sequence, head_sequence = ( await runner._resume_approved_tool_batch( trace=await self.store.get_trace("child-a"), history=[assistant.to_llm_dict()], goal_tree=None, config=resumed_config, sequence=5, head_seq=4, side_branch_ctx=None, runtime_tool_names={"manage_candidate"}, ) ) completed_batch = await self.store.get_tool_approval_batch("child-a") completed_ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual("completed", completed_batch.status) self.assertEqual(6, next_sequence) self.assertEqual(5, head_sequence) self.assertEqual(1, len(completed_ledger.candidates)) self.assertEqual(1, len(self.repository.contents)) self.assertEqual(2, self.repository.put_calls) # Crash window: the Tool Result is durable, while the last approval # update still shows the call as executing. Recovery must finalize the # existing message instead of replaying the candidate command. completed_batch.status = "executing" completed_batch.calls[0].execution_status = "executing" completed_batch.calls[0].result_message_id = None await self.store.replace_tool_approval_batch("child-a", completed_batch) await self.store.update_trace("child-a", head_sequence=4) self.assertTrue( await runner._recover_candidate_approval_if_durable( "child-a", completed_batch, ) ) reconciled = await self.store.get_tool_approval_batch("child-a") trace_after_reconcile = await self.store.get_trace("child-a") matching_results = [ message for message in await self.store.get_trace_messages("child-a") if message.role == "tool" and message.tool_call_id == "approved-recoverable-candidate" ] self.assertEqual("completed", reconciled.status) self.assertEqual("executed", reconciled.calls[0].execution_status) self.assertEqual( matching_results[0].message_id, reconciled.calls[0].result_message_id, ) self.assertEqual( matching_results[0].sequence, trace_after_reconcile.head_sequence, ) self.assertEqual(1, len(matching_results)) self.assertEqual(2, self.repository.put_calls) async def test_failed_outcome_publish_failure_remains_orphan_replayable(self): await self._create_child("child-a") self.repository.bad_artifact = True registry = get_tool_registry().clone(["manage_candidate"]) runner = AgentRunner( trace_store=self.store, tool_registry=registry, llm_call=lambda **_kwargs: None, application_binding=self.binding, candidate_service=self.service, ) config = RunConfig( tools=["manage_candidate"], tool_groups=[], auto_execute_tools=True, ) request = { "request": { "operation": "create", "content": {"text": "invalid artifact candidate"}, "parent_refs": [], }, } assistant = TraceMessage.create( trace_id="child-a", role="assistant", sequence=4, content={ "text": "", "tool_calls": [{ "id": "failed-publish-candidate", "type": "function", "function": { "name": "manage_candidate", "arguments": json.dumps(request), }, }], }, ) await self.store.add_message(assistant) await self.store.update_trace("child-a", head_sequence=4, last_sequence=4) trace = await self.store.get_trace("child-a") original_replace = self.store.replace_candidate_ledger writes = 0 async def fail_first_failed_publish(root_trace_id, raw_ledger): nonlocal writes writes += 1 if writes == 2: raise OSError("failed outcome publish interrupted") await original_replace(root_trace_id, raw_ledger) with patch.object( self.store, "replace_candidate_ledger", side_effect=fail_first_failed_publish, ): with self.assertRaises(RecoverableToolExecutionError): await registry.execute( "manage_candidate", request, context=runner._build_tool_context( config=config, trace=trace, trace_id=trace.trace_id, goal_id=None, goal_tree=None, sequence=5, head_sequence=4, tool_call_id="failed-publish-candidate", side_branch_ctx=None, trigger_event=None, ), allowed_tool_names={"manage_candidate"}, tool_call_id="failed-publish-candidate", ) pending = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual("pending", pending.operations[-1].status) self.assertFalse(any( item.role == "tool" for item in await self.store.get_trace_messages("child-a") )) healed, next_sequence = await runner._heal_orphaned_tool_calls( [assistant], await self.store.get_trace("child-a"), None, config, 5, ) terminal = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(6, next_sequence) self.assertEqual("tool", healed[-1].role) self.assertEqual("failed-publish-candidate", healed[-1].tool_call_id) self.assertEqual("failed", terminal.operations[-1].status) self.assertEqual(0, len(terminal.candidates)) async def test_pending_recovery_survives_transient_repository_error(self): await self._create_child("child-a") command_id = "transient-recovery-candidate" request = { "operation": "create", "content": {"text": "eventually consistent candidate"}, "parent_refs": [], "effective_at_sequence": 5, "command_id": command_id, } original_replace = self.store.replace_candidate_ledger writes = 0 async def fail_first_terminal_publish(root_trace_id, raw_ledger): nonlocal writes writes += 1 if writes == 2: raise OSError("candidate terminal publish interrupted") await original_replace(root_trace_id, raw_ledger) with patch.object( self.store, "replace_candidate_ledger", side_effect=fail_first_terminal_publish, ): with self.assertRaises(RecoverableToolExecutionError): await self.service.manage("child-a", **request) with patch.object( self.repository, "put_version", side_effect=ValueError("repository temporarily unavailable"), ): with self.assertRaises(RecoverableToolExecutionError): await self.service.manage("child-a", **request) still_pending = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual("pending", still_pending.operations[-1].status) with patch.object( self.repository, "resolve", side_effect=RuntimeError("artifact resolver temporarily unavailable"), ): with self.assertRaises(RecoverableToolExecutionError): await self.service.manage("child-a", **request) still_pending = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual("pending", still_pending.operations[-1].status) recovered = await self.service.manage("child-a", **request) completed = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(recovered, completed.candidates[-1]) self.assertEqual("completed", completed.operations[-1].status) self.assertEqual(1, len(self.repository.contents)) async def test_first_repository_timeout_after_commit_remains_replayable(self): await self._create_child("child-a") request = { "operation": "create", "content": {"text": "committed before response timeout"}, "parent_refs": [], "effective_at_sequence": 5, "command_id": "first-call-timeout-candidate", } original_put = self.repository.put_version async def commit_then_timeout(version_request): await original_put(version_request) raise TimeoutError("response was lost after commit") with patch.object( self.repository, "put_version", side_effect=commit_then_timeout, ): with self.assertRaises(RecoverableToolExecutionError): await self.service.manage("child-a", **request) pending = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual("pending", pending.operations[-1].status) self.assertEqual(1, len(self.repository.contents)) recovered = await self.service.manage("child-a", **request) completed = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual(recovered, completed.candidates[-1]) self.assertEqual("completed", completed.operations[-1].status) self.assertEqual(1, len(self.repository.contents)) async def test_typed_repository_rejection_is_terminal(self): await self._create_child("child-a") with patch.object( self.repository, "put_version", side_effect=CandidateRepositoryRejected("content policy rejected"), ): with self.assertRaisesRegex( CandidateStateError, "content policy rejected", ): await self.service.manage( "child-a", operation="create", content={"text": "rejected candidate"}, parent_refs=[], effective_at_sequence=5, command_id="deterministic-rejection-candidate", ) failed = CandidateLedger.model_validate( await self.store.get_candidate_ledger("root-a") ) self.assertEqual("failed", failed.operations[-1].status) self.assertEqual(0, len(failed.candidates)) self.assertEqual(0, len(self.repository.contents)) async def test_descendant_adoption_blocks_ancestor_rewind_without_mapping(self): await self._create_child("child-a") await self.store.create_trace(Trace( trace_id="grandchild-a", mode="agent", agent_type="writer", uid="user-1", parent_trace_id="child-a", context=self._context( "root-a", task_brief=TaskBrief( objective="write descendant candidate", reason="the child needs an option", completion_criteria=["candidate is complete"], expected_outputs=["one candidate"], ), ), )) candidate = await self.service.manage( "grandchild-a", operation="create", content={"text": "nested candidate"}, parent_refs=[], effective_at_sequence=4, ) await self._record_validation("grandchild-a", candidate) await self.service.apply_review_actions( "child-a", "grandchild-a", report_refs=[candidate], actions=[CandidateReviewAction( action="adopt", candidate_ref=candidate, reason="commit descendant output", )], effective_at_sequence=9, ) with self.assertRaisesRegex(CandidateStateError, "committed"): await self.service.assert_rewind_allowed("root-a", 100) 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 == 3: 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()