Agent/cyber_agent/application/runtime.py

"""Application registration, immutable binding, and AgentRunner assembly."""

from __future__ import annotations

from dataclasses import dataclass
import logging
from types import MappingProxyType
from typing import Any, Callable, Mapping

from cyber_agent.application.models import (
    AgentApplication,
    AgentRole,
    ApplicationRef,
    RunLimits,
    application_manifest,
    role_manifest,
    stable_hash,
)
from cyber_agent.application.ports import TaskContextProvider
from cyber_agent.core.artifacts import ArtifactResolver
from cyber_agent.core.validation import ValidationPolicy
from cyber_agent.tools.registry import ToolRegistry


logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class ApplicationServices:
    """Non-serializable providers and private tool implementations."""

    tool_registry: ToolRegistry
    context_provider: TaskContextProvider | None = None
    artifact_resolver: ArtifactResolver | None = None
    candidate_repository: Any | None = None
    quality_provider: Any | None = None
    event_projector: Any | None = None
    runtime_validator: Callable[[], None] | None = None


@dataclass(frozen=True)
class RoleBinding:
    role: AgentRole
    role_hash: str
    tool_names: tuple[str, ...]
    delegated_tool_names: tuple[str, ...]
    effective_limits: RunLimits
    system_prompt: str


@dataclass(frozen=True)
class ApplicationBinding:
    application: AgentApplication
    application_ref: ApplicationRef
    roles: Mapping[str, RoleBinding]
    tool_registry: ToolRegistry
    services: ApplicationServices
    validation_policy: ValidationPolicy

    def role(self, role_id: str) -> RoleBinding:
        try:
            return self.roles[role_id]
        except KeyError as exc:
            raise ValueError(f"Unknown application role: {role_id}") from exc

    def configure_run_config(self, config: Any, role_id: str) -> RoleBinding:
        """Apply an already-bound role to an existing RunConfig instance."""
        binding = self.role(role_id)
        role = binding.role
        config.model = role.model
        config.temperature = role.temperature
        config.max_iterations = binding.effective_limits.max_iterations
        config.extra_llm_params = dict(role.model_parameters)
        config.tools = list(binding.tool_names)
        config.tool_groups = []
        config.exclude_tools = []
        config.agent_type = role.role_id
        config.system_prompt = binding.system_prompt
        config.skills = []
        config.max_parallel_children = binding.effective_limits.max_parallel_children
        config.application_ref = self.application_ref
        config.role_id = role.role_id
        config.role_hash = binding.role_hash
        config.effective_run_limits = binding.effective_limits.model_dump(mode="json")
        return binding


