fix: make headroom wrap readiness probe timeout configurable for slow ML imports (#581)

## Summary

Makes `headroom wrap` wait long enough for slow proxy startups instead
of failing at a fixed readiness window, with an ML-aware default and an
env-var override.

## Why

`headroom wrap` failed when the proxy took longer than a fixed startup
window to bind its port. Issue #195 reports that on ML-heavy setups the
proxy imports large libraries (torch, sentence_transformers, spacy) at
startup and routinely exceeds the hardcoded window, so `wrap` aborts on
a working proxy and the failure message gives no way to extend the wait.

## Description

`headroom wrap` now lets slow proxy startups finish instead of failing
at a fixed window. The readiness probe in `headroom/cli/wrap.py` reads a
`HEADROOM_WRAP_PROXY_TIMEOUT` environment variable, and when no value is
set it picks the default automatically: 90 seconds when an ML stack
(torch, sentence_transformers, spacy) is detected via
`importlib.util.find_spec` without importing it, otherwise 45 seconds.
The failure message now names the active timeout and the env var to
raise it.

Fixes #195

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- Resolve the wrap proxy readiness window from
`HEADROOM_WRAP_PROXY_TIMEOUT` in `_start_proxy`, falling back to an
ML-aware default.
- Detect optional ML extras with `importlib.util.find_spec` so the check
itself does not pay the cold-import cost the issue describes.
- Include the configured timeout and the env var name in the
`RuntimeError` raised when the proxy genuinely never binds the port.

## Testing

Describe the tests you ran to verify your changes:

- [x] Unit tests pass (`pytest`)
- [ ] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed

New cases in `tests/test_cli_proxy_env.py` cover the default window, an
extended window via the env var, an invalid value raising a clear error,
and the failure message naming the configured timeout. Covered by the
new tests in this PR; full suite runs in CI.

## Test Output

```
# Paste relevant test output here
pytest -v tests/test_cli_proxy_env.py
```

The new `tests/test_cli_proxy_env.py` cases pass locally; the full suite
runs in CI.

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable (N/A: no
CHANGELOG.md is maintained in this repo)

## Screenshots (if applicable)

N/A. This is a CLI startup-timeout fix with no visual surface.

## Additional Notes

The default is conservative: 90s only when an ML stack is detected via
`importlib.util.find_spec` (no import cost), otherwise 45s.
`HEADROOM_WRAP_PROXY_TIMEOUT` overrides both, and the failure message
now names the active timeout and the env var to raise it.

Co-authored-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>

headroom/cli/wrap.py

@@ -14,6 +14,7 @@ Usage:
 
 from __future__ import annotations
 
+import importlib.util
 import io
 import json
 import os
@@ -87,6 +88,10 @@ _CONTEXT_TOOL_ENV = "HEADROOM_CONTEXT_TOOL"
 _CONTEXT_TOOL_RTK = "rtk"
 _CONTEXT_TOOL_LEAN_CTX = "lean-ctx"
 _VALID_CONTEXT_TOOLS = {_CONTEXT_TOOL_RTK, _CONTEXT_TOOL_LEAN_CTX}
+_WRAP_PROXY_TIMEOUT_ENV = "HEADROOM_WRAP_PROXY_TIMEOUT"
+_WRAP_PROXY_TIMEOUT_DEFAULT_SECONDS = 45
+_WRAP_PROXY_TIMEOUT_ML_DEFAULT_SECONDS = 90
+_WRAP_PROXY_TIMEOUT_ML_MODULES = ("torch", "sentence_transformers", "spacy")
 
 # Issue #746: Claude Code disables on-demand tool loading (deferral) when
 # ANTHROPIC_BASE_URL is a custom host and ENABLE_TOOL_SEARCH is unset, which
@@ -171,6 +176,49 @@ def _selected_context_tool() -> str:
     return raw
 
 
+def _module_available(module_name: str) -> bool:
+    """Return whether an optional module is installed without importing it."""
+
+    try:
+        return importlib.util.find_spec(module_name) is not None
+    except (ImportError, ModuleNotFoundError, ValueError):
+        return False
+
+
+def _ml_wrap_extras_detected() -> bool:
+    """Detect slow optional ML stacks without triggering their import cost."""
+
+    return any(_module_available(module_name) for module_name in _WRAP_PROXY_TIMEOUT_ML_MODULES)
+
+
+def _default_wrap_proxy_timeout_seconds() -> int:
+    """Return the default wrap proxy startup timeout for this environment."""
+
+    if _ml_wrap_extras_detected():
+        return _WRAP_PROXY_TIMEOUT_ML_DEFAULT_SECONDS
+    return _WRAP_PROXY_TIMEOUT_DEFAULT_SECONDS
+
+
+def _resolve_wrap_proxy_timeout_seconds() -> int:
+    """Resolve the wrap proxy readiness timeout from env or defaults."""
+
+    raw = os.environ.get(_WRAP_PROXY_TIMEOUT_ENV, "").strip()
+    if not raw:
+        return _default_wrap_proxy_timeout_seconds()
+
+    try:
+        timeout_seconds = int(raw)
+    except ValueError:
+        raise RuntimeError(
+            f"{_WRAP_PROXY_TIMEOUT_ENV} must be a positive integer number of seconds (got {raw!r})"
+        ) from None
+    if timeout_seconds <= 0:
+        raise RuntimeError(
+            f"{_WRAP_PROXY_TIMEOUT_ENV} must be a positive integer number of seconds (got {raw!r})"
+        )
+    return timeout_seconds
+
+
 def _print_telemetry_notice() -> None:
     """Print a telemetry notice when anonymous telemetry is enabled.
 
@@ -288,6 +336,7 @@ def _start_proxy(
     if anthropic_api_url:
         cmd.extend(["--anthropic-api-url", anthropic_api_url])
 
+    timeout_seconds = _resolve_wrap_proxy_timeout_seconds()
     log_path = _get_log_path()
     log_file = open(log_path, "a")  # noqa: SIM115
 
@@ -318,10 +367,10 @@ def _start_proxy(
         start_new_session=os.name == "posix",
     )
 
-    # Wait for proxy to be ready (up to 45 seconds).
+    # Wait for proxy to be ready.
     # ML components (Kompress, Magika, Tree-sitter) load synchronously before
     # uvicorn binds the port. On slower machines this can take 20-30 seconds.
-    for _i in range(45):
+    for _i in range(timeout_seconds):
         time.sleep(1)
         if _check_proxy(port):
             click.echo(f"  Logs: {log_path}")
@@ -338,7 +387,10 @@ def _start_proxy(
 
     proc.kill()
     log_file.close()
-    raise RuntimeError(f"Proxy failed to start on port {port} within 45 seconds")
+    raise RuntimeError(
+        f"Proxy failed to start on port {port} within {timeout_seconds} seconds. "
+        f"Set {_WRAP_PROXY_TIMEOUT_ENV} to a larger number of seconds for slow startup."
+    )
 
 
 def _setup_rtk(verbose: bool = False) -> Path | None:

tests/test_cli_proxy_env.py

@@ -3,6 +3,7 @@
 Verifies that:
 1. Provider target URL env vars are read by `headroom proxy`
 2. litellm-* backends are accepted by both CLI and argparse paths
+3. HEADROOM_WRAP_PROXY_TIMEOUT controls `headroom wrap` proxy readiness waits
 """
 
 import os
@@ -15,6 +16,7 @@ pytest.importorskip("fastapi")
 
 from click.testing import CliRunner  # noqa: E402
 
+from headroom.cli import wrap as wrap_mod  # noqa: E402
 from headroom.cli.main import main  # noqa: E402
 
 
@@ -23,6 +25,98 @@ def runner():
     return CliRunner()
 
 
+class _FakeProxyProcess:
+    returncode = None
+
+    def __init__(self):
+        self.killed = False
+
+    def poll(self):
+        return None
+
+    def kill(self):
+        self.killed = True
+
+
+class TestCLIWrapProxyTimeout:
+    """Test wrap proxy readiness timeout configuration."""
+
+    def test_default_timeout_stays_current_without_ml_extras(self, monkeypatch):
+        monkeypatch.delenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, raising=False)
+        monkeypatch.setattr(wrap_mod, "_ml_wrap_extras_detected", lambda: False)
+
+        assert (
+            wrap_mod._resolve_wrap_proxy_timeout_seconds()
+            == wrap_mod._WRAP_PROXY_TIMEOUT_DEFAULT_SECONDS
+        )
+
+    def test_default_timeout_is_longer_when_ml_extras_detected(self, monkeypatch):
+        monkeypatch.delenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, raising=False)
+        monkeypatch.setattr(wrap_mod, "_ml_wrap_extras_detected", lambda: True)
+
+        assert (
+            wrap_mod._resolve_wrap_proxy_timeout_seconds()
+            == wrap_mod._WRAP_PROXY_TIMEOUT_ML_DEFAULT_SECONDS
+        )
+
+    def test_start_proxy_succeeds_when_ready_within_default_timeout(self, monkeypatch, tmp_path):
+        fake_proc = _FakeProxyProcess()
+        sleeps = []
+
+        monkeypatch.delenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, raising=False)
+        monkeypatch.setattr(wrap_mod, "_ml_wrap_extras_detected", lambda: False)
+        monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
+        monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: True)
+        monkeypatch.setattr(wrap_mod.time, "sleep", lambda seconds: sleeps.append(seconds))
+        monkeypatch.setattr(wrap_mod.subprocess, "Popen", lambda *args, **kwargs: fake_proc)
+
+        proc = wrap_mod._start_proxy(8787, agent_type="codex")
+
+        assert proc is fake_proc
+        assert sleeps == [1]
+        assert fake_proc.killed is False
+
+    def test_env_timeout_allows_slow_start_proxy_to_succeed(self, monkeypatch, tmp_path):
+        fake_proc = _FakeProxyProcess()
+        sleeps = []
+        checks = []
+
+        monkeypatch.setenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, "4")
+        monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
+        monkeypatch.setattr(wrap_mod.time, "sleep", lambda seconds: sleeps.append(seconds))
+        monkeypatch.setattr(wrap_mod.subprocess, "Popen", lambda *args, **kwargs: fake_proc)
+
+        def ready_on_fourth_check(port):
+            checks.append(port)
+            return len(checks) == 4
+
+        monkeypatch.setattr(wrap_mod, "_check_proxy", ready_on_fourth_check)
+
+        proc = wrap_mod._start_proxy(8787, agent_type="codex")
+
+        assert proc is fake_proc
+        assert checks == [8787, 8787, 8787, 8787]
+        assert sleeps == [1, 1, 1, 1]
+        assert fake_proc.killed is False
+
+    def test_timeout_error_names_configured_timeout_and_env_var(self, monkeypatch, tmp_path):
+        fake_proc = _FakeProxyProcess()
+
+        monkeypatch.setenv(wrap_mod._WRAP_PROXY_TIMEOUT_ENV, "2")
+        monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
+        monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: False)
+        monkeypatch.setattr(wrap_mod.time, "sleep", lambda _seconds: None)
+        monkeypatch.setattr(wrap_mod.subprocess, "Popen", lambda *args, **kwargs: fake_proc)
+
+        with pytest.raises(RuntimeError) as excinfo:
+            wrap_mod._start_proxy(8787, agent_type="codex")
+
+        message = str(excinfo.value)
+        assert "within 2 seconds" in message
+        assert wrap_mod._WRAP_PROXY_TIMEOUT_ENV in message
+        assert fake_proc.killed is True
+
+
 class TestCLIProxyEnvVars:
     """Test that the CLI proxy command reads API URL env vars."""