turnstone/tests/test_server_node_models_metadata.py

"""Tests for the per-node ``models`` metadata pipeline.

Two helpers in ``server.py`` carry the load:

- ``_collect_node_models_metadata`` projects the live ``ModelRegistry``
  into the node_metadata row shape ``[{alias, provider, healthy}, ...]``.
- ``_publish_models_metadata`` short-circuits redundant writes via a
  payload cache on ``app_state`` and is the helper called from both
  the heartbeat loop and ``internal_model_reload``.

These tests pin the projection shape, the health-flag wiring, the
cache short-circuit, and the model-reload integration.
"""

from __future__ import annotations

import json
from types import SimpleNamespace
from unittest.mock import MagicMock

from turnstone.core.healthcheck import HealthTrackerRegistry
from turnstone.core.model_registry import ModelConfig, ModelRegistry
from turnstone.server import (
    _collect_node_models_metadata,
    _publish_models_metadata,
)


def _registry(*aliases_with_url: tuple[str, str]) -> ModelRegistry:
    """Build a registry from ``(alias, base_url)`` pairs.

    Two aliases sharing a ``base_url`` deliberately share a tracker —
    that's the contract the cluster-level health surface needs to
    preserve, and it's worth pinning in a test.
    """
    models = {
        alias: ModelConfig(alias=alias, base_url=url, api_key="k", model=alias, provider="openai")
        for alias, url in aliases_with_url
    }
    default = aliases_with_url[0][0]
    return ModelRegistry(models, default=default)


def test_returns_none_when_registry_missing():
    state = SimpleNamespace()
    assert _collect_node_models_metadata(state) is None


def test_projects_all_aliases_with_default_healthy_when_no_tracker():
    """Without a ``health_registry`` (or before any request has flowed
    through a backend), every alias surfaces as ``healthy=True`` —
    operators shouldn't get an empty ``models`` list on a freshly
    started node just because the backends haven't been exercised."""
    reg = _registry(("a", "http://x"), ("b", "http://y"))
    state = SimpleNamespace(registry=reg)
    entry = _collect_node_models_metadata(state)
    assert entry is not None
    key, value, source = entry
    assert key == "models"
    assert source == "auto"
    rows = json.loads(value)
    assert len(rows) == 2
    aliases = {r["alias"] for r in rows}
    assert aliases == {"a", "b"}
    assert all(r["healthy"] is True for r in rows)
    assert all(r["provider"] == "openai" for r in rows)
    # Provider-side model identifier intentionally omitted — coords
    # kept passing it as ``spawn_workstream(model=...)`` when they
    # should have passed the local alias.  Lock the projected keys
    # so a future contributor doesn't reintroduce the footgun.
    for row in rows:
        assert set(row.keys()) == {"alias", "provider", "healthy"}


def test_health_flag_reflects_tracker_state():
    reg = _registry(("a", "http://x"), ("b", "http://y"))
    health_reg = HealthTrackerRegistry(failure_threshold=2)
    # Seed the tracker for "a"'s backend and drive it into the degraded
    # state — two consecutive failures cross the threshold.
    bad_tracker = health_reg.get_tracker(provider="openai", base_url="http://x")
    bad_tracker.record_failure()
    bad_tracker.record_failure()
    assert bad_tracker.is_degraded
    # "b" gets a tracker that has only seen successes.
    good_tracker = health_reg.get_tracker(provider="openai", base_url="http://y")
    good_tracker.record_success()
    state = SimpleNamespace(registry=reg, health_registry=health_reg)
    rows = json.loads(_collect_node_models_metadata(state)[1])
    by_alias = {r["alias"]: r for r in rows}
    assert by_alias["a"]["healthy"] is False
    assert by_alias["b"]["healthy"] is True


def test_two_aliases_sharing_a_backend_share_a_tracker():
    """Two aliases that point at the same ``(provider, base_url)``
    share a single :class:`BackendHealthTracker` — degrading one is
    expected to surface as degraded on the other.  The list_nodes
    projection should respect that, otherwise a coord could see
    ``alias-a`` healthy and ``alias-b`` degraded for the same
    backend."""
    reg = _registry(("alpha", "http://shared"), ("beta", "http://shared"))
    health_reg = HealthTrackerRegistry(failure_threshold=1)
    tracker = health_reg.get_tracker(provider="openai", base_url="http://shared")
    tracker.record_failure()  # threshold=1 — degraded immediately
    state = SimpleNamespace(registry=reg, health_registry=health_reg)
    rows = json.loads(_collect_node_models_metadata(state)[1])
    assert {r["alias"]: r["healthy"] for r in rows} == {"alpha": False, "beta": False}


