fix(ai): send the API key when waiting for oMLX to start

stacklets/ai/hooks/on_install.py

@@ -108,6 +108,22 @@ def _install_omlx_formula(ctx) -> None:
     ctx.shell_live(f"{BREW_INSTALL} omlx --with-grammar")
 
 
+def wait_for_omlx(ctx, base_url: str, key: str, timeout: float = 30) -> bool:
+    """Wait until oMLX lists its models; False when it does not in time.
+
+    The request carries the key the install wrote into oMLX's settings:
+    oMLX answers 401 to a request without it.
+    """
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        try:
+            ctx.shell(f'curl -sf --max-time 2 -H "Authorization: Bearer {key}" {base_url}/models')
+            return True
+        except RuntimeError:
+            time.sleep(2)
+    return False
+
+
 def install_omlx(ctx, state_dir: Path):
     section("oMLX", "MLX inference (Metal GPU)")
 
@@ -153,14 +169,7 @@ def install_omlx(ctx, state_dir: Path):
     (state_dir / "omlx-managed").touch()
 
     ctx.step("Waiting for oMLX to start...")
-    deadline = time.time() + 30
-    while time.time() < deadline:
-        try:
-            ctx.shell(f'curl -sf --max-time 2 http://localhost:{OMLX_PORT}/v1/models')
-            break
-        except RuntimeError:
-            time.sleep(2)
-    else:
+    if not wait_for_omlx(ctx, f"http://localhost:{OMLX_PORT}/v1", settings["auth"]["api_key"]):
         warn(f"oMLX not responding on port {OMLX_PORT} — it may still be starting")
 
     done(f"oMLX running on port {OMLX_PORT}")

tests/unit/stacklets/test_ai_omlx_install.py

@@ -115,3 +115,36 @@ def test_a_homebrew_without_a_trust_gate_still_installs(self):
         ctx = FakeCtx(no_trust_subcommand=True)
         _install_omlx_formula(ctx)  # must not raise
         assert any("brew install omlx" in c for c in ctx.commands)
+
+
+class ShellCtx:
+    """Runs the hook's shell commands for real, as `ctx.shell` does."""
+
+    def shell(self, cmd):
+        import subprocess
+        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+        if result.returncode != 0:
+            raise RuntimeError(cmd)
+        return result.stdout
+
+
+class TestWaitingForOmlx:
+    """oMLX 0.7 answers `/v1/models` with 401 unless the request carries the key.
+
+    The install writes that key into oMLX's settings, so its readiness
+    check has to send it, or it waits out its whole timeout on a server
+    that is up and warns that it is not.
+    """
+
+    def test_a_server_that_wants_the_key_counts_as_up(self, httpserver):
+        httpserver.expect_request(
+            "/v1/models", headers={"Authorization": "Bearer local"},
+        ).respond_with_json({"data": []})
+        httpserver.expect_request("/v1/models").respond_with_data("API key required", status=401)
+
+        assert _ai_on_install.wait_for_omlx(ShellCtx(), httpserver.url_for("/v1"), "local", timeout=5)
+
+    def test_a_server_that_refuses_the_key_is_not_up(self, httpserver):
+        httpserver.expect_request("/v1/models").respond_with_data("API key required", status=401)
+
+        assert not _ai_on_install.wait_for_omlx(ShellCtx(), httpserver.url_for("/v1"), "local", timeout=1)