diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a0f2d09f..11a13f1b1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -103,6 +103,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Features +* **wrap:** add `headroom wrap omp` / `headroom unwrap omp` for Oh My Pi — points omp's built-in `anthropic` provider at the local proxy via a marker-fenced `providers.anthropic.baseUrl` override in `~/.omp/agent/models.yml`, snapshotting the pre-wrap file byte-for-byte and restoring it on unwrap. omp resolves its Anthropic chat endpoint from models.yml (`ANTHROPIC_BASE_URL` only feeds its web-search helper), and a same-ID override keeps omp's bundled model catalog and stored credentials ([#1149](https://github.com/headroomlabs-ai/headroom/issues/1149)) * **compress:** expose `frozen_message_count` in library-mode `compress()` via a new `CompressConfig` field (default `0`, unchanged behavior). `read_lifecycle.apply()` already skips stale-Read replacements inside a frozen message prefix, but only the proxy handlers could pass it — `ContentRouter` reads it from transform kwargs and the public API never forwarded it. Library-mode callers that manage their own conversation loop can now stop transforms from rewriting messages already anchored in the provider's prompt cache, which would otherwise convert 0.1x cached prefix reads into full-price cache writes ([#2178](https://github.com/headroomlabs-ai/headroom/pull/2178)). * **proxy:** report a new-content-relative input savings rate in `/stats`: `tokens.new_input_tokens` (provider-billed non-cache-read input: uncached + cache-write tokens, from response usage) and `tokens.new_input_savings_percent` (savings as a fraction of new input plus the tokens compression removed before they could be billed). The existing whole-request ratios recount the full transcript on every turn, so a 200-turn session counts its history 200x into the denominator and long-running cached sessions (especially 1M-context models, which never compact) dilute toward ~0% regardless of how well compression performs on content newly entering context. Purely additive; existing fields unchanged. Reports 0 when no cache usage data exists (e.g. providers without cache metrics) rather than dividing savings by themselves. diff --git a/README.md b/README.md index 4d8a52605..8ca5259ec 100644 --- a/README.md +++ b/README.md @@ -49,7 +49,7 @@ Headroom compresses everything your AI agent reads — tool outputs, logs, RAG c - **Library** — `compress(messages)` in Python or TypeScript, inline in any app - **Proxy** — `headroom proxy --port 8787`, zero code changes, any language -- **Agent wrap** — `headroom wrap claude|codex|grok|copilot|cursor|aider|opencode|cline|continue|goose|openhands|openclaw|vibe|zcode` in one command; undo with `headroom unwrap ` +- **Agent wrap** — `headroom wrap claude|codex|grok|copilot|cursor|aider|opencode|cline|continue|goose|openhands|openclaw|vibe|omp|zcode` in one command; undo with `headroom unwrap ` - **MCP server** — `headroom_compress`, `headroom_retrieve`, `headroom_stats` for any MCP client - **Cross-agent memory** — shared store across Claude, Codex, Gemini, Grok, auto-dedup - **`headroom learn`** — mines failed sessions, writes corrections to `CLAUDE.local.md` (default, gitignored) or `CLAUDE.md` / `AGENTS.md` / `GEMINI.md` / `GROK.md` @@ -238,11 +238,12 @@ shows an **Output Tokens Saved** card next to input compression, labelled | Goose | ✅ | starts proxy + launches | | OpenHands | ✅ | starts proxy + launches | | Mistral Vibe | ✅ | starts proxy + launches | +| Oh My Pi | ✅ | injects config · starts proxy + launches | | Cortex Code | Library only | 60–65% savings (library mode; no `wrap`) | | ZCode | ✅ | starts proxy and prints base URLs for ZCode settings | Any OpenAI-compatible client works via `headroom proxy`. MCP-native: `headroom mcp install`. -Undo durable wrapping with `headroom unwrap ` (supports: `claude`, `copilot`, `codex`, `grok`, `opencode`, `openclaw`, `zcode`). +Undo durable wrapping with `headroom unwrap ` (supports: `claude`, `copilot`, `codex`, `grok`, `omp`, `opencode`, `openclaw`, `zcode`). Registry authors can use the canonical [`server.json`](server.json) in the repo root instead of reconstructing the `headroom mcp serve` contract from prose. ### GitHub Copilot CLI subscription mode diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index a2af2ecee..13ca9b05c 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -111,6 +111,10 @@ from headroom.providers.copilot import ( from headroom.providers.cursor import render_setup_lines as _render_cursor_setup_lines from headroom.providers.grok import build_launch_env as _build_grok_launch_env from headroom.providers.mistral_vibe import build_launch_env as _build_mistral_vibe_launch_env +from headroom.providers.omp import build_launch_env as _build_omp_launch_env +from headroom.providers.omp import inject_models_override as _inject_omp_models_override +from headroom.providers.omp import models_yml_path as _omp_models_yml_path +from headroom.providers.omp import restore_models_override as _restore_omp_models_override from headroom.providers.openclaw import ( OPENCLAW_NPM_PACKAGE, ) @@ -3865,6 +3869,8 @@ def wrap() -> None: headroom wrap openhands # OpenHands CLI headroom wrap openclaw # OpenClaw plugin bootstrap headroom wrap opencode # OpenCode CLI + headroom wrap omp # Oh My Pi CLI + headroom wrap zcode # ZCode desktop app setup \b `wrap` vs `proxy`: @@ -6904,6 +6910,142 @@ def unwrap_codex(port: int, no_stop_proxy: bool) -> None: click.echo() +# ============================================================================= +# Oh My Pi (omp) +# ============================================================================= + + +@wrap.command(context_settings={"ignore_unknown_options": True}) +@click.option( + "--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)" +) +@click.option( + "--no-context-tool", + "--no-rtk", + "no_rtk", + is_flag=True, + help="Skip CLI context-tool setup", +) +@click.option( + "--code-graph", + is_flag=True, + help="Enable code graph indexing via codebase-memory-mcp (optional)", +) +@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)") +@click.option("--learn", is_flag=True, help="Enable live traffic learning") +@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory") +@click.option("--verbose", "-v", is_flag=True, help="Verbose output") +@click.option("--prepare-only", is_flag=True, hidden=True) +@click.argument("omp_args", nargs=-1, type=click.UNPROCESSED) +def omp( + port: int, + no_rtk: bool, + code_graph: bool, + no_proxy: bool, + learn: bool, + memory: bool, + verbose: bool, + prepare_only: bool, + omp_args: tuple, +) -> None: + """Launch Oh My Pi (omp) through Headroom proxy. + + \b + Points omp's built-in `anthropic` provider at Headroom by injecting a + marker-fenced `providers.anthropic.baseUrl` override into + ~/.omp/agent/models.yml (pre-wrap file backed up byte-for-byte; undo with + `headroom unwrap omp`). omp resolves its Anthropic chat endpoint from + models.yml — ANTHROPIC_BASE_URL only affects its web-search helper — and a + same-ID override keeps omp's bundled model catalog and stored credentials. + omp's other providers (OpenAI-direct, Gemini, ...) keep their normal + endpoints; route those via your own custom provider in models.yml. + + \b + Examples: + headroom wrap omp # Start proxy + context tool + omp + headroom wrap omp -- -p "fix the bug" # omp in non-interactive print mode + headroom wrap omp -- --model opus # Pick a model (fuzzy match) + headroom wrap omp --no-context-tool # Skip CLI context-tool setup + headroom unwrap omp # Restore pre-wrap models.yml + """ + # Setup CLI context tool for omp — it reads AGENTS.md from the project root. + if not no_rtk: + if _selected_context_tool() == _CONTEXT_TOOL_LEAN_CTX: + click.echo(" Setting up lean-ctx for omp...") + _setup_lean_ctx_agent("omp", verbose=verbose) + else: + click.echo(" Setting up rtk for omp...") + rtk_path = _ensure_rtk_binary(verbose=verbose) + if rtk_path: + agents_md = Path.cwd() / "AGENTS.md" + _inject_rtk_instructions(agents_md, verbose=verbose) + + if prepare_only: + _inject_omp_models_override(port, _project_name_from_cwd()) + return + + omp_bin = shutil.which("omp") + if not omp_bin: + click.echo("Error: 'omp' not found in PATH.") + click.echo("Install Oh My Pi: npm install -g @oh-my-pi/pi-coding-agent") + raise SystemExit(1) + + env, env_vars_display = _build_omp_launch_env( + port, os.environ, project=_project_name_from_cwd() + ) + + # Durable endpoint redirect (survives omp-spawned child sessions, which + # re-read models.yml rather than inheriting a parent env) — same durable + # wrap + backup + unwrap contract as the Codex config.toml injection. + models_file, _ = _inject_omp_models_override(port, _project_name_from_cwd()) + click.echo(f" models.yml override written: {models_file}") + + _launch_tool( + binary=omp_bin, + args=omp_args, + env=env, + port=port, + no_proxy=no_proxy, + tool_label="OMP", + env_vars_display=env_vars_display, + learn=learn, + memory=memory, + agent_type="omp", + code_graph=code_graph, + ) + + +@unwrap.command("omp") +@click.option( + "--port", "-p", default=8787, type=click.IntRange(1, 65535), help="Proxy port (default: 8787)" +) +@click.option("--no-stop-proxy", is_flag=True, help="Do not stop the local Headroom proxy") +def unwrap_omp(port: int, no_stop_proxy: bool) -> None: + """Undo ``headroom wrap omp`` edits to omp's models.yml. + + Restores the byte-for-byte pre-wrap backup when one exists, or removes the + wrap-created file when there was no models.yml before wrapping. A + models.yml the wrap does not manage is never touched. Also removes the + marker-fenced rtk guidance from the current project's AGENTS.md. + """ + status = _restore_omp_models_override() + if status == "restored": + click.echo(f" Restored pre-wrap models.yml from backup: {_omp_models_yml_path()}") + elif status == "removed": + click.echo(f" Removed wrap-created models.yml: {_omp_models_yml_path()}") + else: + click.echo(" No Headroom-managed models.yml found — nothing to restore.") + + if _remove_rtk_instructions(Path.cwd() / "AGENTS.md"): + click.echo(" Removed Headroom rtk instructions from AGENTS.md.") + + click.echo() + click.echo("✓ omp is no longer routed through the Headroom proxy.") + if not no_stop_proxy and status != "noop": + _echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(port), port) + click.echo() + + # ============================================================================= # Grok CLI (unwrap) # ============================================================================= diff --git a/headroom/providers/omp/__init__.py b/headroom/providers/omp/__init__.py new file mode 100644 index 000000000..61cda94a3 --- /dev/null +++ b/headroom/providers/omp/__init__.py @@ -0,0 +1,23 @@ +"""Oh My Pi (omp)-specific provider helpers.""" + +from .runtime import ( + MANAGED_MARKER, + backup_path, + build_launch_env, + inject_models_override, + is_managed, + models_yml_path, + proxy_anthropic_base_url, + restore_models_override, +) + +__all__ = [ + "MANAGED_MARKER", + "backup_path", + "build_launch_env", + "inject_models_override", + "is_managed", + "models_yml_path", + "proxy_anthropic_base_url", + "restore_models_override", +] diff --git a/headroom/providers/omp/runtime.py b/headroom/providers/omp/runtime.py new file mode 100644 index 000000000..6ede7c0ee --- /dev/null +++ b/headroom/providers/omp/runtime.py @@ -0,0 +1,166 @@ +"""Runtime helpers for Oh My Pi (omp) integrations. + +omp resolves its Anthropic chat endpoint from the model registry +(``providers.anthropic.baseUrl`` in ``~/.omp/agent/models.yml``), not from +``ANTHROPIC_BASE_URL`` — that env var only feeds omp's web-search helper. +Verified empirically: with ``ANTHROPIC_BASE_URL`` pointed at a local probe +server, omp's chat traffic still went to the real Anthropic endpoint; with a +``models.yml`` same-ID override, every ``/v1/messages`` request arrived at the +probe. A same-ID override keeps omp's bundled Anthropic model catalog and +stored credentials (both keyed by provider id ``anthropic``), so only the +endpoint moves. + +The wrap therefore injects a marker-fenced ``providers.anthropic.baseUrl`` +override into ``models.yml``, snapshotting the pre-wrap file byte-for-byte +first — the same durable-wrap + backup + ``headroom unwrap`` contract the +Codex wrap uses for ``config.toml``. +""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from pathlib import Path + +from headroom.providers.claude import proxy_base_url as claude_proxy_base_url +from headroom.proxy.project_context import with_project_prefix + +MANAGED_MARKER = "# managed by `headroom wrap omp`" +_MANAGED_HEADER = ( + f"{MANAGED_MARKER} — do not hand-edit while wrapped.\n" + "# `headroom unwrap omp` restores the pre-wrap file (or removes this one\n" + "# if it did not exist). The original is kept at .\n" +) +BACKUP_SUFFIX = ".headroom-backup" + + +def models_yml_path() -> Path: + """Path to omp's model/provider registry file. + + ``PI_CODING_AGENT_DIR`` relocates omp's ``~/.omp/agent`` state directory + (per omp's environment-variable reference); ``models.yml`` moves with it. + """ + base = os.environ.get("PI_CODING_AGENT_DIR", "").strip() + agent_dir = Path(base).expanduser() if base else Path.home() / ".omp" / "agent" + return agent_dir / "models.yml" + + +def backup_path(models_file: Path) -> Path: + """Backup location for the pre-wrap ``models.yml`` snapshot.""" + return models_file.with_name(models_file.name + BACKUP_SUFFIX) + + +def proxy_anthropic_base_url(port: int, project: str | None = None) -> str: + """Proxy base URL omp's ``anthropic`` provider is pointed at. + + ``project`` (the wrap launch directory) is encoded as a ``/p/`` + base-URL prefix — same as the Aider wrap — so the proxy attributes + savings per project. omp appends ``/v1/messages`` after any path segments + (its docs show path-carrying Anthropic base URLs, e.g. the Cloudflare AI + Gateway example), and the proxy strips the prefix on arrival. + """ + return with_project_prefix(claude_proxy_base_url(port), project) + + +def is_managed(models_file: Path) -> bool: + """Whether ``models_file`` is currently a wrap-managed override.""" + if not models_file.exists(): + return False + try: + head = models_file.read_bytes() + except OSError: + return False + return MANAGED_MARKER.encode("utf-8") in head + + +def inject_models_override(port: int, project: str | None = None) -> tuple[Path, str]: + """Point ``providers.anthropic.baseUrl`` at the local proxy. + + Returns ``(models_file, base_url)``. + + * First injection snapshots the user's pre-wrap file byte-for-byte to + ``models.yml.headroom-backup`` so unwrap can restore it exactly. A file + that is already wrap-managed is never re-snapshotted (that would clobber + the pristine backup — same guard as the Codex config snapshot). + * The managed file is regenerated from the backup (or from scratch when + the user had no ``models.yml``) on every call, so re-running with a + different ``--port`` updates the override idempotently. + * Any user-defined providers/models from the pre-wrap file are preserved: + the override only deep-sets ``providers.anthropic.baseUrl``. + """ + import yaml # type: ignore[import-untyped] # PyYAML ships no stubs; lint env installs no deps + + models_file = models_yml_path() + backup = backup_path(models_file) + base_url = proxy_anthropic_base_url(port, project) + + original_bytes: bytes | None = None + if backup.exists(): + original_bytes = backup.read_bytes() + elif models_file.exists(): + if is_managed(models_file): + # Managed file without a backup: we created it from scratch — + # regenerate from scratch rather than merging our own output. + original_bytes = None + else: + # Snapshot as raw bytes so the unwrap restore is byte-for-byte + # even for non-UTF-8 or newline-sensitive files (text-mode I/O + # would translate newlines on Windows). + backup.parent.mkdir(parents=True, exist_ok=True) + original_bytes = models_file.read_bytes() + backup.write_bytes(original_bytes) + + data: dict = {} + if original_bytes: + loaded = yaml.safe_load(original_bytes.decode("utf-8", errors="replace")) + if isinstance(loaded, dict): + data = loaded + providers = data.setdefault("providers", {}) + if not isinstance(providers, dict): # malformed user value: keep it in backup only + providers = {} + data["providers"] = providers + anthropic = providers.setdefault("anthropic", {}) + if not isinstance(anthropic, dict): + anthropic = {} + providers["anthropic"] = anthropic + anthropic["baseUrl"] = base_url + + models_file.parent.mkdir(parents=True, exist_ok=True) + rendered = yaml.safe_dump(data, sort_keys=False, default_flow_style=False) + models_file.write_text(_MANAGED_HEADER + rendered, encoding="utf-8", newline="\n") + return models_file, base_url + + +def restore_models_override() -> str: + """Undo :func:`inject_models_override`. + + Returns one of ``"restored"`` (backup moved back), ``"removed"`` + (wrap-created file deleted), or ``"noop"`` (nothing wrap-managed found). + Never touches a ``models.yml`` the wrap does not manage. + """ + models_file = models_yml_path() + backup = backup_path(models_file) + + if backup.exists(): + backup.replace(models_file) + return "restored" + if is_managed(models_file): + models_file.unlink() + return "removed" + return "noop" + + +def build_launch_env( + port: int, + environ: Mapping[str, str] | None = None, + project: str | None = None, +) -> tuple[dict[str, str], list[str]]: + """Build the launch environment and display lines for the omp wrap. + + The endpoint redirect itself lives in ``models.yml`` (see module + docstring), so the environment passes through unchanged; the display + lines surface where the override went. + """ + env = dict(environ or os.environ) + base_url = proxy_anthropic_base_url(port, project) + return env, [f"models.yml: providers.anthropic.baseUrl={base_url}"] diff --git a/headroom/telemetry/context.py b/headroom/telemetry/context.py index 0a724d35f..e68ec5e75 100644 --- a/headroom/telemetry/context.py +++ b/headroom/telemetry/context.py @@ -22,7 +22,7 @@ logger = logging.getLogger(__name__) _KNOWN_WRAP_AGENTS = frozenset( - {"claude", "copilot", "codex", "aider", "cursor", "openclaw", "opencode"} + {"claude", "copilot", "codex", "aider", "cursor", "omp", "openclaw", "opencode"} ) # Stack slugs must start with a letter and contain only [a-z0-9_], max 64 chars. diff --git a/llms.txt b/llms.txt index ffc120031..84cc27da9 100644 --- a/llms.txt +++ b/llms.txt @@ -21,7 +21,7 @@ The canonical, always-current documentation index lives at the docs site below. - TypeScript / Node: `npm install headroom-ai` (or `pnpm add headroom-ai`, `bun add headroom-ai`) - Docker: `docker run -p 8787:8787 ghcr.io/chopratejas/headroom:latest` - Run the proxy: `headroom proxy --port 8787` then point any client at `http://127.0.0.1:8787` -- Wrap an agent in one command: `headroom wrap claude` (also: `codex`, `copilot`, `cursor`, `aider`, `opencode`, `cline`, `continue`, `goose`, `openhands`, `openclaw`, `vibe`) +- Wrap an agent in one command: `headroom wrap claude` (also: `codex`, `copilot`, `cursor`, `aider`, `opencode`, `cline`, `continue`, `goose`, `openhands`, `openclaw`, `vibe`, `omp`) ## Entry points diff --git a/pyproject.toml b/pyproject.toml index 692306f19..274ecd568 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,6 +58,7 @@ dependencies = [ "rich>=13.0.0", # Rich terminal output "opentelemetry-api>=1.24.0", # Safe no-op OTEL API for instrumentation "ast-grep-cli>=0.30.0", # AST-aware code slicing (CodeCompressor); binary wheel + "pyyaml>=6.0", # omp wrap: parse/merge omp's models.yml registry "tomli>=2.0.0; python_version < '3.11'", # tomllib backport for helper scripts ] diff --git a/tests/test_cli/test_wrap_omp.py b/tests/test_cli/test_wrap_omp.py new file mode 100644 index 000000000..47064db02 --- /dev/null +++ b/tests/test_cli/test_wrap_omp.py @@ -0,0 +1,324 @@ +"""Tests for `headroom wrap omp` / `headroom unwrap omp`. + +Covers the omp runtime override contract (fresh-create vs merge-preserving +injection, pristine backups, re-injection idempotency, restore statuses) and +the CLI wiring that drives it. Every test isolates omp's agent directory via +``PI_CODING_AGENT_DIR`` and runs from a tmp cwd so the real ``~/.omp`` is +never touched. +""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import patch + +import pytest +import yaml +from click.testing import CliRunner + +from headroom.cli.main import main +from headroom.cli.wrap import _inject_rtk_instructions +from headroom.providers.omp import ( + MANAGED_MARKER, + backup_path, + build_launch_env, + inject_models_override, + models_yml_path, + restore_models_override, +) + + +@pytest.fixture +def runner() -> CliRunner: + return CliRunner() + + +@pytest.fixture +def omp_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + """Isolate omp's agent dir under tmp_path and run from a tmp cwd. + + Returns the ``models.yml`` path the runtime resolves to. + """ + agent_dir = tmp_path / "omp-agent" + agent_dir.mkdir() + monkeypatch.setenv("PI_CODING_AGENT_DIR", str(agent_dir)) + monkeypatch.chdir(tmp_path) + return agent_dir / "models.yml" + + +# --------------------------------------------------------------------------- +# runtime: path resolution +# --------------------------------------------------------------------------- + + +def test_models_yml_path_honors_pi_coding_agent_dir( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + custom = tmp_path / "relocated-agent" + monkeypatch.setenv("PI_CODING_AGENT_DIR", str(custom)) + monkeypatch.chdir(tmp_path) + assert models_yml_path() == custom / "models.yml" + + +# --------------------------------------------------------------------------- +# runtime: injection +# --------------------------------------------------------------------------- + + +def test_inject_fresh_create_writes_managed_marker_and_no_backup(omp_home: Path) -> None: + models_file, base_url = inject_models_override(8787, "proj") + + assert models_file == omp_home + assert base_url == "http://127.0.0.1:8787/p/proj" + + text = models_file.read_text(encoding="utf-8") + assert MANAGED_MARKER in text + assert yaml.safe_load(text)["providers"]["anthropic"]["baseUrl"] == base_url + + # Nothing pre-existed, so there is nothing to snapshot. + assert not backup_path(models_file).exists() + + +def test_inject_over_existing_backs_up_pristine_and_merges(omp_home: Path) -> None: + original = ( + "providers:\n" + " anthropic:\n" + " apiKey: sk-user-secret\n" + " openai:\n" + " baseUrl: https://api.openai.com/v1\n" + "models:\n" + " - id: my-custom-model\n" + ) + omp_home.write_bytes(original.encode("utf-8")) + + _, base_url = inject_models_override(8787, "proj") + + # Pre-wrap file is snapshotted byte-for-byte. + assert backup_path(omp_home).read_bytes() == original.encode("utf-8") + + merged = yaml.safe_load(omp_home.read_text(encoding="utf-8")) + assert merged["providers"]["anthropic"]["baseUrl"] == base_url + # Only anthropic.baseUrl is set; every other user key survives the merge. + assert merged["providers"]["anthropic"]["apiKey"] == "sk-user-secret" + assert merged["providers"]["openai"]["baseUrl"] == "https://api.openai.com/v1" + assert merged["models"] == [{"id": "my-custom-model"}] + assert MANAGED_MARKER in omp_home.read_text(encoding="utf-8") + + +def test_reinject_new_port_regenerates_from_pristine_backup(omp_home: Path) -> None: + original = "providers:\n anthropic:\n apiKey: sk-user-secret\n" + omp_home.write_bytes(original.encode("utf-8")) + backup = backup_path(omp_home) + + inject_models_override(8787, "proj") + assert backup.read_bytes() == original.encode("utf-8") + + _, base_url_9999 = inject_models_override(9999, "proj") + + # Re-injection never clobbers the pristine pre-wrap backup. + assert backup.read_bytes() == original.encode("utf-8") + + merged = yaml.safe_load(omp_home.read_text(encoding="utf-8")) + assert base_url_9999 == "http://127.0.0.1:9999/p/proj" + assert merged["providers"]["anthropic"]["baseUrl"] == base_url_9999 + # Regenerated from the backup, so user creds still survive the new port. + assert merged["providers"]["anthropic"]["apiKey"] == "sk-user-secret" + + +# --------------------------------------------------------------------------- +# runtime: restore +# --------------------------------------------------------------------------- + + +def test_restore_restores_pristine_and_removes_backup(omp_home: Path) -> None: + original = "providers:\n anthropic:\n apiKey: sk-user-secret\n" + omp_home.write_bytes(original.encode("utf-8")) + inject_models_override(8787, "proj") + + assert restore_models_override() == "restored" + assert omp_home.read_bytes() == original.encode("utf-8") + assert not backup_path(omp_home).exists() + + +def test_restore_removes_wrap_created_file(omp_home: Path) -> None: + inject_models_override(8787, "proj") # fresh create → no backup + assert omp_home.exists() + + assert restore_models_override() == "removed" + assert not omp_home.exists() + assert not backup_path(omp_home).exists() + + +def test_restore_noop_when_nothing_managed(omp_home: Path) -> None: + assert restore_models_override() == "noop" + + +def test_restore_leaves_unmanaged_file_untouched(omp_home: Path) -> None: + user_content = "providers:\n anthropic:\n apiKey: sk-user-secret\n" + omp_home.write_bytes(user_content.encode("utf-8")) + + assert restore_models_override() == "noop" + # A models.yml the wrap does not manage is never modified or deleted. + assert omp_home.read_bytes() == user_content.encode("utf-8") + assert not backup_path(omp_home).exists() + + +# --------------------------------------------------------------------------- +# runtime: launch env +# --------------------------------------------------------------------------- + + +def test_build_launch_env_passes_env_through_and_emits_display(omp_home: Path) -> None: + source = {"PATH": "/usr/bin", "ANTHROPIC_BASE_URL": "https://api.anthropic.com"} + env, display = build_launch_env(8787, source, project="proj") + + # The redirect lives in models.yml, so env is a verbatim copy — notably + # ANTHROPIC_BASE_URL is NOT rewritten to the proxy. + assert env == source + assert env is not source # a copy, so caller's environ can't be mutated + assert display == ["models.yml: providers.anthropic.baseUrl=http://127.0.0.1:8787/p/proj"] + + +# --------------------------------------------------------------------------- +# CLI: wrap omp +# --------------------------------------------------------------------------- + + +def test_wrap_omp_missing_binary_exits_with_install_hint(runner: CliRunner, omp_home: Path) -> None: + with patch("headroom.cli.wrap.shutil.which", return_value=None): + result = runner.invoke(main, ["wrap", "omp", "--no-rtk"]) + + assert result.exit_code == 1 + assert "npm install -g @oh-my-pi/pi-coding-agent" in result.output + # Fail fast before mutating omp's config. + assert not omp_home.exists() + + +def test_wrap_omp_happy_path_injects_before_launch(runner: CliRunner, omp_home: Path) -> None: + captured: dict[str, object] = {} + + def fake_launch_tool(**kwargs: object) -> None: + captured.update(kwargs) + # Prove models.yml is on disk BEFORE omp is launched. + captured["models_text_at_launch"] = ( + omp_home.read_text(encoding="utf-8") if omp_home.exists() else None + ) + + with ( + patch("headroom.cli.wrap.shutil.which", return_value="omp"), + patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool), + ): + result = runner.invoke(main, ["wrap", "omp", "--no-rtk", "--", "-p", "fix the bug"]) + + assert result.exit_code == 0, result.output + assert captured["tool_label"] == "OMP" + assert captured["agent_type"] == "omp" + assert captured["args"] == ("-p", "fix the bug") + + text_at_launch = captured["models_text_at_launch"] + assert isinstance(text_at_launch, str) + assert MANAGED_MARKER in text_at_launch + base_url = yaml.safe_load(text_at_launch)["providers"]["anthropic"]["baseUrl"] + assert base_url.startswith("http://127.0.0.1:8787/p/") + + display = captured["env_vars_display"] + assert isinstance(display, list) + assert f"models.yml: providers.anthropic.baseUrl={base_url}" in display + + +def test_wrap_omp_no_rtk_skips_agents_md(runner: CliRunner, omp_home: Path, tmp_path: Path) -> None: + with ( + patch("headroom.cli.wrap.shutil.which", return_value="omp"), + patch("headroom.cli.wrap._launch_tool"), + ): + result = runner.invoke(main, ["wrap", "omp", "--no-rtk"]) + + assert result.exit_code == 0, result.output + assert not (tmp_path / "AGENTS.md").exists() + + +def test_wrap_omp_rtk_injects_into_cwd_agents_md( + runner: CliRunner, omp_home: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("HEADROOM_CONTEXT_TOOL", "rtk") + with ( + patch("headroom.cli.wrap.shutil.which", return_value="omp"), + patch("headroom.cli.wrap._launch_tool"), + patch("headroom.cli.wrap._ensure_rtk_binary", return_value=tmp_path / "rtk"), + ): + result = runner.invoke(main, ["wrap", "omp"]) + + assert result.exit_code == 0, result.output + agents_md = tmp_path / "AGENTS.md" + assert agents_md.exists() + assert "headroom:rtk-instructions" in agents_md.read_text(encoding="utf-8") + + +# --------------------------------------------------------------------------- +# CLI: unwrap omp +# --------------------------------------------------------------------------- + + +def test_unwrap_omp_restored_and_cleans_agents_md( + runner: CliRunner, omp_home: Path, tmp_path: Path +) -> None: + original = "providers:\n anthropic:\n apiKey: sk-user-secret\n" + omp_home.write_bytes(original.encode("utf-8")) + inject_models_override(8787, "proj") + + agents_md = tmp_path / "AGENTS.md" + agents_md.write_text("# My project rules\n\nBe nice.\n", encoding="utf-8") + _inject_rtk_instructions(agents_md) + + stopped: list[int] = [] + with patch( + "headroom.cli.wrap._stop_local_proxy_for_unwrap", + side_effect=lambda port: stopped.append(port) or "not_running", + ): + result = runner.invoke(main, ["unwrap", "omp"]) + + assert result.exit_code == 0, result.output + assert "Restored pre-wrap models.yml" in result.output + assert omp_home.read_bytes() == original.encode("utf-8") + assert not backup_path(omp_home).exists() + + # Only the marker-fenced rtk block is scrubbed; user content survives. + remaining = agents_md.read_text(encoding="utf-8") + assert "headroom:rtk-instructions" not in remaining + assert "Be nice." in remaining + + # A real restore (not a noop) attempts to stop the proxy on the given port. + assert stopped == [8787] + + +def test_unwrap_omp_removes_wrap_created_file(runner: CliRunner, omp_home: Path) -> None: + inject_models_override(8787, "proj") # fresh create → no backup + + stopped: list[int] = [] + with patch( + "headroom.cli.wrap._stop_local_proxy_for_unwrap", + side_effect=lambda port: stopped.append(port) or "not_running", + ): + result = runner.invoke(main, ["unwrap", "omp", "--port", "9191"]) + + assert result.exit_code == 0, result.output + assert "Removed wrap-created models.yml" in result.output + assert not omp_home.exists() + assert stopped == [9191] + + +def test_unwrap_omp_noop_leaves_unmanaged_and_skips_proxy_stop( + runner: CliRunner, omp_home: Path +) -> None: + user_content = "providers:\n anthropic:\n apiKey: sk-user-secret\n" + omp_home.write_bytes(user_content.encode("utf-8")) + + with patch("headroom.cli.wrap._stop_local_proxy_for_unwrap") as stop_proxy: + result = runner.invoke(main, ["unwrap", "omp"]) + + assert result.exit_code == 0, result.output + assert "nothing to restore" in result.output + # Unmanaged file is left exactly as the user had it. + assert omp_home.read_bytes() == user_content.encode("utf-8") + # noop status → the proxy is left running (the `status != "noop"` guard). + stop_proxy.assert_not_called() diff --git a/uv.lock b/uv.lock index 28b84d856..81067e869 100644 --- a/uv.lock +++ b/uv.lock @@ -1465,6 +1465,7 @@ dependencies = [ { name = "litellm", marker = "python_full_version < '3.14'" }, { name = "opentelemetry-api" }, { name = "pydantic" }, + { name = "pyyaml" }, { name = "rich" }, { name = "tiktoken" }, { name = "tomli", marker = "python_full_version < '3.11'" }, @@ -1760,6 +1761,7 @@ requires-dist = [ { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.21.0" }, { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=4.0.0" }, + { name = "pyyaml", specifier = ">=6.0" }, { name = "qdrant-client", marker = "extra == 'memory-stack'", specifier = ">=1.9.0,<2.0" }, { name = "rapidocr", marker = "python_full_version >= '3.13' and extra == 'image'", specifier = ">=3.0,<4" }, { name = "rapidocr-onnxruntime", marker = "python_full_version < '3.13' and extra == 'image'", specifier = ">=1.4.0,<2" },