From cc484864b2f86a5cd1fcc83ac08d2f560e05e1ca Mon Sep 17 00:00:00 2001 From: Parideboy Date: Sun, 23 Aug 2026 18:12:37 +0200 Subject: [PATCH 01/18] docs(troubleshooting): note server-managed settings skip custom ANTHROPIC_BASE_URL (#3118) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Documents a known Claude Code client-side limitation: server-managed settings (delivered from the claude.ai admin console) are silently skipped whenever `ANTHROPIC_BASE_URL` is non-default — which is exactly the condition Headroom wrapping creates. Fixes #3074 by explaining the root cause is upstream, not a Headroom bug, and pointing affected users at the endpoint-managed alternative. ## Type of Change - [x] Documentation - [ ] Bug fix - [ ] New feature - [ ] Breaking change - [ ] Refactor / chore ## Changes Made - Added a new "Server-managed settings unavailable through custom ANTHROPIC_BASE_URL" section to `docs/content/docs/troubleshooting.mdx`, immediately after the existing "Remote Control unavailable through custom ANTHROPIC_BASE_URL" section (same class of Claude-side gate, same Symptom/Cause/Fix format). - Explains why Headroom has no endpoint to implement here (per Anthropic's docs, Claude Code skips the fetch client-side before any request is sent) and distinguishes this from the unrelated, unaffected OS-level `managed-settings.json` file. - Links to Anthropic's official docs and to #3074. ## Testing - [x] Verified locally - [ ] Added/updated automated tests - [ ] N/A ``` $ python3 -c " import re text = open('docs/content/docs/troubleshooting.mdx', encoding='utf-8').read() headings = re.findall(r'^##\s+.*$', text, re.MULTILINE) start = text.index('## Server-managed settings') end = text.index('## Compression Too Aggressive') section = text[start:end] print('backticks even:', section.count(chr(96)) % 2 == 0) print('brackets balanced:', section.count('[') == section.count(']')) print('parens balanced:', section.count('(') == section.count(')')) " backticks even: True brackets balanced: True parens balanced: True ``` ## Real Behavior Proof - Environment: Docs-only change (MDX prose, no code path). `docs/` npm deps are not installed in this sandbox, so the Next.js docs build (`npm run build`) was not run. - Exact command / steps: Diffed the new section against the file's existing neighboring section (git diff), and ran a Python script validating heading structure and backtick/bracket/paren balance within the new section (shown above). - Observed result: New `##` heading inserted cleanly between the two existing sections with no structural changes elsewhere in the file; markdown syntax (bold labels, inline code, links) mirrors the adjacent "Remote Control" section exactly, and is balanced/well-formed. - Not tested: The actual Next.js docs site build/render (`npm run build` in `docs/`) — no network/npm install available in this sandbox. No functional/runtime behavior is affected by this change. ## Runtime Rollout Safety - Rollout-managed feature(s): None - Minimum rollout channel: N/A - Stable/default behavior changed: No - Kill switch / disable path: N/A - Unsafe override required: No - Qualification impact: None - Rollback path: Revert the commit; no state or config is introduced ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review Co-authored-by: Claude Sonnet 5 --- docs/content/docs/troubleshooting.mdx | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/content/docs/troubleshooting.mdx b/docs/content/docs/troubleshooting.mdx index 22ef73ae8..5b796b371 100644 --- a/docs/content/docs/troubleshooting.mdx +++ b/docs/content/docs/troubleshooting.mdx @@ -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 ` 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. From 455f4f263ce2638ff387c52d42601760f5cf0784 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Mon, 24 Aug 2026 00:21:33 +0530 Subject: [PATCH 02/18] fix(cache/semantic): don't semantic-match an empty query across contexts (#3226) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `SemanticCache.get()` matches on the **embedding of the last user message** whenever an `embedding_fn` is wired. That query is empty (`""`) for the overwhelming majority of agent/tool turns — a `tool_result` continuation carries no text block, so `SemanticCacheLayer._extract_query` returns `""`. A real sentence embedder maps `""` to a fixed **non-zero** vector, so every empty-query turn is ~identical to every other in embedding space. The exact `messages_hash` guard (correctly chosen so `"continue"`/`"yes"` turns in different contexts don't collide) is then bypassed by the semantic path: an empty-query request misses on its unique hash, falls through to embedding matching, and hits a **different conversation's** stored response. Reproduction (realistic embedder, non-zero for `""`): ```python c = SemanticCache(embedding_fn=embed) c.put(query="", response={"answer": "A"}, messages_hash="ctxA") # conversation A c.get(query="", messages_hash="ctxB") # conversation B, different context # -> returned A's response (cross-context false hit) ``` Measured on 330 real Claude Code transcripts (28,441 requests): **95.7% have an empty extracted query**, so this is the dominant case, not a corner case. The exact-hash path is unaffected; only the embedding-similarity path is. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/cache/semantic.py`: - `get()`: gate the semantic-similarity branch on `query.strip()` — an empty/blank query can only ever hit via its exact `messages_hash` (context-complete), never via embedding similarity. - `put()`: store no embedding for an empty/blank query, so such an entry is skipped by `_find_similar` (which ignores entries with no embedding) and can never be a match target. - `tests/test_cache/test_semantic.py`: added `test_empty_query_never_semantic_matches` (cross-context empty-query miss, exact-hash still hits, whitespace treated as empty). ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality ### Test Output ```text tests/test_cache/test_semantic.py -> 22 passed in 2.39s uvx ruff@0.16.2 check headroom/cache/semantic.py tests/test_cache/test_semantic.py -> All checks passed! uvx mypy@1.20.2 headroom/cache/semantic.py -> Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12.11, project venv, pytest 9.1.1, ruff 0.16.2 and mypy 1.20.2 via uvx. - Exact command / steps: before the fix, two different-context empty-query requests (`ctxA` then `ctxB`) returned `ctxA`'s response via the embedding path. After the fix, the second returns `None`, while `ctxA`'s own exact-hash lookup still returns its response, and a legitimate non-empty semantic hit (`"What is the weather today?"` -> `"How is the weather?"`) still works. - Observed result: empty/blank queries no longer semantic-match across contexts; exact-hash and non-empty semantic matching are unchanged. - Not tested: no live embedder model wired (the current client wires none — the embedding path is exercised with an injected `embedding_fn`, which is the documented usage). ## Runtime Rollout Safety - Rollout-managed feature(s): none. `SemanticCache` is an SDK-side cache (`headroom.cache`), not a rollout-channel-gated runtime feature; semantic matching only runs when a caller injects an `embedding_fn`. - Minimum rollout channel: N/A. - Stable/default behavior changed: no. Exact-hash matching and non-empty semantic matching are unchanged; only empty/blank-query semantic matching (a false-hit source) is removed. - Kill switch / disable path: N/A. - Unsafe override required: no. - Qualification impact: none; correctness-only. - Rollback path: revert this PR. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation (N/A: internal behavior) - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` --- headroom/cache/semantic.py | 22 ++++++++++++++++++---- tests/test_cache/test_semantic.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/headroom/cache/semantic.py b/headroom/cache/semantic.py index 0ea2effba..946fada32 100644 --- a/headroom/cache/semantic.py +++ b/headroom/cache/semantic.py @@ -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() diff --git a/tests/test_cache/test_semantic.py b/tests/test_cache/test_semantic.py index c24d261f9..f5923f1cb 100644 --- a/tests/test_cache/test_semantic.py +++ b/tests/test_cache/test_semantic.py @@ -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.""" From 7550efb68fde6acf0674244b03eb268b976ffced Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Sun, 23 Aug 2026 14:52:50 -0400 Subject: [PATCH 03/18] fix(mcp): add explicit Serena reconciliation (#3222) ## Description Headroom repeatedly warns about user-managed Serena drift but has no scoped remediation command. Add a Claude-only read-only mcp reconcile command with explicit --adopt consent, using the canonical Serena spec and existing Claude registrar. Adoption validates every relevant ledger and Claude config root before mutation, writes only the Serena entry, and records ownership after the config write succeeds. Automatic wrap migration and ordinary install remain unchanged. Closes #3054 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (feature that would cause existing behavior to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Add Claude-only `headroom mcp reconcile`, read-only by default, with `--adopt` as its only mutation action. - Reuse the shared `CLAUDE_SERENA_CONTEXT` and canonical Claude Serena spec builder. - Fail closed on malformed or unreadable ledger/config state before adoption. - Preserve automatic wrap recovery, user-managed warnings, ordinary `mcp install --force`, unrelated Claude config, and corrupt-ledger tolerance outside explicit adoption. - Record Headroom ownership only after a successful registrar write. ## Testing - [x] Unit tests pass (`uv run pytest tests/test_cli/test_mcp_reconcile.py tests/test_cli/test_serena_reconcile.py tests/test_mcp_registry/test_ledger.py`) - [x] Linting passes (`uv run ruff check .`) - [ ] Type checking passes (`uv run mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed through the file-backed Claude registrar ### Test Output ```text uv run pytest tests/test_cli/test_mcp_reconcile.py tests/test_mcp_registry/test_ledger.py tests/test_cli/test_serena_reconcile.py tests/test_mcp_registry/test_claude_registrar.py tests/test_mcp_registry/test_install.py -q 102 passed in 0.70s uv run ruff check headroom/mcp_registry/ledger.py headroom/cli/wrap.py tests/test_mcp_registry/test_ledger.py tests/test_cli/test_serena_reconcile.py tests/test_cli/test_mcp_reconcile.py All checks passed! uv run ruff format --check headroom/mcp_registry/ledger.py headroom/cli/wrap.py tests/test_mcp_registry/test_ledger.py tests/test_cli/test_serena_reconcile.py tests/test_cli/test_mcp_reconcile.py 5 files already formatted git diff --check ``` ## Real Behavior Proof - Environment: Windows, file-backed Claude configuration and isolated MCP ledger. - Exact command / steps: run the stale user-managed Serena fixture from `tests/fixtures/headroom-issue-3054.json`; run read-only reconcile; run `mcp reconcile --adopt`; rerun wrap and ordinary `mcp install --force`; exercise malformed JSON, non-dict `mcpServers`, null ledger agents, and unreadable-ledger adoption. - Observed result: read-only reconciliation leaves config and ledger bytes unchanged; adoption updates only Claude Serena and records ownership after a successful write; automatic wrap remains lenient; unsafe adoption inputs leave all files unchanged; ordinary install does not adopt Serena. - Not tested: live Claude CLI acceptance and Serena stdio handshake ## Runtime Rollout Safety - Rollout-managed feature(s): None; explicit `mcp reconcile --adopt` is the only mutation path. - Minimum rollout channel: Stable; no staged rollout mechanism exists for this command. - Stable/default behavior changed: No, read-only reconcile is the default and automatic wrap plus ordinary install remain unchanged. - Kill switch / disable path: Do not invoke `--adopt` or revert the release commit. - Unsafe override required: No. - Qualification impact: None. - Rollback path: Revert the release commit. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes The changelog is generated by the release pipeline. This change is limited to Claude Serena reconciliation and does not add a new persistent acknowledgement state or a multi-provider adoption route. --- docs/content/docs/mcp.mdx | 2 + headroom/cli/mcp.py | 46 ++++ headroom/cli/wrap.py | 23 +- headroom/mcp_registry/__init__.py | 5 +- headroom/mcp_registry/claude.py | 32 +++ headroom/mcp_registry/install.py | 1 + headroom/mcp_registry/ledger.py | 53 ++++- tests/fixtures/headroom-issue-3054.json | 26 +++ tests/test_cli/test_mcp_reconcile.py | 291 ++++++++++++++++++++++++ tests/test_cli/test_serena_reconcile.py | 124 ++++++++++ tests/test_mcp_registry/test_ledger.py | 100 +++++--- 11 files changed, 660 insertions(+), 43 deletions(-) create mode 100644 tests/fixtures/headroom-issue-3054.json create mode 100644 tests/test_cli/test_mcp_reconcile.py create mode 100644 tests/test_cli/test_serena_reconcile.py diff --git a/docs/content/docs/mcp.mdx b/docs/content/docs/mcp.mdx index af6302825..7fe121ff8 100644 --- a/docs/content/docs/mcp.mdx +++ b/docs/content/docs/mcp.mdx @@ -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. diff --git a/headroom/cli/mcp.py b/headroom/cli/mcp.py index 7563822bb..40e8c25ca 100644 --- a/headroom/cli/mcp.py +++ b/headroom/cli/mcp.py @@ -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. diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 638d9e56a..1c7176e2c 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -2018,17 +2018,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 +2042,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: @@ -4932,11 +4939,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, diff --git a/headroom/mcp_registry/__init__.py b/headroom/mcp_registry/__init__.py index 5f1ff4f8f..f19ad1ba7 100644 --- a/headroom/mcp_registry/__init__.py +++ b/headroom/mcp_registry/__init__.py @@ -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", diff --git a/headroom/mcp_registry/claude.py b/headroom/mcp_registry/claude.py index 7acbc02fb..de7ee308c 100644 --- a/headroom/mcp_registry/claude.py +++ b/headroom/mcp_registry/claude.py @@ -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: diff --git a/headroom/mcp_registry/install.py b/headroom/mcp_registry/install.py index e550f6ded..4d1fbeb72 100644 --- a/headroom/mcp_registry/install.py +++ b/headroom/mcp_registry/install.py @@ -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]: diff --git a/headroom/mcp_registry/ledger.py b/headroom/mcp_registry/ledger.py index 6bb0160e4..e802227e7 100644 --- a/headroom/mcp_registry/ledger.py +++ b/headroom/mcp_registry/ledger.py @@ -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: diff --git a/tests/fixtures/headroom-issue-3054.json b/tests/fixtures/headroom-issue-3054.json new file mode 100644 index 000000000..e5eb4f12b --- /dev/null +++ b/tests/fixtures/headroom-issue-3054.json @@ -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" + ] +} diff --git a/tests/test_cli/test_mcp_reconcile.py b/tests/test_cli/test_mcp_reconcile.py new file mode 100644 index 000000000..519aa0591 --- /dev/null +++ b/tests/test_cli/test_mcp_reconcile.py @@ -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 diff --git a/tests/test_cli/test_serena_reconcile.py b/tests/test_cli/test_serena_reconcile.py new file mode 100644 index 000000000..f65cbea1a --- /dev/null +++ b/tests/test_cli/test_serena_reconcile.py @@ -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 diff --git a/tests/test_mcp_registry/test_ledger.py b/tests/test_mcp_registry/test_ledger.py index 73100a32b..aa7ff08dc 100644 --- a/tests/test_mcp_registry/test_ledger.py +++ b/tests/test_mcp_registry/test_ledger.py @@ -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) From 7784bb184614bdcd7ba1afd93a23e8d1458ecc96 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ra=C3=BAl?= Date: Sun, 23 Aug 2026 20:53:17 +0200 Subject: [PATCH 04/18] fix(transforms): stop folding datetime-prefixed user messages as search results (#3221) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Interactive `headroom wrap copilot` sessions intermittently lose the user's message: the model answers "How can I help you today?" to a real task prompt. Root cause: Copilot CLI prepends `…` to every interactive user turn; the ISO-8601 timestamp matches the grep `file:line:` detector, so a datetime + one-line prompt (1 match / 2 non-empty lines = 50% ≥ 30%) classifies as `SEARCH_RESULTS`, and `SearchCompressor` — which keeps only detector-matching lines — deletes the prompt before upstream. On the OpenAI chat streaming path there is no retrieval tool, so the loss is unrecoverable. Fix: `_try_detect_search` now (a) requires the pre-colon segment to look like a file path (no `<`, `>`, `=`), and (b) requires at least two matching lines, so one coincidental `word:digits:` line can no longer classify a whole payload. A genuine one-line grep result loses nothing: all its lines match, so the compressor would have kept it verbatim anyway. Closes #3220 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `headroom/transforms/content_detector.py`: new `_is_search_result_line` helper (path-like prefix gate); `_try_detect_search` gains a two-matching-line absolute floor. - `tests/test_transforms_content_detection.py`: regression tests — datetime-prefixed one-liner not search; two-line floor; tag-like / `key=value` prefixes rejected; genuine grep output still detected. - `tests/test_transforms_content_router.py`: router-level regression — the incident payload never routes to SEARCH and the prose survives `ContentRouter().compress()`. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ .venv/bin/python -m pytest tests/test_transforms_content_detection.py tests/test_transforms_content_router.py tests/test_mixed_content_sections.py tests/test_text_compressors.py tests/test_transforms_tabular.py -q 135 passed in 21.06s $ .venv/bin/ruff check headroom/transforms/content_detector.py tests/test_transforms_content_detection.py tests/test_transforms_content_router.py All checks passed! $ .venv/bin/mypy headroom/transforms/content_detector.py Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: macOS 26.5 arm64, Python 3.13, editable source build 0.37.0-dev; upstream `api.githubcopilot.com`, cheapest subscription model `kimi-k2.7-code`. - Exact command / steps: standalone copilot-routed proxy (`OPENAI_TARGET_API_URL=https://api.githubcopilot.com headroom proxy --port 8899`) + `.overlay/e2e-copilot-content-probe.sh --port 8899 --model kimi-k2.7-code`, which sends the real interactive wire shape (`…` + one-line sentinel prompt, streaming) and a multi-line control. - Observed result: BEFORE the fix, probe 1 FAIL — model replied "Hello! I see the current datetime is … How can I assist you today?" with proxy log `transforms=router:search:0.50` (prompt deleted). AFTER the fix, both probes PASS — the sentinel echoes verbatim, proving the user message reached upstream intact. - Not tested: other harnesses' interactive wrappers (claude/droid/auggie send different shapes; the detector fix is generic); the mixed-content section splitter has its own grep pattern (out of scope — its 1-line "search" sections are kept verbatim, no data loss). ## Runtime Rollout Safety - Rollout-managed feature(s): none. - Minimum rollout channel: N/A (no flag). - Stable/default behavior changed: content with exactly one `path:line:`-shaped line no longer classifies as search results (stays uncompressed instead — safe direction; compression only ever engages on ≥2 matching lines now). - Kill switch / disable path: N/A. - Unsafe override required: no. - Qualification impact: none. - Rollback path: revert; prior behavior restores (with the bug). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` — it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Screenshots (if applicable) N/A — proxy transform change; no UI. ## Additional Notes Detection-precision tradeoff is documented in code comments: single-line genuine grep output is no longer folded (no data loss either way — the compressor keeps all-matching content verbatim). A residual edge (prose with ≥2 coincidental `x:1:` lines in ≤6 lines) is accepted and documented in the issue. --- headroom/transforms/content_detector.py | 27 +++++++++++++- tests/test_transforms_content_detection.py | 43 ++++++++++++++++++++++ tests/test_transforms_content_router.py | 20 ++++++++++ 3 files changed, 88 insertions(+), 2 deletions(-) diff --git a/headroom/transforms/content_detector.py b/headroom/transforms/content_detector.py index 5e364fed6..bb8325e42 100644 --- a/headroom/transforms/content_detector.py +++ b/headroom/transforms/content_detector.py @@ -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 ``…`` 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 diff --git a/tests/test_transforms_content_detection.py b/tests/test_transforms_content_detection.py index 1201c2a60..e74776292 100644 --- a/tests/test_transforms_content_detection.py +++ b/tests/test_transforms_content_detection.py @@ -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 ``…`` 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 = ( + "2026-08-23T09:57:59.792+02:00\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('started\nstopped') + 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( [ diff --git a/tests/test_transforms_content_router.py b/tests/test_transforms_content_router.py index 4f76c6386..74ce9b5c9 100644 --- a/tests/test_transforms_content_router.py +++ b/tests/test_transforms_content_router.py @@ -1785,3 +1785,23 @@ def test_detect_content_overrides_html_misroute_for_grep_and_logs( "" ) 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 ``…`` 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 = ( + "2026-08-23T09:57:59.792+02:00\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 From 701e4616d92d6d4b110bf5356675b6db409802e8 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Sun, 23 Aug 2026 22:55:14 -0400 Subject: [PATCH 05/18] fix(kimi): route managed Kimi Code through the proxy (#3223) ## Description Managed Kimi Code reads KIMI_CODE_BASE_URL while headroom wrap kimi previously supplied only KIMI_BASE_URL. The managed client can therefore keep its direct endpoint while the wrapper appears healthy. Emit both provider-owned keys and recompute them through the existing launch callback at the proxy's actual port. Preserve the legacy route and unrelated wrappers. Closes #3207 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (feature that would cause existing behavior to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Set KIMI_CODE_BASE_URL and KIMI_BASE_URL from one project-aware proxy URL. - Recompute both values and their display lines through the Kimi configure_launch callback after port fallback. - Remove the generic display rewrite from _launch_tool so other wrappers retain their base behavior. - Add production-boundary child, fallback-port, legacy-preservation, and non-Kimi negative-space tests. ## Testing - [x] Unit tests pass (`uv run pytest tests/test_cli/test_wrap_kimi.py -q`) - [x] Linting passes (`uv run ruff check .`) - [ ] Type checking passes (`uv run mypy headroom`) - [x] New tests added for the regression - [x] Manual testing performed through the production subprocess boundary ### Test Output ```text uv run pytest tests/test_cli/test_wrap_kimi.py -q 10 passed in 0.40s uv run pytest tests/test_cli/test_wrap_grok.py -q 2 passed uv run ruff check . All checks passed! uv run ruff format . --check 1534 files already formatted git diff --check ``` ## Real Behavior Proof - Environment: Windows, isolated Kimi wrapper subprocess harness. - Exact command / steps: launch a contract-compatible child through the Kimi wrapper; exercise requested and fallback ports, project prefixes, legacy selection, and a non-Kimi wrapper. - Observed result: the child receives the effective project-aware proxy URL in both Kimi keys; the displayed URL matches it after fallback; legacy and non-Kimi behavior remain unchanged. - Not tested: live authenticated Kimi Code managed request ## Runtime Rollout Safety - Rollout-managed feature(s): None; managed Kimi Code routing is selected by the existing wrapper mode. - Minimum rollout channel: Stable; no staged rollout mechanism exists for this wrapper path. - Stable/default behavior changed: Yes, managed Kimi Code launches now receive the effective proxy URL in both provider-owned keys. - Kill switch / disable path: Stop using the managed Kimi wrapper path or revert the release commit. - Unsafe override required: No. - Qualification impact: None. - Rollback path: Revert the release commit. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes Kimi Code owns OAuth credentials and the /login flow. Headroom does not read or modify Kimi config or credential files. The changelog is generated by the release pipeline. --- headroom/cli/wrap.py | 27 +++++-- headroom/providers/kimi/runtime.py | 16 ++-- tests/test_cli/test_wrap_kimi.py | 114 ++++++++++++++++++++++++++++- 3 files changed, 141 insertions(+), 16 deletions(-) diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 1c7176e2c..2a6c7d936 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -4393,7 +4393,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}") @@ -6302,9 +6302,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: @@ -6322,9 +6323,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, @@ -6339,6 +6351,7 @@ def kimi( agent_type="kimi", code_graph=code_graph, openai_api_url=kimi_api_url, + configure_launch=configure_kimi_launch, ) diff --git a/headroom/providers/kimi/runtime.py b/headroom/providers/kimi/runtime.py index fa2a30678..1d89e7a8f 100644 --- a/headroom/providers/kimi/runtime.py +++ b/headroom/providers/kimi/runtime.py @@ -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/`` 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}", + ] diff --git a/tests/test_cli/test_wrap_kimi.py b/tests/test_cli/test_wrap_kimi.py index 01f94b45c..05ce3cc61 100644 --- a/tests/test_cli/test_wrap_kimi.py +++ b/tests/test_cli/test_wrap_kimi.py @@ -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, From f27f235032d8aef522efbcb4daf548787692c7e4 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Sun, 23 Aug 2026 22:36:17 -0700 Subject: [PATCH 06/18] fix(wrap): stop concurrent wrap sessions clobbering settings.local.json (#3232) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Several `headroom wrap` sessions in one project each write the proxy URL into `.claude/settings.local.json` and restore it on exit. That read-modify-write was unsynchronised. The write itself is atomic so the file never tears, but the updates were still lost against each other: - **Live sessions were silently unrouted.** The first session to exit deleted the key while its siblings were still running. They kept working, but their traffic stopped going through the proxy — no error, no warning, no savings. - **A dead proxy was written back into the project.** A session that started second captured the *first* session's proxy URL as "the original", so its exit restored a URL pointing at a port that was already gone. Every later session in that project then failed to connect. - **SIGTERM/SIGHUP never ran the restore at all.** `cleanup` was registered as the handler, but a Python signal handler that returns normally does not unwind the stack — under PEP 475 the interrupted `waitpid` is simply retried. The `finally` block that restores `settings.local.json` never ran, while the handler had already terminated the proxy underneath a child that was still alive. Closes #3205 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - **`_wrap_settings_lock`** — an exclusive OS lock (flock / `msvcrt.locking`) held across the settings read-modify-write. A workspace that cannot hold lock state degrades to the previous behaviour rather than failing, matching `_proxy_start_lock`. - **`.headroom_wrap_owners.json`** — a sidecar recording, per env key, the true pre-wrap `original` plus the live sessions holding it. The first writer records the original; later writers inherit it and are flagged `inherited`, so no session restores a value it did not observe first-hand. A session exits without restoring while a sibling still holds the key. Dead holders are pruned with the same conservative PID+identity liveness the proxy-client markers use, so a SIGKILLed session cannot wedge the key. - **`unwrap` passes `force=True`** — unwrap is the user explicitly asking for their settings back, so it drops every claim instead of deferring to a live sibling and silently printing success while leaving the proxy URL in the file. - **The #2221 self-heal passes `dead_ports`** — a wrapper process can outlive its proxy (proxy alone SIGKILLed). Its claim would otherwise veto the self-heal and leave `ANTHROPIC_BASE_URL` pointing at a port just proven dead. - **`_rehome_wrap_marker`** — the wrap marker has one slot, won by the last writer. When that writer exits while a sibling still owns the key, the marker is rewritten to describe the survivor (carrying the record's true original), so the survivor keeps its #2221 self-heal record instead of being left with a marker describing a dead process. - **`_exit_on_signal`** replaces `cleanup` as the SIGTERM/SIGHUP handler. Raising `SystemExit` unwinds, so the settings restore actually runs and cleanup happens exactly once from `finally`. - **`_proxy_start_lock` now shares `_locked_file`** with the new settings lock rather than carrying a second verbatim copy of the platform branches. ## Testing - [x] Unit tests pass (`pytest`) — full suite, 11518 passed / 588 skipped - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed `tests/test_wrap_concurrent_settings.py` (14 tests) covers: a sibling exit leaving survivors routed, the last session out restoring the true original, a pre-existing user URL surviving the whole cycle, three sessions in every exit order, a crashed session not wedging the key, forced unwrap past a live session, a holder that outlived its proxy not vetoing the self-heal, marker rehoming, and the signal-handler unwind. ### Test Output ```text $ uv run pytest tests/test_cli/test_unwrap_claude.py tests/test_cli/test_wrap_claude_base_url.py \ tests/test_cli/test_wrap_claude_finally_unbound.py tests/test_cli/test_wrap_claude_vertex_proxy_env.py \ tests/test_cli/test_wrap_claude.py tests/test_cli/test_wrap_dead_marker_selfheal.py \ tests/test_cli/test_wrap_helpers.py tests/test_cli/test_wrap_stale_marker.py \ tests/test_cli/test_wrap_persistent.py tests/test_wrap_concurrent_settings.py tests/test_cli_doctor.py -q tests/test_wrap_concurrent_settings.py .............. [ 72%] tests/test_cli_doctor.py ............................................... [ 89%] ............................... [100%] ============================= 285 passed in 3.01s ============================== $ uv run pytest tests/ -q ======== 11518 passed, 588 skipped, 6036 warnings in 1831.34s (0:30:31) ======== $ uv run ruff check . All checks passed! $ uv run mypy headroom Success: no issues found in 527 source files ``` ## Real Behavior Proof - **Environment:** macOS 15 (Darwin 25.4.0), Python 3.12.13, repo venv, Claude provider path (`ANTHROPIC_BASE_URL` in `.claude/settings.local.json`). - **Exact command / steps:** a script spawning **two real OS processes** — no mocks, real PIDs, real files — that call the same `_write_claude_wrap_base_url` / `_restore_claude_wrap_base_url` helpers `wrap claude` uses. The project starts with a real user gateway already set. Session A (port 8787) starts, session B (port 8788) starts 0.7s later, A exits while B is still running, then B exits. Run identically on `main` and on this branch. **Before (on `main`) — both bugs visible:** ```text start : {"ANTHROPIC_BASE_URL": "https://my-gateway.example.com"} session port=8787 started, remembers previous='https://my-gateway.example.com' session port=8788 started, remembers previous='http://127.0.0.1:8787' both sessions running : {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8788"} session port=8787 exited after FIRST session exits : {"ANTHROPIC_BASE_URL": "https://my-gateway.example.com"} session port=8788 exited after LAST session exits : {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787"} ``` Session B is still running, but after A exits the proxy URL is gone from under it — B is unrouted with no error. And the final state is `http://127.0.0.1:8787`: a dead proxy left permanently in the user's project, with their real gateway lost. **After (this branch):** ```text start : {"ANTHROPIC_BASE_URL": "https://my-gateway.example.com"} session port=8787 started, remembers previous='https://my-gateway.example.com' session port=8788 started, remembers previous='http://127.0.0.1:8787' both sessions running : {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8788"} session port=8787 exited after FIRST session exits : {"ANTHROPIC_BASE_URL": "http://127.0.0.1:8788"} session port=8788 exited after LAST session exits : {"ANTHROPIC_BASE_URL": "https://my-gateway.example.com"} ``` B stays routed after A exits, and the last session out restores the user's real gateway. - **Observed result:** matches the intent on both counts — no unrouting, no dead proxy residue, user's pre-existing URL preserved. - **Not tested:** Windows (`msvcrt.locking`) — the lock and dead-holder pruning are exercised on POSIX only; the Windows branch is the same code path `_proxy_start_lock` has shipped with. No live end-to-end run against a real Anthropic endpoint with two concurrent `claude` CLIs; the proof above drives the same helpers out of two real processes instead. Foundry/Vertex key variants are covered by unit tests, not by a live run. Real SIGTERM/SIGHUP delivery to a running `wrap claude` was not exercised end to end — the handler's unwind is covered by a unit test, and full signal delivery would need a spawned and killed subprocess, which the existing #1768 test also declined to do. ## Runtime Rollout Safety - **Rollout-managed feature(s):** none — this is an unconditional correctness fix on the wrap settings path. - **Minimum rollout channel:** n/a. - **Stable/default behavior changed:** yes, three ways. (1) A wrap session exiting while a sibling holds the key now leaves the key in place instead of removing it. (2) SIGTERM/SIGHUP now unwinds, so the child CLI is terminated by `subprocess.run`'s cleanup rather than being left running against a torn-down proxy. (3) Two new sidecar files appear next to `settings.local.json`: `.headroom_wrap_owners.json` (removed when the last holder exits) and `.headroom_wrap_settings.lock` (retained by design — deleting a live lock file creates an inode-replacement race). - **Kill switch / disable path:** none. A workspace where the lock file cannot be created degrades to the previous unsynchronised behaviour automatically. - **Unsafe override required:** no. - **Qualification impact:** none beyond the wrap settings path. - **Rollback path:** revert the commit; the sidecar files are ignored by older versions and can be deleted safely. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` ## Additional Notes - The ownership record is keyed per env key, so `ANTHROPIC_BASE_URL`, the Foundry/Vertex variants and the tool-search entry are tracked independently. - Documentation: the behaviour is documented in the helper docstrings rather than user-facing docs — the sidecar files are internal state a user never configures. - Follow-up worth considering: `.headroom_wrap_settings.lock` is intentionally never deleted (matching `_proxy_start_lock`'s retention rationale), so it stays in `.claude/` after `unwrap`. Removing it safely needs a separate think about the inode-replacement race. Co-authored-by: Tejas Chopra Co-authored-by: Claude Opus 5 --- headroom/cli/wrap.py | 459 +++++++++++++++++++---- tests/test_cli/test_unwrap_claude.py | 6 + tests/test_cli/test_wrap_stale_marker.py | 20 +- tests/test_wrap_concurrent_settings.py | 254 +++++++++++++ 4 files changed, 660 insertions(+), 79 deletions(-) create mode 100644 tests/test_wrap_concurrent_settings.py diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 2a6c7d936..cc87ba6e0 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -36,7 +36,7 @@ from collections.abc import Callable 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( @@ -4104,39 +4420,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) @@ -4330,6 +4615,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, @@ -4361,7 +4660,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() @@ -4833,11 +5132,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: @@ -5222,6 +5521,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 diff --git a/tests/test_cli/test_unwrap_claude.py b/tests/test_cli/test_unwrap_claude.py index a562e05a3..e3b688196 100644 --- a/tests/test_cli/test_unwrap_claude.py +++ b/tests/test_cli/test_unwrap_claude.py @@ -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, }, ] diff --git a/tests/test_cli/test_wrap_stale_marker.py b/tests/test_cli/test_wrap_stale_marker.py index 966c551b2..6fb50fa08 100644 --- a/tests/test_cli/test_wrap_stale_marker.py +++ b/tests/test_cli/test_wrap_stale_marker.py @@ -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 diff --git a/tests/test_wrap_concurrent_settings.py b/tests/test_wrap_concurrent_settings.py new file mode 100644 index 000000000..f0ecfe76e --- /dev/null +++ b/tests/test_wrap_concurrent_settings.py @@ -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" From b9d7dcc3da3ec67a968d0ac1e35bffba9b7cf1c2 Mon Sep 17 00:00:00 2001 From: inix <62450194+inix-x@users.noreply.github.com> Date: Mon, 24 Aug 2026 13:43:58 +0800 Subject: [PATCH 07/18] fix(proxy): make output-savings flush atomic and keep it off the event loop (#3231) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description The output-shaper's periodic savings-ledger flush ran synchronously on the asyncio event loop: every 25th shaped request, `emit_request_outcome` performed a full ledger reload (file read + `json.loads`) followed by a `json.dumps` + in-place `write_text`, with no await or executor. The write was also non-atomic, so a crash mid-write truncated the existing ledger, and `SavingsLedger.load()` silently swallowed the resulting decode error — corrupted history was indistinguishable from no history yet. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [x] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `emit_request_outcome` now runs `record_from_labels` + `estimate_request_savings` together on a worker thread via one `asyncio.to_thread` call — both take the recorder lock, and the periodic flush holds that lock across disk I/O, so nothing touches it from the event loop anymore. - `SavingsLedger.save()` writes through the existing `headroom.fsutil.write_text` helper (temp file in the target directory, fsync, atomic `os.replace`, temp cleanup on failure) instead of a truncating in-place write. - `SavingsLedger.load()` logs a warning naming the unreadable ledger file and still fails open with an empty ledger. - Added `TestFlushDurability` to `tests/test_output_savings.py`: failed-save intactness (+ no temp residue), corrupt-file warning, and off-loop-thread assertions. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ python -m pytest tests/test_output_savings.py tests/test_output_shaping_rollup.py tests/test_output_savings_cli.py -q ============================== 49 passed in 1.38s ============================== $ ruff check headroom/proxy/output_savings.py headroom/proxy/outcome.py tests/test_output_savings.py All checks passed! $ ruff format --check headroom/proxy/output_savings.py headroom/proxy/outcome.py tests/test_output_savings.py 3 files already formatted $ mypy headroom/proxy/output_savings.py headroom/proxy/outcome.py Success: no issues found in 2 source files Fail-before evidence (the three files checked out at upstream/main, fix reverted): $ python -m pytest tests/test_output_savings.py::TestFlushDurability -q FAILED tests/test_output_savings.py::TestFlushDurability::test_crash_mid_write_leaves_previous_ledger_intact - KeyError: 'opus|code|m|tools' FAILED tests/test_output_savings.py::TestFlushDurability::test_corrupt_ledger_warns_and_starts_empty - AssertionError: corrupt ledger was swallowed silently FAILED tests/test_output_savings.py::TestFlushDurability::test_emit_request_outcome_flushes_off_the_loop_thread - assert False 3 failed in 0.49s ``` ## Real Behavior Proof - Environment: macOS 15 (arm64), CPython 3.13.13, project venv; branch `fix/output-savings-atomic-offload` = upstream/main `7784bb18` + the single fix commit. - Exact command / steps: With the three touched files reverted to upstream/main: `python -m pytest tests/test_output_savings.py::TestFlushDurability -q` → all 3 new tests fail (torn write destroys the prior ledger; no warning on a corrupt file; flush observed on the loop thread). Re-applied the commit and re-ran the same command plus ruff/format/mypy as pasted under Test Output. - Observed result: All 3 fail-before cases now pass — a save failure before rename leaves the previous ledger loadable with no `*.tmp` residue, a corrupt ledger logs a warning and still fails open empty, and the flush triggered through `emit_request_outcome` runs on a worker thread distinct from the event-loop thread; 49 recorder/rollup/CLI tests pass. - Not tested: Windows behavior of the atomic rename (covered by `fsutil.write_text`, exercised only on POSIX here), a full local suite run (unrelated pre-existing native hangs on macOS), and the dashboard rendering of the ledger. ## Runtime Rollout Safety - Rollout-managed feature(s): None changed. The shaper itself is opt-in; this PR only changes how and where its ledger persistence happens. - Minimum rollout channel: Stable. - Stable/default behavior changed: No. With output shaping disabled the funnel never reaches this code path; when enabled, identical data is persisted — written atomically instead of truncating, and off the event loop. - Kill switch / disable path: Unset `HEADROOM_OUTPUT_SHAPER` (or disable the `proxy_output_shaper` rollout flag); the recorder then neither records nor flushes. - Unsafe override required: No. - Qualification impact: None. - Rollback path: Revert this single commit; the on-disk ledger format is unchanged, so no data migration is involved either way. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` — it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Additional Notes - Two failures seen while running neighbouring suites locally (`tests/test_stateless_writers.py::test_memory_disabled_under_stateless`, `tests/test_5xx_accounting_all_providers.py::test_gemini_count_tokens_handler_threads_real_529_onto_outcome`) reproduce on pristine upstream/main without this patch — pre-existing, not introduced here. - Known adjacent shapes deliberately left out of scope: `get_recorder().estimate()` in the `/stats` payload also reads the ledger file inline (already exception-guarded there), and `SavingsRecorder.flush()` is not yet wired into graceful shutdown. --- headroom/proxy/outcome.py | 16 +++-- headroom/proxy/output_savings.py | 15 ++++- tests/test_output_savings.py | 103 +++++++++++++++++++++++++++++-- 3 files changed, 122 insertions(+), 12 deletions(-) diff --git a/headroom/proxy/outcome.py b/headroom/proxy/outcome.py index c2b78b1dc..ba94feee7 100644 --- a/headroom/proxy/outcome.py +++ b/headroom/proxy/outcome.py @@ -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 diff --git a/headroom/proxy/output_savings.py b/headroom/proxy/output_savings.py index b0ee6baf3..5fe317cb0 100644 --- a/headroom/proxy/output_savings.py +++ b/headroom/proxy/output_savings.py @@ -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() diff --git a/tests/test_output_savings.py b/tests/test_output_savings.py index e9d67a826..f62bc3e37 100644 --- a/tests/test_output_savings.py +++ b/tests/test_output_savings.py @@ -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) From 4408e881064741a5113be69d631f336ce92631f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ra=C3=BAl?= Date: Mon, 24 Aug 2026 13:03:03 +0200 Subject: [PATCH 08/18] fix(proxy): protect file reads from lossy compression on the Responses API path (Copilot view + HEADROOM_PROTECT_READS) (#3238) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description On the OpenAI Responses API path (used by `headroom wrap copilot` and Codex), fresh file reads were lossy-compressed one turn after production, so the model saw its own just-read file content garbled (Kompress word-dropping) and had to re-read it — the exact turn inflation `HEADROOM_PROTECT_READS` was built to prevent on the chat/Anthropic path. Two gaps combined: 1. Copilot CLI's `view` tool (its file-read tool) was not in `DEFAULT_EXCLUDE_TOOLS` — the set only covered Claude-Code names (`Read`, `Write`, …). 2. `_compress_openai_responses_live_text_units_with_router` (`headroom/proxy/handlers/openai.py`) protected only excluded tool *names* and never implemented the `HEADROOM_PROTECT_READS` read-command detection that `ContentRouter.apply()` has — so `bash` reads like `nl -ba FILE | sed -n '1,75p'` were lossy-compressed even with the `coding` profile's `protect_reads=True`. Fix (design adversarially reviewed with gpt-5.6-sol before implementation; verdict "correct with modifications" — all modifications adopted): - `view` added to **both** `DEFAULT_EXCLUDE_TOOLS` and `DEFAULT_VERBATIM_EXCLUDE_TOOLS` — byte-exact contract: no lossy compression, no lossless JSON rewrite, no cross-turn dedup fold. - Responses units path now ports the read-command guard: the producing command is normalized from both wire shapes (`function_call.arguments`, `local_shell_call.action` argv/string) via the shared `_tool_call_command_text`; each output is content-gated by `_read_output_should_be_protected` (lockfiles/JSON/logs/search stay compressible); protected ids are unioned into the dedup protection set. - Shared `read_protection_enabled()` env helper extracted in `content_router.py`, used by both paths. - Latent debug-path defect fixed (unbound `fold` when an excluded tool's output is a content-part list and debug logging is enabled). Follow-up (not in scope): Rust Responses path (`crates/headroom-core/src/transforms/live_zone.rs`) currently only protects `headroom_retrieve` — needs parity before that runtime becomes default. Closes #3237 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - `headroom/config.py` — `view` in both exclusion sets. - `headroom/proxy/handlers/openai.py` — read-command protection for the Responses units path; dedup shield; debug-path fix. - `headroom/transforms/content_router.py` — shared `read_protection_enabled()` helper (both paths). - `tests/test_openai_responses_read_protection.py` — 16 regression tests. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text # Before fix (first commit on this branch, repro-only): 2 failed, 1 passed # view read compressed; bash nl|sed read compressed despite HEADROOM_PROTECT_READS=1 # After fix: $ uv run pytest tests/test_openai_responses_read_protection.py -q 16 passed (incl. content-gate release, string-form local_shell_call, debug paths, scan robustness) $ uv run pytest tests/test_openai_responses_compression_units.py tests/test_responses_cross_turn_dedup.py \ tests/test_lossless_excluded_compaction.py tests/test_observed_wire_shapes.py \ tests/test_content_router_exclude_tools.py tests/test_content_router_compact_json.py -q 69 passed in 3.85s $ uv run ruff check && uv run ruff format --check All checks passed! 4 files already formatted $ uv run mypy headroom/config.py headroom/transforms/content_router.py headroom/proxy/handlers/openai.py Success: no issues found ``` ## Real Behavior Proof - Environment: macOS (Darwin), Python 3.13.7, headroom proxy 0.37.0-dev, `HEADROOM_STACK=wrap_copilot`, savings profile `coding` (effective per proxy banner), model `gpt-5.6-luna` via GitHub Copilot API. - Exact command / steps: incident forensics on Copilot CLI session `5487d36f-3e7d-48b0-a56a-a92e4969c17b` (events.jsonl tool results byte-matched to proxy log compression units), then the failing→passing repro above. - Observed result: (pre-fix proxy log `~/.headroom/logs/proxy-8794.log`) ```text 08:46:34 [hr_1787553986_000014] WS /v1/responses slow compression unit … strategy=text … bytes=3857 … tokens_saved=235 08:46:34 [hr_1787553986_000014] … strategy=text … bytes=6852 … tokens_saved=413 08:46:34 [hr_1787553986_000014] … strategy=text … bytes=3079 … tokens_saved=178 08:46:47 [hr_1787554001_000015] … strategy=text … bytes=6924 … tokens_saved=478 08:46:47 [hr_1787554001_000015] … strategy=text … bytes=4418 … tokens_saved=305 ``` Byte sizes match the session's `view` (3857/6852/3079) and `nl|sed` (6924/4418) tool results exactly. Post-fix, those payload shapes are byte-exact through `_compress_openai_responses_live_text_units_with_router` (asserted by the regression tests over the same wire shapes). - Not tested: full `pytest tests/` run (upstream suite has pre-existing order-dependent failures — 7 failed on clean `main` under `-k "content_router or read or protect"` — and a pre-existing `litellm` import error in `tests/test_memory_eval.py`; the 5 additional failures in that selection with my branch pass in isolation and also fail on clean main under the same selection); live end-to-end with a running Copilot wrap (unit-level wire-shape coverage instead); Rust core path (follow-up). ## Runtime Rollout Safety - Rollout-managed feature(s): none - Minimum rollout channel: N/A - Stable/default behavior changed: yes — `view` outputs and `HEADROOM_PROTECT_READS`-covered bash read outputs stay verbatim on the Responses path (fidelity improvement; slightly fewer tokens saved) - Kill switch / disable path: `HEADROOM_PROTECT_READS=0` restores old bash-read behavior; `HEADROOM_EXCLUDE_TOOLS` overrides tool exclusion - Unsafe override required: no - Qualification impact: none - Rollback path: revert the commit; no state/migration ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation (N/A, no user-facing docs for this internal guard) - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` --- headroom/config.py | 7 + headroom/proxy/handlers/openai.py | 64 ++- headroom/transforms/content_router.py | 28 +- .../test_openai_responses_read_protection.py | 531 ++++++++++++++++++ 4 files changed, 616 insertions(+), 14 deletions(-) create mode 100644 tests/test_openai_responses_read_protection.py diff --git a/headroom/config.py b/headroom/config.py index 123cc9e93..b1dc607eb 100644 --- a/headroom/config.py +++ b/headroom/config.py @@ -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", } ) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index cb2f70766..fd928d110 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -2116,6 +2116,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 +2198,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 +2237,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 +2249,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 +2257,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 +2266,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 +2682,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 diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 6f1fba323..090408ab5 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -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. @@ -4827,12 +4841,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 +4863,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 diff --git a/tests/test_openai_responses_read_protection.py b/tests/test_openai_responses_read_protection.py new file mode 100644 index 000000000..36ed3691d --- /dev/null +++ b/tests/test_openai_responses_read_protection.py @@ -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" From 6262c28a4834f07bac14143e3347d7b2cdc43937 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Tue, 25 Aug 2026 11:40:54 +0530 Subject: [PATCH 09/18] fix(memory/graph): skip a corrupt row instead of aborting a whole graph scan (#3239) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `SQLiteGraphStore._row_to_entity` and `_row_to_relationship` parse stored text back into objects with no error handling: ```python properties=json.loads(row["properties"]), created_at=datetime.fromisoformat(row["created_at"]), metadata=json.loads(row["metadata"]), ``` These run inside row loops in the multi-row scans — `get_relationships` and `query_subgraph` (both the relationship loop and neighbour-entity expansion). A single unparseable row — from a partial write, a manual edit, or a bad migration — raises `ValueError` (`JSONDecodeError`/bad ISO timestamp) *inside the loop*, aborting the **entire** query and taking unrelated, perfectly good edges/nodes down with it. Reproduction (A→B and A→C both valid; corrupt only A→B's `properties`): ```python # corrupt one row out-of-band con.execute("UPDATE relationships SET properties='{oops' WHERE target_id=?", (b.id,)) # BEFORE: both of these raise JSONDecodeError, even though A->C is fine: await store.get_relationships(a.id) await store.query_subgraph([a.id], max_hops=1, direction=OUTGOING) ``` This is the same "one bad row breaks the whole scan" robustness gap already fixed for the CCR store (`cache/backends/sqlite.py`) and the vector adapter (`memory/adapters/sqlite_vector.py`); the graph adapter was the remaining store with unguarded row parsing. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/memory/adapters/sqlite_graph.py`: - `_row_to_entity` / `_row_to_relationship` now return `... | None`, wrapping construction in `except (ValueError, TypeError, KeyError)` and returning `None` (with a `logger.warning`) on a corrupt row. - Multi-row call sites skip `None`: `get_relationships`, `query_subgraph` (initial entities, relationship loop, neighbour expansion), and the per-user entity listing. The single-row `get_entity` / `get_entity_by_name` already return `Entity | None`, so a corrupt row now reads as "not found" rather than raising. - Added a module `logger`. - `tests/test_sqlite_graph_store.py`: added `test_one_corrupt_row_does_not_abort_a_multi_row_scan` — corrupts one relationship row out-of-band and asserts `get_relationships` returns the one good edge and `query_subgraph` completes with `{A, C}`. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality ### Test Output ```text tests/test_sqlite_graph_store.py::...one_corrupt_row_does_not_abort_a_multi_row_scan -> passes with fix, FAILS without it (verified via git stash) uvx ruff@0.16.2 check headroom/memory/adapters/sqlite_graph.py tests/test_sqlite_graph_store.py -> All checks passed! uvx mypy@1.20.2 headroom/memory/adapters/sqlite_graph.py -> Success: no issues found in 1 source file ``` (Note: this test file has pre-existing, unrelated failures/errors on `main` on Windows — `TestSQLiteGraphStoreMemoryTrackerIntegration` plus temp-file teardown `WinError 32` in the `NamedTemporaryFile`-based fixtures. Verified identical counts before and after this change; my new test uses `tmp_path` and is unaffected.) ## Real Behavior Proof - Environment: Windows 11, Python 3.12.11, project venv, pytest 9.1.1, ruff 0.16.2 and mypy 1.20.2 via uvx. - Exact command / steps: built A→B and A→C edges, corrupted A→B's `properties` to invalid JSON via a direct sqlite connection, then called `get_relationships(A)` and `query_subgraph([A], OUTGOING)`. Before the fix both raised `JSONDecodeError`; after the fix `get_relationships` returns just the A→C edge and `query_subgraph` returns entities `{A, C}` with one relationship, skipping the corrupt row. - Observed result: corrupt rows are skipped (with a warning log); valid rows in the same scan are returned normally. - Not tested: no corruption occurs in normal operation; the corrupt row is produced out-of-band to exercise the guard (matching the real triggers: partial write, manual edit, migration). ## Runtime Rollout Safety - Rollout-managed feature(s): none. This is a SQLite graph-store read path in the memory subsystem, not a rollout-channel-gated runtime feature. - Minimum rollout channel: N/A. - Stable/default behavior changed: no for well-formed data — every valid row parses and is returned exactly as before. Only the previously-crashing corrupt-row case changes, from an aborted query to a skipped row. - Kill switch / disable path: N/A. - Unsafe override required: no. - Qualification impact: none. - Rollback path: revert this PR. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation (N/A: internal behavior) - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` --- headroom/memory/adapters/sqlite_graph.py | 97 ++++++++++++++++-------- tests/test_sqlite_graph_store.py | 43 +++++++++++ 2 files changed, 110 insertions(+), 30 deletions(-) diff --git a/headroom/memory/adapters/sqlite_graph.py b/headroom/memory/adapters/sqlite_graph.py index 2c2ac432c..8c74bde6f 100644 --- a/headroom/memory/adapters/sqlite_graph.py +++ b/headroom/memory/adapters/sqlite_graph.py @@ -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.""" diff --git a/tests/test_sqlite_graph_store.py b/tests/test_sqlite_graph_store.py index 9fb5a90b0..b3db37e62 100644 --- a/tests/test_sqlite_graph_store.py +++ b/tests/test_sqlite_graph_store.py @@ -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.""" From c2fbb4eed0c47973f8baaa88e45be0e51e38e2e2 Mon Sep 17 00:00:00 2001 From: gglucass Date: Wed, 26 Aug 2026 02:58:11 +0200 Subject: [PATCH 10/18] test(agno): follow the metrics dataclass move in agno 3.0.0 (#3260) ## Description agno 3.0.0 (released 2026-08-24) removed the `agno.models.metrics` module; the per-message usage dataclass now lives at `agno.metrics` under the name `MessageMetrics`. The mock fixtures in `tests/test_integrations/agno/test_model.py` import the old path inline, and the `test-agno` CI job installs `wheel[dev,agno]` with an unpinned `agno>=1.0.0`, so it now resolves agno 3.0.0 and fails on every branch - including `main` (see the CI run for #3239's merge commit) and currently-open PRs. This resolves the class once at module level: prefer the pre-3 location, fall back to `MessageMetrics` on agno >= 3. `MessageMetrics` exists under both names in 2.x and the constructor kwargs the fixtures use (`input_tokens`, `output_tokens`, `total_tokens`) are unchanged, so both major versions stay green. Tests-only change; the runtime integration (`headroom/integrations/agno/`) never imported the removed module - the other 76 agno tests already pass on 3.0.0. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - Replace the two inline `from agno.models.metrics import Metrics` imports in the `mock_agno_model` fixture with one module-level compat resolution that tries `agno.models.metrics.Metrics` (agno < 3) and falls back to `agno.metrics.MessageMetrics as Metrics` (agno >= 3). ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) - [ ] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ uv run --frozen --extra dev --extra agno --with agno==3.0.0 pytest tests/test_integrations/agno/ -q ================== 79 passed, 5 skipped, 1 warning in 19.87s =================== $ uv run --frozen --extra dev --extra agno --with agno==2.9.0 pytest tests/test_integrations/agno/ -q ================== 79 passed, 5 skipped, 1 warning in 11.60s =================== $ ruff check tests/test_integrations/agno/test_model.py All checks passed! $ ruff format --check tests/test_integrations/agno/test_model.py 1 file already formatted ``` ## Real Behavior Proof - Environment: macOS 15 (arm64), CPython 3.12, uv-managed venv; branch = upstream/main `6262c28a` + this one test commit. - Exact command / steps: On pristine upstream/main, `uv run --frozen --extra dev --extra agno --with agno==3.0.0 pytest tests/test_integrations/agno/ -q` reproduces the CI failure: 3 failed (`test_response_applies_optimization`, `test_response_stream_applies_optimization`, `test_model_wrapper_real_optimization`), all `ModuleNotFoundError: No module named 'agno.models.metrics'` - the same three failures as the `test-agno` job on current PRs. Applied this commit and re-ran the same command under agno 3.0.0 and agno 2.9.0. - Observed result: 79 passed / 5 skipped under both agno versions; the three fail-before tests pass. - Not tested: agno 1.x (the extra's floor); `mypy headroom` not re-run - no runtime module is touched. ## Runtime Rollout Safety - Rollout-managed feature(s): None - tests only. - Minimum rollout channel: Stable. - Stable/default behavior changed: No. No shipped code changes. - Kill switch / disable path: Not applicable (test-only change). - Unsafe override required: No. - Qualification impact: None. - Rollback path: Revert the single commit. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` - it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Additional Notes - No new regression test: the three existing tests are the regression - they fail on agno 3.0.0 without this change and pass with it. Docs untouched (test-only fix). - An alternative was pinning `agno<3` in the extra; not taken, since the runtime integration works unmodified on 3.0.0 and a pin would block users already on agno 3. Co-authored-by: Claude Fable 5 --- tests/test_integrations/agno/test_model.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/test_integrations/agno/test_model.py b/tests/test_integrations/agno/test_model.py index 26545996c..5a0bf990e 100644 --- a/tests/test_integrations/agno/test_model.py +++ b/tests/test_integrations/agno/test_model.py @@ -19,6 +19,11 @@ try: AGNO_AVAILABLE = True except ImportError: AGNO_AVAILABLE = False +else: + try: # agno < 3: the per-message usage dataclass lived at agno.models.metrics + from agno.models.metrics import Metrics + except ImportError: # agno >= 3 moved it to agno.metrics, renamed MessageMetrics + from agno.metrics import MessageMetrics as Metrics from headroom import HeadroomConfig, HeadroomMode @@ -50,8 +55,6 @@ 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", @@ -73,8 +76,6 @@ 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...", From 36cc800162eae83aceffce705cc4ecba05f6ec02 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Tue, 25 Aug 2026 21:37:12 -0500 Subject: [PATCH 11/18] fix(copilot): honor corporate TLS for token refresh (#3246) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Copilot OAuth/device-auth, user-info, and short-lived token exchange requests used `urllib.request.urlopen` directly, bypassing the corporate CA and X.509 strictness configuration already applied to Headroom's upstream HTTP client. Reuse that TLS resolver for every Copilot GitHub request so token refresh works behind TLS inspection. Closes #3244 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Added a `urlopen` adapter for Headroom's existing corporate TLS resolver. - Routed Copilot device authorization, user-info, and token exchange through it. - Added a regression test proving token exchange receives the configured TLS context. ## Testing - [ ] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text pytest tests/test_copilot_auth.py tests/test_ssl_context.py tests/test_copilot_vscode_completions_routing.py -q 202 passed in 2.45s ruff check . --exclude .codex-worktrees All checks passed! ruff format --check . --exclude .codex-worktrees 1449 files already formatted mypy headroom/copilot_auth.py headroom/proxy/ssl_context.py Success: no issues found in 2 source files ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12.13, OpenSSL 3.5.0; local HTTPS server signed by a private test CA; `REQUESTS_CA_BUNDLE` set to that CA. The exercised request path is the same adapter used by Copilot token exchange. - Exact command / steps: generated a one-day localhost certificate, started an in-process TLS HTTP server, set only `REQUESTS_CA_BUNDLE` to the private CA, and called `headroom.copilot_auth._urlopen(Request(local_https_url), timeout=5)`. - Observed result: `corporate_ca_https_status=200` and `response_body=ok`. - Not tested: a real Cisco/Zscaler interception appliance, macOS, or a live GitHub Copilot Business token (no corporate network/account is available locally). ## Runtime Rollout Safety - Rollout-managed feature(s): None. - Minimum rollout channel: N/A. - Stable/default behavior changed: Only Copilot GitHub requests when a custom CA or `HEADROOM_TLS_STRICT=0` produces an explicit TLS context; default `urlopen` behavior remains unchanged otherwise. - Kill switch / disable path: Unset `SSL_CERT_FILE`, `REQUESTS_CA_BUNDLE`, or `NODE_EXTRA_CA_CERTS` and leave `HEADROOM_TLS_STRICT` enabled. - Unsafe override required: No. - Qualification impact: Restores existing documented corporate TLS settings for Copilot authentication traffic. - Rollback path: Revert this commit. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing relevant unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` — it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Screenshots (if applicable) N/A. ## Additional Notes The issue attributes token exchange to the Rust extension, but current `main` performs it in Python via `urllib`. The direct `urllib` path was the trust-configuration gap. Full-suite execution was also started locally; unrelated environment-dependent failures appeared outside the changed Copilot/TLS scope, while all focused tests pass. --- headroom/copilot_auth.py | 18 +++++++++--- headroom/proxy/ssl_context.py | 13 +++++++++ tests/test_copilot_auth.py | 34 ++++++++++++++++++++++ tests/test_integrations/agno/test_model.py | 30 ++++++++++++------- 4 files changed, 81 insertions(+), 14 deletions(-) diff --git a/headroom/copilot_auth.py b/headroom/copilot_auth.py index 2603577b5..fd9a5504c 100644 --- a/headroom/copilot_auth.py +++ b/headroom/copilot_auth.py @@ -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 {} diff --git a/headroom/proxy/ssl_context.py b/headroom/proxy/ssl_context.py index 7ecba7c31..058c70d1f 100644 --- a/headroom/proxy/ssl_context.py +++ b/headroom/proxy/ssl_context.py @@ -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. diff --git a/tests/test_copilot_auth.py b/tests/test_copilot_auth.py index dd377c987..797b70568 100644 --- a/tests/test_copilot_auth.py +++ b/tests/test_copilot_auth.py @@ -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 diff --git a/tests/test_integrations/agno/test_model.py b/tests/test_integrations/agno/test_model.py index 5a0bf990e..a3fe45c40 100644 --- a/tests/test_integrations/agno/test_model.py +++ b/tests/test_integrations/agno/test_model.py @@ -31,6 +31,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).""" @@ -59,11 +77,7 @@ def mock_agno_model(): 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) @@ -79,11 +93,7 @@ def mock_agno_model(): 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) From 632cb81dbe7400cf77a20d0d395d4cfb0ea245e7 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Tue, 25 Aug 2026 21:40:12 -0500 Subject: [PATCH 12/18] fix(learn): surface Codex analysis failures (#3016) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `headroom learn` could invoke Codex CLI from a non-Git working directory without Codex’s required bypass flag. The resulting backend error was then swallowed by the analyzer and rendered as “No actionable patterns found” with exit code 0. This fixes both coupled defects so Codex can run from discovered project locations and genuine analysis failures remain visible and machine-detectable. Closes #3008 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Added `--skip-git-repo-check` to the Codex CLI analysis backend command. - Added an explicit `analysis_error` result field instead of conflating backend failure with an empty recommendation set. - Kept multi-project analysis best-effort, while returning exit code 1 after any project analysis fails. - Prevented failed analysis from printing a misleading no-pattern success message. - Added analyzer and CLI regression coverage for the command and failure-propagation contracts. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text uv run pytest -q tests/test_learn/test_analyzer.py tests/test_cli_learn.py 102 passed in 2.34s uv run pytest -q tests/test_learn tests/test_cli_learn.py 257 passed, 7 skipped in 3.11s uv run mypy headroom Success: no issues found in 520 source files uv run ruff check All checks passed! uv run ruff format --check 5 files already formatted uv run pytest tests scripts/tests --splits 4 --group N --tb=short -q shard 1: 2766 passed, 140 skipped in 174.08s shard 2: 2699 passed, 207 skipped in 60.00s shard 3: 2822 passed, 84 skipped in 76.10s shard 4: 2734 passed, 172 skipped in 80.29s ``` ## Real Behavior Proof - Environment: macOS arm64, Python 3.13, Codex CLI 0.147.0-compatible command surface, current `main` including #2996. - Exact command / steps: verified `codex exec --help`; exercised `_call_cli_llm` with a captured subprocess command; invoked the Click command with a simulated Codex nonzero backend result. - Observed result: the subprocess command is `codex exec --skip-git-repo-check`; backend failure text is printed as `Analysis failed`, the misleading no-pattern message is absent, and the CLI exits 1. - Not tested: live paid Codex analysis against production account credentials; subprocess and CLI behavior are covered deterministically. ## Runtime Rollout Safety - Rollout-managed feature(s): none; this is CLI-only failure handling. - Minimum rollout channel: normal patch release after exact-head CI is entirely green. - Stable/default behavior changed: failed LLM analysis now exits nonzero instead of reporting success; successful and genuinely empty analyses are unchanged. - Kill switch / disable path: select another backend with `HEADROOM_LEARN_CLI` or `--model` if Codex CLI is unavailable. - Unsafe override required: none. - Qualification impact: all four Python CI shards, static checks, security checks, and command-level regression tests must pass. - Rollback path: fix forward through a human-reviewed corrective PR; no persisted data or migration is involved. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project’s style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation — inline result-contract documentation; no separate user guide change is required - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` — it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Screenshots (if applicable) Not applicable; command-line backend and exit semantics only. ## Additional Notes Human review only. No merge or auto-merge is configured. This corrects the root failure and exit semantics without extending any timeout. --- headroom/cli/learn.py | 10 ++++++++++ headroom/learn/analyzer.py | 6 ++++-- headroom/learn/models.py | 1 + tests/test_cli_learn.py | 29 +++++++++++++++++++++++++++++ tests/test_learn/test_analyzer.py | 3 ++- 5 files changed, 46 insertions(+), 3 deletions(-) diff --git a/headroom/cli/learn.py b/headroom/cli/learn.py index b17a339ec..6ea29f2ad 100644 --- a/headroom/cli/learn.py +++ b/headroom/cli/learn.py @@ -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. diff --git a/headroom/learn/analyzer.py b/headroom/learn/analyzer.py index 670b5c2b6..5fdb44a0e 100644 --- a/headroom/learn/analyzer.py +++ b/headroom/learn/analyzer.py @@ -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 diff --git a/headroom/learn/models.py b/headroom/learn/models.py index 5af68211e..6852f38f8 100644 --- a/headroom/learn/models.py +++ b/headroom/learn/models.py @@ -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: diff --git a/tests/test_cli_learn.py b/tests/test_cli_learn.py index 0142e8fc9..cdb92e123 100644 --- a/tests/test_cli_learn.py +++ b/tests/test_cli_learn.py @@ -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: diff --git a/tests/test_learn/test_analyzer.py b/tests/test_learn/test_analyzer.py index e271d2226..69308f4e3 100644 --- a/tests/test_learn/test_analyzer.py +++ b/tests/test_learn/test_analyzer.py @@ -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): From 997a47992c490e7a13db34ef39eccf3dfb006488 Mon Sep 17 00:00:00 2001 From: JD Davis Date: Tue, 25 Aug 2026 21:46:31 -0500 Subject: [PATCH 13/18] fix(copilot): preserve native enterprise model routing (#2998) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description GitHub Copilot Enterprise/Business users without a BYOK provider key were routed through Copilot CLI's single-model provider override. Native model aliases and runtime `/model` switches were therefore forwarded literally to the override and rejected with `400 model not supported`. This change routes implicit GitHub OAuth through Copilot's native API surface while retaining explicit subscription and provider-key behavior. Closes #1910 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Added explicit `--native` routing and made it automatic for implicit GitHub OAuth without BYOK. - Clears every Copilot BYOK variable before native launch. - Routes both OpenAI and Anthropic protocol targets through the resolved tenant Copilot host. - Preserves Enterprise/Business native aliases and runtime model switching. - Rejects BYOK-only options when native routing is selected. - Refuses known Copilot bundles that do not reference `COPILOT_API_URL`, avoiding silent proxy bypass. - Preserves explicit `--subscription` and provider-key BYOK semantics. - Added coverage for unreadable and unverifiable Copilot CLI bundles. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text 884 passed, 4 skipped in 103.11s ruff check .: All checks passed ruff format --check .: 1412 files already formatted mypy headroom/providers/copilot/wrap.py headroom/cli/wrap.py: Success: no issues found in 2 source files ``` Exact-head CI is entirely green on `0aca48c096485cb7825aa2239d27db7783962f9d`. ## Real Behavior Proof - Environment: macOS arm64/Python 3.13 locally; GitHub-hosted macOS and Ubuntu native-wrap jobs. - Exact command / steps: invoke `headroom wrap copilot` with implicit OAuth and an Enterprise model alias; inspect the captured child/proxy environment and resolved target URLs; exercise explicit native conflicts and bundle-support probes. - Observed result: native launch uses `COPILOT_API_URL`, clears all BYOK state, and points both protocol targets at the tenant host. Native-wrap jobs are green on macOS and Ubuntu for the refreshed head. - Not tested: live request against a real Enterprise tenant; the repository has no organization Enterprise credential available to CI. ## Runtime Rollout Safety - Rollout-managed feature(s): implicit native Copilot routing for GitHub OAuth sessions without BYOK. - Minimum rollout channel: normal patch release. - Stable/default behavior changed: implicit OAuth now uses native routing; explicit subscription and BYOK paths are unchanged. - Kill switch / disable path: use an explicit supported provider-key BYOK configuration; native mode also fails closed when CLI support is known absent. - Unsafe override required: none. - Qualification impact: native-wrap macOS/Ubuntu, Docker wrapper, full Python matrix, and Copilot focused suites must pass. - Rollback path: human revert of this PR restores the fixed-wire OAuth behavior; no configuration migration is persisted. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation — CLI help and inline routing documentation; no separate guide required - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` — it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Screenshots (if applicable) Not applicable; CLI routing change. ## Additional Notes Human review only. No merge or auto-merge is configured. Refreshed from main after #2996; the MCP cap `mcp>=1.28.1,<2.0.0` is preserved. --- headroom/cli/wrap.py | 191 +++++++++++++++++++--------- headroom/providers/copilot/wrap.py | 67 ++++++++++ tests/test_cli/test_wrap_copilot.py | 25 ++-- tests/test_copilot_native_mode.py | 148 +++++++++++++++++++++ 4 files changed, 362 insertions(+), 69 deletions(-) create mode 100644 tests/test_copilot_native_mode.py diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index cc87ba6e0..966d150ac 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -32,7 +32,7 @@ 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 @@ -3943,6 +3943,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, @@ -4646,6 +4660,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, @@ -4682,6 +4697,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, @@ -5607,6 +5623,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( @@ -5619,6 +5643,7 @@ def copilot( wire_api: str | None, subscription: bool, memory: bool, + native: bool, verbose: bool, copilot_args: tuple[str, ...], ) -> None: @@ -5653,6 +5678,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) @@ -5663,6 +5689,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"): @@ -5690,12 +5727,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 @@ -5708,7 +5755,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 @@ -5716,7 +5791,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( @@ -5725,57 +5800,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/ 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/ 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, @@ -5798,7 +5874,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) @@ -5835,6 +5911,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, diff --git a/headroom/providers/copilot/wrap.py b/headroom/providers/copilot/wrap.py index db1610a3b..e73ec3110 100644 --- a/headroom/providers/copilot/wrap.py +++ b/headroom/providers/copilot/wrap.py @@ -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, diff --git a/tests/test_cli/test_wrap_copilot.py b/tests/test_cli/test_wrap_copilot.py index 98f261f09..9772b3bdb 100644 --- a/tests/test_cli/test_wrap_copilot.py +++ b/tests/test_cli/test_wrap_copilot.py @@ -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 diff --git a/tests/test_copilot_native_mode.py b/tests/test_copilot_native_mode.py new file mode 100644 index 000000000..72a7d7f77 --- /dev/null +++ b/tests/test_copilot_native_mode.py @@ -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 From f4119c3bc047c0fc6077f273f266fdc8a2c9750f Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Wed, 26 Aug 2026 14:44:20 +0530 Subject: [PATCH 14/18] test(agno): drop dead module-level Metrics import (#3269) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Main's lint gate is red: `tests/test_integrations/agno/test_model.py:26` fails ruff F401 (`MessageMetrics as Metrics` imported but unused) after recent changes left the module-level import dead — `_response_usage()` already resolves the metrics dataclass locally for both Agno 2.x and 3.x layouts. This deletes the dead try/except block. Every open PR inherits this failure through its merge ref (it blocked #3261's lint check), so this unblocks the queue. - `ruff check` + `ruff format --check`: clean - Test file behavior unchanged (the block was unused) 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01EWKCmcH47hvvoQ35wftXhE Co-authored-by: Claude Fable 5 --- tests/test_integrations/agno/test_model.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/tests/test_integrations/agno/test_model.py b/tests/test_integrations/agno/test_model.py index a3fe45c40..d1aac8982 100644 --- a/tests/test_integrations/agno/test_model.py +++ b/tests/test_integrations/agno/test_model.py @@ -19,11 +19,6 @@ try: AGNO_AVAILABLE = True except ImportError: AGNO_AVAILABLE = False -else: - try: # agno < 3: the per-message usage dataclass lived at agno.models.metrics - from agno.models.metrics import Metrics - except ImportError: # agno >= 3 moved it to agno.metrics, renamed MessageMetrics - from agno.metrics import MessageMetrics as Metrics from headroom import HeadroomConfig, HeadroomMode From 826b600c9b5ac80dae5930e1b7aa376a463d28d6 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Wed, 26 Aug 2026 15:04:43 +0530 Subject: [PATCH 15/18] feat(proxy): self-limiting session state for the compression-cache registry (#3261) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem The per-session `CompressionCache` registry (the map that replays previously-compressed messages byte-identically so the provider prefix cache stays warm) had no lifetime management: - Idle/dead sessions lived forever until the hardcoded 500-session cap was hit. - At capacity, eviction dropped the oldest-**created** quarter — which could wipe the busiest long-lived session (busting every one of its prefixes at once) while dead sessions survived. - Neither the cap nor any TTL was tunable, which blocks gateway deployments (e.g. Kong sidecar/pool) fanning many concurrent sessions into one process. ## Changes - **Idle-TTL sweep**: sessions idle longer than `HEADROOM_COMPRESSION_CACHE_TTL_SECONDS` (default 3900s) are evicted by a lazy sweep, at most once per 60s, piggybacked on `_get_compression_cache` — same pattern as `PrefixCacheTrackerRegistry._maybe_cleanup`, no background task. `last_seen` refreshes on **every** access, so an active session never expires. - **LRU capacity eviction**: the registry is now an access-ordered `OrderedDict`; capacity pressure sheds the *idlest* quarter, never a busy session. - **Tunable cap**: `HEADROOM_COMPRESSION_CACHE_MAX_SESSIONS` (default 500, floor 1). ## Why 3900s Eviction is bust-free only once the provider's own prompt cache has lapsed. Providers don't expose their cache TTLs, and the risk is one-sided (late eviction costs a few MB; early eviction *causes* the bust this state exists to prevent), so the default is the upper bound of documented lifetimes across providers — Anthropic's 1h extended breakpoint, OpenAI's "up to an hour off-peak", Gemini's 60-min default — plus 5m grace. A parse-time floor of 600s keeps the TTL from ever dropping below the prefix tracker's session TTL: after the tracker expires, the byte-identical swap is the only remaining protection for a still-live provider prefix. Read-hit signals are untouched: they govern the freeze boundary, never eviction — `read_hits == 0` usually means cold start or TTL lapse, where the map was just (re)written into the provider cache and deleting it would guarantee a second bust. ## Behavior impact - Steady state (any session active within the TTL): zero change — same instances, same bytes, same freeze behavior. - A session returning after >65 min idle now finds its map evicted — but every provider had already forgotten its prefix by then, so that turn was paying the cache-write price regardless (fail-open, no failed requests). - Capacity eviction now protects busy sessions instead of punishing them. ## Testing - New `tests/test_compression_cache_registry.py`: LRU-not-FIFO capacity eviction, small-cap edge case, TTL sweep eviction, access-refreshes-clock, sweep rate limiting. - 386 tests pass across compression-cache, cache-stability (Anthropic + OpenAI), prefix-overlay, cold-start, cache-mode, and Bedrock-tracker suites; ruff check/format clean. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01EWKCmcH47hvvoQ35wftXhE --------- Co-authored-by: Claude Fable 5 --- headroom/proxy/handlers/anthropic.py | 18 ++- headroom/proxy/helpers.py | 48 +++++- headroom/proxy/server.py | 95 ++++++++++-- headroom/transforms/content_router.py | 14 +- ...est_anthropic_pre_upstream_backpressure.py | 71 +++++++++ tests/test_compression_cache_registry.py | 134 ++++++++++++++++ tests/test_org_scale_limits.py | 143 ++++++++++++++++++ ...thropic_no_optimize_history_passthrough.py | 49 ++++-- 8 files changed, 540 insertions(+), 32 deletions(-) create mode 100644 tests/test_compression_cache_registry.py create mode 100644 tests/test_org_scale_limits.py diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 9b0c9d502..00c719d2e 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2074,7 +2074,16 @@ class AnthropicHandlerMixin: # 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: @@ -2089,15 +2098,10 @@ class AnthropicHandlerMixin: optimized_messages = _ov optimized_tokens = tokenizer.count_messages(optimized_messages) 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 diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 40efb6571..76b2700d1 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -1245,8 +1245,52 @@ 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). + COMPRESSION_CACHE_TTL_SECONDS = max( + 600.0, + float(os.environ.get("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", "3900")), + ) +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 # --------------------------------------------------------------------------- diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index a392e9fe1..b526493ad 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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,49 @@ 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 + expired = [ + sid + for sid, seen in self._compression_cache_last_seen.items() + if now - seen > COMPRESSION_CACHE_TTL_SECONDS + ] + 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 _get_compression_cache(self, session_id: str) -> CompressionCache: """Get or create a CompressionCache for a session. @@ -1588,27 +1641,43 @@ 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. 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] + evict_count = min( + max(1, MAX_COMPRESSION_CACHE_SESSIONS // 4), + len(self._compression_caches), + ) + for _ in range(evict_count): + sid, _evicted = self._compression_caches.popitem(last=False) + self._compression_cache_last_seen.pop(sid, None) logger.info( - "Evicted %d compression caches (exceeded %d max sessions)", - len(oldest_keys), + "Evicted %d least-recently-used compression caches " + "(exceeded %d max sessions)", + evict_count, 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. diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 090408ab5..fa445edc9 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -1947,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) diff --git a/tests/test_anthropic_pre_upstream_backpressure.py b/tests/test_anthropic_pre_upstream_backpressure.py index be51e48a0..383157bd5 100644 --- a/tests/test_anthropic_pre_upstream_backpressure.py +++ b/tests/test_anthropic_pre_upstream_backpressure.py @@ -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) diff --git a/tests/test_compression_cache_registry.py b/tests/test_compression_cache_registry.py new file mode 100644 index 000000000..f510ff8d8 --- /dev/null +++ b/tests/test_compression_cache_registry.py @@ -0,0 +1,134 @@ +"""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 diff --git a/tests/test_org_scale_limits.py b/tests/test_org_scale_limits.py new file mode 100644 index 000000000..716b1cc5e --- /dev/null +++ b/tests/test_org_scale_limits.py @@ -0,0 +1,143 @@ +"""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 + + +# --------------------------------------------------------------------------- # +# 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) diff --git a/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py b/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py index c79071b3b..1b79fdf40 100644 --- a/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py +++ b/tests/test_proxy/test_anthropic_no_optimize_history_passthrough.py @@ -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(): From 4fa88026d9ff64092feaa750789a28be3ed0ace8 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Wed, 26 Aug 2026 15:34:23 +0530 Subject: [PATCH 16/18] feat(compress): session-aware /v1/compress (sidecar mode) + /v1/usage relay (#3270) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > Replaces #3262 (same changeset, squashed to one conventional commit — the stacked branch's history could not pass commitlint after #3261's squash-merge broke ancestry, and force-pushing the original branch was not permitted). All review findings from the two max-effort reviews are already incorporated; #3261 is merged. ## Why Gateways that own routing (e.g. Kong as the upstream caller) can't use Headroom's proxy path, and the stateless `/v1/compress` pushes all byte-replay bookkeeping onto the caller. This PR moves that state into the endpoint: **the caller sends the raw conversation + a session id every turn, forwards the returned bytes verbatim, and gets a byte-identical prefix — provider prompt cache preserved, no forwarding through Headroom.** ## Design - **Session pre-work** mirrors the proxy's Zone 1: content-addressed swap of previously-computed compressed bytes, then freeze the **entire locally-replayable prefix** (`compute_frozen_count`). - **Freeze posture deliberately differs from the proxy's `min(tracker, cache)`**: in sidecar mode, whatever this endpoint previously returned *is* the provider's cache contract — recompressing an already-returned message (even into a smaller form) is a bust. Over-freezing only forgoes tail compression; it can never bust. (A test caught exactly this: recompression drift produced a smaller form, and `overlay_cached_prefix`'s non-inflation guard then couldn't repair it.) - **`PrefixCacheTracker.record_returned()`** — the sidecar equivalent of "last forwarded", captured at return time because whatever is returned is what the caller forwards. - **`POST /v1/usage`** (same loopback exposure policy): the caller relays the provider's usage block; `update_from_response` makes freeze decisions provider-confirmed. Optional — skipping it degrades freeze precision, never correctness. - Sessions are NUL-namespaced (`compress\x00`, unspoofable via HTTP headers); the registry's TTL/LRU lifecycle from #3261 applies automatically. No session id ⇒ stateless contract byte-for-byte unchanged. ## Hardening (from two max-effort code reviews, all applied) - `compress_user_messages` + session_id → 400 (user-message rewrites are not content-addressed → guaranteed later bust). - Session-mode timeout / lock-busy → 503 `compression_timeout` / `session_busy` with retry semantics, instead of failing open with raw bytes (desync bust). - Header-based session ids gated behind `HEADROOM_COMPRESS_SESSION_FROM_HEADER` (default off). - `/v1/usage` validation: unknown/expired session → 404; both cache fields absent → 400; single-present-zero → `{"applied": false, "reason": "no_cache_signal"}` (never wipes freeze state). - Warm-turn savings recomputed from the raw payload (honest `tokens_saved`), all CPU work in the executor under a per-session turn lock. ## Caller contract (Kong) 1. Send raw history + `config.session_id` (or `x-headroom-session-id` with the env gate on) every turn. 2. Forward the returned `messages` to the provider **verbatim**. 3. Optionally relay the provider's usage block to `/v1/usage`. ## Testing 20 cases in `tests/test_compress_session_mode.py`: stateless regression + no state leakage, invalid-id rejection, 2-turn and 3-turn whole-prefix byte-stability, tracker-loss stability, spoof-resistance, header gating, lock-busy 503s, usage validation and no-signal handling, unknown/expired-session 404, TTL-eviction fail-open, explicit `frozen_message_count` precedence. Plus the full local suite green (11k+ tests). ## Phase 2 (follow-up) #3263 migrates the proxy request path onto this same session engine so both modes share one compression/state codepath. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01EWKCmcH47hvvoQ35wftXhE Co-authored-by: Claude Fable 5 --- headroom/cache/compression_cache.py | 6 + headroom/cache/prefix_tracker.py | 30 ++ headroom/proxy/handlers/openai.py | 362 ++++++++++++++++++++-- headroom/proxy/helpers.py | 10 +- headroom/proxy/server.py | 7 + tests/test_compress_session_mode.py | 427 ++++++++++++++++++++++++++ tests/test_org_scale_limits.py | 22 ++ tests/test_proxy_compress_endpoint.py | 18 +- 8 files changed, 845 insertions(+), 37 deletions(-) create mode 100644 tests/test_compress_session_mode.py diff --git a/headroom/cache/compression_cache.py b/headroom/cache/compression_cache.py index 00abf07b0..dcbe0c0ee 100644 --- a/headroom/cache/compression_cache.py +++ b/headroom/cache/compression_cache.py @@ -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 diff --git a/headroom/cache/prefix_tracker.py b/headroom/cache/prefix_tracker.py index 200adddb8..9b145c333 100644 --- a/headroom/cache/prefix_tracker.py +++ b/headroom/cache/prefix_tracker.py @@ -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,16 @@ 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 tracker for ``session_id`` if one exists, 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`. + """ + return self._trackers.get(session_id) + 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() diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index fd928d110..18b6bfb33 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -9591,6 +9591,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) @@ -9655,6 +9658,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. @@ -9701,23 +9762,121 @@ 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. + from headroom.cache.prefix_tracker import overlay_cached_prefix + + with comp_cache.session_turn_lock: + prev_original = session_tracker.get_last_original_messages() + prev_returned = session_tracker.get_last_forwarded_messages() + derived_frozen = comp_cache.compute_frozen_count(messages) + session_frozen = max(derived_frozen, frozen_message_count or 0) + comp_cache.mark_stable_from_messages(messages, session_frozen) + pipeline_input = comp_cache.apply_cached(messages) + pipeline_kwargs["frozen_message_count"] = session_frozen + result = pipeline.apply(messages=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. + final = overlay_cached_prefix( + result.messages, messages, prev_original, prev_returned + ) + replayed = final != result.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: # nosec B110 - fall back to pipeline counts + 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": replayed, + } + return result, final, raw_tokens_before, final_tokens_after, info + + # 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( @@ -9749,28 +9908,56 @@ 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, + ) + 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( @@ -9820,6 +10007,119 @@ 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 session is + # answered without leaving a footprint, and the session keeps the + # provider its compress call inferred rather than a default from here. + tracker = self.session_tracker_store.peek(f"compress\x00{session_id}") + last_returned = tracker.get_last_forwarded_messages() if tracker is not None else [] + if not last_returned: + # 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 — it only means this telemetry 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." + ), + } + }, + ) + + 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 JSONResponse( + { + "session_id": session_id, + "frozen_message_count": tracker.get_frozen_message_count(), + } + ) + async def _maybe_compress_passthrough_responses( self, body: bytes, *, client: str | None = None ) -> bytes: diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 76b2700d1..8d2a2f43c 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -1271,10 +1271,12 @@ 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). - COMPRESSION_CACHE_TTL_SECONDS = max( - 600.0, - float(os.environ.get("HEADROOM_COMPRESSION_CACHE_TTL_SECONDS", "3900")), - ) + # 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 diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index b526493ad..dac2fe6e5 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -5252,6 +5252,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 diff --git a/tests/test_compress_session_mode.py b/tests/test_compress_session_mode.py new file mode 100644 index 000000000..01380fa32 --- /dev/null +++ b/tests/test_compress_session_mode.py @@ -0,0 +1,427 @@ +"""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 diff --git a/tests/test_org_scale_limits.py b/tests/test_org_scale_limits.py index 716b1cc5e..d9ceb4879 100644 --- a/tests/test_org_scale_limits.py +++ b/tests/test_org_scale_limits.py @@ -75,6 +75,28 @@ def test_compression_cache_entry_cap_env_parsing(monkeypatch) -> None: 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. # # --------------------------------------------------------------------------- # diff --git a/tests/test_proxy_compress_endpoint.py b/tests/test_proxy_compress_endpoint.py index c6434e4ed..9ab8084f5 100644 --- a/tests/test_proxy_compress_endpoint.py +++ b/tests/test_proxy_compress_endpoint.py @@ -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()) From d12ea501222b92f371b5427595c14ec2466a3f39 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Wed, 26 Aug 2026 15:57:31 +0530 Subject: [PATCH 17/18] feat(proxy): unify proxy and sidecar compression on one session engine (#3271) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > Replaces #3263 (same changeset, squashed to one conventional commit — after the base PRs squash-merged, the stacked branch's commit history could not pass the commitlint gate against main, and force-pushing the original branch was not permitted). #3261 and #3270 (which replaced #3262) are merged; this is the last piece of the stack. ## Goal One brain. The cache-management tier — freeze computation, Zone-1 byte swap, cached-prefix overlay — previously existed twice: inline in the proxy request handlers, and (as of #3270) in the `/v1/compress` sidecar path. This PR extracts it into **`headroom/proxy/session_engine.py`**, invoked by BOTH. Every future cache-management fix lands in both modes by construction. ## Design **`prepare_turn(...)` → `TurnPrep`** — freeze + `mark_stable_from_messages` + `apply_cached`, with two *deliberately different, documented* freeze policies: - `FREEZE_POLICY_CONFIRMED_CLAMP`: `min(tracker_frozen, cache_count)` — never freeze past provider-confirmed (the #327 posture). The Anthropic proxy passes its already-composed tracker/strict-override value, reproducing the previous `min()` byte-for-byte. - `FREEZE_POLICY_REPLAYABLE`: `max(cache_count, explicit)` — freeze everything locally replayable, because whatever was previously returned *is* the provider's cache contract; recompressing it (even "better") busts. **`finalize_turn(...)` → `TurnFinal`** — the byte-identical cached-prefix replay (`overlay_cached_prefix`) + conditional token recount hook. Run as a **strictly behavior-preserving extraction**: the bar was every pre-existing test passing *unmodified*, and it held. ## What migrated | Path | Status | |---|---| | `/v1/compress` sidecar turn | ✅ engine (REPLAYABLE); lock, executor offload, savings accounting, record_returned unchanged | | `anthropic.py` token-mode pre-block + overlay | ✅ engine (CONFIRMED_CLAMP); background compression, cold-start fast pass, `_cold_recompact_active` skip preserved | | `openai.py` proxy token-mode pre-block + overlay | ✅ engine (REPLAYABLE — formula-identical to the old bare `compute_frozen_count`); the added `mark_stable` call means the freeze now survives entry-level LRU eviction (test-pinned); the router's `_frozen_verdicts` remains the boundary-message protection | | `openai.py` cache-mode branch | ⏸ keeps bare `apply_cached` — cache mode keeps the latest observation mutable by design | Also fixed for BOTH handlers: overlay replay now runs under backpressure (shedding it busted every gated session's prompt cache exactly at peak load), and the inflation guard exempts replayed prefixes. ## Hardening (max-effort review, all applied) `/v1/usage` applies on the executor under the per-session turn lock with a timed acquire (503 `session_busy`); registry eviction skips sessions mid-turn; `peek()` is expiry-aware; silent fallbacks log warnings; RequestOutcome recorded on session 503s. ## Testing - `tests/test_session_engine.py`: 13 direct unit tests — both policies, explicit-pin precedence, REPLAYABLE-without-pin ≡ bare `compute_frozen_count`, overlay fires/doesn't, recount only on replay, freeze-survives-entry-eviction. - Parity bar: full pre-existing suites pass unmodified — cache-stability (Anthropic + OpenAI), overlay, backpressure (incl. replay-under-saturation regression), cold-start fast pass, cache-mode, session-mode byte-stability, compress-API, org-scale, registry. Full local suite: 11k+ green. - ruff check/format clean (CI's ruff 0.16.3). 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01EWKCmcH47hvvoQ35wftXhE Co-authored-by: Claude Fable 5 --- headroom/cache/prefix_tracker.py | 14 +- headroom/proxy/handlers/anthropic.py | 70 +++--- headroom/proxy/handlers/openai.py | 289 ++++++++++++++++++----- headroom/proxy/server.py | 46 +++- headroom/proxy/session_engine.py | 185 +++++++++++++++ tests/test_compress_session_mode.py | 141 +++++++++++ tests/test_compression_cache_registry.py | 51 ++++ tests/test_session_engine.py | 225 ++++++++++++++++++ 8 files changed, 923 insertions(+), 98 deletions(-) create mode 100644 headroom/proxy/session_engine.py create mode 100644 tests/test_session_engine.py diff --git a/headroom/cache/prefix_tracker.py b/headroom/cache/prefix_tracker.py index 9b145c333..ea47417e3 100644 --- a/headroom/cache/prefix_tracker.py +++ b/headroom/cache/prefix_tracker.py @@ -1270,14 +1270,24 @@ class SessionTrackerStore: self._lineage_counter = itertools.count(1) def peek(self, session_id: str) -> PrefixCacheTracker | None: - """Return the tracker for ``session_id`` if one exists, else 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. """ - return self._trackers.get(session_id) + 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.""" diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 00c719d2e..538cc2b6d 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -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,10 +2071,8 @@ 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 @@ -2087,16 +2091,18 @@ class AnthropicHandlerMixin: 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: logger.debug( "[%s] Cached-prefix replay skipped: reason=%s", diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 18b6bfb33..42ac1a8c0 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -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. @@ -3753,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( @@ -3854,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" @@ -9815,24 +9860,47 @@ class OpenAIHandlerMixin: # 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. - from headroom.cache.prefix_tracker import overlay_cached_prefix + # 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, + ) - with comp_cache.session_turn_lock: + # 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() - derived_frozen = comp_cache.compute_frozen_count(messages) - session_frozen = max(derived_frozen, frozen_message_count or 0) - comp_cache.mark_stable_from_messages(messages, session_frozen) - pipeline_input = comp_cache.apply_cached(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=pipeline_input, model=model, **pipeline_kwargs) + 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. - final = overlay_cached_prefix( - result.messages, messages, prev_original, prev_returned - ) - replayed = final != result.messages + 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 @@ -9843,7 +9911,21 @@ class OpenAIHandlerMixin: _tok = get_tokenizer(model_name) raw_tokens_before = _tok.count_messages(messages) final_tokens_after = _tok.count_messages(final) - except Exception: # nosec B110 - fall back to pipeline counts + 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) @@ -9855,9 +9937,11 @@ class OpenAIHandlerMixin: info = { "id": session_id, "frozen_message_count": session_frozen, - "cached_prefix_replayed": replayed, + "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 @@ -9942,6 +10026,32 @@ class OpenAIHandlerMixin: 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={ @@ -10081,45 +10191,112 @@ class OpenAIHandlerMixin: # 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 session is - # answered without leaving a footprint, and the session keeps the - # provider its compress call inferred rather than a default from here. - tracker = self.session_tracker_store.peek(f"compress\x00{session_id}") - last_returned = tracker.get_last_forwarded_messages() if tracker is not None else [] - if not last_returned: - # 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 — it only means this telemetry landed - # nowhere. + # 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( - status_code=404, + { + "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": "unknown_session", + "type": "session_busy", "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." + "A compress turn for this session is in flight; retry the usage relay." ), } }, ) - - tracker.update_from_response( - cache_read_tokens=cache_read, - cache_write_tokens=cache_write, - messages=last_returned, - original_messages=tracker.get_last_original_messages(), - ) + if frozen_count is None: + return self._compress_usage_unknown_session(session_id) return JSONResponse( { "session_id": session_id, - "frozen_message_count": tracker.get_frozen_message_count(), + "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: diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index dac2fe6e5..233a0367c 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -1617,10 +1617,18 @@ class HeadroomProxy( ): 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) @@ -1633,6 +1641,16 @@ class HeadroomProxy( 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. @@ -1657,20 +1675,32 @@ class HeadroomProxy( # 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: evict_count = min( max(1, MAX_COMPRESSION_CACHE_SESSIONS // 4), len(self._compression_caches), ) - for _ in range(evict_count): - sid, _evicted = self._compression_caches.popitem(last=False) + 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) - logger.info( - "Evicted %d least-recently-used compression caches " - "(exceeded %d max sessions)", - evict_count, - MAX_COMPRESSION_CACHE_SESSIONS, - ) + if evictable: + logger.info( + "Evicted %d least-recently-used compression caches " + "(exceeded %d max sessions)", + len(evictable), + MAX_COMPRESSION_CACHE_SESSIONS, + ) cache = CompressionCache(max_entries=COMPRESSION_CACHE_MAX_ENTRIES) self._compression_caches[session_id] = cache diff --git a/headroom/proxy/session_engine.py b/headroom/proxy/session_engine.py new file mode 100644 index 000000000..4864ae08a --- /dev/null +++ b/headroom/proxy/session_engine.py @@ -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) diff --git a/tests/test_compress_session_mode.py b/tests/test_compress_session_mode.py index 01380fa32..cc3a23563 100644 --- a/tests/test_compress_session_mode.py +++ b/tests/test_compress_session_mode.py @@ -425,3 +425,144 @@ def test_explicit_frozen_count_still_wins_when_larger() -> None: 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" diff --git a/tests/test_compression_cache_registry.py b/tests/test_compression_cache_registry.py index f510ff8d8..da3aba4e8 100644 --- a/tests/test_compression_cache_registry.py +++ b/tests/test_compression_cache_registry.py @@ -132,3 +132,54 @@ def test_sweep_is_rate_limited(monkeypatch) -> None: 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() diff --git a/tests/test_session_engine.py b/tests/test_session_engine.py new file mode 100644 index 000000000..a29ae4c10 --- /dev/null +++ b/tests/test_session_engine.py @@ -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 From 1e448b55039b6968caedb8afb810ba70a3abe0c1 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Wed, 26 Aug 2026 22:43:11 +0530 Subject: [PATCH 18/18] fix(providers): route Claude requests to Copilot when the OpenAI target is a Copilot host (#3258) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Through `headroom wrap vscode` / `wrap copilot --subscription`, GitHub Copilot **GPT** models work but **Claude** models fail with `Invalid bearer token` (issue #3247). The logs tell the story: ```text # GPT — works: event=outbound_request path=https://api.githubcopilot.com/chat/completions status=200 # Claude — fails: event=outbound_request path=https://api.anthropic.com/v1/messages status=401 ``` GitHub Copilot serves **both** surfaces from the same host: its OpenAI surface (`/chat/completions`, `/responses`) and its Anthropic surface for Claude models (`/v1/messages`) — `build_copilot_upstream_url` already documents and handles this. But `resolve_api_targets` resolves each provider target independently: when the Copilot flow points the **OpenAI** target at a Copilot host (so GPT works), the **Anthropic** target is left at its default `https://api.anthropic.com`. Claude-model requests are therefore forwarded to the real Anthropic API carrying the GitHub Copilot bearer, which Anthropic rejects with `Invalid bearer token`. ## Fix In `resolve_api_targets`, when the resolved OpenAI target is a Copilot upstream host **and no explicit Anthropic target was configured**, default the Anthropic target to that same Copilot host. Claude requests then reach `https://api.githubcopilot.com/v1/messages` — the surface that serves them, where the Copilot bearer is valid. An explicit `ANTHROPIC_TARGET_API_URL` always wins (only a `None` override is filled in), and non-Copilot OpenAI targets are untouched, so direct-Anthropic setups are unaffected. Reproduction: ```python resolve_api_targets(ProviderApiOverrides(openai="https://api.githubcopilot.com", anthropic=None, ...)) # BEFORE: targets.anthropic == "https://api.anthropic.com" -> Copilot bearer 401s there # AFTER: targets.anthropic == "https://api.githubcopilot.com" ``` ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/providers/registry.py`: `resolve_api_targets` now fills a `None` Anthropic override with the OpenAI target when that target is a Copilot host (`is_copilot_upstream_url`). Explicit overrides and non-Copilot targets are unchanged. - `tests/test_provider_registry.py`: added three tests — Copilot OpenAI target routes Anthropic to Copilot; an explicit Anthropic override wins; a non-Copilot OpenAI target leaves the Anthropic default alone. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality ### Test Output ```text tests/test_provider_registry.py tests/test_provider_registry_extended.py tests/test_banner_upstream_targets.py -> 37 passed in 12.11s (the new Copilot test FAILS on pre-fix code — verified via git stash) uvx ruff@0.16.2 check headroom/providers/registry.py tests/test_provider_registry.py -> All checks passed! uvx mypy@1.20.2 headroom/providers/registry.py -> Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12.11, project venv, pytest 9.1.1, ruff 0.16.2 and mypy 1.20.2 via uvx. - Exact command / steps: `resolve_api_targets` with `openai="https://api.githubcopilot.com"` (and the `api.business.githubcopilot.com` variant) and `anthropic=None` returned `anthropic="https://api.anthropic.com"` before the fix and the Copilot host after; an explicit `anthropic="https://api.anthropic.com"` is preserved; `openai="https://api.openai.com"` leaves `anthropic` at the default. - Observed result: Claude-model requests now resolve to the Copilot host that serves them; OpenAI/direct-Anthropic behavior is unchanged. - Not tested: no live macOS/VS Code Copilot round trip (environment-specific); the target-resolution seam that decides the upstream host is exercised directly. `is_copilot_upstream_url` already recognizes the github.com Copilot hosts (verified). ## Runtime Rollout Safety - Rollout-managed feature(s): none. This is upstream target resolution in the provider registry, not a rollout-channel-gated runtime feature. - Minimum rollout channel: N/A. - Stable/default behavior changed: only the broken case changes — a Copilot OpenAI target with no Anthropic override now sends Claude to Copilot instead of 401ing against api.anthropic.com. Explicit Anthropic targets and non-Copilot OpenAI targets are byte-for-byte unchanged. - Kill switch / disable path: set `ANTHROPIC_TARGET_API_URL` explicitly to opt out of the default. - Unsafe override required: no. - Qualification impact: none for non-Copilot deployments. - Rollback path: revert this PR. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation (N/A: internal behavior) - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` ## Additional Notes Fixes the routing/auth mismatch at the resolution layer so it applies uniformly across the Copilot config paths (`wrap vscode`, `wrap copilot --subscription`) that set the OpenAI target to a Copilot host. If a specific deploy sets neither target to a Copilot host (relying solely on path-based passthrough routing for OpenAI), configuring `ANTHROPIC_TARGET_API_URL` to the Copilot host remains the explicit escape hatch. --- headroom/providers/registry.py | 22 +++++++++++++-- tests/test_provider_registry.py | 49 +++++++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 2 deletions(-) diff --git a/headroom/providers/registry.py b/headroom/providers/registry.py index 04c212a4d..a8b816f3b 100644 --- a/headroom/providers/registry.py +++ b/headroom/providers/registry.py @@ -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), diff --git a/tests/test_provider_registry.py b/tests/test_provider_registry.py index 0b0ff5b04..07310cd41 100644 --- a/tests/test_provider_registry.py +++ b/tests/test_provider_registry.py @@ -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",