hermes-agent/tests/hermes_cli/test_load_progress.py

"""Model-load progress: SSE events -> composite percent -> wait notices.

The 40-second problem: a cold local model streams 16-21 GB of weights
before the first token, and the chat rendered that as the generic
"waiting on <model>" long-wait notice. llama-server's child
emits real per-tensor progress which the router relays over /models/sse
ONLY — these tests pin the consumer that turns that stream into the
status route's `loading` field and the chat's load notice."""

from __future__ import annotations

import json

import hermes_cli.local_runtime.load_progress as lp


def setup_function(_fn):
    with lp._lock:
        lp._snapshot.clear()


# ── composite percent ────────────────────────────────────────


def test_composite_percent_text_stage_dominates():
    stages = ["text_model", "spec_model", "mmproj_model"]
    # Text model owns [0, 85): halfway through it reads ~42%.
    assert lp._composite_percent(stages, "text_model", 0.5) == 42  # 0.5*85
    # Extras start where text ends and never regress below it.
    assert lp._composite_percent(stages, "spec_model", 0.0) == 85
    assert lp._composite_percent(stages, "mmproj_model", 1.0) == 100


def test_composite_percent_monotone_across_stage_walk():
    """Walking the stages in llama-server's real order never moves the
    bar backwards — the property that makes the bar trustworthy."""
    stages = ["text_model", "spec_model", "mmproj_model"]
    walk = [("text_model", v / 10) for v in range(11)] + \
           [("spec_model", v / 10) for v in range(11)] + \
           [("mmproj_model", v / 10) for v in range(11)]
    seen = [lp._composite_percent(stages, s, v) for s, v in walk]
    assert seen == sorted(seen)
    assert seen[0] == 0 and seen[-1] == 100


def test_composite_percent_single_stage_is_plain():
    assert lp._composite_percent(["text_model"], "text_model", 0.4) == 40


# ── event application ────────────────────────────────────────


def _loading_event(value: float, current: str = "text_model") -> dict:
    return {"status": "loading",
            "progress": {"stages": ["text_model", "mmproj_model"],
                         "current": current, "value": value}}


def test_loading_events_build_snapshot_and_terminal_clears():
    lp._apply_event("m1", "status_change", _loading_event(0.5))
    snap = lp.get_loading_progress()
    assert "m1" in snap
    assert snap["m1"]["percent"] == 42  # 0.5 * 85 within text stage
    assert snap["m1"]["stage"] == "text_model"

    lp._apply_event("m1", "status_change", {"status": "loaded", "info": {}})
    assert lp.get_loading_progress() == {}


def test_unload_and_failure_clear_too():
    lp._apply_event("m1", "status_change", _loading_event(0.2))
    lp._apply_event("m1", "status_change", {"status": "unloaded", "exit_code": 1})
    assert lp.get_loading_progress() == {}

    lp._apply_event("m2", "status_change", _loading_event(0.9))
    lp._apply_event("m2", "model_remove", {})
    assert lp.get_loading_progress() == {}


def test_progressless_loading_event_keeps_entry_alive():
    """The router's first model_status event says just {status: loading} —
    it must register the load (indeterminate) without inventing a percent."""
    lp._apply_event("m1", "model_status", {"status": "loading"})
    snap = lp.get_loading_progress()
    assert snap["m1"]["percent"] == 0


def test_stale_entries_expire():
    lp._apply_event("m1", "status_change", _loading_event(0.5))
    with lp._lock:
        lp._snapshot["m1"]["ts"] -= lp._STALE_ENTRY_TTL_S + 1
    assert lp.get_loading_progress() == {}


# ── chat wait-notice ─────────────────────────────────────────


def test_load_notice_for_managed_model(tmp_path, monkeypatch):
    from agent.chat_completion_local import _managed_local_load_notice

    state = tmp_path / "server.json"
    state.write_text(json.dumps({"base_url": "http://127.0.0.1:18434/v1",
                                 "api_key": "k"}), encoding="utf-8")
    monkeypatch.setattr("hermes_cli.local_runtime.supervisor.state_path",
                        lambda: state)
    lp._apply_event("Qwen-Test", "status_change", _loading_event(0.5))
    monkeypatch.setattr(lp, "_ensure_watcher", lambda: None)

    class _Agent:
        base_url = "http://127.0.0.1:18434/v1"

    notice = _managed_local_load_notice(_Agent(), {"model": "Qwen-Test"})
    assert notice is not None
    assert notice.startswith("⏳ loading Qwen-Test into memory — 42%")

    # Different endpoint (user's own server): never claim its loads.
    class _Other:
        base_url = "http://127.0.0.1:9999/v1"

    assert _managed_local_load_notice(_Other(), {"model": "Qwen-Test"}) is None
    # Managed endpoint but a model that isn't loading: no notice.
    assert _managed_local_load_notice(_Agent(), {"model": "Elsewhere"}) is None




