From 0638a126f75550b98df0eec3bbf22344edfc5b2c Mon Sep 17 00:00:00 2001 From: Arthur Date: Tue, 6 Oct 2026 15:12:53 +0200 Subject: [PATCH] fix(ai): send the API key when waiting for oMLX to start oMLX 0.7 answers `/v1/models` with 401 unless the request carries the API key. The install's readiness check sent none, so every managed install waited the full 30 seconds, warned that oMLX was not responding, and then reported it running. The check now sends the key the install writes into oMLX's settings. --- stacklets/ai/hooks/on_install.py | 25 ++++++++++----- tests/unit/stacklets/test_ai_omlx_install.py | 33 ++++++++++++++++++++ 2 files changed, 50 insertions(+), 8 deletions(-) diff --git a/stacklets/ai/hooks/on_install.py b/stacklets/ai/hooks/on_install.py index 6263a1d5..18dad130 100644 --- a/stacklets/ai/hooks/on_install.py +++ b/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}") diff --git a/tests/unit/stacklets/test_ai_omlx_install.py b/tests/unit/stacklets/test_ai_omlx_install.py index 3be260d3..8d0b50a3 100644 --- a/tests/unit/stacklets/test_ai_omlx_install.py +++ b/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)