pallas/pallas/registry.py

"""
Agent Registry

Serves GET /.well-known/mcp/server.json for agent discovery.
Reads agent topology from agents.yaml and model capabilities from
fastagent.config.yaml in the working directory.

Also exposes standard health and observability endpoints:
    GET /live     — liveness probe (always 200 while process is up)
    GET /ready    — readiness probe (200 when all configured agents reachable)
    GET /metrics  — Prometheus metrics in text exposition format
"""

import asyncio
import logging
import os
from datetime import datetime, timezone
from pathlib import Path

import httpx
import yaml
from prometheus_client import CONTENT_TYPE_LATEST, generate_latest
from starlette.applications import Starlette

from starlette.requests import Request
from starlette.responses import JSONResponse, PlainTextResponse, Response
from starlette.routing import Route

from pallas.metrics import REGISTRY as _metrics_registry, set_agent_info
from pallas.server import DEFAULT_CONTEXT_WINDOW, DEFAULT_MAX_OUTPUT_TOKENS

logger = logging.getLogger(__name__)


def _config_root() -> Path:
    """Return the working directory where agents.yaml and fastagent configs live."""
    return Path.cwd()


def _load_deployment_config() -> dict:
    """Load agents.yaml — single source of truth for deployment topology."""
    config_path = _config_root() / os.environ.get("PALLAS_AGENTS_CONFIG", "agents.yaml")
    if not config_path.exists():
        return {}

    with open(config_path) as f:
        return yaml.safe_load(f) or {}


def _load_fastagent_defaults() -> tuple[str, dict]:
    """Read ``(default_model, model_capabilities)`` from fastagent.config.yaml."""
    config_path = _config_root() / "fastagent.config.yaml"
    if not config_path.exists():
        return "", {}

    with open(config_path) as f:
        config = yaml.safe_load(f) or {}

    return (
        config.get("default_model", "") or "",
        config.get("model_capabilities", {}) or {},
    )


def _resolve_capabilities(
    agent: dict, default_model: str, default_capabilities: dict
) -> dict | None:
    """Resolve the capabilities Pallas actually registered for one agent.

    Mirrors ``server._register_unknown_models``: an agent's own ``model``
    in agents.yaml wins over the global ``default_model``, its own
    ``model_capabilities`` block wins over the global one, and the same
    defaults apply — so the registry advertises the effective values rather
    than nulls.  Returns ``None`` when no model is known for this agent.
    """
    model_spec = agent.get("model") or default_model
    if not model_spec:
        return None

    capabilities = agent.get("model_capabilities") or default_capabilities
    model_name = model_spec.split(".", 1)[-1] if "." in model_spec else model_spec

    return {
        "model": model_name,
        "vision": capabilities.get("vision", False),
        "context_window": capabilities.get("context_window", DEFAULT_CONTEXT_WINDOW),
        "max_output_tokens": capabilities.get(
            "max_output_tokens", DEFAULT_MAX_OUTPUT_TOKENS
        ),
    }


def _build_registry(config: dict) -> dict:
    """Build the registry JSON from agents.yaml + fastagent.config.yaml."""
    now = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    default_model, default_capabilities = _load_fastagent_defaults()

    host = config.get("host", "localhost")
    namespace = config.get("namespace", "")
    version = config.get("version", "1.0.0")
    agents = config.get("agents", {})

    entries = []
    for name, agent in agents.items():
        # Build registry name: namespace/slug (e.g. ca.helu.mentor/jarvis)
        slug = name.replace("_", "-")
        registry_name = f"{namespace}/{slug}" if namespace else slug

        server_entry: dict = {
            "$schema": "https://static.modelcontextprotocol.io/schemas/2025-12-11/server.schema.json",
            "name": registry_name,
            "title": agent.get("title", name.title()),
            "description": agent.get("description", ""),
            "version": version,
            "remotes": [
                {
                    "type": "streamable-http",
                    "url": f"http://{host}:{agent['port']}/mcp",
                }
            ],
        }
        capabilities = _resolve_capabilities(agent, default_model, default_capabilities)
        if capabilities:
            server_entry["capabilities"] = capabilities

        entries.append(
            {
                "server": server_entry,
                "_meta": {
                    "io.modelcontextprotocol.registry/official": {
                        "status": "active",
                        "updatedAt": now,
                        "isLatest": True,
                    }
                },
            }
        )

    return {"servers": entries}


# ── Route handlers ────────────────────────────────────────────────────────────

_deployment_config = _load_deployment_config()
set_agent_info(_deployment_config.get("agents", {}))


async def server_json(request: Request) -> JSONResponse:
    return JSONResponse(_build_registry(_deployment_config))


async def live(request: Request) -> JSONResponse:
    """Liveness probe — always 200 while the process is running."""
    return JSONResponse({"status": "alive"})


async def ready(request: Request) -> Response:
    """Readiness probe — 200 when all configured agents are reachable."""
    agents = _deployment_config.get("agents", {})
    if not agents:
        return JSONResponse({"status": "ready"})

    missing: list[str] = []
    async with httpx.AsyncClient(timeout=2.0) as client:
        checks = await asyncio.gather(
            *(_probe_agent(client, name, agent["port"]) for name, agent in agents.items()),
            return_exceptions=True,
        )

    for name, result in zip(agents.keys(), checks):
        if isinstance(result, Exception) or result is False:
            missing.append(name)

    if missing:
        return Response(
            content=_json_bytes({"status": "unavailable", "missing": missing}),
            status_code=503,
            media_type="application/json",
        )
    return JSONResponse({"status": "ready"})


async def _probe_agent(client: httpx.AsyncClient, name: str, port: int) -> bool:
    """Return True if the agent's MCP port is accepting connections."""
    try:
        await client.get(f"http://127.0.0.1:{port}/mcp")
        return True
    except Exception:
        return False


async def metrics(request: Request) -> Response:
    """Prometheus metrics in text exposition format."""
    data = generate_latest(_metrics_registry)
    return Response(content=data, media_type=CONTENT_TYPE_LATEST)


def _json_bytes(obj: dict) -> bytes:
    import json
    return json.dumps(obj).encode()


# ── Starlette app ─────────────────────────────────────────────────────────────

app = Starlette(
    routes=[
        Route("/.well-known/mcp/server.json", server_json),
        Route("/live", live),
        Route("/ready", ready),
        Route("/metrics", metrics),
    ],
)


async def run_registry(
    host: str = "0.0.0.0",
    port: int = 24200,
) -> None:
    """Run the registry server."""
    import uvicorn

    deploy_name = _deployment_config.get("name", "pallas")
    logger.info(
        "Registry started: %s, port %d, %d agent(s)",
        deploy_name,
        port,
        len(_deployment_config.get("agents", {})),
    )

    config = uvicorn.Config(app, host=host, port=port, log_level="warning")
    server = uvicorn.Server(config)
    await server.serve()