LangBot/src/langbot/pkg/agent/runner/registry.py

"""Agent runner registry for discovering and caching runner descriptors."""

from __future__ import annotations

import typing
import asyncio

from langbot_plugin.api.entities.builtin.runner.manifest import (
    RunnerManifest,
)

from ...core import app
from ...api.http.context import ExecutionContext
from ...api.http.service.tenant import TenantContext
from .descriptor import RunnerDescriptor
from .id import parse_runner_id, format_runner_id
from .errors import RunnerNotFoundError, RunnerNotAuthorizedError


class RunnerRegistry:
    """Registry for discovering and managing agent runners.

    Responsibilities:
    - Discover runners from plugin runtime via LIST_RUNNERS
    - Validate runner manifests (kind, metadata, spec)
    - Cache discovered runners for performance
    - Filter runners by bound plugins
    - Handle manifest errors gracefully (log warning, skip runner)
    """

    ap: app.Application

    _cache: dict[tuple[str, str, int], dict[str, RunnerDescriptor]]
    """Runner descriptors keyed by immutable Workspace execution scope."""

    _cache_lock: asyncio.Lock
    """Lock for cache refresh operations"""

    def __init__(self, ap: app.Application):
        self.ap = ap
        self._cache = {}
        self._cache_lock = asyncio.Lock()

    @staticmethod
    def _cache_key(context: ExecutionContext) -> tuple[str, str, int]:
        return (
            context.instance_uuid,
            context.workspace_uuid,
            context.placement_generation,
        )

    async def _resolve_context(self, context: TenantContext) -> ExecutionContext:
        return await self.ap.plugin_connector.require_workspace_context(context)

    async def _discover_runners(self) -> dict[str, RunnerDescriptor]:
        """Discover runners from plugin runtime.

        Always discovers ALL runners (no bound_plugins filter).
        The cache should contain unfiltered discovery results.

        Returns:
            Dict of runner descriptors keyed by runner ID
        """
        if not self.ap.plugin_connector.is_enable_plugin:
            return {}

        runners: dict[str, RunnerDescriptor] = {}

        try:
            # Always list all runners (bound_plugins=None)
            plugin_runners = await self.ap.plugin_connector.list_runners(None)

            for runner_data in plugin_runners:
                try:
                    descriptor = self._validate_and_build_descriptor(runner_data)
                    if descriptor is not None:
                        runners[descriptor.id] = descriptor
                except Exception as e:
                    plugin_author = runner_data.get('plugin_author', 'unknown')
                    plugin_name = runner_data.get('plugin_name', 'unknown')
                    runner_name = runner_data.get('runner_name', 'unknown')
                    self.ap.logger.warning(
                        f'Invalid runner manifest for plugin:{plugin_author}/{plugin_name}/{runner_name}: {e}'
                    )
                    continue

        except Exception as e:
            self.ap.logger.warning(f'Failed to list agent runners from plugin runtime: {e}')
            return {}

        return runners

    def _validate_and_build_descriptor(self, runner_data: dict[str, typing.Any]) -> RunnerDescriptor | None:
        """Validate runner manifest and build descriptor.

        Args:
            runner_data: Raw runner data from plugin runtime with fields:
                - plugin_author, plugin_name, runner_name
                - manifest (typed RunnerManifest)

        Returns:
            RunnerDescriptor if valid, None if invalid
        """
        plugin_author = runner_data.get('plugin_author', '')
        plugin_name = runner_data.get('plugin_name', '')
        runner_name = runner_data.get('runner_name', '')

        if not plugin_author or not plugin_name or not runner_name:
            return None

        manifest = runner_data.get('manifest', {})
        runner_id = format_runner_id(
            source='plugin',
            plugin_author=plugin_author,
            plugin_name=plugin_name,
            runner_name=runner_name,
        )

        typed_manifest = RunnerManifest.model_validate(manifest)
        config_schema = [item.model_dump(mode='json') for item in typed_manifest.config_schema]

        descriptor = RunnerDescriptor(
            id=runner_id,
            component_kind=typed_manifest.component_kind,
            usages=typed_manifest.usages,
            supported_event_patterns=typed_manifest.supported_event_patterns,
            source='plugin',
            label=typed_manifest.label,
            description=typed_manifest.description,
            plugin_author=plugin_author,
            plugin_name=plugin_name,
            runner_name=runner_name,
            plugin_version=runner_data.get('plugin_version'),
            config_schema=config_schema,
            capabilities=typed_manifest.capabilities,
            permissions=typed_manifest.permissions,
            raw_manifest=manifest,
        )
        return descriptor

    async def refresh(self, context: TenantContext) -> None:
        """Refresh runner cache.

        Always discovers ALL runners (no bound_plugins filter).
        The cache contains unfiltered discovery results.
        """
        execution_context = await self._resolve_context(context)
        runners = await self._discover_runners()
        async with self._cache_lock:
            self._cache[self._cache_key(execution_context)] = runners

    async def list_runners(
        self,
        context: TenantContext,
        bound_plugins: list[str] | None = None,
        use_cache: bool = True,
        usage: typing.Literal['agent', 'event'] | None = 'agent',
    ) -> list[RunnerDescriptor]:
        """List available runners.

        Args:
            bound_plugins: Optional filter for bound plugins (applied locally)
            use_cache: Use cached data if available

        Returns:
            List of runner descriptors
        """
        execution_context = await self._resolve_context(context)
        cache_key = self._cache_key(execution_context)
        cached = self._cache.get(cache_key)
        if use_cache and cached:
            # Filter from cache. Do not treat an empty cache as final because the
            # plugin runtime may still be launching installed plugins when the
            # first metadata request arrives.
            return [
                r
                for r in self._filter_runners_by_bound_plugins(cached, bound_plugins)
                if usage is None or usage in r.usages
            ]

        # Discover fresh (always full list)
        runners = await self._discover_runners()

        # Update cache (full list, unfiltered)
        async with self._cache_lock:
            self._cache[cache_key] = runners

        # Filter locally
        return [
            r
            for r in self._filter_runners_by_bound_plugins(runners, bound_plugins)
            if usage is None or usage in r.usages
        ]

    def _filter_runners_by_bound_plugins(
        self,
        runners: dict[str, RunnerDescriptor],
        bound_plugins: list[str] | None,
    ) -> list[RunnerDescriptor]:
        """Filter runners by bound plugins.

        Args:
            runners: Dict of runner descriptors
            bound_plugins: Optional filter (None means all plugins allowed)

        Returns:
            Filtered list of runner descriptors
        """
        if bound_plugins is None:
            # All plugins allowed
            return list(runners.values())

        allowed_plugin_ids = set(bound_plugins)
        filtered = []
        for descriptor in runners.values():
            plugin_id = descriptor.get_plugin_id()
            if plugin_id in allowed_plugin_ids:
                filtered.append(descriptor)

        return filtered

    async def get(
        self,
        context: TenantContext,
        runner_id: str,
        bound_plugins: list[str] | None = None,
    ) -> RunnerDescriptor:
        """Get a specific runner descriptor.

        Args:
            runner_id: Runner ID to lookup
            bound_plugins: Optional bound plugins filter

        Returns:
            RunnerDescriptor

        Raises:
            RunnerNotFoundError: If runner not found
            RunnerNotAuthorizedError: If runner not in bound plugins
        """
        # Parse and validate runner ID format
        try:
            parse_runner_id(runner_id)
        except ValueError as e:
            raise RunnerNotFoundError(runner_id) from e

        runners = await self.list_runners(context, bound_plugins=None, usage=None)
        descriptor = next((item for item in runners if item.id == runner_id), None)
        if descriptor is None:
            # The runtime launches installed plugins asynchronously, so an
            # early non-empty discovery can still be only a partial snapshot.
            runners = await self.list_runners(
                context,
                bound_plugins=None,
                use_cache=False,
                usage=None,
            )
            descriptor = next((item for item in runners if item.id == runner_id), None)
        if descriptor is None:
            raise RunnerNotFoundError(runner_id)

        # Check authorization
        if bound_plugins is not None:
            plugin_id = descriptor.get_plugin_id()
            if plugin_id not in bound_plugins:
                raise RunnerNotAuthorizedError(runner_id, bound_plugins)

        return descriptor

    async def get_runner_metadata_for_pipeline(
        self,
        context: TenantContext,
    ) -> tuple[list[dict[str, typing.Any]], list[dict[str, typing.Any]]]:
        """Get runner metadata for pipeline configuration UI.

        Returns runner options and their config schemas for the DynamicForm.
        """
        # Get all runners (no bound plugin filter for metadata listing)
        runners = await self.list_runners(
            context,
            bound_plugins=None,
            use_cache=False,
        )

        options = []
        stages = []

        for descriptor in runners:
            config_schema = []
            for index, config_item in enumerate(descriptor.config_schema):
                item = dict(config_item)
                if not item.get('id'):
                    item_name = item.get('name') or str(index)
                    item['id'] = f'{descriptor.id}.{item_name}'
                config_schema.append(item)

            # Add runner option
            options.append(
                {
                    'name': descriptor.id,
                    'label': descriptor.label,
                    'description': descriptor.description,
                }
            )

            # Add config schema as stage if not empty
            if descriptor.config_schema:
                stages.append(
                    {
                        'name': descriptor.id,
                        'label': descriptor.label,
                        'description': descriptor.description,
                        'config': config_schema,
                    }
                )

        return options, stages