Agent/tests/test_application_events.py

import tempfile
import unittest
from unittest.mock import patch

from cyber_agent.application import (
    AgentApplication,
    AgentRole,
    ApplicationRegistry,
    ApplicationRuntime,
    ApplicationServices,
    ProviderRef,
)
from cyber_agent.application.events import RunEvent
from cyber_agent.core.task_protocol import TaskProgress
from cyber_agent.tools.registry import ToolRegistry
from cyber_agent.trace.store import FileSystemTraceStore


async def unused_llm(**kwargs):
    raise AssertionError(f"LLM must not run: {kwargs}")


class TransactionalProjector:
    def __init__(self):
        self.cursor = 0
        self.view = {}
        self.calls = []
        self.fail_key = None
        self.failed_once = False

    async def load_cursor(self, application_ref, root_trace_id):
        del application_ref, root_trace_id
        return self.cursor

    async def project(self, event, expected_cursor):
        if expected_cursor != self.cursor:
            raise ValueError("stale projector cursor")
        self.calls.append(event.event_key)
        self.view.setdefault(event.event_key, event.payload)
        if event.event_key == self.fail_key and not self.failed_once:
            self.failed_once = True
            raise RuntimeError("failed after business data write")
        self.cursor = event.event_id


class ApplicationEventsTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.store = FileSystemTraceStore(self.temp.name)
        self.projector = TransactionalProjector()
        application = AgentApplication(
            application_id="events.reference",
            application_version="1",
            root_role="writer",
            roles=(AgentRole(
                role_id="writer",
                model="fake",
                system_prompt="write",
            ),),
            event_projector_ref=ProviderRef(
                provider_id="projector",
                provider_version="1",
            ),
        )
        self.registry = ApplicationRegistry()
        self.binding = self.registry.register(
            application,
            ApplicationServices(
                tool_registry=ToolRegistry(),
                event_projector=self.projector,
            ),
        )
        self.runtime = ApplicationRuntime(
            registry=self.registry,
            trace_store=self.store,
            llm_call=unused_llm,
        )
        self.runner, config = self.runtime.new_run(
            "events.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 test_projection_stops_in_order_and_replay_is_idempotent(self):
        service = self.runner.event_service
        await service.emit(
            source_trace_id=self.trace.trace_id,
            event_type="test.first",
            event_key="first",
            effective_at_sequence=1,
            payload={"value": 1},
        )
        self.projector.fail_key = "second"
        await service.emit(
            source_trace_id=self.trace.trace_id,
            event_type="test.second",
            event_key="second",
            effective_at_sequence=2,
            payload={"value": 2},
        )
        self.assertEqual({"first", "second"}, set(self.projector.view))
        self.projector.failed_once = False
        await service.emit(
            source_trace_id=self.trace.trace_id,
            event_type="test.third",
            event_key="third",
            effective_at_sequence=3,
            payload={"value": 3},
        )
        self.assertEqual({"first", "second"}, set(self.projector.view))
        self.assertNotIn("third", self.projector.view)

        self.projector.fail_key = None
        await service.pump.pump(self.trace.trace_id)
        self.assertEqual({"first", "second", "third"}, set(self.projector.view))
        self.assertEqual(1, list(self.projector.view).count("second"))
        self.assertEqual(3, self.projector.calls.count("second"))

    async def test_unknown_schema_and_event_id_gap_are_safely_projected(self):
        await self.store.append_event(
            self.trace.trace_id,
            "ordinary",
            {"ignored": True},
        )
        event_id = await self.store.append_event(
            self.trace.trace_id,
            "run_event",
            {
                "schema_version": 99,
                "event_type": "future.event",
                "event_key": "future",
                "application_ref": self.binding.application_ref.model_dump(
                    mode="json"
                ),
                "root_trace_id": self.trace.trace_id,
                "source_trace_id": self.trace.trace_id,
                "effective_at_sequence": 1,
                "payload": {"future": True},
            },
        )
        await self.runner.event_service.pump.pump(self.trace.trace_id)
        self.assertEqual(event_id, self.projector.cursor)
        self.assertEqual({"future": True}, self.projector.view["future"])

    async def test_event_append_failure_preserves_state_and_reconcile_repairs_it(self):
        service = self.runner.event_service
        original_append = self.store.append_event

        async def fail_run_event(trace_id, event_type, payload):
            if event_type == "run_event":
                raise OSError("journal unavailable")
            return await original_append(trace_id, event_type, payload)

        with patch.object(self.store, "append_event", side_effect=fail_run_event):
            revision = await self.runner.task_protocol_service.update_progress(
                self.trace.trace_id,
                expected_revision=1,
                progress=TaskProgress(
                    phase="investigating",
                    decision_rationale="verify one fact before writing",
                ),
                effective_at_sequence=2,
            )
        self.assertEqual(2, revision.revision)
        events = await self.store.get_events(self.trace.trace_id)
        self.assertFalse(any(
            item.get("event_key", "").endswith(":2") for item in events
        ))

        await self.runtime.reconcile_events(self.trace.trace_id)
        events = await self.store.get_events(self.trace.trace_id)
        repaired = [
            RunEvent.model_validate({
                key: item[key] for key in RunEvent.model_fields if key in item
            })
            for item in events
            if item.get("event_key", "").endswith(":2")
        ]
        self.assertEqual(1, len(repaired))
        self.assertEqual("task.progress.revised", repaired[0].event_type)