Agent/tests/test_candidate_quality.py

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)