mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(codex): preserve wrapped sessions and recover state (#2160)
## Description Closes #2159. Codex wrappers currently launch against a disposable `CODEX_HOME`, so session state created during a wrapped run can disappear when that temporary directory is removed. This change launches Codex against its durable home, keeps proxy routing process-local, and adds recovery for retained temporary homes and pinned recovery sources. ## 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) - [x] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Launch Codex against its durable `CODEX_HOME` and apply routing through process-local config overrides after the actual proxy port is resolved. - Preserve custom provider identity and reject providers that cannot be redirected safely. - Detect dangling temporary Codex homes before interactive wraps and offer recovery. - Add `headroom recover codex` with automatic discovery, repeatable `--source`, preview, confirmation, retained backups, and rollback on failure. - Search Python's temp root, `$TMPDIR`, `/tmp`, `/private/tmp`, and macOS `/private/var/folders/*/*/T` for retained `headroom-codex-home-*` directories. - Reuse `source-pinned/` copies left by interrupted or failed recovery attempts after the original temporary home has disappeared. - Report deleted temporary homes still referenced by SQLite rollout paths without treating paths pasted into prompts or errors as filesystem evidence. - Audit the durable thread index, rollout files, and history when no source remains, including indexed chat counts and history-only orphan records. - Normalize legacy localhost `headroom` providers in both SQLite thread rows and rollout `session_meta`, including retries after an earlier broken recovery, while preserving user-defined remote providers named `headroom`. - Merge compatible config, JSONL, rollout, SQLite, credential, and regular-file state without propagating deletions or runtime artifacts. - Rewrite recovered thread rollout paths to the durable home and restore legacy Headroom thread providers to the active provider. - Validate SQLite schemas, SQLx migration checksums, integrity, and foreign keys, and quarantine malformed JSONL. - Preserve failed targets with an atomic rename before rollback, avoiding recursive-deletion races with live SQLite runtime files. - Document discovery, migration, retained backups, rollback behavior, and the limits of deleted-source recovery. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy`) - [x] New tests added for new functionality - [x] Manual testing performed in isolated Docker containers ### Test Output ```text $ uv run pytest tests/test_cli/test_wrap_codex.py tests/test_cli/test_recover_codex.py -q 122 passed $ uv run ruff check headroom/cli/recover.py headroom/providers/codex/recovery.py tests/test_cli/test_recover_codex.py All checks passed! $ uv run ruff format --check headroom/cli/recover.py headroom/providers/codex/recovery.py tests/test_cli/test_recover_codex.py 3 files already formatted $ uv run mypy headroom/cli/recover.py headroom/providers/codex/recovery.py Success: no issues found in 2 source files ``` All validation ran in `ghcr.io/astral-sh/uv:python3.12-bookworm` against a writable disposable copy of a read-only source mount. Codex was not installed or launched, and no real user Codex state was read or modified. The tests cover multi-root discovery, deleted-reference reporting, retained pinned-source recovery, durable SQLite path relocation, SQLite and rollout provider normalization, idempotent repair after an earlier broken recovery, remote provider preservation, unrelated dangling target rows, backup retention, atomic rollback, malformed-state quarantine, SQLite validation, and Windows-safe handle closure. The repository shim E2E was not launched locally because this recovery work intentionally avoids launching Codex. Upstream CI exercises wrapper E2E in isolated environments. ## Real Behavior Proof - Environment: `ghcr.io/astral-sh/uv:python3.12-bookworm`, Python 3.12, a writable disposable checkout copied from a read-only source mount, at head `2d89ecec`. - Exact command / steps: Run `pytest -q tests/test_cli/test_wrap_codex.py tests/test_cli/test_recover_codex.py`, then run `ruff check` and `ruff format --check` against `headroom/cli/wrap.py`, `headroom/cli/recover.py`, `headroom/providers/codex/recovery.py`, `tests/test_cli/test_wrap_codex.py`, and `tests/test_cli/test_recover_codex.py`. - Observed result: `122 passed in 10.08s`; Ruff reported `All checks passed!` and `5 files already formatted`. - Not tested: Launching a real Codex process or modifying a real user `CODEX_HOME`; these were intentionally excluded to protect live user state. ## 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 - [x] I have commented my code where the behavior is hard to understand - [x] I have made corresponding documentation changes - [x] My changes generate no new warnings - [x] I have added tests that prove the fix is effective - [x] New and existing focused unit tests pass with my changes - [x] I have updated `CHANGELOG.md` if applicable ## Additional Notes The temporary-home behavior was introduced by #1507 in `ad9d086f43`. Related context: #730, #731, #961, #1034, #1050, #1349, #1853, #1889, #2103, and #2104. A temporary home that macOS or `TemporaryDirectory` already deleted cannot be reconstructed unless a retained `source-pinned/` copy exists. Recovery identifies genuine dangling SQLite paths, audits surviving durable history, and recovers any retained pinned source it can find. Prompt text without a rollout cannot reconstruct a full transcript. The unchecked changelog item is not applicable because this repository does not require a changelog entry for this fix. --------- Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
Rudimar Ronsoni
GitHub
parent
57e8dcb425
commit
dec60de976
13 changed files with 2195 additions and 250 deletions
79
docs/content/docs/codex-recovery.mdx
Normal file
79
docs/content/docs/codex-recovery.mdx
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
---
|
||||
title: Recover Codex State
|
||||
description: Merge sessions and configuration left in a temporary Headroom Codex home back into the durable Codex home.
|
||||
---
|
||||
|
||||
Older versions of `headroom wrap codex` could run Codex with a temporary `CODEX_HOME` named `headroom-codex-home-*`. Chats and configuration created during that wrapped session stayed in the temporary home, so they disappeared from the normal Codex history after the wrapper exited. This recovery migrates the retained state back into the durable Codex home without deleting either source.
|
||||
|
||||
This procedure can recover temporary homes that still exist and pinned sources retained by an interrupted recovery. A temporary home that was already deleted cannot be reconstructed unless one of those retained copies exists. Headroom reports deleted temporary homes still referenced by the Codex thread database. It does not treat paths pasted into prompts or error messages as recovery sources. See [issue #2159](https://github.com/headroomlabs-ai/headroom/issues/2159) for the regression details.
|
||||
|
||||
## Before you recover
|
||||
|
||||
Close Codex and any process using the temporary or durable Codex homes. Recovery detects changes made while creating each backup and aborts, but quiet homes are required for a consistent migration.
|
||||
|
||||
The target defaults to `$CODEX_HOME` when it is set, otherwise `~/.codex`. Check that this is the durable home you normally use before confirming the migration.
|
||||
|
||||
## Automatic recovery
|
||||
|
||||
The first interactive `headroom wrap codex` with the fixed version searches Python's temporary directory, `$TMPDIR`, `/tmp`, `/private/tmp`, and the macOS `/private/var/folders/*/*/T` roots for non-empty `headroom-codex-home-*` directories. If it finds any, it lists them and offers to back up and recover them before Codex starts.
|
||||
|
||||
Declining the prompt changes nothing. You can run the manual command later.
|
||||
|
||||
## Manual recovery
|
||||
|
||||
Let Headroom discover retained temporary homes and `source-pinned/` copies left by interrupted or failed recovery attempts, then preview the migration:
|
||||
|
||||
```bash
|
||||
headroom recover codex
|
||||
```
|
||||
|
||||
If no recoverable copy remains, the command audits the durable Codex thread database, rollout files, and `history.jsonl`. It reports indexed active and archived chats, surviving rollouts missing from the thread index, and history-only records whose rollouts no longer exist. History-only prompt text cannot reconstruct a full chat transcript.
|
||||
|
||||
Codex filters the default resume picker by the current working directory. Run `codex resume --all` yourself to display indexed chats from every working directory. Headroom does not launch Codex during recovery.
|
||||
|
||||
To select a known temporary home and an explicit durable target:
|
||||
|
||||
```bash
|
||||
headroom recover codex \
|
||||
--source /path/to/headroom-codex-home-12345 \
|
||||
--target "${CODEX_HOME:-$HOME/.codex}"
|
||||
```
|
||||
|
||||
Repeat `--source` to merge multiple homes. Headroom shows the target and every source before asking for confirmation. Use `--yes` only in automation where those paths have already been reviewed.
|
||||
|
||||
## What gets migrated
|
||||
|
||||
When a source used the localhost `headroom` model provider injected by the old wrapper, recovery rewrites that provider in both SQLite thread rows and rollout `session_meta` records to the active target provider. This also repairs newer target records left by an earlier broken recovery. A user-defined remote provider named `headroom` is preserved.
|
||||
|
||||
- `history.jsonl` and other JSONL indexes are combined without duplicate records. Malformed input is excluded from the result and copied to the backup quarantine.
|
||||
- Session rollouts under `sessions/` and `archived_sessions/` keep the newer file when the same path exists in both homes.
|
||||
- SQLite databases are merged only when their tables, indexes, triggers, views, and migration checksums are compatible. Primary-key conflicts keep rows from the newer database. Recovered thread rows are rewritten to the durable rollout path, including rows restored from a pinned source after the original temporary home was deleted.
|
||||
- `config.toml` tables are merged recursively. Values from the newer config win, and localhost Headroom routing injected by the old wrapper is removed. A user-defined remote provider is preserved.
|
||||
- Other files, including credentials and user settings, keep the newer copy. The durable target wins equal modification-time ties, and missing source files never delete target files.
|
||||
|
||||
Sockets, lock files, SQLite journals, FIFOs, and other runtime-only artifacts are recorded but not copied.
|
||||
|
||||
## Backups and rollback
|
||||
|
||||
Every source is pinned before migration. The existing durable home, including its current Codex configuration, is backed up before any merge begins. Backups are owner-only and retained next to the target:
|
||||
|
||||
```text
|
||||
<target-parent>/.headroom-codex-recovery/<timestamp-pid>/
|
||||
├── source-pinned/
|
||||
├── target-before/
|
||||
├── target-failed/ # only when a merge is rolled back
|
||||
├── manifest.json
|
||||
└── quarantine/ # only when malformed input is found
|
||||
```
|
||||
|
||||
`target-before/` is absent when the target did not exist before recovery. `manifest.json` records copied, merged, quarantined, and skipped paths.
|
||||
|
||||
If configuration parsing, SQLite schema validation, integrity checks, foreign-key checks, or filesystem writes fail, Headroom atomically renames the failed target to `target-failed/` before restoring `target-before/`. This avoids recursive-deletion races with SQLite runtime files and preserves the failed merge for inspection. When the target did not exist before recovery, the partial target is retained only as `target-failed/`. The pinned source and recovery backup remain available for inspection.
|
||||
|
||||
Recovery also refuses overlapping source and target paths, target symlink traversal, and a source that changes while it is being pinned.
|
||||
|
||||
## After recovery
|
||||
|
||||
The command prints the retained backup path after each successful source merge. Review its `manifest.json`, then start Codex normally and confirm that the recovered chats and settings appear. Keep the backup until you have verified the durable history.
|
||||
|
||||
Fixed versions of `headroom wrap codex` run Codex against the durable home and apply proxy routing only to the launched process, so new wrapped sessions remain visible after Headroom exits.
|
||||
|
|
@ -19,6 +19,8 @@ headroom wrap codex --prepare-only
|
|||
|
||||
Supported values are `rtk` and `lean-ctx`; unset defaults to `rtk`.
|
||||
|
||||
If Codex history disappeared after using an older wrapper, see [Recover Codex State](/docs/codex-recovery) before wrapping Codex again.
|
||||
|
||||
The proxy reads RTK lifetime savings with global scope by default so a shared
|
||||
daemon reports savings across the operator's projects. Set
|
||||
`HEADROOM_RTK_GAIN_SCOPE=project` to query `rtk gain --project` from the
|
||||
|
|
|
|||
|
|
@ -40,6 +40,7 @@
|
|||
"mcp",
|
||||
"---Configuration---",
|
||||
"configuration",
|
||||
"codex-recovery",
|
||||
"pipeline-extensions",
|
||||
"filesystem-contract",
|
||||
"---Observability---",
|
||||
|
|
|
|||
|
|
@ -581,8 +581,19 @@ def verify_codex_wrap(
|
|||
|
||||
config_path = Path(base_env["HOME"]) / ".codex" / "config.toml"
|
||||
assert_true(
|
||||
not config_path.exists(),
|
||||
"Codex wrap should leave ~/.codex/config.toml untouched during normal launch",
|
||||
config_path.exists(),
|
||||
"Codex wrap should persist its MCP setup in the durable Codex home",
|
||||
)
|
||||
durable_config = config_path.read_text(encoding="utf-8")
|
||||
assert_true(
|
||||
"[mcp_servers.headroom]" in durable_config,
|
||||
"Codex wrap should persist the Headroom MCP server",
|
||||
)
|
||||
assert_true(
|
||||
'model_provider = "headroom"' not in durable_config
|
||||
and "[model_providers.headroom]" not in durable_config
|
||||
and "openai_base_url" not in durable_config,
|
||||
"Codex wrap should keep proxy routing out of the durable config",
|
||||
)
|
||||
|
||||
entries = read_jsonl(log_dir / "codex.jsonl")
|
||||
|
|
@ -591,49 +602,41 @@ def verify_codex_wrap(
|
|||
session_home = env_vars.get("CODEX_HOME")
|
||||
assert_true(
|
||||
isinstance(session_home, str) and session_home,
|
||||
"Codex wrap should launch the child with a session-scoped CODEX_HOME",
|
||||
"Codex wrap should launch the child with CODEX_HOME",
|
||||
)
|
||||
assert_true(
|
||||
Path(session_home) != config_path.parent,
|
||||
"Codex wrap should not point the child at the real ~/.codex home",
|
||||
Path(session_home) == config_path.parent,
|
||||
"Codex wrap should point the child at the durable ~/.codex home",
|
||||
)
|
||||
config = entries[-1].get("session_config")
|
||||
assert_true(
|
||||
isinstance(config, str) and config,
|
||||
"Codex wrap should capture the session-scoped config during launch",
|
||||
"Codex wrap should capture the durable config during launch",
|
||||
)
|
||||
project_prefix = f"/p/{quote(project_dir.name, safe='')}"
|
||||
expected_base_url = f"http://127.0.0.1:{port}{project_prefix}/v1"
|
||||
argv = entries[-1].get("argv")
|
||||
assert_true(
|
||||
f'openai_base_url = "http://127.0.0.1:{port}/v1"' in config,
|
||||
"Codex wrap should inject openai_base_url into the session config",
|
||||
)
|
||||
assert_true(
|
||||
f'base_url = "http://127.0.0.1:{port}/v1"' in config,
|
||||
"Codex wrap should inject the headroom provider base_url into the session config",
|
||||
isinstance(argv, list) and f'openai_base_url="{expected_base_url}"' in argv,
|
||||
"Codex wrap should pass openai_base_url as a process-local override",
|
||||
)
|
||||
assert_true(
|
||||
'env_key = "OPENAI_API_KEY"' not in config,
|
||||
"Codex wrap should preserve OAuth and never inject env_key into the session config",
|
||||
)
|
||||
# Bug 3 (#406): requires_openai_auth must be absent from headroom provider blocks.
|
||||
assert_true(
|
||||
"requires_openai_auth" not in config,
|
||||
"Codex wrap must NOT inject requires_openai_auth into the session headroom provider block",
|
||||
)
|
||||
assert_true(
|
||||
"supports_websockets = true" in config,
|
||||
"Codex wrap missing 'supports_websockets = true' in the session config",
|
||||
)
|
||||
|
||||
assert_true(
|
||||
env_vars.get("OPENAI_BASE_URL") == f"http://127.0.0.1:{port}/v1",
|
||||
env_vars.get("OPENAI_BASE_URL") == expected_base_url,
|
||||
"Codex wrap should set OPENAI_BASE_URL",
|
||||
)
|
||||
assert_true(
|
||||
entries[-1]["probes"]
|
||||
== [
|
||||
{"url": f"http://127.0.0.1:{port}/v1/models", "status": 200},
|
||||
{"url": f"http://127.0.0.1:{port}/v1/chat/completions", "status": 200},
|
||||
{"url": f"http://127.0.0.1:{port}/stats", "status": 200},
|
||||
{"url": f"{expected_base_url}/models", "status": 200},
|
||||
{"url": f"{expected_base_url}/chat/completions", "status": 200},
|
||||
{
|
||||
"url": f"http://127.0.0.1:{port}{project_prefix}/stats",
|
||||
"status": 200,
|
||||
},
|
||||
],
|
||||
"Codex shim should prove OPENAI_BASE_URL points at a live proxy and that Headroom logged the wrapped message",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from . import ( # noqa: F401
|
|||
mcp,
|
||||
perf,
|
||||
proxy,
|
||||
recover,
|
||||
tools,
|
||||
update,
|
||||
wrap,
|
||||
|
|
|
|||
|
|
@ -74,6 +74,7 @@ def _register_commands() -> None:
|
|||
output_savings, # noqa: F401
|
||||
perf, # noqa: F401
|
||||
proxy, # noqa: F401
|
||||
recover, # noqa: F401
|
||||
savings, # noqa: F401
|
||||
tools, # noqa: F401
|
||||
update, # noqa: F401
|
||||
|
|
|
|||
94
headroom/cli/recover.py
Normal file
94
headroom/cli/recover.py
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
"""Recovery commands for state left behind by interrupted wrappers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import click
|
||||
|
||||
from headroom.cli.main import main
|
||||
from headroom.providers.codex.recovery import (
|
||||
audit_codex_history,
|
||||
discover_dangling_homes,
|
||||
discover_referenced_temp_homes,
|
||||
discover_retained_sources,
|
||||
recover_codex_home,
|
||||
)
|
||||
|
||||
|
||||
@main.group()
|
||||
def recover() -> None:
|
||||
"""Recover agent state left in a temporary Headroom home."""
|
||||
|
||||
|
||||
@recover.command("codex")
|
||||
@click.option(
|
||||
"sources",
|
||||
"--source",
|
||||
type=click.Path(path_type=Path, exists=True, file_okay=False, resolve_path=True),
|
||||
multiple=True,
|
||||
help="Temporary Codex home to merge. Repeat to merge more than one.",
|
||||
)
|
||||
@click.option(
|
||||
"--target",
|
||||
type=click.Path(path_type=Path, file_okay=False, resolve_path=True),
|
||||
help="Active Codex home. Defaults to CODEX_HOME or ~/.codex.",
|
||||
)
|
||||
@click.option("--yes", is_flag=True, help="Apply the recovery without prompting.")
|
||||
def recover_codex(sources: tuple[Path, ...], target: Path | None, yes: bool) -> None:
|
||||
"""Merge sessions and configuration from dangling Codex homes."""
|
||||
target = target or Path(os.environ.get("CODEX_HOME", Path.home() / ".codex"))
|
||||
selected_sources = list(sources) or [
|
||||
*discover_dangling_homes(),
|
||||
*discover_retained_sources(target),
|
||||
]
|
||||
if not selected_sources:
|
||||
deleted_sources = [
|
||||
source for source in discover_referenced_temp_homes(target) if not source.exists()
|
||||
]
|
||||
if deleted_sources:
|
||||
click.echo("Referenced temporary Codex homes were already deleted:")
|
||||
for source in deleted_sources:
|
||||
click.echo(f" {source}")
|
||||
audit = audit_codex_history(target)
|
||||
if audit is not None:
|
||||
click.echo(
|
||||
f"Durable Codex history: {audit.indexed} indexed chats "
|
||||
f"({audit.active} active, {audit.archived} archived)."
|
||||
)
|
||||
if audit.unindexed_rollouts:
|
||||
click.echo(
|
||||
"Surviving rollout files missing from the thread database: "
|
||||
f"{len(audit.unindexed_rollouts)}"
|
||||
)
|
||||
for session_id in audit.unindexed_rollouts:
|
||||
click.echo(f" {session_id}")
|
||||
if audit.history_without_rollout:
|
||||
click.echo(
|
||||
"History-only records without a surviving rollout: "
|
||||
f"{len(audit.history_without_rollout)}"
|
||||
)
|
||||
for session_id in audit.history_without_rollout:
|
||||
click.echo(f" {session_id}")
|
||||
click.echo("Their full transcripts cannot be restored without a retained rollout.")
|
||||
if audit.indexed:
|
||||
click.echo("Run `codex resume --all` to show chats from every working directory.")
|
||||
click.echo("No recoverable Headroom Codex homes were found.")
|
||||
return
|
||||
|
||||
click.echo(f"Target Codex home: {target}")
|
||||
click.echo("Sources:")
|
||||
for source in selected_sources:
|
||||
click.echo(f" {source}")
|
||||
click.echo("Both the current target and each source will be backed up before merging.")
|
||||
if not yes and not click.confirm("Recover these Codex homes?", default=True):
|
||||
click.echo("Recovery cancelled. No Codex state was changed.")
|
||||
return
|
||||
|
||||
for source in selected_sources:
|
||||
try:
|
||||
report = recover_codex_home(source=source, target=target)
|
||||
except (OSError, RuntimeError, ValueError) as exc:
|
||||
raise click.ClickException(f"Codex recovery failed: {exc}") from exc
|
||||
click.echo(f"Recovery complete. Backup retained at {report.backup_dir}")
|
||||
|
|
@ -32,7 +32,6 @@ import tempfile
|
|||
import time
|
||||
import urllib.parse
|
||||
from collections.abc import Callable
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
|
||||
|
|
@ -46,6 +45,11 @@ if sys.platform == "win32" and hasattr(sys.stdout, "buffer"):
|
|||
|
||||
import click
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
import tomllib
|
||||
else: # pragma: no cover - exercised only on Python 3.10
|
||||
import tomli as tomllib # type: ignore[no-redef]
|
||||
|
||||
from headroom import fsutil
|
||||
from headroom._version import __version__ as _HEADROOM_VERSION
|
||||
from headroom._version import normalize_release_version as _normalize_release_version
|
||||
|
|
@ -1789,32 +1793,135 @@ def _codex_home_dir() -> Path:
|
|||
return Path.home() / ".codex"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _codex_session_home_overlay() -> Any:
|
||||
"""Seed a temporary Codex home from the active home and point the process at it."""
|
||||
source_home = _codex_home_dir()
|
||||
original_codex_home = os.environ.get("CODEX_HOME")
|
||||
def _codex_profile_from_args(codex_args: tuple[str, ...]) -> str | None:
|
||||
"""Return the profile selected by Codex CLI arguments, if any."""
|
||||
for index, argument in enumerate(codex_args):
|
||||
if argument.startswith("--profile="):
|
||||
return argument.partition("=")[2]
|
||||
if argument in {"--profile", "-p"} and index + 1 < len(codex_args):
|
||||
return codex_args[index + 1]
|
||||
return None
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="headroom-codex-home-") as tmp_dir:
|
||||
session_home = Path(tmp_dir)
|
||||
if source_home.exists():
|
||||
shutil.copytree(
|
||||
source_home,
|
||||
session_home,
|
||||
dirs_exist_ok=True,
|
||||
ignore=lambda directory, names: [
|
||||
name for name in names if (Path(directory) / name).is_socket()
|
||||
],
|
||||
|
||||
def _codex_model_provider_from_args(codex_args: tuple[str, ...]) -> str | None:
|
||||
"""Return the last top-level model_provider CLI override, if present."""
|
||||
provider: str | None = None
|
||||
for index, argument in enumerate(codex_args):
|
||||
override: str | None = None
|
||||
if argument.startswith("--config=") or argument.startswith("-c="):
|
||||
override = argument.partition("=")[2]
|
||||
elif argument in {"--config", "-c"} and index + 1 < len(codex_args):
|
||||
override = codex_args[index + 1]
|
||||
if override is None:
|
||||
continue
|
||||
key, separator, value = override.partition("=")
|
||||
if separator and key.strip() == "model_provider":
|
||||
provider = value.strip().strip("\"'")
|
||||
return provider
|
||||
|
||||
|
||||
def _codex_toml_value(value: Any) -> str:
|
||||
"""Serialize values accepted by Codex's TOML command-line overrides."""
|
||||
if isinstance(value, str):
|
||||
return json.dumps(value)
|
||||
if isinstance(value, bool):
|
||||
return "true" if value else "false"
|
||||
if isinstance(value, (int, float)):
|
||||
return str(value)
|
||||
raise TypeError(f"unsupported Codex config override: {type(value).__name__}")
|
||||
|
||||
|
||||
def _codex_dotted_key(*parts: str) -> str:
|
||||
return ".".join(json.dumps(part) for part in parts)
|
||||
|
||||
|
||||
def _codex_session_launch_settings(
|
||||
*, port: int, codex_args: tuple[str, ...], environ: dict[str, str]
|
||||
) -> tuple[tuple[str, ...], dict[str, str], list[str]]:
|
||||
"""Build process-local routing while preserving the selected provider id."""
|
||||
config_file, _ = _codex_config_paths()
|
||||
try:
|
||||
config = tomllib.loads(_read_text(config_file)) if config_file.exists() else {}
|
||||
except (OSError, tomllib.TOMLDecodeError) as exc:
|
||||
raise click.ClickException(
|
||||
f"could not read Codex config for session routing: {exc}"
|
||||
) from exc
|
||||
|
||||
profile_name = _codex_profile_from_args(codex_args)
|
||||
profiles = config.get("profiles", {})
|
||||
profile = profiles.get(profile_name, {}) if profile_name and isinstance(profiles, dict) else {}
|
||||
provider = _codex_model_provider_from_args(codex_args) or (
|
||||
profile.get("model_provider")
|
||||
if isinstance(profile, dict) and profile.get("model_provider")
|
||||
else config.get("model_provider", "openai")
|
||||
)
|
||||
provider = str(provider)
|
||||
|
||||
project = _project_name_from_cwd()
|
||||
proxy_url = _with_project_prefix(f"http://127.0.0.1:{port}/v1", project)
|
||||
overrides: list[str] = []
|
||||
env = dict(environ)
|
||||
display = [f"OPENAI_BASE_URL={proxy_url}"]
|
||||
env["OPENAI_BASE_URL"] = proxy_url
|
||||
|
||||
if provider == "openai":
|
||||
overrides.append(f"openai_base_url={_codex_toml_value(proxy_url)}")
|
||||
else:
|
||||
providers = config.get("model_providers", {})
|
||||
provider_config = providers.get(provider) if isinstance(providers, dict) else None
|
||||
if not isinstance(provider_config, dict):
|
||||
raise click.ClickException(
|
||||
f"Codex provider {provider!r} cannot be redirected without changing its identity"
|
||||
)
|
||||
upstream = provider_config.get("base_url")
|
||||
if not isinstance(upstream, str) or not upstream.strip():
|
||||
raise click.ClickException(
|
||||
f"Codex custom provider {provider!r} has no upstream base_url"
|
||||
)
|
||||
prefix = ("model_providers", provider)
|
||||
overrides.extend(
|
||||
(
|
||||
f"{_codex_dotted_key(*prefix, 'base_url')}={_codex_toml_value(proxy_url)}",
|
||||
f"{_codex_dotted_key(*prefix, 'supports_websockets')}=true",
|
||||
)
|
||||
)
|
||||
env[_UPSTREAM_BASE_URL_ENV_VAR] = upstream.rstrip("/")
|
||||
display.append(f"{_UPSTREAM_BASE_URL_ENV_VAR}={upstream.rstrip('/')}")
|
||||
overrides.append(
|
||||
f"{_codex_dotted_key(*prefix, 'env_http_headers', _UPSTREAM_BASE_URL_HEADER_NAME)}="
|
||||
f"{_codex_toml_value(_UPSTREAM_BASE_URL_ENV_VAR)}"
|
||||
)
|
||||
|
||||
os.environ["CODEX_HOME"] = str(session_home)
|
||||
try:
|
||||
yield session_home
|
||||
finally:
|
||||
if original_codex_home is None:
|
||||
os.environ.pop("CODEX_HOME", None)
|
||||
else:
|
||||
os.environ["CODEX_HOME"] = original_codex_home
|
||||
if project and "HEADROOM_PROJECT" not in env:
|
||||
env["HEADROOM_PROJECT"] = project
|
||||
config_args = tuple(item for override in overrides for item in ("--config", override))
|
||||
return (*config_args, *codex_args), env, display
|
||||
|
||||
|
||||
def _offer_dangling_codex_recovery(active_home: Path) -> None:
|
||||
"""Offer recovery before an interactive wrap creates more Codex state."""
|
||||
if not sys.stdin.isatty():
|
||||
return
|
||||
from headroom.providers.codex.recovery import (
|
||||
discover_dangling_homes,
|
||||
recover_codex_home,
|
||||
)
|
||||
|
||||
candidates = [path for path in discover_dangling_homes() if path != active_home]
|
||||
if not candidates:
|
||||
return
|
||||
click.echo("\nFound Codex state left by an earlier Headroom temporary home:")
|
||||
for candidate in candidates:
|
||||
click.echo(f" {candidate}")
|
||||
if not click.confirm(
|
||||
"Back up both homes and recover this state before launching Codex?",
|
||||
default=True,
|
||||
):
|
||||
click.echo("Skipped recovery. Run `headroom recover codex` to recover it later.")
|
||||
return
|
||||
for candidate in candidates:
|
||||
report = recover_codex_home(source=candidate, target=active_home)
|
||||
click.echo(f"Recovered Codex state. Backup retained at {report.backup_dir}")
|
||||
|
||||
|
||||
def _codex_config_paths() -> tuple[Path, Path]:
|
||||
|
|
@ -3825,6 +3932,11 @@ def _launch_tool(
|
|||
copilot_api_token: str | None = None,
|
||||
copilot_refresh_oauth_token: str | None = None,
|
||||
copilot_api_token_expires_at: float | None = None,
|
||||
configure_launch: Callable[
|
||||
[int, tuple, dict[str, str], list[str]],
|
||||
tuple[tuple, dict[str, str], list[str]],
|
||||
]
|
||||
| None = None,
|
||||
) -> None:
|
||||
"""Common logic: start proxy, launch tool, clean up."""
|
||||
proxy_holder: list[subprocess.Popen | None] = [None]
|
||||
|
|
@ -3868,6 +3980,9 @@ def _launch_tool(
|
|||
for k, v in dict(env).items():
|
||||
env[k] = v.replace(f"127.0.0.1:{port}", f"127.0.0.1:{actual_port}")
|
||||
|
||||
if configure_launch is not None:
|
||||
args, env, env_vars_display = configure_launch(actual_port, args, env, env_vars_display)
|
||||
|
||||
if code_graph:
|
||||
_setup_code_graph(verbose=False)
|
||||
|
||||
|
|
@ -4996,23 +5111,18 @@ def _prepare_codex_wrap_state(
|
|||
memory: bool,
|
||||
verbose: bool,
|
||||
rtk_home: Path | None = None,
|
||||
) -> str | None:
|
||||
"""Prepare the active Codex home for a wrap or prepare-only invocation.
|
||||
|
||||
Returns the custom upstream base URL detected in the user's Codex config
|
||||
(or None). Callers that launch Codex must export this into
|
||||
``HEADROOM_CODEX_UPSTREAM_BASE_URL`` so the injected provider's
|
||||
``X-Headroom-Base-Url`` header carries it; otherwise the proxy falls back to
|
||||
``api.openai.com`` and the user's gateway key is sent to the wrong host.
|
||||
"""
|
||||
persistent_routing: bool = True,
|
||||
) -> None:
|
||||
"""Prepare the active Codex home for a wrap or prepare-only invocation."""
|
||||
# Snapshot Codex config.toml BEFORE any wrap-time mutation so
|
||||
# `headroom unwrap codex` can restore the user's pre-wrap state
|
||||
# byte-for-byte. The snapshot is a no-op if the backup already exists
|
||||
# or if the file already has Headroom markers, so this is safe to
|
||||
# call repeatedly. Crucially this must run before MCP install, which
|
||||
# writes its marker block to the same file.
|
||||
_codex_config_file, _codex_backup_file = _codex_config_paths()
|
||||
_snapshot_codex_config_if_unwrapped(_codex_config_file, _codex_backup_file)
|
||||
if persistent_routing:
|
||||
_codex_config_file, _codex_backup_file = _codex_config_paths()
|
||||
_snapshot_codex_config_if_unwrapped(_codex_config_file, _codex_backup_file)
|
||||
|
||||
# Setup CLI context tool for Codex.
|
||||
if not no_rtk:
|
||||
|
|
@ -5099,10 +5209,10 @@ def _prepare_codex_wrap_state(
|
|||
# transport unless a custom provider declares supports_websockets = true.
|
||||
# NOTE: this must run BEFORE _inject_memory_mcp_config because it rewrites
|
||||
# the config file. Re-inject MCP config after if memory is enabled.
|
||||
custom_upstream = _inject_codex_provider_config(port)
|
||||
if memory:
|
||||
_inject_memory_mcp_config(os.environ.get("USER", os.environ.get("USERNAME", "default")))
|
||||
return custom_upstream
|
||||
if persistent_routing:
|
||||
_inject_codex_provider_config(port)
|
||||
if memory:
|
||||
_inject_memory_mcp_config(os.environ.get("USER", os.environ.get("USERNAME", "default")))
|
||||
|
||||
|
||||
def _run_codex_wrap(
|
||||
|
|
@ -5124,7 +5234,7 @@ def _run_codex_wrap(
|
|||
prepare_only: bool,
|
||||
codex_args: tuple,
|
||||
) -> None:
|
||||
"""Execute the Codex wrap flow with the session overlay when launching."""
|
||||
"""Execute the Codex wrap flow against the durable Codex home."""
|
||||
if prepare_only:
|
||||
_prepare_codex_wrap_state(
|
||||
port=port,
|
||||
|
|
@ -5145,55 +5255,53 @@ def _run_codex_wrap(
|
|||
raise SystemExit(1)
|
||||
|
||||
active_codex_home = _codex_home_dir()
|
||||
with _codex_session_home_overlay() as session_codex_home:
|
||||
custom_upstream = _prepare_codex_wrap_state(
|
||||
port=port,
|
||||
no_rtk=no_rtk,
|
||||
no_mcp=no_mcp,
|
||||
no_tokensave=no_tokensave,
|
||||
serena=serena,
|
||||
no_serena=no_serena,
|
||||
memory=memory,
|
||||
verbose=verbose,
|
||||
rtk_home=active_codex_home,
|
||||
_offer_dangling_codex_recovery(active_codex_home)
|
||||
_prepare_codex_wrap_state(
|
||||
port=port,
|
||||
no_rtk=no_rtk,
|
||||
no_mcp=no_mcp,
|
||||
no_tokensave=no_tokensave,
|
||||
serena=serena,
|
||||
no_serena=no_serena,
|
||||
memory=memory,
|
||||
verbose=verbose,
|
||||
rtk_home=active_codex_home,
|
||||
persistent_routing=False,
|
||||
)
|
||||
|
||||
env, env_vars_display = _build_codex_launch_env(port, os.environ)
|
||||
env["CODEX_HOME"] = str(active_codex_home)
|
||||
|
||||
def configure_codex_launch(
|
||||
actual_port: int,
|
||||
current_args: tuple,
|
||||
current_env: dict[str, str],
|
||||
current_display: list[str],
|
||||
) -> tuple[tuple, dict[str, str], list[str]]:
|
||||
del current_display
|
||||
return _codex_session_launch_settings(
|
||||
port=actual_port,
|
||||
codex_args=current_args,
|
||||
environ=current_env,
|
||||
)
|
||||
|
||||
env, env_vars_display = _build_codex_launch_env(port, os.environ)
|
||||
|
||||
# Export the detected custom upstream so Codex actually emits the
|
||||
# X-Headroom-Base-Url header the injected provider declares. Without
|
||||
# this the proxy falls back to api.openai.com and a user with an
|
||||
# OpenAI-compatible gateway in ~/.codex/config.toml has their gateway
|
||||
# key sent to OpenAI (regression of #1614). A user-set value wins.
|
||||
if custom_upstream and _UPSTREAM_BASE_URL_ENV_VAR not in env:
|
||||
env[_UPSTREAM_BASE_URL_ENV_VAR] = custom_upstream
|
||||
env_vars_display.append(f"{_UPSTREAM_BASE_URL_ENV_VAR}={custom_upstream}")
|
||||
|
||||
# Per-project savings attribution: the injected provider config maps the
|
||||
# X-Headroom-Project header to HEADROOM_PROJECT via env_http_headers, so
|
||||
# Codex sends it only when this var is set. A user-set value wins.
|
||||
_codex_project = _project_name_from_cwd()
|
||||
if _codex_project and "HEADROOM_PROJECT" not in env:
|
||||
env["HEADROOM_PROJECT"] = _codex_project
|
||||
|
||||
env["CODEX_HOME"] = str(session_codex_home)
|
||||
|
||||
_launch_tool(
|
||||
binary=codex_bin,
|
||||
args=codex_args,
|
||||
env=env,
|
||||
port=port,
|
||||
no_proxy=no_proxy,
|
||||
tool_label="CODEX",
|
||||
env_vars_display=env_vars_display,
|
||||
learn=learn,
|
||||
memory=memory,
|
||||
agent_type="codex",
|
||||
code_graph=code_graph,
|
||||
backend=backend,
|
||||
anyllm_provider=anyllm_provider,
|
||||
region=region,
|
||||
)
|
||||
_launch_tool(
|
||||
binary=codex_bin,
|
||||
args=codex_args,
|
||||
env=env,
|
||||
port=port,
|
||||
no_proxy=no_proxy,
|
||||
tool_label="CODEX",
|
||||
env_vars_display=env_vars_display,
|
||||
learn=learn,
|
||||
memory=memory,
|
||||
agent_type="codex",
|
||||
code_graph=code_graph,
|
||||
backend=backend,
|
||||
anyllm_provider=anyllm_provider,
|
||||
region=region,
|
||||
configure_launch=configure_codex_launch,
|
||||
)
|
||||
|
||||
|
||||
@wrap.command(context_settings={"ignore_unknown_options": True})
|
||||
|
|
|
|||
748
headroom/providers/codex/recovery.py
Normal file
748
headroom/providers/codex/recovery.py
Normal file
|
|
@ -0,0 +1,748 @@
|
|||
"""Transactional recovery of Codex state left in a temporary Headroom home."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gc
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sqlite3
|
||||
import stat
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import tomlkit
|
||||
|
||||
_TEMP_HOME_PREFIX = "headroom-codex-home-"
|
||||
_RECOVERY_DIR = ".headroom-codex-recovery"
|
||||
_SQLITE_SUFFIXES = {".sqlite", ".db"}
|
||||
_RUNTIME_NAMES = {".DS_Store"}
|
||||
_RUNTIME_SUFFIXES = {".lock", ".sock", ".socket", "-shm", "-wal", "-journal"}
|
||||
|
||||
|
||||
def _latest_state_mtime_ns(home: Path) -> int:
|
||||
timestamps = [entry.lstat().st_mtime_ns for entry in home.rglob("*") if not entry.is_dir()]
|
||||
return max(timestamps, default=home.stat().st_mtime_ns)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecoveryReport:
|
||||
source: Path
|
||||
target: Path
|
||||
backup_dir: Path
|
||||
copied: list[str] = field(default_factory=list)
|
||||
merged: list[str] = field(default_factory=list)
|
||||
quarantined: list[str] = field(default_factory=list)
|
||||
skipped_runtime: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CodexHistoryAudit:
|
||||
indexed: int
|
||||
active: int
|
||||
archived: int
|
||||
unindexed_rollouts: tuple[str, ...]
|
||||
history_without_rollout: tuple[str, ...]
|
||||
|
||||
|
||||
def discover_dangling_homes(temp_root: Path | None = None) -> list[Path]:
|
||||
"""Find non-empty Headroom temporary Codex homes, newest first."""
|
||||
if temp_root is not None:
|
||||
roots = [temp_root]
|
||||
else:
|
||||
roots = [Path(tempfile.gettempdir())]
|
||||
if env_tmpdir := os.environ.get("TMPDIR"):
|
||||
roots.append(Path(env_tmpdir))
|
||||
roots.extend((Path("/tmp"), Path("/private/tmp")))
|
||||
roots.extend(Path("/private/var/folders").glob("*/*/T"))
|
||||
candidates: list[Path] = []
|
||||
seen: set[Path] = set()
|
||||
for root in roots:
|
||||
try:
|
||||
paths = root.glob(f"{_TEMP_HOME_PREFIX}*")
|
||||
for path in paths:
|
||||
resolved = path.resolve()
|
||||
if (
|
||||
resolved not in seen
|
||||
and path.is_dir()
|
||||
and not path.is_symlink()
|
||||
and any(path.iterdir())
|
||||
):
|
||||
seen.add(resolved)
|
||||
candidates.append(path)
|
||||
except OSError:
|
||||
continue
|
||||
return sorted(candidates, key=_latest_state_mtime_ns, reverse=True)
|
||||
|
||||
|
||||
def discover_referenced_temp_homes(target: Path) -> list[Path]:
|
||||
"""Find temporary Codex homes referenced by retained thread rows."""
|
||||
references: set[Path] = set()
|
||||
for database in target.rglob("*"):
|
||||
if not database.is_file() or database.suffix not in _SQLITE_SUFFIXES:
|
||||
continue
|
||||
connection: sqlite3.Connection | None = None
|
||||
try:
|
||||
connection = sqlite3.connect(f"file:{database}?mode=ro", uri=True)
|
||||
tables = _database_schema(connection, "main")
|
||||
if "threads" not in tables:
|
||||
continue
|
||||
columns = {str(column[1]) for column in _table_columns(connection, "main", "threads")}
|
||||
if "rollout_path" not in columns:
|
||||
continue
|
||||
for (rollout_path,) in connection.execute("SELECT rollout_path FROM threads"):
|
||||
path = Path(str(rollout_path))
|
||||
for index, part in enumerate(path.parts):
|
||||
if part.startswith(_TEMP_HOME_PREFIX):
|
||||
references.add(Path(*path.parts[: index + 1]))
|
||||
break
|
||||
except sqlite3.Error:
|
||||
continue
|
||||
finally:
|
||||
if connection is not None:
|
||||
connection.close()
|
||||
return sorted(references)
|
||||
|
||||
|
||||
def _history_session_ids(history: Path) -> set[str]:
|
||||
session_ids: set[str] = set()
|
||||
if not history.is_file():
|
||||
return session_ids
|
||||
for line in history.read_text(encoding="utf-8", errors="replace").splitlines():
|
||||
try:
|
||||
record = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
session_id = record.get("session_id") or record.get("thread_id") or record.get("id")
|
||||
if session_id:
|
||||
session_ids.add(str(session_id))
|
||||
return session_ids
|
||||
|
||||
|
||||
def _rollout_session_ids(target: Path) -> set[str]:
|
||||
session_ids: set[str] = set()
|
||||
for directory in (target / "sessions", target / "archived_sessions"):
|
||||
if not directory.is_dir():
|
||||
continue
|
||||
for rollout in directory.rglob("rollout-*.jsonl"):
|
||||
try:
|
||||
with rollout.open(encoding="utf-8", errors="replace") as handle:
|
||||
first_line = handle.readline()
|
||||
record = json.loads(first_line)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
continue
|
||||
if record.get("type") != "session_meta":
|
||||
continue
|
||||
payload = record.get("payload", {})
|
||||
session_id = payload.get("id") or payload.get("session_id")
|
||||
if session_id:
|
||||
session_ids.add(str(session_id))
|
||||
return session_ids
|
||||
|
||||
|
||||
def audit_codex_history(target: Path) -> CodexHistoryAudit | None:
|
||||
"""Compare durable history, rollout files, and the canonical thread index."""
|
||||
database = target / "state_5.sqlite"
|
||||
if not database.is_file():
|
||||
return None
|
||||
connection: sqlite3.Connection | None = None
|
||||
try:
|
||||
connection = sqlite3.connect(f"{database.resolve().as_uri()}?mode=ro", uri=True)
|
||||
tables = _database_schema(connection, "main")
|
||||
if "threads" not in tables:
|
||||
return None
|
||||
columns = {str(column[1]) for column in _table_columns(connection, "main", "threads")}
|
||||
if "id" not in columns:
|
||||
return None
|
||||
if "archived" in columns:
|
||||
rows = connection.execute("SELECT id, archived FROM threads").fetchall()
|
||||
else:
|
||||
rows = [(thread_id, 0) for (thread_id,) in connection.execute("SELECT id FROM threads")]
|
||||
except sqlite3.Error:
|
||||
return None
|
||||
finally:
|
||||
if connection is not None:
|
||||
connection.close()
|
||||
thread_ids = {str(row[0]) for row in rows}
|
||||
archived = sum(bool(row[1]) for row in rows)
|
||||
rollout_ids = _rollout_session_ids(target)
|
||||
history_ids = _history_session_ids(target / "history.jsonl")
|
||||
return CodexHistoryAudit(
|
||||
indexed=len(thread_ids),
|
||||
active=len(thread_ids) - archived,
|
||||
archived=archived,
|
||||
unindexed_rollouts=tuple(sorted(rollout_ids - thread_ids)),
|
||||
history_without_rollout=tuple(sorted(history_ids - thread_ids - rollout_ids)),
|
||||
)
|
||||
|
||||
|
||||
def discover_retained_sources(target: Path) -> list[Path]:
|
||||
"""Find pinned sources retained by interrupted or failed recovery attempts."""
|
||||
recovery_root = target.parent / _RECOVERY_DIR
|
||||
candidates: list[Path] = []
|
||||
try:
|
||||
attempts = recovery_root.iterdir()
|
||||
for attempt in attempts:
|
||||
pinned = attempt / "source-pinned"
|
||||
if (
|
||||
attempt.is_dir()
|
||||
and not (attempt / "manifest.json").exists()
|
||||
and pinned.is_dir()
|
||||
and not pinned.is_symlink()
|
||||
and any(pinned.iterdir())
|
||||
):
|
||||
candidates.append(pinned)
|
||||
except OSError:
|
||||
return []
|
||||
return sorted(candidates, key=_latest_state_mtime_ns, reverse=True)
|
||||
|
||||
|
||||
def home_fingerprint(home: Path) -> str:
|
||||
"""Return a content-independent fingerprint used to detect concurrent writes."""
|
||||
digest = hashlib.sha256()
|
||||
if not home.exists():
|
||||
return digest.hexdigest()
|
||||
for path in sorted(home.rglob("*")):
|
||||
relative = path.relative_to(home)
|
||||
try:
|
||||
metadata = path.lstat()
|
||||
except FileNotFoundError:
|
||||
continue
|
||||
digest.update(os.fsencode(str(relative)))
|
||||
digest.update(str(metadata.st_mode).encode())
|
||||
digest.update(str(metadata.st_size).encode())
|
||||
digest.update(str(metadata.st_mtime_ns).encode())
|
||||
if stat.S_ISLNK(metadata.st_mode):
|
||||
digest.update(os.fsencode(os.readlink(path)))
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _is_runtime_artifact(path: Path) -> bool:
|
||||
name = path.name
|
||||
return name in _RUNTIME_NAMES or any(name.endswith(suffix) for suffix in _RUNTIME_SUFFIXES)
|
||||
|
||||
|
||||
def _secure_tree(path: Path) -> None:
|
||||
if not path.exists():
|
||||
return
|
||||
path.chmod(0o700)
|
||||
for entry in path.rglob("*"):
|
||||
if entry.is_symlink():
|
||||
continue
|
||||
entry.chmod(0o700 if entry.is_dir() else 0o600)
|
||||
|
||||
|
||||
def _copy_home(source: Path, destination: Path, skipped: list[str]) -> None:
|
||||
destination.mkdir(parents=True, exist_ok=False)
|
||||
destination.chmod(0o700)
|
||||
for entry in source.rglob("*"):
|
||||
relative = entry.relative_to(source)
|
||||
output = destination / relative
|
||||
try:
|
||||
metadata = entry.lstat()
|
||||
except FileNotFoundError:
|
||||
continue
|
||||
if stat.S_ISSOCK(metadata.st_mode):
|
||||
skipped.append(str(relative))
|
||||
continue
|
||||
if entry.is_symlink():
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.symlink_to(os.readlink(entry))
|
||||
elif entry.is_dir():
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
elif entry.is_file():
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(entry, output, follow_symlinks=False)
|
||||
else:
|
||||
skipped.append(str(relative))
|
||||
_secure_tree(destination)
|
||||
|
||||
|
||||
def _new_backup_dir(target: Path) -> Path:
|
||||
root = target.parent / _RECOVERY_DIR
|
||||
root.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
stamp = time.strftime("%Y%m%d-%H%M%S")
|
||||
stem = f"{stamp}-{os.getpid()}"
|
||||
for counter in range(1000):
|
||||
suffix = "" if counter == 0 else f"-{counter}"
|
||||
candidate = root / f"{stem}{suffix}"
|
||||
try:
|
||||
candidate.mkdir(mode=0o700)
|
||||
except FileExistsError:
|
||||
continue
|
||||
return candidate
|
||||
raise RuntimeError("could not allocate a Codex recovery backup directory")
|
||||
|
||||
|
||||
def _uses_legacy_headroom_routing(document: Any) -> bool:
|
||||
providers = document.get("model_providers")
|
||||
headroom_provider = providers.get("headroom") if providers is not None else None
|
||||
local_provider = headroom_provider is not None and str(
|
||||
headroom_provider.get("base_url", "")
|
||||
).startswith("http://127.0.0.1:")
|
||||
local_openai_url = str(document.get("openai_base_url", "")).startswith("http://127.0.0.1:")
|
||||
return document.get("model_provider") == "headroom" and (local_provider or local_openai_url)
|
||||
|
||||
|
||||
def _clean_managed_codex_config(document: Any) -> None:
|
||||
providers = document.get("model_providers")
|
||||
headroom_provider = providers.get("headroom") if providers is not None else None
|
||||
local_provider = headroom_provider is not None and str(
|
||||
headroom_provider.get("base_url", "")
|
||||
).startswith("http://127.0.0.1:")
|
||||
local_openai_url = str(document.get("openai_base_url", "")).startswith("http://127.0.0.1:")
|
||||
managed_routing = local_provider or local_openai_url
|
||||
if managed_routing and document.get("model_provider") == "headroom":
|
||||
del document["model_provider"]
|
||||
if providers is not None and local_provider:
|
||||
del providers["headroom"]
|
||||
if not providers:
|
||||
del document["model_providers"]
|
||||
if managed_routing and local_openai_url:
|
||||
del document["openai_base_url"]
|
||||
|
||||
|
||||
def _merge_toml_table(target: Any, source: Any, *, source_wins: bool) -> None:
|
||||
for key, source_value in source.items():
|
||||
if key not in target:
|
||||
target[key] = source_value
|
||||
continue
|
||||
target_value = target[key]
|
||||
if hasattr(target_value, "items") and hasattr(source_value, "items"):
|
||||
_merge_toml_table(target_value, source_value, source_wins=source_wins)
|
||||
elif source_wins:
|
||||
target[key] = source_value
|
||||
|
||||
|
||||
def _merge_config(source: Path, target: Path) -> None:
|
||||
source_document = tomlkit.parse(source.read_text(encoding="utf-8"))
|
||||
_clean_managed_codex_config(source_document)
|
||||
if target.exists():
|
||||
target_document = tomlkit.parse(target.read_text(encoding="utf-8"))
|
||||
_clean_managed_codex_config(target_document)
|
||||
source_wins = source.stat().st_mtime_ns > target.stat().st_mtime_ns
|
||||
_merge_toml_table(target_document, source_document, source_wins=source_wins)
|
||||
else:
|
||||
target_document = source_document
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text(tomlkit.dumps(target_document), encoding="utf-8")
|
||||
|
||||
|
||||
def _read_jsonl(path: Path, quarantine: Path, report: RecoveryReport) -> list[str]:
|
||||
lines: list[str] = []
|
||||
malformed = False
|
||||
for raw_line in path.read_text(encoding="utf-8", errors="replace").splitlines():
|
||||
if not raw_line.strip():
|
||||
continue
|
||||
try:
|
||||
json.loads(raw_line)
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
malformed = True
|
||||
continue
|
||||
lines.append(raw_line)
|
||||
if malformed:
|
||||
destination = quarantine / path.name
|
||||
counter = 1
|
||||
while destination.exists():
|
||||
destination = quarantine / f"{path.stem}-{counter}{path.suffix}"
|
||||
counter += 1
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(path, destination)
|
||||
report.quarantined.append(str(path))
|
||||
return lines
|
||||
|
||||
|
||||
def _merge_jsonl(source: Path, target: Path, quarantine: Path, report: RecoveryReport) -> None:
|
||||
existing = _read_jsonl(target, quarantine, report) if target.exists() else []
|
||||
incoming = _read_jsonl(source, quarantine, report)
|
||||
merged = list(existing)
|
||||
seen = set(existing)
|
||||
for line in incoming:
|
||||
if line not in seen:
|
||||
seen.add(line)
|
||||
merged.append(line)
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text("".join(f"{line}\n" for line in merged), encoding="utf-8")
|
||||
|
||||
|
||||
def _merge_rollout(
|
||||
source: Path,
|
||||
target: Path,
|
||||
quarantine: Path,
|
||||
report: RecoveryReport,
|
||||
replacement_provider: str | None,
|
||||
) -> None:
|
||||
incoming = _read_jsonl(source, quarantine, report)
|
||||
keep_target = target.exists() and source.stat().st_mtime_ns <= target.stat().st_mtime_ns
|
||||
lines = _read_jsonl(target, quarantine, report) if keep_target else incoming
|
||||
provider_replaced = False
|
||||
if replacement_provider is not None:
|
||||
for index, line in enumerate(lines):
|
||||
record = json.loads(line)
|
||||
if not isinstance(record, dict):
|
||||
continue
|
||||
payload = record.get("payload")
|
||||
if (
|
||||
record.get("type") == "session_meta"
|
||||
and isinstance(payload, dict)
|
||||
and payload.get("model_provider") == "headroom"
|
||||
):
|
||||
payload["model_provider"] = replacement_provider
|
||||
lines[index] = json.dumps(record, separators=(",", ":"))
|
||||
provider_replaced = True
|
||||
if keep_target and not provider_replaced:
|
||||
return
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text("".join(f"{line}\n" for line in lines), encoding="utf-8")
|
||||
|
||||
|
||||
def _quote(identifier: str) -> str:
|
||||
return '"' + identifier.replace('"', '""') + '"'
|
||||
|
||||
|
||||
def _database_schema(connection: sqlite3.Connection, schema: str) -> dict[str, str]:
|
||||
rows = connection.execute(
|
||||
f"SELECT name, sql FROM {_quote(schema)}.sqlite_master "
|
||||
"WHERE type = 'table' AND name NOT LIKE 'sqlite_%' ORDER BY name"
|
||||
).fetchall()
|
||||
return dict(rows)
|
||||
|
||||
|
||||
def _database_schema_objects(
|
||||
connection: sqlite3.Connection, schema: str
|
||||
) -> list[tuple[str, str, str, str]]:
|
||||
return connection.execute(
|
||||
f"SELECT type, name, tbl_name, sql FROM {_quote(schema)}.sqlite_master "
|
||||
"WHERE name NOT LIKE 'sqlite_%' AND sql IS NOT NULL "
|
||||
"ORDER BY type, name"
|
||||
).fetchall()
|
||||
|
||||
|
||||
def _table_columns(
|
||||
connection: sqlite3.Connection, schema: str, table: str
|
||||
) -> list[tuple[Any, ...]]:
|
||||
return connection.execute(f"PRAGMA {_quote(schema)}.table_info({_quote(table)})").fetchall()
|
||||
|
||||
|
||||
def _active_model_provider(config_file: Path) -> str:
|
||||
if not config_file.is_file():
|
||||
return "openai"
|
||||
document = tomlkit.parse(config_file.read_text(encoding="utf-8"))
|
||||
return str(document.get("model_provider", "openai"))
|
||||
|
||||
|
||||
def _recovered_rollout_relative_path(rollout_path: Path, source_home: Path) -> Path | None:
|
||||
try:
|
||||
return rollout_path.relative_to(source_home)
|
||||
except ValueError:
|
||||
for index, part in enumerate(rollout_path.parts):
|
||||
if part.startswith(_TEMP_HOME_PREFIX):
|
||||
return Path(*rollout_path.parts[index + 1 :])
|
||||
return None
|
||||
|
||||
|
||||
def _thread_rollout_paths(connection: sqlite3.Connection, schema: str) -> dict[Any, str]:
|
||||
tables = _database_schema(connection, schema)
|
||||
if "threads" not in tables:
|
||||
return {}
|
||||
columns = {str(column[1]) for column in _table_columns(connection, schema, "threads")}
|
||||
if not {"id", "rollout_path"}.issubset(columns):
|
||||
return {}
|
||||
return dict(
|
||||
connection.execute(f"SELECT id, rollout_path FROM {_quote(schema)}.threads").fetchall()
|
||||
)
|
||||
|
||||
|
||||
def _normalize_recovered_threads(
|
||||
connection: sqlite3.Connection,
|
||||
*,
|
||||
source_home: Path,
|
||||
target_home: Path,
|
||||
replacement_provider: str | None,
|
||||
source_rollout_paths: dict[Any, str],
|
||||
) -> None:
|
||||
tables = _database_schema(connection, "main")
|
||||
if "threads" not in tables:
|
||||
return
|
||||
columns = {str(column[1]) for column in _table_columns(connection, "main", "threads")}
|
||||
if not {"id", "rollout_path"}.issubset(columns):
|
||||
return
|
||||
select_columns = "id, rollout_path"
|
||||
if "model_provider" in columns:
|
||||
select_columns += ", model_provider"
|
||||
rows = connection.execute(f"SELECT {select_columns} FROM threads").fetchall()
|
||||
for row in rows:
|
||||
thread_id, rollout_path = row[:2]
|
||||
source_rollout_path = source_rollout_paths.get(thread_id)
|
||||
if source_rollout_path is None:
|
||||
continue
|
||||
relative = _recovered_rollout_relative_path(Path(source_rollout_path), source_home)
|
||||
if relative is None:
|
||||
continue
|
||||
durable_rollout = target_home / relative
|
||||
if str(rollout_path) not in {source_rollout_path, str(durable_rollout)}:
|
||||
continue
|
||||
if not durable_rollout.is_file():
|
||||
raise RuntimeError(f"recovered rollout is missing for thread {thread_id}")
|
||||
connection.execute(
|
||||
"UPDATE threads SET rollout_path = ? WHERE id = ?",
|
||||
(str(durable_rollout), thread_id),
|
||||
)
|
||||
if replacement_provider is not None and len(row) == 3 and row[2] == "headroom":
|
||||
connection.execute(
|
||||
"UPDATE threads SET model_provider = ? WHERE id = ?",
|
||||
(replacement_provider, thread_id),
|
||||
)
|
||||
|
||||
|
||||
def _merge_database(
|
||||
source: Path,
|
||||
target: Path,
|
||||
*,
|
||||
source_home: Path,
|
||||
target_home: Path,
|
||||
replacement_provider: str | None,
|
||||
) -> None:
|
||||
if not target.exists():
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(source, target)
|
||||
connection = sqlite3.connect(target)
|
||||
try:
|
||||
source_rollout_paths = _thread_rollout_paths(connection, "main")
|
||||
_normalize_recovered_threads(
|
||||
connection,
|
||||
source_home=source_home,
|
||||
target_home=target_home,
|
||||
replacement_provider=replacement_provider,
|
||||
source_rollout_paths=source_rollout_paths,
|
||||
)
|
||||
connection.commit()
|
||||
except Exception:
|
||||
connection.rollback()
|
||||
raise
|
||||
finally:
|
||||
connection.close()
|
||||
return
|
||||
source_is_newer = source.stat().st_mtime_ns > target.stat().st_mtime_ns
|
||||
connection = sqlite3.connect(target)
|
||||
try:
|
||||
connection.execute("PRAGMA foreign_keys = OFF")
|
||||
connection.execute("ATTACH DATABASE ? AS incoming", (str(source),))
|
||||
target_schema = _database_schema(connection, "main")
|
||||
source_schema = _database_schema(connection, "incoming")
|
||||
if target_schema != source_schema or _database_schema_objects(
|
||||
connection, "main"
|
||||
) != _database_schema_objects(connection, "incoming"):
|
||||
raise RuntimeError(f"SQLite schema mismatch for {target.name}")
|
||||
source_rollout_paths = _thread_rollout_paths(connection, "incoming")
|
||||
if "_sqlx_migrations" in target_schema:
|
||||
migration_columns = [
|
||||
str(column[1]) for column in _table_columns(connection, "main", "_sqlx_migrations")
|
||||
]
|
||||
if "version" in migration_columns and "checksum" in migration_columns:
|
||||
target_migrations = dict(
|
||||
connection.execute(
|
||||
"SELECT version, checksum FROM main._sqlx_migrations"
|
||||
).fetchall()
|
||||
)
|
||||
source_migrations = dict(
|
||||
connection.execute(
|
||||
"SELECT version, checksum FROM incoming._sqlx_migrations"
|
||||
).fetchall()
|
||||
)
|
||||
for version in target_migrations.keys() & source_migrations.keys():
|
||||
if target_migrations[version] != source_migrations[version]:
|
||||
raise RuntimeError(
|
||||
f"SQLite migration mismatch for {target.name} at version {version}"
|
||||
)
|
||||
connection.execute("BEGIN IMMEDIATE")
|
||||
try:
|
||||
for table in target_schema:
|
||||
columns = _table_columns(connection, "main", table)
|
||||
source_columns = _table_columns(connection, "incoming", table)
|
||||
if columns != source_columns:
|
||||
raise RuntimeError(f"SQLite schema mismatch for {target.name}:{table}")
|
||||
column_names = [str(column[1]) for column in columns]
|
||||
primary_key = [
|
||||
str(column[1])
|
||||
for column in sorted(columns, key=lambda row: row[5])
|
||||
if column[5]
|
||||
]
|
||||
quoted_columns = ", ".join(_quote(name) for name in column_names)
|
||||
rows = connection.execute(
|
||||
f"SELECT {quoted_columns} FROM incoming.{_quote(table)}"
|
||||
).fetchall()
|
||||
placeholders = ", ".join("?" for _ in column_names)
|
||||
verb = (
|
||||
"INSERT OR REPLACE" if primary_key and source_is_newer else "INSERT OR IGNORE"
|
||||
)
|
||||
connection.executemany(
|
||||
f"{verb} INTO {_quote(table)} ({quoted_columns}) VALUES ({placeholders})",
|
||||
rows,
|
||||
)
|
||||
_normalize_recovered_threads(
|
||||
connection,
|
||||
source_home=source_home,
|
||||
target_home=target_home,
|
||||
replacement_provider=replacement_provider,
|
||||
source_rollout_paths=source_rollout_paths,
|
||||
)
|
||||
connection.commit()
|
||||
integrity = connection.execute("PRAGMA integrity_check").fetchone()
|
||||
if not integrity or integrity[0] != "ok":
|
||||
raise RuntimeError(f"SQLite integrity check failed for {target.name}")
|
||||
violations = connection.execute("PRAGMA foreign_key_check").fetchall()
|
||||
if violations:
|
||||
raise RuntimeError(f"SQLite foreign key check failed for {target.name}")
|
||||
except Exception:
|
||||
connection.rollback()
|
||||
raise
|
||||
finally:
|
||||
connection.execute("DETACH DATABASE incoming")
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
def _merge_pinned_home(
|
||||
pinned: Path,
|
||||
target: Path,
|
||||
report: RecoveryReport,
|
||||
*,
|
||||
source_home: Path,
|
||||
) -> None:
|
||||
quarantine = report.backup_dir / "quarantine"
|
||||
source_config = pinned / "config.toml"
|
||||
replace_legacy_provider = source_config.is_file() and _uses_legacy_headroom_routing(
|
||||
tomlkit.parse(source_config.read_text(encoding="utf-8"))
|
||||
)
|
||||
replacement_provider = (
|
||||
_active_model_provider(target / "config.toml") if replace_legacy_provider else None
|
||||
)
|
||||
for source in sorted(pinned.rglob("*")):
|
||||
relative = source.relative_to(pinned)
|
||||
destination = target / relative
|
||||
if source.is_dir() or source.is_symlink():
|
||||
if source.is_symlink() and not destination.exists() and not destination.is_symlink():
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
destination.symlink_to(os.readlink(source))
|
||||
report.copied.append(str(relative))
|
||||
continue
|
||||
if _is_runtime_artifact(source):
|
||||
report.skipped_runtime.append(str(relative))
|
||||
continue
|
||||
if source.name == "config.toml":
|
||||
_merge_config(source, destination)
|
||||
report.merged.append(str(relative))
|
||||
elif source.suffix == ".jsonl" and relative.parts[0] in {
|
||||
"sessions",
|
||||
"archived_sessions",
|
||||
}:
|
||||
_merge_rollout(
|
||||
source,
|
||||
destination,
|
||||
quarantine,
|
||||
report,
|
||||
replacement_provider,
|
||||
)
|
||||
report.merged.append(str(relative))
|
||||
elif source.suffix == ".jsonl":
|
||||
_merge_jsonl(source, destination, quarantine, report)
|
||||
report.merged.append(str(relative))
|
||||
elif source.suffix in _SQLITE_SUFFIXES:
|
||||
_merge_database(
|
||||
source,
|
||||
destination,
|
||||
source_home=source_home,
|
||||
target_home=target,
|
||||
replacement_provider=replacement_provider,
|
||||
)
|
||||
report.merged.append(str(relative))
|
||||
elif not destination.exists() or source.stat().st_mtime_ns > destination.stat().st_mtime_ns:
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy2(source, destination)
|
||||
report.copied.append(str(relative))
|
||||
|
||||
|
||||
def _restore_target(target: Path, target_backup: Path, target_existed: bool) -> None:
|
||||
# Windows can retain SQLite file handles until statement finalizers run.
|
||||
gc.collect()
|
||||
if target.exists():
|
||||
target.replace(target_backup.parent / "target-failed")
|
||||
if target_existed:
|
||||
shutil.copytree(target_backup, target, symlinks=True)
|
||||
|
||||
|
||||
def _capture_modes(home: Path) -> dict[str, int]:
|
||||
modes = {".": stat.S_IMODE(home.stat().st_mode)}
|
||||
for entry in home.rglob("*"):
|
||||
if not entry.is_symlink():
|
||||
modes[str(entry.relative_to(home))] = stat.S_IMODE(entry.stat().st_mode)
|
||||
return modes
|
||||
|
||||
|
||||
def _restore_modes(home: Path, modes: dict[str, int]) -> None:
|
||||
for relative, mode in modes.items():
|
||||
path = home if relative == "." else home / relative
|
||||
if path.exists() and not path.is_symlink():
|
||||
path.chmod(mode)
|
||||
|
||||
|
||||
def _reject_target_symlink_traversal(source: Path, target: Path) -> None:
|
||||
if not target.exists():
|
||||
return
|
||||
for entry in source.rglob("*"):
|
||||
if entry.is_dir() or entry.is_symlink():
|
||||
continue
|
||||
destination = target
|
||||
for part in entry.relative_to(source).parts:
|
||||
destination /= part
|
||||
if destination.is_symlink():
|
||||
raise ValueError(f"recovery would write through target symlink: {destination}")
|
||||
|
||||
|
||||
def recover_codex_home(*, source: Path, target: Path) -> RecoveryReport:
|
||||
"""Merge one quiet temporary Codex home into the active home transactionally."""
|
||||
source = source.expanduser().resolve()
|
||||
target = target.expanduser().resolve()
|
||||
if not source.is_dir() or source == target:
|
||||
raise ValueError("source must be an existing Codex home different from target")
|
||||
if source in target.parents or target in source.parents or target == Path(target.anchor):
|
||||
raise ValueError("source and target Codex homes must not overlap")
|
||||
_reject_target_symlink_traversal(source, target)
|
||||
before = home_fingerprint(source)
|
||||
backup_dir = _new_backup_dir(target)
|
||||
report = RecoveryReport(source=source, target=target, backup_dir=backup_dir)
|
||||
pinned = backup_dir / "source-pinned"
|
||||
target_backup = backup_dir / "target-before"
|
||||
_copy_home(source, pinned, report.skipped_runtime)
|
||||
if home_fingerprint(source) != before:
|
||||
raise RuntimeError("source Codex home changed while it was being pinned")
|
||||
target_existed = target.exists()
|
||||
target_modes: dict[str, int] = {}
|
||||
if target_existed:
|
||||
target_modes = _capture_modes(target)
|
||||
target_fingerprint = home_fingerprint(target)
|
||||
_copy_home(target, target_backup, report.skipped_runtime)
|
||||
if home_fingerprint(target) != target_fingerprint:
|
||||
raise RuntimeError("target Codex home changed while it was being backed up")
|
||||
else:
|
||||
target.mkdir(mode=0o700, parents=True)
|
||||
try:
|
||||
_merge_pinned_home(pinned, target, report, source_home=source)
|
||||
except Exception:
|
||||
_restore_target(target, target_backup, target_existed)
|
||||
if target_existed:
|
||||
_restore_modes(target, target_modes)
|
||||
raise
|
||||
manifest = asdict(report)
|
||||
manifest.update(source=str(source), target=str(target), backup_dir=str(backup_dir))
|
||||
manifest_file = backup_dir / "manifest.json"
|
||||
manifest_file.write_text(
|
||||
json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8"
|
||||
)
|
||||
manifest_file.chmod(0o600)
|
||||
return report
|
||||
|
|
@ -60,6 +60,7 @@ dependencies = [
|
|||
"ast-grep-cli>=0.30.0", # AST-aware code slicing (CodeCompressor); binary wheel
|
||||
"pyyaml>=6.0", # omp wrap: parse/merge omp's models.yml registry
|
||||
"tomli>=2.0.0; python_version < '3.11'", # tomllib backport for helper scripts
|
||||
"tomlkit>=0.13.0,<1.0", # Loss-minimizing Codex config recovery
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
|
|
|
|||
866
tests/test_cli/test_recover_codex.py
Normal file
866
tests/test_cli/test_recover_codex.py
Normal file
|
|
@ -0,0 +1,866 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import errno
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import sqlite3
|
||||
import stat
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
import headroom.providers.codex.recovery as codex_recovery
|
||||
from headroom.cli.main import main
|
||||
from headroom.providers.codex.recovery import discover_dangling_homes, recover_codex_home
|
||||
|
||||
|
||||
def _write_db(path: Path, rows: list[tuple[str, str]]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
connection = sqlite3.connect(path)
|
||||
try:
|
||||
connection.execute("CREATE TABLE threads (id TEXT PRIMARY KEY, title TEXT NOT NULL)")
|
||||
connection.executemany("INSERT INTO threads VALUES (?, ?)", rows)
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
def _write_sqlx_db(path: Path, checksum: bytes) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
connection = sqlite3.connect(path)
|
||||
try:
|
||||
connection.execute(
|
||||
"CREATE TABLE _sqlx_migrations (version INTEGER PRIMARY KEY, checksum BLOB NOT NULL)"
|
||||
)
|
||||
connection.execute("INSERT INTO _sqlx_migrations VALUES (1, ?)", (checksum,))
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
def _write_thread_db(
|
||||
path: Path,
|
||||
rows: list[tuple[str, str, str]],
|
||||
) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with sqlite3.connect(path) as connection:
|
||||
connection.execute(
|
||||
"CREATE TABLE threads ("
|
||||
"id TEXT PRIMARY KEY, rollout_path TEXT NOT NULL, model_provider TEXT NOT NULL"
|
||||
")"
|
||||
)
|
||||
connection.executemany("INSERT INTO threads VALUES (?, ?, ?)", rows)
|
||||
|
||||
|
||||
def test_discover_dangling_homes_only_returns_codex_homes(tmp_path: Path) -> None:
|
||||
candidate = tmp_path / "headroom-codex-home-abc"
|
||||
candidate.mkdir()
|
||||
(candidate / "config.toml").write_text('model = "gpt-5"\n', encoding="utf-8")
|
||||
(tmp_path / "headroom-codex-home-empty").mkdir()
|
||||
(tmp_path / "other").mkdir()
|
||||
outside = tmp_path / "outside"
|
||||
outside.mkdir()
|
||||
(outside / "history.jsonl").write_text("{}\n", encoding="utf-8")
|
||||
(tmp_path / "headroom-codex-home-linked").symlink_to(outside, target_is_directory=True)
|
||||
|
||||
assert discover_dangling_homes(tmp_path) == [candidate]
|
||||
|
||||
|
||||
def test_discover_dangling_homes_uses_newest_state_not_directory_mtime(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
newest_state = tmp_path / "headroom-codex-home-newest-state"
|
||||
newest_directory = tmp_path / "headroom-codex-home-newest-directory"
|
||||
newest_state.mkdir()
|
||||
newest_directory.mkdir()
|
||||
newest_state_file = newest_state / "history.jsonl"
|
||||
newest_directory_file = newest_directory / "history.jsonl"
|
||||
newest_state_file.write_text('{"session_id":"newest"}\n', encoding="utf-8")
|
||||
newest_directory_file.write_text('{"session_id":"older"}\n', encoding="utf-8")
|
||||
os.utime(newest_state_file, ns=(400, 400))
|
||||
os.utime(newest_directory_file, ns=(300, 300))
|
||||
os.utime(newest_state, ns=(100, 100))
|
||||
os.utime(newest_directory, ns=(500, 500))
|
||||
|
||||
assert discover_dangling_homes(tmp_path) == [newest_state, newest_directory]
|
||||
|
||||
|
||||
def test_discover_dangling_homes_searches_tmpdir_and_python_temp_root(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
env_root = tmp_path / "env-tmp"
|
||||
python_root = tmp_path / "python-tmp"
|
||||
env_candidate = env_root / "headroom-codex-home-env"
|
||||
python_candidate = python_root / "headroom-codex-home-python"
|
||||
env_candidate.mkdir(parents=True)
|
||||
python_candidate.mkdir(parents=True)
|
||||
(env_candidate / "history.jsonl").write_text("{}\n", encoding="utf-8")
|
||||
(python_candidate / "history.jsonl").write_text("{}\n", encoding="utf-8")
|
||||
os.utime(env_candidate / "history.jsonl", ns=(100, 100))
|
||||
os.utime(python_candidate / "history.jsonl", ns=(200, 200))
|
||||
monkeypatch.setenv("TMPDIR", str(env_root))
|
||||
monkeypatch.setattr(tempfile, "gettempdir", lambda: str(python_root))
|
||||
|
||||
assert discover_dangling_homes() == [python_candidate, env_candidate]
|
||||
|
||||
|
||||
def test_recovery_merges_files_config_and_sqlite_with_backups(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
(target / "config.toml").write_text(
|
||||
'model = "target-model"\n[features]\nexisting = true\n', encoding="utf-8"
|
||||
)
|
||||
(source / "config.toml").write_text(
|
||||
'model = "source-model"\n[features]\nfrom_wrap = true\n', encoding="utf-8"
|
||||
)
|
||||
os.utime(target / "config.toml", ns=(100, 100))
|
||||
os.utime(source / "config.toml", ns=(200, 200))
|
||||
rollout = source / "sessions" / "2026" / "07" / "14" / "rollout.jsonl"
|
||||
rollout.parent.mkdir(parents=True)
|
||||
rollout.write_text('{"type":"session_meta"}\n', encoding="utf-8")
|
||||
_write_db(target / "sqlite" / "state_5.sqlite", [("target", "Target")])
|
||||
_write_db(source / "sqlite" / "state_5.sqlite", [("source", "Source")])
|
||||
|
||||
report = recover_codex_home(source=source, target=target)
|
||||
|
||||
config = (target / "config.toml").read_text(encoding="utf-8")
|
||||
assert 'model = "source-model"' in config
|
||||
assert "existing = true" in config
|
||||
assert "from_wrap = true" in config
|
||||
assert rollout.relative_to(source).with_name("rollout.jsonl")
|
||||
assert (target / rollout.relative_to(source)).read_text(encoding="utf-8") == (
|
||||
'{"type":"session_meta"}\n'
|
||||
)
|
||||
with sqlite3.connect(target / "sqlite" / "state_5.sqlite") as connection:
|
||||
assert connection.execute("SELECT id, title FROM threads ORDER BY id").fetchall() == [
|
||||
("source", "Source"),
|
||||
("target", "Target"),
|
||||
]
|
||||
assert report.backup_dir.is_dir()
|
||||
assert (report.backup_dir / "target-before").is_dir()
|
||||
assert (report.backup_dir / "source-pinned").is_dir()
|
||||
assert (report.backup_dir / "manifest.json").is_file()
|
||||
|
||||
|
||||
def test_recovery_relocates_thread_rollout_paths_to_durable_home(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
relative_rollout = Path("sessions/2026/07/14/rollout-2026-07-14T10-00-00-thread-1.jsonl")
|
||||
source_rollout = source / relative_rollout
|
||||
source_rollout.parent.mkdir(parents=True)
|
||||
source_rollout.write_text('{"type":"session_meta"}\n', encoding="utf-8")
|
||||
_write_thread_db(target / "state_5.sqlite", [])
|
||||
_write_thread_db(
|
||||
source / "state_5.sqlite",
|
||||
[("thread-1", str(source_rollout), "openai")],
|
||||
)
|
||||
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
durable_rollout = target / relative_rollout
|
||||
assert durable_rollout.is_file()
|
||||
with sqlite3.connect(target / "state_5.sqlite") as connection:
|
||||
assert connection.execute(
|
||||
"SELECT rollout_path FROM threads WHERE id = 'thread-1'"
|
||||
).fetchone() == (str(durable_rollout),)
|
||||
|
||||
|
||||
def test_recovery_ignores_unrelated_dangling_target_thread(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
unrelated_rollout = Path("/private/tmp/headroom-codex-home-deleted/sessions/unrelated.jsonl")
|
||||
relative_rollout = Path("sessions/2026/07/14/rollout-recovered.jsonl")
|
||||
source_rollout = source / relative_rollout
|
||||
source_rollout.parent.mkdir(parents=True)
|
||||
source_rollout.write_text('{"type":"session_meta"}\n', encoding="utf-8")
|
||||
_write_thread_db(
|
||||
target / "state_5.sqlite",
|
||||
[("unrelated", str(unrelated_rollout), "openai")],
|
||||
)
|
||||
_write_thread_db(
|
||||
source / "state_5.sqlite",
|
||||
[("recovered", str(source_rollout), "openai")],
|
||||
)
|
||||
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
with sqlite3.connect(target / "state_5.sqlite") as connection:
|
||||
rows = dict(connection.execute("SELECT id, rollout_path FROM threads"))
|
||||
assert rows == {
|
||||
"unrelated": str(unrelated_rollout),
|
||||
"recovered": str(target / relative_rollout),
|
||||
}
|
||||
|
||||
|
||||
def test_recovery_restores_legacy_headroom_threads_to_active_provider(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
(target / "config.toml").write_text(
|
||||
'model_provider = "azure"\n'
|
||||
"[model_providers.azure]\n"
|
||||
'base_url = "https://azure.example/v1"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
(source / "config.toml").write_text(
|
||||
'model_provider = "headroom"\n'
|
||||
"[model_providers.headroom]\n"
|
||||
'base_url = "http://127.0.0.1:8787/v1"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
relative_rollout = Path("sessions/2026/07/14/rollout-thread-1.jsonl")
|
||||
source_rollout = source / relative_rollout
|
||||
source_rollout.parent.mkdir(parents=True)
|
||||
source_rollout.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session_meta",
|
||||
"payload": {
|
||||
"id": "thread-1",
|
||||
"model_provider": "headroom",
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
_write_thread_db(target / "state_5.sqlite", [])
|
||||
_write_thread_db(
|
||||
source / "state_5.sqlite",
|
||||
[("thread-1", str(source_rollout), "headroom")],
|
||||
)
|
||||
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
with sqlite3.connect(target / "state_5.sqlite") as connection:
|
||||
assert connection.execute(
|
||||
"SELECT model_provider FROM threads WHERE id = 'thread-1'"
|
||||
).fetchone() == ("azure",)
|
||||
session_meta = json.loads((target / relative_rollout).read_text(encoding="utf-8"))
|
||||
assert session_meta["payload"]["model_provider"] == "azure"
|
||||
|
||||
|
||||
def test_recovery_repairs_legacy_provider_after_a_previous_broken_recovery(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
(target / "config.toml").write_text(
|
||||
'model_provider = "azure"\n'
|
||||
"[model_providers.azure]\n"
|
||||
'base_url = "https://azure.example/v1"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
(source / "config.toml").write_text(
|
||||
'model_provider = "headroom"\n'
|
||||
"[model_providers.headroom]\n"
|
||||
'base_url = "http://127.0.0.1:8787/v1"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
relative_rollout = Path("sessions/2026/06/01/rollout-thread-1.jsonl")
|
||||
source_rollout = source / relative_rollout
|
||||
target_rollout = target / relative_rollout
|
||||
source_rollout.parent.mkdir(parents=True)
|
||||
target_rollout.parent.mkdir(parents=True)
|
||||
session_meta = json.dumps(
|
||||
{
|
||||
"type": "session_meta",
|
||||
"payload": {"id": "thread-1", "model_provider": "headroom"},
|
||||
}
|
||||
)
|
||||
source_rollout.write_text(session_meta + "\n", encoding="utf-8")
|
||||
response_item = '{"type":"response_item","payload":{"text":"kept"}}'
|
||||
target_rollout.write_text(session_meta + "\n" + response_item + "\n", encoding="utf-8")
|
||||
os.utime(source_rollout, ns=(1, 1))
|
||||
os.utime(target_rollout, ns=(2, 2))
|
||||
target_db = target / "state_5.sqlite"
|
||||
source_db = source / "state_5.sqlite"
|
||||
_write_thread_db(target_db, [("thread-1", str(target_rollout), "headroom")])
|
||||
_write_thread_db(source_db, [("thread-1", str(source_rollout), "headroom")])
|
||||
os.utime(source_db, ns=(1, 1))
|
||||
os.utime(target_db, ns=(2, 2))
|
||||
|
||||
recover_codex_home(source=source, target=target)
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
with sqlite3.connect(target_db) as connection:
|
||||
assert connection.execute(
|
||||
"SELECT model_provider, rollout_path FROM threads WHERE id = 'thread-1'"
|
||||
).fetchone() == ("azure", str(target_rollout))
|
||||
recovered_meta = json.loads(target_rollout.read_text(encoding="utf-8").splitlines()[0])
|
||||
assert recovered_meta["payload"]["model_provider"] == "azure"
|
||||
assert target_rollout.read_text(encoding="utf-8").splitlines()[1] == response_item
|
||||
|
||||
|
||||
def test_recovery_preserves_nonlocal_provider_named_headroom(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-remote"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
(target / "config.toml").write_text('model = "gpt-5"\n', encoding="utf-8")
|
||||
(source / "config.toml").write_text(
|
||||
'model_provider = "headroom"\n'
|
||||
"[model_providers.headroom]\n"
|
||||
'base_url = "https://gateway.example/v1"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
relative_rollout = Path("archived_sessions/rollout-thread-1.jsonl")
|
||||
source_rollout = source / relative_rollout
|
||||
source_rollout.parent.mkdir(parents=True)
|
||||
source_rollout.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session_meta",
|
||||
"payload": {"id": "thread-1", "model_provider": "headroom"},
|
||||
}
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
_write_thread_db(target / "state_5.sqlite", [])
|
||||
_write_thread_db(
|
||||
source / "state_5.sqlite",
|
||||
[("thread-1", str(source_rollout), "headroom")],
|
||||
)
|
||||
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
recovered_meta = json.loads((target / relative_rollout).read_text(encoding="utf-8"))
|
||||
assert recovered_meta["payload"]["model_provider"] == "headroom"
|
||||
with sqlite3.connect(target / "state_5.sqlite") as connection:
|
||||
assert connection.execute(
|
||||
"SELECT model_provider FROM threads WHERE id = 'thread-1'"
|
||||
).fetchone() == ("headroom",)
|
||||
|
||||
|
||||
def test_recovery_rolls_back_when_sqlite_schema_differs(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
original = 'model = "target"\n'
|
||||
(target / "config.toml").write_text(original, encoding="utf-8")
|
||||
target.chmod(0o755)
|
||||
(target / "config.toml").chmod(0o644)
|
||||
_write_db(target / "sqlite" / "state_5.sqlite", [("target", "Target")])
|
||||
source_db = source / "sqlite" / "state_5.sqlite"
|
||||
source_db.parent.mkdir(parents=True)
|
||||
with sqlite3.connect(source_db) as connection:
|
||||
connection.execute("CREATE TABLE threads (id TEXT PRIMARY KEY, title BLOB)")
|
||||
|
||||
with pytest.raises(RuntimeError, match="schema mismatch"):
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
assert (target / "config.toml").read_text(encoding="utf-8") == original
|
||||
if os.name != "nt":
|
||||
assert stat.S_IMODE(target.stat().st_mode) == 0o755
|
||||
assert stat.S_IMODE((target / "config.toml").stat().st_mode) == 0o644
|
||||
with sqlite3.connect(target / "sqlite" / "state_5.sqlite") as connection:
|
||||
assert connection.execute("SELECT id, title FROM threads").fetchall() == [
|
||||
("target", "Target")
|
||||
]
|
||||
|
||||
|
||||
def test_recovery_rollback_does_not_delete_live_target_recursively(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
(target / "config.toml").write_text('model = "target"\n', encoding="utf-8")
|
||||
(source / "config.toml").write_text('model = "source"\n', encoding="utf-8")
|
||||
_write_db(target / "state_5.sqlite", [("target", "Target")])
|
||||
with sqlite3.connect(source / "state_5.sqlite") as connection:
|
||||
connection.execute("CREATE TABLE threads (id TEXT PRIMARY KEY, title BLOB)")
|
||||
|
||||
def fail_recursive_delete(path: Path) -> None:
|
||||
raise OSError(errno.ENOTEMPTY, "Directory not empty", path)
|
||||
|
||||
monkeypatch.setattr(codex_recovery.shutil, "rmtree", fail_recursive_delete)
|
||||
|
||||
with pytest.raises(RuntimeError, match="SQLite schema mismatch"):
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
assert (target / "config.toml").read_text(encoding="utf-8") == 'model = "target"\n'
|
||||
failed_targets = list((tmp_path / ".headroom-codex-recovery").glob("*/target-failed"))
|
||||
assert len(failed_targets) == 1
|
||||
|
||||
|
||||
def test_recover_codex_cli_previews_then_merges(tmp_path: Path) -> None:
|
||||
home = tmp_path / "home"
|
||||
target = home / ".codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir(parents=True)
|
||||
source.mkdir()
|
||||
(source / "history.jsonl").write_text('{"session_id":"new"}\n', encoding="utf-8")
|
||||
runner = CliRunner()
|
||||
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["recover", "codex", "--source", str(source), "--target", str(target), "--yes"],
|
||||
env={"HOME": str(home)},
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Recovery complete" in result.output
|
||||
assert json.loads((target / "history.jsonl").read_text(encoding="utf-8"))["session_id"] == (
|
||||
"new"
|
||||
)
|
||||
|
||||
|
||||
def test_recover_codex_cli_audits_history_without_treating_prompt_text_as_paths(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
target = tmp_path / "codex"
|
||||
target.mkdir()
|
||||
deleted = Path("/private/tmp/headroom-codex-home-deleted")
|
||||
rollout = target / "sessions/2026/07/14/rollout-indexed.jsonl"
|
||||
rollout.parent.mkdir(parents=True)
|
||||
rollout.write_text('{"type":"session_meta"}\n', encoding="utf-8")
|
||||
_write_thread_db(
|
||||
target / "state_5.sqlite",
|
||||
[("indexed", str(rollout), "openai")],
|
||||
)
|
||||
with sqlite3.connect(target / "state_5.sqlite") as connection:
|
||||
connection.execute("ALTER TABLE threads ADD COLUMN archived INTEGER NOT NULL DEFAULT 0")
|
||||
(target / "history.jsonl").write_text(
|
||||
json.dumps({"session_id": "indexed", "text": "surviving chat"})
|
||||
+ "\n"
|
||||
+ json.dumps(
|
||||
{
|
||||
"session_id": "orphaned",
|
||||
"text": f"pasted error referenced {deleted}/sessions/x",
|
||||
}
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
main,
|
||||
["recover", "codex", "--target", str(target), "--yes"],
|
||||
env={"TMPDIR": str(tmp_path / "empty-tmp")},
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Referenced temporary Codex homes were already deleted:" not in result.output
|
||||
assert str(deleted) not in result.output
|
||||
assert "Durable Codex history: 1 indexed chats (1 active, 0 archived)." in result.output
|
||||
assert "History-only records without a surviving rollout: 1" in result.output
|
||||
assert "orphaned" in result.output
|
||||
assert "codex resume --all" in result.output
|
||||
assert "No recoverable Headroom Codex homes were found." in result.output
|
||||
|
||||
|
||||
def test_recover_codex_cli_reuses_source_pinned_by_failed_recovery(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
target.mkdir()
|
||||
pinned = tmp_path / ".headroom-codex-recovery" / "interrupted-attempt" / "source-pinned"
|
||||
pinned.mkdir(parents=True)
|
||||
(pinned / "history.jsonl").write_text(
|
||||
json.dumps({"session_id": "recovered", "text": "retained"}) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
relative_rollout = Path("sessions/2026/07/14/rollout-retained.jsonl")
|
||||
pinned_rollout = pinned / relative_rollout
|
||||
pinned_rollout.parent.mkdir(parents=True)
|
||||
pinned_rollout.write_text('{"type":"session_meta"}\n', encoding="utf-8")
|
||||
deleted_source = Path("/private/tmp/headroom-codex-home-deleted")
|
||||
_write_thread_db(
|
||||
pinned / "state_5.sqlite",
|
||||
[("retained", str(deleted_source / relative_rollout), "openai")],
|
||||
)
|
||||
|
||||
result = CliRunner().invoke(
|
||||
main,
|
||||
["recover", "codex", "--target", str(target), "--yes"],
|
||||
env={"TMPDIR": str(tmp_path / "empty-tmp")},
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert str(pinned) in result.output
|
||||
assert "Recovery complete." in result.output
|
||||
assert '"session_id": "recovered"' in (target / "history.jsonl").read_text(encoding="utf-8")
|
||||
with sqlite3.connect(target / "state_5.sqlite") as connection:
|
||||
assert connection.execute(
|
||||
"SELECT rollout_path FROM threads WHERE id = 'retained'"
|
||||
).fetchone() == (str(target / relative_rollout),)
|
||||
|
||||
|
||||
def test_recover_codex_cli_decline_changes_nothing(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
target_history = target / "history.jsonl"
|
||||
target_history.write_text('{"session_id":"target"}\n', encoding="utf-8")
|
||||
(source / "history.jsonl").write_text('{"session_id":"source"}\n', encoding="utf-8")
|
||||
|
||||
result = CliRunner().invoke(
|
||||
main,
|
||||
["recover", "codex", "--source", str(source), "--target", str(target)],
|
||||
input="n\n",
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Recovery cancelled. No Codex state was changed." in result.output
|
||||
assert target_history.read_text(encoding="utf-8") == '{"session_id":"target"}\n'
|
||||
assert not (tmp_path / ".headroom-codex-recovery").exists()
|
||||
|
||||
|
||||
def test_recover_codex_cli_reports_malformed_config_and_removes_new_target(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
source.mkdir()
|
||||
(source / "config.toml").write_text("[invalid\n", encoding="utf-8")
|
||||
|
||||
result = CliRunner().invoke(
|
||||
main,
|
||||
["recover", "codex", "--source", str(source), "--target", str(target), "--yes"],
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert "Error: Codex recovery failed:" in result.output
|
||||
assert not target.exists()
|
||||
|
||||
|
||||
def test_recovery_never_writes_through_target_symlinks(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
outside = tmp_path / "outside"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
outside.mkdir()
|
||||
(target / "sessions").symlink_to(outside, target_is_directory=True)
|
||||
source_session = source / "sessions" / "rollout.jsonl"
|
||||
source_session.parent.mkdir()
|
||||
source_session.write_text('{"type":"session_meta"}\n', encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match="symlink"):
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
assert not (outside / "rollout.jsonl").exists()
|
||||
assert (target / "sessions").is_symlink()
|
||||
|
||||
|
||||
def test_recovery_rolls_back_when_sqlite_indexes_differ(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
target_db = target / "sqlite" / "state_5.sqlite"
|
||||
source_db = source / "sqlite" / "state_5.sqlite"
|
||||
_write_db(target_db, [("target", "Target")])
|
||||
connection = sqlite3.connect(target_db)
|
||||
try:
|
||||
connection.execute("CREATE UNIQUE INDEX thread_title ON threads(title)")
|
||||
connection.commit()
|
||||
finally:
|
||||
connection.close()
|
||||
_write_db(source_db, [("source", "Source")])
|
||||
|
||||
with pytest.raises(RuntimeError, match="schema mismatch"):
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
with sqlite3.connect(target_db) as connection:
|
||||
assert connection.execute("SELECT id, title FROM threads").fetchall() == [
|
||||
("target", "Target")
|
||||
]
|
||||
|
||||
|
||||
def test_recovery_rolls_back_when_sqlx_checksums_differ(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
target_db = target / "sqlite" / "state_5.sqlite"
|
||||
source_db = source / "sqlite" / "state_5.sqlite"
|
||||
_write_sqlx_db(target_db, b"target-checksum")
|
||||
_write_sqlx_db(source_db, b"source-checksum")
|
||||
|
||||
with pytest.raises(RuntimeError, match="migration mismatch"):
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
with sqlite3.connect(target_db) as connection:
|
||||
assert connection.execute("SELECT version, checksum FROM _sqlx_migrations").fetchall() == [
|
||||
(1, b"target-checksum")
|
||||
]
|
||||
|
||||
|
||||
def test_recovery_rolls_back_when_source_sqlite_is_corrupt(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
original_config = 'model = "target"\n'
|
||||
(target / "config.toml").write_text(original_config, encoding="utf-8")
|
||||
(source / "config.toml").write_text('model = "source"\n', encoding="utf-8")
|
||||
target_db = target / "sqlite" / "state_5.sqlite"
|
||||
_write_db(target_db, [("target", "Target")])
|
||||
source_db = source / "sqlite" / "state_5.sqlite"
|
||||
source_db.parent.mkdir(parents=True)
|
||||
source_db.write_bytes(b"not a sqlite database")
|
||||
|
||||
with pytest.raises(sqlite3.DatabaseError):
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
assert (target / "config.toml").read_text(encoding="utf-8") == original_config
|
||||
with sqlite3.connect(target_db) as connection:
|
||||
assert connection.execute("SELECT id, title FROM threads").fetchall() == [
|
||||
("target", "Target")
|
||||
]
|
||||
|
||||
|
||||
def test_recovery_rolls_back_when_source_sqlite_breaks_foreign_keys(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
target_db = target / "sqlite" / "state_5.sqlite"
|
||||
source_db = source / "sqlite" / "state_5.sqlite"
|
||||
for database in (target_db, source_db):
|
||||
database.parent.mkdir(parents=True, exist_ok=True)
|
||||
with sqlite3.connect(database) as connection:
|
||||
connection.executescript(
|
||||
"CREATE TABLE parents (id TEXT PRIMARY KEY);"
|
||||
"CREATE TABLE children ("
|
||||
"id TEXT PRIMARY KEY, parent_id TEXT REFERENCES parents(id)"
|
||||
");"
|
||||
)
|
||||
with sqlite3.connect(target_db) as connection:
|
||||
connection.execute("INSERT INTO parents VALUES ('target-parent')")
|
||||
with sqlite3.connect(source_db) as connection:
|
||||
connection.execute("INSERT INTO children VALUES ('orphan', 'missing-parent')")
|
||||
|
||||
with pytest.raises(RuntimeError, match="foreign key check failed"):
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
with sqlite3.connect(target_db) as connection:
|
||||
assert connection.execute("SELECT id FROM parents").fetchall() == [("target-parent",)]
|
||||
assert connection.execute("SELECT id, parent_id FROM children").fetchall() == []
|
||||
|
||||
|
||||
def test_recovery_removes_legacy_headroom_routing_from_config(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
(target / "config.toml").write_text('model = "gpt-5"\n', encoding="utf-8")
|
||||
(source / "config.toml").write_text(
|
||||
'model_provider = "headroom"\n'
|
||||
'openai_base_url = "http://127.0.0.1:8787/v1"\n'
|
||||
"[model_providers.headroom]\n"
|
||||
'base_url = "http://127.0.0.1:8787/v1"\n'
|
||||
"[features]\nfrom_wrapped_session = true\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
config = (target / "config.toml").read_text(encoding="utf-8")
|
||||
assert "headroom" not in config
|
||||
assert "127.0.0.1:8787" not in config
|
||||
assert "from_wrapped_session = true" in config
|
||||
|
||||
|
||||
def test_recovery_preserves_user_defined_remote_headroom_provider(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
(source / "config.toml").write_text(
|
||||
'model_provider = "headroom"\n'
|
||||
"[model_providers.headroom]\n"
|
||||
'base_url = "https://gateway.example/v1"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
config = (target / "config.toml").read_text(encoding="utf-8")
|
||||
assert 'model_provider = "headroom"' in config
|
||||
assert "[model_providers.headroom]" in config
|
||||
assert 'base_url = "https://gateway.example/v1"' in config
|
||||
|
||||
|
||||
def test_recovery_quarantines_malformed_jsonl_and_keeps_valid_records(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
(target / "history.jsonl").write_text('{"session_id":"target"}\n', encoding="utf-8")
|
||||
(source / "history.jsonl").write_text(
|
||||
'{"session_id":"source-1"}\nnot-json\n{"session_id":"source-2"}\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
report = recover_codex_home(source=source, target=target)
|
||||
|
||||
recovered = [
|
||||
json.loads(line)["session_id"]
|
||||
for line in (target / "history.jsonl").read_text(encoding="utf-8").splitlines()
|
||||
]
|
||||
assert recovered == ["target", "source-1", "source-2"]
|
||||
assert report.quarantined == [str(report.backup_dir / "source-pinned" / "history.jsonl")]
|
||||
assert "not-json" in (report.backup_dir / "quarantine" / "history.jsonl").read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
|
||||
|
||||
def test_recovery_keeps_newest_divergent_rollout_and_backs_up_both(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
relative = Path("sessions/2026/07/14/rollout.jsonl")
|
||||
target_rollout = target / relative
|
||||
source_rollout = source / relative
|
||||
target_rollout.parent.mkdir(parents=True)
|
||||
source_rollout.parent.mkdir(parents=True)
|
||||
target_rollout.write_text('{"thread":"newer-target"}\n', encoding="utf-8")
|
||||
source_rollout.write_text('{"thread":"older-source"}\n', encoding="utf-8")
|
||||
os.utime(source_rollout, ns=(100, 100))
|
||||
os.utime(target_rollout, ns=(200, 200))
|
||||
|
||||
report = recover_codex_home(source=source, target=target)
|
||||
|
||||
assert target_rollout.read_text(encoding="utf-8") == '{"thread":"newer-target"}\n'
|
||||
assert (report.backup_dir / "source-pinned" / relative).read_text(
|
||||
encoding="utf-8"
|
||||
) == '{"thread":"older-source"}\n'
|
||||
assert (report.backup_dir / "target-before" / relative).read_text(
|
||||
encoding="utf-8"
|
||||
) == '{"thread":"newer-target"}\n'
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.platform == "win32" or not hasattr(socket, "AF_UNIX"),
|
||||
reason="requires POSIX Unix domain sockets",
|
||||
)
|
||||
def test_recovery_records_sockets_and_secures_both_backups(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir(mode=0o755)
|
||||
source.mkdir(mode=0o755)
|
||||
source_history = source / "history.jsonl"
|
||||
source_history.write_text('{"session_id":"source"}\n', encoding="utf-8")
|
||||
source_history.chmod(0o644)
|
||||
socket_path = source / "codex.sock"
|
||||
fifo_path = source / "codex.pipe"
|
||||
os.mkfifo(fifo_path)
|
||||
|
||||
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as codex_socket:
|
||||
codex_socket.bind(str(socket_path))
|
||||
report = recover_codex_home(source=source, target=target)
|
||||
|
||||
pinned = report.backup_dir / "source-pinned"
|
||||
target_backup = report.backup_dir / "target-before"
|
||||
assert "codex.sock" in report.skipped_runtime
|
||||
assert "codex.pipe" in report.skipped_runtime
|
||||
assert not (pinned / "codex.sock").exists()
|
||||
assert not (pinned / "codex.pipe").exists()
|
||||
assert stat.S_IMODE(report.backup_dir.stat().st_mode) == 0o700
|
||||
assert stat.S_IMODE(pinned.stat().st_mode) == 0o700
|
||||
assert stat.S_IMODE(target_backup.stat().st_mode) == 0o700
|
||||
assert stat.S_IMODE((pinned / "history.jsonl").stat().st_mode) == 0o600
|
||||
assert stat.S_IMODE((report.backup_dir / "manifest.json").stat().st_mode) == 0o600
|
||||
|
||||
|
||||
def test_recovery_never_propagates_source_deletions(tmp_path: Path) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target_rule = target / "rules" / "user.rules"
|
||||
target_rule.parent.mkdir(parents=True)
|
||||
source.mkdir()
|
||||
target_rule.write_text("allow user setting\n", encoding="utf-8")
|
||||
(source / "history.jsonl").write_text('{"session_id":"source"}\n', encoding="utf-8")
|
||||
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
assert target_rule.read_text(encoding="utf-8") == "allow user setting\n"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("source_mtime", "target_mtime", "expected_token"),
|
||||
[(100, 200, "target-token"), (300, 200, "source-token"), (200, 200, "target-token")],
|
||||
)
|
||||
def test_recovery_uses_newest_credentials_with_target_winning_ties(
|
||||
tmp_path: Path,
|
||||
source_mtime: int,
|
||||
target_mtime: int,
|
||||
expected_token: str,
|
||||
) -> None:
|
||||
target = tmp_path / "codex"
|
||||
source = tmp_path / "headroom-codex-home-broken"
|
||||
target.mkdir()
|
||||
source.mkdir()
|
||||
target_auth = target / "auth.json"
|
||||
source_auth = source / "auth.json"
|
||||
target_auth.write_text('{"token":"target-token"}\n', encoding="utf-8")
|
||||
source_auth.write_text('{"token":"source-token"}\n', encoding="utf-8")
|
||||
os.utime(target_auth, ns=(target_mtime, target_mtime))
|
||||
os.utime(source_auth, ns=(source_mtime, source_mtime))
|
||||
|
||||
recover_codex_home(source=source, target=target)
|
||||
|
||||
assert json.loads(target_auth.read_text(encoding="utf-8"))["token"] == expected_token
|
||||
|
||||
|
||||
def test_recover_codex_cli_retains_distinct_backups_for_multiple_sources(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
target = tmp_path / "codex"
|
||||
first = tmp_path / "headroom-codex-home-first"
|
||||
second = tmp_path / "headroom-codex-home-second"
|
||||
first.mkdir()
|
||||
second.mkdir()
|
||||
(first / "history.jsonl").write_text('{"session_id":"first"}\n', encoding="utf-8")
|
||||
(second / "history.jsonl").write_text('{"session_id":"second"}\n', encoding="utf-8")
|
||||
|
||||
result = CliRunner().invoke(
|
||||
main,
|
||||
[
|
||||
"recover",
|
||||
"codex",
|
||||
"--source",
|
||||
str(first),
|
||||
"--source",
|
||||
str(second),
|
||||
"--target",
|
||||
str(target),
|
||||
"--yes",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert result.output.count("Recovery complete") == 2
|
||||
backup_root = tmp_path / ".headroom-codex-recovery"
|
||||
assert len([path for path in backup_root.iterdir() if path.is_dir()]) == 2
|
||||
assert [
|
||||
json.loads(line)["session_id"]
|
||||
for line in (target / "history.jsonl").read_text(encoding="utf-8").splitlines()
|
||||
] == ["first", "second"]
|
||||
|
|
@ -10,16 +10,19 @@ way a user would from the shell.
|
|||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
import socket
|
||||
import sqlite3
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import tomllib
|
||||
from click.testing import CliRunner
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
import tomllib
|
||||
else: # pragma: no cover - exercised in the Python 3.10 test job
|
||||
import tomli as tomllib
|
||||
|
||||
from headroom.cli import wrap as wrap_mod
|
||||
from headroom.cli.main import main
|
||||
from headroom.mcp_registry.install import build_headroom_spec
|
||||
|
|
@ -840,8 +843,6 @@ class TestInjectAvoidsDuplicateTopLevelKeys:
|
|||
def test_inject_does_not_create_duplicate_model_provider(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
import tomllib # Python 3.11+ stdlib
|
||||
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
config_dir = tmp_path / ".codex"
|
||||
config_dir.mkdir()
|
||||
|
|
@ -903,8 +904,6 @@ class TestInjectAvoidsDuplicateTopLevelKeys:
|
|||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""Idempotent re-wrap on a config that already has top-level keys."""
|
||||
import tomllib
|
||||
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
config_dir = tmp_path / ".codex"
|
||||
config_dir.mkdir()
|
||||
|
|
@ -939,8 +938,6 @@ class TestInjectAvoidsDuplicateTopLevelKeys:
|
|||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""Existing headroom provider table must not create duplicate TOML keys."""
|
||||
import tomllib # Python 3.11+ stdlib
|
||||
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
config_dir = tmp_path / ".codex"
|
||||
config_dir.mkdir()
|
||||
|
|
@ -1055,67 +1052,7 @@ def test_wrap_codex_prepare_only_respects_codex_home(
|
|||
assert not (tmp_path / ".codex" / "config.toml").exists()
|
||||
|
||||
|
||||
def test_codex_session_home_overlay_seeds_active_home_and_cleans_up(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
codex_home = tmp_path / "custom-codex-home"
|
||||
codex_home.mkdir()
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
|
||||
config_file = codex_home / "config.toml"
|
||||
auth_file = codex_home / "auth.json"
|
||||
original_config = '[profiles.default]\nmodel = "gpt-4o"\n'
|
||||
original_auth = '{"auth_mode": "apikey"}'
|
||||
config_file.write_text(original_config, encoding="utf-8")
|
||||
auth_file.write_text(original_auth, encoding="utf-8")
|
||||
|
||||
with wrap_mod._codex_session_home_overlay() as session_home:
|
||||
seeded_config = (session_home / "config.toml").read_text(encoding="utf-8")
|
||||
seeded_auth = (session_home / "auth.json").read_text(encoding="utf-8")
|
||||
assert seeded_config == original_config
|
||||
assert seeded_auth == original_auth
|
||||
(session_home / "config.toml").write_text('model_provider = "headroom"\n', encoding="utf-8")
|
||||
assert config_file.read_text(encoding="utf-8") == original_config
|
||||
|
||||
assert not session_home.exists()
|
||||
assert config_file.read_text(encoding="utf-8") == original_config
|
||||
assert auth_file.read_text(encoding="utf-8") == original_auth
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.platform == "win32" or not hasattr(socket, "AF_UNIX"),
|
||||
reason="requires POSIX Unix domain sockets",
|
||||
)
|
||||
def test_codex_session_home_overlay_skips_unix_sockets(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
codex_home = tmp_path / "custom-codex-home"
|
||||
socket_dir = codex_home / "vendor_imports" / "skills" / ".git"
|
||||
socket_dir.mkdir(parents=True)
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
monkeypatch.chdir(codex_home)
|
||||
|
||||
head_file = socket_dir / "HEAD"
|
||||
head_file.write_text("ref: refs/heads/main\n", encoding="utf-8")
|
||||
socket_file = socket_dir / "fsmonitor--daemon.ipc"
|
||||
relative_socket_file = socket_file.relative_to(codex_home)
|
||||
|
||||
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as fsmonitor_socket:
|
||||
fsmonitor_socket.bind(str(relative_socket_file))
|
||||
|
||||
with wrap_mod._codex_session_home_overlay() as session_home:
|
||||
session_socket_dir = session_home / socket_dir.relative_to(codex_home)
|
||||
assert (session_socket_dir / "HEAD").read_text(encoding="utf-8") == (
|
||||
"ref: refs/heads/main\n"
|
||||
)
|
||||
assert not (session_socket_dir / socket_file.name).exists()
|
||||
|
||||
assert socket_file.is_socket()
|
||||
|
||||
|
||||
def test_wrap_codex_launch_uses_session_scoped_codex_home(
|
||||
def test_wrap_codex_launch_uses_durable_codex_home(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
|
|
@ -1131,7 +1068,7 @@ def test_wrap_codex_launch_uses_session_scoped_codex_home(
|
|||
auth_file.write_text(original_auth, encoding="utf-8")
|
||||
|
||||
launch_env: dict[str, str] = {}
|
||||
session_home_seen: list[Path] = []
|
||||
rollout = codex_home / "sessions" / "2026" / "07" / "14" / "rollout-thread.jsonl"
|
||||
|
||||
def fake_launch(
|
||||
*,
|
||||
|
|
@ -1147,15 +1084,9 @@ def test_wrap_codex_launch_uses_session_scoped_codex_home(
|
|||
del args, port, no_proxy, tool_label, env_vars_display, kwargs
|
||||
assert binary == "/fake/codex"
|
||||
launch_env.update(env)
|
||||
session_home = Path(env["CODEX_HOME"])
|
||||
session_home_seen.append(session_home)
|
||||
assert session_home.exists()
|
||||
seeded_config = (session_home / "config.toml").read_text(encoding="utf-8")
|
||||
assert original_config in seeded_config
|
||||
assert 'model_provider = "headroom"' in seeded_config
|
||||
assert 'base_url = "http://127.0.0.1:8787/v1"' in seeded_config
|
||||
assert "[mcp_servers.headroom]" in seeded_config
|
||||
assert (session_home / "auth.json").read_text(encoding="utf-8") == original_auth
|
||||
assert Path(env["CODEX_HOME"]) == codex_home
|
||||
rollout.parent.mkdir(parents=True)
|
||||
rollout.write_text('{"type":"session_meta"}\n', encoding="utf-8")
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
with patch(
|
||||
|
|
@ -1176,48 +1107,94 @@ def test_wrap_codex_launch_uses_session_scoped_codex_home(
|
|||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert session_home_seen
|
||||
assert launch_env["CODEX_HOME"] == str(session_home_seen[0])
|
||||
assert launch_env["CODEX_HOME"] == str(codex_home)
|
||||
assert launch_env["OPENAI_BASE_URL"] == "http://127.0.0.1:8787/v1"
|
||||
assert config_file.read_text(encoding="utf-8") == original_config
|
||||
persisted_config = config_file.read_text(encoding="utf-8")
|
||||
assert original_config in persisted_config
|
||||
assert "[mcp_servers.headroom]" in persisted_config
|
||||
assert 'model_provider = "headroom"' not in persisted_config
|
||||
assert "[model_providers.headroom]" not in persisted_config
|
||||
assert auth_file.read_text(encoding="utf-8") == original_auth
|
||||
assert not session_home_seen[0].exists()
|
||||
assert rollout.read_text(encoding="utf-8") == '{"type":"session_meta"}\n'
|
||||
|
||||
|
||||
def test_wrap_codex_launches_use_distinct_session_homes_per_port(
|
||||
def test_codex_session_launch_settings_keep_routing_process_local(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
codex_home = tmp_path / "custom-codex-home"
|
||||
codex_home.mkdir()
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
monkeypatch.setattr(wrap_mod, "_project_name_from_cwd", lambda: None)
|
||||
config_file = codex_home / "config.toml"
|
||||
original_config = 'model = "gpt-5"\n'
|
||||
config_file.write_text(original_config, encoding="utf-8")
|
||||
|
||||
args, env, display = wrap_mod._codex_session_launch_settings(
|
||||
port=9898,
|
||||
codex_args=("exec", "hello"),
|
||||
environ={"CODEX_HOME": str(codex_home)},
|
||||
)
|
||||
|
||||
assert args == (
|
||||
"--config",
|
||||
'openai_base_url="http://127.0.0.1:9898/v1"',
|
||||
"exec",
|
||||
"hello",
|
||||
)
|
||||
assert env["CODEX_HOME"] == str(codex_home)
|
||||
assert env["OPENAI_BASE_URL"] == "http://127.0.0.1:9898/v1"
|
||||
assert display == ["OPENAI_BASE_URL=http://127.0.0.1:9898/v1"]
|
||||
assert config_file.read_text(encoding="utf-8") == original_config
|
||||
|
||||
|
||||
def test_codex_session_launch_settings_preserve_custom_provider_identity(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
codex_home = tmp_path / "custom-codex-home"
|
||||
codex_home.mkdir()
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
monkeypatch.setattr(wrap_mod, "_project_name_from_cwd", lambda: None)
|
||||
config_file = codex_home / "config.toml"
|
||||
original_config = (
|
||||
'[profiles.work]\nmodel_provider = "company"\n\n'
|
||||
'[model_providers.company]\nbase_url = "https://api.example.test/v1"\n'
|
||||
)
|
||||
config_file.write_text(original_config, encoding="utf-8")
|
||||
|
||||
args, env, _ = wrap_mod._codex_session_launch_settings(
|
||||
port=9898,
|
||||
codex_args=("--profile", "work"),
|
||||
environ={"CODEX_HOME": str(codex_home)},
|
||||
)
|
||||
|
||||
assert "model_provider=headroom" not in " ".join(args)
|
||||
assert '"model_providers"."company"."base_url"="http://127.0.0.1:9898/v1"' in args
|
||||
assert env[wrap_mod._UPSTREAM_BASE_URL_ENV_VAR] == "https://api.example.test/v1"
|
||||
assert config_file.read_text(encoding="utf-8") == original_config
|
||||
|
||||
|
||||
def test_wrap_codex_rejects_custom_provider_without_upstream_base_url(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
codex_home = tmp_path / "custom-codex-home"
|
||||
codex_home.mkdir()
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
(codex_home / "config.toml").write_text(
|
||||
'model_provider = "company"\n[model_providers.company]\nname = "Company"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
config_file = codex_home / "config.toml"
|
||||
auth_file = codex_home / "auth.json"
|
||||
original_config = '[profiles.default]\nmodel = "gpt-4o"\n'
|
||||
original_auth = '{"auth_mode": "apikey"}'
|
||||
config_file.write_text(original_config, encoding="utf-8")
|
||||
auth_file.write_text(original_auth, encoding="utf-8")
|
||||
|
||||
launch_records: list[tuple[int, Path, str]] = []
|
||||
|
||||
def fake_launch(
|
||||
*,
|
||||
binary: str,
|
||||
args: tuple,
|
||||
env: dict[str, str],
|
||||
port: int,
|
||||
no_proxy: bool,
|
||||
tool_label: str,
|
||||
env_vars_display: list[str],
|
||||
**kwargs: object,
|
||||
) -> None:
|
||||
del args, no_proxy, tool_label, env_vars_display, kwargs
|
||||
assert binary == "/fake/codex"
|
||||
session_home = Path(env["CODEX_HOME"])
|
||||
assert session_home.exists()
|
||||
launch_records.append(
|
||||
(port, session_home, (session_home / "config.toml").read_text(encoding="utf-8"))
|
||||
def fake_launch(**kwargs: object) -> None:
|
||||
configure_launch = kwargs["configure_launch"]
|
||||
assert callable(configure_launch)
|
||||
configure_launch(
|
||||
8787,
|
||||
kwargs["args"],
|
||||
kwargs["env"],
|
||||
kwargs["env_vars_display"],
|
||||
)
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
|
|
@ -1226,25 +1203,72 @@ def test_wrap_codex_launches_use_distinct_session_homes_per_port(
|
|||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch):
|
||||
first = runner.invoke(
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["wrap", "codex", "--port", "8787", "--no-tokensave", "--no-serena"],
|
||||
)
|
||||
second = runner.invoke(
|
||||
main,
|
||||
["wrap", "codex", "--port", "9898", "--no-tokensave", "--no-serena"],
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--port",
|
||||
"8787",
|
||||
"--no-rtk",
|
||||
"--no-mcp",
|
||||
"--no-tokensave",
|
||||
"--no-serena",
|
||||
],
|
||||
)
|
||||
|
||||
assert first.exit_code == 0, first.output
|
||||
assert second.exit_code == 0, second.output
|
||||
assert len(launch_records) == 2
|
||||
assert launch_records[0][1] != launch_records[1][1]
|
||||
assert 'base_url = "http://127.0.0.1:8787/v1"' in launch_records[0][2]
|
||||
assert 'base_url = "http://127.0.0.1:9898/v1"' in launch_records[1][2]
|
||||
assert config_file.read_text(encoding="utf-8") == original_config
|
||||
assert auth_file.read_text(encoding="utf-8") == original_auth
|
||||
assert not launch_records[0][1].exists()
|
||||
assert not launch_records[1][1].exists()
|
||||
assert result.exit_code != 0
|
||||
assert "custom provider 'company' has no upstream base_url" in result.output
|
||||
|
||||
|
||||
def test_wrap_codex_routes_model_provider_selected_by_config_argument(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
codex_home = tmp_path / "custom-codex-home"
|
||||
codex_home.mkdir()
|
||||
monkeypatch.setenv("CODEX_HOME", str(codex_home))
|
||||
monkeypatch.setattr(wrap_mod, "_project_name_from_cwd", lambda: None)
|
||||
(codex_home / "config.toml").write_text(
|
||||
'[model_providers.company]\nbase_url = "https://api.example.test/v1"\n',
|
||||
encoding="utf-8",
|
||||
)
|
||||
configured_env: dict[str, str] = {}
|
||||
|
||||
def fake_launch(**kwargs: object) -> None:
|
||||
configure_launch = kwargs["configure_launch"]
|
||||
assert callable(configure_launch)
|
||||
_, env, _ = configure_launch(
|
||||
8787,
|
||||
kwargs["args"],
|
||||
kwargs["env"],
|
||||
kwargs["env_vars_display"],
|
||||
)
|
||||
configured_env.update(env)
|
||||
|
||||
with patch("headroom.cli.wrap._ensure_rtk_binary", return_value=None):
|
||||
with patch(
|
||||
"headroom.cli.wrap.shutil.which",
|
||||
side_effect=lambda cmd: "/fake/codex" if cmd == "codex" else None,
|
||||
):
|
||||
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
[
|
||||
"wrap",
|
||||
"codex",
|
||||
"--no-rtk",
|
||||
"--no-mcp",
|
||||
"--no-tokensave",
|
||||
"--no-serena",
|
||||
"--",
|
||||
"--config",
|
||||
'model_provider="company"',
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert configured_env[wrap_mod._UPSTREAM_BASE_URL_ENV_VAR] == ("https://api.example.test/v1")
|
||||
|
||||
|
||||
def test_wrap_codex_injects_rtk_globally_without_changing_project_agents(
|
||||
|
|
@ -1968,22 +1992,28 @@ class TestCodexLaunchExportsCustomUpstream:
|
|||
the wrong host (regression of #1614)."""
|
||||
|
||||
def _launch_env(self, monkeypatch, tmp_path, *, custom_upstream):
|
||||
from contextlib import contextmanager
|
||||
|
||||
captured: dict = {}
|
||||
|
||||
monkeypatch.setattr(wrap_mod.shutil, "which", lambda name: "/usr/bin/codex")
|
||||
monkeypatch.setattr(wrap_mod, "_codex_home_dir", lambda: tmp_path)
|
||||
monkeypatch.setattr(wrap_mod, "_offer_dangling_codex_recovery", lambda active_home: None)
|
||||
monkeypatch.setattr(wrap_mod, "_prepare_codex_wrap_state", lambda **kwargs: None)
|
||||
if custom_upstream:
|
||||
(tmp_path / "config.toml").write_text(
|
||||
"\n".join(
|
||||
(
|
||||
'model_provider = "gateway"',
|
||||
"[model_providers.gateway]",
|
||||
f'base_url = "{custom_upstream}"',
|
||||
)
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
@contextmanager
|
||||
def _fake_overlay():
|
||||
yield tmp_path / "session"
|
||||
|
||||
monkeypatch.setattr(wrap_mod, "_codex_session_home_overlay", _fake_overlay)
|
||||
# Stand in for the heavy prepare step; only its return value matters here.
|
||||
monkeypatch.setattr(wrap_mod, "_prepare_codex_wrap_state", lambda **kwargs: custom_upstream)
|
||||
|
||||
def _fake_launch(*, env, **kwargs):
|
||||
def _fake_launch(*, env, port, configure_launch, args=(), env_vars_display=(), **kwargs):
|
||||
if configure_launch is not None:
|
||||
_args, env, _display = configure_launch(port, args, env, list(env_vars_display))
|
||||
captured["env"] = env
|
||||
|
||||
monkeypatch.setattr(wrap_mod, "_launch_tool", _fake_launch)
|
||||
|
|
|
|||
11
uv.lock
generated
11
uv.lock
generated
|
|
@ -1469,6 +1469,7 @@ dependencies = [
|
|||
{ name = "rich" },
|
||||
{ name = "tiktoken" },
|
||||
{ name = "tomli", marker = "python_full_version < '3.11'" },
|
||||
{ name = "tomlkit" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
|
|
@ -1781,6 +1782,7 @@ requires-dist = [
|
|||
{ name = "strands-agents", marker = "extra == 'strands'", specifier = ">=0.1.0" },
|
||||
{ name = "tiktoken", specifier = ">=0.5.0" },
|
||||
{ name = "tomli", marker = "python_full_version < '3.11'", specifier = ">=2.0.0" },
|
||||
{ name = "tomlkit", specifier = ">=0.13.0,<1.0" },
|
||||
{ name = "torch", marker = "sys_platform == 'darwin' and extra == 'pytorch-mps'", specifier = ">=2.12.1" },
|
||||
{ name = "torch", marker = "(platform_machine != 'x86_64' and extra == 'ml') or (sys_platform != 'darwin' and extra == 'ml')", specifier = ">=2.12.1" },
|
||||
{ name = "torch", marker = "(platform_machine != 'x86_64' and extra == 'voice') or (sys_platform != 'darwin' and extra == 'voice')", specifier = ">=2.12.1" },
|
||||
|
|
@ -5923,6 +5925,15 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/23/d1/136eb2cb77520a31e1f64cbae9d33ec6df0d78bdf4160398e86eec8a8754/tomli-2.4.0-py3-none-any.whl", hash = "sha256:1f776e7d669ebceb01dee46484485f43a4048746235e683bcdffacdf1fb4785a", size = 14477, upload-time = "2026-01-11T11:22:37.446Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tomlkit"
|
||||
version = "0.15.0"
|
||||
source = { registry = "https://pypi.org/simple/" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/51/db/03eaf4331631ef6b27d6e3c9b68c54dc6f0d63d87201fed600cc409307fd/tomlkit-0.15.0.tar.gz", hash = "sha256:7d1a9ecba3086638211b13814ea79c90dd54dd11993564376f3aa92271f5c7a3", size = 161875, upload-time = "2026-05-10T07:38:22.245Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/6a/43/8bd850ee71a191bf072e31302c73a66be413fecdd98fdcd111ecbcce13ca/tomlkit-0.15.0-py3-none-any.whl", hash = "sha256:4dbc8f0fc024412b57ced8757ac7461305126a648ff8c2c807fcb8e133a78738", size = 41328, upload-time = "2026-05-10T07:38:23.517Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "torch"
|
||||
version = "2.13.0"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue