Merge branch 'main' into pr/3208-launchd-exec-handoff

This commit is contained in:
JD Davis 2026-08-26 12:43:04 -05:00 committed by GitHub
commit 0b28c58f77
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
58 changed files with 5273 additions and 409 deletions

View file

@ -187,6 +187,8 @@ Install Headroom so it's globally on PATH — `uv tool install "headroom-ai[mcp]
## Architecture
For user-managed Serena drift, run `headroom mcp reconcile` to inspect the current recommendation. Add `--adopt` only when you want Headroom to replace the Serena entry.
### MCP only (no proxy)
The LLM calls `headroom_compress` on demand. Compression happens locally in the MCP process. Originals are stored in a local `CompressionStore` with 1-hour TTL.

View file

@ -223,6 +223,18 @@ See [issue #746](https://github.com/headroomlabs-ai/headroom/issues/746) for the
`ENABLE_TOOL_SEARCH` is unaffected and can stay enabled for context-window savings while routing through Headroom.
## Server-managed settings unavailable through custom ANTHROPIC_BASE_URL
**Symptom**: Settings pushed from **Admin Settings > Claude Code > Managed settings** in the claude.ai console (server-managed settings) don't apply to sessions running through Headroom, even though they apply fine without the proxy.
**Cause**: This is a Claude-side gate, not a Headroom limitation. Per Anthropic's docs, server-managed settings require a direct connection to `api.anthropic.com`; if `ANTHROPIC_BASE_URL` is set to any non-default host — which is exactly what wrapping via Headroom does — Claude Code skips the settings fetch entirely for that session. The request never reaches Headroom, so there is no endpoint for Headroom to implement or proxy.
This is separate from the OS-level `managed-settings.json` file (macOS `/Library/Application Support/ClaudeCode/`, Linux `/etc/claude-code/`, Windows `C:\Program Files\ClaudeCode\`): that file is read straight from local disk at startup and is unaffected by `ANTHROPIC_BASE_URL` or Headroom. If that file isn't taking effect, the cause is unrelated to proxying (path, permissions, or JSON syntax) — check `claude --debug-file <path>` and search the log for `Remote settings`.
**Fix**: None available on the Headroom side — this is an intentional Anthropic security boundary (a proxy in the path could otherwise forge org policy). If your org relies on server-managed settings, deploy the same policy as [endpoint-managed settings](https://code.claude.com/docs/en/settings#settings-files) (MDM profile, Windows registry, or a local `managed-settings.json`) instead, since those are read locally and unaffected by proxying.
See [Server-managed settings platform availability](https://code.claude.com/docs/en/server-managed-settings#platform-availability) and [issue #3074](https://github.com/headroomlabs-ai/headroom/issues/3074).
## Compression Too Aggressive
**Symptom**: LLM responses are missing information that was in tool outputs.

View file

@ -123,6 +123,12 @@ class CompressionCache:
# `RLock` (not `Lock`) so future code can call locked methods from
# inside another locked method without self-deadlock.
self._lock = threading.RLock()
# Serializes one sidecar-mode compress turn per session (pre-work,
# pipeline, post-work run as one block on an executor thread). The
# sidecar contract is sequential turns per conversation; this lock
# keeps a contract-violating concurrent pair from interleaving and
# tearing the tracker's prev-original/prev-returned snapshots.
self.session_turn_lock = threading.Lock()
self._cache: OrderedDict[str, _CacheEntry] = OrderedDict()
# `_stable_hashes` is CONTENT-KEYED, not positional. It records "we
# have seen this content before and it is known not to compress

View file

@ -965,6 +965,26 @@ class PrefixCacheTracker:
def get_last_forwarded_messages(self) -> list[dict[str, Any]]:
return copy.deepcopy(self._last_forwarded_messages)
def record_returned(
self,
original_messages: list[dict[str, Any]],
returned_messages: list[dict[str, Any]],
) -> None:
"""Record the compressed form handed back to a compress-only caller.
Sidecar mode (session-aware ``/v1/compress``): Headroom does not
forward upstream, but whatever it RETURNS is what the caller forwards
the same fact ``update_from_response`` records in proxy mode, just
captured at return time instead of send time. Only the transcript
snapshots and the activity clock move here; frozen-prefix counts are
left untouched because no provider response has confirmed anything
yet they advance when the caller relays usage via ``/v1/usage``
(``update_from_response``), or stay at their conservative local value.
"""
self._last_activity = time.time()
self._last_original_messages = copy.deepcopy(original_messages)
self._last_forwarded_messages = copy.deepcopy(returned_messages)
def resolved_cache_ttl_seconds(self) -> int:
"""Effective prompt-cache lifetime for this session's provider."""
if self.config.cache_ttl_seconds is not None:
@ -1249,6 +1269,26 @@ class SessionTrackerStore:
self._lineage_affinities: dict[str, str | None] = {}
self._lineage_counter = itertools.count(1)
def peek(self, session_id: str) -> PrefixCacheTracker | None:
"""Return the live tracker for ``session_id``, else None.
Never creates: lookup paths that must not leave a footprint (e.g. the
``/v1/usage`` unknown-session check, where ``get_or_create`` would let
a flood of novel ids grow the store unboundedly within each TTL
window) use this instead of :meth:`get_or_create`.
A TTL-expired-but-unswept tracker answers None too: the sweep runs
lazily from get_or_create at 60s granularity, so without this check an
expired session would keep answering with stale pre-expiry state and
a caller that then touched it (``update_from_response`` stamps
``_last_activity``) would resurrect the dead tracker indefinitely,
making the documented 404-on-expired contract nondeterministic.
"""
tracker = self._trackers.get(session_id)
if tracker is None or tracker.is_expired:
return None
return tracker
def get_or_create(self, session_id: str, provider: str) -> PrefixCacheTracker:
"""Get existing tracker or create a new one for this session."""
self._maybe_cleanup()

View file

@ -160,8 +160,18 @@ class SemanticCache:
self._hits += 1
return entry
# Try semantic similarity if we have embedding function
if self._embedding_fn:
# Try semantic similarity if we have embedding function.
#
# Only for a NON-EMPTY query: the query is the last user message, and in
# agent/tool traffic the overwhelming majority of turns are tool_result
# continuations whose extracted query is "" (no text block). Embedding
# matching on "" makes every such turn ~identical to every other (a real
# sentence embedder maps "" to a fixed non-zero vector), so an empty
# query would false-hit and serve one conversation's response to an
# unrelated one — precisely the cross-context collision the messages_hash
# key is chosen to avoid. An empty query may still hit via the exact
# messages_hash above, which is context-complete and safe.
if self._embedding_fn and query.strip():
query_embedding = self._embedding_fn(query)
best_match, best_similarity = self._find_similar(query_embedding)
@ -208,9 +218,13 @@ class SemanticCache:
while key not in self._cache and len(self._cache) >= self.config.max_entries:
self._evict_oldest()
# Generate embedding if available
# Generate embedding if available — but never for an empty/blank query.
# A stored empty-query entry with an embedding would be a false-match
# target for the semantic get() path; leaving its embedding empty makes
# _find_similar skip it (it ignores entries with no embedding), so an
# empty-query entry is reachable only by its exact messages_hash.
embedding: list[float] = []
if self._embedding_fn:
if self._embedding_fn and query.strip():
embedding = self._embedding_fn(query)
now = time.time()

View file

@ -225,6 +225,7 @@ def learn(
total_projects = 0
total_failures = 0
total_recommendations = 0
total_analysis_failures = 0
matched_projects = 0
available_projects: list[tuple[str, Path]] = []
@ -299,6 +300,12 @@ def learn(
f"Failures: {result_data.total_failures} ({result_data.failure_rate:.1%})"
)
analysis_error = getattr(result_data, "analysis_error", None)
if analysis_error:
total_analysis_failures += 1
click.echo(f" Analysis failed: {analysis_error}", err=True)
continue
if result_data.failure_rate == 0 and not result_data.recommendations:
click.echo(" No failures or patterns found.")
continue
@ -350,6 +357,9 @@ def learn(
f"{total_recommendations} recommendations"
)
if total_analysis_failures:
raise SystemExit(1)
def _make_llm_judge(model: str) -> Any:
"""Build an LLM judge callable for verbosity, or None if unavailable.

View file

@ -216,6 +216,52 @@ def mcp_uninstall() -> None:
click.echo("Headroom MCP is not configured. Nothing to uninstall.")
@mcp.command("reconcile")
@click.option("--adopt", is_flag=True, help="Replace only the Serena entry with Headroom's spec.")
def mcp_reconcile(adopt: bool) -> None:
"""Inspect or explicitly reconcile a user-managed Serena MCP entry."""
from headroom.mcp_registry import (
CLAUDE_SERENA_CONTEXT,
ClaudeConfigMutationError,
ClaudeRegistrar,
RegisterStatus,
build_serena_spec,
)
from headroom.mcp_registry.ledger import (
LedgerMutationError,
record_install,
validate_ledger_for_mutation,
)
registrar = ClaudeRegistrar()
if not registrar.detect():
raise click.ClickException("claude is not detected")
recommended = build_serena_spec(CLAUDE_SERENA_CONTEXT)
observed = registrar.get_server("serena")
if adopt:
try:
registrar.validate_configs_for_mutation()
validate_ledger_for_mutation()
except (ClaudeConfigMutationError, LedgerMutationError) as exc:
raise click.ClickException(str(exc)) from exc
if adopt:
result = registrar.register_server(recommended, force=True)
if result.status not in (RegisterStatus.REGISTERED, RegisterStatus.ALREADY):
raise click.ClickException(result.detail or "could not adopt Serena configuration")
record_install("claude", recommended)
click.echo(
"Adopted Headroom's Serena configuration for Claude; unrelated config preserved."
)
return
click.echo("Serena reconciliation for Claude")
click.echo(f" observed: {'absent' if observed is None else 'present'}")
click.echo(f" recommendation: {recommended.command} {' '.join(recommended.args)}")
if observed is not None and observed != recommended:
click.echo(" action: use --adopt to replace it")
@mcp.command("status")
def mcp_status() -> None:
"""Check Headroom MCP configuration status.

View file

@ -32,11 +32,11 @@ import subprocess
import sys
import time
import urllib.parse
from collections.abc import Callable
from collections.abc import Callable, Mapping
from contextlib import contextmanager
from functools import wraps
from pathlib import Path
from typing import Any, cast
from typing import Any, NamedTuple, cast
from headroom._subprocess import pid_alive, run
@ -1192,6 +1192,235 @@ def _wrap_marker_path(settings_path: Path) -> Path:
return settings_path.parent / ".headroom_wrap_marker.json"
def _wrap_owners_path(settings_path: Path) -> Path:
"""Sidecar recording which live wrap sessions own each settings env key.
Separate from ``.headroom_wrap_marker.json`` on purpose: that marker
describes a single writer and is consumed by doctor, unwrap and the
staleness checks. Concurrency ownership is additive state, so it lives in
its own file rather than changing a shape those readers depend on.
"""
return settings_path.parent / ".headroom_wrap_owners.json"
def _wrap_settings_lock(settings_path: Path) -> Any:
"""Serialize settings read-modify-write across concurrent wrap sessions.
Writing the proxy URL into ``settings.local.json`` is a read-modify-write,
and several ``headroom wrap`` sessions in one project run it concurrently.
The write itself is atomic, so the file never tears -- but without this the
updates are still lost against each other (#3205).
"""
from contextlib import nullcontext
lock_path = settings_path.parent / ".headroom_wrap_settings.lock"
try:
lock_path.parent.mkdir(parents=True, exist_ok=True)
lock_file = open(lock_path, "a+b") # noqa: SIM115
except OSError:
# Matches _proxy_start_lock: a workspace that cannot hold lock state is
# degraded, not unusable.
return nullcontext()
return _locked_file(lock_file)
@contextmanager
def _locked_file(lock_file: Any) -> Any:
"""Hold an exclusive OS lock on an already-open file for the block.
Shared by ``_proxy_start_lock`` and ``_wrap_settings_lock`` -- the two
differ only in which file they lock, and an OS-lock dance duplicated per
call site is one place for the platform branches to drift apart.
"""
with lock_file:
if sys.platform == "win32":
import msvcrt
# msvcrt.locking operates on bytes from the current file position.
lock_file.seek(0)
if lock_file.read(1) == b"":
lock_file.seek(0)
lock_file.write(b"0")
lock_file.flush()
lock_file.seek(0)
# LK_LOCK has implementation-dependent retry limits, and a holder
# may legitimately take longer than that (a proxy loading ML
# components), so use the non-blocking primitive in a loop.
while True:
try:
msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1)
break
except OSError:
time.sleep(0.05)
try:
yield
finally:
lock_file.seek(0)
msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1)
else:
import fcntl
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
try:
yield
finally:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
def _read_wrap_owners(settings_path: Path) -> dict[str, Any]:
try:
rec = json.loads(_read_text(_wrap_owners_path(settings_path)))
except (OSError, ValueError):
return {}
return rec if isinstance(rec, dict) else {}
def _write_wrap_owners(settings_path: Path, owners: dict[str, Any]) -> None:
target = _wrap_owners_path(settings_path)
try:
if not owners:
target.unlink(missing_ok=True)
return
_write_text(target, json.dumps(owners, indent=2) + "\n")
except OSError:
pass
def _live_holders(entry: Any, *, dead_ports: frozenset[int] = frozenset()) -> list[dict[str, Any]]:
"""Holders in *entry* whose process is still provably alive.
Reuses the same conservative liveness the proxy-client markers use: a PID
that is gone, or that is now provably a different process, is dropped. Any
uncertainty keeps the holder, because dropping a live owner is what causes
a running session to be unrouted.
``dead_ports`` additionally drops holders whose proxy port the caller has
*proven* dead. A wrapper process outlives its proxy after a hard reboot or
SIGKILL of the proxy alone, and such a holder routes nothing; left in place
it would block the #2221 self-heal from clearing a base_url that now points
at nothing.
"""
if not isinstance(entry, dict):
return []
holders = entry.get("holders")
if not isinstance(holders, list):
return []
live: list[dict[str, Any]] = []
for holder in holders:
if not isinstance(holder, dict):
continue
pid = holder.get("pid")
if not isinstance(pid, int) or not _pid_alive(pid):
continue
if _identity_mismatch(holder.get("start_src"), holder.get("start_time"), pid):
continue
port = holder.get("port")
if isinstance(port, int) and port in dead_ports:
continue
live.append(holder)
return live
def _self_holder(port: int | None) -> dict[str, Any]:
ident = _proc_identity(os.getpid())
return {
"pid": os.getpid(),
"start_src": ident[0] if ident else None,
"start_time": ident[1] if ident else None,
"port": port,
}
def _claim_wrap_key(
settings_path: Path,
key: str,
current_value: str | None,
*,
port: int | None = None,
) -> None:
"""Register this process as an owner of *key*, recording the true original.
The first live owner records ``original``; later owners inherit it and are
flagged ``inherited`` so their exit knows the value they happened to
observe was not the pre-wrap one. Without that, a second wrap session
captures the *first session's* proxy URL as the value to restore, and puts
a dead proxy back into the file on exit (#3205).
"""
owners = _read_wrap_owners(settings_path)
entry = owners.get(key)
live = _live_holders(entry)
inherited = bool(live) and isinstance(entry, dict) and "original" in entry
original = entry.get("original") if inherited and isinstance(entry, dict) else current_value
me = _self_holder(port)
me["inherited"] = inherited
live = [h for h in live if h.get("pid") != me["pid"]]
live.append(me)
owners[key] = {"original": original, "holders": live}
_write_wrap_owners(settings_path, owners)
class _KeyRelease(NamedTuple):
"""Outcome of dropping this process's claim on a settings env key."""
should_restore: bool
original: str | None
trust_caller: bool
survivor: dict[str, Any] | None
def _release_wrap_key(
settings_path: Path,
key: str,
*,
force: bool = False,
dead_ports: frozenset[int] = frozenset(),
) -> _KeyRelease:
"""Drop this process's claim on *key*.
``should_restore`` is False while another live wrap session still owns the
key -- restoring then silently unroutes a running session. ``force`` is for
``unwrap``, where the user is explicitly asking for their settings back:
every claim is dropped and the restore happens regardless.
``trust_caller`` says whether the caller's remembered ``previous`` is its
own first-hand observation of the pre-wrap value. True when there is no
owner record at all (unwrap of a pre-upgrade session, and the legacy
callers that pass the value directly), and when this process founded the
record. False for an inheriting holder -- it remembers the *first
session's* proxy URL, so honouring it writes a dead proxy back, the exact
bug #3205 is about -- and false for a caller with no claim of its own,
whose marker-derived value is second-hand where the record is not.
``survivor`` is a still-live holder the caller can re-point the
single-slot wrap marker at, so an exiting session does not take the
surviving one's #2221 self-heal record with it.
"""
owners = _read_wrap_owners(settings_path)
entry = owners.get(key)
if not isinstance(entry, dict):
return _KeyRelease(True, None, True, None)
me = os.getpid()
remaining = [h for h in _live_holders(entry, dead_ports=dead_ports) if h.get("pid") != me]
original = entry.get("original")
# Look this process's own claim up in the raw holder list, never the
# liveness-filtered one: the caller is by definition running, and its claim
# is what says whether the value it remembers is first-hand.
raw = entry.get("holders")
mine = (
next((h for h in raw if isinstance(h, dict) and h.get("pid") == me), None)
if isinstance(raw, list)
else None
)
trust_caller = mine is not None and not mine.get("inherited")
if remaining and not force:
owners[key] = {"original": original, "holders": remaining}
_write_wrap_owners(settings_path, owners)
return _KeyRelease(False, original, trust_caller, remaining[0])
owners.pop(key, None)
_write_wrap_owners(settings_path, owners)
return _KeyRelease(True, original, trust_caller, None)
def _write_wrap_marker(settings_path: Path, *, port: int, key: str, previous: str | None) -> None:
"""Best-effort record of which (pid, port, key) wrote the base_url entry.
@ -1214,6 +1443,53 @@ def _write_wrap_marker(settings_path: Path, *, port: int, key: str, previous: st
pass
def _rehome_wrap_marker(
settings_path: Path,
*,
key: str,
survivor: dict[str, Any] | None,
original: str | None,
) -> None:
"""Hand this session's wrap marker to a session that is still running.
The marker has one slot and the last writer wins it. When that writer exits
while a sibling still owns the key, leaving the marker describes a dead
process, and deleting it strips the survivor of the #2221 dead-proxy
self-heal record. Rewrite it to describe the survivor instead, carrying the
owner record's ``original`` as the value to restore -- the marker's own
``previous`` may be an earlier session's proxy URL (#3205).
Only ever touches a marker this process wrote; a sibling's marker is
already accurate.
"""
marker_path = _wrap_marker_path(settings_path)
marker = _read_wrap_marker(settings_path)
if marker is None or marker.get("key") != key or marker.get("pid") != os.getpid():
return
port = survivor.get("port") if survivor is not None else None
try:
if survivor is None or not isinstance(port, int):
# No survivor to hand it to, or one whose port we never recorded:
# a marker without a usable port is worse than none.
marker_path.unlink(missing_ok=True)
return
_write_text(
marker_path,
json.dumps(
{
"pid": survivor.get("pid"),
"start_src": survivor.get("start_src"),
"start_time": survivor.get("start_time"),
"port": port,
"key": key,
"previous": original,
}
),
)
except OSError:
pass
def _read_wrap_marker(settings_path: Path) -> dict[str, Any] | None:
marker = _wrap_marker_path(settings_path)
try:
@ -1337,7 +1613,15 @@ def _check_and_clear_dead_wrap_marker(settings_path: Path, *, key: str) -> str |
f"running (issue #2221); restoring prior value",
err=True,
)
_restore_claude_wrap_base_url(previous, settings_path=settings_path, _key_override=key)
_restore_claude_wrap_base_url(
previous,
settings_path=settings_path,
_key_override=key,
# The wrapper process can outlive its proxy (the proxy alone was
# SIGKILLed). Its ownership claim would otherwise veto this restore and
# leave the base_url pointing at a port proven dead just above (#3205).
dead_ports=frozenset({port}) if isinstance(port, int) else frozenset(),
)
return previous
@ -1503,16 +1787,21 @@ def _write_claude_wrap_base_url(
detected and self-healed (issue #1768).
"""
path = settings_path or (Path.cwd() / ".claude" / "settings.local.json")
payload = _read_settings_for_write(path)
env_map = dict(payload.get("env") or {}) if isinstance(payload.get("env"), dict) else {}
key = _claude_wrap_base_url_env_key(foundry_mode=foundry_mode, vertex_mode=vertex_mode)
previous = env_map.get(key)
env_map[key] = proxy_url
payload["env"] = env_map
path.parent.mkdir(parents=True, exist_ok=True)
_write_text(path, json.dumps(payload, indent=2) + "\n")
if port is not None:
_write_wrap_marker(path, port=port, key=key, previous=previous)
with _wrap_settings_lock(path):
payload = _read_settings_for_write(path)
env_map = dict(payload.get("env") or {}) if isinstance(payload.get("env"), dict) else {}
previous = env_map.get(key)
# Claim before writing, so the recorded original is the value that was
# there before *any* wrap session touched it -- not the previous
# session's proxy URL (#3205).
_claim_wrap_key(path, key, previous, port=port)
env_map[key] = proxy_url
payload["env"] = env_map
_write_text(path, json.dumps(payload, indent=2) + "\n")
if port is not None:
_write_wrap_marker(path, port=port, key=key, previous=previous)
return previous
@ -1525,13 +1814,15 @@ def _write_claude_wrap_tool_search(value: str, *, settings_path: Path | None = N
process, and is restored transactionally when the wrap session exits.
"""
path = settings_path or (Path.cwd() / ".claude" / "settings.local.json")
payload = _read_settings_for_write(path)
env_map = dict(payload.get("env") or {}) if isinstance(payload.get("env"), dict) else {}
previous = env_map.get(_TOOL_SEARCH_ENV)
env_map[_TOOL_SEARCH_ENV] = value
payload["env"] = env_map
path.parent.mkdir(parents=True, exist_ok=True)
_write_text(path, json.dumps(payload, indent=2) + "\n")
with _wrap_settings_lock(path):
payload = _read_settings_for_write(path)
env_map = dict(payload.get("env") or {}) if isinstance(payload.get("env"), dict) else {}
previous = env_map.get(_TOOL_SEARCH_ENV)
_claim_wrap_key(path, _TOOL_SEARCH_ENV, previous)
env_map[_TOOL_SEARCH_ENV] = value
payload["env"] = env_map
_write_text(path, json.dumps(payload, indent=2) + "\n")
return previous
@ -1553,6 +1844,8 @@ def _restore_claude_wrap_base_url(
vertex_mode: bool = False,
settings_path: Path | None = None,
_key_override: str | None = None,
force: bool = False,
dead_ports: frozenset[int] = frozenset(),
) -> None:
"""Restore (or remove) the env key written by _write_claude_wrap_base_url.
@ -1561,40 +1854,63 @@ def _restore_claude_wrap_base_url(
``previous`` is None the key is removed; when it has a value it is
restored preserving any URL the project already had set. Also clears
this key's sidecar wrap marker, if any (issue #1768).
Concurrency (#3205): while another live wrap session still owns the key,
this is a no-op restoring underneath a running session unroutes it. Set
``force`` when the user has explicitly asked for their settings back
(``unwrap``), and ``dead_ports`` to name proxy ports already proven dead so
holders that outlived their proxy stop counting as live.
"""
path = settings_path or (Path.cwd() / ".claude" / "settings.local.json")
key = _key_override or _claude_wrap_base_url_env_key(
foundry_mode=foundry_mode, vertex_mode=vertex_mode
)
if not path.exists():
_clear_wrap_marker(path, key=key)
return
try:
payload = json.loads(_read_text(path))
except (OSError, json.JSONDecodeError):
return
if not isinstance(payload, dict):
return
env_map = payload.get("env")
if not isinstance(env_map, dict):
return
if previous is None:
if key not in env_map:
with _wrap_settings_lock(path):
# Another live wrap session in this project may still be using the key.
# Restoring underneath it silently unroutes a running session -- traffic
# bypasses the proxy with no error anywhere (#3205).
release = _release_wrap_key(path, key, force=force, dead_ports=dead_ports)
if not release.should_restore:
# The value stays, but this session's marker must not linger
# describing a process that is gone: hand the slot to a survivor.
_rehome_wrap_marker(path, key=key, survivor=release.survivor, original=release.original)
return
# The owner record holds the value from before *any* wrap session wrote.
# Prefer the caller's own value only when the caller observed it
# first-hand; a session that started second remembers the first
# session's (now dead) proxy URL, and so does the marker an unwrap or a
# self-heal reads it from.
restore_to = previous if release.trust_caller else release.original
if not path.exists():
_clear_wrap_marker(path, key=key)
return
del env_map[key]
if env_map:
payload["env"] = env_map
try:
payload = json.loads(_read_text(path))
except (OSError, json.JSONDecodeError):
return
if not isinstance(payload, dict):
return
env_map = payload.get("env")
if not isinstance(env_map, dict):
return
if restore_to is None:
if key not in env_map:
_clear_wrap_marker(path, key=key)
return
del env_map[key]
if env_map:
payload["env"] = env_map
else:
payload.pop("env", None)
else:
payload.pop("env", None)
else:
env_map[key] = previous
payload["env"] = env_map
if payload:
_write_text(path, json.dumps(payload, indent=2) + "\n")
else:
path.unlink(missing_ok=True)
_clear_wrap_marker(path, key=key)
env_map[key] = restore_to
payload["env"] = env_map
if payload:
_write_text(path, json.dumps(payload, indent=2) + "\n")
else:
path.unlink(missing_ok=True)
_clear_wrap_marker(path, key=key)
def _setup_headroom_mcp(
@ -2018,17 +2334,18 @@ def _setup_serena_mcp(
spec = build_serena_spec(context)
result = registrar.register_server(spec, force=force)
owned_drift = (
result.status == RegisterStatus.MISMATCH
and not force
and headroom_installed_matching(registrar.name, registrar.get_server("serena"))
)
# Migrate a stale Headroom-installed entry. register_server won't overwrite
# a differing spec without force, so an older Headroom Serena entry would
# otherwise persist across re-wraps. Force-update it only when the ledger
# proves Headroom installed the entry that's currently on disk — never a
# user-managed Serena.
if (
result.status == RegisterStatus.MISMATCH
and not force
and headroom_installed_matching(registrar.name, registrar.get_server("serena"))
):
if result.status == RegisterStatus.MISMATCH and not force and owned_drift:
result = registrar.register_server(spec, force=True)
if result.status == RegisterStatus.REGISTERED:
click.echo(" Serena MCP: migrated previously-installed entry to current spec")
@ -2041,7 +2358,13 @@ def _setup_serena_mcp(
result,
label="Serena MCP",
verbose=verbose,
overwrite_hint="update or remove the existing serena MCP entry, then rerun headroom wrap",
overwrite_hint=(
"run headroom wrap again"
if owned_drift
else "run headroom mcp reconcile --adopt"
if registrar.name == "claude"
else "update or remove the existing serena MCP entry, then rerun headroom wrap"
),
restart_hint=f"restart {registrar.display_name} if it was already running",
)
if line is not None:
@ -3638,6 +3961,20 @@ def _copilot_default_wire_api_for_model(model: str | None) -> str:
return _copilot_default_wire_api_for_model_impl(model)
def _build_copilot_native_launch_env(
*, port: int, environ: dict[str, str], project: str | None
) -> tuple[dict[str, str], list[str]]:
from headroom.providers.copilot.wrap import build_native_launch_env
return build_native_launch_env(port=port, environ=environ, project=project)
def _native_api_url_supported(*, environ: Mapping[str, str] | None = None) -> bool | None:
from headroom.providers.copilot.wrap import native_api_url_supported
return native_api_url_supported(environ=environ)
def _should_use_copilot_oauth(
*,
backend: str | None,
@ -4115,39 +4452,8 @@ def _proxy_start_lock(port: int) -> Any:
# environment.
yield
return
with lock_file:
if sys.platform == "win32":
import msvcrt
# msvcrt.locking operates on bytes from the current file position.
lock_file.seek(0)
if lock_file.read(1) == b"":
lock_file.seek(0)
lock_file.write(b"0")
lock_file.flush()
lock_file.seek(0)
# LK_LOCK has implementation-dependent retry limits. A proxy may
# legitimately take longer than that to load ML components, so
# use the non-blocking primitive in a loop instead.
while True:
try:
msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1)
break
except OSError:
time.sleep(0.05)
try:
yield
finally:
lock_file.seek(0)
msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1)
else:
import fcntl
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
try:
yield
finally:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
with _locked_file(lock_file):
yield
@wraps(_ensure_proxy_unlocked)
@ -4341,6 +4647,20 @@ def _ignore_child_sigint(signum: int | None = None, frame: Any = None) -> None:
return None
def _exit_on_signal(signum: int | None = None, frame: Any = None) -> None:
"""Unwind on SIGTERM/SIGHUP so the ``finally`` block actually runs.
Registering ``cleanup`` itself as the handler did not achieve what its call
site documented. A Python signal handler that returns normally does not
unwind the stack -- under PEP 475 the interrupted ``waitpid`` is simply
retried -- so the ``finally`` that restores ``settings.local.json`` never
ran, while the handler had already terminated the proxy underneath a child
that was still alive. Raising SystemExit reverses that: the settings are
restored and cleanup runs exactly once, from ``finally`` (#3205).
"""
raise SystemExit(128 + int(signum or 0))
def _launch_tool(
binary: str,
args: tuple,
@ -4358,6 +4678,7 @@ def _launch_tool(
anyllm_provider: str | None = None,
region: str | None = None,
openai_api_url: str | None = None,
anthropic_api_url: str | None = None,
copilot_api_token: str | None = None,
copilot_refresh_oauth_token: str | None = None,
copilot_api_token_expires_at: float | None = None,
@ -4372,7 +4693,7 @@ def _launch_tool(
port_holder: list[int] = [port]
cleanup = _make_cleanup(proxy_holder, port_holder)
signal.signal(signal.SIGINT, _ignore_child_sigint)
signal.signal(signal.SIGTERM, cleanup)
signal.signal(signal.SIGTERM, _exit_on_signal)
try:
click.echo()
@ -4394,6 +4715,7 @@ def _launch_tool(
anyllm_provider=anyllm_provider,
region=region,
openai_api_url=openai_api_url,
anthropic_api_url=anthropic_api_url,
copilot_api_token=copilot_api_token,
copilot_refresh_oauth_token=copilot_refresh_oauth_token,
copilot_api_token_expires_at=copilot_api_token_expires_at,
@ -4404,7 +4726,7 @@ def _launch_tool(
port_holder[0] = actual_port
_push_runtime_env(actual_port, no_proxy)
# If port fell back, update env URLs to point at the actual port
# If port fell back, update environment URLs to point at the actual port.
if actual_port != port:
for k, v in dict(env).items():
env[k] = v.replace(f"127.0.0.1:{port}", f"127.0.0.1:{actual_port}")
@ -4844,11 +5166,11 @@ def claude(
)
cleanup = _make_cleanup(proxy_holder, port_holder)
signal.signal(signal.SIGINT, _ignore_child_sigint)
signal.signal(signal.SIGTERM, cleanup)
signal.signal(signal.SIGTERM, _exit_on_signal)
if hasattr(signal, "SIGHUP"):
# Terminal close / tmux kill-session sends SIGHUP, not SIGTERM — without
# this, the finally block's base_url restore never runs (issue #1768).
signal.signal(signal.SIGHUP, cleanup)
signal.signal(signal.SIGHUP, _exit_on_signal)
# Memory sync BEFORE proxy startup — sync headroom DB ↔ Claude's files
if memory:
@ -4950,11 +5272,11 @@ def claude(
click.echo(" Skipping MCP retrieve tool (--no-mcp)")
# Coding-task compressor: Serena (retires any legacy tokensave entry).
from headroom.mcp_registry import ClaudeRegistrar
from headroom.mcp_registry import CLAUDE_SERENA_CONTEXT, ClaudeRegistrar
_setup_coding_compressor(
ClaudeRegistrar(),
serena_context="claude-code",
serena_context=CLAUDE_SERENA_CONTEXT,
serena=serena,
no_serena=no_serena,
no_tokensave=no_tokensave,
@ -5233,6 +5555,10 @@ def unwrap_claude(
foundry_mode=_foundry,
vertex_mode=_vertex,
settings_path=_unwrap_settings_path,
# unwrap is the user asking for their settings back, so it drops
# every wrap session's claim rather than deferring to a live
# sibling and silently doing nothing (#3205).
force=True,
)
# Issue #2238: unwrap restores settings.local.json, but a proxy URL that was
@ -5315,6 +5641,14 @@ def _require_copilot_subscription_resolution() -> CopilotSubscriptionTokenResolu
),
)
@click.option("--memory", is_flag=True, help="Enable persistent cross-session memory")
@click.option(
"--native",
is_flag=True,
help=(
"Route Copilot's own GitHub-authenticated API through Headroom instead of "
"the single-model BYOK override. Keeps native model aliases and /model switching."
),
)
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
@click.argument("copilot_args", nargs=-1, type=click.UNPROCESSED)
def copilot(
@ -5327,6 +5661,7 @@ def copilot(
wire_api: str | None,
subscription: bool,
memory: bool,
native: bool,
verbose: bool,
copilot_args: tuple[str, ...],
) -> None:
@ -5361,6 +5696,7 @@ def copilot(
)
raise SystemExit(1)
explicit_subscription = subscription
effective_backend = backend or os.environ.get("HEADROOM_BACKEND")
if _check_proxy(port):
running_backend = _detect_running_proxy_backend(port)
@ -5371,6 +5707,17 @@ def copilot(
)
effective_backend = running_backend or effective_backend
if native:
subscription = True
if provider_type == "anthropic":
raise click.ClickException(
"--native does not use the BYOK provider override; drop --provider-type anthropic."
)
if wire_api is not None:
raise click.ClickException(
"--native selects the wire per request; drop the BYOK-only --wire-api option."
)
effective_provider_type = _resolve_copilot_provider_type(effective_backend, provider_type)
if subscription:
if effective_backend not in (None, "", "anthropic"):
@ -5398,12 +5745,22 @@ def copilot(
copilot_api_token_expires_at: float | None = None
client_bearer: str | None = None
subscription_resolution: CopilotSubscriptionTokenResolution | None = None
if _should_use_copilot_oauth(
anthropic_api_url: str | None = None
use_copilot_oauth = _should_use_copilot_oauth(
backend=effective_backend,
provider_type=provider_type,
env=env,
force_subscription=subscription,
):
)
# Without a provider key, the old implicit OAuth lane still configured
# Copilot as a one-model BYOK client. Native aliases (and runtime /model
# switches) were then forwarded literally and rejected by GitHub (#1910).
# Explicit --subscription remains on its existing fixed-wire behavior;
# implicit GitHub OAuth uses Copilot's own routing automatically.
if use_copilot_oauth and not explicit_subscription:
native = True
if use_copilot_oauth:
if subscription:
subscription_resolution = _require_copilot_subscription_resolution()
client_bearer = subscription_resolution.token
@ -5416,7 +5773,35 @@ def copilot(
"GITHUB_COPILOT_TOKEN / GITHUB_COPILOT_GITHUB_TOKEN."
)
selected_model = _copilot_model_from_args(copilot_args, env)
if native:
openai_api_url = (
subscription_resolution.api_url
if subscription_resolution is not None
else resolve_copilot_api_url(client_bearer)
)
env, env_vars_display = _build_copilot_native_launch_env(
port=port,
environ=env,
project=_project_name_from_cwd(),
)
env["GITHUB_COPILOT_API_URL"] = openai_api_url
env["OPENAI_TARGET_API_URL"] = openai_api_url
env["ANTHROPIC_TARGET_API_URL"] = openai_api_url
anthropic_api_url = openai_api_url
copilot_proxy_token = client_bearer
if subscription_resolution is not None:
copilot_refresh_oauth_token = subscription_resolution.refresh_oauth_token
copilot_api_token_expires_at = subscription_resolution.api_token_expires_at
support = _native_api_url_supported(environ=os.environ)
if support is False:
raise click.ClickException(
"This Copilot CLI build does not reference COPILOT_API_URL; refusing "
"a native launch that could silently bypass Headroom."
)
if support is None and verbose:
click.echo(" Note: could not verify this Copilot CLI's COPILOT_API_URL support.")
else:
selected_model = _copilot_model_from_args(copilot_args, env)
# ``--model auto`` is a Copilot-internal routing token that the BYOK
# API rejects with ``400 The requested model is not supported``. In
@ -5424,7 +5809,7 @@ def copilot(
# Copilot's own native auto-selection works fine — we just need to
# strip the ``--model auto`` flag before launch so Copilot doesn't
# forward it to the provider endpoint.
if _is_auto_model(selected_model):
if not native and _is_auto_model(selected_model):
copilot_args = _strip_auto_model_args(copilot_args)
selected_model = None
click.echo(
@ -5433,57 +5818,58 @@ def copilot(
"automatic model selection."
)
env_wire_api = env.get("COPILOT_PROVIDER_WIRE_API")
effective_wire_api = wire_api or (
env_wire_api
if env_wire_api in {"completions", "responses"}
else _copilot_default_wire_api_for_model(selected_model)
)
env["COPILOT_PROVIDER_TYPE"] = "openai"
# Per-project savings: the Copilot CLI cannot send custom headers, so
# the project rides as a /p/<name> base-URL prefix the proxy strips.
env["COPILOT_PROVIDER_BASE_URL"] = _with_project_prefix(
f"http://127.0.0.1:{port}/v1", _project_name_from_cwd()
)
env["COPILOT_PROVIDER_WIRE_API"] = effective_wire_api
env["COPILOT_PROVIDER_BEARER_TOKEN"] = client_bearer
env["GITHUB_COPILOT_USE_TOKEN_EXCHANGE"] = "false"
env.pop("COPILOT_PROVIDER_API_KEY", None)
# Hand the exact token we resolved (and, for --subscription, validated
# against GitHub) to the proxy explicitly via copilot_proxy_token below.
# The proxy pins it as GITHUB_COPILOT_API_TOKEN, so upstream auth is
# deterministic instead of the proxy re-running unvalidated discovery
# (read_cached_oauth_token returns the *first* candidate, which may not
# be the one the wrapper approved → environment-dependent 401s). Passing
# it as a launch argument — rather than mutating this process's global
# os.environ — keeps the token off shared state and out of unrelated
# code paths.
copilot_proxy_token = client_bearer
if subscription_resolution is not None:
copilot_refresh_oauth_token = subscription_resolution.refresh_oauth_token
copilot_api_token_expires_at = subscription_resolution.api_token_expires_at
env_vars_display = [
"COPILOT_PROVIDER_TYPE=openai",
f"COPILOT_PROVIDER_BASE_URL={env['COPILOT_PROVIDER_BASE_URL']}",
f"COPILOT_PROVIDER_WIRE_API={effective_wire_api}",
(
"COPILOT_AUTH_MODE=github-subscription-experimental"
if subscription
else "COPILOT_AUTH_MODE=github-oauth"
),
]
# Non-subscription OAuth keeps upstream's generic-host policy from
# #610. Subscription mode can use the endpoint returned by the Copilot
# token exchange, which is how Business accounts advertise their API
# host without requiring users to configure it manually.
openai_api_url = (
subscription_resolution.api_url
if subscription_resolution is not None
else resolve_copilot_api_url(client_bearer)
)
env["GITHUB_COPILOT_API_URL"] = openai_api_url
env["OPENAI_TARGET_API_URL"] = openai_api_url
env_vars_display.append(f"COPILOT_PROVIDER_API_URL={openai_api_url}")
if not native:
env_wire_api = env.get("COPILOT_PROVIDER_WIRE_API")
effective_wire_api = wire_api or (
env_wire_api
if env_wire_api in {"completions", "responses"}
else _copilot_default_wire_api_for_model(selected_model)
)
env["COPILOT_PROVIDER_TYPE"] = "openai"
# Per-project savings: the Copilot CLI cannot send custom headers, so
# the project rides as a /p/<name> base-URL prefix the proxy strips.
env["COPILOT_PROVIDER_BASE_URL"] = _with_project_prefix(
f"http://127.0.0.1:{port}/v1", _project_name_from_cwd()
)
env["COPILOT_PROVIDER_WIRE_API"] = effective_wire_api
env["COPILOT_PROVIDER_BEARER_TOKEN"] = client_bearer
env["GITHUB_COPILOT_USE_TOKEN_EXCHANGE"] = "false"
env.pop("COPILOT_PROVIDER_API_KEY", None)
# Hand the exact token we resolved (and, for --subscription, validated
# against GitHub) to the proxy explicitly via copilot_proxy_token below.
# The proxy pins it as GITHUB_COPILOT_API_TOKEN, so upstream auth is
# deterministic instead of the proxy re-running unvalidated discovery
# (read_cached_oauth_token returns the *first* candidate, which may not
# be the one the wrapper approved → environment-dependent 401s). Passing
# it as a launch argument — rather than mutating this process's global
# os.environ — keeps the token off shared state and out of unrelated
# code paths.
copilot_proxy_token = client_bearer
if subscription_resolution is not None:
copilot_refresh_oauth_token = subscription_resolution.refresh_oauth_token
copilot_api_token_expires_at = subscription_resolution.api_token_expires_at
env_vars_display = [
"COPILOT_PROVIDER_TYPE=openai",
f"COPILOT_PROVIDER_BASE_URL={env['COPILOT_PROVIDER_BASE_URL']}",
f"COPILOT_PROVIDER_WIRE_API={effective_wire_api}",
(
"COPILOT_AUTH_MODE=github-subscription-experimental"
if subscription
else "COPILOT_AUTH_MODE=github-oauth"
),
]
# Non-subscription OAuth keeps upstream's generic-host policy from
# #610. Subscription mode can use the endpoint returned by the Copilot
# token exchange, which is how Business accounts advertise their API
# host without requiring users to configure it manually.
openai_api_url = (
subscription_resolution.api_url
if subscription_resolution is not None
else resolve_copilot_api_url(client_bearer)
)
env["GITHUB_COPILOT_API_URL"] = openai_api_url
env["OPENAI_TARGET_API_URL"] = openai_api_url
env_vars_display.append(f"COPILOT_PROVIDER_API_URL={openai_api_url}")
else:
env, env_vars_display = _build_copilot_launch_env(
port=port,
@ -5506,7 +5892,7 @@ def copilot(
)
raise SystemExit(1)
if not subscription and not _copilot_model_configured(copilot_args, env):
if not subscription and not native and not _copilot_model_configured(copilot_args, env):
# Distinguish between "--model auto" (wrong model for BYOK) and
# genuinely missing model (no --model flag at all).
raw_model = _copilot_model_from_args(copilot_args, env)
@ -5543,6 +5929,7 @@ def copilot(
anyllm_provider=anyllm_provider,
region=region,
openai_api_url=openai_api_url,
anthropic_api_url=anthropic_api_url,
copilot_api_token=copilot_proxy_token,
copilot_refresh_oauth_token=copilot_refresh_oauth_token,
copilot_api_token_expires_at=copilot_api_token_expires_at,
@ -6313,9 +6700,10 @@ def kimi(
"""Launch Kimi CLI through Headroom proxy.
\b
Sets KIMI_BASE_URL to route Kimi's OpenAI-compatible /chat/completions
traffic through Headroom. Kimi's own OAuth bearer is forwarded upstream,
so no extra login is required run `kimi` once to authenticate first.
Sets KIMI_CODE_BASE_URL for managed Kimi Code and KIMI_BASE_URL for legacy
kimi-cli to route OpenAI-compatible /chat/completions traffic through
Headroom. Managed Kimi Code needs one `/login` after the proxy URL changes
so its OAuth slot matches that URL; legacy kimi-cli keeps its existing login.
\b
Examples:
@ -6333,9 +6721,20 @@ def kimi(
click.echo("Install Kimi CLI: https://github.com/MoonshotAI/kimi-cli")
raise SystemExit(1)
env, env_vars_display = _build_kimi_launch_env(
port, os.environ, project=_project_name_from_cwd()
)
project = _project_name_from_cwd()
env, env_vars_display = _build_kimi_launch_env(port, os.environ, project=project)
def configure_kimi_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
updated_env, updated_display = _build_kimi_launch_env(
actual_port, current_env, project=project
)
return current_args, updated_env, updated_display
_launch_tool(
binary=kimi_bin,
@ -6350,6 +6749,7 @@ def kimi(
agent_type="kimi",
code_graph=code_graph,
openai_api_url=kimi_api_url,
configure_launch=configure_kimi_launch,
)

View file

@ -228,6 +228,9 @@ DEFAULT_EXCLUDE_TOOLS: frozenset[str] = frozenset(
"WebSearch",
"WebFetch",
"headroom_retrieve",
# Copilot CLI's file-read tool (its `Read` equivalent): raw file bytes
# the model byte-patches against.
"view",
# Lowercase variants for case-insensitive matching
"read",
"glob",
@ -253,6 +256,10 @@ DEFAULT_VERBATIM_EXCLUDE_TOOLS: frozenset[str] = frozenset(
"web_search",
"web_fetch",
"headroom_retrieve",
# `view` (Copilot CLI file read) must stay BYTE-EXACT: the model produces
# line/byte-precise edits against it, and even "lossless" JSON rewrites
# or cross-turn dedup folds break old_str matching and force re-reads.
"view",
}
)

View file

@ -25,6 +25,7 @@ from headroom import paths
from headroom._subprocess import run
from headroom.copilot_linux_secret import read_copilot_oauth_token as read_linux_secret_token
from headroom.copilot_macos_keychain import read_copilot_oauth_token as read_macos_keychain_token
from headroom.proxy import ssl_context as proxy_ssl_context
logger = logging.getLogger(__name__)
@ -76,6 +77,15 @@ _OAUTH_TOKEN_KEYS = (
_EXPIRY_KEYS = ("expires_at", "expiresAt", "expiry", "expires")
def _urlopen(request: urllib_request.Request, *, timeout: float) -> Any:
"""Open a GitHub request with Headroom's configured corporate trust roots."""
context = proxy_ssl_context.build_urlopen_context()
if context is not None:
return urllib_request.urlopen(request, timeout=timeout, context=context)
return urllib_request.urlopen(request, timeout=timeout)
@dataclass(frozen=True)
class CopilotAPIToken:
"""Short-lived API token exchanged from a GitHub OAuth token."""
@ -662,7 +672,7 @@ def start_copilot_device_authorization(
},
method="POST",
)
with urllib_request.urlopen(request, timeout=timeout) as response:
with _urlopen(request, timeout=timeout) as response:
payload = json.loads(response.read().decode("utf-8", errors="replace"))
if not isinstance(payload, dict):
raise RuntimeError("GitHub device authorization returned an invalid response.")
@ -700,7 +710,7 @@ def poll_copilot_device_authorization(
},
method="POST",
)
with urllib_request.urlopen(request, timeout=timeout) as response:
with _urlopen(request, timeout=timeout) as response:
payload = json.loads(response.read().decode("utf-8", errors="replace"))
if not isinstance(payload, dict):
raise RuntimeError("GitHub device authorization returned an invalid response.")
@ -1341,7 +1351,7 @@ def _fetch_copilot_user_info(token: str) -> dict[str, Any] | None:
headers = _copilot_token_exchange_headers(token)
request = urllib_request.Request(_user_info_url(), headers=headers, method="GET")
try:
with urllib_request.urlopen(request, timeout=10.0) as response:
with _urlopen(request, timeout=10.0) as response:
payload = json.loads(response.read().decode("utf-8"))
except Exception as exc:
logger.debug("Unable to resolve Copilot API URL from user info: %s", exc)
@ -1457,7 +1467,7 @@ class CopilotTokenProvider:
def _exchange_token_sync(headers: dict[str, str]) -> dict[str, Any]:
request = urllib_request.Request(_token_exchange_url(), headers=headers, method="GET")
try:
with urllib_request.urlopen(request, timeout=10.0) as response:
with _urlopen(request, timeout=10.0) as response:
payload = json.loads(response.read().decode("utf-8"))
if not isinstance(payload, dict):
return {}

View file

@ -56,7 +56,7 @@ _MAX_DIGEST_TOKENS = 80_000 # Budget for the digest (leave room for prompt + ou
_CLI_BACKENDS: list[tuple[str, str, list[str]]] = [
("claude", "claude-cli", ["claude", "-p", "--output-format", "stream-json", "--verbose"]),
("gemini", "gemini-cli", ["gemini", "-p"]),
("codex", "codex-cli", ["codex", "exec"]),
("codex", "codex-cli", ["codex", "exec", "--skip-git-repo-check"]),
]
# Set of valid CLI model identifiers, derived from _CLI_BACKENDS.
@ -202,7 +202,9 @@ class SessionAnalyzer:
result.recommendations.sort(key=lambda r: r.estimated_tokens_saved, reverse=True)
except Exception as e:
logger.warning("LLM analysis failed: %s", e)
# Return result with stats but no recommendations
# Preserve the stats so multi-project runs can continue, but retain
# the failure so the CLI cannot report an empty result as success.
result.analysis_error = str(e) or type(e).__name__
return result

View file

@ -174,6 +174,7 @@ class AnalysisResult:
total_calls: int = 0
total_failures: int = 0
recommendations: list[Recommendation] = field(default_factory=list)
analysis_error: str | None = None
@property
def failure_rate(self) -> float:

View file

@ -14,11 +14,12 @@ without changing the calling code.
from __future__ import annotations
from .base import MCPRegistrar, RegisterResult, RegisterStatus, ServerSpec
from .claude import ClaudeRegistrar
from .claude import ClaudeConfigMutationError, ClaudeRegistrar
from .codex import CodexRegistrar
from .display import any_succeeded, format_result, format_results
from .grok import GrokRegistrar
from .install import (
CLAUDE_SERENA_CONTEXT,
DEFAULT_PROXY_URL,
build_headroom_spec,
build_serena_spec,
@ -30,6 +31,8 @@ from .server_json import build_server_json, render_server_json
__all__ = [
"DEFAULT_PROXY_URL",
"CLAUDE_SERENA_CONTEXT",
"ClaudeConfigMutationError",
"ClaudeRegistrar",
"CodexRegistrar",
"GrokRegistrar",

View file

@ -26,6 +26,10 @@ from .base import MCPRegistrar, RegisterResult, RegisterStatus, ServerSpec
logger = logging.getLogger(__name__)
class ClaudeConfigMutationError(ValueError):
"""Raised when a Claude config cannot be safely changed."""
class ClaudeRegistrar(MCPRegistrar):
"""Register MCP servers with Claude Code."""
@ -84,6 +88,34 @@ class ClaudeRegistrar(MCPRegistrar):
return entry
return None
def validate_configs_for_mutation(self) -> None:
"""Validate every Claude config root before an explicit mutation."""
seen: set[Path] = set()
for config_path in (self._modern_config, self._legacy_config):
if config_path in seen or not config_path.exists():
continue
seen.add(config_path)
try:
raw = config_path.read_text(encoding="utf-8")
except OSError as exc:
raise ClaudeConfigMutationError(
f"could not read Claude config {config_path}: {exc}"
) from exc
try:
config = json.loads(raw)
except json.JSONDecodeError as exc:
raise ClaudeConfigMutationError(
f"Claude config {config_path} is not valid JSON; refusing to mutate"
) from exc
if not isinstance(config, dict):
raise ClaudeConfigMutationError(
f"Claude config {config_path} must contain a JSON object"
)
if "mcpServers" in config and not isinstance(config["mcpServers"], dict):
raise ClaudeConfigMutationError(
f"Claude config {config_path} has a non-object mcpServers; refusing to mutate"
)
def register_server(self, spec: ServerSpec, *, force: bool = False) -> RegisterResult:
existing = self.get_server(spec.name)
if existing is not None:

View file

@ -14,6 +14,7 @@ from .opencode import OpencodeRegistrar
#: Default proxy URL used when none is given.
DEFAULT_PROXY_URL = "http://127.0.0.1:8787"
CLAUDE_SERENA_CONTEXT = "claude-code"
def get_all_registrars() -> list[MCPRegistrar]:

View file

@ -21,6 +21,10 @@ from .base import ServerSpec
_LEDGER_FILE = "mcp_installs.json"
class LedgerMutationError(ValueError):
"""Raised when a ledger cannot be safely updated."""
def ledger_path() -> Path:
"""Return the Headroom MCP install ledger path."""
return paths.workspace_dir() / _LEDGER_FILE
@ -41,9 +45,17 @@ def spec_fingerprint(spec: ServerSpec) -> str:
def record_install(agent: str, spec: ServerSpec, *, path: Path | None = None) -> None:
"""Record that Headroom installed ``spec`` for ``agent``."""
ledger_file = path or ledger_path()
# Automatic installs must recover from a stale or damaged ledger. The
# explicit reconcile route performs strict validation before config writes.
data = _read_ledger(ledger_file)
agents = data.setdefault("agents", {})
agent_entry = agents.setdefault(agent, {})
agents = data.get("agents")
if not isinstance(agents, dict):
agents = {}
data["agents"] = agents
agent_entry = agents.get(agent)
if not isinstance(agent_entry, dict):
agent_entry = {}
agents[agent] = agent_entry
agent_entry[spec.name] = {
"fingerprint": spec_fingerprint(spec),
"installed_at": datetime.now(timezone.utc).isoformat(),
@ -89,16 +101,45 @@ def headroom_installed_matching(
return entry.get("fingerprint") == spec_fingerprint(current_spec)
def _read_ledger(path: Path) -> dict[str, Any]:
def validate_ledger_for_mutation(path: Path | None = None) -> None:
"""Reject malformed ledger structure before a config mutation."""
_read_ledger(path or ledger_path(), for_mutation=True)
def _read_ledger(path: Path, *, for_mutation: bool = False) -> dict[str, Any]:
try:
raw = path.read_text(encoding="utf-8")
except OSError:
except FileNotFoundError:
return {}
except OSError as exc:
if for_mutation:
raise LedgerMutationError(f"MCP install ledger is unreadable: {path}") from exc
return {}
try:
data = json.loads(raw)
except json.JSONDecodeError:
except json.JSONDecodeError as exc:
if for_mutation:
raise LedgerMutationError(f"MCP install ledger is invalid JSON: {path}") from exc
return {}
return data if isinstance(data, dict) else {}
if not isinstance(data, dict):
if for_mutation:
raise LedgerMutationError("MCP install ledger must contain a JSON object")
return {}
if for_mutation:
for section in ("agents",):
section_data = data.get(section)
if not isinstance(section_data, dict) or any(
not isinstance(agent_entry, dict)
or any(
not isinstance(server_entry, dict)
or not isinstance(server_entry.get("fingerprint"), str)
or not isinstance(server_entry.get("installed_at"), str)
for server_entry in agent_entry.values()
)
for agent_entry in section_data.values()
):
raise LedgerMutationError(f"MCP install ledger section {section!r} is malformed")
return data
def _write_ledger(path: Path, data: dict[str, Any]) -> None:

View file

@ -14,6 +14,7 @@ This is a drop-in replacement for InMemoryGraphStore that:
from __future__ import annotations
import json
import logging
import sqlite3
from collections import deque
from datetime import datetime
@ -26,6 +27,8 @@ from .graph_models import Entity, Relationship, RelationshipDirection, Subgraph
if TYPE_CHECKING:
from ..tracker import ComponentStats
logger = logging.getLogger(__name__)
class SQLiteGraphStore:
"""SQLite-based graph store implementing the GraphStore protocol.
@ -165,19 +168,31 @@ class SQLiteGraphStore:
"metadata": json.dumps(entity.metadata),
}
def _row_to_entity(self, row: sqlite3.Row) -> Entity:
"""Convert database row to Entity object."""
return Entity(
id=row["id"],
user_id=row["user_id"],
name=row["name"],
entity_type=row["entity_type"],
description=row["description"],
properties=json.loads(row["properties"]),
created_at=datetime.fromisoformat(row["created_at"]),
updated_at=datetime.fromisoformat(row["updated_at"]),
metadata=json.loads(row["metadata"]),
)
def _row_to_entity(self, row: sqlite3.Row) -> Entity | None:
"""Convert a database row to an Entity, or None if the row is corrupt.
``properties``/``metadata`` (JSON) and ``created_at``/``updated_at``
(ISO timestamps) are parsed from stored text. A single unparseable row
from a partial write, a manual edit, or a bad migration must not abort
an entire multi-row scan (``query_subgraph``, neighbour expansion): one
corrupt edge would otherwise make an unrelated part of the graph
unqueryable. Skip the bad row instead.
"""
try:
return Entity(
id=row["id"],
user_id=row["user_id"],
name=row["name"],
entity_type=row["entity_type"],
description=row["description"],
properties=json.loads(row["properties"]),
created_at=datetime.fromisoformat(row["created_at"]),
updated_at=datetime.fromisoformat(row["updated_at"]),
metadata=json.loads(row["metadata"]),
)
except (ValueError, TypeError, KeyError) as exc:
logger.warning("skipping corrupt entity row %r: %s", row["id"], exc)
return None
def _relationship_to_row(self, relationship: Relationship) -> dict[str, Any]:
"""Convert Relationship object to row dict for insertion."""
@ -193,19 +208,29 @@ class SQLiteGraphStore:
"metadata": json.dumps(relationship.metadata),
}
def _row_to_relationship(self, row: sqlite3.Row) -> Relationship:
"""Convert database row to Relationship object."""
return Relationship(
id=row["id"],
user_id=row["user_id"],
source_id=row["source_id"],
target_id=row["target_id"],
relation_type=row["relation_type"],
weight=row["weight"],
properties=json.loads(row["properties"]),
created_at=datetime.fromisoformat(row["created_at"]),
metadata=json.loads(row["metadata"]),
)
def _row_to_relationship(self, row: sqlite3.Row) -> Relationship | None:
"""Convert a database row to a Relationship, or None if the row is corrupt.
Same contract as :meth:`_row_to_entity`: a single unparseable relationship
row (bad ``properties``/``metadata`` JSON or ``created_at`` timestamp) must
not abort a whole ``get_relationships`` / ``query_subgraph`` scan and take
unrelated edges down with it. Skip the bad row instead.
"""
try:
return Relationship(
id=row["id"],
user_id=row["user_id"],
source_id=row["source_id"],
target_id=row["target_id"],
relation_type=row["relation_type"],
weight=row["weight"],
properties=json.loads(row["properties"]),
created_at=datetime.fromisoformat(row["created_at"]),
metadata=json.loads(row["metadata"]),
)
except (ValueError, TypeError, KeyError) as exc:
logger.warning("skipping corrupt relationship row %r: %s", row["id"], exc)
return None
# =========================================================================
# Entity Operations
@ -373,7 +398,9 @@ class SQLiteGraphStore:
params,
)
return [self._row_to_relationship(row) for row in cursor]
return [
rel for row in cursor if (rel := self._row_to_relationship(row)) is not None
]
async def delete_relationship(self, relationship_id: str) -> bool:
"""Delete a single relationship.
@ -436,9 +463,12 @@ class SQLiteGraphStore:
)
row = cursor.fetchone()
if row is not None:
entity = self._row_to_entity(row)
if entity is None:
continue
queue.append((entity_id, 0))
visited.add(entity_id)
collected_entities[entity_id] = self._row_to_entity(row)
collected_entities[entity_id] = entity
# BFS traversal
while queue:
@ -470,6 +500,8 @@ class SQLiteGraphStore:
for rel_row in cursor:
rel = self._row_to_relationship(rel_row)
if rel is None:
continue
# Add relationship
collected_relationships[rel.id] = rel
@ -496,8 +528,11 @@ class SQLiteGraphStore:
)
neighbor_row = neighbor_cursor.fetchone()
if neighbor_row is not None:
neighbor = self._row_to_entity(neighbor_row)
if neighbor is None:
continue
visited.add(neighbor_id)
collected_entities[neighbor_id] = self._row_to_entity(neighbor_row)
collected_entities[neighbor_id] = neighbor
queue.append((neighbor_id, depth + 1))
return Subgraph(
@ -651,7 +686,9 @@ class SQLiteGraphStore:
"SELECT * FROM entities WHERE user_id = ?",
(user_id,),
)
return [self._row_to_entity(row) for row in cursor]
return [
entity for row in cursor if (entity := self._row_to_entity(row)) is not None
]
async def clear(self) -> None:
"""Clear all data from the store."""

View file

@ -156,6 +156,73 @@ def provider_key_source(provider_type: str) -> str:
return "ANTHROPIC_API_KEY" if provider_type == "anthropic" else "OPENAI_API_KEY"
COPILOT_NATIVE_API_URL_ENV = "COPILOT_API_URL"
# Any survivor keeps Copilot in its single-model BYOK lane, defeating native
# model routing while making the launch look superficially successful.
COPILOT_BYOK_ENV_VARS: tuple[str, ...] = (
"COPILOT_PROVIDER_BASE_URL",
"COPILOT_PROVIDER_TYPE",
"COPILOT_PROVIDER_API_KEY",
"COPILOT_PROVIDER_BEARER_TOKEN",
"COPILOT_PROVIDER_WIRE_API",
"COPILOT_PROVIDER_TRANSPORT",
"COPILOT_PROVIDER_AZURE_API_VERSION",
"COPILOT_PROVIDER_MODEL_ID",
"COPILOT_PROVIDER_WIRE_MODEL",
"COPILOT_PROVIDER_MODEL_LIMITS_ID",
"COPILOT_PROVIDER_MAX_PROMPT_TOKENS",
"COPILOT_PROVIDER_MAX_OUTPUT_TOKENS",
"COPILOT_PROVIDER_HEADERS",
)
def build_native_launch_env(
*,
port: int,
environ: Mapping[str, str] | None = None,
project: str | None = None,
) -> tuple[dict[str, str], list[str]]:
"""Redirect Copilot's native API surface through Headroom, not BYOK."""
env = dict(environ if environ is not None else os.environ)
base_url = with_project_prefix(f"http://127.0.0.1:{port}", project)
env[COPILOT_NATIVE_API_URL_ENV] = base_url
for variable in COPILOT_BYOK_ENV_VARS:
env.pop(variable, None)
return env, [
f"{COPILOT_NATIVE_API_URL_ENV}={base_url}",
"COPILOT_AUTH_MODE=github-native",
]
def native_api_url_supported(*, environ: Mapping[str, str] | None = None) -> bool | None:
"""Best-effort tri-state probe for the CLI's native API URL override."""
env = environ if environ is not None else os.environ
local = env.get("LOCALAPPDATA") or env.get("HOME") or os.path.expanduser("~")
roots = (
os.path.join(local, "copilot", "pkg"),
os.path.join(os.path.expanduser("~"), ".local", "share", "copilot", "pkg"),
)
found_bundle = False
for root in roots:
if not os.path.isdir(root):
continue
for dirpath, _dirnames, filenames in os.walk(root):
if "app.js" not in filenames:
continue
found_bundle = True
try:
with open(
os.path.join(dirpath, "app.js"), encoding="utf-8", errors="replace"
) as bundle:
while chunk := bundle.read(1 << 20):
if COPILOT_NATIVE_API_URL_ENV in chunk:
return True
except OSError:
continue
return False if found_bundle else None
def build_launch_env(
*,
port: int,

View file

@ -18,11 +18,13 @@ def build_launch_env(
Kimi CLI (``kimi`` / ``kimi-cli``) talks to its managed coding endpoint with
an OpenAI-compatible ``/chat/completions`` client (``kosong``'s ``Kimi``
provider wraps ``AsyncOpenAI``). Its base URL is overridable via the
``KIMI_BASE_URL`` environment variable, so we point it at the local proxy.
provider wraps ``AsyncOpenAI``). Its base URL is overridable via
``KIMI_CODE_BASE_URL`` for the managed client and ``KIMI_BASE_URL`` for
legacy clients, so both point at the local proxy.
The proxy forwards the request including Kimi's own OAuth ``Authorization``
bearer (passthrough auth mode) to the real upstream configured by
``--openai-api-url`` (``https://api.kimi.com/coding/v1``).
bearer after the managed client completes its proxy-scoped ``/login`` to
the real upstream configured by ``--openai-api-url``
(``https://api.kimi.com/coding/v1``).
``project`` (the wrap launch directory) is encoded as a ``/p/<name>``
base-URL prefix because the Kimi base-URL override cannot carry custom
@ -30,5 +32,9 @@ def build_launch_env(
"""
env = dict(environ or os.environ)
base_url = with_project_prefix(codex_proxy_base_url(port), project)
env["KIMI_CODE_BASE_URL"] = base_url
env["KIMI_BASE_URL"] = base_url
return env, [f"KIMI_BASE_URL={base_url}"]
return env, [
f"KIMI_CODE_BASE_URL={base_url}",
f"KIMI_BASE_URL={base_url}",
]

View file

@ -162,9 +162,27 @@ def resolve_extra_headers(
def resolve_api_targets(overrides: ProviderApiOverrides) -> ProviderApiTargets:
"""Resolve normalized upstream provider targets from configured overrides."""
from headroom.copilot_auth import is_copilot_upstream_url
openai = _normalize_api_url(overrides.openai, default=DEFAULT_OPENAI_API_URL)
# GitHub Copilot serves BOTH its OpenAI surface (``/chat/completions``,
# ``/responses``) and its Anthropic surface (``/v1/messages``, for Claude
# models) from the same host. When the OpenAI target is a Copilot host
# (``wrap copilot --subscription`` / ``wrap vscode`` both point it there so
# GPT models work) but no Anthropic target was set, Claude-model requests
# fell back to ``DEFAULT_ANTHROPIC_API_URL`` (api.anthropic.com) and 401'd
# with the Copilot bearer — "Invalid bearer token" (#3247). Default the
# Anthropic target to the same Copilot host so those requests reach the
# surface that actually serves them. An explicit ``ANTHROPIC_TARGET_API_URL``
# still wins (only a ``None`` override is filled in here).
anthropic_override = overrides.anthropic
if anthropic_override is None and is_copilot_upstream_url(openai):
anthropic_override = openai
return ProviderApiTargets(
anthropic=_normalize_api_url(overrides.anthropic, default=DEFAULT_ANTHROPIC_API_URL),
openai=_normalize_api_url(overrides.openai, default=DEFAULT_OPENAI_API_URL),
anthropic=_normalize_api_url(anthropic_override, default=DEFAULT_ANTHROPIC_API_URL),
openai=openai,
gemini=_normalize_api_url(overrides.gemini, default=DEFAULT_GEMINI_API_URL),
cloudcode=_normalize_api_url(overrides.cloudcode, default=DEFAULT_CLOUDCODE_API_URL),
vertex=_normalize_api_url(overrides.vertex, default=DEFAULT_VERTEX_API_URL),

View file

@ -1685,36 +1685,42 @@ class AnthropicHandlerMixin:
if is_token_mode(self.config.mode):
comp_cache = self._get_compression_cache(session_id)
# Re-freeze boundary: consecutive stable messages from start.
# Safety: never freeze beyond provider-confirmed cached prefix.
# `prefix_tracker.frozen_message_count` (set above) is the
# AUTHORITATIVE positional truth — derived from Anthropic's
# `cache_read_input_tokens` response. `compute_frozen_count`
# provides a defensive lower bound from local cache state.
# Use the smaller; never extend past what Anthropic actually
# has cached.
# Freeze + stable marking + Zone-1 swap now live in the
# shared session engine (PROXY policy: clamp by BOTH the
# provider-confirmed count and the locally-replayable
# bound — see session_engine.py's module docstring).
# `frozen_message_count` here has already been through
# tracker + strict-override logic above, so it is the
# AUTHORITATIVE positional truth derived from Anthropic's
# `cache_read_input_tokens` response.
#
# Issue #327: a previous version walked past
# `prefix_tracker.frozen_message_count` whenever an upcoming
# tool_result's content-hash matched `_stable_hashes` or
# `should_defer_compression` returned True. That conflated
# content equality with positional cache membership: the
# prefix cache is positional (bytes 0..K cached, anything
# past K is fresh), but `_stable_hashes` is content-keyed
# and grows unbounded. On long Claude Code sessions where
# tool_result content rhymes across turns (repeated system
# prompts, repeated file reads, etc.), the walker advanced
# Issue #327 (history kept at the call site): a previous
# version walked past `prefix_tracker.frozen_message_count`
# whenever an upcoming tool_result's content-hash matched
# `_stable_hashes` or `should_defer_compression` returned
# True. That conflated content equality with positional
# cache membership: the prefix cache is positional (bytes
# 0..K cached, anything past K is fresh), but
# `_stable_hashes` is content-keyed and grows unbounded.
# On long Claude Code sessions where tool_result content
# rhymes across turns, the walker advanced
# `frozen_message_count` to `len(messages)` and the
# pipeline produced `transforms_applied=[]` on 73% of
# requests. The walker has been removed; trust
# `prefix_tracker` clamped by `compute_frozen_count`.
cache_frozen_count = comp_cache.compute_frozen_count(messages)
frozen_message_count = min(frozen_message_count, cache_frozen_count)
# Record all tool_results in the verified frozen prefix as stable
comp_cache.mark_stable_from_messages(messages, frozen_message_count)
from headroom.proxy.session_engine import (
FREEZE_POLICY_CONFIRMED_CLAMP,
prepare_turn,
)
# Zone 1: Swap cached compressed versions into working copy
working_messages = comp_cache.apply_cached(messages)
_prep = prepare_turn(
comp_cache,
messages,
policy=FREEZE_POLICY_CONFIRMED_CLAMP,
tracker_frozen=frozen_message_count,
)
frozen_message_count = _prep.frozen_message_count
working_messages = _prep.pipeline_input
if (
getattr(self, "_background_compression_enabled", False)
and frozen_message_count == 0
@ -2065,39 +2071,43 @@ class AnthropicHandlerMixin:
# previously-forwarded prefix keeps it byte-identical → cache hits.
# Append-only-guarded and idempotent (cache mode already replays), so
# it is safe to run unconditionally here.
from headroom.cache.prefix_tracker import (
normalize_message_cache_control,
overlay_cached_prefix,
)
from headroom.cache.prefix_tracker import normalize_message_cache_control
from headroom.proxy.session_engine import finalize_turn
_overlay_replayed = False
# On a confirmed-cold turn we deliberately do NOT replay the previously
# forwarded prefix: the cache is dead (nothing to keep byte-identical for)
# and the replay would clobber the whole-prefix recompaction we just did.
if _decision.should_compress and not _skip_compression_for_backpressure:
#
# Backpressure skips the compression PIPELINE but must NOT skip this
# replay: on the saturated path `optimized_messages` is the raw
# originals, which mismatch the compressed prefix the provider cached
# — so every gated request busted its session's prompt cache exactly
# when traffic (and the re-write cost) peaked. The overlay itself is
# O(prefix) comparisons plus one token recount only when it actually
# replays, which is far cheaper than the whole-prefix cache re-write
# it prevents, so it stays on even under backpressure.
if _decision.should_compress:
if _cold_recompact_active:
_overlay_replayed = False
else:
_ov = overlay_cached_prefix(
_final = finalize_turn(
optimized_messages,
original_client_messages,
previous_original_messages,
previous_forwarded_messages,
count_tokens=tokenizer.count_messages,
)
_overlay_replayed = _ov != optimized_messages
_overlay_replayed = _final.replayed
if _overlay_replayed:
optimized_messages = _ov
optimized_tokens = tokenizer.count_messages(optimized_messages)
optimized_messages = _final.messages
if _final.tokens is not None:
optimized_tokens = _final.tokens
else:
replay_skip_reason = (
"pre_upstream_backpressure"
if _skip_compression_for_backpressure
else _decision.passthrough_reason
)
logger.debug(
"[%s] Cached-prefix replay skipped: reason=%s",
request_id,
replay_skip_reason,
_decision.passthrough_reason,
)
# Own cache_control placement: the client moves the breakpoint each

View file

@ -1593,6 +1593,14 @@ WS_FIRST_FRAME_TIMEOUT_SECONDS = 60.0
# "lossless" would otherwise look like it worked).
COMPRESS_MODES = ("ccr", "lossy_inline", "lossless_then_lossy")
# Max wait for a sidecar session's turn lock, on the executor. MUST stay
# well below COMPRESSION_TIMEOUT_SECONDS: with an untimed acquire, a slow
# turn's 503-driven retries would park executor workers blocked on the lock
# doing no work, each recording timeout debt toward the compression
# quarantine. Failing the acquire raises TimeoutError, which maps to the
# session-mode 503 retry path.
_SESSION_TURN_LOCK_TIMEOUT_SECONDS = 10.0
def _extract_codex_handshake_headers(upstream: Any) -> list[tuple[str, str]]:
"""Return the ``x-codex-*`` headers from an upstream WS handshake response.
@ -2116,6 +2124,45 @@ class OpenAIHandlerMixin:
if is_tool_excluded(fn_name, DEFAULT_VERBATIM_EXCLUDE_TOOLS)
}
# Read protection (HEADROOM_PROTECT_READS) — parity with the
# chat/Anthropic path (ContentRouter.apply). Output of a file-READ
# command (cat/nl/sed -n/head/tail/…) must stay verbatim: the agent
# byte-patches against it, and lossy reads caused re-reads /
# turn-inflation + resolve loss on SWE-bench. The Responses wire carries
# the producing command in two shapes, both normalized by the shared
# _tool_call_command_text helper:
# - function_call.arguments (Copilot bash, Codex exec_command, …)
# - local_shell_call.action (native Responses shell; argv or string)
# Content is gated per-output by _read_output_should_be_protected so
# confidently non-code DATA reads (lockfiles, JSON, logs, search) stay
# compressible, exactly like the chat path.
from headroom.transforms.content_router import (
_is_read_command,
_read_output_should_be_protected,
_tool_call_command_text,
read_protection_enabled,
)
read_command_by_call_id: dict[str, str] = {}
if read_protection_enabled():
for item in items:
if not isinstance(item, dict):
continue
item_type = item.get("type")
if item_type == "function_call":
command = _tool_call_command_text(item.get("arguments"))
elif item_type == "local_shell_call":
command = _tool_call_command_text(item.get("action"))
else:
continue
call_id = item.get("call_id")
if command and isinstance(call_id, str) and call_id and _is_read_command(command):
read_command_by_call_id[call_id] = command
# Outputs protected by read-command detection. Also unioned into the
# cross-turn dedup protection set below: a [↑…] fold of a read would
# break the exact-bytes contract just like lossy compression would.
read_protected_call_ids: set[str] = set()
timing_sink: dict[str, float] = timing if timing is not None else {}
def _add_timing(name: str, started_at: float) -> None:
@ -2159,6 +2206,24 @@ class OpenAIHandlerMixin:
}
)
continue
if isinstance(call_id, str) and call_id in read_command_by_call_id:
# Finalize by CONTENT (same gate as ContentRouter.apply):
# protect unless the output is confidently non-code DATA.
if _read_output_should_be_protected(_responses_part_text(item.get("output"))):
read_protected_call_ids.add(call_id)
if debug_enabled:
extraction_debug.append(
{
"index": idx,
"eligible": False,
"reason": "read_command_protected",
"item_type": item_type,
"call_id": call_id,
"command": read_command_by_call_id[call_id],
"item": item,
}
)
continue
if isinstance(call_id, str) and call_id in excluded_call_ids:
if call_id in verbatim_excluded_call_ids:
if debug_enabled:
@ -2180,6 +2245,7 @@ class OpenAIHandlerMixin:
# Note: when output is a content-part array, fold each text part
# individually using ("output_part", index) slots to preserve the
# array structure (non-text parts like images are left untouched).
excluded_folded = False
raw_output = item.get("output")
if isinstance(raw_output, list):
for pidx, part in enumerate(raw_output):
@ -2191,6 +2257,7 @@ class OpenAIHandlerMixin:
part_text = part["text"]
pf = router._lossless_compact_excluded(part_text)
if pf is not None:
excluded_folded = True
lossless_excluded.append(
(idx, ("output_part", pidx), pf[0], part_text)
)
@ -2198,6 +2265,7 @@ class OpenAIHandlerMixin:
excl_out = _responses_part_text(raw_output)
fold = router._lossless_compact_excluded(excl_out) if excl_out else None
if fold is not None:
excluded_folded = True
lossless_excluded.append((idx, ("output", None), fold[0], excl_out))
if debug_enabled:
extraction_debug.append(
@ -2206,7 +2274,7 @@ class OpenAIHandlerMixin:
"eligible": False,
"reason": (
"exclude_tools_lossless_fold"
if fold is not None
if excluded_folded
else "exclude_tools_protected"
),
"item_type": item_type,
@ -2622,7 +2690,7 @@ class OpenAIHandlerMixin:
updated_items,
self.OPENAI_RESPONSES_OUTPUT_TYPES,
tokenizer.count_text,
protected_call_ids=verbatim_excluded_call_ids,
protected_call_ids=verbatim_excluded_call_ids | read_protected_call_ids,
)
if dd_folded:
modified = True
@ -3693,17 +3761,45 @@ class OpenAIHandlerMixin:
if is_token_mode(self.config.mode):
comp_cache = self._get_compression_cache(openai_session_id)
# Zone 1: Swap cached compressed versions
working_messages = comp_cache.apply_cached(messages)
# Re-freeze boundary. Token mode can use the compression
# cache's positional frozen count. Cache mode must keep the
# latest observation mutable even when the compression
# cache has no compressible entry for it yet; otherwise
# OpenAI-compatible tool-call clients freeze the entire
# conversation and report near-zero savings.
if not is_cache_mode(self.config.mode):
openai_frozen_count = comp_cache.compute_frozen_count(messages)
# Token mode: shared engine, REPLAYABLE policy — its
# formula with no explicit pin is exactly this path's
# historical freeze (compute_frozen_count alone; the
# tracker count feeds cache mode below, never token
# mode). The engine also runs
# mark_stable_from_messages, which this path skipped:
# that marks tool_results INSIDE the frozen prefix as
# stable — redundant in the common case (an in-prefix
# tool_result is already stable via its cache entry)
# but it keeps `_stable_hashes` bookkeeping identical
# across all three paths, e.g. preserving stability
# across cache-entry LRU turnover. Note it can never
# mark the BOUNDARY tool_result that stopped the
# count (it sits outside messages[:frozen]) — the
# protection against re-compressing a passthrough
# boundary tool_result under rising context pressure
# is the router-level `_frozen_verdicts` pin, on every
# path, unchanged by this migration.
from headroom.proxy.session_engine import (
FREEZE_POLICY_REPLAYABLE,
prepare_turn,
)
_prep = prepare_turn(
comp_cache,
messages,
policy=FREEZE_POLICY_REPLAYABLE,
)
working_messages = _prep.pipeline_input
openai_frozen_count = _prep.frozen_message_count
else:
# Cache mode: Zone-1 swap only. The latest observation
# must stay mutable even when the compression cache
# has no entry for it yet (otherwise OpenAI-compatible
# tool-call clients freeze the entire conversation and
# report near-zero savings), so the freeze comes from
# the tracker (set above), never from the cache count.
working_messages = comp_cache.apply_cached(messages)
result = await self._run_compression_in_executor(
lambda: self.openai_pipeline.apply(
@ -3794,21 +3890,30 @@ class OpenAIHandlerMixin:
# Cache-safety (ALL modes): forward the previously-cached (compressed)
# prefix byte-identical, so freezing can't bust the prompt cache. See the
# matching guard in the Anthropic handler for the full rationale. Append-
# only-guarded and idempotent (cache mode already replays).
from headroom.cache.prefix_tracker import overlay_cached_prefix
# only-guarded and idempotent (cache mode already replays). Shared
# implementation: session_engine.finalize_turn.
from headroom.proxy.session_engine import finalize_turn
_ov = overlay_cached_prefix(
_final = finalize_turn(
optimized_messages,
original_client_messages,
openai_prefix_tracker.get_last_original_messages(),
openai_prefix_tracker.get_last_forwarded_messages(),
count_tokens=tokenizer.count_messages,
)
if _ov != optimized_messages:
optimized_messages = _ov
optimized_tokens = tokenizer.count_messages(optimized_messages)
if _final.replayed:
optimized_messages = _final.messages
if _final.tokens is not None:
optimized_tokens = _final.tokens
# Guard: if "optimization" inflated tokens, revert to originals
if optimized_tokens > original_tokens:
# Guard: if "optimization" inflated tokens, revert to originals.
# NEVER after the overlay replayed (same exemption as the Anthropic
# handler): the replayed prefix is the exact bytes the provider
# cached, and reverting to raw originals re-forwards the uncompressed
# prefix — trading a 90% read discount for a full cache re-write. The
# nominal "inflation" there is an artifact of comparing the cached
# (compressed) forwarding against the raw original count.
if optimized_tokens > original_tokens and not _final.replayed:
logger.warning(
f"[{request_id}] Optimization inflated tokens "
f"({original_tokens} -> {optimized_tokens}), reverting to original messages"
@ -9531,6 +9636,9 @@ class OpenAIHandlerMixin:
headers = dict(request.headers)
tags = extract_tags(headers)
client = classify_client(headers)
# Initialized before the try so the TimeoutError handler can branch on
# it even if the failure happened before session parsing.
session_id = None
try:
# Use OpenAI pipeline (messages are in OpenAI format from TS SDK)
@ -9595,6 +9703,64 @@ class OpenAIHandlerMixin:
}
},
)
# Session-aware sidecar mode (opt-in): with a session id the
# endpoint keeps the byte-replay state ITSELF — the same
# per-session machinery the proxy path uses (compression cache +
# prefix tracker, with the registry's TTL/LRU lifecycle) — so a
# gateway that owns routing (e.g. Kong) can resend the RAW
# conversation every turn and still get a byte-identical prefix
# back. Contract: the caller forwards the returned messages
# verbatim, and may relay provider usage via POST /v1/usage for
# telemetry/attribution. Without a session id, behaviour is the
# stateless contract, unchanged.
session_id = compress_config.get("session_id")
# The x-headroom-session-id header is honored only behind an
# explicit env opt-in: deployments whose gateways already stamp
# that header on ALL traffic (it is the documented proxy-path
# session key) would otherwise silently flip stateless callers
# into session mode on upgrade — and a header value shared across
# conversations (Claude Code subagents do exactly this) would
# blend unrelated conversations into one replay state.
if session_id is None and os.environ.get(
"HEADROOM_COMPRESS_SESSION_FROM_HEADER", ""
).lower() in ("1", "true"):
session_id = request.headers.get("x-headroom-session-id")
if session_id is not None and (
not isinstance(session_id, str) or not session_id.strip() or len(session_id) > 256
):
return JSONResponse(
status_code=400,
content={
"error": {
"type": "invalid_request",
"message": (
f"Invalid config.session_id: {session_id!r}. "
"Expected a non-empty string of at most 256 characters."
),
}
},
)
if session_id is not None and compress_user_messages:
# User/assistant rewrites are not content-addressed (the
# session cache replays tool_result content only), so once the
# tracker's overlay snapshots expire a rewritten user message
# would come back in RAW form — a guaranteed prefix bust inside
# the tracker-TTL/cache-TTL window. Refuse the combination
# rather than bust later.
return JSONResponse(
status_code=400,
content={
"error": {
"type": "invalid_request",
"message": (
"config.compress_user_messages is not supported with "
"config.session_id: user-message rewrites cannot be "
"byte-replayed across turns, which would bust the "
"provider prompt cache."
),
}
},
)
# Mode selection. Default is marker-free (see _no_ccr_pipeline):
# no caller of this route can resolve a CCR marker unless it opts in
# with mode="ccr", which restores the full marker + store behaviour.
@ -9641,23 +9807,160 @@ class OpenAIHandlerMixin:
if frozen_message_count is not None:
pipeline_kwargs["frozen_message_count"] = frozen_message_count
# Offload the CPU-bound pipeline to the bounded compression executor
# (mirrors the request handlers above). Running apply() inline blocked
# the single event loop on a large payload, so even GET /health stalled
# until it finished (#718). The executor also enforces a timeout so a
# too-large body fails fast instead of hanging forever.
result = await self._run_compression_in_executor(
lambda: pipeline.apply(
messages=messages,
model=model,
**pipeline_kwargs,
),
# Sidecar session pre-work: swap in previously-computed compressed
# bytes (Zone 1), then freeze the ENTIRE locally-replayable prefix
# (`compute_frozen_count`). This deliberately differs from the
# proxy path's `min(tracker, cache)` posture: in sidecar mode,
# whatever this endpoint previously RETURNED is the provider's
# cache contract, so every already-returned message must come back
# byte-identical — recompressing it (even "better") is a bust.
# Over-freezing relative to the provider's actual cache only
# forgoes tail compression; it can never bust. The tracker's
# /v1/usage-fed freeze count is deliberately NOT a freeze floor —
# freezing a message whose cache entry was evicted would forward
# raw original bytes. An explicit config.frozen_message_count
# still wins when larger: the caller may know more about the
# provider cache than local state does.
comp_cache = None
session_tracker = None
if session_id:
# Namespaced with a NUL separator so sidecar sessions can
# never collide with proxy-path session ids: NUL cannot
# appear in an HTTP header value, so no client-supplied
# x-headroom-session-id on the proxy path can spoof its way
# into a sidecar session's tracker or replay cache (the same
# trick SessionTrackerStore uses for its synthetic lineage
# keys). A plain "compress:" string prefix was spoofable.
_session_key = f"compress\x00{session_id}"
_tracker_provider = (
"anthropic"
if ("claude" in model_name.lower() or "anthropic" in model_name.lower())
else "openai"
)
comp_cache = self._get_compression_cache(_session_key)
session_tracker = self.session_tracker_store.get_or_create(
_session_key, _tracker_provider
)
def _run_stateless():
result = pipeline.apply(messages=messages, model=model, **pipeline_kwargs)
return (
result,
result.messages,
result.tokens_before,
result.tokens_after,
None,
)
def _run_session_turn():
# One sidecar turn as a single executor-side block: every step
# here is CPU-bound (content hashing, deep compares, token
# counts, full-transcript deepcopies) and must stay off the
# event loop for the same reason pipeline.apply does (#718).
# The per-session lock serializes contract-violating
# concurrent turns so an older in-flight turn cannot tear or
# overwrite a newer turn's tracker snapshots mid-flight.
# Cache management (freeze + swap + overlay) lives in the
# shared session engine — one brain for this path and the
# proxy request paths.
from headroom.proxy.session_engine import (
FREEZE_POLICY_REPLAYABLE,
finalize_turn,
prepare_turn,
)
# TIMED acquire, strictly shorter than the executor timeout:
# an untimed `with lock:` here lets one slow session's
# 503-driven retries park executor workers doing no work —
# each blocked worker records timeout debt and can arm the
# compression quarantine for ALL traffic. Failing fast maps
# to the same TimeoutError → session-mode 503 → retry path.
if not comp_cache.session_turn_lock.acquire(
timeout=_SESSION_TURN_LOCK_TIMEOUT_SECONDS
):
raise TimeoutError(
f"session turn lock busy for {session_id!r} "
"(a previous turn for this session is still running)"
)
try:
prev_original = session_tracker.get_last_original_messages()
prev_returned = session_tracker.get_last_forwarded_messages()
prep = prepare_turn(
comp_cache,
messages,
policy=FREEZE_POLICY_REPLAYABLE,
explicit_frozen=frozen_message_count,
)
session_frozen = prep.frozen_message_count
pipeline_kwargs["frozen_message_count"] = session_frozen
result = pipeline.apply(
messages=prep.pipeline_input, model=model, **pipeline_kwargs
)
# Replay last turn's exact returned prefix over any drift
# the pipeline introduced — byte-identical is the contract
# the caller forwards on.
turn = finalize_turn(result.messages, messages, prev_original, prev_returned)
final = turn.messages
# Savings are reported against the caller's RAW payload,
# not the cache-swapped pipeline input: on a warm turn the
# swap has already shrunk the input before the pipeline
# counts it, which made every warm turn report ~0 saved.
try:
from headroom.tokenizers import get_tokenizer
_tok = get_tokenizer(model_name)
raw_tokens_before = _tok.count_messages(messages)
final_tokens_after = _tok.count_messages(final)
except Exception as e:
# Fail-open, but LOUD: this fallback reverts to the
# pipeline's counts of the cache-swapped input, which
# silently resurrects the ~0-saved warm-turn bug the
# raw recount exists to fix — per-model, so it can
# hide indefinitely without this log.
logger.warning(
"[compress:%s] raw-payload token recount failed for "
"model %s (%s: %s); savings for this turn are "
"reported against the cache-swapped input",
session_id,
model_name,
type(e).__name__,
e,
)
raw_tokens_before = result.tokens_before
final_tokens_after = result.tokens_after
comp_cache.update_from_result(messages, final)
# Record this turn's result as the new "last returned" —
# the sidecar equivalent of "last forwarded", captured at
# return time because whatever we hand back IS what the
# caller sends upstream.
session_tracker.record_returned(messages, final)
info = {
"id": session_id,
"frozen_message_count": session_frozen,
"cached_prefix_replayed": turn.replayed,
}
return result, final, raw_tokens_before, final_tokens_after, info
finally:
comp_cache.session_turn_lock.release()
# Offload the CPU-bound work to the bounded compression executor
# (mirrors the request handlers above). Running it inline blocked
# the single event loop on a large payload, so even GET /health
# stalled until it finished (#718). The executor also enforces a
# timeout so a too-large body fails fast instead of hanging.
(
result,
final_messages,
tokens_before,
tokens_after,
session_info,
) = await self._run_compression_in_executor(
_run_session_turn if session_id else _run_stateless,
timeout=COMPRESSION_TIMEOUT_SECONDS,
)
ccr_hashes = _response_ccr_hashes(result.messages, result.markers_inserted)
tokens_before = result.tokens_before
tokens_after = result.tokens_after
ccr_hashes = _response_ccr_hashes(final_messages, result.markers_inserted)
tokens_saved = max(0, tokens_before - tokens_after)
latency_ms = (time.time() - start_time) * 1000
await self._record_request_outcome(
@ -9689,28 +9992,82 @@ class OpenAIHandlerMixin:
)
)
return JSONResponse(
{
"messages": result.messages,
"tokens_before": result.tokens_before,
"tokens_after": result.tokens_after,
"tokens_saved": result.tokens_before - result.tokens_after,
"compression_ratio": (
result.tokens_after / result.tokens_before
if result.tokens_before > 0
else 1.0
),
"transforms_applied": result.transforms_applied,
"transforms_summary": result.transforms_summary,
"ccr_hashes": ccr_hashes,
}
)
_payload = {
"messages": final_messages,
"tokens_before": tokens_before,
"tokens_after": tokens_after,
# Clamped like the telemetry above: the overlay's byte-replay
# can legitimately return a slightly larger prefix than the
# pipeline's best effort, and a negative "saved" here while
# telemetry records 0 would be two answers for one number.
"tokens_saved": tokens_saved,
"compression_ratio": (tokens_after / tokens_before if tokens_before > 0 else 1.0),
"transforms_applied": result.transforms_applied,
"transforms_summary": result.transforms_summary,
"ccr_hashes": ccr_hashes,
}
if session_info is not None:
_payload["session"] = session_info
return JSONResponse(_payload)
except TimeoutError:
self.metrics.record_compression_failed("timeout")
if session_id:
# Fail-open-with-originals is WRONG for a session call: the
# timed-out worker cannot be cancelled and may still finish
# and record its compressed result as "last returned" — while
# the caller, handed the originals, forwards those instead.
# The desynced snapshot then busts the next turn. A 503 tells
# the gateway to retry; the retry lands on whatever state the
# straggler recorded and replays it consistently.
logger.warning(
"Compression timed out after %.0fs for session %r; "
"returning 503 (session mode cannot fail open without "
"desyncing replay state)",
COMPRESSION_TIMEOUT_SECONDS,
session_id,
)
# Same outcome recording as the stateless timeout path below:
# session timeouts hit the largest transcripts, and skipping
# the RequestOutcome here under-counts exactly those requests
# when dashboards reconcile failure counters against outcomes.
_timeout_latency_ms = (time.time() - start_time) * 1000
await self._record_request_outcome(
RequestOutcome(
request_id=(
await self._next_request_id()
if hasattr(self, "_next_request_id")
else f"compress_{int(time.time())}"
),
provider="compress",
model=model if isinstance(model, str) else str(model),
original_tokens=0,
optimized_tokens=0,
output_tokens=0,
tokens_saved=0,
attempted_input_tokens=0,
total_latency_ms=_timeout_latency_ms,
overhead_ms=_timeout_latency_ms,
num_messages=len(messages) if isinstance(messages, list) else 0,
tags=tags,
client=client,
)
)
return JSONResponse(
status_code=503,
content={
"error": {
"type": "compression_timeout",
"message": (
"Compression timed out; retry this turn. "
"Session replay state remains consistent."
),
}
},
)
logger.warning(
"Compression timed out after %.0fs; failing open with original messages",
COMPRESSION_TIMEOUT_SECONDS,
)
self.metrics.record_compression_failed("timeout")
latency_ms = (time.time() - start_time) * 1000
await self._record_request_outcome(
RequestOutcome(
@ -9760,6 +10117,186 @@ class OpenAIHandlerMixin:
},
)
async def handle_compress_usage(self, request: Request) -> JSONResponse:
"""Relay of the provider's usage block for a sidecar compress session.
POST /v1/usage
Body: {"session_id": "...",
"usage": {"cache_read_input_tokens": N,
"cache_creation_input_tokens": N}}
The session-aware ``/v1/compress`` never sees the provider's response
(the caller owns routing). This relay feeds the provider-confirmed
numbers into the session's tracker — the same signal the proxy path
reads from the response itself powering cache-hit/miss attribution,
idle-vs-prefix-change classification, and savings accounting for
sidecar sessions.
Deliberately NOT a freeze input: the compress path freezes exactly the
locally-replayable prefix (``compute_frozen_count``), and raising that
to a provider-confirmed count could freeze a message whose cache entry
was evicted which would forward raw original bytes and bust the very
prefix the count vouched for. Optional: skipping this call costs
telemetry fidelity, never correctness.
"""
from fastapi.responses import JSONResponse
from headroom.proxy.helpers import _read_request_json
def _invalid(message: str) -> JSONResponse:
return JSONResponse(
status_code=400,
content={"error": {"type": "invalid_request", "message": message}},
)
try:
body = await _read_request_json(request)
except Exception:
return _invalid("Invalid JSON in request body.")
session_id = body.get("session_id")
if not isinstance(session_id, str) or not session_id.strip() or len(session_id) > 256:
return _invalid(
"Missing or invalid session_id: expected a non-empty string "
"of at most 256 characters."
)
usage = body.get("usage")
if not isinstance(usage, dict):
return _invalid("Missing or invalid usage: expected an object.")
# A usage block carrying NEITHER cache field is a no-signal relay (an
# OpenAI-style {"prompt_tokens": N} forwarded verbatim, for example).
# Defaulting the absent fields to 0 would make update_from_response
# treat it as a provider-confirmed fully-cold turn and wipe the
# tracker's cached-prefix state — so absence of both is a 400, not 0.
if "cache_read_input_tokens" not in usage and "cache_creation_input_tokens" not in usage:
return _invalid(
"usage must carry cache_read_input_tokens and/or "
"cache_creation_input_tokens; a block with neither carries no "
"cache signal and is not accepted."
)
def _token_field(name: str) -> int | None:
value = usage.get(name, 0)
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
return None
return value
cache_read = _token_field("cache_read_input_tokens")
cache_write = _token_field("cache_creation_input_tokens")
if cache_read is None or cache_write is None:
return _invalid(
"usage.cache_read_input_tokens and usage.cache_creation_input_tokens "
"must be non-negative integers when present."
)
# Same NUL-separated namespace as handle_compress: unspoofable from
# any HTTP header. peek() (never get_or_create) so a flood of novel
# session ids cannot grow the tracker store — an unknown or expired
# session is answered without leaving a footprint, and the session
# keeps the provider its compress call inferred rather than a default
# from here.
_session_key = f"compress\x00{session_id}"
tracker = self.session_tracker_store.peek(_session_key)
if tracker is None:
return self._compress_usage_unknown_session(session_id)
# No create, no LRU bump: the cache is only needed for its turn lock.
comp_cache = self._peek_compression_cache(_session_key)
# A relay whose only present field is zero carries no positive cache
# signal (an OpenAI-mapped gateway naturally sends
# {"cache_read_input_tokens": 0} with no write field — OpenAI has no
# write signal). Applying it would hit update_from_response's
# total_cached == 0 branch and wipe the tracker's cached-prefix
# state — a "provider-confirmed fully cold" reset the relay never
# actually asserted. Only a relay with BOTH fields present may claim
# a genuine fully-cold turn.
_both_present = (
"cache_read_input_tokens" in usage and "cache_creation_input_tokens" in usage
)
if cache_read + cache_write == 0 and not _both_present:
return JSONResponse(
{
"session_id": session_id,
"frozen_message_count": tracker.get_frozen_message_count(),
"applied": False,
"reason": "no_cache_signal",
}
)
def _apply_usage():
# Off the event loop (full-transcript deepcopies + per-message
# token estimation live in update_from_response), and under the
# session turn lock: an unlocked update here races the
# executor-side compress turn — record_returned installs turn
# N+1's snapshots, then this write would roll them back to turn
# N's copies and the next overlay would refuse to replay.
lock = comp_cache.session_turn_lock if comp_cache is not None else None
if lock is not None and not lock.acquire(timeout=_SESSION_TURN_LOCK_TIMEOUT_SECONDS):
raise TimeoutError(f"session turn lock busy for {session_id!r}")
try:
last_returned = tracker.get_last_forwarded_messages()
if not last_returned:
return None
tracker.update_from_response(
cache_read_tokens=cache_read,
cache_write_tokens=cache_write,
messages=last_returned,
original_messages=tracker.get_last_original_messages(),
)
return tracker.get_frozen_message_count()
finally:
if lock is not None:
lock.release()
try:
frozen_count = await self._run_compression_in_executor(
_apply_usage, timeout=COMPRESSION_TIMEOUT_SECONDS
)
except TimeoutError:
return JSONResponse(
status_code=503,
content={
"error": {
"type": "session_busy",
"message": (
"A compress turn for this session is in flight; retry the usage relay."
),
}
},
)
if frozen_count is None:
return self._compress_usage_unknown_session(session_id)
return JSONResponse(
{
"session_id": session_id,
"frozen_message_count": frozen_count,
"applied": True,
}
)
@staticmethod
def _compress_usage_unknown_session(session_id: str):
from fastapi.responses import JSONResponse
# No compress state for this session: never seen, or the tracker's
# session TTL reclaimed it. Note the byte-replay cache lives longer
# than the tracker, so a 404 here does NOT mean the next /v1/compress
# loses replay — only this telemetry relay landed nowhere.
return JSONResponse(
status_code=404,
content={
"error": {
"type": "unknown_session",
"message": (
f"No usage-tracking state for session {session_id!r} "
"(never seen, or expired). Compression replay for the "
"session may still be active; only this telemetry "
"relay landed nowhere."
),
}
},
)
async def _maybe_compress_passthrough_responses(
self, body: bytes, *, client: str | None = None
) -> bytes:

View file

@ -1245,8 +1245,54 @@ try:
except ValueError:
EAGER_PRELOAD_TIMEOUT_SECONDS = 120.0
# Maximum compression cache sessions (prevents unbounded memory growth)
MAX_COMPRESSION_CACHE_SESSIONS = 500
# Maximum compression cache sessions (prevents unbounded memory growth).
# Overridable via HEADROOM_COMPRESSION_CACHE_MAX_SESSIONS for gateway
# deployments (e.g. Kong sidecars) that fan many concurrent sessions into one
# proxy process. Falls back to 500 on an unparseable value; floor of 1.
try:
MAX_COMPRESSION_CACHE_SESSIONS = max(
1, int(os.environ.get("HEADROOM_COMPRESSION_CACHE_MAX_SESSIONS", "500"))
)
except ValueError:
MAX_COMPRESSION_CACHE_SESSIONS = 500
# Idle TTL for per-session compression caches. Eviction is bust-free only
# once the provider's own prompt cache has lapsed, so this must exceed the
# LONGEST provider cache TTL Headroom serves — Anthropic's 1h extended
# breakpoint (3600s), not just the common 5m ephemeral cache. Evicting
# earlier would itself cause the bust this state exists to prevent: the
# session returns, the provider still holds the old bytes, but the map that
# replays them is gone. The cache must also outlive the prefix TRACKER's
# session TTL (600s): after the tracker expires, `apply_cached`'s
# byte-identical swap is the only thing still protecting the provider
# prefix. Default 3900s = 1h + 5m grace. Deployments that never opt into
# the 1h breakpoint can lower it via HEADROOM_COMPRESSION_CACHE_TTL_SECONDS.
try:
# Floor of 600s: never below the prefix tracker's session TTL, or the
# sweep could reclaim the byte-identical swap map while it is the only
# remaining protection for a still-live provider prefix (see above).
# Non-finite floats ("nan"/"inf") parse but poison every idle comparison,
# so they are rejected like any other unparseable value.
_ttl_env = float(os.environ.get("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", "3900"))
if _ttl_env != _ttl_env or _ttl_env in (float("inf"), float("-inf")):
raise ValueError("non-finite TTL")
COMPRESSION_CACHE_TTL_SECONDS = max(600.0, _ttl_env)
except ValueError:
COMPRESSION_CACHE_TTL_SECONDS = 3900.0
# Entries per session compression cache. 10k covers a single conversation with
# ~2x headroom even at a 1M-token context (a compressible tool_result is at
# least a few hundred tokens, so at most ~5k can be live at once). Raise via
# HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES only for workloads that fan many
# concurrent conversations into ONE session id (shared fallback ids, heavy
# subagent fan-out) — entry LRU is hit-refreshed, so an undersized cap shows
# up as misses on still-live entries, i.e. prefix-cache busts. Floor of 100.
try:
COMPRESSION_CACHE_MAX_ENTRIES = max(
100, int(os.environ.get("HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES", "10000"))
)
except ValueError:
COMPRESSION_CACHE_MAX_ENTRIES = 10000
# ---------------------------------------------------------------------------

View file

@ -25,6 +25,7 @@ actually reports.
from __future__ import annotations
import asyncio
import logging
from dataclasses import dataclass, field
from datetime import datetime, timezone
@ -455,10 +456,17 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
from headroom.proxy.output_savings import get_recorder
_rec = get_recorder()
_rec.record_from_labels(outcome.transforms_applied, outcome.output_tokens)
output_tokens_saved_est = _rec.estimate_request_savings(
outcome.transforms_applied, outcome.output_tokens
)
def _record_and_estimate() -> int:
_rec.record_from_labels(outcome.transforms_applied, outcome.output_tokens)
return _rec.estimate_request_savings(
outcome.transforms_applied, outcome.output_tokens
)
# Both calls take the recorder lock, and the every-Nth record also
# does a full read-modify-write of the ledger file — run them
# together off the event loop (#18) so a slow flush can't stall it.
output_tokens_saved_est = await asyncio.to_thread(_record_and_estimate)
except Exception: # pragma: no cover - defensive
pass

View file

@ -39,6 +39,7 @@ Pure module: no I/O except explicit ``load``/``save``.
from __future__ import annotations
import json
import logging
import math
from dataclasses import asdict, dataclass, field
from typing import Any
@ -68,6 +69,8 @@ from .output_savings_policy import (
stratum_label as stratum_label,
)
logger = logging.getLogger(__name__)
@dataclass
class _Accum:
@ -328,9 +331,13 @@ class SavingsLedger:
def save(self, path: Any) -> None:
from pathlib import Path
from headroom import fsutil
p = Path(path)
p.parent.mkdir(parents=True, exist_ok=True)
p.write_text(json.dumps(self.to_dict(), separators=(",", ":")))
# fsutil.write_text is atomic (temp file + os.replace), so a crash
# mid-write cannot truncate the ledger already on disk (#18).
fsutil.write_text(p, json.dumps(self.to_dict(), separators=(",", ":")))
@classmethod
def load(cls, path: Any) -> SavingsLedger:
@ -341,7 +348,11 @@ class SavingsLedger:
return cls()
try:
return cls.from_dict(json.loads(p.read_text()))
except (json.JSONDecodeError, ValueError, OSError):
except (json.JSONDecodeError, ValueError, OSError) as exc:
# Fail open (empty ledger), but surface the loss — silently
# swallowing a corrupt file made lost history indistinguishable
# from no history yet (#18).
logger.warning("output-savings ledger %s unreadable, starting empty: %s", p, exc)
return cls()

View file

@ -36,6 +36,7 @@ import os
import sys
import threading
import time
from collections import OrderedDict
from collections.abc import Callable
from dataclasses import fields, is_dataclass, replace
from datetime import datetime, timezone
@ -128,6 +129,8 @@ from headroom.proxy.cost import (
merge_cost_stats, # noqa: F401
)
from headroom.proxy.helpers import (
COMPRESSION_CACHE_MAX_ENTRIES,
COMPRESSION_CACHE_TTL_SECONDS,
COMPRESSION_TIMEOUT_SECONDS, # noqa: F401
EAGER_PRELOAD_TIMEOUT_SECONDS,
MAX_COMPRESSION_CACHE_SESSIONS, # noqa: F401
@ -1028,7 +1031,14 @@ class HeadroomProxy(
# `CompressionCache` instances have their own internal lock guarding
# `_cache`/`_stable_hashes`/`_first_seen` against concurrent
# async-dispatched requests for the same session.
self._compression_caches: dict[str, CompressionCache] = {}
# Ordered by last access: `_get_compression_cache` moves a session to
# the end on every hit, so capacity eviction drops the idlest sessions
# — whose provider prefix cache has lapsed anyway — never a busy
# long-lived one. `_compression_cache_last_seen` drives the idle-TTL
# sweep in `_maybe_cleanup_compression_caches`.
self._compression_caches: OrderedDict[str, CompressionCache] = OrderedDict()
self._compression_cache_last_seen: dict[str, float] = {}
self._compression_caches_last_cleanup: float = time.time()
self._compression_caches_lock = threading.RLock()
self.logger = (
@ -1580,6 +1590,67 @@ class HeadroomProxy(
loop = asyncio.get_running_loop()
return await loop.run_in_executor(self._background_compression_executor, fn)
# How often the lazy TTL sweep in `_get_compression_cache` may run.
_COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS = 60.0
def _maybe_cleanup_compression_caches(self, now: float) -> None:
"""Evict per-session compression caches idle past their TTL.
Caller must hold `_compression_caches_lock`. Piggybacked on
`_get_compression_cache` (the same lazy-sweep pattern as
`PrefixCacheTrackerRegistry._maybe_cleanup`) so no background task is
needed: any traffic at all keeps memory tracking the active-session
window, and a fully idle process has no memory pressure worth a timer.
A session idle longer than `COMPRESSION_CACHE_TTL_SECONDS` has
outlived the provider prompt cache its entries protect the default
exceeds Anthropic's 1h extended breakpoint, the longest provider TTL
served so evicting it cannot bust anything: the provider already
forgot the prefix. If the session does return, the cost is one
cache-write turn (fail-open), which it was going to pay regardless.
The TTL must never be set below the prefix tracker's session TTL:
after the tracker expires, this cache's byte-identical swap is the
only remaining protection for a still-live provider prefix.
"""
if now - self._compression_caches_last_cleanup < (
self._COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS
):
return
self._compression_caches_last_cleanup = now
# Skip sessions with a turn in flight (session_turn_lock held): popping
# one would hand its retry a FRESH cache with a NEW lock — straggler
# and retry then run unserialized against the same tracker, and the
# retry's empty cache recompresses previously-returned content into
# different bytes. An in-flight session is by definition not idle; it
# will be swept on a later pass once genuinely quiet.
expired = [
sid
for sid, seen in self._compression_cache_last_seen.items()
if now - seen > COMPRESSION_CACHE_TTL_SECONDS
and (cache := self._compression_caches.get(sid)) is not None
and not cache.session_turn_lock.locked()
]
for sid in expired:
self._compression_caches.pop(sid, None)
self._compression_cache_last_seen.pop(sid, None)
if expired:
logger.info(
"Evicted %d compression caches idle > %.0fs (%d sessions remain)",
len(expired),
COMPRESSION_CACHE_TTL_SECONDS,
len(self._compression_caches),
)
def _peek_compression_cache(self, session_id: str) -> CompressionCache | None:
"""Return the session's cache if one exists — no create, no LRU bump.
For lookup paths that must not leave a footprint or distort access
recency (e.g. /v1/usage taking the session turn lock): an unknown
session answers None instead of allocating an empty cache.
"""
with self._compression_caches_lock:
return self._compression_caches.get(session_id)
def _get_compression_cache(self, session_id: str) -> CompressionCache:
"""Get or create a CompressionCache for a session.
@ -1588,27 +1659,55 @@ class HeadroomProxy(
for the same conversation) must return the **same** instance,
otherwise the per-session cache state splits and the two halves
diverge across requests.
Every access refreshes both the LRU position and the idle-TTL clock,
so eviction capacity or TTL only ever hits sessions that have
gone quiet. Losing one costs at most a single cache-write turn
upstream; it never fails a request.
"""
with self._compression_caches_lock:
if session_id not in self._compression_caches:
now = time.time()
self._maybe_cleanup_compression_caches(now)
cache = self._compression_caches.get(session_id)
if cache is None:
from headroom.cache.compression_cache import CompressionCache
# Evict oldest caches if at capacity
# Evict the least-recently-used quarter at capacity. The
# OrderedDict is maintained in access order, so the front is
# always the idlest session — never a busy long-lived one.
# Sessions with a turn in flight (session_turn_lock held) are
# skipped: popping one splits its lock across two cache
# instances and desyncs the straggler from its retry (see the
# TTL sweep's comment). If every candidate is mid-turn, no
# eviction happens this round — briefly exceeding the cap is
# cheaper than a guaranteed prefix bust.
if len(self._compression_caches) >= MAX_COMPRESSION_CACHE_SESSIONS:
# Remove oldest quarter to amortize cleanup cost
oldest_keys = list(self._compression_caches.keys())[
: MAX_COMPRESSION_CACHE_SESSIONS // 4
]
for key in oldest_keys:
del self._compression_caches[key]
logger.info(
"Evicted %d compression caches (exceeded %d max sessions)",
len(oldest_keys),
MAX_COMPRESSION_CACHE_SESSIONS,
evict_count = min(
max(1, MAX_COMPRESSION_CACHE_SESSIONS // 4),
len(self._compression_caches),
)
evictable = [
sid
for sid, c in self._compression_caches.items()
if not c.session_turn_lock.locked()
][:evict_count]
for sid in evictable:
del self._compression_caches[sid]
self._compression_cache_last_seen.pop(sid, None)
if evictable:
logger.info(
"Evicted %d least-recently-used compression caches "
"(exceeded %d max sessions)",
len(evictable),
MAX_COMPRESSION_CACHE_SESSIONS,
)
self._compression_caches[session_id] = CompressionCache()
return self._compression_caches[session_id]
cache = CompressionCache(max_entries=COMPRESSION_CACHE_MAX_ENTRIES)
self._compression_caches[session_id] = cache
else:
self._compression_caches.move_to_end(session_id)
self._compression_cache_last_seen[session_id] = now
return cache
def _setup_code_aware(self, config: ProxyConfig, transforms: list) -> str:
"""Set up code-aware compression if enabled.
@ -5183,6 +5282,13 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
async def compress_messages(request: Request):
return await proxy.handle_compress(request)
# Sidecar-mode usage relay: same exposure policy as /v1/compress — the two
# form one contract (compress returns the bytes, usage reports what the
# provider said about them), so they must be reachable from the same place.
@app.post("/v1/usage", dependencies=_compress_dependencies)
async def compress_usage(request: Request):
return await proxy.handle_compress_usage(request)
register_provider_routes(app, proxy)
return app

View file

@ -0,0 +1,185 @@
"""Session-turn engine — the single cache-management brain for both modes.
One conversation turn, from the cache's point of view, is always the same
three-step dance regardless of who owns the upstream call:
1. **Prepare** (:func:`prepare_turn`): decide how many leading messages are
frozen, mark the stable prefix, and swap previously-computed compressed
bytes into the working copy (``apply_cached`` "Zone 1").
2. Run the compression pipeline over the prepared input (owned by the
caller: the proxy handlers wrap it in background/cold-start/backpressure
orchestration, the sidecar path runs it inline on the executor).
3. **Finalize** (:func:`finalize_turn`): replay last turn's exact
previously-forwarded/returned prefix over any residual drift the pipeline
introduced (``overlay_cached_prefix``), so the bytes that leave the
process are byte-identical to what the provider already cached.
Historically the proxy request handlers (anthropic + openai token mode) and
the sidecar ``/v1/compress`` session path each carried their own inline copy
of steps 1 and 3. This module is the shared implementation: a
cache-management fix landed here reaches BOTH modes at once.
Freeze policies
---------------
The one deliberate behavioural difference between the modes lives in step 1,
and it is a *policy parameter*, not a fork of the code:
``FREEZE_POLICY_CONFIRMED_CLAMP`` ``min(tracker_frozen, cache_count)``.
The proxy sees the provider's responses, so ``tracker_frozen`` is the
provider-confirmed cached prefix (from ``cache_read_input_tokens``).
Freezing is clamped by BOTH bounds: never past what the provider
actually has cached (freezing more would forgo compression of content
that is not yet cache-protected the #327 posture), and never past what
the local cache can byte-replay (freezing a message whose entry was
evicted would pass through raw original bytes).
``FREEZE_POLICY_REPLAYABLE`` ``max(cache_count, explicit_frozen or 0)``.
Freeze everything the local cache can byte-replay. Used by callers with
no provider-confirmed count to clamp against: the sidecar ``/v1/compress``
endpoint (it never sees the provider's response — whatever it previously
RETURNED is the provider's cache contract, so every already-returned
message must come back byte-identical), and the OpenAI proxy token path
(its tracker feeds cache mode, not token mode). Recompressing an
already-returned message even into a *smaller* form is a bust: the
drift was observed in practice, and ``overlay_cached_prefix``'s
non-inflation guard cannot repair a shrunken form (replaying the larger
original bytes would "inflate" the candidate). Freezing the entire
locally-replayable prefix eliminates that recompression outright.
Over-freezing relative to the provider's real cache only forgoes tail
compression; it can never bust. An explicit ``frozen_message_count``
from the caller still wins when larger the caller may know more about
the provider cache than local state does.
Why the Anthropic proxy path cannot simply adopt the replayable posture: its
provider-confirmed clamp deliberately KEEPS not-yet-cached content
compressible, and its overlay inputs (tracker snapshots) are refreshed on
every response, so drift repair is reliable there. Each posture is correct
for the information its mode actually has.
"""
from __future__ import annotations
import logging
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any
from headroom.cache.prefix_tracker import overlay_cached_prefix
logger = logging.getLogger(__name__)
FREEZE_POLICY_CONFIRMED_CLAMP = "confirmed_clamp"
FREEZE_POLICY_REPLAYABLE = "replayable"
_FREEZE_POLICIES = (FREEZE_POLICY_CONFIRMED_CLAMP, FREEZE_POLICY_REPLAYABLE)
@dataclass(frozen=True)
class TurnPrep:
"""Result of :func:`prepare_turn`.
``frozen_message_count`` is what the pipeline must be told to skip;
``pipeline_input`` is the working copy with previously-compressed bytes
swapped in (never the caller's list — ``apply_cached`` copies).
"""
frozen_message_count: int
pipeline_input: list[dict[str, Any]]
@dataclass(frozen=True)
class TurnFinal:
"""Result of :func:`finalize_turn`.
``messages`` are the bytes to forward/return; ``replayed`` says whether
the overlay restored last turn's prefix over pipeline drift; ``tokens``
is the recount of ``messages`` when a ``count_tokens`` hook was supplied
and the overlay actually fired (None otherwise the pipeline's own
count is still valid when nothing was replaced).
"""
messages: list[dict[str, Any]]
replayed: bool
tokens: int | None = None
def prepare_turn(
comp_cache: Any,
messages: list[dict[str, Any]],
*,
policy: str,
tracker_frozen: int | None = None,
explicit_frozen: int | None = None,
) -> TurnPrep:
"""Freeze decision + stable marking + cached-byte swap for one turn.
Args:
comp_cache: the session's ``CompressionCache``.
messages: the caller's RAW message list (never mutated).
policy: ``FREEZE_POLICY_CONFIRMED_CLAMP`` or ``FREEZE_POLICY_REPLAYABLE``
see the module docstring for why they differ.
tracker_frozen: provider-confirmed frozen count (proxy policy only;
``None`` means "nothing confirmed" and freezes 0 there).
explicit_frozen: caller-pinned frozen count (sidecar policy only;
wins when larger than the locally-derived bound).
"""
if policy not in _FREEZE_POLICIES:
raise ValueError(f"unknown freeze policy: {policy!r}")
cache_count = comp_cache.compute_frozen_count(messages)
if policy == FREEZE_POLICY_CONFIRMED_CLAMP:
# Never freeze past the provider-confirmed prefix, and never past
# what local state can byte-replay.
frozen = min(tracker_frozen or 0, cache_count)
else:
# Freeze the entire locally-replayable prefix; an explicit caller
# pin may extend it (the caller vouches the provider cached those
# exact raw bytes, so passing them through untouched is correct).
frozen = max(cache_count, explicit_frozen or 0)
comp_cache.mark_stable_from_messages(messages, frozen)
return TurnPrep(
frozen_message_count=frozen,
pipeline_input=comp_cache.apply_cached(messages),
)
def finalize_turn(
result_messages: list[dict[str, Any]],
original_messages: list[dict[str, Any]],
prev_original: list[dict[str, Any]] | None,
prev_returned: list[dict[str, Any]] | None,
*,
count_tokens: Callable[[list[dict[str, Any]]], int] | None = None,
) -> TurnFinal:
"""Replay last turn's exact forwarded/returned prefix over pipeline drift.
``overlay_cached_prefix`` self-guards (positional alignment, append-only
shape, non-inflation), so calling this is always safe: when replay is not
provably correct it returns the pipeline's own output unchanged.
``count_tokens`` is invoked only when the overlay actually replaced
bytes the pipeline's own token count is still accurate otherwise. A
failing hook falls back to "no recount" rather than failing the turn.
"""
final = overlay_cached_prefix(result_messages, original_messages, prev_original, prev_returned)
replayed = final != result_messages
tokens: int | None = None
if replayed and count_tokens is not None:
try:
tokens = count_tokens(final)
except Exception as e:
# Fail-open: the turn still forwards, but the caller keeps the
# pipeline's count of messages that are NOT being forwarded —
# tokens_saved accounting is stale for this turn. Loud, not
# silent: a tokenizer that cannot count the replayed form is a
# bug worth surfacing even though it must not fail the request.
logger.warning(
"finalize_turn: token recount of replayed prefix failed "
"(%s: %s); keeping the pipeline's pre-overlay count",
type(e).__name__,
e,
)
tokens = None
return TurnFinal(messages=final, replayed=replayed, tokens=tokens)

View file

@ -188,6 +188,19 @@ def build_httpx_verify() -> ssl.SSLContext | bool:
return True
def build_urlopen_context() -> ssl.SSLContext | None:
"""Return Headroom's configured TLS context for ``urllib.request.urlopen``.
``urlopen`` already handles Python's default trust configuration when no
explicit context is passed. Return only a custom context here so callers
retain that default while sharing Headroom's corporate CA and strict-mode
handling when it is configured.
"""
verify = build_httpx_verify()
return verify if isinstance(verify, ssl.SSLContext) else None
def apply_global_tls_relaxation() -> bool:
"""Strip ``VERIFY_X509_STRICT`` from urllib3's context builder when opted in.

View file

@ -457,6 +457,23 @@ def _try_detect_html(content: str) -> DetectionResult | None:
)
def _is_search_result_line(line: str) -> bool:
"""True when a line looks like ``path:line:content`` grep output.
The bare ``^[^\\s:]+:\\d+:`` shape also matches ISO-8601 timestamps
(``T09:57:59``) and XML-ish wrappers harnesses prepend to user turns
(Copilot CLI's ``<current_datetime>…`` line), which misroutes prose to
the SearchCompressor and that compressor keeps only matching lines,
deleting the rest. So the pre-colon segment must additionally look like
a file path: no angle brackets and no ``=`` (rules out markup tags and
``key=value:12:`` log lines).
"""
if not _SEARCH_RESULT_PATTERN.match(line):
return False
prefix = line.split(":", 1)[0]
return "<" not in prefix and ">" not in prefix and "=" not in prefix
def _try_detect_search(content: str) -> DetectionResult | None:
"""Try to detect grep/ripgrep search results."""
lines = content.split("\n")[:100] # Check first 100 lines
@ -465,10 +482,16 @@ def _try_detect_search(content: str) -> DetectionResult | None:
matching_lines = 0
for line in lines:
if line.strip() and _SEARCH_RESULT_PATTERN.match(line):
if line.strip() and _is_search_result_line(line):
matching_lines += 1
if matching_lines == 0:
# Absolute floor: a single coincidental `word:digits:` line (a timestamp,
# a URL, a time literal inside prose) must not classify a whole payload as
# search results — the SearchCompressor drops every non-matching line, so
# a false positive is data loss. A genuine one-line grep result loses
# nothing by staying uncompressed: all of its lines match, so the
# compressor would have kept it verbatim anyway.
if matching_lines < 2:
return None
# Calculate confidence based on proportion of matching lines

View file

@ -532,6 +532,20 @@ def _tool_call_args_text(raw: Any) -> str:
return " ".join(text.split())[:300]
def read_protection_enabled() -> bool:
"""True when HEADROOM_PROTECT_READS opts into byte-exact file-read protection.
Shared by every request path (chat/Anthropic ``ContentRouter.apply`` and the
OpenAI Responses units path) so the flag means the same thing everywhere.
"""
return os.environ.get("HEADROOM_PROTECT_READS", "0").strip().lower() not in (
"0",
"",
"false",
"no",
)
def _tool_call_command_text(raw: Any) -> str:
"""Extract the raw shell command from a tool call's args, if present.
@ -1933,7 +1947,19 @@ class ContentRouter(Transform):
# we match that posture with a dedicated lock rather than relying on
# GIL atomicity (which would not protect the read-then-evict sequence).
self._frozen_verdicts: dict[int, bool] = {}
self._frozen_verdicts_max = 4096
# The store is process-wide (one router per pipeline, shared by every
# session), so the cap must scale with the number of CONCURRENT
# sessions, not one user's workload: at org scale (many users behind
# one sidecar) 4096 churns in minutes and FIFO eviction lets tightened
# thresholds flip a still-cached block's verdict — a prefix bust.
# Read at construction so tests and multi-tenant deployments can size
# it via HEADROOM_FROZEN_VERDICTS_MAX without a module reload.
try:
self._frozen_verdicts_max = max(
256, int(os.environ.get("HEADROOM_FROZEN_VERDICTS_MAX", "4096"))
)
except ValueError:
self._frozen_verdicts_max = 4096
self._frozen_lock = threading.Lock()
# Reset verdicts whenever the shadowed cache is cleared.
self._cache.register_on_clear(self._clear_frozen_verdicts)
@ -4827,12 +4853,7 @@ class ContentRouter(Transform):
# Type-specific by design: grep/test/ls output stays compressible, so the
# cache-mode delta still compresses whenever the newest turn is NOT a read.
self._protect_read_tool_ids = set()
if os.environ.get("HEADROOM_PROTECT_READS", "0").strip().lower() not in (
"0",
"",
"false",
"no",
):
if read_protection_enabled():
# Use _tool_call_commands (the parsed shell command), NOT
# _tool_call_args (a compact free-text blob that, for OpenAI-style
# JSON-string args, is the raw ``{"command": ...}`` JSON — on which
@ -4854,12 +4875,7 @@ class ContentRouter(Transform):
# cat/sed/head code reads are protected on ANY model/harness, not just
# those that emit tool-call/tool_result blocks.
self._protect_read_msg_indices: set[int] = set()
if os.environ.get("HEADROOM_PROTECT_READS", "0").strip().lower() not in (
"0",
"",
"false",
"no",
):
if read_protection_enabled():
for _idx, _m in enumerate(messages):
if _m.get("role") != "user":
continue

26
tests/fixtures/headroom-issue-3054.json vendored Normal file
View file

@ -0,0 +1,26 @@
{
"issue": 3054,
"url": "https://github.com/headroomlabs-ai/headroom/issues/3054",
"old_serena_args": [
"--from",
"git+https://github.com/oraios/serena",
"serena",
"start-mcp-server",
"--project-from-cwd",
"--context",
"claude-code",
"--open-web-dashboard",
"False"
],
"recommended_serena_args": [
"--from",
"serena-agent",
"serena",
"start-mcp-server",
"--project-from-cwd",
"--context",
"claude-code",
"--open-web-dashboard",
"False"
]
}

View file

@ -1094,3 +1094,74 @@ def test_response_cache_keys_on_lookup_messages_not_mutated():
assert cache.set_messages == cache.get_messages
# And specifically the raw lookup messages, not the scanner's rewrite.
assert cache.set_messages == [{"role": "user", "content": "hello"}]
# --------------------------------------------------------------------------- #
# Backpressure must not bust the provider prompt cache: the compression #
# pipeline is skipped under saturation, but the previously-forwarded #
# (compressed) prefix must still be replayed byte-identical. Forwarding raw #
# originals would mismatch the bytes the provider cached — busting every #
# gated session's prefix exactly when the proxy is busiest. #
# --------------------------------------------------------------------------- #
def test_backpressure_passthrough_replays_cached_prefix(stage_log_capture):
prev_original = [{"role": "user", "content": "ORIGINAL " * 6000}]
prev_forwarded = [{"role": "user", "content": "[compressed-form]"}]
async def _run() -> None:
sem = asyncio.Semaphore(1)
await sem.acquire() # saturate: the request's acquire will time out
handler = _DummyAnthropicHandler(anthropic_pre_upstream_sem=sem)
handler.config.optimize = True
handler.config.anthropic_pre_upstream_acquire_timeout_seconds = 0.01
handler.anthropic_pipeline = SimpleNamespace(apply=MagicMock())
tracker = SimpleNamespace(
_cached_token_count=0,
get_frozen_message_count=lambda: 0,
get_last_original_messages=lambda: copy.deepcopy(prev_original),
get_last_forwarded_messages=lambda: copy.deepcopy(prev_forwarded),
update_from_response=lambda *a, **k: None,
record_request=lambda *a, **k: None,
)
handler.session_tracker_store = SimpleNamespace(
compute_session_id=lambda *a, **k: "sess-1",
get_or_create=lambda *a, **k: tracker,
resolve_tracker=lambda *a, **k: tracker,
)
forwarded_bodies: list[dict] = []
orig_retry = handler._retry_request
async def _capturing_retry(method, url, headers, body, **kw):
forwarded_bodies.append(copy.deepcopy(body))
return await orig_retry(method, url, headers, body, **kw)
handler._retry_request = _capturing_retry
req = _build_request(
{
"model": "claude-3-5-sonnet-latest",
"messages": copy.deepcopy(prev_original)
+ [{"role": "user", "content": "next turn"}],
},
{"authorization": "Bearer sk-ant-api-test"},
)
try:
response = await handler.handle_anthropic_messages(req)
assert response.status_code == 200
# Saturation must still skip the CPU-bound pipeline...
assert not handler.anthropic_pipeline.apply.called
finally:
sem.release()
assert forwarded_bodies, "request never reached upstream"
sent = forwarded_bodies[-1]["messages"]
# ...but the forwarded prefix must be last turn's exact bytes, not the
# raw original (which the provider never cached).
assert sent[0]["content"] == "[compressed-form]"
assert sent[-1]["content"] == "next turn"
with _tokenizer_patch():
anyio.run(_run)

View file

@ -208,6 +208,36 @@ class TestSemanticCache:
entry = cache.get("What time is it?")
assert entry is None
def test_empty_query_never_semantic_matches(self):
"""An empty extracted query must not trigger cross-context false hits.
The query is the last user message; in agent/tool traffic most turns are
tool_result continuations whose extracted query is "". A real embedder
maps "" to a fixed non-zero vector, so without a guard every empty-query
turn would be ~identical to every other and serve one conversation's
response to an unrelated one. An empty query may only ever hit via the
exact messages_hash (which is context-complete).
"""
def const_embedding(text: str) -> list[float]:
# Realistic: a non-zero, identical vector for every input (incl. "").
return [0.5, 0.5, 0.5]
config = SemanticCacheConfig(similarity_threshold=0.9)
cache = SemanticCache(config, embedding_fn=const_embedding)
# Conversation A: an empty-query turn (unique full-context hash).
cache.put("", "response-A", messages_hash="ctxA")
# Conversation B: a different empty-query turn — must NOT get A's answer.
assert cache.get("", messages_hash="ctxB") is None
# Its own exact hash still works.
assert cache.get("", messages_hash="ctxA").response == "response-A"
# A whitespace-only query is treated the same as empty.
cache.put(" \n\t", "response-C", messages_hash="ctxC")
assert cache.get(" ", messages_hash="ctxD") is None
class TestSemanticCacheLayer:
"""Test SemanticCacheLayer functionality."""

View file

@ -0,0 +1,291 @@
from __future__ import annotations
import json
import os
import sys
from pathlib import Path
import pytest
from click.testing import CliRunner
from headroom.cli.main import main
from headroom.mcp_registry import ClaudeRegistrar, build_serena_spec
from headroom.mcp_registry.ledger import headroom_installed_matching
FIXTURE = Path(__file__).parents[1] / "fixtures" / "headroom-issue-3054.json"
def _setup(monkeypatch, tmp_path: Path):
config = tmp_path / ".claude.json"
config.write_text(
json.dumps(
{
"oauthAccount": {"email": "user@example.com"},
"mcpServers": {
"serena": {
"command": "uvx",
"args": json.loads(FIXTURE.read_text())["old_serena_args"],
},
"other": {"command": "other", "args": []},
},
"projects": {"/repo": {"trust": True}},
}
)
)
registrar = ClaudeRegistrar(claude_cli=None, home_dir=tmp_path)
monkeypatch.setattr("headroom.mcp_registry.ClaudeRegistrar", lambda: registrar)
ledger = tmp_path / "ledger.json"
monkeypatch.setattr("headroom.mcp_registry.ledger.ledger_path", lambda: ledger)
return config, ledger
def test_issue_fixture_reconcile_is_base_fail_head_pass(monkeypatch, tmp_path: Path):
config, _ = _setup(monkeypatch, tmp_path)
fixture = json.loads(FIXTURE.read_text())
recommended = build_serena_spec("claude-code")
assert list(recommended.args) == fixture["recommended_serena_args"]
assert CliRunner().invoke(main, ["mcp", "reconcile"]).exit_code == 0
adopted = CliRunner().invoke(main, ["mcp", "reconcile", "--adopt"])
assert adopted.exit_code == 0, adopted.output
assert json.loads(config.read_text())["mcpServers"]["serena"]["args"] == list(recommended.args)
def test_read_only_preserves_config_and_ledger_bytes_and_mtimes(monkeypatch, tmp_path: Path):
config, ledger = _setup(monkeypatch, tmp_path)
ledger.write_text("not json")
before = (
config.read_bytes(),
ledger.read_bytes(),
os.stat(config).st_mtime_ns,
os.stat(ledger).st_mtime_ns,
)
result = CliRunner().invoke(main, ["mcp", "reconcile"])
assert result.exit_code == 0, result.output
after = (
config.read_bytes(),
ledger.read_bytes(),
os.stat(config).st_mtime_ns,
os.stat(ledger).st_mtime_ns,
)
assert after == before
assert "--adopt" in result.output
def test_adopt_preserves_unrelated_config_and_records_ownership(monkeypatch, tmp_path: Path):
config, ledger = _setup(monkeypatch, tmp_path)
result = CliRunner().invoke(main, ["mcp", "reconcile", "--adopt"])
assert result.exit_code == 0, result.output
data = json.loads(config.read_text())
assert data["oauthAccount"] == {"email": "user@example.com"}
assert data["projects"] == {"/repo": {"trust": True}}
assert data["mcpServers"]["other"] == {"command": "other", "args": []}
assert data["mcpServers"]["serena"]["args"] == list(build_serena_spec("claude-code").args)
assert json.loads(ledger.read_text())["agents"]["claude"]["serena"]["fingerprint"]
@pytest.mark.parametrize(
"contents",
[
"not json",
"[]",
'{"agents": null}',
'{"agents": []}',
'{"agents": {"claude": null}}',
'{"agents": {"claude": []}}',
'{"agents": {"claude": {"serena": null}}}',
],
)
def test_malformed_ledger_blocks_adopt_before_config_write(
monkeypatch, tmp_path: Path, contents: str
):
config, ledger = _setup(monkeypatch, tmp_path)
before = config.read_bytes()
ledger.write_text(contents)
result = CliRunner().invoke(main, ["mcp", "reconcile", "--adopt"])
assert result.exit_code != 0
assert "ledger" in result.output.lower()
assert config.read_bytes() == before
def test_corrupt_ledger_is_tolerated_by_read_only(monkeypatch, tmp_path: Path):
_, ledger = _setup(monkeypatch, tmp_path)
ledger.write_text('{"agents": []}')
result = CliRunner().invoke(main, ["mcp", "reconcile"])
assert result.exit_code == 0, result.output
def test_reconcile_rejects_absent_claude(monkeypatch, tmp_path: Path):
_, _ = _setup(monkeypatch, tmp_path)
registrar = ClaudeRegistrar(claude_cli=None, home_dir=tmp_path)
monkeypatch.setattr(registrar, "detect", lambda: False)
monkeypatch.setattr("headroom.mcp_registry.ClaudeRegistrar", lambda: registrar)
result = CliRunner().invoke(main, ["mcp", "reconcile", "--adopt"])
assert result.exit_code != 0
assert "claude is not detected" in result.output
def test_reconcile_adopt_preserves_malformed_config(monkeypatch, tmp_path: Path):
config, _ = _setup(monkeypatch, tmp_path)
config.write_text("not json")
before = config.read_bytes()
result = CliRunner().invoke(main, ["mcp", "reconcile", "--adopt"])
assert result.exit_code != 0
assert config.read_bytes() == before
def test_adopt_rejects_malformed_modern_before_touching_valid_legacy(monkeypatch, tmp_path: Path):
modern = tmp_path / ".claude.json"
legacy = tmp_path / ".claude" / "mcp.json"
legacy.parent.mkdir()
modern.write_text("not json")
legacy.write_text(
json.dumps(
{
"mcpServers": {
"serena": {"command": "uvx", "args": ["--from", "user"]},
"other": {"command": "other"},
}
}
)
)
registrar = ClaudeRegistrar(claude_cli=None, home_dir=tmp_path)
monkeypatch.setattr("headroom.mcp_registry.ClaudeRegistrar", lambda: registrar)
ledger = tmp_path / "ledger.json"
monkeypatch.setattr("headroom.mcp_registry.ledger.ledger_path", lambda: ledger)
before = (modern.read_bytes(), legacy.read_bytes())
result = CliRunner().invoke(main, ["mcp", "reconcile", "--adopt"])
assert result.exit_code != 0
assert "not valid JSON" in result.output
assert (modern.read_bytes(), legacy.read_bytes()) == before
def test_adopt_rejects_non_dict_mcp_servers_in_legacy_root(monkeypatch, tmp_path: Path):
modern, _ = _setup(monkeypatch, tmp_path)
legacy = tmp_path / ".claude" / "mcp.json"
legacy.parent.mkdir()
legacy.write_text(json.dumps({"mcpServers": []}))
before = (modern.read_bytes(), legacy.read_bytes())
result = CliRunner().invoke(main, ["mcp", "reconcile", "--adopt"])
assert result.exit_code != 0
assert "non-object mcpServers" in result.output
assert (modern.read_bytes(), legacy.read_bytes()) == before
def test_unreadable_ledger_blocks_adopt_without_partial_mutation(monkeypatch, tmp_path: Path):
config, ledger = _setup(monkeypatch, tmp_path)
ledger.write_text(json.dumps({"agents": {}}))
before = (config.read_bytes(), ledger.read_bytes())
original_read_text = Path.read_text
def unreadable(path: Path, *args, **kwargs):
if path == ledger:
raise PermissionError("test unreadable ledger")
return original_read_text(path, *args, **kwargs)
monkeypatch.setattr(Path, "read_text", unreadable)
result = CliRunner().invoke(main, ["mcp", "reconcile", "--adopt"])
assert result.exit_code != 0
assert "unreadable" in result.output
assert (config.read_bytes(), ledger.read_bytes()) == before
@pytest.mark.parametrize("state", ["absent", "matching", "user-drift", "headroom-drift"])
@pytest.mark.parametrize("adopt", [False, True])
def test_reconcile_state_matrix(monkeypatch, tmp_path: Path, state: str, adopt: bool):
config, ledger = _setup(monkeypatch, tmp_path)
data = json.loads(config.read_text())
recommended = build_serena_spec("claude-code")
owned_spec = None
if state == "absent":
del data["mcpServers"]["serena"]
elif state == "matching":
data["mcpServers"]["serena"] = {
"command": recommended.command,
"args": list(recommended.args),
}
elif state == "user-drift":
data["mcpServers"]["serena"]["args"] = ["--from", "user-managed"]
elif state == "headroom-drift":
from headroom.mcp_registry.ledger import record_install
stale = build_serena_spec("claude-code")
stale.args = ("--from", "headroom-installed-old")
owned_spec = stale
data["mcpServers"]["serena"] = {
"command": stale.command,
"args": list(stale.args),
}
record_install("claude", stale, path=ledger)
config.write_text(json.dumps(data))
if owned_spec is not None:
assert headroom_installed_matching("claude", owned_spec, path=ledger)
result = CliRunner().invoke(main, ["mcp", "reconcile"] + (["--adopt"] if adopt else []))
assert result.exit_code == 0, result.output
observed = json.loads(config.read_text())["mcpServers"].get("serena")
ownership = observed is not None and headroom_installed_matching(
"claude",
build_serena_spec("claude-code") if observed["args"] == list(recommended.args) else None,
path=ledger,
)
if adopt:
assert observed == {
"command": recommended.command,
"args": list(recommended.args),
}
assert ownership
assert "Adopted Headroom" in result.output
elif state == "headroom-drift":
assert observed["args"] == ["--from", "headroom-installed-old"]
assert headroom_installed_matching("claude", owned_spec, path=ledger)
assert ownership is False
assert "observed: present" in result.output
else:
assert not ownership
assert "Serena reconciliation for Claude" in result.output
def test_only_adopt_is_a_reconcile_mutation(monkeypatch, tmp_path: Path):
_setup(monkeypatch, tmp_path)
result = CliRunner().invoke(main, ["mcp", "reconcile", "--help"])
assert result.exit_code == 0
assert "--adopt" in result.output
for option in ("--acknowledge", "--clear", "--agent", "--server"):
assert option not in result.output
def test_ordinary_install_does_not_adopt_serena(monkeypatch, tmp_path: Path):
config, _ = _setup(monkeypatch, tmp_path)
before = config.read_bytes()
monkeypatch.setitem(sys.modules, "mcp", object())
registrar = ClaudeRegistrar(claude_cli=None, home_dir=tmp_path)
monkeypatch.setattr("headroom.mcp_registry.install.get_all_registrars", lambda: [registrar])
result = CliRunner().invoke(main, ["mcp", "install", "--agent", "claude"])
assert result.exit_code == 0, result.output
after = json.loads(config.read_text())
before_data = json.loads(before)
assert after["mcpServers"]["serena"] == before_data["mcpServers"]["serena"]
assert after["mcpServers"]["headroom"]["args"] == ["mcp", "serve"]
assert "mcp reconcile --adopt" not in result.output
def test_mcp_install_force_preserves_user_managed_serena(monkeypatch, tmp_path: Path):
config, _ = _setup(monkeypatch, tmp_path)
before = json.loads(config.read_text())["mcpServers"]["serena"]
monkeypatch.setitem(sys.modules, "mcp", object())
registrar = ClaudeRegistrar(claude_cli=None, home_dir=tmp_path)
monkeypatch.setattr("headroom.mcp_registry.install.get_all_registrars", lambda: [registrar])
result = CliRunner().invoke(main, ["mcp", "install", "--agent", "claude", "--force"])
assert result.exit_code == 0, result.output
assert json.loads(config.read_text())["mcpServers"]["serena"] == before

View file

@ -0,0 +1,124 @@
from __future__ import annotations
from pathlib import Path
from headroom.cli import wrap as wrap_cli
from headroom.mcp_registry import build_serena_spec
from headroom.mcp_registry.base import RegisterResult, RegisterStatus, ServerSpec
from headroom.mcp_registry.ledger import headroom_installed_matching, record_install
class _Registrar:
display_name = "Claude Code"
def __init__(self, current: ServerSpec | None, *, name: str = "claude"):
self.name = name
self.current = current
self.force_calls: list[bool] = []
def detect(self) -> bool:
return True
def get_server(self, name: str) -> ServerSpec | None:
return self.current if name == "serena" else None
def register_server(self, spec: ServerSpec, *, force: bool = False) -> RegisterResult:
self.force_calls.append(force)
if self.current == spec:
return RegisterResult(RegisterStatus.ALREADY, "matches")
if self.current is not None and not force:
return RegisterResult(RegisterStatus.MISMATCH, "different")
self.current = spec
return RegisterResult(RegisterStatus.REGISTERED, "updated")
def _quiet(monkeypatch):
monkeypatch.setattr(wrap_cli, "_ensure_serena_dashboard_disabled", lambda **kwargs: None)
monkeypatch.setattr(wrap_cli, "_inject_serena_instructions", lambda *args, **kwargs: None)
monkeypatch.setattr(wrap_cli, "_serena_project_skip_reason", lambda root: "test")
monkeypatch.setattr(wrap_cli, "_index_serena_project", lambda **kwargs: None)
monkeypatch.setattr(wrap_cli.shutil, "which", lambda name: "uvx" if name == "uvx" else None)
def test_automatic_wrap_migrates_owned_drift_and_recurs_to_noop(
monkeypatch, tmp_path: Path, capsys
):
_quiet(monkeypatch)
monkeypatch.setattr(
"headroom.mcp_registry.ledger.ledger_path", lambda: tmp_path / "ledger.json"
)
stale = ServerSpec("serena", "uvx", ("--from", "old"))
record_install("claude", stale)
registrar = _Registrar(stale)
wrap_cli._setup_serena_mcp(registrar, context="claude-code", verbose=True)
assert registrar.current == build_serena_spec("claude-code")
assert registrar.force_calls == [False, True]
assert headroom_installed_matching("claude", registrar.current)
capsys.readouterr()
wrap_cli._setup_serena_mcp(registrar, context="claude-code", verbose=True)
assert registrar.force_calls == [False, True, False]
def test_automatic_wrap_owned_drift_suggests_rerun_wrap(monkeypatch, tmp_path: Path, capsys):
_quiet(monkeypatch)
monkeypatch.setattr(
"headroom.mcp_registry.ledger.ledger_path", lambda: tmp_path / "ledger.json"
)
stale = ServerSpec("serena", "uvx", ("--from", "old"))
record_install("claude", stale)
class _FailedMigrationRegistrar(_Registrar):
def register_server(self, spec, *, force=False):
if force:
self.force_calls.append(force)
return RegisterResult(RegisterStatus.MISMATCH, "still different")
return super().register_server(spec, force=force)
wrap_cli._setup_serena_mcp(
_FailedMigrationRegistrar(stale), context="claude-code", verbose=True
)
output = capsys.readouterr().out
assert "run headroom wrap again" in output
assert "mcp reconcile --adopt" not in output
def test_automatic_wrap_preserves_user_managed_warning(monkeypatch, tmp_path: Path, capsys):
_quiet(monkeypatch)
monkeypatch.setattr(
"headroom.mcp_registry.ledger.ledger_path", lambda: tmp_path / "ledger.json"
)
user = ServerSpec("serena", "uvx", ("--from", "user"))
registrar = _Registrar(user)
wrap_cli._setup_serena_mcp(registrar, context="claude-code", verbose=True)
assert registrar.current == user
assert registrar.force_calls == [False]
assert "existing config differs" in capsys.readouterr().out
def test_automatic_wrap_recovers_from_malformed_ledger(monkeypatch, tmp_path: Path):
_quiet(monkeypatch)
ledger = tmp_path / "ledger.json"
ledger.write_text("not json")
monkeypatch.setattr("headroom.mcp_registry.ledger.ledger_path", lambda: ledger)
registrar = _Registrar(None)
wrap_cli._setup_serena_mcp(registrar, context="claude-code", verbose=True)
current = registrar.get_server("serena")
assert current == build_serena_spec("claude-code")
assert headroom_installed_matching("claude", current)
def test_non_claude_wrap_keeps_usable_remediation_hint(monkeypatch, tmp_path: Path, capsys):
_quiet(monkeypatch)
monkeypatch.setattr(
"headroom.mcp_registry.ledger.ledger_path", lambda: tmp_path / "ledger.json"
)
registrar = _Registrar(ServerSpec("serena", "uvx", ("--from", "user")), name="codex")
wrap_cli._setup_serena_mcp(registrar, context="codex", verbose=True)
output = capsys.readouterr().out
assert "update or remove the existing serena MCP entry" in output
assert "mcp reconcile --adopt" not in output

View file

@ -243,18 +243,24 @@ def test_unwrap_claude_restores_all_base_url_modes(runner: CliRunner) -> None:
"foundry_mode": False,
"vertex_mode": False,
"settings_path": settings_path,
# unwrap is the user asking for their settings back, so it drops
# every wrap session's ownership claim instead of deferring to a
# live sibling and silently doing nothing (#3205).
"force": True,
},
{
"previous": None,
"foundry_mode": True,
"vertex_mode": False,
"settings_path": settings_path,
"force": True,
},
{
"previous": None,
"foundry_mode": False,
"vertex_mode": True,
"settings_path": settings_path,
"force": True,
},
]

View file

@ -266,17 +266,16 @@ def test_wrap_copilot_prefers_existing_oauth_session(
assert result.exit_code == 0, result.output
env = captured["env"]
assert isinstance(env, dict)
assert env["COPILOT_PROVIDER_TYPE"] == "openai"
assert env["COPILOT_PROVIDER_BASE_URL"] == (
f"http://127.0.0.1:8787{_expected_project_prefix()}/v1"
)
assert env["COPILOT_PROVIDER_WIRE_API"] == "completions"
assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "gho-existing"
assert env["COPILOT_API_URL"] == f"http://127.0.0.1:8787{_expected_project_prefix()}"
assert "COPILOT_PROVIDER_TYPE" not in env
assert "COPILOT_PROVIDER_BASE_URL" not in env
assert "COPILOT_PROVIDER_WIRE_API" not in env
assert "COPILOT_PROVIDER_BEARER_TOKEN" not in env
assert env["GITHUB_COPILOT_API_URL"] == DEFAULT_API_URL
assert env["OPENAI_TARGET_API_URL"] == DEFAULT_API_URL
assert "COPILOT_PROVIDER_API_KEY" not in env
assert captured["openai_api_url"] == DEFAULT_API_URL
assert f"COPILOT_PROVIDER_API_URL={DEFAULT_API_URL}" in captured["env_vars_display"]
assert "COPILOT_AUTH_MODE=github-native" in captured["env_vars_display"]
@pytest.mark.parametrize(
@ -293,7 +292,7 @@ def test_wrap_copilot_oauth_defaults_wire_api_for_selected_model(
model: str,
expected_wire_api: str,
) -> None:
"""OAuth sessions use the same model-aware wire API default as subscriptions."""
"""Implicit OAuth leaves wire selection to Copilot's native router."""
_wrap_cli, main = wrap_modules
_clear_copilot_env(monkeypatch)
captured: dict[str, object] = {}
@ -315,8 +314,8 @@ def test_wrap_copilot_oauth_defaults_wire_api_for_selected_model(
assert result.exit_code == 0, result.output
env = captured["env"]
assert isinstance(env, dict)
assert env["COPILOT_PROVIDER_WIRE_API"] == expected_wire_api
assert f"COPILOT_PROVIDER_WIRE_API={expected_wire_api}" in captured["env_vars_display"]
assert "COPILOT_PROVIDER_WIRE_API" not in env
assert env["COPILOT_API_URL"].startswith("http://127.0.0.1:8787")
@pytest.mark.parametrize("wire_api", ["completions", "responses"])
@ -348,7 +347,8 @@ def test_wrap_copilot_oauth_honors_existing_wire_api(
assert result.exit_code == 0, result.output
env = captured["env"]
assert isinstance(env, dict)
assert env["COPILOT_PROVIDER_WIRE_API"] == wire_api
assert "COPILOT_PROVIDER_WIRE_API" not in env
assert env["COPILOT_API_URL"].startswith("http://127.0.0.1:8787")
def test_wrap_copilot_subscription_uses_github_auth_without_provider_key(
@ -869,7 +869,8 @@ def test_wrap_copilot_oauth_keeps_generic_endpoint_when_account_advertised(
assert result.exit_code == 0, result.output
env = captured["env"]
assert isinstance(env, dict)
assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "gho-oauth"
assert "COPILOT_PROVIDER_BEARER_TOKEN" not in env
assert env["COPILOT_API_URL"].startswith("http://127.0.0.1:8787")
assert captured["openai_api_url"] == DEFAULT_API_URL
assert env["OPENAI_TARGET_API_URL"] == DEFAULT_API_URL
assert env["GITHUB_COPILOT_API_URL"] == DEFAULT_API_URL

View file

@ -2,6 +2,8 @@
from __future__ import annotations
import os
import sys
from pathlib import Path
from typing import Any
from unittest.mock import patch
@ -18,6 +20,68 @@ def runner() -> CliRunner:
return CliRunner()
def test_managed_route_reproduction(
runner: CliRunner,
capfd: pytest.CaptureFixture[str],
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
) -> None:
"""The production launcher passes the managed endpoint to an exact-contract child."""
direct = "https://api.kimi.com/coding/v1"
monkeypatch.setenv("KIMI_BASE_URL", direct)
monkeypatch.setenv("KIMI_TEST_UNRELATED", "preserved")
monkeypatch.setattr(wrap_mod, "_project_name_from_cwd", lambda: "repo")
captured: dict[str, Any] = {}
def fake_launch_tool(**kwargs: Any) -> None: # noqa: ANN003
captured.update(kwargs)
with patch.object(wrap_mod.shutil, "which", return_value=sys.executable):
with patch.object(wrap_mod, "_launch_tool", side_effect=fake_launch_tool):
result = runner.invoke(main, ["wrap", "kimi", "--port", "8787"])
assert result.exit_code == 0, result.output
env = captured["env"]
display = captured["env_vars_display"]
configure_launch = captured["configure_launch"]
child_result = tmp_path / "kimi-child.txt"
child = (
"import os, sys; from pathlib import Path; Path(r'"
f"{child_result}"
"').write_text('CHILD|' + os.environ['KIMI_CODE_BASE_URL'] + '|' + "
"os.environ['KIMI_BASE_URL'] + '|' + os.environ['KIMI_TEST_UNRELATED'] + '|' + sys.argv[1])"
)
with (
patch.object(wrap_mod, "_make_cleanup", return_value=lambda: None),
patch.object(wrap_mod.signal, "signal"),
patch.object(wrap_mod, "_register_proxy_client"),
patch.object(wrap_mod, "_ensure_proxy", return_value=(None, 9001)),
patch.object(wrap_mod, "_unregister_proxy_client"),
patch.object(wrap_mod, "_push_runtime_env"),
patch.object(wrap_mod, "_configure_quiet_cli_env", return_value=[]),
):
with pytest.raises(SystemExit) as raised:
wrap_mod._launch_tool(
binary=os.fspath(Path(sys.executable)),
args=("-c", child, "child-arg"),
env=env,
port=8787,
no_proxy=False,
tool_label="KIMI",
env_vars_display=display,
configure_launch=configure_launch,
)
assert raised.value.code == 0
output = result.output + capfd.readouterr().out
expected = "http://127.0.0.1:9001/p/repo/v1"
assert f"KIMI_CODE_BASE_URL={expected}" in output
assert f"KIMI_BASE_URL={expected}" in output
assert child_result.read_text() == f"CHILD|{expected}|{expected}|preserved|child-arg"
assert direct not in output
def test_wrap_kimi_launch(
runner: CliRunner,
tmp_path: Path,
@ -46,15 +110,17 @@ def test_wrap_kimi_launch(
assert captured["agent_type"] == "kimi"
assert captured["args"] == ("-m", "kimi-for-coding")
assert captured["openai_api_url"] == "https://api.kimi.com/coding/v1"
assert callable(captured["configure_launch"])
assert env["KIMI_CODE_BASE_URL"] == "http://127.0.0.1:9000/v1"
assert env["KIMI_BASE_URL"] == "http://127.0.0.1:9000/v1"
def test_wrap_kimi_with_project_name(
def test_project_name(
runner: CliRunner,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Project name is encoded in KIMI_BASE_URL when run from a project directory."""
"""Project name is encoded in both Kimi endpoint variables."""
project_dir = tmp_path / "my-project"
project_dir.mkdir()
monkeypatch.chdir(project_dir)
@ -71,6 +137,7 @@ def test_wrap_kimi_with_project_name(
assert result.exit_code == 0, result.output
env = captured["env"]
assert env["KIMI_CODE_BASE_URL"] == "http://127.0.0.1:7000/p/my-project/v1"
assert env["KIMI_BASE_URL"] == "http://127.0.0.1:7000/p/my-project/v1"
@ -117,12 +184,12 @@ def test_wrap_kimi_not_found(
assert "https://github.com/MoonshotAI/kimi-cli" in result.output
def test_wrap_kimi_custom_port(
def test_port_fallback(
runner: CliRunner,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Custom --port is passed to _launch_tool and appears in KIMI_BASE_URL."""
"""Custom --port is passed to _launch_tool and appears in both URLs."""
monkeypatch.chdir(tmp_path)
monkeypatch.delenv("HEADROOM_CONTEXT_TOOL", raising=False)
@ -138,9 +205,48 @@ def test_wrap_kimi_custom_port(
assert result.exit_code == 0, result.output
assert captured["port"] == 9999
assert captured["env"]["KIMI_CODE_BASE_URL"] == "http://127.0.0.1:9999/v1"
assert captured["env"]["KIMI_BASE_URL"] == "http://127.0.0.1:9999/v1"
def test_non_kimi_fallback_display_is_unchanged(
capfd: pytest.CaptureFixture[str], tmp_path: Path
) -> None:
env = {**os.environ, "OTHER_BASE_URL": "http://127.0.0.1:8787/v1"}
display = ["OTHER_BASE_URL=http://127.0.0.1:8787/v1"]
child_result = tmp_path / "other-child.txt"
child = (
"import os; from pathlib import Path; Path(r'"
f"{child_result}"
"').write_text('CHILD|' + os.environ['OTHER_BASE_URL'])"
)
with (
patch.object(wrap_mod, "_make_cleanup", return_value=lambda: None),
patch.object(wrap_mod.signal, "signal"),
patch.object(wrap_mod, "_register_proxy_client"),
patch.object(wrap_mod, "_ensure_proxy", return_value=(None, 9001)),
patch.object(wrap_mod, "_unregister_proxy_client"),
patch.object(wrap_mod, "_push_runtime_env"),
patch.object(wrap_mod, "_configure_quiet_cli_env", return_value=[]),
):
with pytest.raises(SystemExit) as raised:
wrap_mod._launch_tool(
binary=os.fspath(Path(sys.executable)),
args=("-c", child),
env=env,
port=8787,
no_proxy=False,
tool_label="OTHER",
env_vars_display=display,
)
assert raised.value.code == 0
output = capfd.readouterr().out
assert "OTHER_BASE_URL=http://127.0.0.1:8787/v1" in output
assert child_result.read_text() == "CHILD|http://127.0.0.1:9001/v1"
def test_wrap_kimi_custom_api_url(
runner: CliRunner,
tmp_path: Path,

View file

@ -1,8 +1,11 @@
from __future__ import annotations
import json
import signal
from pathlib import Path
import pytest
from headroom.cli import doctor as doctor_cli
from headroom.cli import wrap as wrap_cli
@ -49,4 +52,19 @@ def test_claude_command_registers_sighup_next_to_sigterm() -> None:
src = inspect.getsource(wrap_cli.claude.callback)
assert 'hasattr(signal, "SIGHUP")' in src
assert "signal.signal(signal.SIGHUP, cleanup)" in src
assert "signal.signal(signal.SIGHUP, _exit_on_signal)" in src
assert "signal.signal(signal.SIGTERM, _exit_on_signal)" in src
def test_signal_handler_unwinds_so_the_restore_can_run() -> None:
"""Registering `cleanup` directly never achieved what #1768 wanted.
A Python signal handler that returns normally does not unwind the stack --
under PEP 475 the interrupted `waitpid` is simply retried -- so the finally
block that restores settings.local.json never ran, while the handler had
already torn the proxy down under a live child. The handler must raise.
"""
with pytest.raises(SystemExit) as excinfo:
wrap_cli._exit_on_signal(signal.SIGHUP, None)
assert excinfo.value.code == 128 + signal.SIGHUP

View file

@ -370,6 +370,35 @@ def test_learn_handles_empty_sessions_and_no_pattern_outputs(
assert "No actionable patterns found." in result.output
def test_learn_surfaces_analysis_failure_and_exits_nonzero(
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
) -> None:
project = SimpleNamespace(name="broken", project_path=tmp_path / "broken")
plugin = FakePlugin("codex", "Codex", [project])
class FailingAnalyzer(FakeAnalyzer):
def analyze(self, project, sessions): # noqa: ANN001, ANN201
self.calls.append((project, sessions))
return SimpleNamespace(
total_sessions=1,
total_calls=3,
total_failures=1,
failure_rate=1 / 3,
recommendations=[],
analysis_error="codex CLI failed (exit 1): Not inside a trusted directory",
)
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "codex-cli")
monkeypatch.setattr("headroom.learn.registry.get_plugin", lambda name: plugin)
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", FailingAnalyzer)
result = runner.invoke(main, ["learn", "--agent", "codex", "--all"])
assert result.exit_code == 1
assert "Analysis failed: codex CLI failed (exit 1)" in result.output
assert "No actionable patterns found." not in result.output
def test_learn_main_only_flag_threads_to_scanner(
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
) -> None:

View file

@ -0,0 +1,568 @@
"""Session-aware /v1/compress (sidecar mode) + the /v1/usage relay.
Contract under test: a gateway that owns routing (e.g. Kong) sends the RAW
conversation plus a session id every turn; Headroom keeps the byte-replay
state itself and returns a byte-identical prefix; the gateway forwards the
result verbatim and may relay provider usage via POST /v1/usage to make
freeze decisions exact.
The critical property is byte-stability: content already returned for a
session must come back byte-for-byte identical on later turns, or the
provider prompt cache busts.
"""
from __future__ import annotations
import json
import pytest
pytest.importorskip("fastapi")
from fastapi.testclient import TestClient # noqa: E402
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
def _make_client() -> TestClient:
config = ProxyConfig(
optimize=True,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
image_optimize=False,
)
app = create_app(config)
client = TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345))
return client
def _big_tool_history() -> list[dict]:
"""A conversation whose tool result is large enough to be compressed."""
items = [
{
"id": i,
"score": 0.99 if i % 30 == 0 else 0.6,
"msg": f"Result {i:03d}{' error' if i % 30 == 0 else ' ok'}",
"blob": f"payload-{i:04d}-" + "".join(chr(97 + (i * 7 + j) % 26) for j in range(240)),
}
for i in range(200)
]
return [
{"role": "user", "content": "Get items"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{"id": "c1", "type": "function", "function": {"name": "get", "arguments": "{}"}}
],
},
{"role": "tool", "tool_call_id": "c1", "content": json.dumps(items)},
]
def _compress(client: TestClient, messages: list[dict], **config) -> dict:
resp = client.post(
"/v1/compress",
json={"model": "gpt-4o", "messages": messages, "config": config},
)
assert resp.status_code == 200, resp.text
return resp.json()
# --------------------------------------------------------------------------- #
# Stateless behaviour is unchanged (regression guard). #
# --------------------------------------------------------------------------- #
# The NUL separator makes the namespace unspoofable from any HTTP header.
SESSION_KEY_PREFIX = "compress\x00"
def test_no_session_id_stays_stateless() -> None:
with _make_client() as client:
body = _compress(client, _big_tool_history())
assert "session" not in body
# And nothing session-shaped leaked into the registry.
proxy = client.app.state.proxy
assert not any(k.startswith(SESSION_KEY_PREFIX) for k in proxy._compression_caches)
def test_invalid_session_id_is_rejected() -> None:
with _make_client() as client:
for bad in ["", " ", "x" * 300, 42]:
resp = client.post(
"/v1/compress",
json={
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
"config": {"session_id": bad},
},
)
assert resp.status_code == 400, f"session_id {bad!r} was not rejected"
def test_compress_user_messages_rejected_with_session() -> None:
"""User-message rewrites are not content-addressed, so they cannot be
byte-replayed after tracker state expires the combination is a latent
prefix-cache bust and must be refused up front."""
with _make_client() as client:
resp = client.post(
"/v1/compress",
json={
"model": "gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
"config": {"session_id": "conv-x", "compress_user_messages": True},
},
)
assert resp.status_code == 400
assert "compress_user_messages" in resp.json()["error"]["message"]
def test_session_key_is_not_spoofable_via_string_prefix() -> None:
"""A caller passing 'compress:...' (or similar) as its session id must
land on a key that no proxy-path header value can also produce."""
with _make_client() as client:
_compress(client, _big_tool_history(), session_id="compress:sneaky")
proxy = client.app.state.proxy
keys = [k for k in proxy._compression_caches if "sneaky" in k]
assert keys == [f"{SESSION_KEY_PREFIX}compress:sneaky"]
# NUL cannot appear in an HTTP header value, so no x-headroom-session-id
# on the proxy path can collide with this key.
assert all("\x00" in k for k in keys)
# --------------------------------------------------------------------------- #
# The core sidecar property: turn 2 replays turn 1's exact bytes. #
# --------------------------------------------------------------------------- #
def test_second_turn_replays_first_turn_bytes() -> None:
with _make_client() as client:
history = _big_tool_history()
turn1 = _compress(client, history, session_id="conv-1")
assert turn1["session"]["id"] == "conv-1"
# The tool result must actually have been compressed, otherwise the
# byte-stability assertion below is vacuous.
t1_tool_content = turn1["messages"][2]["content"]
assert t1_tool_content != history[2]["content"]
assert turn1["tokens_saved"] > 0
# Turn 2: the caller resends the RAW history (as real clients do) plus
# the new turns. Headroom must return the OLD prefix byte-identical to
# what it handed back on turn 1 — that is what the provider cached.
turn2_history = history + [
{"role": "assistant", "content": "The top items are listed above."},
{"role": "user", "content": "Now sort them by score."},
]
turn2 = _compress(client, turn2_history, session_id="conv-1")
assert turn2["messages"][2]["content"] == t1_tool_content
# The WHOLE turn-1 prefix, not just the tool result: any drifted byte
# anywhere in the leading messages is a provider-cache bust.
assert turn2["messages"][: len(turn1["messages"])] == turn1["messages"]
assert turn2["messages"][-1]["content"] == "Now sort them by score."
assert turn2["session"]["id"] == "conv-1"
# Savings must be reported against the RAW payload the caller sent —
# the warm turn still saved the caller ~everything turn 1 saved, even
# though the pipeline itself only saw the already-swapped input.
assert turn2["tokens_saved"] > 0
assert turn2["tokens_before"] > turn2["tokens_after"]
def test_third_turn_still_byte_stable() -> None:
"""The WHOLE returned prefix — every message, byte for byte — must be
stable across N turns. Checking only the tool result would let drift in
any other message (a mutated plain message, a moved marker) bust the
provider cache while the test stayed green.
"""
with _make_client() as client:
history = _big_tool_history()
turn1 = _compress(client, history, session_id="conv-multi")
history2 = history + [{"role": "user", "content": "next"}]
turn2 = _compress(client, history2, session_id="conv-multi")
# Turn 2's leading messages must be exactly turn 1's returned bytes.
assert turn2["messages"][: len(turn1["messages"])] == turn1["messages"]
history3 = history2 + [
{"role": "assistant", "content": "ok"},
{"role": "user", "content": "and again"},
]
turn3 = _compress(client, history3, session_id="conv-multi")
# And turn 3's leading messages must be exactly turn 2's.
assert turn3["messages"][: len(turn2["messages"])] == turn2["messages"]
def test_prefix_stable_even_after_tracker_state_loss() -> None:
"""The overlay's tracker snapshots live shorter (600s session TTL) than
the compression cache (3900s). In that window the frozen+swap path is the
ONLY protection this test kills the tracker between turns and demands
whole-prefix byte stability from frozen+swap alone.
"""
with _make_client() as client:
history = _big_tool_history()
turn1 = _compress(client, history, session_id="conv-trackerloss")
proxy = client.app.state.proxy
# Simulate the tracker registry's TTL sweep reclaiming the session
# while the compression cache (longer TTL) survives.
store = proxy.session_tracker_store
removed = [k for k in list(store._trackers) if "conv-trackerloss" in k]
for k in removed:
del store._trackers[k]
assert removed, "tracker was never created for the session"
assert any("conv-trackerloss" in k for k in proxy._compression_caches)
turn2 = _compress(
client,
history + [{"role": "user", "content": "after tracker loss"}],
session_id="conv-trackerloss",
)
assert turn2["messages"][: len(turn1["messages"])] == turn1["messages"]
def test_header_session_id_ignored_by_default() -> None:
"""Deployments whose gateways stamp x-headroom-session-id on ALL traffic
must not silently flip stateless /v1/compress callers into session mode
(or blend conversations sharing one header value into one replay state)."""
with _make_client() as client:
resp = client.post(
"/v1/compress",
json={"model": "gpt-4o", "messages": _big_tool_history(), "config": {}},
headers={"x-headroom-session-id": "conv-header"},
)
assert resp.status_code == 200, resp.text
assert "session" not in resp.json()
def test_header_session_id_works_with_env_opt_in(monkeypatch) -> None:
monkeypatch.setenv("HEADROOM_COMPRESS_SESSION_FROM_HEADER", "1")
with _make_client() as client:
resp = client.post(
"/v1/compress",
json={"model": "gpt-4o", "messages": _big_tool_history(), "config": {}},
headers={"x-headroom-session-id": "conv-header"},
)
assert resp.status_code == 200, resp.text
assert resp.json()["session"]["id"] == "conv-header"
def test_sessions_are_isolated() -> None:
with _make_client() as client:
history = _big_tool_history()
a1 = _compress(client, history, session_id="conv-a")
b1 = _compress(client, history, session_id="conv-b")
# Same content in, same compressed form out — but through separate
# session state. Interleave new turns and re-check both replay.
a2 = _compress(
client,
history + [{"role": "user", "content": "a follow-up"}],
session_id="conv-a",
)
b2 = _compress(
client,
history + [{"role": "user", "content": "b follow-up"}],
session_id="conv-b",
)
assert a2["messages"][2]["content"] == a1["messages"][2]["content"]
assert b2["messages"][2]["content"] == b1["messages"][2]["content"]
assert a2["messages"][-1]["content"] == "a follow-up"
assert b2["messages"][-1]["content"] == "b follow-up"
# --------------------------------------------------------------------------- #
# /v1/usage: telemetry relay for sidecar sessions. Deliberately NOT a freeze #
# input — freeze stays the locally-replayable bound (see handler docstring). #
# --------------------------------------------------------------------------- #
def test_usage_relay_is_recorded_and_freeze_stays_local() -> None:
with _make_client() as client:
history = _big_tool_history()
_compress(client, history, session_id="conv-usage")
resp = client.post(
"/v1/usage",
json={
"session_id": "conv-usage",
"usage": {
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 50_000,
},
},
)
assert resp.status_code == 200, resp.text
# The tracker recorded the provider-confirmed prefix (telemetry).
assert resp.json()["frozen_message_count"] >= 1
# The next compress freezes from the LOCAL replayable bound, which
# covers the whole previously-returned prefix here.
turn2 = _compress(
client,
history + [{"role": "user", "content": "next"}],
session_id="conv-usage",
)
assert turn2["session"]["frozen_message_count"] >= 1
# An absurdly large confirmed count must never drag freezing past
# what local state can actually replay (that would forward raw bytes
# for evicted entries — the bust this design refuses).
resp2 = client.post(
"/v1/usage",
json={
"session_id": "conv-usage",
"usage": {"cache_read_input_tokens": 10_000_000},
},
)
assert resp2.status_code == 200
turn3 = _compress(
client,
history
+ [
{"role": "user", "content": "next"},
{"role": "assistant", "content": "done"},
{"role": "user", "content": "more"},
],
session_id="conv-usage",
)
# Freeze is capped by message count minus the trailing message — it
# can never exceed what exists, regardless of relayed numbers.
assert turn3["session"]["frozen_message_count"] < 6
def test_usage_unknown_session_is_404_and_leaves_no_footprint() -> None:
with _make_client() as client:
proxy = client.app.state.proxy
before = len(proxy.session_tracker_store._trackers)
for i in range(20):
resp = client.post(
"/v1/usage",
json={
"session_id": f"never-seen-{i}",
"usage": {"cache_read_input_tokens": 100},
},
)
assert resp.status_code == 404
assert resp.json()["error"]["type"] == "unknown_session"
# A flood of novel ids must not grow the tracker store (peek, never
# get_or_create): each ghost tracker would otherwise live a full TTL.
assert len(proxy.session_tracker_store._trackers) == before
def test_usage_without_cache_fields_is_rejected_not_treated_as_cold() -> None:
"""A usage block with NEITHER cache field (e.g. an OpenAI-style
{'prompt_tokens': N} relayed verbatim) carries no cache signal. Treating
the absent fields as 0 would tell the tracker 'provider confirmed fully
cold' and wipe its cached-prefix state on every signal-free relay."""
with _make_client() as client:
_compress(client, _big_tool_history(), session_id="conv-nosignal")
resp = client.post(
"/v1/usage",
json={"session_id": "conv-nosignal", "usage": {"prompt_tokens": 12345}},
)
assert resp.status_code == 400
assert "cache" in resp.json()["error"]["message"]
def test_usage_validation() -> None:
with _make_client() as client:
cases = [
{}, # no session_id
{"session_id": "s"}, # no usage
{"session_id": "s", "usage": "nope"}, # usage not a dict
{"session_id": "s", "usage": {"cache_read_input_tokens": -1}},
{"session_id": "s", "usage": {"cache_read_input_tokens": True}},
]
for body in cases:
resp = client.post("/v1/usage", json=body)
assert resp.status_code == 400, f"body {body!r} was not rejected"
# --------------------------------------------------------------------------- #
# Lifecycle: sidecar sessions ride the registry's TTL/LRU machinery. #
# --------------------------------------------------------------------------- #
def test_session_state_lives_in_registry_and_survives_eviction() -> None:
import time as _time
with _make_client() as client:
history = _big_tool_history()
turn1 = _compress(client, history, session_id="conv-ttl")
proxy = client.app.state.proxy
_key = f"{SESSION_KEY_PREFIX}conv-ttl"
assert _key in proxy._compression_caches
# Simulate the idle-TTL sweep reclaiming the session.
now = _time.time()
proxy._compression_cache_last_seen[_key] = now - 999_999
proxy._compression_caches_last_cleanup = now - 61
proxy._get_compression_cache("unrelated")
assert _key not in proxy._compression_caches
# A post-eviction turn is fail-open: fresh state, valid response, and
# the compressed form is reproducible (deterministic pipeline), even
# though the replay guarantee had to restart from scratch.
turn2 = _compress(
client,
history + [{"role": "user", "content": "after the gap"}],
session_id="conv-ttl",
)
assert turn2["session"]["id"] == "conv-ttl"
assert turn2["messages"][-1]["content"] == "after the gap"
assert isinstance(turn1["messages"][2]["content"], str)
def test_explicit_frozen_count_still_wins_when_larger() -> None:
with _make_client() as client:
history = _big_tool_history()
# First turn with an explicit pin covering the whole tool result: the
# caller asserts the provider already cached it, so it must come back
# byte-for-byte untouched even though no session state exists yet.
turn1 = _compress(client, history, session_id="conv-pin", frozen_message_count=3)
assert turn1["messages"][2]["content"] == history[2]["content"]
assert turn1["session"]["frozen_message_count"] == 3
# --------------------------------------------------------------------------- #
# Review fixes: turn-lock contention, no-signal usage, expired trackers. #
# --------------------------------------------------------------------------- #
def test_compress_503_when_turn_lock_busy(monkeypatch) -> None:
"""A concurrent turn for the same session must fail fast with a 503,
not park an executor worker on an untimed lock acquire."""
import headroom.proxy.handlers.openai as openai_mod
monkeypatch.setattr(openai_mod, "_SESSION_TURN_LOCK_TIMEOUT_SECONDS", 0.05)
with _make_client() as client:
history = _big_tool_history()
_compress(client, history, session_id="conv-lock")
proxy = client.app.state.proxy
lock = proxy._compression_caches[f"{SESSION_KEY_PREFIX}conv-lock"].session_turn_lock
assert lock.acquire(timeout=1), "test could not take the turn lock"
try:
resp = client.post(
"/v1/compress",
json={
"model": "gpt-4o",
"messages": history + [{"role": "user", "content": "blocked"}],
"config": {"session_id": "conv-lock"},
},
)
assert resp.status_code == 503, resp.text
finally:
lock.release()
# With the lock free again the same turn succeeds.
after = _compress(
client,
history + [{"role": "user", "content": "blocked"}],
session_id="conv-lock",
)
assert after["session"]["id"] == "conv-lock"
def test_usage_503_when_turn_lock_busy(monkeypatch) -> None:
"""/v1/usage must take the same turn lock as the compress turn — an
unlocked update races the executor and rolls tracker snapshots back."""
import headroom.proxy.handlers.openai as openai_mod
monkeypatch.setattr(openai_mod, "_SESSION_TURN_LOCK_TIMEOUT_SECONDS", 0.05)
with _make_client() as client:
_compress(client, _big_tool_history(), session_id="conv-ulock")
proxy = client.app.state.proxy
lock = proxy._compression_caches[f"{SESSION_KEY_PREFIX}conv-ulock"].session_turn_lock
assert lock.acquire(timeout=1)
try:
resp = client.post(
"/v1/usage",
json={
"session_id": "conv-ulock",
"usage": {
"cache_read_input_tokens": 100,
"cache_creation_input_tokens": 0,
},
},
)
assert resp.status_code == 503, resp.text
assert resp.json()["error"]["type"] == "session_busy"
finally:
lock.release()
def test_usage_single_zero_field_does_not_wipe_state() -> None:
"""{"cache_read_input_tokens": 0} with no write field (the natural
OpenAI-mapped relay on a cold turn) carries no cache signal it must
not reset the tracker's provider-confirmed prefix state."""
with _make_client() as client:
_compress(client, _big_tool_history(), session_id="conv-zero")
# Establish real provider-confirmed state (both fields present).
resp = client.post(
"/v1/usage",
json={
"session_id": "conv-zero",
"usage": {
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 50_000,
},
},
)
assert resp.status_code == 200
assert resp.json()["applied"] is True
established = resp.json()["frozen_message_count"]
assert established >= 1
# The no-signal relay is acknowledged but NOT applied.
resp2 = client.post(
"/v1/usage",
json={"session_id": "conv-zero", "usage": {"cache_read_input_tokens": 0}},
)
assert resp2.status_code == 200
body = resp2.json()
assert body["applied"] is False
assert body["reason"] == "no_cache_signal"
assert body["frozen_message_count"] == established # state intact
# A relay with BOTH fields zero is a genuine fully-cold assertion
# and IS applied.
resp3 = client.post(
"/v1/usage",
json={
"session_id": "conv-zero",
"usage": {
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 0,
},
},
)
assert resp3.status_code == 200
assert resp3.json()["applied"] is True
def test_usage_404_for_ttl_expired_tracker() -> None:
"""peek() must treat a TTL-expired-but-unswept tracker as gone — a 200
here would resurrect the dead tracker on every relay."""
import time as _time
with _make_client() as client:
_compress(client, _big_tool_history(), session_id="conv-expired")
proxy = client.app.state.proxy
tracker = proxy.session_tracker_store._trackers[f"{SESSION_KEY_PREFIX}conv-expired"]
tracker._last_activity = _time.time() - 999_999
resp = client.post(
"/v1/usage",
json={
"session_id": "conv-expired",
"usage": {"cache_read_input_tokens": 100},
},
)
assert resp.status_code == 404
assert resp.json()["error"]["type"] == "unknown_session"

View file

@ -0,0 +1,185 @@
"""Session-level lifecycle of the compression-cache registry.
Covers the two eviction paths on ``HeadroomProxy._get_compression_cache``:
* capacity eviction must be LRU by *access* (a busy long-lived session
survives; the idlest session goes), not FIFO by creation, and
* the lazy idle-TTL sweep must reclaim sessions whose provider prompt
cache has lapsed, while an access refreshes the clock.
Entry-level LRU/limits inside a single ``CompressionCache`` live in
``test_compression_cache.py``.
"""
from __future__ import annotations
import time
import pytest
pytest.importorskip("fastapi")
def _make_proxy():
from headroom.proxy.server import ProxyConfig, create_app
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
image_optimize=False,
)
app = create_app(config)
return app.state.proxy
def test_capacity_eviction_is_lru_not_fifo(monkeypatch) -> None:
"""At capacity, the idlest session is evicted — not the oldest-created."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "MAX_COMPRESSION_CACHE_SESSIONS", 4)
proxy = _make_proxy()
for sid in ("a", "b", "c", "d"):
proxy._get_compression_cache(sid)
# "a" is the oldest-created; touch it so "b" becomes the LRU.
cache_a = proxy._get_compression_cache("a")
proxy._get_compression_cache("e")
assert "b" not in proxy._compression_caches
assert proxy._get_compression_cache("a") is cache_a
assert "b" not in proxy._compression_cache_last_seen
def test_capacity_eviction_count_respects_small_caps(monkeypatch) -> None:
"""A cap below 4 still evicts at least one session instead of looping."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "MAX_COMPRESSION_CACHE_SESSIONS", 2)
proxy = _make_proxy()
proxy._get_compression_cache("a")
proxy._get_compression_cache("b")
proxy._get_compression_cache("c")
assert len(proxy._compression_caches) == 2
assert "a" not in proxy._compression_caches
def test_idle_ttl_sweep_evicts_expired_sessions(monkeypatch) -> None:
"""A session idle past the TTL is reclaimed by the lazy sweep."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "COMPRESSION_CACHE_TTL_SECONDS", 100.0)
proxy = _make_proxy()
proxy._get_compression_cache("stale")
proxy._get_compression_cache("fresh")
now = time.time()
# Backdate "stale" past the TTL and allow the sweep to run again.
proxy._compression_cache_last_seen["stale"] = now - 101.0
proxy._compression_caches_last_cleanup = (
now - proxy._COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS - 1.0
)
proxy._get_compression_cache("trigger")
assert "stale" not in proxy._compression_caches
assert "stale" not in proxy._compression_cache_last_seen
assert "fresh" in proxy._compression_caches
def test_access_refreshes_ttl_clock(monkeypatch) -> None:
"""Accessing a session resets its idle clock, so it survives the sweep."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "COMPRESSION_CACHE_TTL_SECONDS", 100.0)
proxy = _make_proxy()
proxy._get_compression_cache("busy")
now = time.time()
proxy._compression_cache_last_seen["busy"] = now - 101.0
# Access refreshes last_seen before any sweep can see it as expired.
cache = proxy._get_compression_cache("busy")
proxy._compression_caches_last_cleanup = (
now - proxy._COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS - 1.0
)
proxy._get_compression_cache("trigger")
assert proxy._get_compression_cache("busy") is cache
def test_sweep_is_rate_limited(monkeypatch) -> None:
"""Within the cleanup interval, even an expired session is not swept."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "COMPRESSION_CACHE_TTL_SECONDS", 100.0)
proxy = _make_proxy()
proxy._get_compression_cache("stale")
proxy._compression_cache_last_seen["stale"] = time.time() - 101.0
# _compression_caches_last_cleanup is recent (set in __init__), so the
# sweep must not run yet.
proxy._get_compression_cache("trigger")
assert "stale" in proxy._compression_caches
def test_ttl_sweep_never_evicts_a_session_mid_turn(monkeypatch) -> None:
"""Popping a session whose turn lock is held splits the lock across two
cache instances: the straggler and its retry then run unserialized and
the retry's empty cache recompresses previously-returned bytes."""
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "COMPRESSION_CACHE_TTL_SECONDS", 100.0)
proxy = _make_proxy()
cache = proxy._get_compression_cache("mid-turn")
now = time.time()
proxy._compression_cache_last_seen["mid-turn"] = now - 999.0
assert cache.session_turn_lock.acquire(timeout=1)
try:
proxy._compression_caches_last_cleanup = (
now - proxy._COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS - 1.0
)
proxy._get_compression_cache("trigger-1")
# In-flight: must survive the sweep despite being far past TTL.
assert proxy._compression_caches.get("mid-turn") is cache
finally:
cache.session_turn_lock.release()
# Turn finished: the next sweep may reclaim it.
proxy._compression_caches_last_cleanup = (
time.time() - proxy._COMPRESSION_CACHE_CLEANUP_INTERVAL_SECONDS - 1.0
)
proxy._get_compression_cache("trigger-2")
assert "mid-turn" not in proxy._compression_caches
def test_capacity_eviction_skips_locked_sessions(monkeypatch) -> None:
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "MAX_COMPRESSION_CACHE_SESSIONS", 2)
proxy = _make_proxy()
cache_a = proxy._get_compression_cache("a")
proxy._get_compression_cache("b")
assert cache_a.session_turn_lock.acquire(timeout=1)
try:
# "a" is the LRU but mid-turn — capacity pressure must evict "b".
proxy._get_compression_cache("c")
assert proxy._compression_caches.get("a") is cache_a
assert "b" not in proxy._compression_caches
finally:
cache_a.session_turn_lock.release()

View file

@ -10,6 +10,7 @@ from urllib import error as urllib_error
import pytest
from headroom import copilot_auth
from headroom.proxy import ssl_context
def test_device_authorization_uses_form_encoded_request(monkeypatch: pytest.MonkeyPatch) -> None:
@ -1615,3 +1616,36 @@ def test_exchange_token_sync_returns_payload_on_success(monkeypatch: pytest.Monk
)
assert result == payload
def test_exchange_token_sync_uses_configured_corporate_tls_context(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Copilot refresh must use the same corporate trust config as upstream I/O."""
payload = {"token": "copilot-api", "expires_at": int(time.time()) + 3600}
tls_context = object()
captured: dict[str, object] = {}
class FakeResponse:
def read(self) -> bytes:
return json.dumps(payload).encode()
def __enter__(self):
return self
def __exit__(self, *args):
pass
def fake_urlopen(*args, **kwargs): # noqa: ANN002, ANN003, ANN202
captured.update(kwargs)
return FakeResponse()
monkeypatch.setattr(ssl_context, "build_urlopen_context", lambda: tls_context)
monkeypatch.setattr(copilot_auth.urllib_request, "urlopen", fake_urlopen)
result = copilot_auth.CopilotTokenProvider._exchange_token_sync(
{"Authorization": "Bearer gho_test"} # noqa: S105
)
assert result == payload
assert captured["context"] is tls_context

View file

@ -0,0 +1,148 @@
from __future__ import annotations
import pytest
from click.testing import CliRunner
from headroom.providers.copilot.wrap import (
COPILOT_BYOK_ENV_VARS,
COPILOT_NATIVE_API_URL_ENV,
build_launch_env,
build_native_launch_env,
native_api_url_supported,
)
def test_native_env_redirects_api_and_clears_all_byok_state() -> None:
seeded = dict.fromkeys(COPILOT_BYOK_ENV_VARS, "stale")
seeded["UNRELATED"] = "preserved"
env, _ = build_native_launch_env(port=8890, environ=seeded, project="repo name")
assert env[COPILOT_NATIVE_API_URL_ENV] == "http://127.0.0.1:8890/p/repo%20name"
assert env["UNRELATED"] == "preserved"
assert not any(variable in env for variable in COPILOT_BYOK_ENV_VARS)
def test_byok_builder_remains_disjoint_from_native_mode() -> None:
env, _ = build_launch_env(
port=8787,
provider_type="openai",
wire_api="responses",
environ={},
)
assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1"
assert env["COPILOT_PROVIDER_WIRE_API"] == "responses"
assert COPILOT_NATIVE_API_URL_ENV not in env
def test_native_support_probe_distinguishes_unknown_and_unsupported(tmp_path) -> None:
local = tmp_path / "local"
assert native_api_url_supported(environ={"LOCALAPPDATA": str(local)}) is None
bundle = local / "copilot" / "pkg" / "platform" / "1.0" / "app.js"
bundle.parent.mkdir(parents=True)
bundle.write_text("no override here", encoding="utf-8")
assert native_api_url_supported(environ={"LOCALAPPDATA": str(local)}) is False
bundle.write_text("process.env.COPILOT_API_URL", encoding="utf-8")
assert native_api_url_supported(environ={"LOCALAPPDATA": str(local)}) is True
def test_native_support_probe_skips_unreadable_bundle(monkeypatch, tmp_path) -> None:
local = tmp_path / "local"
bundle = local / "copilot" / "pkg" / "platform" / "1.0" / "app.js"
bundle.parent.mkdir(parents=True)
bundle.write_text("process.env.COPILOT_API_URL", encoding="utf-8")
def _unreadable(*_args, **_kwargs):
raise OSError("synthetic unreadable bundle")
monkeypatch.setattr("builtins.open", _unreadable)
assert native_api_url_supported(environ={"LOCALAPPDATA": str(local)}) is False
def _invoke_native(
monkeypatch: pytest.MonkeyPatch,
extra: list[str] | None = None,
*,
support: bool | None = True,
):
from headroom.cli import wrap as wrap_mod
from headroom.cli.main import main
captured: dict[str, object] = {}
class Resolution:
token = "copilot-token"
api_url = "https://api.business.githubcopilot.com"
refresh_oauth_token = "refresh-token"
api_token_expires_at = 123.0
monkeypatch.setattr(wrap_mod.shutil, "which", lambda _name: "/usr/bin/copilot")
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: False)
monkeypatch.setattr(wrap_mod, "_require_copilot_subscription_resolution", lambda: Resolution())
monkeypatch.setattr(wrap_mod, "_native_api_url_supported", lambda **_kwargs: support)
monkeypatch.setattr(wrap_mod, "_launch_tool", lambda **kwargs: captured.update(kwargs))
result = CliRunner().invoke(
main,
["wrap", "copilot", "--native", "--port", "8890", *(extra or [])],
)
return result, captured
def test_implicit_oauth_uses_native_routing_without_flag(monkeypatch) -> None:
from headroom.cli import wrap as wrap_mod
from headroom.cli.main import main
captured: dict[str, object] = {}
monkeypatch.setattr(wrap_mod.shutil, "which", lambda _name: "/usr/bin/copilot")
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda _port: False)
monkeypatch.setattr(wrap_mod, "has_oauth_auth", lambda: True)
monkeypatch.setattr(wrap_mod, "resolve_client_bearer_token", lambda: "oauth-token")
monkeypatch.setattr(
wrap_mod, "resolve_copilot_api_url", lambda _token: "https://api.githubcopilot.com"
)
monkeypatch.setattr(wrap_mod, "_native_api_url_supported", lambda **_kwargs: True)
monkeypatch.setattr(wrap_mod, "_launch_tool", lambda **kwargs: captured.update(kwargs))
result = CliRunner().invoke(
main,
["wrap", "copilot", "--port", "8890", "--", "--model", "claude-sonnet-5"],
)
assert result.exit_code == 0, result.output
env = captured["env"]
assert isinstance(env, dict)
assert COPILOT_NATIVE_API_URL_ENV in env
assert not any(variable in env for variable in COPILOT_BYOK_ENV_VARS)
def test_native_cli_routes_both_protocols_to_tenant_host(monkeypatch) -> None:
result, captured = _invoke_native(monkeypatch)
assert result.exit_code == 0, result.output
assert captured["openai_api_url"] == "https://api.business.githubcopilot.com"
assert captured["anthropic_api_url"] == "https://api.business.githubcopilot.com"
env = captured["env"]
assert isinstance(env, dict)
assert COPILOT_NATIVE_API_URL_ENV in env
assert not any(variable in env for variable in COPILOT_BYOK_ENV_VARS)
@pytest.mark.parametrize("extra", [["--wire-api", "responses"], ["--provider-type", "anthropic"]])
def test_native_cli_rejects_byok_only_options(monkeypatch, extra) -> None:
result, captured = _invoke_native(monkeypatch, extra)
assert result.exit_code != 0
assert not captured
def test_native_cli_refuses_known_unsupported_bundle(monkeypatch) -> None:
result, captured = _invoke_native(monkeypatch, support=False)
assert result.exit_code != 0
assert "COPILOT_API_URL" in result.output
assert not captured
def test_native_cli_reports_unknown_support_in_verbose_mode(monkeypatch) -> None:
result, captured = _invoke_native(monkeypatch, ["--verbose"], support=None)
assert result.exit_code == 0, result.output
assert "could not verify" in result.output
assert captured

View file

@ -26,6 +26,24 @@ from headroom import HeadroomConfig, HeadroomMode
pytestmark = pytest.mark.skipif(not AGNO_AVAILABLE, reason="Agno not installed")
def _response_usage(input_tokens: int, output_tokens: int, total_tokens: int):
"""Build response usage across Agno 2.x and 3.x module layouts."""
try:
from agno.metrics import MessageMetrics
metrics_type = MessageMetrics
except ImportError:
from agno.models.metrics import Metrics
metrics_type = Metrics
return metrics_type(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=total_tokens,
)
@pytest.fixture
def mock_agno_model():
"""Create a mock Agno model (OpenAIChat-like)."""
@ -50,17 +68,11 @@ def mock_agno_model():
# Mock invoke method (returns ModelResponse for Agno's response() loop)
def mock_invoke(messages, **kwargs):
from agno.models.metrics import Metrics
# Create a proper ModelResponse that Agno's response() can process
return ModelResponse(
role="assistant",
content="Hello! I'm a mock response.",
response_usage=Metrics(
input_tokens=10,
output_tokens=5,
total_tokens=15,
),
response_usage=_response_usage(10, 5, 15),
)
mock.invoke = MagicMock(side_effect=mock_invoke)
@ -73,16 +85,10 @@ def mock_agno_model():
# Mock invoke_stream for streaming
def mock_invoke_stream(messages, **kwargs):
from agno.models.metrics import Metrics
yield ModelResponse(
role="assistant",
content="Streaming...",
response_usage=Metrics(
input_tokens=10,
output_tokens=5,
total_tokens=15,
),
response_usage=_response_usage(10, 5, 15),
)
mock.invoke_stream = MagicMock(side_effect=mock_invoke_stream)

View file

@ -473,6 +473,7 @@ class TestSessionAnalyzer:
assert result.total_calls == 1
assert result.total_failures == 1
assert result.recommendations == []
assert result.analysis_error == "API key not set"
@patch("headroom.learn.analyzer._call_llm")
def test_passes_events_to_digest(self, mock_call_llm: MagicMock):
@ -865,7 +866,7 @@ class TestCallCliLlm:
result = _call_cli_llm("test digest", "codex-cli")
assert result == {"context_file_rules": [], "memory_file_rules": []}
cmd = mock_run.call_args[0][0]
assert cmd == ["codex", "exec"]
assert cmd == ["codex", "exec", "--skip-git-repo-check"]
@patch("headroom.learn.analyzer.subprocess.run")
def test_gemini_cli_uses_p_flag(self, mock_run: MagicMock):

View file

@ -1,53 +1,97 @@
from __future__ import annotations
import json
import pytest
import headroom.mcp_registry.ledger as ledger_module
from headroom.mcp_registry.base import ServerSpec
from headroom.mcp_registry.ledger import (
LedgerMutationError,
clear_install,
headroom_installed_matching,
record_install,
spec_fingerprint,
validate_ledger_for_mutation,
)
def _spec(command: str = "uvx") -> ServerSpec:
return ServerSpec(
name="serena",
command=command,
args=("--from", "git+https://github.com/oraios/serena", "serena"),
)
return ServerSpec("serena", command, ("--from", "serena-agent", "serena"))
def test_ledger_records_matching_install(tmp_path):
def test_ledger_records_and_clears_matching_install(tmp_path):
ledger = tmp_path / "mcp_installs.json"
spec = _spec()
record_install("claude", spec, path=ledger)
assert headroom_installed_matching("claude", spec, path=ledger)
clear_install("claude", "serena", path=ledger)
assert not headroom_installed_matching("claude", spec, path=ledger)
def test_spec_fingerprint_is_stable_for_env_order():
a = ServerSpec("serena", "uvx", env={"B": "2", "A": "1"})
b = ServerSpec("serena", "uvx", env={"A": "1", "B": "2"})
assert spec_fingerprint(a) == spec_fingerprint(b)
@pytest.mark.parametrize(
"value",
[
"not json",
[],
{"agents": None},
{"agents": []},
{"agents": {"claude": None}},
{"agents": {"claude": []}},
{"agents": {"claude": {"serena": None}}},
{"agents": {"claude": {"serena": {"fingerprint": "only"}}}},
],
)
def test_mutation_preflight_rejects_unsafe_shapes(tmp_path, value):
ledger = tmp_path / "mcp_installs.json"
ledger.write_text(value if isinstance(value, str) else json.dumps(value))
with pytest.raises(LedgerMutationError):
validate_ledger_for_mutation(ledger)
def test_mutation_preflight_rejects_unreadable_ledger(monkeypatch, tmp_path):
ledger = tmp_path / "mcp_installs.json"
ledger.write_text('{"agents": {}}')
original_read_text = ledger_module.Path.read_text
def unreadable(path, *args, **kwargs):
if path == ledger:
raise PermissionError("test unreadable ledger")
return original_read_text(path, *args, **kwargs)
monkeypatch.setattr(ledger_module.Path, "read_text", unreadable)
with pytest.raises(LedgerMutationError, match="unreadable"):
validate_ledger_for_mutation(ledger)
def test_read_matching_tolerates_corrupt_ledger(tmp_path):
ledger = tmp_path / "mcp_installs.json"
ledger.write_text("not json")
assert not headroom_installed_matching("claude", _spec(), path=ledger)
def test_record_install_recovers_from_corrupt_ledger(tmp_path):
ledger = tmp_path / "mcp_installs.json"
ledger.write_text("not json")
spec = _spec()
record_install("claude", spec, path=ledger)
assert headroom_installed_matching("claude", spec, path=ledger) is True
assert headroom_installed_matching("claude", spec, path=ledger)
def test_ledger_rejects_changed_spec(tmp_path):
@pytest.mark.parametrize("contents", ['{"agents": null}', '{"agents": {"claude": null}}'])
def test_record_install_recovers_from_unsafe_ledger_shape(tmp_path, contents):
ledger = tmp_path / "mcp_installs.json"
ledger.write_text(contents)
record_install("claude", _spec(), path=ledger)
assert (
headroom_installed_matching("claude", _spec(command="/custom/serena"), path=ledger) is False
)
def test_clear_install_removes_entry(tmp_path):
ledger = tmp_path / "mcp_installs.json"
spec = _spec()
record_install("claude", spec, path=ledger)
clear_install("claude", "serena", path=ledger)
assert headroom_installed_matching("claude", spec, path=ledger) is False
def test_spec_fingerprint_stable_for_env_order():
a = ServerSpec(name="serena", command="uvx", env={"B": "2", "A": "1"})
b = ServerSpec(name="serena", command="uvx", env={"A": "1", "B": "2"})
assert spec_fingerprint(a) == spec_fingerprint(b)
assert headroom_installed_matching("claude", _spec(), path=ledger)

View file

@ -0,0 +1,531 @@
"""Regression tests: file reads over the OpenAI Responses API path must stay verbatim.
Copilot CLI (and other Responses-native harnesses) read files two ways:
1. A first-class ``view`` tool (the Copilot equivalent of Claude Code's ``Read``)
whose output is raw file content the model will byte-patch against.
2. Shell reads through ``bash`` (``cat``/``nl``/``sed -n`` ), which the
chat/Anthropic path protects via ``HEADROOM_PROTECT_READS`` read-command
detection in ``ContentRouter``.
The Responses compression-units path historically protected neither: only
``DEFAULT_EXCLUDE_TOOLS`` names were honored, and ``HEADROOM_PROTECT_READS``
was never consulted. Lossy (Kompress) compression of a fresh file read garbles
exactly the bytes the model needs for line-precise edits, forcing re-reads
(turn inflation) the harm read protection exists to prevent.
"""
from __future__ import annotations
from types import MethodType, SimpleNamespace
from headroom.proxy.handlers.openai import OpenAIHandlerMixin
from headroom.transforms.content_router import (
CompressionStrategy,
ContentRouter,
RouterCompressionResult,
)
class TokenCounter:
def count_text(self, text: str) -> int:
return len(text.split())
def _handler_with_router(router: ContentRouter) -> OpenAIHandlerMixin:
handler = OpenAIHandlerMixin()
handler.openai_pipeline = SimpleNamespace(transforms=[router])
handler.openai_provider = SimpleNamespace(
get_token_counter=lambda _model: TokenCounter(),
)
return handler
def _lossy_router() -> ContentRouter:
"""Router whose compress() always 'lossy-compresses' any candidate it sees."""
router = ContentRouter()
def compress(self, content: str, **_kwargs):
return RouterCompressionResult(
compressed="kept words",
original=content,
strategy_used=CompressionStrategy.KOMPRESS,
)
router.compress = MethodType(compress, router)
return router
def _run(handler: OpenAIHandlerMixin, payload: dict):
return handler._compress_openai_responses_live_text_units_with_router(
payload,
model="gpt-5",
request_id="req_read_protection",
)
_FILE_CONTENT = "\n".join(
f"## Section {i}\nSome roadmap prose line {i} with enough words to matter" for i in range(90)
)
_NL_OUTPUT = "\n".join(
f"{i}\tline {i} of the roadmap file with a handful of words in it" for i in range(1, 110)
)
def test_responses_view_tool_read_stays_verbatim():
"""Copilot's `view` tool returns raw file bytes: never lossy-compress them."""
handler = _handler_with_router(_lossy_router())
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_view",
"name": "view",
"arguments": '{"path": "/repo/ROADMAP.md"}',
},
{
"type": "function_call_output",
"call_id": "call_view",
"output": _FILE_CONTENT,
},
],
}
new_payload, _modified, _saved, _t, _u, _s, _a = _run(handler, payload)
assert new_payload["input"][1]["output"] == _FILE_CONTENT
def test_responses_bash_read_command_stays_verbatim_when_protect_reads(monkeypatch):
"""HEADROOM_PROTECT_READS=1 must cover bash file reads on the Responses path too."""
monkeypatch.setenv("HEADROOM_PROTECT_READS", "1")
handler = _handler_with_router(_lossy_router())
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_bash",
"name": "bash",
"arguments": ('{"command": "nl -ba .overlay/ROADMAP.md | sed -n \'1,75p\'"}'),
},
{
"type": "function_call_output",
"call_id": "call_bash",
"output": _NL_OUTPUT,
},
],
}
new_payload, _modified, _saved, _t, _u, _s, _a = _run(handler, payload)
assert new_payload["input"][1]["output"] == _NL_OUTPUT
def test_responses_excluded_read_tool_stays_verbatim_control():
"""Control: Claude-style `Read` outputs are already protected today."""
handler = _handler_with_router(_lossy_router())
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_read",
"name": "Read",
"arguments": '{"file_path": "/repo/ROADMAP.md"}',
},
{
"type": "function_call_output",
"call_id": "call_read",
"output": _FILE_CONTENT,
},
],
}
new_payload, _modified, _saved, _t, _u, _s, _a = _run(handler, payload)
assert new_payload["input"][1]["output"] == _FILE_CONTENT
def test_responses_bash_read_compresses_when_protect_reads_disabled(monkeypatch):
"""Control: with HEADROOM_PROTECT_READS unset/0, bash reads stay compressible."""
monkeypatch.delenv("HEADROOM_PROTECT_READS", raising=False)
handler = _handler_with_router(_lossy_router())
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_bash",
"name": "bash",
"arguments": '{"command": "cat src/main.py"}',
},
{
"type": "function_call_output",
"call_id": "call_bash",
"output": _NL_OUTPUT,
},
],
}
new_payload, modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert modified is True
assert new_payload["input"][1]["output"] == "kept words"
def test_responses_non_read_bash_command_still_compresses(monkeypatch):
"""Protection is type-specific: test/build/search output stays compressible."""
monkeypatch.setenv("HEADROOM_PROTECT_READS", "1")
handler = _handler_with_router(_lossy_router())
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_test",
"name": "bash",
"arguments": '{"command": "uv run pytest tests/ -q"}',
},
{
"type": "function_call_output",
"call_id": "call_test",
"output": _NL_OUTPUT,
},
],
}
new_payload, modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert modified is True
assert new_payload["input"][1]["output"] == "kept words"
def test_responses_lockfile_read_stays_compressible(monkeypatch):
"""Lockfiles are tool-regenerated, never byte-patched: the command-level
carve-out keeps `cat uv.lock` compressible even with protection on."""
monkeypatch.setenv("HEADROOM_PROTECT_READS", "1")
handler = _handler_with_router(_lossy_router())
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_lock",
"name": "bash",
"arguments": '{"command": "cat uv.lock"}',
},
{
"type": "function_call_output",
"call_id": "call_lock",
"output": _NL_OUTPUT,
},
],
}
new_payload, modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert modified is True
assert new_payload["input"][1]["output"] == "kept words"
def test_responses_local_shell_call_read_stays_verbatim(monkeypatch):
"""Codex native shell: local_shell_call.action.command (argv) read protected."""
monkeypatch.setenv("HEADROOM_PROTECT_READS", "1")
handler = _handler_with_router(_lossy_router())
payload = {
"model": "gpt-5",
"input": [
{
"type": "local_shell_call",
"call_id": "call_lsc",
"action": {"type": "exec", "command": ["nl", "-ba", "ROADMAP.md"]},
},
{
"type": "local_shell_call_output",
"call_id": "call_lsc",
"output": _NL_OUTPUT,
},
],
}
new_payload, _modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert new_payload["input"][1]["output"] == _NL_OUTPUT
def test_responses_view_output_content_part_array_stays_verbatim():
"""`view` output shaped as a content-part array is protected byte-exactly,
including non-text parts."""
handler = _handler_with_router(_lossy_router())
parts = [
{"type": "output_text", "text": _FILE_CONTENT},
{"type": "refusal", "refusal": "n/a"},
]
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_view",
"name": "view",
"arguments": '{"path": "/repo/ROADMAP.md"}',
},
{
"type": "function_call_output",
"call_id": "call_view",
"output": parts,
},
],
}
new_payload, _modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert new_payload["input"][1]["output"] == parts
def test_responses_view_json_shaped_output_stays_byte_exact():
"""Even JSON-shaped `view` output is verbatim: the byte-exact contract beats
the lossless JSON minification other excluded tools accept."""
handler = _handler_with_router(_lossy_router())
pretty_json = "\n".join(
["{"] + [f' "key_{i}": {i},' for i in range(120)] + [' "end": true', "}"]
)
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_view",
"name": "view",
"arguments": '{"path": "/repo/data.json"}',
},
{
"type": "function_call_output",
"call_id": "call_view",
"output": pretty_json,
},
],
}
new_payload, _modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert new_payload["input"][1]["output"] == pretty_json
def test_responses_malformed_arguments_do_not_break_extraction(monkeypatch):
"""Malformed function_call arguments yield no command -> normal compression."""
monkeypatch.setenv("HEADROOM_PROTECT_READS", "1")
handler = _handler_with_router(_lossy_router())
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_bad",
"name": "bash",
"arguments": "{not json at all",
},
{
"type": "function_call_output",
"call_id": "call_bad",
"output": _NL_OUTPUT,
},
],
}
new_payload, modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert modified is True
assert new_payload["input"][1]["output"] == "kept words"
def test_responses_protected_read_survives_cross_turn_dedup(monkeypatch):
"""A repeated protected read must not be replaced by a [↑…] dedup pointer."""
monkeypatch.setenv("HEADROOM_PROTECT_READS", "1")
router = _lossy_router()
router._cross_turn_dedup_enabled = True
handler = _handler_with_router(router)
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_r1",
"name": "bash",
"arguments": '{"command": "nl -ba ROADMAP.md"}',
},
{
"type": "function_call_output",
"call_id": "call_r1",
"output": _NL_OUTPUT,
},
{
"type": "function_call",
"call_id": "call_r2",
"name": "bash",
"arguments": '{"command": "nl -ba ROADMAP.md"}',
},
{
"type": "function_call_output",
"call_id": "call_r2",
"output": _NL_OUTPUT,
},
],
}
new_payload, _modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert new_payload["input"][1]["output"] == _NL_OUTPUT
assert new_payload["input"][3]["output"] == _NL_OUTPUT
def test_responses_debug_path_with_excluded_list_output(monkeypatch):
"""Regression: debug logging over an excluded tool's content-part output must
not raise (latent unbound `fold` variable in the list branch)."""
from headroom.proxy.handlers import openai as openai_handler
monkeypatch.setattr(openai_handler, "_log_codex_compression_debug", lambda *a, **k: None)
handler = _handler_with_router(_lossy_router())
parts = [{"type": "output_text", "text": _FILE_CONTENT}]
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_read",
"name": "Read",
"arguments": '{"file_path": "/repo/ROADMAP.md"}',
},
{
"type": "function_call_output",
"call_id": "call_read",
"output": parts,
},
],
}
new_payload, _modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert new_payload["input"][1]["output"] == parts
def test_responses_read_command_with_releasable_json_output_compresses(monkeypatch):
"""Content gate: a read command whose output is confidently DATA (JSON array)
is released to compression even with HEADROOM_PROTECT_READS=1."""
monkeypatch.setenv("HEADROOM_PROTECT_READS", "1")
handler = _handler_with_router(_lossy_router())
json_output = "[" + ",".join(f'{{"line": {i}, "text": "value {i}"}}' for i in range(60)) + "]"
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_json",
"name": "bash",
"arguments": '{"command": "cat data.json"}',
},
{
"type": "function_call_output",
"call_id": "call_json",
"output": json_output,
},
],
}
new_payload, modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert modified is True
assert new_payload["input"][1]["output"] == "kept words"
def test_responses_local_shell_call_string_command_read_stays_verbatim(monkeypatch):
"""local_shell_call with a string (not argv) command is also covered."""
monkeypatch.setenv("HEADROOM_PROTECT_READS", "1")
handler = _handler_with_router(_lossy_router())
payload = {
"model": "gpt-5",
"input": [
{
"type": "local_shell_call",
"call_id": "call_lsc_str",
"action": {"type": "exec", "command": "cat src/app.py"},
},
{
"type": "local_shell_call_output",
"call_id": "call_lsc_str",
"output": _NL_OUTPUT,
},
],
}
new_payload, _modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert new_payload["input"][1]["output"] == _NL_OUTPUT
def test_responses_debug_path_with_read_protected_output(monkeypatch):
"""Debug logging over a read-protected output records and does not raise."""
from headroom.proxy.handlers import openai as openai_handler
monkeypatch.setattr(openai_handler, "_log_codex_compression_debug", lambda *a, **k: None)
monkeypatch.setenv("HEADROOM_PROTECT_READS", "1")
handler = _handler_with_router(_lossy_router())
payload = {
"model": "gpt-5",
"input": [
{
"type": "function_call",
"call_id": "call_dbg",
"name": "bash",
"arguments": '{"command": "nl -ba ROADMAP.md"}',
},
{
"type": "function_call_output",
"call_id": "call_dbg",
"output": _NL_OUTPUT,
},
],
}
new_payload, _modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert new_payload["input"][1]["output"] == _NL_OUTPUT
def test_responses_read_scan_tolerates_non_dict_and_missing_call_id(monkeypatch):
"""The producer scan must skip non-dict items and calls without a string
call_id without breaking normal compression."""
monkeypatch.setenv("HEADROOM_PROTECT_READS", "1")
handler = _handler_with_router(_lossy_router())
payload = {
"model": "gpt-5",
"input": [
"a bare string item",
{
"type": "function_call",
"name": "bash",
"arguments": '{"command": "cat src/app.py"}',
},
{
"type": "function_call",
"call_id": 42,
"name": "bash",
"arguments": '{"command": "cat src/app.py"}',
},
{
"type": "function_call_output",
"call_id": "call_x",
"output": _NL_OUTPUT,
},
],
}
new_payload, modified, _s, _t, _u, _c, _a = _run(handler, payload)
assert modified is True
assert new_payload["input"][0] == "a bare string item"
assert new_payload["input"][3]["output"] == "kept words"

View file

@ -0,0 +1,165 @@
"""Org-scale sizing knobs: shared-process stores must be tunable and safe.
One Headroom process shared by many users (gateway sidecar/pool) stresses
stores that were sized for a single user's workload:
* the per-session compression-cache entry cap
(``HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES``),
* the process-wide frozen-verdicts store
(``HEADROOM_FROZEN_VERDICTS_MAX``), and
* the session registry under churn (active sessions must survive a flood
of transient ones the LRU property at scale).
Registry TTL/LRU mechanics live in ``test_compression_cache_registry.py``.
"""
from __future__ import annotations
import pytest
pytest.importorskip("fastapi")
def _make_proxy():
from headroom.proxy.server import ProxyConfig, create_app
config = ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
image_optimize=False,
)
app = create_app(config)
return app.state.proxy
# --------------------------------------------------------------------------- #
# Per-session entry cap is plumbed through and env-tunable. #
# --------------------------------------------------------------------------- #
def test_compression_cache_entry_cap_is_plumbed(monkeypatch) -> None:
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "COMPRESSION_CACHE_MAX_ENTRIES", 123)
proxy = _make_proxy()
assert proxy._get_compression_cache("s").max_entries == 123
def test_compression_cache_entry_cap_env_parsing(monkeypatch) -> None:
import importlib
import headroom.proxy.helpers as helpers_mod
monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES", "50000")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_MAX_ENTRIES == 50000
# Floor: an absurdly small value cannot disable the cache.
monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES", "1")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_MAX_ENTRIES == 100
# Garbage falls back to the default.
monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES", "banana")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_MAX_ENTRIES == 10000
monkeypatch.delenv("HEADROOM_COMPRESSION_CACHE_MAX_ENTRIES")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_MAX_ENTRIES == 10000
def test_compression_cache_ttl_env_rejects_non_finite(monkeypatch) -> None:
"""'nan'/'inf' parse as floats but poison every idle comparison — they
must fall back to the default like any other unparseable value."""
import importlib
import headroom.proxy.helpers as helpers_mod
for bad in ("nan", "inf", "-inf"):
monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", bad)
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_TTL_SECONDS == 3900.0, bad
# Below the 600s floor clamps up; above it passes through.
monkeypatch.setenv("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", "60")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_TTL_SECONDS == 600.0
monkeypatch.delenv("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS")
importlib.reload(helpers_mod)
assert helpers_mod.COMPRESSION_CACHE_TTL_SECONDS == 3900.0
# --------------------------------------------------------------------------- #
# Frozen-verdicts store: process-wide, so it must be sizeable per deployment. #
# --------------------------------------------------------------------------- #
def test_frozen_verdicts_cap_env(monkeypatch) -> None:
from headroom.transforms.content_router import ContentRouter, ContentRouterConfig
monkeypatch.setenv("HEADROOM_FROZEN_VERDICTS_MAX", "65536")
assert ContentRouter(ContentRouterConfig())._frozen_verdicts_max == 65536
# Floor: cannot be sized below 256.
monkeypatch.setenv("HEADROOM_FROZEN_VERDICTS_MAX", "1")
assert ContentRouter(ContentRouterConfig())._frozen_verdicts_max == 256
# Garbage falls back to the default.
monkeypatch.setenv("HEADROOM_FROZEN_VERDICTS_MAX", "banana")
assert ContentRouter(ContentRouterConfig())._frozen_verdicts_max == 4096
monkeypatch.delenv("HEADROOM_FROZEN_VERDICTS_MAX")
assert ContentRouter(ContentRouterConfig())._frozen_verdicts_max == 4096
def test_frozen_verdicts_eviction_honors_configured_cap(monkeypatch) -> None:
from headroom.transforms.content_router import ContentRouter, ContentRouterConfig
monkeypatch.setenv("HEADROOM_FROZEN_VERDICTS_MAX", "256")
router = ContentRouter(ContentRouterConfig())
for key in range(300):
router._record_frozen_verdict(key, True)
assert len(router._frozen_verdicts) == 256
# FIFO: the oldest keys were evicted, the newest survive.
assert 0 not in router._frozen_verdicts
assert 299 in router._frozen_verdicts
# --------------------------------------------------------------------------- #
# Session registry under org-scale churn: active sessions always survive a #
# flood of transient ones (the property that keeps busts away at capacity). #
# --------------------------------------------------------------------------- #
def test_active_sessions_survive_transient_flood(monkeypatch) -> None:
import headroom.proxy.server as server_mod
monkeypatch.setattr(server_mod, "MAX_COMPRESSION_CACHE_SESSIONS", 100)
proxy = _make_proxy()
active = [f"active-{i}" for i in range(40)]
active_caches = {sid: proxy._get_compression_cache(sid) for sid in active}
# 400 transient sessions arrive interleaved with active-session traffic —
# 4x the cap, forcing repeated capacity evictions along the way.
for i in range(400):
proxy._get_compression_cache(f"transient-{i}")
if i % 5 == 0: # active sessions keep making requests
for sid in active:
proxy._get_compression_cache(sid)
# Every active session survived with its instance (and therefore its
# byte-replay state) intact; evictions only ever hit transient sessions.
for sid in active:
assert proxy._get_compression_cache(sid) is active_caches[sid], (
f"active session {sid} lost its cache to transient churn"
)
assert len(proxy._compression_caches) <= 100 + len(active)

View file

@ -389,12 +389,7 @@ class TestRecorderBaselineReload:
@staticmethod
def _key() -> str:
return stratum_key(
turn_kind="code",
input_tokens=8000,
model="claude-opus-4-8",
has_tools=True,
)
return SAMPLE_KEY
def test_adopts_baseline_learned_after_start(self, tmp_path):
path = str(tmp_path / "output_savings.json")
@ -482,3 +477,99 @@ class TestRecorderBaselineReload:
relearned.save(path)
assert recorder.estimate().baseline_tokens > baseline_tokens_v1
# ---------------------------------------------------------------------------
# flush durability + event-loop safety
# ---------------------------------------------------------------------------
# Deterministic stratum key shared by the recorder tests below.
SAMPLE_KEY = stratum_key(
turn_kind="code",
input_tokens=8000,
model="claude-opus-4-8",
has_tools=True,
)
class TestFlushDurability:
def test_crash_mid_write_leaves_previous_ledger_intact(self, tmp_path, monkeypatch):
import headroom.fsutil
path = str(tmp_path / "output_savings.json")
key = SAMPLE_KEY
recorder = SavingsRecorder(path, flush_every=1)
recorder.record_from_labels([stratum_label("treatment", key)], 200)
recorder.flush()
assert SavingsLedger.load(path).treatment[key].n == 1
def _die_before_rename(*args, **kwargs):
raise OSError(5, "simulated crash before rename")
monkeypatch.setattr(headroom.fsutil.os, "replace", _die_before_rename)
recorder.record_from_labels([stratum_label("treatment", key)], 210)
recorder.flush() # OSError swallowed by the recorder — fail-open by design
# The pre-crash sample must survive and no temp residue may be left
# behind: a failed save may not corrupt or clutter the ledger.
assert SavingsLedger.load(path).treatment[key].n == 1
assert not list(tmp_path.glob("*.tmp"))
def test_corrupt_ledger_warns_and_starts_empty(self, tmp_path, caplog):
import logging
path = tmp_path / "output_savings.json"
path.write_text("{not json")
with caplog.at_level(logging.WARNING):
SavingsRecorder(str(path))
assert caplog.records, "corrupt ledger was swallowed silently"
def test_emit_request_outcome_flushes_off_the_loop_thread(self, tmp_path, monkeypatch):
import asyncio
import threading
from headroom.proxy.outcome import RequestOutcome, emit_request_outcome
path = str(tmp_path / "output_savings.json")
recorder = SavingsRecorder(path, flush_every=1)
monkeypatch.setattr("headroom.proxy.output_savings.get_recorder", lambda: recorder)
saved_on_threads = []
real_save = SavingsLedger.save
def _spy_save(self, save_path):
saved_on_threads.append(threading.get_ident())
real_save(self, save_path)
monkeypatch.setattr(SavingsLedger, "save", _spy_save)
class _Metrics:
async def record_request(self, **kwargs):
pass
class _Handler:
def __init__(self):
self.metrics = _Metrics()
self.cost_tracker = None
self.logger = None
outcome = RequestOutcome(
request_id="req-shaper",
provider="openai",
model="gpt-5",
status_code=200,
original_tokens=100,
optimized_tokens=80,
output_tokens=50,
tokens_saved=20,
attempted_input_tokens=100,
transforms_applied=(stratum_label("treatment", SAMPLE_KEY),),
)
asyncio.run(emit_request_outcome(_Handler(), outcome))
loop_thread = threading.get_ident()
assert saved_on_threads, "flush never ran"
assert all(t != loop_thread for t in saved_on_threads)

View file

@ -56,6 +56,55 @@ def test_resolve_api_targets_normalizes_trailing_v1() -> None:
assert targets.vertex == "https://vertex.example"
def test_copilot_openai_target_routes_anthropic_to_copilot() -> None:
"""When the OpenAI target is a Copilot host and no Anthropic override is set,
the Anthropic target must default to the same Copilot host.
Copilot serves Claude models via its Anthropic surface (``/v1/messages``) on
the same host. Without this, Claude requests fell back to api.anthropic.com
and 401'd with the Copilot bearer ("Invalid bearer token", #3247).
"""
targets = resolve_api_targets(
ProviderApiOverrides(
anthropic=None,
openai="https://api.githubcopilot.com",
gemini=None,
cloudcode=None,
vertex=None,
)
)
assert targets.openai == "https://api.githubcopilot.com"
assert targets.anthropic == "https://api.githubcopilot.com"
def test_explicit_anthropic_override_wins_over_copilot_default() -> None:
"""An explicit Anthropic target is never overridden by the Copilot default."""
targets = resolve_api_targets(
ProviderApiOverrides(
anthropic="https://api.anthropic.com",
openai="https://api.githubcopilot.com",
gemini=None,
cloudcode=None,
vertex=None,
)
)
assert targets.anthropic == "https://api.anthropic.com"
def test_non_copilot_openai_target_leaves_anthropic_default() -> None:
"""A non-Copilot OpenAI target must not touch the Anthropic default."""
targets = resolve_api_targets(
ProviderApiOverrides(
anthropic=None,
openai="https://api.openai.com",
gemini=None,
cloudcode=None,
vertex=None,
)
)
assert targets.anthropic == "https://api.anthropic.com"
def test_proxy_config_exposes_provider_api_overrides() -> None:
config = ProxyConfig(
anthropic_api_url="https://anthropic.example",

View file

@ -199,7 +199,18 @@ def test_bypass_header_does_not_invoke_cached_prefix_replay(monkeypatch):
assert captured[-1]["messages"] == messages
def test_backpressure_does_not_invoke_cached_prefix_replay(monkeypatch):
def test_backpressure_still_invokes_cached_prefix_replay(monkeypatch):
"""INVERTED from the pre-#3261 contract this test used to pin.
Backpressure sheds the compression PIPELINE (the CPU-heavy stage), but
the byte-identical cached-prefix replay must STILL run: skipping it
forwarded raw originals over a compressed cached prefix, busting every
gated session's prompt cache exactly at peak load (the saturated path
previously emitted `Cached-prefix replay skipped:
reason=pre_upstream_backpressure` that skip was the bug). The replay
self-guards and no-ops here (no previous turn), so the raw messages
still pass through unchanged.
"""
app = create_app(
_config(
optimize=True,
@ -208,12 +219,29 @@ def test_backpressure_does_not_invoke_cached_prefix_replay(monkeypatch):
)
)
def fail_if_called(*args, **kwargs): # noqa: ANN002, ANN003
raise AssertionError("cached-prefix replay must be skipped under backpressure")
from headroom.cache import prefix_tracker as _pt
monkeypatch.setattr("headroom.cache.prefix_tracker.overlay_cached_prefix", fail_if_called)
debug = Mock()
monkeypatch.setattr("headroom.proxy.handlers.anthropic.logger.debug", debug)
real_overlay = _pt.overlay_cached_prefix
overlay_calls: list[int] = []
def spy(*args, **kwargs): # noqa: ANN002, ANN003
overlay_calls.append(1)
return real_overlay(*args, **kwargs)
# Patch every binding of overlay_cached_prefix: the handler historically
# imported it from prefix_tracker per-request, and the shared session
# engine (headroom.proxy.session_engine, later in this stack) binds it
# at module import — cover both so this test holds across the stack.
monkeypatch.setattr("headroom.cache.prefix_tracker.overlay_cached_prefix", spy)
try:
import headroom.proxy.session_engine as _se
monkeypatch.setattr(_se, "overlay_cached_prefix", spy)
except ImportError:
pass
info = Mock()
monkeypatch.setattr("headroom.proxy.handlers.anthropic.logger.info", info)
proxy = app.state.proxy
class _SaturatedSemaphore:
@ -236,10 +264,13 @@ def test_backpressure_does_not_invoke_cached_prefix_replay(monkeypatch):
proxy.anthropic_pre_upstream_sem.release()
assert response.status_code == 200, response.text
assert captured[-1]["messages"] == messages
assert any(
call.args and call.args[-1] == "pre_upstream_backpressure" for call in debug.call_args_list
# Backpressure engaged (pipeline shed)...
assert any("pre_upstream_backpressure" in str(call) for call in info.call_args_list), (
"backpressure did not engage — the test setup no longer saturates"
)
# ...but the replay ran (and, with no previous turn, no-op'd safely).
assert overlay_calls, "cached-prefix replay must run under backpressure"
assert captured[-1]["messages"] == messages
def test_optimize_on_aligned_history_preserves_replay():

View file

@ -303,7 +303,12 @@ class TestCompressEndpointCompression:
transforms_summary={"test_transform": 1},
markers_inserted=[],
)
run_compression = AsyncMock(return_value=result)
# The executor callable returns the 5-tuple contract of
# _run_stateless/_run_session_turn:
# (result, final_messages, tokens_before, tokens_after, session_info).
run_compression = AsyncMock(
return_value=(result, result.messages, result.tokens_before, result.tokens_after, None)
)
record_outcome = AsyncMock()
monkeypatch.setattr(proxy, "_run_compression_in_executor", run_compression)
monkeypatch.setattr(proxy, "_record_request_outcome", record_outcome)
@ -349,7 +354,16 @@ class TestCompressEndpointCompression:
monkeypatch.setattr(
proxy,
"_run_compression_in_executor",
AsyncMock(return_value=result),
# Same 5-tuple contract as _run_stateless (see above).
AsyncMock(
return_value=(
result,
result.messages,
result.tokens_before,
result.tokens_after,
None,
)
),
)
monkeypatch.setattr(proxy, "_record_request_outcome", AsyncMock())

View file

@ -0,0 +1,225 @@
"""Unit tests for the shared session-turn engine (headroom/proxy/session_engine).
The engine is the single cache-management brain for the proxy request paths
and the sidecar /v1/compress path; these tests pin its two freeze policies
and the overlay finalization directly, without an HTTP harness.
"""
from __future__ import annotations
import json
import pytest
from headroom.cache.compression_cache import CompressionCache
from headroom.proxy.session_engine import (
FREEZE_POLICY_CONFIRMED_CLAMP,
FREEZE_POLICY_REPLAYABLE,
finalize_turn,
prepare_turn,
)
def _tool_msg(content: str, call_id: str = "c1") -> dict:
return {"role": "tool", "tool_call_id": call_id, "content": content}
def _history_with_cached_tool(
cache: CompressionCache, original: str, compressed: str
) -> list[dict]:
"""A 3-message history whose tool result has a cached compressed form."""
cache.store_compressed(cache.content_hash(original), compressed, tokens_saved=10)
return [
{"role": "user", "content": "get items"},
{"role": "assistant", "content": "calling"},
_tool_msg(original),
]
# --------------------------------------------------------------------------- #
# prepare_turn: freeze policies #
# --------------------------------------------------------------------------- #
def test_sidecar_policy_freezes_full_replayable_prefix() -> None:
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
# user, assistant, cached tool are all stable; the trailing message is
# always excluded by compute_frozen_count.
assert prep.frozen_message_count == 3
# The swap replaced the tool result with its cached compressed form.
assert prep.pipeline_input[2]["content"] == "[compressed]"
# The caller's list is never mutated.
assert messages[2]["content"].startswith("ORIGINAL")
def test_sidecar_policy_explicit_pin_wins_when_larger() -> None:
cache = CompressionCache()
messages = [
{"role": "user", "content": "a"},
_tool_msg("never seen before " * 50), # not in cache -> derived stops here
{"role": "user", "content": "next"},
]
derived = cache.compute_frozen_count(messages)
assert derived == 1 # only the leading plain message
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE, explicit_frozen=2)
assert prep.frozen_message_count == 2
def test_sidecar_policy_derived_wins_when_explicit_smaller() -> None:
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE, explicit_frozen=1)
assert prep.frozen_message_count == 3
def test_proxy_policy_clamps_by_cache_count() -> None:
"""Provider says 5 messages are cached, but local state can only replay 3:
freezing past the replayable bound would forward raw bytes."""
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
messages.append(_tool_msg("uncached " * 50, "c2"))
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_CONFIRMED_CLAMP, tracker_frozen=5)
assert prep.frozen_message_count == 3
def test_proxy_policy_clamps_by_tracker() -> None:
"""Local state could replay 3, but the provider only confirmed 1: content
past the confirmed prefix stays compressible (the #327 posture)."""
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_CONFIRMED_CLAMP, tracker_frozen=1)
assert prep.frozen_message_count == 1
def test_proxy_policy_none_tracker_freezes_nothing() -> None:
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "ORIGINAL " * 100, "[compressed]")
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_CONFIRMED_CLAMP, tracker_frozen=None)
assert prep.frozen_message_count == 0
def test_unknown_policy_rejected() -> None:
cache = CompressionCache()
with pytest.raises(ValueError):
prepare_turn(cache, [], policy="wat")
def test_prepare_marks_frozen_tool_results_stable() -> None:
cache = CompressionCache()
original = "ORIGINAL " * 100
messages = _history_with_cached_tool(cache, original, "[compressed]")
messages.append({"role": "user", "content": "next"})
prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
assert cache.content_hash(original) in cache._stable_hashes
# --------------------------------------------------------------------------- #
# finalize_turn: overlay + recount hook #
# --------------------------------------------------------------------------- #
def _prev_pair() -> tuple[list[dict], list[dict]]:
prev_original = [
{"role": "user", "content": "ORIGINAL " * 100},
{"role": "assistant", "content": "ok"},
]
prev_returned = [
{"role": "user", "content": "[returned-form]"},
{"role": "assistant", "content": "ok"},
]
return prev_original, prev_returned
def test_finalize_replays_previous_returned_prefix() -> None:
prev_original, prev_returned = _prev_pair()
current = prev_original + [{"role": "user", "content": "next"}]
# The pipeline "drifted": it emitted the raw original for message 0.
drifted = [dict(m) for m in current]
counted: list[int] = []
def _count(msgs: list[dict]) -> int:
counted.append(len(json.dumps(msgs)))
return 42
turn = finalize_turn(drifted, current, prev_original, prev_returned, count_tokens=_count)
assert turn.replayed
assert turn.messages[0]["content"] == "[returned-form]"
assert turn.messages[-1]["content"] == "next"
assert turn.tokens == 42
assert len(counted) == 1
def test_finalize_noop_without_prev_snapshots() -> None:
current = [{"role": "user", "content": "hi"}]
calls: list[int] = []
turn = finalize_turn(current, current, [], [], count_tokens=lambda m: calls.append(1) or 1)
assert not turn.replayed
assert turn.messages == current
assert turn.tokens is None
assert not calls # count_tokens only runs when the overlay fired
def test_finalize_count_hook_failure_falls_back() -> None:
prev_original, prev_returned = _prev_pair()
current = prev_original + [{"role": "user", "content": "next"}]
def _boom(_msgs: list[dict]) -> int:
raise RuntimeError("tokenizer down")
turn = finalize_turn(
[dict(m) for m in current], current, prev_original, prev_returned, count_tokens=_boom
)
assert turn.replayed
assert turn.tokens is None
# --------------------------------------------------------------------------- #
# OpenAI proxy token-path migration: formula identity + marking benefit. #
# --------------------------------------------------------------------------- #
def test_replayable_without_pin_equals_bare_cache_count() -> None:
"""The OpenAI proxy token path historically froze on compute_frozen_count
alone; REPLAYABLE with no explicit pin must be formula-identical, so its
migration onto the engine is a pure extraction."""
cache = CompressionCache()
messages = _history_with_cached_tool(cache, "AAAA " * 50, "[c1]")
messages.append(_tool_msg("uncached content", call_id="c2"))
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
assert prep.frozen_message_count == cache.compute_frozen_count(messages)
# And that count stops at the uncached tool_result (index 3).
assert prep.frozen_message_count == 3
def test_marking_preserves_freeze_across_entry_eviction() -> None:
"""The one real benefit mark_stable_from_messages adds on the migrated
path: an in-prefix tool_result stays stable via `_stable_hashes` even
after its compressed ENTRY is evicted by the per-cache LRU, so the frozen
count does not collapse at that position on the next turn."""
cache = CompressionCache(max_entries=100)
original = "BBBB " * 50
messages = _history_with_cached_tool(cache, original, "[c1]")
messages.append({"role": "user", "content": "next"})
prep = prepare_turn(cache, messages, policy=FREEZE_POLICY_REPLAYABLE)
assert prep.frozen_message_count == 3 # tool in prefix, marked stable
# Simulate entry LRU turnover: the compressed entry disappears.
h = cache.content_hash(original)
with cache._lock:
cache._cache.pop(h, None)
# Without marking, the frozen count would collapse to 2 here; the
# stable-hash record keeps the position frozen.
assert cache.compute_frozen_count(messages) == 3

View file

@ -689,6 +689,49 @@ class TestSQLiteGraphStoreEdgeCases:
assert len(subgraph.entities) == 0
assert len(subgraph.relationships) == 0
@pytest.mark.asyncio
async def test_one_corrupt_row_does_not_abort_a_multi_row_scan(self, tmp_path):
"""A single unparseable row must not take down an entire query.
Regression: ``_row_to_relationship`` / ``_row_to_entity`` parsed stored
JSON/timestamps with no guard, so one corrupt row (partial write, manual
edit, bad migration) raised inside the row loop and aborted the whole
``query_subgraph`` / ``get_relationships`` scan taking unrelated,
perfectly good edges down with it. The corrupt row is now skipped.
"""
import sqlite3
store = SQLiteGraphStore(db_path=str(tmp_path / "graph.db"))
a = Entity(user_id="u", name="A", entity_type="n")
b = Entity(user_id="u", name="B", entity_type="n")
c = Entity(user_id="u", name="C", entity_type="n")
for entity in (a, b, c):
await store.add_entity(entity)
await store.add_relationship(
Relationship(user_id="u", source_id=a.id, target_id=b.id, relation_type="e")
)
await store.add_relationship(
Relationship(user_id="u", source_id=a.id, target_id=c.id, relation_type="e")
)
# Corrupt the A->B relationship row's properties JSON out-of-band.
con = sqlite3.connect(str(store.db_path))
con.execute("UPDATE relationships SET properties = '{oops' WHERE target_id = ?", (b.id,))
con.commit()
con.close()
# get_relationships returns the one good edge instead of raising.
rels = await store.get_relationships(a.id)
assert len(rels) == 1
assert rels[0].target_id == c.id
# query_subgraph completes, skipping the corrupt edge and its node.
subgraph = await store.query_subgraph(
[a.id], max_hops=1, direction=RelationshipDirection.OUTGOING
)
assert {e.name for e in subgraph.entities} == {"A", "C"}
assert len(subgraph.relationships) == 1
@pytest.mark.asyncio
async def test_entity_with_special_characters(self, store):
"""Test entity names with special characters."""

View file

@ -170,6 +170,49 @@ def test_search_detection_uses_match_ratio() -> None:
assert _try_detect_search("\n\n") is None
def test_search_detection_rejects_datetime_prefixed_user_message() -> None:
"""Regression: wrap-copilot ate one-line interactive prompts (2026-08-23).
Copilot CLI prepends ``<current_datetime></current_datetime>`` to every
interactive user turn. The ISO-8601 ``T09:57:59`` matched the grep
``file:line:`` pattern, so a datetime + one-line prompt classified as
SEARCH_RESULTS (1 match / 2 lines = 50% 30%) and the SearchCompressor
deleted the prompt line the model received only the timestamp.
"""
incident = (
"<current_datetime>2026-08-23T09:57:59.792+02:00</current_datetime>\n"
"\n"
"Please update the PR desc and check .overlay/ for hints."
)
assert _try_detect_search(incident) is None
assert detect_content_type(incident).content_type is not ContentType.SEARCH_RESULTS
def test_search_detection_requires_two_matching_lines() -> None:
"""A single coincidental ``word:digits:`` line must not classify prose."""
assert _try_detect_search("src/foo.py:12:def foo():") is None
assert (
_try_detect_search(
"Meeting at 09:30:00 tomorrow.\nBring the reports.\nDo not forget coffee."
)
is None
)
# Two genuine grep lines still classify.
two = "src/foo.py:12:def foo():\nsrc/bar.py:34: foo()"
result = _try_detect_search(two)
assert result is not None
assert result.content_type is ContentType.SEARCH_RESULTS
def test_search_detection_rejects_tag_like_and_key_value_prefixes() -> None:
"""Markup / key=value lines are not file paths even with ``:\\d+:`` inside."""
assert (
_try_detect_search('<log time="10:00:00">started</log>\n<log time="10:00:01">stopped</log>')
is None
)
assert _try_detect_search("timeout=30:12:retried\ntimeout=31:12:retried") is None
def test_log_detection_prefers_build_output_patterns() -> None:
log_output = "\n".join(
[

View file

@ -1785,3 +1785,23 @@ def test_detect_content_overrides_html_misroute_for_grep_and_logs(
"</section></main></body></html>"
)
assert _detect_content(html).content_type is ContentType.HTML
def test_datetime_prefixed_user_prompt_survives_router() -> None:
"""Regression (2026-08-23): interactive wrap-copilot prompts were deleted.
Copilot CLI prepends ``<current_datetime></current_datetime>`` to every
interactive user turn; the ISO timestamp matched the grep ``file:line:``
detector, the one-line prompt classified as SEARCH_RESULTS, and
SearchCompressor kept only the datetime line the model received no
request and answered "How can I help you today?". The router must never
route this shape to the search line-filter and must keep the prose.
"""
prompt = (
"<current_datetime>2026-08-23T09:57:59.792+02:00</current_datetime>\n"
"\n"
"Please update the PR desc and check .overlay/ for hints."
)
result = ContentRouter().compress(prompt)
assert result.strategy_used is not CompressionStrategy.SEARCH
assert "Please update the PR desc" in result.compressed

View file

@ -0,0 +1,254 @@
"""Concurrent `headroom wrap` sessions sharing one project's settings (#3205).
`wrap claude` writes ANTHROPIC_BASE_URL into `.claude/settings.local.json` and
restores it on exit. Several sessions in one project run that read-modify-write
concurrently. The write is atomic so the file never tears, but the updates were
still lost against each other:
* the first session's exit deleted the key while the others were still
running -- they silently stopped routing through the proxy, kept working,
and lost every byte of compression with no error anywhere; and
* a session that started second remembered the *first* session's proxy URL as
"the original", so its exit wrote a dead proxy back into the file, which
every later session in that project then failed to connect to.
"""
from __future__ import annotations
import json
from pathlib import Path
from unittest import mock
import pytest
from headroom.cli import wrap as W
@pytest.fixture
def settings(tmp_path: Path) -> Path:
path = tmp_path / ".claude" / "settings.local.json"
path.parent.mkdir(parents=True)
path.write_text(json.dumps({"env": {"FOO": "bar"}}), encoding="utf-8")
return path
def _env(path: Path) -> dict:
return json.loads(path.read_text(encoding="utf-8")).get("env", {}) if path.exists() else {}
class _Sessions:
"""Drive several wrap sessions with distinct, controllable PIDs."""
def __init__(self, *pids: int) -> None:
self.live = set(pids)
def __enter__(self) -> _Sessions:
self._patches = [
mock.patch.object(W, "_pid_alive", lambda pid: pid in self.live),
mock.patch.object(W, "_identity_mismatch", lambda *a: False),
]
for p in self._patches:
p.start()
return self
def __exit__(self, *exc: object) -> None:
for p in self._patches:
p.stop()
def launch(self, pid: int, url: str, path: Path, port: int | None = None) -> str | None:
with mock.patch("os.getpid", lambda: pid):
return W._write_claude_wrap_base_url(url, settings_path=path, port=port)
def exit(self, pid: int, previous: str | None, path: Path) -> None:
self.live.discard(pid)
with mock.patch("os.getpid", lambda: pid):
W._restore_claude_wrap_base_url(previous, settings_path=path)
def crash(self, pid: int) -> None:
"""Vanish without running cleanup (SIGKILL, hard reboot)."""
self.live.discard(pid)
def test_first_session_exiting_leaves_the_others_routed(settings: Path) -> None:
"""The reported symptom: sessions silently stop routing when a sibling exits."""
with _Sessions(1001, 1002) as s:
a = s.launch(1001, "http://127.0.0.1:8787", settings)
s.launch(1002, "http://127.0.0.1:8788", settings)
s.exit(1001, a, settings)
assert "ANTHROPIC_BASE_URL" in _env(settings), "surviving session was unrouted"
def test_last_session_out_restores_the_true_original(settings: Path) -> None:
"""A later session must not restore an earlier session's dead proxy URL."""
with _Sessions(1001, 1002) as s:
a = s.launch(1001, "http://127.0.0.1:8787", settings)
b = s.launch(1002, "http://127.0.0.1:8788", settings)
s.exit(1001, a, settings)
s.exit(1002, b, settings)
assert _env(settings) == {"FOO": "bar"}, "stale proxy URL left behind"
def test_a_pre_existing_user_base_url_survives_the_whole_cycle(settings: Path) -> None:
"""A URL the project already had is restored, not deleted."""
settings.write_text(
json.dumps({"env": {"ANTHROPIC_BASE_URL": "http://user-proxy:1234"}}), encoding="utf-8"
)
with _Sessions(1001, 1002) as s:
a = s.launch(1001, "http://127.0.0.1:8787", settings)
b = s.launch(1002, "http://127.0.0.1:8788", settings)
s.exit(1001, a, settings)
s.exit(1002, b, settings)
assert _env(settings)["ANTHROPIC_BASE_URL"] == "http://user-proxy:1234"
def test_three_sessions_any_exit_order(settings: Path) -> None:
for order in ([1001, 1002, 1003], [1003, 1001, 1002], [1002, 1003, 1001]):
settings.write_text(json.dumps({"env": {"FOO": "bar"}}), encoding="utf-8")
with _Sessions(*order) as s:
prev = {
pid: s.launch(pid, f"http://127.0.0.1:{8787 + i}", settings)
for i, pid in enumerate(order)
}
for pid in order[:-1]:
s.exit(pid, prev[pid], settings)
assert "ANTHROPIC_BASE_URL" in _env(settings), f"unrouted early in {order}"
s.exit(order[-1], prev[order[-1]], settings)
assert _env(settings) == {"FOO": "bar"}, f"residue after {order}"
def test_a_crashed_session_does_not_wedge_the_key(settings: Path) -> None:
"""A SIGKILLed session never releases; its claim must be pruned as dead."""
with _Sessions(1001, 1002) as s:
s.launch(1001, "http://127.0.0.1:8787", settings)
b = s.launch(1002, "http://127.0.0.1:8788", settings)
s.crash(1001)
s.exit(1002, b, settings)
assert _env(settings) == {"FOO": "bar"}
assert not W._wrap_owners_path(settings).exists()
def test_single_session_behaviour_is_unchanged(settings: Path) -> None:
with _Sessions(1001) as s:
a = s.launch(1001, "http://127.0.0.1:8787", settings)
assert _env(settings)["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:8787"
s.exit(1001, a, settings)
assert _env(settings) == {"FOO": "bar"}
def test_restore_without_an_owner_record_still_honours_the_caller(settings: Path) -> None:
"""unwrap and legacy sessions pass the previous value directly."""
settings.write_text(
json.dumps({"env": {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787"}}), encoding="utf-8"
)
assert not W._wrap_owners_path(settings).exists()
W._restore_claude_wrap_base_url("http://legacy:9999", settings_path=settings)
assert _env(settings)["ANTHROPIC_BASE_URL"] == "http://legacy:9999"
def test_tool_search_key_is_tracked_independently(settings: Path) -> None:
"""Ownership is per key -- the tool-search entry has the same race."""
with _Sessions(1001, 1002) as s:
with mock.patch("os.getpid", lambda: 1001):
a = W._write_claude_wrap_tool_search("auto", settings_path=settings)
with mock.patch("os.getpid", lambda: 1002):
W._write_claude_wrap_tool_search("auto", settings_path=settings)
s.live.discard(1001)
with mock.patch("os.getpid", lambda: 1001):
W._restore_claude_wrap_tool_search(a, settings_path=settings)
assert W._TOOL_SEARCH_ENV in _env(settings), "surviving session lost tool-search"
def test_exit_on_signal_unwinds_so_finally_can_run() -> None:
"""`cleanup` as the handler never unwound; the settings restore never ran."""
with pytest.raises(SystemExit) as excinfo:
W._exit_on_signal(15, None)
assert excinfo.value.code == 143
def test_unwrap_forces_the_restore_past_a_live_session(settings: Path) -> None:
"""`unwrap` is the user asking for their settings back -- it must not no-op.
Deferring to a live sibling is right for a session exiting on its own, but
unwrap deferring means the command prints success while leaving the proxy
URL in the file.
"""
with _Sessions(1001) as s:
s.launch(1001, "http://127.0.0.1:8787", settings)
with mock.patch("os.getpid", lambda: 2002):
W._restore_claude_wrap_base_url(None, settings_path=settings, force=True)
assert _env(settings) == {"FOO": "bar"}, "unwrap left the proxy URL behind"
assert not W._wrap_owners_path(settings).exists(), "unwrap left ownership state behind"
def test_unwrap_restores_the_true_original_not_the_marker_value(settings: Path) -> None:
"""A caller with no claim of its own trusts the record over its marker.
The single-slot marker is won by the *last* writer, whose `previous` is the
first session's proxy URL -- restoring that is the #3205 bug via unwrap.
"""
settings.write_text(
json.dumps({"env": {"ANTHROPIC_BASE_URL": "http://user-proxy:1234"}}), encoding="utf-8"
)
with _Sessions(1001, 1002) as s:
s.launch(1001, "http://127.0.0.1:8787", settings, port=8787)
s.launch(1002, "http://127.0.0.1:8788", settings, port=8788)
with mock.patch("os.getpid", lambda: 2002):
W._restore_claude_wrap_base_url(
"http://127.0.0.1:8787", settings_path=settings, force=True
)
assert _env(settings)["ANTHROPIC_BASE_URL"] == "http://user-proxy:1234"
def test_a_holder_that_outlived_its_proxy_cannot_veto_the_selfheal(settings: Path) -> None:
"""#2221: a wrapper PID can outlive its proxy; its claim must not block."""
with _Sessions(1001) as s:
s.launch(1001, "http://127.0.0.1:8787", settings, port=8787)
# PID 1001 is still alive, but port 8787 has been proven dead.
W._restore_claude_wrap_base_url(None, settings_path=settings, dead_ports=frozenset({8787}))
assert _env(settings) == {"FOO": "bar"}, "dead proxy URL survived the self-heal"
def test_exiting_session_hands_its_marker_to_a_survivor(settings: Path) -> None:
"""The marker has one slot; the leaver must not strand or hijack it."""
with _Sessions(1001, 1002) as s:
s.launch(1001, "http://127.0.0.1:8787", settings, port=8787)
b = s.launch(1002, "http://127.0.0.1:8788", settings, port=8788)
marker = W._read_wrap_marker(settings)
assert marker is not None and marker["pid"] == 1002, "last writer owns the marker"
s.exit(1002, b, settings)
marker = W._read_wrap_marker(settings)
assert marker is not None, "survivor lost its #2221 self-heal record"
assert marker["pid"] == 1001, "marker still describes the exited session"
assert marker["port"] == 8787
assert marker["previous"] is None, "marker must carry the true original"
def test_the_founding_session_still_honours_an_explicit_previous(settings: Path) -> None:
"""A sole writer observed the pre-wrap value first-hand; do not override it."""
with _Sessions(1001) as s:
s.launch(1001, "http://127.0.0.1:8787", settings)
s.exit(1001, "https://existing-gateway.example.com/v1", settings)
assert _env(settings)["ANTHROPIC_BASE_URL"] == "https://existing-gateway.example.com/v1"