def test_alias_with_no_tracker_yet_defaults_to_healthy():
    """An alias the registry knows about but whose backend hasn't been
    invoked yet has no tracker.  Default to healthy so a brand-new
    alias is immediately visible to coordinators rather than waiting
    for the first request to seed a tracker.

    The collector calls ``health_reg.get_tracker(...)`` which mints a
    fresh tracker on first lookup — that's the path under test here.
    The freshly minted tracker reports ``is_healthy=True`` (default
    state), so the projection labels the alias healthy.
    """
    reg = _registry(("a", "http://x"))
    health_reg = HealthTrackerRegistry()  # empty — no trackers seeded
    state = SimpleNamespace(registry=reg, health_registry=health_reg)
    rows = json.loads(_collect_node_models_metadata(state)[1])
    assert rows[0]["healthy"] is True


# ---------------------------------------------------------------------------
# _publish_models_metadata — cache short-circuit + projection wiring
# ---------------------------------------------------------------------------


def _publish_state() -> SimpleNamespace:
    """Build an ``app_state`` with a minimal registry + health surface."""
    reg = _registry(("a", "http://x"))
    return SimpleNamespace(registry=reg, health_registry=HealthTrackerRegistry())


def test_publish_writes_when_payload_changes():
    """First publish has nothing in the cache — write happens; cache
    fills.  Second publish on the same unchanged registry skips the
    write entirely."""
    state = _publish_state()
    storage = MagicMock()
    _publish_models_metadata(state, storage, "node-a")
    assert storage.set_node_metadata_bulk.call_count == 1
    cached = state._last_models_payload
    assert isinstance(cached, str) and "alias" in cached
    # Second call, same registry, same health: cached payload matches
    # — write must be skipped to avoid the per-30s UPSERT churn.
    _publish_models_metadata(state, storage, "node-a")
    assert storage.set_node_metadata_bulk.call_count == 1


def test_publish_records_metric_outcome(monkeypatch):
    """The publish helper feeds ``record_node_models_publish`` so
    Prometheus can expose the hit-rate.  Storage failures must NOT
    record either outcome — counters should reflect actual cache
    decisions, not transient DB errors that will retry.

    Replaces the module-level ``turnstone.server._metrics`` binding
    via string-form monkeypatch (with auto-restore) rather than
    patching an instance attribute on the imported singleton.  Other
    tests in the suite reassign ``srv_mod._metrics`` (some without
    using monkeypatch), so an instance captured at import time can
    diverge from the binding the live ``_publish_models_metadata``
    reads on each call.
    """
    state = _publish_state()
    storage = MagicMock()
    calls: list[bool] = []

    class _FakeMetrics:
        def record_node_models_publish(self, *, written: bool) -> None:
            calls.append(written)

    monkeypatch.setattr("turnstone.server._metrics", _FakeMetrics())

    _publish_models_metadata(state, storage, "node-a")  # first → write
    _publish_models_metadata(state, storage, "node-a")  # second → skip
    assert calls == [True, False]

    # Storage error: no metric recorded.
    storage.set_node_metadata_bulk.side_effect = RuntimeError("db down")
    state._last_models_payload = None  # invalidate cache to force a write attempt
    _publish_models_metadata(state, storage, "node-a")
    assert calls == [True, False]  # unchanged


def test_publish_rewrites_when_health_flips():
    """A health-tracker state change must invalidate the cache and
    drive a fresh write — otherwise the discovery surface would lag
    a flip indefinitely."""
    state = _publish_state()
    storage = MagicMock()
    _publish_models_metadata(state, storage, "node-a")
    assert storage.set_node_metadata_bulk.call_count == 1
    # Drive the only tracker to degraded.
    tracker = state.health_registry.get_tracker(provider="openai", base_url="http://x")
    for _ in range(10):
        tracker.record_failure()
    assert tracker.is_degraded
    _publish_models_metadata(state, storage, "node-a")
    assert storage.set_node_metadata_bulk.call_count == 2


def test_publish_swallows_storage_error_without_updating_cache():
    """A storage failure must NOT poison the cache — the next call
    should retry the write rather than think it succeeded."""
    state = _publish_state()
    storage = MagicMock()
    storage.set_node_metadata_bulk.side_effect = RuntimeError("db down")
    _publish_models_metadata(state, storage, "node-a")
    assert storage.set_node_metadata_bulk.call_count == 1
    assert getattr(state, "_last_models_payload", None) is None
    # Recover: a subsequent successful call writes again.
    storage.set_node_metadata_bulk.side_effect = None
    _publish_models_metadata(state, storage, "node-a")
    assert storage.set_node_metadata_bulk.call_count == 2
    assert state._last_models_payload is not None


def test_publish_skips_when_registry_missing():
    """Without a registry there's nothing to project; nothing should
    be written and the cache must not be set."""
    state = SimpleNamespace()
    storage = MagicMock()
    _publish_models_metadata(state, storage, "node-a")
    assert storage.set_node_metadata_bulk.call_count == 0
    assert getattr(state, "_last_models_payload", None) is None


# ---------------------------------------------------------------------------
# internal_model_reload — integration: registry change must rewrite the row
# ---------------------------------------------------------------------------


def test_model_reload_endpoint_rewrites_models_metadata(monkeypatch, tmp_path):
    """A successful ``internal_model_reload`` must refresh
    ``node_metadata.models`` so a coordinator sees the new alias on
    its next ``list_nodes`` without waiting up to 30s for the
    heartbeat tick.

    The endpoint pulls a fresh registry from
    ``load_model_registry(...)`` and reloads in-place — we stub the
    loader to return a registry with a different alias set so the
    publish-cache invalidation is exercised end-to-end.
    """
    from turnstone.core.storage._sqlite import SQLiteBackend
    from turnstone.server import internal_model_reload

    storage = SQLiteBackend(str(tmp_path / "reload.db"))

    # Old registry — single alias "a".
    old_reg = _registry(("a", "http://x"))
    # New registry that ``load_model_registry`` will return — adds "b".
    new_reg = ModelRegistry(
        {
            "a": ModelConfig(
                alias="a", base_url="http://x", api_key="k", model="a", provider="openai"
            ),
            "b": ModelConfig(
                alias="b", base_url="http://y", api_key="k", model="b", provider="openai"
            ),
        },
        default="a",
    )
    health_reg = HealthTrackerRegistry()

    app_state = SimpleNamespace(
        registry=old_reg,
        health_registry=health_reg,
        cli_model_args={
            "base_url": "",
            "api_key": "",
            "model": "",
            "context_window": 0,
            "provider": "openai",
        },
        config_store=None,
        node_id="node-a",
    )
    request = SimpleNamespace(app=SimpleNamespace(state=app_state))

    # Patch the loader and storage accessors used inside the endpoint.
    # ``internal_model_reload`` does ``from turnstone.core.storage._registry
    # import get_storage`` inline, so patching the symbol on that module
    # is what intercepts the call.
    monkeypatch.setattr("turnstone.core.model_registry.load_model_registry", lambda **_kw: new_reg)
    monkeypatch.setattr("turnstone.core.storage._registry.get_storage", lambda: storage)
    # The endpoint also broadcasts schema refreshes to active sessions
    # — stub this out, it's irrelevant to the metadata-write path.
    monkeypatch.setattr("turnstone.server._broadcast_agent_tool_schema_refresh", lambda _s: None)

    response = internal_model_reload(request)  # type: ignore[arg-type]
    assert response.status_code == 200

    rows = storage.get_node_metadata("node-a")
    by_key = {r["key"]: r for r in rows}
    assert "models" in by_key
    payload = json.loads(by_key["models"]["value"])
    assert {r["alias"] for r in payload} == {"a", "b"}


def test_model_reload_refuses_dynamic_auth_without_key(monkeypatch, tmp_path, caplog):
    """A keyless node cannot acquire a dynamic alias via model-reload: 503,
    a deployment fault, not the 422 bad-arguments exit."""
    from turnstone.core.storage._sqlite import SQLiteBackend
    from turnstone.server import internal_model_reload

    storage = SQLiteBackend(str(tmp_path / "reload.db"))

    old_reg = _registry(("a", "http://x"))
    new_reg = ModelRegistry(
        {
            "a": ModelConfig(
                alias="a", base_url="http://x", api_key="k", model="a", provider="openai"
            ),
            "gw": ModelConfig(
                alias="gw",
                base_url="http://gw",
                api_key="",
                model="m",
                provider="openai",
                auth_mode="entra_obo",
                obo_audience="api://gateway",
            ),
        },
        default="a",
    )
    app_state = SimpleNamespace(
        registry=old_reg,
        health_registry=HealthTrackerRegistry(),
        cli_model_args={
            "base_url": "",
            "api_key": "",
            "model": "",
            "context_window": 0,
            "provider": "openai",
        },
        config_store=None,
        node_id="node-a",
        mcp_token_store=None,
    )
    request = SimpleNamespace(app=SimpleNamespace(state=app_state))

    monkeypatch.setattr("turnstone.core.model_registry.load_model_registry", lambda **_kw: new_reg)
    monkeypatch.setattr("turnstone.core.storage._registry.get_storage", lambda: storage)
    monkeypatch.setattr("turnstone.server._broadcast_agent_tool_schema_refresh", lambda _s: None)

    with caplog.at_level("ERROR", logger="turnstone.server"):
        response = internal_model_reload(request)  # type: ignore[arg-type]

    assert response.status_code == 503
    assert "mcp_token_encryption" in json.loads(response.body)["reason"]
    # Refusal must not mutate: the old registry keeps serving.
    assert not old_reg.has_alias("gw")
    assert not old_reg.has_dynamic_auth()
    assert any("model_auth_key_missing" in r.message for r in caplog.records)


