Agent/tests/test_candidate_service.py

import asyncio
import tempfile
import unittest

from pydantic import ValidationError

from cyber_agent.application import (
    AgentApplication,
    AgentRole,
    ApplicationRegistry,
    ApplicationServices,
    ProviderRef,
)
from cyber_agent.application.candidate import (
    CandidateLedger,
    CandidatePointer,
)
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.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

    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):
        raise NotImplementedError


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 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")


if __name__ == "__main__":
    unittest.main()