Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tests/hermes_cli/test_load_progress.py

@@ -0,0 +1,260 @@
+"""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
+"provider may be slow or overloaded" stall warning. 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 time
+
+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_helpers 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
+
+
+def test_load_notice_matches_desktop_wait_filter():
+    """The notices must pass the desktop's providerWaitText regex and parse
+    under parseModelLoadWait's shapes — pinned here as plain string
+    contracts so the two sides can't drift silently."""
+    import re
+
+    accept = r"^(?:⏳|⚠|↻|⚙)\s*(?:waiting on|loading|processing prompt|no (?:output|response)|model returned)"
+
+    load = "⏳ loading Qwen3.6-35B-A3B-UD-Q4_K_M into memory — 43% (responses start once the model is loaded)"
+    assert re.match(accept, load)
+    m = re.match(r"^⏳\s*loading\s+(.+?)\s+into memory\s+—\s+(\d{1,3})%", load)
+    assert m and m.group(1) == "Qwen3.6-35B-A3B-UD-Q4_K_M" and m.group(2) == "43"
+
+    prefill = "⚙ processing prompt — 31%"
+    assert re.match(accept, prefill)
+    p = re.match(r"^⚙\s*processing prompt(?:\s+—\s+(\d{1,3})%)?", prefill)
+    assert p and p.group(1) == "31"
+
+    bare = "⚙ processing prompt"
+    assert re.match(accept, bare)
+    b = re.match(r"^⚙\s*processing prompt(?:\s+—\s+(\d{1,3})%)?", bare)
+    assert b and b.group(1) is None
+
+
+# ── prefill progress ─────────────────────────────────────────
+
+
+def test_prefill_notice_for_managed_model(tmp_path, monkeypatch):
+    from agent.chat_completion_helpers 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_helpers 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")