mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
b4571cc346
## Summary This PR implements transparent `headroom wrap opencode` support without asking users to edit OpenCode provider URLs, choose an extra CLI flag, or maintain a static provider list. The wrapper now lives at the runtime transport boundary: OpenCode keeps its user/provider config, while Headroom intercepts outbound provider traffic in-process and routes it through the local Headroom proxy. ## What changed ### Transparent OpenCode wrapping - `headroom wrap opencode` injects the `headroom-opencode` plugin through `OPENCODE_CONFIG_CONTENT`. - Existing OpenCode provider URLs are preserved. We do not rewrite user config URLs to point at Headroom. - Existing `OPENAI_BASE_URL` and `ANTHROPIC_BASE_URL` env vars are preserved. - Local OpenCode traffic, localhost traffic, and Headroom proxy traffic bypass the shim to avoid loops. ### Runtime transport interception - Added an OpenCode plugin transport shim that wraps: - `globalThis.fetch` - `http.request` / `http.get` - `https.request` / `https.get` - External provider calls are routed to the local Headroom proxy. - The original upstream origin is passed through `x-headroom-base-url`, so the proxy can forward to the real provider without changing OpenCode config. - External `http2.connect` is blocked loudly instead of allowing direct provider traffic to leak outside Headroom. ### Live provider additions Provider coverage is no longer based on a static config scan. Because routing happens at outbound request time, providers added mid-session are routed through Headroom automatically as long as they use the covered Node transport paths. ### Subagent and child-process coverage - The parent OpenCode plugin sets a packaged Node preload shim through `NODE_OPTIONS=--import=.../hook-shim/handler.js`. - The transport shim patches `child_process.spawn`, `exec`, `execFile`, and `fork` so child Node processes receive the Headroom preload even when OpenCode passes a custom `env`. - The child-process shim fails closed if it loads without `HEADROOM_OPENCODE_TRANSPORT_PROXY_URL`. - This closes the subagent leak path where a child Node process could otherwise start without Headroom transport interception. ## Why this goes beyond PR #1089 PR #1089 improves OpenCode provider registration, but it still focuses on provider config shape. This PR moves the enforcement boundary to runtime transport interception. This PR goes further because: - No provider URL rewriting is required. - New providers added mid-session are covered automatically. - Subagents and child Node processes inherit the Headroom transport shim. - Direct external HTTP/2 paths fail loudly instead of leaking. - The wrap remains transparent to the user's OpenCode provider config. - The wrapper is fail-closed for unsupported child-process preload state. ## Additional robustness fixes While validating the change in Docker, the full Python suite exposed unrelated Linux/container robustness issues. These are fixed in this PR so the suite is green: - Binary cache handling now treats cache paths under a non-writable existing parent as unavailable, including when tests run as root in Docker. - `release_version.py` honors `MANUAL_VER` before git calls so direct script execution works outside a `.git` checkout. - Test logger isolation now resets relevant Headroom child loggers so proxy logging setup cannot poison later `caplog` tests. - The scanner missing-path test now uses a guaranteed missing `tmp_path` child instead of relying on `/nonexistent/path`. ## Validation All implementation validation was run inside Docker. - Full Python suite from a fresh Docker copy: `6605 passed, 523 skipped`. - Ruff on changed Python/OpenCode paths: passed. - OpenCode plugin typecheck: passed. - OpenCode plugin tests: `9 passed`. - OpenCode plugin build: passed. - Hook shim preload smoke test: passed. ## Notes This PR intentionally does not add a CLI option. `headroom wrap opencode` means full wrap. Either Headroom wraps OpenCode transparently, or the path fails loudly instead of silently leaking provider traffic. --------- Co-authored-by: Rudimar Ronsoni <6081613+rudironsoni@users.noreply.github.com>
115 lines
3.6 KiB
Python
115 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
import builtins
|
|
import sys
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import pytest
|
|
|
|
from headroom.models.ml_models import MLModelRegistry
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_ml_model_registry():
|
|
MLModelRegistry.reset()
|
|
yield
|
|
MLModelRegistry.reset()
|
|
|
|
|
|
def test_unload_many_removes_requested_keys_once(monkeypatch) -> None:
|
|
MLModelRegistry.reset()
|
|
registry = MLModelRegistry.get()
|
|
kept_model = object()
|
|
registry._models.update(
|
|
{
|
|
"technique_router:demo": object(),
|
|
"siglip:demo": object(),
|
|
"sentence_transformer:keep": kept_model,
|
|
}
|
|
)
|
|
release = Mock()
|
|
monkeypatch.setattr(MLModelRegistry, "_release_runtime_memory", release)
|
|
|
|
removed = MLModelRegistry.unload_many(["missing", "technique_router:demo", "siglip:demo"])
|
|
|
|
assert removed == ["technique_router:demo", "siglip:demo"]
|
|
assert registry._models == {"sentence_transformer:keep": kept_model}
|
|
release.assert_called_once_with()
|
|
|
|
|
|
def test_unload_many_skips_runtime_cleanup_when_nothing_removed(monkeypatch) -> None:
|
|
MLModelRegistry.reset()
|
|
registry = MLModelRegistry.get()
|
|
registry._models["sentence_transformer:keep"] = object()
|
|
release = Mock()
|
|
monkeypatch.setattr(MLModelRegistry, "_release_runtime_memory", release)
|
|
|
|
removed = MLModelRegistry.unload_many(["missing"])
|
|
|
|
assert removed == []
|
|
assert "sentence_transformer:keep" in registry._models
|
|
release.assert_not_called()
|
|
|
|
|
|
def test_unload_prefix_removes_only_matching_models(monkeypatch) -> None:
|
|
MLModelRegistry.reset()
|
|
registry = MLModelRegistry.get()
|
|
kept_model = object()
|
|
registry._models.update(
|
|
{
|
|
"siglip:a": object(),
|
|
"siglip:b": object(),
|
|
"technique_router:keep": kept_model,
|
|
}
|
|
)
|
|
release = Mock()
|
|
monkeypatch.setattr(MLModelRegistry, "_release_runtime_memory", release)
|
|
|
|
removed = MLModelRegistry.unload_prefix("siglip:")
|
|
|
|
assert removed == ["siglip:a", "siglip:b"]
|
|
assert registry._models == {"technique_router:keep": kept_model}
|
|
release.assert_called_once_with()
|
|
|
|
|
|
def test_unload_delegates_to_unload_many(monkeypatch) -> None:
|
|
unload_many = Mock(return_value=["siglip:demo"])
|
|
monkeypatch.setattr(MLModelRegistry, "unload_many", unload_many)
|
|
|
|
assert MLModelRegistry.unload("siglip:demo") is True
|
|
|
|
unload_many.assert_called_once_with(["siglip:demo"])
|
|
|
|
|
|
def test_release_runtime_memory_handles_missing_torch(monkeypatch) -> None:
|
|
collect = Mock()
|
|
monkeypatch.setattr("headroom.models.ml_models.gc.collect", collect)
|
|
real_import = builtins.__import__
|
|
|
|
def fake_import(name, *args, **kwargs): # noqa: ANN001, ANN202
|
|
if name == "torch":
|
|
raise ImportError("torch unavailable")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", fake_import)
|
|
|
|
MLModelRegistry._release_runtime_memory()
|
|
|
|
collect.assert_called_once_with()
|
|
|
|
|
|
def test_release_runtime_memory_clears_available_torch_caches(monkeypatch) -> None:
|
|
collect = Mock()
|
|
cuda = SimpleNamespace(is_available=Mock(return_value=True), empty_cache=Mock())
|
|
mps = SimpleNamespace(empty_cache=Mock())
|
|
fake_torch = SimpleNamespace(cuda=cuda, mps=mps)
|
|
monkeypatch.setattr("headroom.models.ml_models.gc.collect", collect)
|
|
monkeypatch.setitem(sys.modules, "torch", fake_torch)
|
|
|
|
MLModelRegistry._release_runtime_memory()
|
|
|
|
collect.assert_called_once_with()
|
|
cuda.is_available.assert_called_once_with()
|
|
cuda.empty_cache.assert_called_once_with()
|
|
mps.empty_cache.assert_called_once_with()
|