Source code for taimoe.platform.registry

"""Runtime registry exposed by customer services."""

from __future__ import annotations

from collections.abc import Iterable
from threading import RLock
from typing import Any

from ._sync_types import SyncErrorHandler
from .agent_sync import AgentSyncer
from .cache import RuntimeConfigCache
from .client import TaimoeClient
from .manifest import ManifestBuilder
from .sync import RuntimeSyncer
from .types.manifest import AgentsManifest, HealthStatus, RuntimeManifest




class TaimoeRegistry:
    """Coordinates discovery, sync, and local config access for a runtime service.

    Construction does not touch global SDK state. To wire this registry into
    the global ``taimoe`` handle (so user code can call ``taimoe.agent(name)``
    and resolve against this registry's cache), do it explicitly::

        from taimoe.platform.adapters.adk import taimoe
        registry = TaimoeRegistry(...)
        taimoe.bind_to(registry)

    This keeps tests and multi-tenant setups cleanly separable.
    """

    def __init__(
        self,
        *,
        runtime_name: str | None,
        platform_url: str,
        platform_token: str | None = None,
        runtime_id: str | None = None,
        poll_interval_seconds: float = 300,
        framework: str | None = None,
        framework_version: str | None = None,
        on_sync_error: SyncErrorHandler | None = None,
    ) -> None:
        # runtime_name is optional: when present we run the runtime-scoped
        # batch sync (one HTTP call returns every agent under that runtime).
        # When absent we run in *runtime-less* mode and only pull individual
        # agents via the per-agent AgentSyncer below.
        self.runtime_name = runtime_name
        self.runtime_id = runtime_id or runtime_name or ""
        self.platform_url = platform_url
        self.platform_token = platform_token
        self.cache = RuntimeConfigCache()
        self.client = TaimoeClient(base_url=platform_url, api_key=platform_token)
        self._agents: list[Any] = []
        self._agents_lock = RLock()
        self._manifest_builder = ManifestBuilder(
            runtime_name=runtime_name or "runtime-less",
            framework=framework,
            framework_version=framework_version,
        )
        self._syncer: RuntimeSyncer | None = None
        if runtime_name:
            self._syncer = RuntimeSyncer(
                client=self.client,
                runtime_id=self.runtime_id,
                cache=self.cache,
                poll_interval_seconds=poll_interval_seconds,
                on_error=on_sync_error,
            )
        self.agent_syncer = AgentSyncer(
            client=self.client,
            cache=self.cache,
            poll_interval_seconds=poll_interval_seconds,
            on_error=on_sync_error,
        )



    def register(self, agent: Any) -> None:
        with self._agents_lock:
            self._agents.append(agent)




    def register_many(self, agents: Iterable[Any]) -> None:
        for agent in agents:
            self.register(agent)


    @property
    def registered_agents(self) -> tuple[Any, ...]:
        with self._agents_lock:
            return tuple(self._agents)



    def runtime_manifest(self) -> RuntimeManifest:
        return self._manifest_builder.runtime()




    def agents_manifest(self) -> AgentsManifest:
        return self._manifest_builder.agents(self.registered_agents)




    def health(self) -> HealthStatus:
        return self._manifest_builder.health()




    def sync_once(self, *, ignore_errors: bool = True) -> None:
        """Run one synchronous sync pass on both syncers.

        ``ignore_errors`` defaults to True so a single call site can do
        cold-start init without each branch having to choose: failure logs
        a warning, the SDK enters degraded mode, and the next background
        tick retries. Pass ``ignore_errors=False`` if a caller needs the
        exception to bubble up (e.g. CI / tests).
        """
        if self._syncer is not None:
            self._syncer.sync_once(ignore_errors=ignore_errors)
        self.agent_syncer.sync_once(ignore_errors=ignore_errors)




    def start_sync(self) -> None:
        if self._syncer is not None:
            self._syncer.start()
        self.agent_syncer.start()




    def stop_sync(self, timeout: float | None = None) -> None:
        if self._syncer is not None:
            self._syncer.stop(timeout)
        self.agent_syncer.stop(timeout)




    def fastapi_router(self) -> Any:
        """Return a single combined FastAPI router for discovery endpoints.

        Convenient but must be mounted at the FastAPI app root to keep the
        well-known URI RFC 8615-compliant. Prefer :meth:`install_routes`
        for new code — it wires both routers with the correct mounts.
        """
        from .routes import create_fastapi_router

        return create_fastapi_router(self)




    def install_routes(self, app: Any) -> None:
        """Wire well-known and discovery routes into a FastAPI app.

        Equivalent to calling :func:`taimoe.platform.routes.install_routes`
        with this registry, just spelled as a method for ergonomics.
        """
        from .routes import install_routes

        install_routes(self, app)