# ── prefill progress ─────────────────────────────────────────


def test_prefill_notice_for_managed_model(tmp_path, monkeypatch):
    from agent.chat_completion_local import _managed_local_load_notice

    state = tmp_path / "server.json"
    state.write_text(json.dumps({"base_url": "http://127.0.0.1:18434/v1",
                                 "api_key": "k"}), encoding="utf-8")
    monkeypatch.setattr("hermes_cli.local_runtime.supervisor.state_path",
                        lambda: state)
    monkeypatch.setattr(lp, "_ensure_watcher", lambda: None)
    # No load in flight; a prefill counter is live.
    monkeypatch.setattr(lp, "get_prefill_progress",
                        lambda model: {"processed": 12288})
    import agent.chat_completion_helpers as cch

    monkeypatch.setattr(cch, "estimate_request_context_tokens",
                        lambda kw: 39551)

    class _Agent:
        base_url = "http://127.0.0.1:18434/v1"

    notice = _managed_local_load_notice(_Agent(), {"model": "Qwen-Test"})
    assert notice == "⚙ processing prompt — 31%"

    # Counter past the estimate (estimator undercounted): no honest
    # denominator, so no percent — never >100%.
    monkeypatch.setattr(cch, "estimate_request_context_tokens", lambda kw: 100)
    notice = _managed_local_load_notice(_Agent(), {"model": "Qwen-Test"})
    assert notice == "⚙ processing prompt"


def test_load_notice_outranks_prefill(tmp_path, monkeypatch):
    """While a load entry exists the load notice wins — prefill can't start
    before the model is resident, so a simultaneous claim means the load
    snapshot is authoritative."""
    from agent.chat_completion_local import _managed_local_load_notice

    state = tmp_path / "server.json"
    state.write_text(json.dumps({"base_url": "http://127.0.0.1:18434/v1",
                                 "api_key": "k"}), encoding="utf-8")
    monkeypatch.setattr("hermes_cli.local_runtime.supervisor.state_path",
                        lambda: state)
    monkeypatch.setattr(lp, "_ensure_watcher", lambda: None)
    lp._apply_event("Qwen-Test", "status_change", _loading_event(0.5))
    monkeypatch.setattr(lp, "get_prefill_progress",
                        lambda model: {"processed": 999})

    class _Agent:
        base_url = "http://127.0.0.1:18434/v1"

    notice = _managed_local_load_notice(_Agent(), {"model": "Qwen-Test"})
    assert notice is not None and notice.startswith("⏳ loading")


def test_prefill_progress_reads_busiest_processing_slot(monkeypatch):
    monkeypatch.setattr(lp, "_endpoint", lambda: ("http://127.0.0.1:1", "k"))

    class _Resp:
        def __init__(self, payload):
            self._payload = payload

        def read(self):
            return json.dumps(self._payload).encode()

        def __enter__(self):
            return self

        def __exit__(self, *a):
            return False

    slots = [
        {"id": 0, "is_processing": False, "n_prompt_tokens_processed": 500},
        {"id": 1, "is_processing": True, "n_prompt_tokens_processed": 42},
        {"id": 2, "is_processing": True, "n_prompt_tokens_processed": 32768},
    ]
    monkeypatch.setattr(lp.urllib.request, "urlopen",
                        lambda req, timeout=0: _Resp(slots))
    assert lp.get_prefill_progress("m") == {"processed": 32768}

    # Nothing processing -> None (idle slots' counters are leftovers).
    idle = [{"id": 0, "is_processing": False, "n_prompt_tokens_processed": 500}]
    monkeypatch.setattr(lp.urllib.request, "urlopen",
                        lambda req, timeout=0: _Resp(idle))
    assert lp.get_prefill_progress("m") is None

    # Unreachable server -> None, never an exception.
    def _boom(req, timeout=0):
        raise OSError("refused")

    monkeypatch.setattr(lp.urllib.request, "urlopen", _boom)
    assert lp.get_prefill_progress("m") is None


def test_endpoint_respects_ownership_guard(monkeypatch):
    """The watcher's endpoint MUST come from the ownership-guarded reader.
    Regression: a raw state-file read attached the SSE watcher to a
    foreign install's server on the shared stable port (health answers
    for anyone; only the dead-pid check proves ownership)."""
    import hermes_cli.local_runtime.load_progress as lp

    # Guard says "not ours": no endpoint, regardless of state on disk.
    monkeypatch.setattr("hermes_cli.local_runtime.endpoint._state_endpoint",
                        lambda: None)
    assert lp._endpoint() is None

    monkeypatch.setattr(
        "hermes_cli.local_runtime.endpoint._state_endpoint",
        lambda: {"base_url": "http://127.0.0.1:18434/v1", "api_key": "k"})
    assert lp._endpoint() == ("http://127.0.0.1:18434", "k")