pallas/tests/test_registry.py

"""Tests for pallas.registry — per-agent model capability resolution.

The registry advertises the capabilities Pallas actually registered with
fast-agent, resolved *per agent*: an agent's own ``model`` /
``model_capabilities`` in agents.yaml override the global ``default_model`` /
``model_capabilities`` in fastagent.config.yaml, and the same effective
defaults as ``server._register_one_model`` apply when a field is omitted.

Regression cover for two bugs: a single capabilities dict was previously
built from ``default_model`` and attached to *every* agent entry (so an
agent with a ``model:`` override was advertised under the wrong model), and
``context_window`` / ``max_output_tokens`` were published as ``null`` when
``model_capabilities`` was absent even though the model had been registered
with 131072 / 16384.

``_build_registry`` takes the deployment config as an argument but reads
fastagent.config.yaml from the working directory on every call, so each test
chdirs into a clean temp workspace and writes only the config it needs.
"""
from __future__ import annotations

from pathlib import Path

import pytest
import yaml

from pallas import registry
from pallas.server import DEFAULT_CONTEXT_WINDOW, DEFAULT_MAX_OUTPUT_TOKENS


# ── Helpers ──────────────────────────────────────────────────────────────────


@pytest.fixture
def workspace(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
    """Chdir into a clean temp workspace — _build_registry reads cwd."""
    monkeypatch.chdir(tmp_path)
    return tmp_path


def _write_fastagent(workspace: Path, **config) -> None:
    (workspace / "fastagent.config.yaml").write_text(yaml.safe_dump(config))


def _deployment(**agents) -> dict:
    """Minimal agents.yaml-shaped config; each agent needs at least a port."""
    return {
        "name": "test-project",
        "namespace": "ca.helu.test",
        "host": "test-host",
        "agents": {
            name: {"module": f"agents.{name}", "port": 9000 + i, **overrides}
            for i, (name, overrides) in enumerate(agents.items())
        },
    }


def _entries(config: dict) -> dict[str, dict]:
    """Build the registry and return ``{agent slug: server entry}``."""
    servers = registry._build_registry(config)["servers"]
    return {e["server"]["name"].rsplit("/", 1)[-1]: e["server"] for e in servers}


def _capabilities(config: dict, agent: str = "solo") -> dict | None:
    return _entries(config)[agent].get("capabilities")


# ── Model resolution ─────────────────────────────────────────────────────────


def test_agent_model_overrides_default_model(workspace: Path) -> None:
    """An agents.yaml ``model:`` wins over the global default_model."""
    _write_fastagent(workspace, default_model="openai.global-model")

    caps = _capabilities(_deployment(solo={"model": "anthropic.agent-model"}))

    assert caps is not None
    assert caps["model"] == "agent-model"


def test_falls_back_to_default_model(workspace: Path) -> None:
    """With no per-agent model, the global default_model is advertised."""
    _write_fastagent(workspace, default_model="anthropic.claude-opus-4-7")

    caps = _capabilities(_deployment(solo={}))

    assert caps is not None
    # Provider prefix is stripped — clients get the bare model name.
    assert caps["model"] == "claude-opus-4-7"


def test_model_without_provider_prefix_passes_through(workspace: Path) -> None:
    """A default_model with no ``provider.`` prefix is emitted verbatim."""
    _write_fastagent(workspace, default_model="bare-model")

    assert _capabilities(_deployment(solo={}))["model"] == "bare-model"


# ── Capability resolution ────────────────────────────────────────────────────


def test_agent_capabilities_override_global(workspace: Path) -> None:
    """A per-agent model_capabilities block replaces the global one."""
    _write_fastagent(
        workspace,
        default_model="openai.global-model",
        model_capabilities={"vision": False, "context_window": 200000},
    )

    caps = _capabilities(
        _deployment(
            solo={
                "model": "anthropic.agent-model",
                "model_capabilities": {"vision": True, "context_window": 400000},
            }
        )
    )

    assert caps["vision"] is True
    assert caps["context_window"] == 400000


def test_effective_defaults_when_capabilities_absent(workspace: Path) -> None:
    """Absent model_capabilities publishes the values actually registered.

    Regression: these were previously advertised as ``null`` while
    ``server._register_one_model`` registered 131072 / 16384.
    """
    _write_fastagent(workspace, default_model="openai.some-model")

    caps = _capabilities(_deployment(solo={}))

    assert caps["context_window"] == DEFAULT_CONTEXT_WINDOW
    assert caps["max_output_tokens"] == DEFAULT_MAX_OUTPUT_TOKENS
    assert caps["vision"] is False


def test_global_capabilities_are_published(workspace: Path) -> None:
    """The ordinary case: one global model + capabilities for every agent."""
    _write_fastagent(
        workspace,
        default_model="anthropic.claude-opus-4-7",
        model_capabilities={
            "vision": True,
            "context_window": 200000,
            "max_output_tokens": 50000,
        },
    )

    caps = _capabilities(_deployment(solo={}))

    assert caps == {
        "model": "claude-opus-4-7",
        "vision": True,
        "context_window": 200000,
        "max_output_tokens": 50000,
    }


# ── Per-agent independence ───────────────────────────────────────────────────


def test_agents_resolve_independently(workspace: Path) -> None:
    """Each entry gets its own capabilities — the original shared-dict bug."""
    _write_fastagent(
        workspace,
        default_model="openai.global-model",
        model_capabilities={"vision": False, "context_window": 128000},
    )

    entries = _entries(
        _deployment(
            inherits={},
            overrides={
                "model": "anthropic.special-model",
                "model_capabilities": {"vision": True, "context_window": 500000},
            },
        )
    )

    assert entries["inherits"]["capabilities"]["model"] == "global-model"
    assert entries["inherits"]["capabilities"]["context_window"] == 128000
    assert entries["overrides"]["capabilities"]["model"] == "special-model"
    assert entries["overrides"]["capabilities"]["context_window"] == 500000


# ── Omission ─────────────────────────────────────────────────────────────────


def test_capabilities_omitted_without_any_model(workspace: Path) -> None:
    """No fastagent.config.yaml and no per-agent model → no capabilities key."""
    entry = _entries(_deployment(solo={}))["solo"]

    assert "capabilities" not in entry


def test_agent_model_published_without_fastagent_config(workspace: Path) -> None:
    """An agent's own model is advertised even with no fastagent.config.yaml."""
    caps = _capabilities(_deployment(solo={"model": "anthropic.agent-model"}))

    assert caps["model"] == "agent-model"
    assert caps["context_window"] == DEFAULT_CONTEXT_WINDOW