def test_model_reload_maps_auth_config_error_to_422(monkeypatch, tmp_path):
    """A row whose auth fields the loader rejects exits as the structured
    422 naming the problem, never a bare 500."""
    from turnstone.core.model_registry import ModelAuthConfigError
    from turnstone.core.storage._sqlite import SQLiteBackend
    from turnstone.server import internal_model_reload

    storage = SQLiteBackend(str(tmp_path / "reload.db"))
    old_reg = _registry(("a", "http://x"))
    app_state = SimpleNamespace(
        registry=old_reg,
        health_registry=HealthTrackerRegistry(),
        cli_model_args={
            "base_url": "",
            "api_key": "",
            "model": "",
            "context_window": 0,
            "provider": "openai",
        },
        config_store=None,
        node_id="node-a",
    )
    request = SimpleNamespace(app=SimpleNamespace(state=app_state))

    def _raise_auth_config(**_kw):
        raise ModelAuthConfigError("Model 'gw' requires obo_audience when auth_mode is 'entra_obo'")

    monkeypatch.setattr("turnstone.core.model_registry.load_model_registry", _raise_auth_config)
    monkeypatch.setattr("turnstone.core.storage._registry.get_storage", lambda: storage)

    response = internal_model_reload(request)  # type: ignore[arg-type]

    assert response.status_code == 422
    assert "obo_audience" in json.loads(response.body)["reason"]
    assert old_reg.has_alias("a")


def test_config_reload_maps_dynamic_auth_key_error_to_503(monkeypatch, tmp_path, caplog):
    """Every caller of the reload chokepoint handles its refusal the same
    way: the settings fan-out answers 503, not an unhandled 500."""
    from turnstone.server import config_reload

    old_reg = ModelRegistry(
        {
            "a": ModelConfig(
                alias="a", base_url="http://x", api_key="k", model="a", provider="openai"
            ),
            "gw": ModelConfig(
                alias="gw",
                base_url="http://gw",
                api_key="",
                model="m",
                provider="openai",
                auth_mode="entra_obo",
                obo_audience="api://gateway",
            ),
        },
        default="a",
    )
    config_store = MagicMock()
    config_store.get.side_effect = lambda key, default=None: (
        "gw" if key == "model.default_alias" else default
    )
    app_state = SimpleNamespace(
        registry=old_reg,
        config_store=config_store,
        mcp_token_store=None,
    )
    request = SimpleNamespace(app=SimpleNamespace(state=app_state))

    with caplog.at_level("ERROR", logger="turnstone.server"):
        response = config_reload(request)  # type: ignore[arg-type]

    assert response.status_code == 503
    assert "mcp_token_encryption" in json.loads(response.body)["reason"]
    # Refused without mutating: the routing override did NOT apply.
    assert old_reg.default == "a"
    assert any("model_auth_key_missing" in r.message for r in caplog.records)


# ---------------------------------------------------------------------------
# Shutdown race: heartbeat write must NOT resurrect post-shutdown delete
# ---------------------------------------------------------------------------


def test_heartbeat_write_awaits_before_shutdown_delete():
    """Pin the shutdown-race fix.

    Before the fix, the lifespan shutdown sequence was:

    1. ``_heartbeat_task.cancel()`` — fire-and-forget
    2. ``delete_node_metadata_by_source(node_id, "auto")``

    A heartbeat tick already inside ``asyncio.to_thread(...)`` for
    the ``set_node_metadata_bulk`` call would complete AFTER step 2,
    resurrecting the deleted ``models`` row.  The fix awaits the
    cancelled task with ``contextlib.suppress(...)`` between (1) and
    (2), so the in-flight write lands first.

    We verify the fix by introspecting ``server.py`` source — the
    real lifespan is hard to test deterministically without a full
    Starlette app, but the textual ordering between
    ``_heartbeat_task.cancel()`` and the delete is a stable contract
    that catches the regression cheaply.
    """
    import inspect
    import sys

    src = inspect.getsource(sys.modules[_collect_node_models_metadata.__module__])
    cancel_idx = src.find("_heartbeat_task.cancel()")
    delete_idx = src.find('delete_node_metadata_by_source, _svc_node_id, "auto"')
    await_idx = src.find("await _heartbeat_task", cancel_idx)
    assert cancel_idx != -1
    assert delete_idx != -1
    assert await_idx != -1
    # The fix-line must sit BETWEEN the cancel and the delete.
    assert cancel_idx < await_idx < delete_idx, (
        "Shutdown race regression: "
        "_heartbeat_task.cancel() must be followed by `await _heartbeat_task` "
        "BEFORE delete_node_metadata_by_source(..., 'auto') so an in-flight "
        "set_node_metadata_bulk lands before the delete."
    )