akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(codex): stop pinning Codex memory MCP to one project db (#1269)

Browse source 
## Description

Stop `headroom wrap codex --memory` from pinning the global
`headroom_memory` MCP server to one absolute SQLite path. Today the
wrapper writes `--db <wrap-cwd>/.headroom/memory.db` into
`~/.codex/config.toml`, which makes later Codex sessions either reopen a
stale project-local DB or fail with `unable to open database file` when
that original path disappears. This change lets the MCP server use its
existing per-cwd default again, so each Codex session resolves
`.headroom/memory.db` from the active project instead of a serialized
past cwd. Closes #1147

The current Codex-memory config surface was shaped by
https://github.com/chopratejas/headroom/issues/462 and
https://github.com/chopratejas/headroom/issues/730; this PR keeps that
surface project-scoped again instead of globally pinning one DB.

## Type of Change

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

## Changes Made

- remove the injected `--db` argument from the global `headroom_memory`
Codex MCP block while keeping `--user` intact
- preserve the wrap-time local `.headroom/memory.db` setup and
Claude-memory import path for the current project
- treat only wrap-owned Codex markers as snapshot-suppression and
unwrap-cleanup signals, so pre-existing named MCP blocks still back up
and restore
- log a startup diagnostic from `headroom.memory.mcp_server` that
records the configured DB path, config source, cwd/project root,
resolved storage scope, path existence/readability, and whether the path
was static or cwd-derived
- add a shared MCP SDK test stub so both the memory MCP and CCR MCP test
surfaces still run in CI when `mcp` is absent
- make the shared MCP stub re-import target modules under the stubbed
dependency set and restore any pre-existing target module object plus
dotted parent-package attribute state after cleanup
- add focused regressions and guard coverage for the persisted Codex
config shape, named-MCP marker backup and restore, the no-backup
memory-only unwrap path, the wrap-memory-then-unwrap cleanup path, the
failed-wrap memory-only cleanup path, the startup-diagnostic path
classification, the shared-store CCR retrieval path, and the shared MCP
stub import lifecycle
- add a `CHANGELOG.md` entry for the user-visible Codex memory scoping
fix

## Testing

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

### Test Output

```text
uv run pytest tests/test_ccr_mcp_server.py tests/test_memory/test_mcp_server.py tests/test_cli/test_wrap_codex.py tests/test_mcp_stub.py
======================== 78 passed, 1 warning in 5.96s ========================
Pytest warning:
PytestConfigWarning: Unknown config option: asyncio_mode
Pytest post-success atexit noise:
PermissionError: [WinError 5] Access is denied: 'C:\Users\Rod\AppData\Local\Temp\pytest-of-Rod\pytest-current'

uv run ruff check headroom/cli/wrap.py headroom/memory/mcp_server.py tests/_mcp_stub.py tests/test_ccr_mcp_server.py tests/test_cli/test_wrap_codex.py tests/test_mcp_stub.py tests/test_memory/test_mcp_server.py
All checks passed!

uv run ruff format headroom/cli/wrap.py headroom/memory/mcp_server.py tests/_mcp_stub.py tests/test_ccr_mcp_server.py tests/test_cli/test_wrap_codex.py tests/test_mcp_stub.py tests/test_memory/test_mcp_server.py --check
7 files already formatted
```

## Real Behavior Proof

- Environment: isolated temp project directories, a temp Codex home, the
real `wrap codex` and `unwrap codex` CLI commands under pytest, a mocked
missing-`codex` launch path for the failed-wrap cleanup case, and shared
MCP-SDK stubs for the memory MCP and CCR MCP test modules so CI still
exercises those paths without a real `mcp` install.
- Exact command / steps: run `uv run pytest tests/test_ccr_mcp_server.py
tests/test_memory/test_mcp_server.py tests/test_cli/test_wrap_codex.py
tests/test_mcp_stub.py`; prove the persisted config shape with
`TestCodexMemoryMcpConfig::test_inject_omits_db_and_replaces_existing_memory_block`;
prove prepare-only wrap cleanup with
`test_wrap_codex_memory_prepare_only_unwrap_removes_memory_mcp_without_prior_config`;
prove failed-wrap cleanup with
`test_wrap_codex_memory_launch_failure_unwrap_cleans_memory_only_config`;
guard pre-existing named Codex MCP preservation with
`test_memory_only_wrap_restores_preexisting_named_mcp_block` and
`test_memory_only_wrap_without_backup_preserves_named_mcp_block`; prove
the startup diagnostic classifications with
`test_memory_mcp_startup_context_reports_dynamic_project_db` and
`test_memory_mcp_startup_context_reports_static_external_db`; prove the
shared-store CCR retrieval path with
`test_mcp_uses_shared_singleton_store` and
`test_mcp_retrieves_proxy_stored_content`; prove stub import cleanup
with `test_import_module_with_mcp_stub_imports_target_and_cleans_up`,
`test_import_module_with_mcp_stub_reimports_target_and_restores_originals`,
and
`test_import_module_with_mcp_stub_cleans_up_dotted_target_attribute`.
- Observed result: the persisted global `headroom_memory` block now
keeps `--user` but omits `--db`; prepare-only memory setup still
bootstraps the current project's `.headroom/memory.db`; `headroom unwrap
codex --no-stop-proxy` now removes both the prepare-only generated
config and the failed-wrap memory-only config instead of leaving
`[mcp_servers.headroom_memory]` behind; pre-existing named Codex MCP
blocks remain restorable across both normal and no-backup memory-only
unwrap paths because only wrap-owned markers suppress backups or trigger
named-block cleanup; the memory MCP server now logs whether its DB path
came from the cwd default or an explicit static path, along with the
resolved path and scope it will open; CI can exercise both MCP test
modules even when the `mcp` package is absent from the shard
environment, and the shared stub now re-imports target modules under the
stubbed SDK while restoring both dependency and dotted parent-package
target-module import state after cleanup.
- Not tested: full end-to-end interactive Codex CLI launch.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

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

## Additional Notes

The code change stays narrowly scoped to Codex memory config
persistence, cleanup, and startup observability. It does not widen into
larger memory-routing redesign or startup-failure recovery logic.
This commit is contained in:
Rod Boev 2026-06-23 08:49:07 -04:00 • committed by GitHub
parent 3d59df7be8
commit ad7993bf15
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 785 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

68
.github/actions/headroom-e2e-setup/action.yml vendored
View file

@ -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 <target>`` 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 <target>`` 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

1
.github/workflows/init-native-e2e.yml vendored
View file

@ -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 }}

1
.github/workflows/install-native-e2e.yml vendored
View file

@ -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

1
.github/workflows/wrap-native-e2e.yml vendored
View file

@ -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

1
CHANGELOG.md
View file

@ -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)).

61
headroom/cli/wrap.py
View file

@ -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,6 +1112,7 @@ 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)
if remove_named_mcp:
content = re.sub(
r"(?ms)^# --- Headroom MCP server: [^\n]+ ---\n.*?"
r"^# --- end Headroom MCP server: [^\n]+ ---\n?",
@ -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")),
)

59
headroom/memory/mcp_server.py
View file

@ -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))

96
tests/_mcp_stub.py Normal file
View file

@ -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

8
tests/test_ccr_mcp_server.py
View file

@ -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, "<<small>>")
# 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", "")

301
tests/test_cli/test_wrap_codex.py
View file

@ -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:

102
tests/test_mcp_stub.py Normal file
View file

@ -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")

116
tests/test_memory/test_mcp_server.py
View file

@ -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
)