import json import tempfile import unittest from unittest.mock import AsyncMock, patch from cyber_agent.application import ( AgentApplication, AgentRole, ApplicationRegistry, ApplicationRuntime, ApplicationServices, ProviderRef, QualityRuleSpec, ) from cyber_agent.application.candidate import CandidateLedger from cyber_agent.application.quality import ( QualityCheckOutcome, ValidationSubject, ) from cyber_agent.core.artifacts import ( ArtifactRef, ValidationMaterial, material_content_hash, ) from cyber_agent.core.task_protocol import ensure_task_protocol from cyber_agent.tools.registry import ToolRegistry from cyber_agent.trace.store import FileSystemTraceStore class CandidateRepository: def __init__(self): self.requests = {} async def put_version(self, request): self.requests.setdefault(request.operation_id, request) saved = self.requests[request.operation_id] return ArtifactRef( artifact_id=f"candidate:{saved.candidate_id}:{saved.revision}", version=str(saved.revision), content_hash=material_content_hash(saved.content), kind="candidate.output", mime_type="application/json", ) async def merge(self, request): return await self.put_version(request) async def load_version(self, candidate_ref): return next( item.content for item in self.requests.values() if item.candidate_id == candidate_ref.candidate_id and item.revision == candidate_ref.revision ) async def resolve(self, ref, root_trace_id, uid): request = next( item for item in self.requests.values() if ref.artifact_id == f"candidate:{item.candidate_id}:{item.revision}" ) return ValidationMaterial( **ref.model_dump(mode="json"), root_trace_id=root_trace_id, uid=uid, content=request.content, ) async def commit_adoption(self, request): raise NotImplementedError class QualityProvider: def __init__(self): self.mode = "normal" self.calls = 0 async def check(self, request): self.calls += 1 outcomes = [] for rule in request.rules: failed = "{{placeholder}}" in json.dumps( request.material.content, ensure_ascii=False, ) outcomes.append(QualityCheckOutcome( rule_id=rule.rule_id, rule_version=rule.version, status="failed" if failed else "passed", issue="placeholder remains" if failed else None, )) if self.mode == "missing": return [] if self.mode == "extra": outcomes.append(QualityCheckOutcome( rule_id="unexpected", rule_version="1", status="passed", )) if self.mode == "error": raise RuntimeError("provider unavailable") return outcomes class CandidateQualityTest(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.temp = tempfile.TemporaryDirectory() self.store = FileSystemTraceStore(self.temp.name) self.repository = CandidateRepository() self.quality = QualityProvider() application = AgentApplication( application_id="quality.reference", application_version="1", root_role="writer", roles=(AgentRole( role_id="writer", model="validator-model", system_prompt="write", ),), artifact_resolver_ref=ProviderRef( provider_id="artifacts", provider_version="1", ), candidate_repository_ref=ProviderRef( provider_id="candidates", provider_version="1", ), quality_provider_ref=ProviderRef( provider_id="quality", provider_version="1", ), quality_rules=(QualityRuleSpec( rule_id="no_placeholder", version="1", criterion="Candidate contains no placeholder text", subject_types=("trace", "candidate"), ),), ) registry = ApplicationRegistry() self.binding = registry.register( application, ApplicationServices( tool_registry=ToolRegistry(), artifact_resolver=self.repository, candidate_repository=self.repository, quality_provider=self.quality, ), ) self.llm_calls = 0 async def llm_call(**kwargs): self.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"] if item["scope"] == scope ], "reason": "candidate output is complete", "retry_from": None, }), "tool_calls": [], } runtime = ApplicationRuntime( registry=registry, trace_store=self.store, llm_call=llm_call, ) self.runner, config = runtime.new_run( "quality.reference", "1", uid="user-1", root_task_anchor={ "objective": "write", "completion_criteria": ["complete"], "constraints": [], }, ) with patch.dict("os.environ", {"AGENT_MODE": "recursive"}, clear=False): self.trace, _goal, _sequence = await self.runner._prepare_new_trace( [{"role": "user", "content": "start"}], config, ) async def asyncTearDown(self): self.temp.cleanup() async def _candidate(self, text, sequence): return await self.runner.candidate_service.manage( self.trace.trace_id, operation="create", content={"text": text}, parent_refs=[], effective_at_sequence=sequence, ) async def _bind_current_progress_to_report(self): trace = await self.store.get_trace(self.trace.trace_id) state = ensure_task_protocol(trace.context) state["task_report_progress_revision"] = state[ "task_progress_head_revision" ] await self.store.update_trace(trace.trace_id, context=trace.context) async def test_deterministic_and_llm_checks_are_separate_and_cached(self): candidate = await self._candidate("finished", 2) first = await self.runner.validate_recursive_trace( self.trace.trace_id, candidate_ref=candidate, ) checks = first.result.scope_results[0].checks self.assertEqual("passed", first.result.outcome) self.assertIn("quality.no_placeholder", [item.check_id for item in checks]) self.assertTrue(any(item.check_id.startswith("output.policy") for item in checks)) self.assertEqual(1, self.quality.calls) self.assertEqual(1, self.llm_calls) second = await self.runner.validate_recursive_trace( self.trace.trace_id, candidate_ref=candidate, ) self.assertTrue(second.cached) self.assertEqual(1, self.quality.calls) self.assertEqual(1, self.llm_calls) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger(self.trace.trace_id) ) self.assertEqual(1, len(ledger.validations)) async def test_failed_candidate_does_not_call_llm_or_pollute_peer(self): good = await self._candidate("finished", 2) bad = await self._candidate("{{placeholder}}", 3) good_run = await self.runner.validate_recursive_trace( self.trace.trace_id, candidate_ref=good, ) bad_run = await self.runner.validate_recursive_trace( self.trace.trace_id, candidate_ref=bad, ) self.assertEqual("passed", good_run.result.outcome) self.assertEqual("failed", bad_run.result.outcome) self.assertEqual(1, self.llm_calls) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger(self.trace.trace_id) ) self.assertEqual(2, len(ledger.validations)) async def test_provider_shape_and_exception_fail_closed_with_complete_plan(self): for index, mode in enumerate(("missing", "extra", "error"), start=2): with self.subTest(mode=mode): self.quality.mode = mode candidate = await self._candidate(f"finished-{mode}", index) run = await self.runner.validate_recursive_trace( self.trace.trace_id, candidate_ref=candidate, ) result = run.result.scope_results[0] self.assertEqual("error", result.outcome) self.assertEqual( { item.check_id for item in result.checks }, { "quality.no_placeholder", *{ item.check_id for item in result.checks if item.check_id.startswith("output.policy") }, }, ) self.assertTrue(all( item.status == "unknown" for item in result.checks )) self.assertEqual(0, self.llm_calls) async def test_trace_quality_failure_only_short_circuits_its_output_scope(self): await self._bind_current_progress_to_report() run = await self.runner.validate_recursive_trace( self.trace.trace_id, task_brief={ "objective": "produce checked copy", "reason": "the root needs a complete output", "completion_criteria": ["copy is complete"], "expected_outputs": ["copy"], "validation_scopes": ["output"], }, task_report={ "summary": "draft", "outputs": [{"text": "{{placeholder}}"}], "evidence": [], }, ) by_scope = {item.scope: item for item in run.result.scope_results} self.assertEqual("failed", by_scope["output"].outcome) self.assertEqual("passed", by_scope["task"].outcome) self.assertEqual(1, self.quality.calls) self.assertEqual(1, self.llm_calls) async def test_trace_rule_is_skipped_when_scope_is_not_effective(self): await self._bind_current_progress_to_report() run = await self.runner.validate_recursive_trace( self.trace.trace_id, scope="task", task_report={"summary": "{{placeholder}}"}, ) self.assertEqual("passed", run.result.outcome) self.assertEqual(0, self.quality.calls) self.assertEqual(1, self.llm_calls) async def test_exhausted_quality_budget_blocks_provider_before_call(self): candidate = await self._candidate("finished", 2) material = await self.repository.resolve( candidate.artifact_ref, candidate.root_trace_id, self.trace.uid, ) subject = ValidationSubject( subject_type="candidate", trace_id=self.trace.trace_id, candidate_ref=candidate, ) before = self.quality.calls with patch.object( self.runner, "record_recursive_validation_usage", new=AsyncMock(side_effect=ValueError("budget exhausted")), ): checks, errors = await self.runner._run_quality_check_batch( subject, (material,), self.binding.application.quality_rules, "quality-budget-exhausted", ) self.assertEqual(before, self.quality.calls) self.assertIn("output", errors) self.assertTrue(all(item.status == "unknown" for item in checks.values())) async def test_candidate_scope_checkpoint_survives_final_publish_crash(self): candidate = await self._candidate("finished", 2) original = self.runner.candidate_service.record_validation attempts = 0 async def fail_final_publish_once(*args, **kwargs): nonlocal attempts attempts += 1 if attempts == 1: raise OSError("final candidate validation publish failed") return await original(*args, **kwargs) with patch.object( self.runner.candidate_service, "record_validation", side_effect=fail_final_publish_once, ): with self.assertRaisesRegex(OSError, "final candidate validation"): await self.runner.validate_recursive_trace( self.trace.trace_id, candidate_ref=candidate, ) recovered = await self.runner.validate_recursive_trace( self.trace.trace_id, candidate_ref=candidate, ) self.assertEqual("passed", recovered.result.outcome) self.assertEqual(1, self.quality.calls) self.assertEqual(1, self.llm_calls) ledger = CandidateLedger.model_validate( await self.store.get_candidate_ledger(self.trace.trace_id) ) self.assertEqual(1, len(ledger.validation_checkpoints)) self.assertEqual(1, len(ledger.validations)) async def test_candidate_material_usage_is_idempotent_across_checkpoint_crash(self): candidate = await self._candidate("finished", 2) root_id = self.trace.trace_id before = await self.runner.resource_budget.get_usage(root_id) original = self.runner.candidate_service.update_validation_checkpoint failed_once = False async def fail_material_checkpoint_once(*args, **kwargs): nonlocal failed_once if kwargs.get("material_usage_recorded") and not failed_once: failed_once = True raise OSError("material checkpoint publish failed") return await original(*args, **kwargs) with patch.object( self.runner.candidate_service, "update_validation_checkpoint", side_effect=fail_material_checkpoint_once, ): with self.assertRaisesRegex(OSError, "material checkpoint"): await self.runner.validate_recursive_trace( self.trace.trace_id, candidate_ref=candidate, ) after_failure = await self.runner.resource_budget.get_usage(root_id) recovered = await self.runner.validate_recursive_trace( self.trace.trace_id, candidate_ref=candidate, ) after_retry = await self.runner.resource_budget.get_usage(root_id) self.assertEqual("passed", recovered.result.outcome) self.assertGreater( after_failure.validation_material_chars, before.validation_material_chars, ) self.assertEqual( after_failure.validation_material_chars, after_retry.validation_material_chars, ) self.assertEqual(1, self.quality.calls) self.assertEqual(1, self.llm_calls)