mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Merge branch 'main' into pr/3208-launchd-exec-handoff
This commit is contained in:
commit
0b28c58f77
58 changed files with 5273 additions and 409 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
6
headroom/cache/compression_cache.py
vendored
6
headroom/cache/compression_cache.py
vendored
|
|
@ -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
|
||||
|
|
|
|||
40
headroom/cache/prefix_tracker.py
vendored
40
headroom/cache/prefix_tracker.py
vendored
|
|
@ -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()
|
||||
|
|
|
|||
22
headroom/cache/semantic.py
vendored
22
headroom/cache/semantic.py
vendored
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 {}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
185
headroom/proxy/session_engine.py
Normal file
185
headroom/proxy/session_engine.py
Normal 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)
|
||||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
26
tests/fixtures/headroom-issue-3054.json
vendored
Normal 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"
|
||||
]
|
||||
}
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
291
tests/test_cli/test_mcp_reconcile.py
Normal file
291
tests/test_cli/test_mcp_reconcile.py
Normal 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
|
||||
124
tests/test_cli/test_serena_reconcile.py
Normal file
124
tests/test_cli/test_serena_reconcile.py
Normal 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
|
||||
|
|
@ -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,
|
||||
},
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
568
tests/test_compress_session_mode.py
Normal file
568
tests/test_compress_session_mode.py
Normal 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"
|
||||
185
tests/test_compression_cache_registry.py
Normal file
185
tests/test_compression_cache_registry.py
Normal 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()
|
||||
|
|
@ -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
|
||||
|
|
|
|||
148
tests/test_copilot_native_mode.py
Normal file
148
tests/test_copilot_native_mode.py
Normal 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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
531
tests/test_openai_responses_read_protection.py
Normal file
531
tests/test_openai_responses_read_protection.py
Normal 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"
|
||||
165
tests/test_org_scale_limits.py
Normal file
165
tests/test_org_scale_limits.py
Normal 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)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
||||
|
|
|
|||
225
tests/test_session_engine.py
Normal file
225
tests/test_session_engine.py
Normal 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
|
||||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
254
tests/test_wrap_concurrent_settings.py
Normal file
254
tests/test_wrap_concurrent_settings.py
Normal 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"
|
||||
Loading…
Add table
Add a link
Reference in a new issue