Agent/tests/test_candidate_service.py

import asyncio
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_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()