class ApplicationRegistry:
    """Explicit, in-memory registry keyed by an exact application version."""

    def __init__(self) -> None:
        self._bindings: dict[tuple[str, str], ApplicationBinding] = {}

    def register(
        self,
        application: AgentApplication,
        services: ApplicationServices,
    ) -> ApplicationBinding:
        key = (application.application_id, application.application_version)
        if key in self._bindings:
            raise ValueError(
                f"Application is already registered: {application.application_id}@"
                f"{application.application_version}"
            )
        self._validate_services(application, services)
        selected_specs = [
            spec
            for tool_set in application.tool_sets
            for spec in tool_set.tools
        ]
        tool_names = [item.tool_name for item in selected_specs]
        if len(tool_names) != len(set(tool_names)):
            by_name: dict[str, Any] = {}
            for spec in selected_specs:
                prior = by_name.setdefault(spec.tool_name, spec)
                if prior != spec:
                    raise ValueError(
                        f"Conflicting declarations for tool: {spec.tool_name}"
                    )
        definitions = services.tool_registry.export_definitions(
            sorted(set(tool_names))
        )
        private_registry = services.tool_registry.clone(
            sorted(set(tool_names)),
            freeze=True,
        )
        manifest = application_manifest(application, definitions)
        application_ref = ApplicationRef(
            application_id=application.application_id,
            application_version=application.application_version,
            config_hash=stable_hash(manifest),
        )
        definition_map = {item["tool_name"]: item for item in definitions}
        role_bindings: dict[str, RoleBinding] = {}
        for role in application.roles:
            role_tools = tuple(dict.fromkeys(
                spec.tool_name
                for tool_set_id in role.tool_sets
                for spec in application.tool_set(tool_set_id).tools
            ))
            role_limits = application.run_limits.tighten(role.run_limits)
            delegated_tool_names = tuple(dict.fromkeys(
                tool_name
                for child_role_id in role.allowed_child_roles
                for tool_name in (
                    spec.tool_name
                    for tool_set_id in application.role(child_role_id).tool_sets
                    for spec in application.tool_set(tool_set_id).tools
                )
            ))
            role_bindings[role.role_id] = RoleBinding(
                role=role,
                role_hash=stable_hash(
                    role_manifest(application, role, definition_map)
                ),
                tool_names=role_tools,
                delegated_tool_names=delegated_tool_names,
                effective_limits=role_limits,
                system_prompt=self._compose_system_prompt(role),
            )
        validation_policy = (
            ValidationPolicy.model_validate(application.validation_policy)
            if application.validation_policy
            else ValidationPolicy()
        )
        binding = ApplicationBinding(
            application=application,
            application_ref=application_ref,
            roles=MappingProxyType(role_bindings),
            tool_registry=private_registry,
            services=services,
            validation_policy=validation_policy,
        )
        self._bindings[key] = binding
        return binding

    def resolve(
        self,
        application_id: str,
        application_version: str,
    ) -> ApplicationBinding:
        key = (application_id, application_version)
        try:
            return self._bindings[key]
        except KeyError as exc:
            raise ValueError(
                f"Application is not registered: {application_id}@{application_version}"
            ) from exc

    def resolve_ref(
        self,
        application_ref: ApplicationRef | Mapping[str, Any],
    ) -> ApplicationBinding:
        if not isinstance(application_ref, ApplicationRef):
            application_ref = ApplicationRef.model_validate(application_ref)
        binding = self.resolve(
            application_ref.application_id,
            application_ref.application_version,
        )
        if binding.application_ref != application_ref:
            raise ValueError("Application config hash does not match the registry")
        return binding

    @staticmethod
    def _compose_system_prompt(role: AgentRole) -> str:
        sections = [role.system_prompt.strip()]
        if role.skills:
            sections.append("# Application Skills")
            sections.extend(
                f"## {skill.name}@{skill.version}\n{skill.content.strip()}"
                for skill in role.skills
            )
        return "\n\n".join(sections)

    @staticmethod
    def _validate_services(
        application: AgentApplication,
        services: ApplicationServices,
    ) -> None:
        requirements = (
            (application.context_provider_ref, services.context_provider, "context_provider"),
            (application.artifact_resolver_ref, services.artifact_resolver, "artifact_resolver"),
            (application.candidate_repository_ref, services.candidate_repository, "candidate_repository"),
            (application.quality_provider_ref, services.quality_provider, "quality_provider"),
            (application.event_projector_ref, services.event_projector, "event_projector"),
        )
        missing = [name for ref, service, name in requirements if ref and service is None]
        if missing:
            raise ValueError(f"Application runtime services are missing: {missing}")


