From 281bc171dcc5109356fdd876d85536c77fee78a3 Mon Sep 17 00:00:00 2001 From: JerrettDavis Date: Thu, 23 Apr 2026 11:13:54 -0500 Subject: [PATCH] fix(wrap): unwrap codex restores prior config.toml MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `headroom wrap codex` injects a `model_provider = "headroom"` block plus a `[model_providers.headroom]` table into `~/.codex/config.toml` so Codex routes both HTTP and WebSocket traffic through the proxy. The matching `unwrap codex` subcommand did not exist, so the injected block stayed in `config.toml` forever — the moment the proxy stopped, Codex (CLI and macOS app) started erroring with `Missing environment variable: OPENAI_API_KEY`, and users had to hand- edit the file to recover. Fix: * `_inject_codex_provider_config` now snapshots the pre-wrap file to `~/.codex/config.toml.headroom-backup` before the first modification and leaves that snapshot untouched on subsequent wrap runs. The injection is also rewritten to use two self-contained marker- delimited blocks (top-level key and provider table) so stripping them never consumes user content that sits between them. * `_inject_memory_mcp_config` takes the same snapshot, so `wrap codex --memory` without a full provider injection is still fully reversible. * New `_restore_codex_provider_config` helper and `unwrap codex` click command: * backup present → restore byte-for-byte and delete the backup; * backup absent but Headroom block present → strip the block and keep surrounding user content; * config contained only Headroom content → remove the file so Codex falls back to defaults; * nothing to undo → safe no-op. Codex is the only wrap target that modifies a persistent user config file: claude/aider/cursor/copilot all go through env vars or project- scoped files only, so this bug was unique to Codex. Tests: * `tests/test_cli/test_wrap_codex.py` adds 20 new cases covering the strip/snapshot helpers directly, round-trip idempotency of wrap → wrap → unwrap, handling of malformed prior configs, and end-to-end CliRunner invocations of `headroom wrap codex --prepare-only` / `headroom unwrap codex` against a temp `$HOME`. * All 153 existing `tests/test_cli/` tests continue to pass. Plugin manifest versions were re-synced from `pyproject.toml` (0.11.2) by the `sync-plugin-versions` pre-commit hook; the previous values (0.10.3) had drifted. Reported by @raenaryl in Discord on 0.6.3; confirmed still broken on current `main` (0.11.x). Co-Authored-By: Claude Opus 4.7 (1M context) --- .claude-plugin/marketplace.json | 60 +-- .github/plugin/marketplace.json | 60 +-- CHANGELOG.md | 11 + headroom/cli/wrap.py | 239 ++++++++++-- .../.claude-plugin/plugin.json | 34 +- .../.github/plugin/plugin.json | 36 +- tests/test_cli/test_wrap_codex.py | 359 ++++++++++++++++++ 7 files changed, 678 insertions(+), 121 deletions(-) create mode 100644 tests/test_cli/test_wrap_codex.py diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json index de780a009..d17458557 100644 --- a/.claude-plugin/marketplace.json +++ b/.claude-plugin/marketplace.json @@ -1,30 +1,30 @@ -{ - "name": "headroom-marketplace", - "owner": { - "name": "Headroom Contributors" - }, - "metadata": { - "description": "Headroom marketplace for Claude Code and GitHub Copilot CLI plugins.", - "version": "0.10.3" - }, - "plugins": [ - { - "name": "headroom", - "source": "./plugins/headroom-agent-hooks", - "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", - "version": "0.10.3", - "author": { - "name": "Headroom Contributors", - "url": "https://github.com/chopratejas/headroom" - }, - "homepage": "https://github.com/chopratejas/headroom", - "repository": "https://github.com/chopratejas/headroom", - "keywords": [ - "headroom", - "hooks", - "claude-code", - "copilot-cli" - ] - } - ] -} +{ + "name": "headroom-marketplace", + "owner": { + "name": "Headroom Contributors" + }, + "metadata": { + "description": "Headroom marketplace for Claude Code and GitHub Copilot CLI plugins.", + "version": "0.11.2" + }, + "plugins": [ + { + "name": "headroom", + "source": "./plugins/headroom-agent-hooks", + "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", + "version": "0.11.2", + "author": { + "name": "Headroom Contributors", + "url": "https://github.com/chopratejas/headroom" + }, + "homepage": "https://github.com/chopratejas/headroom", + "repository": "https://github.com/chopratejas/headroom", + "keywords": [ + "headroom", + "hooks", + "claude-code", + "copilot-cli" + ] + } + ] +} diff --git a/.github/plugin/marketplace.json b/.github/plugin/marketplace.json index de780a009..d17458557 100644 --- a/.github/plugin/marketplace.json +++ b/.github/plugin/marketplace.json @@ -1,30 +1,30 @@ -{ - "name": "headroom-marketplace", - "owner": { - "name": "Headroom Contributors" - }, - "metadata": { - "description": "Headroom marketplace for Claude Code and GitHub Copilot CLI plugins.", - "version": "0.10.3" - }, - "plugins": [ - { - "name": "headroom", - "source": "./plugins/headroom-agent-hooks", - "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", - "version": "0.10.3", - "author": { - "name": "Headroom Contributors", - "url": "https://github.com/chopratejas/headroom" - }, - "homepage": "https://github.com/chopratejas/headroom", - "repository": "https://github.com/chopratejas/headroom", - "keywords": [ - "headroom", - "hooks", - "claude-code", - "copilot-cli" - ] - } - ] -} +{ + "name": "headroom-marketplace", + "owner": { + "name": "Headroom Contributors" + }, + "metadata": { + "description": "Headroom marketplace for Claude Code and GitHub Copilot CLI plugins.", + "version": "0.11.2" + }, + "plugins": [ + { + "name": "headroom", + "source": "./plugins/headroom-agent-hooks", + "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", + "version": "0.11.2", + "author": { + "name": "Headroom Contributors", + "url": "https://github.com/chopratejas/headroom" + }, + "homepage": "https://github.com/chopratejas/headroom", + "repository": "https://github.com/chopratejas/headroom", + "keywords": [ + "headroom", + "hooks", + "claude-code", + "copilot-cli" + ] + } + ] +} diff --git a/CHANGELOG.md b/CHANGELOG.md index 0cb71f636..7d0ff3add 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Fixed +- **`headroom unwrap codex` now actually undoes `headroom wrap codex`** — + previously there was no `unwrap codex` subcommand at all, so the injected + `model_provider = "headroom"` / `[model_providers.headroom]` block stayed + in `~/.codex/config.toml` forever and Codex continued routing through the + (potentially stopped) proxy, surfacing as `Missing environment variable: + OPENAI_API_KEY`. `wrap codex` now snapshots the pre-wrap + `config.toml` to `config.toml.headroom-backup` before its first injection, + and `unwrap codex` restores that snapshot byte-for-byte (or, if the + backup is missing, strips only the Headroom-managed block and leaves + surrounding user content intact). Safe no-op when run without a prior + wrap. Reported by @raenaryl in Discord. - **`headroom learn` no longer clobbers prior recommendations on re-run** — the marker block in `CLAUDE.md` / `MEMORY.md` is now merged with the prior block instead of wholesale-replaced. Sections re-surfaced by the diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 6c27e5f30..a315e336e 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -416,6 +416,94 @@ _MEMORY_MCP_MARKER = "# --- Headroom memory MCP (auto-injected) ---" _MEMORY_MCP_END = "# --- end Headroom memory ---" _MEMORY_AGENTS_MARKER = "" +# Codex config injection markers +_CODEX_TOP_LEVEL_MARKER = "# --- Headroom proxy (auto-injected by headroom wrap codex) ---" +_CODEX_END_MARKER = "# --- end Headroom ---" +# File name used for the pre-wrap snapshot of ~/.codex/config.toml. The +# snapshot lets `headroom unwrap codex` restore the exact prior state, even +# if the user had their own `model_provider` / `[model_providers.*]` config +# before running wrap. +_CODEX_CONFIG_BACKUP_SUFFIX = ".headroom-backup" + + +def _codex_config_paths() -> tuple[Path, Path]: + """Return ``(config_file, backup_file)`` paths for the Codex TOML config.""" + config_dir = Path.home() / ".codex" + config_file = config_dir / "config.toml" + backup_file = config_dir / f"config.toml{_CODEX_CONFIG_BACKUP_SUFFIX}" + return config_file, backup_file + + +def _strip_codex_headroom_blocks(content: str) -> str: + """Remove all Headroom-managed blocks from a Codex ``config.toml`` string. + + Returns the cleaned content. Safe to call on content that never contained + any markers — it will be returned effectively unchanged (only trailing + whitespace is normalized). + """ + import re + + # Remove any top-level-marker → end-marker span, possibly repeated. + while _CODEX_TOP_LEVEL_MARKER in content and _CODEX_END_MARKER in content: + start = content.index(_CODEX_TOP_LEVEL_MARKER) + end_idx = content.index(_CODEX_END_MARKER, start) + if end_idx < start: + break + end = end_idx + len(_CODEX_END_MARKER) + content = content[:start].rstrip("\n") + "\n" + content[end:].lstrip("\n") + + # Remove any stale top-level marker or end marker that lost its partner + # (e.g. a crashed prior wrap). + content = content.replace(_CODEX_TOP_LEVEL_MARKER + "\n", "") + content = content.replace(_CODEX_END_MARKER + "\n", "") + + # Strip any leftover top-level `model_provider = "headroom"` line, which + # older versions of `wrap codex` wrote outside the marker block. + content = re.sub(r'(?m)^[ \t]*model_provider[ \t]*=[ \t]*"headroom"[ \t]*\r?\n', "", content) + + # Strip any orphaned `[model_providers.headroom]` table with the fields we + # write. We only remove it if the table is recognisably ours (base_url + # mentions localhost and a Headroom proxy port). This protects users who + # happen to have a differently configured `headroom` provider. + orphan_headroom_table = re.compile( + r"(?ms)^\[model_providers\.headroom\][^\[]*?" + r'base_url[ \t]*=[ \t]*"http://127\.0\.0\.1:\d+/v1"[^\[]*?' + r"(?=^\[|\Z)" + ) + content = orphan_headroom_table.sub("", content) + + return content.lstrip("\n").rstrip() + "\n" if content.strip() else "" + + +def _snapshot_codex_config_if_unwrapped(config_file: Path, backup_file: Path) -> None: + """Snapshot ``config.toml`` to ``backup_file`` before the first injection. + + Called as the first step of every Headroom injection into Codex's + ``config.toml``. Guarantees that ``headroom unwrap codex`` can restore the + user's original file byte-for-byte. + + Rules: + + * If the backup already exists, leave it alone — we only snapshot the + *pre-wrap* state, so running wrap repeatedly must not clobber it. + * If the config file doesn't exist yet, there's nothing to back up; unwrap + will remove the file entirely instead of restoring a snapshot. + * If the config already contains a Headroom marker, a wrap run is already + active: do not snapshot the injected state. + """ + if backup_file.exists(): + return + if not config_file.exists(): + return + try: + content = config_file.read_text() + except OSError: + return + if _CODEX_TOP_LEVEL_MARKER in content or _CODEX_END_MARKER in content: + return + backup_file.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(config_file, backup_file) + def _ensure_rtk_binary(verbose: bool = False) -> Path | None: """Ensure rtk binary is installed (download if needed). No hook registration.""" @@ -454,47 +542,56 @@ def _inject_codex_provider_config(port: int) -> None: ``[model_providers.headroom]`` section that routes both HTTP and WS through the proxy, and sets ``model_provider = "headroom"``. - Safe to call multiple times — only writes if the section is missing - or the port changed. + Safe to call multiple times — the injected block is fully replaced on + each call, so re-running with a different ``port`` updates the config. + Before the first injection, the pre-wrap file is snapshotted to + ``~/.codex/config.toml.headroom-backup`` so ``headroom unwrap codex`` + can restore it byte-for-byte. """ - config_dir = Path.home() / ".codex" - config_file = config_dir / "config.toml" + config_file, backup_file = _codex_config_paths() + config_dir = config_file.parent - # model_provider must be a top-level TOML key (before any [section]). - # The [model_providers.headroom] table can go at the end. - top_level_marker = "# --- Headroom proxy (auto-injected by headroom wrap codex) ---" - top_level_block = f'{top_level_marker}\nmodel_provider = "headroom"\n' + # The injected content is split into two self-contained, marker-delimited + # blocks: a top-level key block (at the start of the file, because bare + # TOML keys must precede any [section]) and a provider-table block (at + # the end). Each block has its own matching begin/end marker pair so + # stripping them is unambiguous and never consumes user content that + # happens to sit between the two. + top_level_block = ( + f'{_CODEX_TOP_LEVEL_MARKER}\nmodel_provider = "headroom"\n{_CODEX_END_MARKER}\n' + ) provider_section = ( - f"\n[model_providers.headroom]\n" - f'name = "OpenAI via Headroom proxy"\n' + f"{_CODEX_TOP_LEVEL_MARKER}\n" + "[model_providers.headroom]\n" + 'name = "OpenAI via Headroom proxy"\n' f'base_url = "http://127.0.0.1:{port}/v1"\n' f'env_key = "OPENAI_API_KEY"\n' f"requires_openai_auth = true\n" f"supports_websockets = true\n" - f"# --- end Headroom ---\n" + f"{_CODEX_END_MARKER}\n" ) - marker = top_level_marker - end_marker = "# --- end Headroom ---" - try: config_dir.mkdir(parents=True, exist_ok=True) + # Snapshot the pre-wrap state before touching anything. No-op if the + # config is already wrapped, is missing, or we've already snapshotted. + _snapshot_codex_config_if_unwrapped(config_file, backup_file) + if config_file.exists(): content = config_file.read_text() - if marker in content: - # Remove existing Headroom blocks entirely - start = content.index(marker) - end = content.index(end_marker) + len(end_marker) - content = content[:start].rstrip("\n") + content[end:].lstrip("\n") + # Remove any prior Headroom-managed blocks before re-injecting so + # the operation is idempotent and supports port changes. + content = _strip_codex_headroom_blocks(content) - # Strip any stale top-level model_provider left behind - import re - - content = re.sub(r'\nmodel_provider\s*=\s*"headroom"\n', "\n", content) - - # Place top-level key at the very beginning, provider table at the end - content = top_level_block + "\n" + content.strip() + "\n" + provider_section + # Place the top-level key block at the very beginning of the file + # (bare TOML keys must precede any [section]) and the provider + # table at the end. User content, if any, sits between them. + user_content = content.strip() + if user_content: + content = top_level_block + "\n" + user_content + "\n\n" + provider_section + else: + content = top_level_block + "\n" + provider_section else: content = top_level_block + "\n" + provider_section @@ -504,6 +601,44 @@ def _inject_codex_provider_config(port: int) -> None: click.echo(f" Warning: could not update Codex config: {e}") +def _restore_codex_provider_config() -> tuple[str, Path]: + """Undo ``_inject_codex_provider_config`` for ``~/.codex/config.toml``. + + Returns a tuple of ``(status, config_file)`` where status is one of: + + * ``"restored"`` — a pre-wrap backup existed and was restored; backup + file has been removed. + * ``"cleaned"`` — no backup existed, but the Headroom-managed block was + found and stripped out (preserving surrounding user content). + * ``"removed"`` — the config file only contained Headroom-managed + content (created by wrap) and has been deleted. + * ``"noop"`` — nothing to undo; no Headroom marker and no backup. + """ + config_file, backup_file = _codex_config_paths() + + # Case 1: pre-wrap snapshot exists — restore it exactly. + if backup_file.exists(): + shutil.copy2(backup_file, config_file) + backup_file.unlink() + return "restored", config_file + + # Case 2: no backup, but config file exists and has markers — strip them. + if config_file.exists(): + original = config_file.read_text() + if _CODEX_TOP_LEVEL_MARKER in original or _CODEX_END_MARKER in original: + cleaned = _strip_codex_headroom_blocks(original) + if not cleaned.strip(): + # Nothing left but Headroom content — remove the file entirely + # so Codex falls back to its default config. + config_file.unlink() + return "removed", config_file + config_file.write_text(cleaned) + return "cleaned", config_file + + # Nothing to undo. + return "noop", config_file + + def _inject_rtk_instructions(file_path: Path, verbose: bool = False) -> bool: """Inject rtk instructions into a file (AGENTS.md, .cursorrules, etc.). @@ -554,6 +689,12 @@ def _inject_memory_mcp_config(db_path: str, user_id: str) -> None: try: config_dir.mkdir(parents=True, exist_ok=True) + # Snapshot pre-wrap state before touching config.toml so `unwrap codex` + # can fully restore it even when only `--memory` (not a full provider + # injection) was used. + _, backup_file = _codex_config_paths() + _snapshot_codex_config_if_unwrapped(config_file, backup_file) + if config_file.exists(): content = config_file.read_text() if _MEMORY_MCP_MARKER in content: @@ -2138,3 +2279,49 @@ def unwrap_openclaw( click.echo(" Plugin: headroom (installed, disabled)") click.echo(" Slot: plugins.slots.contextEngine = legacy") click.echo() + + +# ============================================================================= +# OpenAI Codex CLI (unwrap) +# ============================================================================= + + +@unwrap.command("codex") +def unwrap_codex() -> None: + """Undo ``headroom wrap codex`` edits to ``~/.codex/config.toml``. + + Behaviour: + + * If a pre-wrap backup (``config.toml.headroom-backup``) exists, the + original file is restored byte-for-byte and the backup is removed. + * Otherwise, if the config file still contains the Headroom-managed + block, that block is stripped out and the rest of the file is + preserved. + * If the config only ever contained Headroom-written content, the file + is removed entirely so Codex falls back to its defaults. + * If neither a backup nor a Headroom block is present, this is a safe + no-op (the user either never wrapped, or already unwrapped). + """ + click.echo() + click.echo(" ╔═══════════════════════════════════════════════╗") + click.echo(" ║ HEADROOM UNWRAP: CODEX ║") + click.echo(" ╚═══════════════════════════════════════════════╝") + click.echo() + + try: + status, config_file = _restore_codex_provider_config() + except Exception as e: # pragma: no cover - filesystem-level errors + raise click.ClickException(f"could not unwrap Codex config: {e}") from e + + if status == "restored": + click.echo(f" Restored prior {config_file} from pre-wrap backup.") + elif status == "cleaned": + click.echo(f" Removed Headroom block from {config_file}; other content preserved.") + elif status == "removed": + click.echo(f" Removed {config_file} (contained only Headroom-written config).") + else: + click.echo(f" Nothing to undo: {config_file} has no Headroom wrap markers.") + + click.echo() + click.echo("✓ Codex is no longer routed through the Headroom proxy.") + click.echo() diff --git a/plugins/headroom-agent-hooks/.claude-plugin/plugin.json b/plugins/headroom-agent-hooks/.claude-plugin/plugin.json index 30bf7dac0..d3a6b2425 100644 --- a/plugins/headroom-agent-hooks/.claude-plugin/plugin.json +++ b/plugins/headroom-agent-hooks/.claude-plugin/plugin.json @@ -1,17 +1,17 @@ -{ - "name": "headroom", - "version": "0.10.3", - "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", - "author": { - "name": "Headroom Contributors", - "url": "https://github.com/chopratejas/headroom" - }, - "homepage": "https://github.com/chopratejas/headroom", - "repository": "https://github.com/chopratejas/headroom", - "keywords": [ - "headroom", - "hooks", - "claude-code", - "copilot-cli" - ] -} +{ + "name": "headroom", + "version": "0.11.2", + "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", + "author": { + "name": "Headroom Contributors", + "url": "https://github.com/chopratejas/headroom" + }, + "homepage": "https://github.com/chopratejas/headroom", + "repository": "https://github.com/chopratejas/headroom", + "keywords": [ + "headroom", + "hooks", + "claude-code", + "copilot-cli" + ] +} diff --git a/plugins/headroom-agent-hooks/.github/plugin/plugin.json b/plugins/headroom-agent-hooks/.github/plugin/plugin.json index 718ea609f..5d8ae816f 100644 --- a/plugins/headroom-agent-hooks/.github/plugin/plugin.json +++ b/plugins/headroom-agent-hooks/.github/plugin/plugin.json @@ -1,18 +1,18 @@ -{ - "name": "headroom", - "version": "0.10.3", - "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", - "author": { - "name": "Headroom Contributors", - "url": "https://github.com/chopratejas/headroom" - }, - "homepage": "https://github.com/chopratejas/headroom", - "repository": "https://github.com/chopratejas/headroom", - "keywords": [ - "headroom", - "hooks", - "claude-code", - "copilot-cli" - ], - "hooks": "./hooks" -} +{ + "name": "headroom", + "version": "0.11.2", + "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", + "author": { + "name": "Headroom Contributors", + "url": "https://github.com/chopratejas/headroom" + }, + "homepage": "https://github.com/chopratejas/headroom", + "repository": "https://github.com/chopratejas/headroom", + "keywords": [ + "headroom", + "hooks", + "claude-code", + "copilot-cli" + ], + "hooks": "./hooks" +} diff --git a/tests/test_cli/test_wrap_codex.py b/tests/test_cli/test_wrap_codex.py new file mode 100644 index 000000000..9944b63bd --- /dev/null +++ b/tests/test_cli/test_wrap_codex.py @@ -0,0 +1,359 @@ +"""Tests for `headroom wrap codex` and `headroom unwrap codex`. + +These exercise the Codex-specific ``config.toml`` injection and restoration +helpers that route Codex through the Headroom proxy. They are deliberately +end-to-end-ish: the unit tests call the helpers directly against a temp +``$HOME``, and the integration tests invoke the real Click commands the same +way a user would from the shell. +""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import patch + +import pytest +from click.testing import CliRunner + +from headroom.cli import wrap as wrap_mod +from headroom.cli.main import main + + +def _set_test_home(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + home = str(tmp_path) + monkeypatch.setenv("HOME", home) + monkeypatch.setenv("USERPROFILE", home) + + +@pytest.fixture +def runner() -> CliRunner: + return CliRunner() + + +# --------------------------------------------------------------------------- +# Unit tests: helpers operating on ~/.codex/config.toml +# --------------------------------------------------------------------------- + + +class TestStripCodexHeadroomBlocks: + """Tests for the regex-based cleanup helper.""" + + def test_empty_content_returns_empty(self) -> None: + assert wrap_mod._strip_codex_headroom_blocks("") == "" + + def test_returns_content_unchanged_when_no_markers(self) -> None: + original = '[profiles.default]\nmodel = "gpt-4o"\n' + cleaned = wrap_mod._strip_codex_headroom_blocks(original) + # Trailing whitespace normalization only — semantic content preserved. + assert 'model = "gpt-4o"' in cleaned + assert "[profiles.default]" in cleaned + + def test_removes_complete_headroom_block(self) -> None: + wrapped = ( + f"{wrap_mod._CODEX_TOP_LEVEL_MARKER}\n" + 'model_provider = "headroom"\n' + "\n" + "[model_providers.headroom]\n" + 'base_url = "http://127.0.0.1:8787/v1"\n' + f"{wrap_mod._CODEX_END_MARKER}\n" + ) + assert wrap_mod._strip_codex_headroom_blocks(wrapped) == "" + + def test_preserves_user_content_around_block(self) -> None: + user_pre = '[profiles.default]\nmodel = "gpt-4o"\n' + user_post = '[mcp_servers.foo]\ncommand = "echo"\n' + wrapped = ( + f"{wrap_mod._CODEX_TOP_LEVEL_MARKER}\n" + 'model_provider = "headroom"\n' + f"{wrap_mod._CODEX_END_MARKER}\n" + user_pre + "\n" + f"{wrap_mod._CODEX_TOP_LEVEL_MARKER}\n" + "[model_providers.headroom]\n" + 'base_url = "http://127.0.0.1:8787/v1"\n' + f"{wrap_mod._CODEX_END_MARKER}\n" + user_post + ) + cleaned = wrap_mod._strip_codex_headroom_blocks(wrapped) + assert wrap_mod._CODEX_TOP_LEVEL_MARKER not in cleaned + assert wrap_mod._CODEX_END_MARKER not in cleaned + assert 'model = "gpt-4o"' in cleaned + assert "[mcp_servers.foo]" in cleaned + + def test_removes_stray_top_level_model_provider_line(self) -> None: + # Old wrap versions left `model_provider = "headroom"` outside markers. + content = 'foo = 1\nmodel_provider = "headroom"\nbar = 2\n' + cleaned = wrap_mod._strip_codex_headroom_blocks(content) + assert 'model_provider = "headroom"' not in cleaned + assert "foo = 1" in cleaned + assert "bar = 2" in cleaned + + +class TestSnapshotCodexConfig: + """Tests for ``_snapshot_codex_config_if_unwrapped``.""" + + def test_creates_backup_on_first_call(self, tmp_path: Path) -> None: + config_file = tmp_path / "config.toml" + backup_file = tmp_path / "config.toml.headroom-backup" + config_file.write_text('model = "gpt-4o"\n') + + wrap_mod._snapshot_codex_config_if_unwrapped(config_file, backup_file) + + assert backup_file.exists() + assert backup_file.read_text() == 'model = "gpt-4o"\n' + + def test_does_not_overwrite_existing_backup(self, tmp_path: Path) -> None: + config_file = tmp_path / "config.toml" + backup_file = tmp_path / "config.toml.headroom-backup" + config_file.write_text("second-wrap content\n") + backup_file.write_text("original-pre-wrap content\n") + + wrap_mod._snapshot_codex_config_if_unwrapped(config_file, backup_file) + + # Backup must still contain the *original* pre-wrap content. + assert backup_file.read_text() == "original-pre-wrap content\n" + + def test_no_backup_when_config_missing(self, tmp_path: Path) -> None: + config_file = tmp_path / "config.toml" + backup_file = tmp_path / "config.toml.headroom-backup" + + wrap_mod._snapshot_codex_config_if_unwrapped(config_file, backup_file) + + assert not backup_file.exists() + + def test_no_backup_when_config_already_wrapped(self, tmp_path: Path) -> None: + config_file = tmp_path / "config.toml" + backup_file = tmp_path / "config.toml.headroom-backup" + config_file.write_text( + f"{wrap_mod._CODEX_TOP_LEVEL_MARKER}\n" + 'model_provider = "headroom"\n' + f"{wrap_mod._CODEX_END_MARKER}\n" + ) + + wrap_mod._snapshot_codex_config_if_unwrapped(config_file, backup_file) + + # Pre-wrap snapshot must never snapshot an already-wrapped file. + assert not backup_file.exists() + + +class TestInjectAndRestoreRoundTrip: + """End-to-end wrap → unwrap cycle operating directly on a temp $HOME.""" + + def test_wrap_unwrap_restores_empty_state( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + _set_test_home(monkeypatch, tmp_path) + config_file = tmp_path / ".codex" / "config.toml" + + wrap_mod._inject_codex_provider_config(8787) + assert config_file.exists() + assert 'model_provider = "headroom"' in config_file.read_text() + + status, _ = wrap_mod._restore_codex_provider_config() + # No prior config existed → the injected file is fully removed. + assert status == "removed" + assert not config_file.exists() + assert not (tmp_path / ".codex" / "config.toml.headroom-backup").exists() + + def test_wrap_unwrap_restores_prior_model_provider( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + _set_test_home(monkeypatch, tmp_path) + config_dir = tmp_path / ".codex" + config_dir.mkdir() + config_file = config_dir / "config.toml" + original = ( + 'model_provider = "openai"\n' + "\n" + "[model_providers.openai]\n" + 'name = "OpenAI"\n' + 'base_url = "https://api.openai.com/v1"\n' + ) + config_file.write_text(original) + + wrap_mod._inject_codex_provider_config(8787) + wrapped = config_file.read_text() + assert 'model_provider = "headroom"' in wrapped + assert "[model_providers.headroom]" in wrapped + + status, _ = wrap_mod._restore_codex_provider_config() + assert status == "restored" + assert config_file.read_text() == original + assert not (config_dir / "config.toml.headroom-backup").exists() + + def test_wrap_is_idempotent(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + _set_test_home(monkeypatch, tmp_path) + config_dir = tmp_path / ".codex" + config_dir.mkdir() + config_file = config_dir / "config.toml" + original = '[profiles.default]\nmodel = "gpt-4o"\n' + config_file.write_text(original) + + wrap_mod._inject_codex_provider_config(8787) + wrap_mod._inject_codex_provider_config(8787) + wrap_mod._inject_codex_provider_config(9999) # port change + + content = config_file.read_text() + # Exactly two Headroom blocks — a top-level-key block and the + # provider-table block. Re-wrapping must not duplicate them. + assert content.count(wrap_mod._CODEX_TOP_LEVEL_MARKER) == 2 + assert content.count(wrap_mod._CODEX_END_MARKER) == 2 + # Latest port is honoured. + assert 'base_url = "http://127.0.0.1:9999/v1"' in content + assert 'base_url = "http://127.0.0.1:8787/v1"' not in content + # User's original content is preserved. + assert 'model = "gpt-4o"' in content + + status, _ = wrap_mod._restore_codex_provider_config() + assert status == "restored" + assert config_file.read_text() == original + + def test_unwrap_is_noop_when_never_wrapped( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + _set_test_home(monkeypatch, tmp_path) + + status, _ = wrap_mod._restore_codex_provider_config() + assert status == "noop" + + def test_unwrap_cleans_block_without_backup( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """Handles crash-case where wrap injected but backup was wiped.""" + _set_test_home(monkeypatch, tmp_path) + config_dir = tmp_path / ".codex" + config_dir.mkdir() + config_file = config_dir / "config.toml" + user_content = '[profiles.default]\nmodel = "gpt-4o"\n' + config_file.write_text( + user_content + f"{wrap_mod._CODEX_TOP_LEVEL_MARKER}\n" + 'model_provider = "headroom"\n\n' + "[model_providers.headroom]\n" + 'base_url = "http://127.0.0.1:8787/v1"\n' + f"{wrap_mod._CODEX_END_MARKER}\n" + ) + + status, _ = wrap_mod._restore_codex_provider_config() + assert status == "cleaned" + cleaned = config_file.read_text() + assert wrap_mod._CODEX_TOP_LEVEL_MARKER not in cleaned + assert wrap_mod._CODEX_END_MARKER not in cleaned + assert 'model_provider = "headroom"' not in cleaned + assert 'model = "gpt-4o"' in cleaned + + def test_unwrap_handles_malformed_prior_config( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """Unwrap preserves backup content verbatim — TOML validity isn't required.""" + _set_test_home(monkeypatch, tmp_path) + config_dir = tmp_path / ".codex" + config_dir.mkdir() + config_file = config_dir / "config.toml" + malformed = 'this is not valid toml ][ "" \x00\n' + config_file.write_text(malformed) + + wrap_mod._inject_codex_provider_config(8787) + status, _ = wrap_mod._restore_codex_provider_config() + + assert status == "restored" + assert config_file.read_text() == malformed + + +# --------------------------------------------------------------------------- +# Integration tests: full `headroom wrap codex` / `headroom unwrap codex` +# --------------------------------------------------------------------------- + + +def test_wrap_codex_prepare_only_creates_backup_and_config( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _set_test_home(monkeypatch, tmp_path) + config_file = tmp_path / ".codex" / "config.toml" + config_file.parent.mkdir(parents=True) + original = 'model_provider = "openai"\n' + config_file.write_text(original) + + with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None): + result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"]) + + assert result.exit_code == 0, result.output + assert 'model_provider = "headroom"' in config_file.read_text() + backup = tmp_path / ".codex" / "config.toml.headroom-backup" + assert backup.exists() + assert backup.read_text() == original + + +def test_unwrap_codex_restores_prior_config_end_to_end( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """The bug report, reproduced: wrap → unwrap must round-trip cleanly.""" + _set_test_home(monkeypatch, tmp_path) + config_file = tmp_path / ".codex" / "config.toml" + config_file.parent.mkdir(parents=True) + original = ( + "[profiles.default]\n" + 'model = "gpt-4o"\n' + "\n" + "[model_providers.openai]\n" + 'base_url = "https://api.openai.com/v1"\n' + ) + config_file.write_text(original) + + with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None): + wrap_result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"]) + assert wrap_result.exit_code == 0, wrap_result.output + assert 'model_provider = "headroom"' in config_file.read_text() + + unwrap_result = runner.invoke(main, ["unwrap", "codex"]) + assert unwrap_result.exit_code == 0, unwrap_result.output + + # Config must be byte-for-byte what the user had before wrap, and the + # injected block must be gone — no more "Missing OPENAI_API_KEY" when the + # proxy is stopped. + assert config_file.read_text() == original + assert 'model_provider = "headroom"' not in config_file.read_text() + assert not (tmp_path / ".codex" / "config.toml.headroom-backup").exists() + + +def test_unwrap_codex_is_safe_noop_with_no_prior_wrap( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _set_test_home(monkeypatch, tmp_path) + + result = runner.invoke(main, ["unwrap", "codex"]) + assert result.exit_code == 0, result.output + assert "Nothing to undo" in result.output + assert not (tmp_path / ".codex" / "config.toml").exists() + + +def test_unwrap_codex_removes_headroom_only_config_file( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _set_test_home(monkeypatch, tmp_path) + + with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None): + wrap_result = runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"]) + assert wrap_result.exit_code == 0, wrap_result.output + + config_file = tmp_path / ".codex" / "config.toml" + assert config_file.exists() + + unwrap_result = runner.invoke(main, ["unwrap", "codex"]) + assert unwrap_result.exit_code == 0, unwrap_result.output + assert not config_file.exists() + + +def test_unwrap_codex_preserves_unrelated_sections( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _set_test_home(monkeypatch, tmp_path) + config_file = tmp_path / ".codex" / "config.toml" + config_file.parent.mkdir(parents=True) + # A config with an MCP server the user configured by hand. + original = '[mcp_servers.local_thing]\ncommand = "/usr/local/bin/thing"\nargs = ["--serve"]\n' + config_file.write_text(original) + + with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None): + runner.invoke(main, ["wrap", "codex", "--prepare-only", "--port", "8787"]) + + result = runner.invoke(main, ["unwrap", "codex"]) + assert result.exit_code == 0, result.output + restored = config_file.read_text() + assert restored == original