diff --git a/.github/actions/headroom-e2e-setup/action.yml b/.github/actions/headroom-e2e-setup/action.yml index dc6d951a0..6dd345dd9 100644 --- a/.github/actions/headroom-e2e-setup/action.yml +++ b/.github/actions/headroom-e2e-setup/action.yml @@ -1,15 +1,25 @@ name: Headroom e2e setup description: >- Checkout-agnostic setup shared by native e2e workflows (init, install, wrap). - Installs Python + Rust toolchain, installs headroom in editable mode (which - builds the bundled Rust extension via maturin), and (optionally) drops a - noop shim onto PATH so ``headroom init -g `` can detect a tool - that isn't actually installed on the runner. + Installs Python, optionally installs the Rust toolchain + editable headroom + package, and (optionally) drops PATH shims for the local ``headroom`` CLI and + target binaries so ``headroom init -g `` can detect tools that aren't + actually installed on the runner. inputs: python-version: description: Python version to install required: false default: "3.11" + install-mode: + description: >- + Install strategy. ``editable-proxy`` builds the local package with + ``pip install -e .[proxy]`` and verifies ``headroom._core``. + ``deps-only-proxy`` installs the base + ``[proxy]`` dependency set from + pyproject.toml, then drops a local ``headroom`` launcher that imports + from the checkout without building the package; use this for CLI tests + that do not exercise the Rust extension. + required: false + default: "editable-proxy" shim-target: description: >- Name of the shim to drop on PATH (e.g. ``claude``, ``codex``). Leave @@ -30,16 +40,19 @@ runs: # Single-wheel architecture: `pip install -e .` invokes maturin (declared # in pyproject.toml's build-system) which calls cargo to compile the Rust - # extension. Toolchain has to be set up before the install step. + # extension. Toolchain has to be set up before the editable install path. - name: Install Rust toolchain + if: ${{ inputs.install-mode == 'editable-proxy' }} uses: dtolnay/rust-toolchain@1.95.0 - name: Cache cargo registry + build + if: ${{ inputs.install-mode == 'editable-proxy' }} uses: Swatinem/rust-cache@v2 with: workspaces: ". -> target" - name: Install headroom (editable, with proxy extras — builds Rust extension) + if: ${{ inputs.install-mode == 'editable-proxy' }} shell: bash run: | python -m pip install --upgrade pip @@ -49,6 +62,51 @@ runs: pip install -e ".[proxy]" python -c "from headroom._core import DiffCompressor; print('headroom._core OK:', DiffCompressor)" + - name: Install base + proxy dependencies without building headroom + if: ${{ inputs.install-mode == 'deps-only-proxy' }} + shell: bash + run: | + python -m pip install --upgrade pip + python - <<'PY' + import subprocess + import sys + import tomllib + from pathlib import Path + + project = tomllib.loads(Path("pyproject.toml").read_text(encoding="utf-8")) + requirements = list(project["project"]["dependencies"]) + requirements.extend(project["project"]["optional-dependencies"]["proxy"]) + subprocess.check_call( + [sys.executable, "-m", "pip", "install", "--retries", "10", "--timeout", "60", *requirements] + ) + PY + python -c "from headroom.cli.main import main; print('headroom CLI OK:', main)" + + - name: Drop local headroom launcher (POSIX) + if: ${{ inputs.install-mode == 'deps-only-proxy' && runner.os != 'Windows' }} + shell: bash + run: | + shim_dir="${RUNNER_TEMP}/headroom-local-bin" + mkdir -p "$shim_dir" + cat > "$shim_dir/headroom" <<'SH' + #!/usr/bin/env bash + exec python -m headroom.cli "$@" + SH + chmod +x "$shim_dir/headroom" + echo "$shim_dir" >> "$GITHUB_PATH" + + - name: Drop local headroom launcher (Windows) + if: ${{ inputs.install-mode == 'deps-only-proxy' && runner.os == 'Windows' }} + shell: pwsh + run: | + $shimDir = Join-Path $env:RUNNER_TEMP "headroom-local-bin" + New-Item -ItemType Directory -Force -Path $shimDir | Out-Null + @" + @echo off + python -m headroom.cli %* + "@ | Out-File -FilePath (Join-Path $shimDir "headroom.cmd") -Encoding ascii + Add-Content -Path $env:GITHUB_PATH -Value $shimDir + - name: Drop shim (POSIX) if: ${{ inputs.shim-target != '' && runner.os != 'Windows' }} id: shim-posix diff --git a/.github/workflows/init-native-e2e.yml b/.github/workflows/init-native-e2e.yml index 652a615f0..49db1f68b 100644 --- a/.github/workflows/init-native-e2e.yml +++ b/.github/workflows/init-native-e2e.yml @@ -60,6 +60,7 @@ jobs: - name: Setup (shim=${{ matrix.target }}) uses: ./.github/actions/headroom-e2e-setup with: + install-mode: deps-only-proxy python-version: "3.11" shim-target: ${{ matrix.target }} diff --git a/.github/workflows/install-native-e2e.yml b/.github/workflows/install-native-e2e.yml index ed8992a00..0349013f8 100644 --- a/.github/workflows/install-native-e2e.yml +++ b/.github/workflows/install-native-e2e.yml @@ -44,6 +44,7 @@ jobs: - name: Setup uses: ./.github/actions/headroom-e2e-setup with: + install-mode: deps-only-proxy python-version: "3.11" - name: Install pytest diff --git a/.github/workflows/wrap-native-e2e.yml b/.github/workflows/wrap-native-e2e.yml index ce41b8b4c..8091c2f2e 100644 --- a/.github/workflows/wrap-native-e2e.yml +++ b/.github/workflows/wrap-native-e2e.yml @@ -49,6 +49,7 @@ jobs: - name: Setup uses: ./.github/actions/headroom-e2e-setup with: + install-mode: deps-only-proxy python-version: "3.11" - name: Install pytest diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ec13e1e7..53398f69d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * **proxy:** force Responses API `store=true` when Headroom injects memory tools so `previous_response_id` continuations work after memory tool calls from clients that requested `store=false` ([#1103](https://github.com/chopratejas/headroom/pull/1103)). * **proxy:** build SSL contexts for custom CA bundles so enterprise/private PKI roots work with Python/OpenSSL strict verification. +* **codex:** stop persisting a project-specific `--db` path in the global `headroom_memory` MCP config, so `headroom wrap codex --memory` falls back to the active cwd's `.headroom/memory.db` at runtime while keeping the current project's local bootstrap work scoped correctly ([#1147](https://github.com/chopratejas/headroom/issues/1147)). * **proxy:** route Codex OAuth image generation and edit requests through the ChatGPT Codex image backend, while preserving OpenAI API-key image passthrough ([#1215](https://github.com/chopratejas/headroom/pull/1215)). * **wrap (codex):** keep RTK guidance in the global Codex `AGENTS.md` instead of modifying the shared project `AGENTS.md` ([#1235](https://github.com/chopratejas/headroom/issues/1235)). * **proxy:** enable SSO credential resolution in the native Bedrock route via the `aws-config` `sso` feature flag, making the credential chain match what `docs/bedrock.md` already documented ([#999](https://github.com/chopratejas/headroom/pull/999)). diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index d0bc0f76d..5c6f94c7e 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -1080,7 +1080,12 @@ def _codex_config_paths() -> tuple[Path, Path]: return config_file, backup_file -def _strip_codex_headroom_blocks(content: str, *, remove_mcp: bool = False) -> str: +def _strip_codex_headroom_blocks( + content: str, + *, + remove_mcp: bool = False, + remove_named_mcp: bool = True, +) -> str: """Remove all Headroom-managed blocks from a Codex ``config.toml`` string. Returns the cleaned content. Safe to call on content that never contained @@ -1107,12 +1112,13 @@ def _strip_codex_headroom_blocks(content: str, *, remove_mcp: bool = False) -> s if remove_mcp: # Remove Headroom-managed MCP blocks written by `wrap codex`. content = _remove_marker_span(content, _CODEX_MCP_MARKER, _CODEX_MCP_END) - content = re.sub( - r"(?ms)^# --- Headroom MCP server: [^\n]+ ---\n.*?" - r"^# --- end Headroom MCP server: [^\n]+ ---\n?", - "", - content, - ) + if remove_named_mcp: + content = re.sub( + r"(?ms)^# --- Headroom MCP server: [^\n]+ ---\n.*?" + r"^# --- end Headroom MCP server: [^\n]+ ---\n?", + "", + content, + ) content = _remove_marker_span(content, _MEMORY_MCP_MARKER, _MEMORY_MCP_END) # Strip any leftover top-level keys that older (or crashed) versions of @@ -1201,6 +1207,17 @@ def _has_redirectable_top_level_key(content: str, key: str) -> bool: return pattern.search(content) is not None +def _codex_config_has_headroom_markers(content: str) -> bool: + """Return whether a Codex config already contains wrap-owned markers.""" + managed_markers = ( + _CODEX_TOP_LEVEL_MARKER, + _CODEX_END_MARKER, + _CODEX_MCP_MARKER, + _MEMORY_MCP_MARKER, + ) + return any(marker in content for marker in managed_markers) + + def _snapshot_codex_config_if_unwrapped(config_file: Path, backup_file: Path) -> None: """Snapshot ``config.toml`` to ``backup_file`` before the first injection. @@ -1214,8 +1231,8 @@ def _snapshot_codex_config_if_unwrapped(config_file: Path, backup_file: Path) -> *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 the config already contains any Headroom-managed Codex marker, a wrap + run is already active: do not snapshot the injected state. """ if backup_file.exists(): return @@ -1225,7 +1242,7 @@ def _snapshot_codex_config_if_unwrapped(config_file: Path, backup_file: Path) -> content = config_file.read_text() except OSError: return - if _CODEX_TOP_LEVEL_MARKER in content or _CODEX_END_MARKER in content: + if _codex_config_has_headroom_markers(content): return backup_file.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(config_file, backup_file) @@ -1469,8 +1486,23 @@ def _restore_codex_provider_config() -> tuple[str, Path]: # 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, remove_mcp=True) + if _codex_config_has_headroom_markers(original): + # Without a backup, only remove named MCP blocks when this file + # also carries wrap-owned provider markers from a full wrap. + remove_named_mcp = any( + marker in original + for marker in ( + _CODEX_TOP_LEVEL_MARKER, + _CODEX_END_MARKER, + _CODEX_MCP_MARKER, + _CODEX_MCP_END, + ) + ) + cleaned = _strip_codex_headroom_blocks( + original, + remove_mcp=True, + remove_named_mcp=remove_named_mcp, + ) if not cleaned.strip(): # Nothing left but Headroom content — remove the file entirely # so Codex falls back to its default config. @@ -1691,7 +1723,7 @@ def _remove_rtk_instructions(file_path: Path) -> bool: return True -def _inject_memory_mcp_config(db_path: str, user_id: str) -> None: +def _inject_memory_mcp_config(user_id: str) -> None: """Register headroom memory as an MCP server in Codex's config.toml. Idempotent — replaces existing section if present. @@ -1704,12 +1736,11 @@ def _inject_memory_mcp_config(db_path: str, user_id: str) -> None: # Use forward slashes in TOML paths (works on all platforms, avoids # backslash escaping issues on Windows) python_bin = sys.executable.replace("\\", "/") - db_path_toml = db_path.replace("\\", "/") mcp_section = ( f"\n{_MEMORY_MCP_MARKER}\n" f"[mcp_servers.headroom_memory]\n" f'command = "{python_bin}"\n' - f'args = ["-m", "headroom.memory.mcp_server", "--db", "{db_path_toml}", "--user", "{user_id}"]\n' + f'args = ["-m", "headroom.memory.mcp_server", "--user", "{user_id}"]\n' f"startup_timeout_sec = 30\n" f"tool_timeout_sec = 30\n" f"{_MEMORY_MCP_END}\n" @@ -3813,7 +3844,7 @@ def codex( mem_user = os.environ.get("USER", os.environ.get("USERNAME", "default")) # Register MCP server in Codex config - _inject_memory_mcp_config(db_path, mem_user) + _inject_memory_mcp_config(mem_user) # Inject memory guidance into project AGENTS.md agents_md = Path.cwd() / "AGENTS.md" @@ -3871,11 +3902,7 @@ def codex( # the config file. Re-inject MCP config after if memory is enabled. _inject_codex_provider_config(port) if memory: - mem_dir = Path.cwd() / ".headroom" - _inject_memory_mcp_config( - str(mem_dir / "memory.db"), - os.environ.get("USER", os.environ.get("USERNAME", "default")), - ) + _inject_memory_mcp_config(os.environ.get("USER", os.environ.get("USERNAME", "default"))) _launch_tool( binary=codex_bin, @@ -4991,9 +5018,8 @@ def opencode( click.echo(" Setting up memory for OpenCode...") mem_dir = Path.cwd() / ".headroom" mem_dir.mkdir(parents=True, exist_ok=True) - db_path = str(mem_dir / "memory.db") mem_user = os.environ.get("USER", os.environ.get("USERNAME", "default")) - _inject_memory_mcp_config(db_path, mem_user) + _inject_memory_mcp_config(mem_user) agents_md = Path.cwd() / "AGENTS.md" _inject_memory_agents_md(agents_md) @@ -5016,7 +5042,6 @@ def opencode( if memory: mem_dir = Path.cwd() / ".headroom" _inject_memory_mcp_config( - str(mem_dir / "memory.db"), os.environ.get("USER", os.environ.get("USERNAME", "default")), ) diff --git a/headroom/memory/mcp_server.py b/headroom/memory/mcp_server.py index e767c8efa..d5fecab97 100644 --- a/headroom/memory/mcp_server.py +++ b/headroom/memory/mcp_server.py @@ -20,7 +20,8 @@ Usage: # Registered in Codex config.toml (done by `headroom wrap codex --memory`): [mcp_servers.headroom_memory] command = "python" - args = ["-m", "headroom.memory.mcp_server", "--db", ".headroom/memory.db"] + args = ["-m", "headroom.memory.mcp_server", "--user", "alice"] + # When --db is omitted, the server resolves .headroom/memory.db from cwd. """ from __future__ import annotations @@ -344,6 +345,42 @@ async def _run(db_path: str, user_id: str) -> None: await server.run(read_stream, write_stream, server.create_initialization_options()) +def _memory_mcp_startup_context( + configured_db: str, cwd: Path, db_flag_present: bool +) -> dict[str, str | bool]: + """Describe the DB path the memory MCP server will try to open.""" + configured_path = Path(configured_db).expanduser() + resolved_path = ( + configured_path if configured_path.is_absolute() else (cwd / configured_path) + ).resolve(strict=False) + active_project_db = (cwd / ".headroom" / "memory.db").resolve(strict=False) + if not db_flag_present: + config_source = "cwd-default" + resolution = "dynamic-cwd" + else: + config_source = "cli-flag" + resolution = "static-cli" + if resolved_path == active_project_db: + storage_scope = "active-project" + elif resolved_path.name == "memory.db": + storage_scope = "external-memory-db" + else: + storage_scope = "custom-db-path" + path_exists = resolved_path.exists() + path_readable = path_exists and os.access(resolved_path, os.R_OK) + return { + "configured_db": str(configured_path), + "resolved_db": str(resolved_path), + "config_source": config_source, + "cwd": str(cwd), + "project_root": str(cwd), + "storage_scope": storage_scope, + "path_exists": path_exists, + "path_readable": path_readable, + "resolution": resolution, + } + + def main() -> None: parser = argparse.ArgumentParser(description="Headroom Memory MCP Server") parser.add_argument( @@ -370,6 +407,26 @@ def main() -> None: format="%(name)s: %(message)s", ) + startup = _memory_mcp_startup_context( + configured_db=args.db, + cwd=Path.cwd(), + db_flag_present=any(arg == "--db" or arg.startswith("--db=") for arg in sys.argv[1:]), + ) + logger.info( + "Memory MCP startup: configured_db=%s, resolved_db=%s, config_source=%s, " + "cwd=%s, project_root=%s, storage_scope=%s, path_exists=%s, " + "path_readable=%s, resolution=%s", + startup["configured_db"], + startup["resolved_db"], + startup["config_source"], + startup["cwd"], + startup["project_root"], + startup["storage_scope"], + startup["path_exists"], + startup["path_readable"], + startup["resolution"], + ) + asyncio.run(_run(args.db, args.user)) diff --git a/tests/_mcp_stub.py b/tests/_mcp_stub.py new file mode 100644 index 000000000..e788feefb --- /dev/null +++ b/tests/_mcp_stub.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +import importlib +import sys +from types import ModuleType + +_MCP_MODULE_NAMES = ( + "mcp", + "mcp.server", + "mcp.server.stdio", + "mcp.types", +) + + +def _build_mcp_sdk_stub() -> dict[str, ModuleType]: + mcp_module = type(sys)("mcp") + mcp_server_module = type(sys)("mcp.server") + mcp_stdio_module = type(sys)("mcp.server.stdio") + mcp_types_module = type(sys)("mcp.types") + + class DummyServer: + def __init__(self, name: str) -> None: + self.name = name + + def list_tools(self): + return lambda fn: fn + + def call_tool(self): + return lambda fn: fn + + def create_initialization_options(self): + return {} + + class DummyTool: + def __init__(self, **kwargs) -> None: + self.kwargs = kwargs + + class DummyTextContent: + def __init__(self, **kwargs) -> None: + self.kwargs = kwargs + + async def dummy_stdio_server(): + raise RuntimeError("stdio_server should not run in unit tests") + + mcp_server_module.Server = DummyServer + mcp_stdio_module.stdio_server = dummy_stdio_server + mcp_types_module.TextContent = DummyTextContent + mcp_types_module.Tool = DummyTool + + return { + "mcp": mcp_module, + "mcp.server": mcp_server_module, + "mcp.server.stdio": mcp_stdio_module, + "mcp.types": mcp_types_module, + } + + +def import_module_with_mcp_stub(module_name: str): + original_target_module = sys.modules.get(module_name) + parent_name, _, child_name = module_name.rpartition(".") + original_parent_module = sys.modules.get(parent_name) if parent_name else None + original_parent_attr_exists = bool( + original_parent_module and child_name and hasattr(original_parent_module, child_name) + ) + original_parent_attr = ( + getattr(original_parent_module, child_name) if original_parent_attr_exists else None + ) + original_modules = {name: sys.modules.get(name) for name in _MCP_MODULE_NAMES} + stub_modules = _build_mcp_sdk_stub() + + sys.modules.pop(module_name, None) + for name, module in stub_modules.items(): + sys.modules[name] = module + + try: + return importlib.import_module(module_name) + finally: + if original_target_module is None: + sys.modules.pop(module_name, None) + else: + sys.modules[module_name] = original_target_module + if child_name: + current_parent_module = sys.modules.get(parent_name) or original_parent_module + if current_parent_module is not None: + if original_parent_attr_exists: + setattr(current_parent_module, child_name, original_parent_attr) + else: + try: + delattr(current_parent_module, child_name) + except AttributeError: + pass + for name, original_module in original_modules.items(): + if original_module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = original_module diff --git a/tests/test_ccr_mcp_server.py b/tests/test_ccr_mcp_server.py index 7a353b1e2..c4a41848d 100644 --- a/tests/test_ccr_mcp_server.py +++ b/tests/test_ccr_mcp_server.py @@ -9,7 +9,9 @@ from headroom.cache.compression_store import ( get_compression_store, reset_compression_store, ) -from headroom.ccr import mcp_server +from tests._mcp_stub import import_module_with_mcp_stub + +mcp_server = import_module_with_mcp_stub("headroom.ccr.mcp_server") def test_shared_stats_work_without_fcntl(monkeypatch, tmp_path) -> None: @@ -47,7 +49,6 @@ def fresh_store(): def test_mcp_uses_shared_singleton_store(fresh_store) -> None: """MCP's store is the global singleton, not a private instance.""" - pytest.importorskip("mcp", reason="MCP SDK required") server = mcp_server.HeadroomMCPServer(check_proxy=False) assert server._get_local_store() is get_compression_store() @@ -56,7 +57,6 @@ def test_mcp_retrieves_proxy_stored_content(fresh_store) -> None: """Content stored via the singleton (as the proxy does) is retrievable through MCP's local-store path. The HTTP fallback is disabled so this passes only via the shared store.""" - pytest.importorskip("mcp", reason="MCP SDK required") original = '{"some": "original proxy-compressed content"}' hash_key = get_compression_store().store(original, '{"compressed": true}') @@ -95,7 +95,6 @@ def test_mcp_retrieve_with_nonmatching_query_returns_full_content(fresh_store) - """A query that matches no item above the relevance floor must still return the stored entry (it exists and is unexpired) rather than the "Content not found" error, which is reserved for genuine misses.""" - pytest.importorskip("mcp", reason="MCP SDK required") original = "the the the the the the the the the the\n" * 5 hash_key = get_compression_store().store(original, "<>") # Precondition: the query genuinely matches nothing above the BM25 floor. @@ -112,7 +111,6 @@ def test_mcp_retrieve_with_nonmatching_query_returns_full_content(fresh_store) - def test_mcp_retrieve_missing_hash_still_errors(fresh_store) -> None: """A genuinely missing hash must still report "Content not found".""" - pytest.importorskip("mcp", reason="MCP SDK required") server = mcp_server.HeadroomMCPServer(check_proxy=False) result = asyncio.run(server._retrieve_content("nonexistent_hash", query="anything")) assert "Content not found" in result.get("error", "") diff --git a/tests/test_cli/test_wrap_codex.py b/tests/test_cli/test_wrap_codex.py index 1aaa0173f..ceb4091f9 100644 --- a/tests/test_cli/test_wrap_codex.py +++ b/tests/test_cli/test_wrap_codex.py @@ -9,6 +9,7 @@ way a user would from the shell. from __future__ import annotations +import shutil import sqlite3 from pathlib import Path from unittest.mock import patch @@ -111,6 +112,25 @@ class TestStripCodexHeadroomBlocks: assert "[mcp_servers.headroom_memory]" not in cleaned assert 'model = "gpt-4o"' in cleaned + def test_preserves_named_mcp_blocks_when_remove_named_mcp_false(self) -> None: + content = ( + "# --- Headroom MCP server: serena ---\n" + "[mcp_servers.serena]\n" + 'command = "uvx"\n' + "# --- end Headroom MCP server: serena ---\n\n" + f"{wrap_mod._MEMORY_MCP_MARKER}\n" + "[mcp_servers.headroom_memory]\n" + 'command = "python"\n' + f"{wrap_mod._MEMORY_MCP_END}\n" + ) + + cleaned = wrap_mod._strip_codex_headroom_blocks( + content, remove_mcp=True, remove_named_mcp=False + ) + + assert "[mcp_servers.serena]" in cleaned + assert "[mcp_servers.headroom_memory]" not in cleaned + class TestSnapshotCodexConfig: """Tests for ``_snapshot_codex_config_if_unwrapped``.""" @@ -158,6 +178,65 @@ class TestSnapshotCodexConfig: # Pre-wrap snapshot must never snapshot an already-wrapped file. assert not backup_file.exists() + def test_no_backup_when_config_already_contains_memory_mcp_block(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._MEMORY_MCP_MARKER}\n" + "[mcp_servers.headroom_memory]\n" + 'command = "python"\n' + 'args = ["-m", "headroom.memory.mcp_server", "--user", "codex-user"]\n' + f"{wrap_mod._MEMORY_MCP_END}\n" + ) + + wrap_mod._snapshot_codex_config_if_unwrapped(config_file, backup_file) + + assert not backup_file.exists() + + def test_backup_when_config_contains_named_mcp_marker(self, tmp_path: Path) -> None: + config_file = tmp_path / "config.toml" + backup_file = tmp_path / "config.toml.headroom-backup" + original = ( + "# --- Headroom MCP server: headroom ---\n" + "[mcp_servers.headroom]\n" + 'command = "headroom"\n' + "# --- end Headroom MCP server: headroom ---\n" + ) + config_file.write_text(original) + + wrap_mod._snapshot_codex_config_if_unwrapped(config_file, backup_file) + + assert backup_file.exists() + assert backup_file.read_text() == original + + +class TestCodexMemoryMcpConfig: + """Tests for the persisted Codex memory MCP block.""" + + def test_inject_omits_db_and_replaces_existing_memory_block( + self, 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) + config_file.write_text( + '[profiles.default]\nmodel = "gpt-4o"\n\n' + f"{wrap_mod._MEMORY_MCP_MARKER}\n" + "[mcp_servers.headroom_memory]\n" + 'command = "python"\n' + 'args = ["-m", "headroom.memory.mcp_server", "--db", "/tmp/project-a/.headroom/memory.db", "--user", "old-user"]\n' + f"{wrap_mod._MEMORY_MCP_END}\n" + ) + + wrap_mod._inject_memory_mcp_config("codex-user") + + content = config_file.read_text() + assert content.count(wrap_mod._MEMORY_MCP_MARKER) == 1 + assert "[mcp_servers.headroom_memory]" in content + assert '"--user", "codex-user"' in content + assert "--db" not in content + assert 'model = "gpt-4o"' in content + class TestInjectAndRestoreRoundTrip: """End-to-end wrap → unwrap cycle operating directly on a temp $HOME.""" @@ -312,6 +391,52 @@ class TestInjectAndRestoreRoundTrip: assert 'model = "gpt-4o"' in cleaned assert "headroom" not in cleaned + def test_memory_only_wrap_restores_preexisting_named_mcp_block( + 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 = ( + "# --- Headroom MCP server: headroom ---\n" + "[mcp_servers.headroom]\n" + 'command = "headroom"\n' + "# --- end Headroom MCP server: headroom ---\n" + ) + config_file.write_text(original) + + wrap_mod._inject_memory_mcp_config("codex-user") + + status, _ = wrap_mod._restore_codex_provider_config() + + assert status == "restored" + assert config_file.read_text() == original + + def test_memory_only_wrap_without_backup_preserves_named_mcp_block( + 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" + backup_file = config_dir / "config.toml.headroom-backup" + original = ( + "# --- Headroom MCP server: headroom ---\n" + "[mcp_servers.headroom]\n" + 'command = "headroom"\n' + "# --- end Headroom MCP server: headroom ---\n" + ) + config_file.write_text(original) + + wrap_mod._inject_memory_mcp_config("codex-user") + backup_file.unlink() + + status, _ = wrap_mod._restore_codex_provider_config() + + assert status == "cleaned" + assert config_file.read_text() == original + def test_unwrap_handles_malformed_prior_config( self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: @@ -879,6 +1004,71 @@ def test_wrap_codex_prepare_only_updates_stale_mcp_proxy_url( assert "http://127.0.0.1:9000" not in content +def test_wrap_codex_memory_prepare_only_uses_local_db_without_persisting_it( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _set_test_home(monkeypatch, tmp_path) + monkeypatch.setenv("USER", "codex-user") + project_dir = tmp_path / "project-a" + project_dir.mkdir() + monkeypatch.chdir(project_dir) + backend_paths: list[str] = [] + imported_users: list[str] = [] + + class FakeBackend: + async def _ensure_initialized(self) -> None: + return None + + async def close(self) -> None: + return None + + class FakeClaudeCodeAdapter: + def __init__(self, memory_dir: Path) -> None: + self.memory_dir = memory_dir + + def fake_build_sync_backend(db_path: str) -> FakeBackend: + backend_paths.append(db_path) + return FakeBackend() + + async def fake_sync_import( + backend: FakeBackend, adapter: FakeClaudeCodeAdapter, user_id: str + ) -> int: + imported_users.append(user_id) + return 0 + + with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None): + with patch("headroom.memory.sync._build_sync_backend", side_effect=fake_build_sync_backend): + with patch("headroom.memory.sync.sync_import", side_effect=fake_sync_import): + with patch( + "headroom.memory.sync_adapters.claude_code.ClaudeCodeAdapter", + FakeClaudeCodeAdapter, + ): + with patch( + "headroom.memory.sync_adapters.claude_code.get_claude_memory_dir", + return_value=tmp_path / "claude-memory", + ): + result = runner.invoke( + main, + [ + "wrap", + "codex", + "--memory", + "--prepare-only", + "--no-mcp", + "--no-serena", + ], + ) + + assert result.exit_code == 0, result.output + assert backend_paths == [str(project_dir / ".headroom" / "memory.db")] + assert imported_users == ["codex-user"] + + content = (tmp_path / ".codex" / "config.toml").read_text() + assert "[mcp_servers.headroom_memory]" in content + assert '"--user", "codex-user"' in content + assert "--db" not in content + + def test_wrap_codex_prepare_only_registers_serena_when_uvx_exists( runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: @@ -969,6 +1159,117 @@ def test_unwrap_codex_no_stop_proxy_leaves_proxy_alone( stop_proxy.assert_not_called() +def test_wrap_codex_memory_prepare_only_unwrap_removes_memory_mcp_without_prior_config( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _set_test_home(monkeypatch, tmp_path) + monkeypatch.setenv("USER", "codex-user") + project_dir = tmp_path / "project-a" + project_dir.mkdir() + monkeypatch.chdir(project_dir) + + class FakeBackend: + async def _ensure_initialized(self) -> None: + return None + + async def close(self) -> None: + return None + + async def fake_sync_import(backend: FakeBackend, adapter: object, user_id: str) -> int: + return 0 + + with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None): + with patch("headroom.memory.sync._build_sync_backend", return_value=FakeBackend()): + with patch("headroom.memory.sync.sync_import", side_effect=fake_sync_import): + with patch( + "headroom.memory.sync_adapters.claude_code.ClaudeCodeAdapter", + autospec=True, + ): + with patch( + "headroom.memory.sync_adapters.claude_code.get_claude_memory_dir", + return_value=tmp_path / "claude-memory", + ): + wrap_result = runner.invoke( + main, + [ + "wrap", + "codex", + "--memory", + "--prepare-only", + "--no-mcp", + "--no-serena", + ], + ) + + assert wrap_result.exit_code == 0, wrap_result.output + config_file = tmp_path / ".codex" / "config.toml" + content = config_file.read_text() + assert "[mcp_servers.headroom_memory]" in content + assert '"--user", "codex-user"' in content + + with patch("headroom.cli.wrap._stop_local_proxy_for_unwrap") as stop_proxy: + unwrap_result = runner.invoke(main, ["unwrap", "codex", "--no-stop-proxy"]) + + assert unwrap_result.exit_code == 0, unwrap_result.output + assert not config_file.exists() + assert not (tmp_path / ".codex" / "config.toml.headroom-backup").exists() + stop_proxy.assert_not_called() + + +def test_wrap_codex_memory_launch_failure_unwrap_cleans_memory_only_config( + runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + _set_test_home(monkeypatch, tmp_path) + monkeypatch.setenv("USER", "codex-user") + project_dir = tmp_path / "project-a" + project_dir.mkdir() + monkeypatch.chdir(project_dir) + + class FakeBackend: + async def _ensure_initialized(self) -> None: + return None + + async def close(self) -> None: + return None + + async def fake_sync_import(backend: FakeBackend, adapter: object, user_id: str) -> int: + return 0 + + def fake_which(cmd: str) -> str | None: + return None if cmd == "codex" else shutil.which(cmd) + + with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None): + with patch("headroom.cli.wrap.shutil.which", side_effect=fake_which): + with patch("headroom.memory.sync._build_sync_backend", return_value=FakeBackend()): + with patch("headroom.memory.sync.sync_import", side_effect=fake_sync_import): + with patch( + "headroom.memory.sync_adapters.claude_code.ClaudeCodeAdapter", + autospec=True, + ): + with patch( + "headroom.memory.sync_adapters.claude_code.get_claude_memory_dir", + return_value=tmp_path / "claude-memory", + ): + wrap_result = runner.invoke( + main, + ["wrap", "codex", "--memory", "--no-mcp", "--no-serena"], + ) + + assert wrap_result.exit_code == 1 + config_file = tmp_path / ".codex" / "config.toml" + content = config_file.read_text() + assert "[mcp_servers.headroom_memory]" in content + assert wrap_mod._CODEX_TOP_LEVEL_MARKER not in content + + with patch("headroom.cli.wrap._stop_local_proxy_for_unwrap") as stop_proxy: + unwrap_result = runner.invoke(main, ["unwrap", "codex", "--no-stop-proxy"]) + + assert unwrap_result.exit_code == 0, unwrap_result.output + assert not config_file.exists() + assert not (tmp_path / ".codex" / "config.toml.headroom-backup").exists() + stop_proxy.assert_not_called() + + def test_stop_local_proxy_for_unwrap_kills_identified_headroom_proxy( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/test_mcp_stub.py b/tests/test_mcp_stub.py new file mode 100644 index 000000000..b3b1a72f4 --- /dev/null +++ b/tests/test_mcp_stub.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import asyncio +import sys +from types import ModuleType + +import pytest + +from tests import _mcp_stub as mcp_stub + + +def test_build_mcp_sdk_stub_exposes_minimum_sdk_contract() -> None: + modules = mcp_stub._build_mcp_sdk_stub() + + server = modules["mcp.server"].Server("headroom") + assert server.name == "headroom" + + sentinel = object() + assert server.list_tools()(sentinel) is sentinel + assert server.call_tool()(sentinel) is sentinel + assert server.create_initialization_options() == {} + + tool = modules["mcp.types"].Tool(name="search") + text = modules["mcp.types"].TextContent(text="payload") + assert tool.kwargs == {"name": "search"} + assert text.kwargs == {"text": "payload"} + + with pytest.raises(RuntimeError, match="should not run"): + asyncio.run(modules["mcp.server.stdio"].stdio_server()) + + +def test_import_module_with_mcp_stub_imports_target_and_cleans_up(monkeypatch) -> None: + original_mcp = ModuleType("mcp") + monkeypatch.setitem(sys.modules, "mcp", original_mcp) + for name in mcp_stub._MCP_MODULE_NAMES[1:]: + sys.modules.pop(name, None) + + imported = ModuleType("fake_target") + + def fake_import_module(module_name: str) -> ModuleType: + assert module_name == "fake_target" + for name in mcp_stub._MCP_MODULE_NAMES: + assert name in sys.modules + return imported + + monkeypatch.setattr(mcp_stub.importlib, "import_module", fake_import_module) + + result = mcp_stub.import_module_with_mcp_stub("fake_target") + + assert result is imported + assert sys.modules["mcp"] is original_mcp + for name in mcp_stub._MCP_MODULE_NAMES[1:]: + assert name not in sys.modules + + +def test_import_module_with_mcp_stub_reimports_target_and_restores_originals(monkeypatch) -> None: + original_modules = {} + for name in mcp_stub._MCP_MODULE_NAMES: + module = ModuleType(f"original::{name}") + original_modules[name] = module + monkeypatch.setitem(sys.modules, name, module) + + existing = ModuleType("fake_target") + monkeypatch.setitem(sys.modules, "fake_target", existing) + imported = ModuleType("fake_target") + + def fake_import_module(module_name: str) -> ModuleType: + assert module_name == "fake_target" + assert "fake_target" not in sys.modules + for name in mcp_stub._MCP_MODULE_NAMES: + assert sys.modules[name] is not original_modules[name] + return imported + + monkeypatch.setattr(mcp_stub.importlib, "import_module", fake_import_module) + + result = mcp_stub.import_module_with_mcp_stub("fake_target") + + assert result is imported + assert sys.modules["fake_target"] is existing + for name, module in original_modules.items(): + assert sys.modules[name] is module + + +def test_import_module_with_mcp_stub_cleans_up_dotted_target_attribute(monkeypatch) -> None: + parent = ModuleType("fakepkg") + monkeypatch.setitem(sys.modules, "fakepkg", parent) + for name in mcp_stub._MCP_MODULE_NAMES: + sys.modules.pop(name, None) + + imported = ModuleType("fakepkg.fake_target") + + def fake_import_module(module_name: str) -> ModuleType: + assert module_name == "fakepkg.fake_target" + parent.fake_target = imported + return imported + + monkeypatch.setattr(mcp_stub.importlib, "import_module", fake_import_module) + + result = mcp_stub.import_module_with_mcp_stub("fakepkg.fake_target") + + assert result is imported + assert not hasattr(parent, "fake_target") diff --git a/tests/test_memory/test_mcp_server.py b/tests/test_memory/test_mcp_server.py index 85207d2f4..bd56174ad 100644 --- a/tests/test_memory/test_mcp_server.py +++ b/tests/test_memory/test_mcp_server.py @@ -1,17 +1,16 @@ +import asyncio from types import SimpleNamespace from unittest.mock import AsyncMock import numpy as np -import pytest -pytest.importorskip("mcp") - -from headroom.memory.mcp_server import _warm_up_backend from headroom.memory.models import Memory +from tests._mcp_stub import import_module_with_mcp_stub + +mcp_server_mod = import_module_with_mcp_stub("headroom.memory.mcp_server") -@pytest.mark.asyncio -async def test_warm_up_backend_batches_embedding_and_indexing() -> None: +def test_warm_up_backend_batches_embedding_and_indexing() -> None: """Warm-up should batch missing embeddings and vector indexing.""" warmup_embedding = np.ones(384, dtype=np.float32) batch_embeddings = [ @@ -49,7 +48,7 @@ async def test_warm_up_backend_batches_embedding_and_indexing() -> None: get_user_memories=AsyncMock(return_value=memories), ) - await _warm_up_backend(backend, "alice") + asyncio.run(mcp_server_mod._warm_up_backend(backend, "alice")) backend._ensure_initialized.assert_awaited_once() backend.get_user_memories.assert_awaited_once_with("alice", limit=500) @@ -61,3 +60,106 @@ async def test_warm_up_backend_batches_embedding_and_indexing() -> None: vector_index.index_batch.assert_awaited_once_with(memories) assert np.array_equal(memory_without_embedding_a.embedding, batch_embeddings[0]) assert np.array_equal(memory_without_embedding_b.embedding, batch_embeddings[1]) + + +def test_memory_mcp_startup_context_reports_dynamic_project_db(tmp_path) -> None: + project_dir = tmp_path / "project-a" + project_dir.mkdir() + configured_db = str(project_dir / ".headroom" / "memory.db") + + context = mcp_server_mod._memory_mcp_startup_context( + configured_db, + project_dir, + db_flag_present=False, + ) + + assert context == { + "configured_db": configured_db, + "resolved_db": configured_db, + "config_source": "cwd-default", + "cwd": str(project_dir), + "project_root": str(project_dir), + "storage_scope": "active-project", + "path_exists": False, + "path_readable": False, + "resolution": "dynamic-cwd", + } + + +def test_memory_mcp_startup_context_reports_static_external_db(tmp_path) -> None: + project_dir = tmp_path / "project-a" + project_dir.mkdir() + external_db = tmp_path / "shared-memory" / "memory.db" + external_db.parent.mkdir() + external_db.write_text("sqlite placeholder") + + context = mcp_server_mod._memory_mcp_startup_context( + str(external_db), + project_dir, + db_flag_present=True, + ) + + assert context == { + "configured_db": str(external_db), + "resolved_db": str(external_db.resolve(strict=False)), + "config_source": "cli-flag", + "cwd": str(project_dir), + "project_root": str(project_dir), + "storage_scope": "external-memory-db", + "path_exists": True, + "path_readable": True, + "resolution": "static-cli", + } + + +def test_memory_mcp_startup_context_reports_custom_db_path(tmp_path) -> None: + project_dir = tmp_path / "project-a" + project_dir.mkdir() + custom_db = tmp_path / "queries.db" + + context = mcp_server_mod._memory_mcp_startup_context( + str(custom_db), + project_dir, + db_flag_present=True, + ) + + assert context["storage_scope"] == "custom-db-path" + assert context["config_source"] == "cli-flag" + assert context["path_exists"] is False + assert context["path_readable"] is False + + +def test_main_logs_memory_mcp_startup_context(monkeypatch, tmp_path, caplog) -> None: + project_dir = tmp_path / "project-a" + project_dir.mkdir() + monkeypatch.chdir(project_dir) + monkeypatch.setenv("USER", "codex-user") + monkeypatch.setattr(mcp_server_mod.logging, "basicConfig", lambda **kwargs: None) + monkeypatch.setattr(mcp_server_mod.sys, "argv", ["memory-mcp"]) + + captured_run_payloads: list[object] = [] + monkeypatch.setattr( + mcp_server_mod, + "_run", + lambda db_path, user_id: ("run", db_path, user_id), + ) + monkeypatch.setattr( + mcp_server_mod.asyncio, + "run", + lambda payload: captured_run_payloads.append(payload), + ) + + caplog.set_level("INFO", logger="headroom.memory.mcp") + + mcp_server_mod.main() + + assert captured_run_payloads == [ + ("run", str(project_dir / ".headroom" / "memory.db"), "codex-user") + ] + assert any( + "Memory MCP startup: configured_db=" in record.message + and "config_source=cwd-default" in record.message + and "storage_scope=active-project" in record.message + and "resolution=dynamic-cwd" in record.message + for record in caplog.records + )