class ApplicationRuntime:
    """Assemble the existing AgentRunner for new and restored application runs."""

    def __init__(
        self,
        *,
        registry: ApplicationRegistry,
        trace_store: Any,
        llm_call: Callable[..., Any],
        utility_llm_call: Callable[..., Any] | None = None,
    ) -> None:
        self.registry = registry
        self.trace_store = trace_store
        self.llm_call = llm_call
        self.utility_llm_call = utility_llm_call

    def build_runner(self, binding: ApplicationBinding):
        from cyber_agent.core.runner import AgentRunner
        from cyber_agent.application.events import RunEventService

        event_service = RunEventService(
            store=self.trace_store,
            application_ref=binding.application_ref,
            projector=binding.services.event_projector,
        )

        candidate_service = None
        if binding.services.candidate_repository is not None:
            from cyber_agent.application.candidate_service import CandidateService

            if binding.services.artifact_resolver is None:
                raise ValueError(
                    "CandidateRepository requires an ArtifactResolver"
                )
            candidate_service = CandidateService(
                store=self.trace_store,
                application_binding=binding,
                repository=binding.services.candidate_repository,
                artifact_resolver=binding.services.artifact_resolver,
                event_service=event_service,
            )

        return AgentRunner(
            trace_store=self.trace_store,
            tool_registry=binding.tool_registry,
            llm_call=self.llm_call,
            utility_llm_call=self.utility_llm_call,
            validation_policy=binding.validation_policy,
            artifact_resolver=binding.services.artifact_resolver,
            application_binding=binding,
            context_provider=binding.services.context_provider,
            candidate_service=candidate_service,
            event_service=event_service,
        )

    async def reconcile_events(self, root_trace_id: str) -> int:
        trace = await self.trace_store.get_trace(root_trace_id)
        if trace is None:
            raise ValueError(f"Trace not found: {root_trace_id}")
        binding = self.registry.resolve_ref(trace.context.get("application_ref"))
        from cyber_agent.application.events import RunEventService

        return await RunEventService(
            store=self.trace_store,
            application_ref=binding.application_ref,
            projector=binding.services.event_projector,
        ).reconcile_root(root_trace_id)

    async def reconcile_all_events(self) -> int:
        roots = {
            item.context.get("root_trace_id") or item.trace_id
            for item in await self.trace_store.list_traces(limit=10_000)
            if item.context.get("application_ref") is not None
        }
        total = 0
        for root_trace_id in sorted(roots):
            try:
                total += await self.reconcile_events(root_trace_id)
            except Exception:
                logger.exception(
                    "Application event reconciliation failed for root %s",
                    root_trace_id,
                )
        return total

    def new_run(
        self,
        application_id: str,
        application_version: str,
        *,
        uid: str | None = None,
        name: str | None = None,
        root_task_anchor: Any = None,
        max_iterations: int | None = None,
    ):
        from cyber_agent.core.runner import RunConfig

        binding = self.registry.resolve(application_id, application_version)
        self._validate_runtime_services(binding)
        role = binding.role(binding.application.root_role)
        if max_iterations is not None and max_iterations > role.effective_limits.max_iterations:
            raise ValueError(
                "max_iterations may only tighten the application run limit"
            )
        config = RunConfig(
            uid=uid,
            name=name,
            root_task_anchor=root_task_anchor,
            enable_research_flow=False,
        )
        binding.configure_run_config(config, role.role.role_id)
        effective = self._compile_deployment_limits(role.effective_limits)
        config.effective_run_limits = effective.model_dump(mode="json")
        config.max_iterations = effective.max_iterations
        config.max_parallel_children = effective.max_parallel_children
        if max_iterations is not None:
            config.max_iterations = max_iterations
            config.effective_run_limits = {
                **config.effective_run_limits,
                "max_iterations": max_iterations,
            }
        return self.build_runner(binding), config

    @staticmethod
    def _compile_deployment_limits(limits: RunLimits) -> RunLimits:
        from cyber_agent.core.resource_budget import ResourceBudget

        deployment = ResourceBudget.from_environment()
        return limits.model_copy(update={
            "max_total_agents": min(
                limits.max_total_agents,
                deployment.max_total_agents,
            ),
            "max_llm_calls": min(
                limits.max_llm_calls,
                deployment.max_llm_calls,
            ),
            "max_total_tokens": min(
                limits.max_total_tokens,
                deployment.max_total_tokens,
            ),
            "max_total_cost_usd": min(
                limits.max_total_cost_usd,
                deployment.max_total_cost_usd,
            ),
            "max_duration_seconds": min(
                limits.max_duration_seconds,
                deployment.max_duration_seconds,
            ),
            "max_validation_tool_calls": min(
                limits.max_validation_tool_calls,
                deployment.max_validation_tool_calls,
            ),
            "max_validation_material_chars": min(
                limits.max_validation_material_chars,
                deployment.max_validation_material_chars,
            ),
        })

    async def restore(self, trace_id: str):
        from cyber_agent.core.run_snapshot import (
            RunConfigSnapshotV2,
            load_run_config_snapshot,
        )
        from cyber_agent.core.runner import RunConfig

        trace = await self.trace_store.get_trace(trace_id)
        if trace is None:
            raise ValueError(f"Trace not found: {trace_id}")
        snapshot = load_run_config_snapshot(trace.context)
        if not isinstance(snapshot, RunConfigSnapshotV2):
            raise ValueError("RunConfigSnapshotV1 is not an application run")
        binding = self.registry.resolve_ref(snapshot.application_ref)
        self._validate_runtime_services(binding)
        role = binding.role(snapshot.role_id)
        if role.role_hash != snapshot.role_hash:
            raise ValueError("Application role hash does not match the registry")
        config = RunConfig(trace_id=trace_id)
        binding.configure_run_config(config, snapshot.role_id)
        config.apply_snapshot(snapshot)
        config.system_prompt = role.system_prompt
        config.skills = []
        runner = self.build_runner(binding)
        if snapshot.uid != trace.uid or snapshot.agent_type != (
            trace.agent_type or "default"
        ):
            raise ValueError(
                "Application snapshot identity does not match Trace metadata"
            )
        runner._validate_application_snapshot(snapshot, trace)
        await self.reconcile_events(
            trace.context.get("root_trace_id") or trace.trace_id
        )
        return runner, config

    @staticmethod
    def _validate_runtime_services(binding: ApplicationBinding) -> None:
        validator = binding.services.runtime_validator
        if validator is not None:
            validator()