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

Merge d73278a4e4 into 7c0b886004

Browse source
This commit is contained in:
Dennis Alexis Valin Dittrich 2026-08-27 18:48:27 +09:00 • committed by GitHub
commit 490a17a1b0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
65 changed files with 16905 additions and 159 deletions
Download patch file Download diff file Expand all files Collapse all files

35
.github/workflows/ci.yml vendored
View file

@ -565,6 +565,41 @@ jobs:
- name: Run native installer wrapper tests
run: pytest tests/test_install/test_native_installers.py -q
agy-windows:
# agy's TLS-MITM transport is selective-host MITM with a process-scoped CA;
# it must run on Windows (a supported developer platform). The main test
# shards run on Linux, so this lane is the only Windows coverage for the
# agy CA / dispatch / terminator / retrieve slice and the wrap command.
needs: changes
if: needs.changes.outputs.code == 'true'
runs-on: windows-latest
timeout-minutes: 30
steps:
- uses: actions/checkout@v6
- uses: actions/setup-python@v6
with:
python-version: "3.12"
- uses: dtolnay/rust-toolchain@1.96.0
- uses: astral-sh/setup-uv@v5
- name: Run agy suite on Windows
shell: bash
# This lane has no prefetch-model dependency and no HuggingFace cache
# restore, unlike the Linux `test` shards. Without the offline flags
# those shards already set, building the proxy app here lazily fetches
# the embedding model from the Hub mid-test: the run log showed
# `huggingface_hub/file_download.py ... xet_get(` attributed to
# test_dispatch_server_tls_and_route, whose 10s guard then tripped
# while the same test passed on Linux. None of the agy tests need the
# model — verified with a cold HF_HOME plus these flags, 38 passed.
env:
HF_HUB_OFFLINE: "1"
TRANSFORMERS_OFFLINE: "1"
run: |
uv run --extra proxy --with pytest --with pytest-asyncio python -m pytest \
tests/test_agy_ca.py tests/test_agy_dispatch.py tests/test_agy_terminator.py \
tests/test_agy_retrieve.py tests/test_agy_stats.py tests/test_agy_registrar.py \
tests/test_proxy_google_cloudcode_route_aliases.py tests/test_wrap_agy.py -q
macos-native-wrapper:
needs: changes
if: needs.changes.outputs.native == 'true'

6
.gitignore vendored
View file

@ -270,3 +270,9 @@ uv.lock
.tokensave
.codebase-memory/
# Embedded-database sidecars written into the repo root by local issue-tracker
# tooling (Dolt-backed). Regenerated on demand; never part of a change set.
/embeddeddolt/
/backup/
/.local_version

151
README.md
View file

@ -242,6 +242,7 @@ shows an **Output Tokens Saved** card next to input compression, labelled
| OpenHands | ✅ | starts proxy + launches |
| Mistral Vibe | ✅ | starts proxy + launches |
| Oh My Pi | ✅ | injects config · starts proxy + launches |
| agy | ✅ | MITM/CA-trust transport (no base-URL override); see [Using headroom with agy](#using-headroom-with-agy) |
| Cortex Code | Library only | 60–65% savings (library mode; no `wrap`) |
| Kimi CLI | ✅ | OAuth bearer forwarded — log in once |
| ZCode | ✅ | starts proxy and prints base URLs for ZCode settings |
@ -331,6 +332,156 @@ See the
for verification, configuration paths, custom profiles, remote development, and
troubleshooting.
### Using headroom with agy
`agy` (Google Antigravity CLI) has no base-URL override, so Headroom wraps it via a selective
TLS-MITM transport rather than a base-URL redirect.
#### Quickstart
```bash
headroom wrap agy # start with MITM transport
headroom wrap agy -- --help # pass args to agy
headroom wrap agy -- --print "ping"
```
#### TLS interception disclosure
Headroom intercepts TLS only for the Cloud Code backend hosts
`daily-cloudcode-pa.googleapis.com` and `cloudcode-pa.googleapis.com`.
All other CONNECT tunnels are byte-spliced unchanged — no certificate, no inspection.
At launch, `headroom wrap agy` prints (one line per intercepted host):
```
┌─ TLS INTERCEPTION DISCLOSURE ──────────────────
│ Headroom terminates TLS for: cloudcode-pa.googleapis.com
│ Headroom terminates TLS for: daily-cloudcode-pa.googleapis.com
│ A process-local CA mints leaf certificates for those hosts.
│ This CA is NEVER added to the OS trust store.
│ Compression and context injection are applied on the decrypted stream.
│
│ To opt out of interception: headroom wrap agy --no-intercept
│ To revert all changes: headroom unwrap agy
└────────────────────────────────────────────────
```
The process-local CA is stored under `~/.headroom/ca/` (directory mode `0700`, key mode `0600`).
It is injected into the child `agy` process only via three environment variables:
```
SSL_CERT_FILE=~/.headroom/combined-ca-bundle.pem
CACERT_PATH=~/.headroom/combined-ca-bundle.pem
NODE_EXTRA_CA_CERTS=~/.headroom/combined-ca-bundle.pem
```
The combined bundle is the system CA bundle plus the Headroom root CA certificate.
No OS trust store is modified.
#### Compression value and mechanism
The compression value — reduced token count on requests to the Cloud Code backend — is
identical to other supported agents. The mechanism differs: instead of a base-URL redirect,
Headroom uses an in-process HTTP CONNECT terminator that splices the accepted connection to a
loopback hypercorn server, which terminates TLS, negotiates HTTP/2 and SSE natively, and
routes decrypted requests through the existing `handle_google_cloudcode_stream` handler. The
request is re-originated to the host `agy` opened the tunnel to, not to a fixed default.
A measured session: `agy` reading a 753 KB log file through the transport compressed 21 tool
results from 23,392 to 567 tokens, with the model's answer unchanged and no fail-open requests.
**Compression of tool output needs the retrieve MCP.** Gemini carries tool results as
`functionResponse` parts, which Headroom compresses into `[Retrieve more: hash=…]` markers.
Those markers are only recoverable when the `headroom` MCP server is registered, so
`--no-mcp` (and any run where registration fails) downgrades tool-output compression to a
lossless mode that saves close to nothing rather than shipping a marker nothing can resolve.
`headroom wrap agy` prints a warning naming the cause whenever that downgrade happens.
Auth headers (`Authorization`, `x-goog-api-key`) are visible to the Headroom process after
TLS termination. They pass through the existing `redact_for_wire_debug` redactor and are not
persisted in the semantic cache.
`agy` is classified as `Subscription` auth mode (UA prefix `antigravity/`). See
[docs/auth-modes.md](docs/auth-modes.md) for the full auth-mode policy table.
#### Opt-out: `--no-intercept`
```bash
headroom wrap agy --no-intercept
```
Launches `agy` with no environment modifications and no TLS interception.
Headroom does not compress traffic in this mode.
#### Reverting: `headroom unwrap agy`
```bash
headroom unwrap agy
```
Removes all Headroom-added persistent configuration: any leftover `GEMINI.md` block from an
older install (markers `<!-- headroom:agy-instructions -->` /
`<!-- /headroom:agy-instructions -->`; current versions write none),
the Headroom MCP retrieve-tool entry from `~/.gemini/config/mcp_config.json` (agy 1.1.x
read-path, shared with the Antigravity IDE; if registered via `headroom mcp install`), and
any Headroom-installed Serena MCP entry. User-managed and IDE `mcp_config.json` entries are preserved.
#### Enterprise / Zero-Trust environments
If `HTTPS_PROXY` is set in the parent environment before running `headroom wrap agy`,
non-allowlisted CONNECT tunnels are chained through that corporate proxy. The terminator
reads the parent's `HTTPS_PROXY` directly; only the child `agy` process receives the
overridden value pointing at the Headroom terminator. Corporate CA certificates (from
`SSL_CERT_FILE` or `NODE_EXTRA_CA_CERTS`) are merged into the combined bundle so the
real internet continues to validate. Only PEM objects with `basicConstraints CA:TRUE`
are merged. The launch banner redacts proxy credentials before printing: it shows only
`scheme://host:port`, never the `user:pass@` userinfo.
If `HTTPS_PROXY` carries `user:pass@` userinfo, it is percent-decoded and sent to the
corporate proxy as an HTTP Basic `Proxy-Authorization` header on every chained CONNECT —
only when the proxy scheme is `http`/`https` (never to a `socks5://` proxy). The
URL-derived credential takes precedence over any `Proxy-Authorization` header the child
sends inbound; the inbound header is used only when the URL carries no userinfo. This is
sent in cleartext when the upstream scheme is `http://`, matching curl, Go and requests
behavior against a plain-HTTP proxy — the credential already lives in an env var every
tool on the box can read, and refusing to send it would break most corporate proxies.
Chaining is not pre-flighted: a broken upstream proxy surfaces per connection, as a `403`
(the upstream proxy is a loopback address — refused, so the terminator cannot chain into
itself) or a `502` (the upstream proxy could not be reached), logged as
`event=self_loop_blocked_proxy` / `event=tunnel_connect_failed`.
An `https://` upstream proxy is dialled over TLS — `ssl.create_default_context()` with
default certificate validation, SNI set to the proxy's own hostname, and ALPN pinned to
`http/1.1` — instead of plaintext on `:443`. There is no bypass knob. An operator whose
corporate TLS proxy presents a certificate from an internal/private CA (not merged via
`SSL_CERT_FILE`/`NODE_EXTRA_CA_CERTS` above) will see chaining fail closed with a `502`
instead of the prior silent-plaintext behavior; add that CA to the merged bundle to fix it.
#### Fail-open and known limits
On compression or dispatch errors, the Headroom terminator fails open (forwards original
bytes) so `agy` continues working. A session-level fail-open warning (first occurrence) and
an end-of-session compression summary are shipped — see the "Compression fail-open
observability" row in [docs/agy-parity-matrix.md](docs/agy-parity-matrix.md).
The Headroom MCP retrieve tool (persistent, ledger-recorded, resolves markers from the on-disk
store) and Serena code memory are wired via `AgyRegistrar`. `--no-mcp` skips the retrieve
server and `--no-serena` skips Serena, matching `wrap claude` and `wrap opencode`; both leave
MCP entries that Headroom did not install untouched. `--code-graph` starts the proxy's live
code-graph watcher, exactly as it does for every other wrapped agent. MCP registration in
`--print`/`-p`/`--prompt` mode requires agy `>= 1.0.16`; older or undetectable agy versions
skip registration and purge any stale entries — see
[docs/agy-parity-matrix.md](docs/agy-parity-matrix.md) for the full parity table.
The following features available on other agents have no agy equivalent in v1:
- `--memory` — no equivalent persistent memory API in agy
- `--learn` — requires a stable dispatch endpoint (headroom-2i0)
The Rust proxy backend is not supported for `agy`; `headroom wrap agy` hard-fails
with a clear message if `HEADROOM_BACKEND=rust` is set.
## When to use · When to skip
**Great fit if you…**

9
crates/headroom-core/tests/auth_mode.rs
View file

@ -149,6 +149,15 @@ fn antigravity_ua_classified_subscription() {
assert_eq!(classify(&h), AuthMode::Subscription);
}
#[test]
fn agy_cli_ua_classified_subscription() {
// agy (Google Antigravity CLI) real UA: lowercase `antigravity/<semver>`.
// Captured from live agy traffic; the prefix list matches on the
// lowercased UA so this is the canonical form agy actually sends.
let h = headers(&[("user-agent", "antigravity/1.0.5")]);
assert_eq!(classify(&h), AuthMode::Subscription);
}
// ── Performance ──────────────────────────────────────────────────
/// Smoke perf check — a strict bench lives at

417
docs/adr/0001-agy-mitm-transport.md Normal file
View file

@ -0,0 +1,417 @@
# ADR 0001 — Transport for compressing Google Antigravity CLI (`agy`) traffic
- Status: Accepted (design-review-gate PASSED — PM/Architect/Designer/Security/CTO all APPROVED, 2026-06-15)
- Date: 2026-06-15
- Epic: `headroom-30y` · Task: `headroom-30y.1`
## Context
Headroom wraps coding agents by pointing them at its local proxy via a base-URL
environment variable (Claude Code → `ANTHROPIC_BASE_URL`, Codex → `config.toml`,
etc.) and compressing the JSON bodies that flow through.
`agy` (Google Antigravity CLI) cannot be wrapped this way. Verified empirically:
- `agy` is a stripped **Go** binary (not Node), config dir `~/.gemini/antigravity-cli/`.
- It exposes **no base-URL override**: `CODE_ASSIST_ENDPOINT`, `GOOGLE_GEMINI_BASE_URL`,
`GOOGLE_CLOUD_CODE_ENDPOINT` are absent from the binary and are **ignored at runtime**
(live test: `agy --print` returned correct output with all three pointed at a dead port).
- It **honors** Go proxy vars (`HTTPS_PROXY`/`HTTP_PROXY`) and CA-trust vars
(`SSL_CERT_FILE`/`CACERT_PATH`/`NODE_EXTRA_CA_CERTS`).
- Backend: reached via HTTP **CONNECT** then TLS + **HTTP/2**, REST JSON
`POST /v1internal:streamGenerateContent?alt=sse` (SSE response). No TLS pinning
(a mitmproxy CA was accepted in the capture spike).
- The request body (`{model, project, request:{contents:[{parts:[{text}]}]}}`) is **already**
what `headroom/proxy/handlers/gemini.py:handle_google_cloudcode_stream` compresses.
So compression value is reachable, but only by intercepting `agy`'s TLS — Headroom has
no forward-proxy / CONNECT / certificate-minting capability today (only a reverse proxy
and upstream CA-trust discovery in `ssl_context.py`).
### Two distinct hosts (do not conflate)
- **Allowlist host** = the host `agy` opens `CONNECT` to (capture-verified:
`daily-cloudcode-pa.googleapis.com`). The terminator matches on this.
- **Upstream host** = where the existing handler re-originates the request. Today that is
the (wrong) constant `ANTIGRAVITY_DAILY_API_URL = "https://daily-cloudcode-pa.sandbox.googleapis.com"`
(`gemini.py:28`), corrected under `headroom-30y.4`. These are separate values.
## Decision
**Selective single-host embedded MITM, hosted in the Python proxy.**
A loopback (`127.0.0.1`-only) forward-proxy listener — a **separate `asyncio.start_server`
listener inside the same process** as the FastAPI/uvicorn app (uvicorn does not accept
`CONNECT`), so "one process" holds.
1. It accepts `CONNECT` and normalizes the target host (`normalize_host`: lowercase, strip a
trailing root dot and any `:port`). If that host is in the **cloudcode allowlist**, the
terminator answers `200 Connection Established` and byte-splices the raw connection to the
in-process **hypercorn** dispatch server on loopback. It does not terminate TLS itself.
2. For **every other** `CONNECT`, it performs a raw bidirectional **byte-splice** — no TLS
termination, no certificate, no inspection.
Both paths splice bytes; only the destination differs. `AgyDispatchServer` terminates TLS for
the allowlisted host, minting a leaf per SNI from the local root CA (`mint_leaf`, cached in
`_LeafCache`), negotiating **h2 or http/1.1** via ALPN, and serving the **existing FastAPI
app** — so the decrypted request reaches the same `/v1internal:streamGenerateContent` route →
`handle_google_cloudcode_stream`, and compression and upstream origination are unchanged.
Serving the app under hypercorn rather than hand-rolling server-side HTTP/2 framing also
removes the h2-vs-http/1.1 unknown (an http/1.1-downgrade live test was inconclusive: agy's
OAuth token had expired, and mitmproxy over-terminates the non-selective auth path). New dep:
`hypercorn`.
*(Superseded 2026-07-31: earlier revisions had the terminator mint a leaf and terminate TLS
in-process, with the hypercorn dispatch server recorded here as a later amendment. Production
always tunnelled to dispatch, so that code path survived only in tests and has been deleted
along with `_upgrade_to_tls_server`, `_build_server_ssl_context` and `DispatchCallback`.)*
### Host normalization is one invariant, not four checks
Four places compare a host against the allowlist: the `CONNECT` target, the dispatch SNI
callback, the post-handshake `Host` guard, and `cloudcode_host_base` on the passthrough path.
All four compare the output of `normalize_host`. When they disagreed, `CloudCode-PA.googleapis.com`
passed the SNI and Host guards but failed the exact-match `CONNECT` check, so the connection
fell through to the blind tunnel: the request still worked, but skipped termination and
compression with no signal that anything had been bypassed. Silent bypass is worse than a
hard failure, which is why the normalization belongs in one function that all four call.
`normalize_host` lowercases, strips a trailing root dot, and strips `:port` **only when the
suffix is all digits**. `example.com:abc` names no port, so it stays whole and fails the
allowlist as it should; requiring exactly one colon leaves IPv6 literals such as `::1` alone.
### Blind-tunnel targets are restricted
The terminator is an unauthenticated `CONNECT` proxy on loopback for the life of an agy
session, so anything running as the user can drive it. `_resolve_tunnel_target` refuses two
destinations and returns the vetted address the tunnel then dials:
- **the terminator's own port** — `CONNECT 127.0.0.1:<terminator_port>` makes it tunnel into
itself, costing two file descriptors per nesting level until they run out;
- **link-local addresses** — `169.254.0.0/16` carries the cloud instance-metadata service.
The check runs on the resolved addresses rather than the literal, so a name that resolves to
`127.0.0.1` is caught too, and dialling the vetted address means no second lookup can
substitute another. Other loopback ports stay reachable on purpose: a local process can open
them directly, so refusing them would buy nothing and break plain local tunnelling.
### Upstream-origination ownership (single connection)
The terminator (A2) is **agy-facing only**. It does **not** dial upstream for the allowlist
host. The dispatch adapter (T2) wraps the decrypted request as a Starlette `Request`
(ASGI scope: method/path/query/headers + a `receive()` yielding the decrypted body — the
seam the handler needs, since it reads `_read_request_json(request)`,
`dict(request.headers.items())`, `request.url.query`) and invokes the **existing**
`handle_google_cloudcode_stream`, which remains the **sole** upstream originator (it already
opens the upstream connection via `self.http_client.send(..., stream=True)`). The terminator
splices the handler's `StreamingResponse` (SSE) back over the terminated socket. Exactly one
upstream TLS session per request; the OAuth token is sent upstream once.
**The request goes back to the host agy chose.** `_resolve_cloudcode_base_url` takes the
`CONNECT` host (carried through as the `Host` header) and re-originates to that same host when
it is allowlisted, falling back to the default backend otherwise. This matters because the
allowlist holds two hosts: resolving every antigravity request to one default sent a request
addressed to `cloudcode-pa.googleapis.com` — and the OAuth bearer with it — to
`daily-cloudcode-pa.googleapis.com` instead. An explicit `HEADROOM_ANTIGRAVITY_API_URL` still
wins over both, since an operator setting it is choosing the backend deliberately.
### Module invariant (acyclic)
`ca-lifecycle (A1) ← terminator (A2) ← dispatch (T2) → existing handler`. Imports point one
way; the dispatch adapter never reaches back into transport.
`agy` is wrapped by injecting `HTTPS_PROXY=127.0.0.1:<port>` plus a combined CA bundle into
`SSL_CERT_FILE`/`CACERT_PATH`/`NODE_EXTRA_CA_CERTS`.
### Transparency & consent (required)
Wrapping `agy` terminates TLS on its AI connection and makes plaintext `Authorization` /
`x-goog-api-key` visible to the Headroom process. This is categorically different from
base-URL wrapping. Therefore:
- `headroom wrap agy` MUST print a clear one-line disclosure at launch, **before**
`subprocess.run` and on all non-early-exit paths (via the `env_vars_display` banner): that
Headroom is intercepting `agy`'s TLS to the **named** cloudcode host
(`daily-cloudcode-pa.googleapis.com`) via a local, process-scoped CA.
- The docs (`headroom-30y.6`) MUST state this plainly (value-parity, MITM mechanism).
- A `--no-intercept` / `--no-mitm` escape hatch runs `agy` through Headroom in
byte-splice-only mode (no compression) for users who decline interception.
- `headroom unwrap agy` MUST exist (agy is the first **durable** wrap-only command — it
writes `mcp_config.json` / `GEMINI.md`; `goose`/`openhands` write nothing and have no
unwrap). Unwrap removes only Headroom-added entries (merge semantics).
### Enterprise / corporate-proxy coexistence (required, v1 = chain)
`agy` honors a single `HTTPS_PROXY` and one CA bundle, which Headroom overwrites. **v1 commits
to chaining** (not documented-unsupported):
- detect a pre-existing user `HTTPS_PROXY` and **chain** to it — the terminator forwards
non-allowlist CONNECTs through the corporate proxy (never TLS-terminating the chained
leg), instead of dialing direct. The child is handed a userinfo-free loopback URL and
sends no `Proxy-Authorization` header of its own, so a chained CONNECT reached `407`
until this ADR's v1.1 update: the terminator now derives `Proxy-Authorization` from
`HTTPS_PROXY`'s own `user:pass@` userinfo (percent-decoded, `http`/`https` schemes
only) and sends it, taking precedence over any inbound header. This is sent in
cleartext to an `http://` upstream proxy, matching curl/Go/requests; and
- merge any pre-existing corporate CA (from the user's `SSL_CERT_FILE`/`NODE_EXTRA_CA_CERTS`
or system store) into the combined bundle so the real internet still validates. Only x509
objects with `basicConstraints CA:TRUE` are merged (do not blindly concatenate arbitrary
user-pointed PEM, which would widen `agy`'s trust beyond intended roots).
Chaining failures are reported per connection (`403` for a loopback upstream proxy, `502`
when it cannot be reached) rather than pre-flighted at launch.
An `https://` upstream proxy is chained to over TLS — `ssl.create_default_context()`
(default certificate validation, no bypass knob), SNI set to the proxy's own hostname
(the tunnelled target's TLS handshake and SNI travel separately, inside the tunnel), and
ALPN pinned to `http/1.1` so a proxy that would otherwise negotiate `h2` cannot leave the
terminator writing a CONNECT frame into an HTTP/2 connection. There is deliberately no
override: an operator whose corporate TLS proxy presents an internal-CA certificate that
isn't merged into the combined bundle goes from working (accidentally, over plaintext) to
a `502`, surfaced per connection rather than silently downgraded.
### Fail-open observability (required)
Failing open (forward original bytes on compression/dispatch error) keeps `agy` working, but
must never silently nullify the product's value. The design MUST:
- emit a one-line **stderr warning on the first** fail-open occurrence per session
(compression degraded to passthrough), and
- print an **end-of-session summary** (compressed exchanges vs passthrough count / observed
token-compression ratio).
The live smoke (T12) already asserts compression is *observed*, not merely error-free; these
signals extend that to the user's normal runtime.
### Properties
- **Performance:** exactly one TLS termination, only on the AI host; all other traffic is a
zero-parse byte-splice. No second process, no double-TLS, no double-HTTP/2 reframe, no
extra network hop. Existing handler reused.
- **Security:** see threat model. Interception surface limited to the AI host; root CA
process-scoped and never in the OS trust store; the **upstream** (Google-facing) leg keeps
**full** certificate verification against system roots — MITM on the agy-facing side never
implies trust-anything upstream.
- **Stability:** fail-open — any compression/dispatch error forwards the original bytes so
`agy` never breaks; fail-fast on security-critical setup (CA generation, port bind).
## Alternatives considered
| # | Alternative | Verdict | Reason |
|---|---|---|---|
| A | **Embedded single-process MITM, Python** | **CHOSEN** | One process, reuses the Starlette-coupled handler; `cryptography` + `h2` available. Lowest effort-adjusted cost. |
| B | Embedded MITM in the Rust proxy (`crates/headroom-proxy`) | **N/A (resolved, headroom-30y.11)** | The Rust proxy crate is a standalone port that **no `wrap` command launches** — every agent (claude/codex/aider/goose/openhands/openclaw/gemini/agy) runs through the Python proxy (`_start_proxy` → `python -m headroom.cli proxy`). The crate is also client-only (no rustls server / `rcgen` / CONNECT acceptor). Porting the MITM stack to a proxy that carries no wrap traffic is effort for a dead path; agy MITM is **Python-only by design**. The `wrap agy` Rust-backend hard-fail (below) is the enforced contract. No silent drift: documented here. (The Rust **core** — `headroom-core` smart_crusher + `auth_mode` agy classification — already has its agy parity via PyO3.) |
| C | Single-host reverse target via `HTTPS_PROXY`, no per-host MITM | Rejected | The capture shows `agy` uses `CONNECT` + TLS; a passive reverse target without TLS termination cannot read the body. |
| D | `mitmproxy` sidecar | Rejected | Second process + double TLS termination + double HTTP/2 reframe per SSE request + heavyweight dep — a middleman that erodes the latency value proposition. |
| — | Full dynamic per-SNI MITM (intercept all hosts) | Rejected | Needless interception surface / security risk; only one upstream host matters. |
## CA threat model
- Root CA generated once, stored `~/.headroom/ca/` (dir `0700`, key `0600`), regenerated on
expiry; `basicConstraints` CA:TRUE, `pathlen:0`. On regeneration, old leaf certs and the
old combined bundle are deleted.
- The CA is **never** added to the OS/system trust store. Injected **only** into the wrapped
`agy` process environment.
- The combined bundle (= system roots + Headroom CA cert + any pre-existing corporate CA;
public certs only, no key) is written under `~/.headroom` with `0600` perms (not a
predictable world-writable temp path); perms asserted after write.
- Leaf certs minted for the cloudcode allowlist host(s), validity ≤ 72h, SAN/EKU
constrained to that host + `serverAuth` only, cached (bound = allowlist size + 1 — the
extra slot holds the `headroom.internal` placeholder leaf, below). A non-served placeholder
leaf is minted once at dispatch start to satisfy `ssl.SSLContext.load_cert_chain` before the
SNI callback exists; it is never put on the wire (see dispatch trust-boundary enforcement).
- **Dispatch trust-boundary enforcement (allowlist at the SNI + authority layer).** The
dispatch hypercorn listener is itself a loopback HTTPS port; a local process could connect
directly and request a leaf for any SNI. Enforced in two layers: (1) the per-SNI
`set_servername_callback` rejects any `server_name` that is `None` or (lowercased) not in
the allowlist with `ssl.ALERT_DESCRIPTION_UNRECOGNIZED_NAME` **before** any mint/cache/swap;
(2) a mandatory post-handshake ASGI `host`/`:authority` guard (`make_host_guard`) returns
421 for absent/duplicate/non-allowlisted Host — covering the no-SNI/placeholder path where
OpenSSL may skip the SNI callback. The dispatch allowlist is the same single value wired
into the CONNECT terminator (no drift).
- **Leaf private key handling:** `load_cert_chain_in_memory` (`headroom/proxy/agy_ca.py`) is
used at both `load_cert_chain` call sites (dispatch placeholder init; dispatch
`_sni_callback`). The terminator has none: it byte-splices to the dispatch server and never
terminates TLS itself. Primary path (Linux, `os.memfd_create`
available): combined cert+key PEM is written into an anonymous `memfd_create("hr_leaf")`
file descriptor and loaded via `/proc/self/fd/{fd}`; the fd is closed after load so no file
ever exists on a filesystem. Fallback path (`memfd_create` absent or `/proc` inaccessible,
e.g., certain containers): `tempfile.mkstemp` creates a 0600 temp file; perms are asserted
via `_assert_perms`; `load_cert_chain` reads it; `os.unlink` removes it in a `finally`
block even if load raises. Leaf private keys are **never** added to any trust store and
**never** persist beyond the single `load_cert_chain` call.
- `~/.headroom` (the bundle's parent dir) is `0700`; the CA store `~/.headroom/ca/` is `0700`
with key `0600`; the combined bundle file is `0600`. All perms asserted after write.
- Listener bound to `127.0.0.1` only; `NO_PROXY=127.0.0.1,localhost` loop-guard so the
terminator can never CONNECT to itself.
- **No TLS-verification bypass:** an earlier revision of this work added an
`HEADROOM_SSL_VERIFY=false` switch that blanked `SSL_CERT_FILE`/`CURL_CA_BUNDLE` and set
`NODE_TLS_REJECT_UNAUTHORIZED=0` for launched agents, with `agy` exempted so the injected
bundle survived. It has been removed: upstream ships no such switch, and a PR that adds
TLS interception must not also add a way to turn verification off.
- Plaintext `Authorization` / `x-goog-api-key` post-termination are routed only through the
existing `redact_for_wire_debug` redactor (helpers.py — covers both keys); the request auth
is not persisted in the semantic cache (verified: cache keys on messages+model, stores
response headers only). No parallel log sink is introduced.
## Files touched (regression-audit surface)
- New: `headroom/proxy/` CA-lifecycle, terminator, dispatch-adapter modules.
- Edited (shared): `headroom/cli/wrap.py` (`agy()` + `unwrap agy` +
`_launch_tool` threading); `headroom/proxy/handlers/gemini.py:28`
(host const + resolver, via T4). Handler `gemini.py:740` reused, not modified internally.
## Consequences
- `agy` shipped wrap-only (like `goose`/`openhands`), not added to `ToolTarget`; but it is the
first wrap-only command with durable on-disk state, so it gains an `unwrap` command.
- HTTP/2 negotiated on the agy-facing side (`h2` sans-io server); upstream leg uses the
handler's existing httpx h2 client.
- If the Rust proxy is the active backend, `headroom wrap agy` hard-fails with a clear
"unsupported on Rust backend" message rather than mis-route. This is the enforced contract.
- The Rust proxy port (`crates/headroom-proxy`) gets no `agy` support — **resolved N/A**
(headroom-30y.11): it carries no `wrap` traffic for any agent, so agy MITM is Python-only by
design. Documented, not silently dropped.
## Retrieve MCP transport: stdio child, not url-MCP
agy 1.0.10 added `url` support in `mcp_config.json`, allowing an MCP server to be addressed
by HTTP URL instead of a stdio subprocess. The headroom retrieve server (`AgyRetrieveServer`,
`headroom/proxy/agy_retrieve.py`) is `AgyDispatchServer(plain_http=True)`: the same hypercorn
lifecycle serving the same FastAPI app, minus the SNI TLS context and the Host guard. It
answers plain HTTP on loopback and does **not** implement the MCP-over-HTTP (streamable HTTP)
transport. Registering it as a `url`-type entry
would require adding an MCP-HTTP transport layer to the retrieve server for **zero added
capability** — the stdio child (`headroom mcp serve`) already satisfies all retrieve use cases,
and the per-run ephemeral listener is reverted on teardown with no dead pointer left in
`mcp_config.json`.
**Decision:** keep the retrieve integration as a stdio child; do not add an MCP-HTTP transport
to `AgyRetrieveServer`. Revisit only if agy deprecates stdio MCP support.
## Cross-platform status
The agy slice runs on Windows and is **CI-gated** on it: the `agy-windows` job
(`.github/workflows/ci.yml`, `windows-latest`) runs the full agy suite plus
`test_wrap_agy.py` on every code change.
Platform specifics:
- `_assert_perms` is a no-op on non-POSIX platforms (no `os.chmod`/`stat` crash on Windows).
- Atomic bundle writes use `os.replace`; `_write_secure` ORs in `os.O_BINARY`
(0 on POSIX) so PEM bytes are written verbatim, not CRLF-translated, on Windows.
- System trust source: POSIX/macOS read the detected on-disk CA bundle; Windows has
no single bundle file, so `_system_trust_pem()` enumerates the ROOT+CA cert stores
via stdlib `ssl.enum_certificates`, run through the same CA:TRUE filter (no leaf
trusted as an anchor; no `certifi` dependency).
- Loopback sockets set `SO_REUSEADDR` only on POSIX; on Windows that flag would let
another local process bind the same port and intercept decrypted traffic, so Windows
uses `SO_EXCLUSIVEADDRUSE` instead.
**Leaf private-key posture differs by platform (security-relevant):**
- **Linux:** the leaf key is loaded from an anonymous `memfd` and **never touches the
filesystem**.
- **Windows / macOS (no `memfd`):** the leaf key is written to a `mkstemp` file and
unlinked immediately after `load_cert_chain`. On POSIX the file is `0600`; on Windows
POSIX mode bits are not enforceable, so protection comes from the temp directory's
ACL. **Verified** on `windows-latest` via `icacls`: an `hr_leaf_*.pem` mkstemp file in
`%LOCALAPPDATA%\Temp` grants Full control only to the owning user, `NT AUTHORITY\SYSTEM`,
and `BUILTIN\Administrators` — no `Users`/`Everyone`/`Authenticated Users` entry, i.e.
user-scoped, not world-readable (Administrators can read any file on any OS — unavoidable).
The guarantee is therefore "owner-only ACL (inherited from `%TEMP%`) + immediate unlink",
**not** the Linux "never on disk" invariant. The residual exposure is the brief on-disk
window, mitigated by the immediate unlink; this is a deliberate, documented degradation.
## Savings & dashboard integration (per-project attribution)
`headroom wrap agy` runs its selective-MITM dispatch as an in-process `create_app()`
inside the wrap process — a **separate OS process** from the long-running shared proxy
that renders the savings dashboard. The dashboard reads that shared process's *in-memory*
metrics (`m.tokens_saved_total`, `m.savings_tracker.stats_preview()`), so agy's savings,
recorded in agy's own process, never reached it. Two consequences were reported on
PR #1044: no agy savings on the dashboard, and no agy project row in Per-Project Savings.
Resolution (two parts):
- **Per-project attribution.** agy is a Go binary with no header/base-URL knob, so the
project label cannot be injected via the child's env (as it is for Claude/Codex). It is
injected at the MITM boundary instead: `make_host_guard` stamps `x-headroom-project`
(the launch-directory basename, computed once) onto every intercepted request *after*
the Host allowlist check — the trust boundary is unchanged.
- **Cross-process savings via a durable event inbox.** agy does not write shared savings
state directly. In the agy process, `HEADROOM_SAVINGS_PATH`, `HEADROOM_SAVINGS_EVENTS_PATH`,
and `HEADROOM_OTEL_METRICS_ENABLED=0` are redirected to a throwaway temp dir, and each
request emits one event file into `~/.headroom/savings.d/` carrying the exact
`PrometheusMetrics.record_request` arguments. The shared proxy drains that inbox (a
periodic task plus an opportunistic drain on `/stats`) and **replays each event through
its own `record_request` funnel** — the single funnel that already updates every
dashboard surface (token/$ heroes, per-project rows, history, CSV, ledger). Because the
proxy replay is the sole writer of shared savings state, each agy request is counted
once. Delivery is **at-least-once** with a best-effort processed-id journal: savings are
estimates, so a rare double-count in the crash window between record and unlink is
accepted rather than paying for a transactional store. For users who never run agy the
inbox is empty and the dashboard is byte-identical to before.
## Third-party tool parity (code memory)
`headroom wrap agy` sets up the same code memory as every other client: **Serena is the
engine**, registered via `AgyRegistrar` with a verify-then-remove `initialize` handshake and
a ledger record so `unwrap agy` removes it cleanly; user-managed entries are preserved.
tokensave and the CLI context tools (rtk, lean-ctx) were retired upstream, so `wrap agy`
installs neither — it only *removes* what earlier releases left behind
(`_disable_tokensave_mcp`, `headroom.context_tool_cleanup`). `--no-tokensave` survives as a
hidden no-op flag; `--code-graph` no longer registers an MCP server for agy at all, it
forwards to the proxy's live code-graph watcher exactly as `wrap claude` does.
**MCP parity in all modes.** An earlier build of agy (~1.0.5) hung indefinitely in
`--print` mode whenever any MCP server was configured, so print mode used to register no MCP.
That hang was **fixed in agy 1.0.16** (re-verified 2026-07-05: Serena and the headroom
retrieve server both answer in ~4s in print mode). agy therefore now wires MCP tooling
**identically in print and interactive mode** — Serena plus the headroom retrieve MCP —
giving agy first-class MCP parity in every mode, like any
other client. Live-verified: `wrap agy -p` wires Serena + retrieve
(handshake-verified) and completes in ~10s. Because the fix is agy-side, `wrap agy` still
runs a runtime `agy --version` preflight before wiring print-mode MCP (headroom-37g.37): an
agy older than 1.0.16, or one whose version can't be detected, is treated as unsafe by
default, so print-mode MCP registration is suppressed and any previously-persisted entries
are purged for that run.
## functionResponse bulk compression (CCR) — where the savings actually come from
The savings-plumbing above only surfaces savings that a compressor produced; for agy the
compressor initially produced ~zero. Root cause: agy's request bulk lives in
`contents[].parts[].functionResponse.response` — the tool-output leaves the coding agent
resends every turn (file reads, greps, command output). Headroom's message-level
compressors never touched those leaves, so a large agy session compressed almost nothing
(PR #1044: "704 → 718" — compression *inflated* tokens and reverted).
**Design — uniform, deterministic, recoverable.** Every `functionResponse.response` string
leaf across ALL of `contents` (history + tail) above a marker-derived token floor is
replaced by a deterministic CCR marker; the original is cached under
`SHA-256(original)[:24]` and recovered on demand via the injected `headroom_retrieve` MCP
tool. Key properties:
- **Uniform, not live-tail-only.** agy is a MITM that never rewrites the client's own
history, and it resends the full history each turn. A live-zone/recency boundary (compress
cold history, keep the tail verbatim) is therefore **cache-incoherent** here: the same leaf
appears compressed in one turn and verbatim the next, so the model re-diffs it every turn.
Compressing every leaf identically each turn keeps the cross-turn byte-image stable.
- **Retrieved-content exemption (anti-thrash).** A leaf whose hash the model already fetched
this turn (a `headroom_retrieve` / `call_mcp_tool` call carrying that 24-hex hash in its
args) is left verbatim — otherwise the re-sent, just-expanded original would be
re-compressed into the same marker and the model would retrieve it forever. This mirrors
the retrieve-call suppression the OpenAI/Anthropic paths already do (keyed by hash, since
agy has no call_id).
- **Envelope exemption (name-independent).** The hash-in-args exemption above recognizes the
retrieve call by name/args; on agy the retrieve-result functionResponse carries an opaque
name and args of just `{hash}`, so that path misses it and the retrieve *output*
(`{hash, source, original_content, …}`) would itself re-compress into a marker the model
re-retrieves. The compressor therefore also detects the retrieve-result envelope by
**content** — value-bearing `hash` (24-hex) + `source` (`local`/`proxy`) anchors,
independent of tool name — and never compresses that leaf (L1), plus adds its resolved hash
to the retrieved set so the resent original is exempt too (L2). Verified live: one retrieval
per hash, no thrash.
- **Default + escape hatch.** `HEADROOM_AGY_FR_MODE` selects `ccr` (default, real savings)
or `lossless` (a safety floor that never emits markers). The WU4 efficacy trial gated the
default: ccr ships because it delivers material savings while `headroom_retrieve` is wired;
a lossless downgrade warns loudly if retrieve is not wired so markers can never become
unrecoverable silently.
- **Revert-independent accounting.** Savings are recorded from the compression decision, not
from whether the upstream later reverts — each turn independently avoided sending those
bytes.
## SSE output-token accounting (Cloud Code Assist response-envelope unwrap)
Cloud Code Assist wraps streaming responses in a `response` envelope
(`{"response": {"usageMetadata": {…}}}`), mirroring the request-side wrap. Both gemini SSE
usage parsers read `usageMetadata` at the top level only, so agy's `candidatesTokenCount`
never parsed and every turn logged "Could not parse output_tokens from SSE, estimating N
from B bytes" — output tokens on the dashboard/ledger were a `bytes//40` estimate. The
gemini branches now unwrap the envelope when top-level `usageMetadata` is absent; native
Gemini (top-level) chunks are unaffected.

32
docs/agy-parity-matrix.md Normal file
View file

@ -0,0 +1,32 @@
# agy Parity Matrix
Claude feature parity table for `headroom wrap agy`.
Status: **WIRED** = actively wired and tested; **N/A** = not wired (concrete evidence given); **DEFERRED** = ticket filed.
| Feature | Status | Mechanism / Evidence |
|---------|--------|----------------------|
| **CLI context tools (rtk, lean-ctx)** | **REMOVED (upstream)** | Headroom no longer ships CLI context tools for any agent (upstream PR #2677), so `wrap agy` wires none. The `--context-tool` / `--no-context-tool` flags and `HEADROOM_CONTEXT_TOOL` are rejected with an explanatory error rather than silently ignored (`_retired_context_tool_callback`, `wrap.py`), and every non-`selfheal`, non-`--help` wrap invocation runs `headroom.context_tool_cleanup.purge_context_tool_artifacts` to uninstall binaries, hook scripts, config backups and MCP entries the old integration left behind. |
| **Context-instructions (GEMINI.md)** | **N/A (cleanup only)** | The only block `wrap agy` ever wrote into `~/.gemini/GEMINI.md` carried the rtk context-tool instructions, so nothing is injected any more. `unwrap_agy` still calls `_remove_gemini_md_block` (`wrap.py`) to delete a block a pre-removal install left behind; user content outside the `<!-- headroom:agy-instructions -->` markers is preserved verbatim. |
| **Headroom MCP retrieve tool (persistent)** | **WIRED (persistent, ledger-recorded; registration version-gated)** | Resolves `[Retrieve more: hash=…]` markers produced by the HTTPS dispatch MITM. The marker cache is the **process-global** `get_compression_store()` singleton (`headroom/cache/compression_store.py`), so a SECOND in-process `create_app()` shares it. Because the dispatch server is HTTPS with a Cloud-Code-SNI leaf only (a `headroom mcp serve` stdio child can't reach it over loopback), `wrap agy` stands up a **second loopback listener — PLAIN HTTP, no TLS — on an ephemeral port** (`AgyRetrieveServer`, `headroom/proxy/agy_retrieve.py`) alongside the terminator+dispatch, via `_start_agy_servers(..., start_retrieve=True)` (`wrap.py`). **The listener now starts unconditionally on every `wrap agy` run** (interactive or print mode) — only the MCP *registration* pointing at it is gated. Registration uses a **stable, port-independent spec** (`build_headroom_spec()` → `env={}`, `install.py`, via `AgyRegistrar`, smoke-verified by `_setup_headroom_retrieve_mcp_agy`); the `headroom mcp serve` child resolves markers from the **on-disk CCR store** (`ccr.mcp_server._retrieve_content`, local-first) so no live proxy or per-run port is needed (the loopback listener above stays as an in-session HTTP fallback only). The entry is registered **persistently** and **recorded in the install ledger** — like Serena/CBM it is **NOT reverted on teardown**, which is what lets agy discover, cache, and **expose** `headroom_retrieve` across sessions (the exposure the `HEADROOM_AGY_RETRIEVE_WIRED` gate checks before keeping ccr on — headroom-h76.5). Print-mode version preflight is unchanged: interactive always wired; print mode requires agy `>= 1.0.16`, else registration is skipped and any `headroom` entry is purged **and its ledger record cleared** (`_purge_agy_mcp_entries`) — old agy hangs on any persisted MCP entry in print mode; re-registration happens on the next compatible wrap. `unwrap_agy` removes the entry **ledger-gated** (`_remove_headroom_installed_retrieve_mcp`), leaving user- or `mcp install`-managed `headroom` entries untouched. **Caveat:** the listener + retrieve HTTP endpoint are **headless-testable** (load-bearing test: a hash stored via `get_compression_store()` resolves over plain HTTP `GET /v1/retrieve/{hash}` from the second `create_app()` — `tests/test_agy_retrieve.py`), but agy actually **invoking** the tool mid-conversation is live-verified, not headless-proven. Ref: **headroom-2i0**. `--no-mcp` skips registration entirely (parity with `wrap claude` / `wrap opencode`), leaving `HEADROOM_AGY_RETRIEVE_WIRED` unset so the handler ships no unrecoverable markers. |
| **Headroom MCP retrieve tool (mcp install / stable proxy)** | **WIRED (install fleet)** | `AgyRegistrar` registered in `get_all_registrars()` (`install.py`). `headroom mcp install` will register the spec into `~/.gemini/config/mcp_config.json` (agy 1.1.x read-path, migrated from `~/.gemini/antigravity-cli/mcp_config.json`; shared with the Antigravity IDE). Merge-not-clobber: other `mcpServers` entries preserved. This fleet path does **not** write the install ledger, so `unwrap_agy`'s now **ledger-gated** removal (`_remove_headroom_installed_retrieve_mcp`) **leaves a `mcp install` entry in place** (deliberate fleet install respected); it removes only the persistent entry that `wrap agy` recorded. |
| **tokensave** | **RETIRED (upstream)** | tokensave is no longer installed for any agent; Serena is the code-memory engine. `--no-tokensave` is accepted but ignored (hidden, deprecated). Every `wrap agy` run calls `_disable_tokensave_mcp` so a tokensave entry a previous release recorded in the ledger is actively removed; user-managed entries are left alone. |
| **Serena MCP** | **WIRED (code memory; version-gated)** | Serena is a generic `uvx` stdio MCP server (`build_serena_spec`, `install.py`) with no proxy-URL/ephemeral-port dependency, so it persists cleanly in `mcp_config.json`. Registered as agy's code-memory engine via `_setup_serena_mcp(AgyRegistrar(), context="ide-assistant", force=True)` (Antigravity is an IDE agent → Serena's generic IDE profile), gated on the same print-mode version preflight. `--no-serena` actively removes a prior Headroom entry via `_disable_serena_mcp`. Reverted by ledger-gated `_remove_headroom_installed_serena_mcp(AgyRegistrar())` in `unwrap_agy` — preserves user-managed Serena entries. |
| **Print-mode MCP suppression (scope)** | **VERSION-GATED (wired when agy >= 1.0.16; else suppressed + purged)** | MCP wiring for agy is gated on a runtime `agy --version` preflight (`_agy_print_mode_mcp_allowed` / `_detect_agy_version`, `wrap.py`), not unconditionally suppressed in print mode. **Interactive** `wrap agy` runs are always wired — no version check. **Print-mode** runs (`--print`/`-p`/`--prompt`) are wired identically to interactive once `_detect_agy_version` finds agy `>= _AGY_PRINT_MODE_MCP_MIN_VERSION` (1.0.16) — re-verified 2026-07-05 that Serena and the headroom retrieve server both answer the `initialize` handshake in ~4s on 1.0.16, so the earlier unconditional print-mode hang no longer applies. When the detected version is older than 1.0.16, or can't be determined at all (no binary, non-zero exit, unparseable output, or a hung `--version` call — treated as unsafe, "safe-by-default"), MCP wiring is skipped for that run **and** `_purge_agy_mcp_entries` actively removes all **4** Headroom-owned MCP surfaces a prior interactive (or newer-agy) run may have persisted in `mcp_config.json`: tokensave and Serena via ledger-aware disable (`_disable_tokensave_mcp` / `_disable_serena_mcp`), plus a legacy codebase-memory-mcp entry and the persistent `headroom` retrieve entry via `registrar.unregister_server(...)` (the retrieve entry **also clears its ledger record** so the next compatible-agy run re-registers cleanly rather than treating the now-absent entry as still-installed). Merely skipping new registration is not enough — a stale entry from an earlier run would still hang this print-mode invocation. All purge calls are idempotent (no-op when the entry is already absent). A user's own pre-existing, non-Headroom-managed MCP servers are never touched by the purge. **The retrieve LISTENER is unaffected by this gate** — it starts unconditionally every run (see the retrieve-tool row above); only MCP *registration* is version-gated. |
| **Code-graph (`--code-graph`)** | **WIRED (opt-in; proxy-side, upstream semantics)** | `--code-graph` is forwarded to `_ensure_proxy(..., code_graph=code_graph)` exactly as every other `wrap` subcommand does, so it starts the proxy's live code-graph watcher (`headroom/graph/watcher.py`, incremental reindex via `codebase-memory-mcp`). agy registers **no** code-graph MCP entry of its own — the earlier agy-only `build_codegraph_spec` / `_setup_code_graph` path was removed when upstream repurposed the flag. Default OFF. `unwrap agy` still unregisters a legacy `codebase-memory-mcp` entry an older build wrote, mirroring `unwrap claude`. Headless tests: `tests/test_wrap_agy_proxy_wiring.py` (`test_code_graph_flag_forwards_to_proxy_watcher`, `test_code_graph_defaults_off`). |
| **functionResponse CCR compression** | **WIRED** | agy's per-turn bulk lives in `contents[].parts[].functionResponse.response` string leaves (tool-output the coding agent resends every turn — file reads, greps, command output), which the existing message-level compressors never touched (those non-text-carrying parts were routed into `preserved_indices` and restored verbatim by `_rebuild_gemini_contents`), so a large agy session compressed almost nothing (PR #1044: "704 → 718, reverting"). `compress_function_response_leaves` (`headroom/transforms/agy_fr_compressor.py`) replaces every `functionResponse.response` string leaf — historical and tail, uniformly — above a marker-derived token floor with a deterministic, SHA-256[:24] CCR marker (`default_ccr_hash`) resolved on demand by `headroom_retrieve`. Because headroom is an in-flight MITM that never rewrites agy's local history, agy re-sends the same original bytes every turn, so the deterministic transform yields a byte-stable compressed prefix that re-hits the Cloud Code Assist server-side cache. `GeminiHandlerMixin._compress_agy_function_responses` delegates to `compress_function_response_leaves` (moved out for standalone unit testing without booting the FastAPI app — headroom-37g.36). Recoverable by construction, never a lossy summary — the model reads functionResponse back as its own prior tool results, so a fabricated summary would corrupt multi-turn reasoning. Default `ccr` mode with a lossless floor. |
| **SSE output-token accounting** | **WIRED** | Cloud Code Assist streams responses wrapped in a response envelope; the SSE usage-metadata reader was not unwrapping it, so agy's `output_tokens` were derived from a byte-length estimate instead of the real upstream value. `_gemini_usage_meta` (`headroom/proxy/handlers/streaming.py`) now unwraps the Cloud Code Assist response envelope so `output_tokens` parse from the actual upstream usage metadata on both the SSE streaming paths that call it. |
| **Compression fail-open observability** | **WIRED** | Ref: **headroom-30y.15**. Two observability surfaces added, agy-scoped only (no edits to gemini.py/transport/compression_store). (1) **First-warning notice:** `FailOpenWarnHandler` (`headroom/providers/agy/stats.py`) is installed on logger `"headroom.proxy"` — the logger gemini.py actually emits on (NOT the child `"headroom.proxy.handlers.gemini"`; Python logging propagates child→parent, so a handler on the child would never receive parent records) — before the agy dispatch thread starts. On the FIRST log record whose message contains `"Cloud Code Assist optimization failed"` it prints ONE user-facing notice to stderr; the message-content filter ignores unrelated `headroom.proxy` warnings; all subsequent fail-open occurrences are counted silently. The handler is removed in `agy()`'s `finally` to avoid leaking into the click process. Thread-safe: one-shot flag and counter share a single `threading.Lock`. (2) **Session compression summary:** `AgySessionStats` (`headroom/providers/agy/stats.py`) snapshots `get_compression_store().get_stats()` at session start (before the dispatch thread); at session end it prints a one-line stderr summary of the delta — new entries added, original→compressed tokens, the compressed share of the original ("N% of original", divide-by-zero guarded), and the observed fail-open count. Idempotent: `print_summary()` prints exactly once regardless of invocation count; safe for the `finally`+SIGTERM double-path. **Honest caveat:** the token counts are a store-based proxy — they reflect the delta in the process-global compression store singleton, not a per-request audit log. Store entries can be evicted between start and end snapshots (delta clamped to ≥ 0). The summary covers all compression in the agy session's process (which only does agy dispatch compression, so this is accurate in practice). Wiring: `wrap.py` (session start: `snapshot_start` + `install_fail_open_handler`; `_agy_sigterm`: `print_summary` + `remove_fail_open_handler`; `finally`: same two calls). |
| **Savings dashboard + Per-Project row (shared 8787 proxy)** | **WIRED (live-smoke VERIFIED)** | Ref: **headroom-508.1** (reported by external testers on PR #1044). agy's in-process MITM dispatch emits savings events to `~/.headroom/savings.d/` (`agy_savings_inbox.emit_event`, `outcome.py`); the shared Headroom proxy's periodic DRAIN loop (`_drain_agy_savings_periodically`, `server.py`) replays them into the dashboard $/token hero + Per-Project Savings row (project via `x-headroom-project` injected at the MITM boundary). Previously `wrap agy` never started/ensured that shared proxy, so the inbox was never drained → empty dashboard. Now `agy()` takes `--port` (default 8787; no `-p` alias — agy's own `-p` is `--print`, headroom-r9k) + `--no-proxy` and calls `_ensure_proxy(port, no_proxy, agent_type="agy", code_graph=code_graph)` **before** the agy-only `os.environ` mutations, so a freshly-spawned shared proxy is NOT poisoned by `HEADROOM_AGY_INBOX_EMIT` / `HEADROOM_SAVINGS_PATH` / `HEADROOM_SAVINGS_EVENTS_PATH` / `HEADROOM_OTEL_METRICS_ENABLED` (which `_start_proxy` would otherwise inherit via `os.environ.copy()`). Coexists with the MITM dispatch (agy has no base-URL knob → NO `_push_runtime_env` redirect). Teardown reuses the refcounted `_make_cleanup`/`_register_proxy_client` (stops only a proxy wrap-agy started, never a pre-existing user proxy), wired into agy's existing `_agy_sigterm` + `finally` (no second `signal.signal`). **Headless tests:** `tests/test_wrap_agy_proxy_wiring.py` (env-ordering guard + `agent_type="agy"` + `--no-proxy` passthrough + cleanup-on-teardown); refcount correctness in `tests/test_cli/test_wrap_helpers.py`. **Live smoke (headroom-90k, CLOSED):** a live `wrap agy` run confirmed the dashboard $/token hero (\$0.158, 52,601 tokens saved) and the Per-Project Savings row both surfaced correctly. |
| **--memory** | **N/A** | `--memory` wires `ClaudeCodeAdapter` writing to `claude_memory_dir` (`wrap.py`). agy has no equivalent persistent memory directory API exposed to external processes. Deferred: **headroom-2i0** notes this as follow-on. |
| **--learn** | **N/A** | `--learn` wires the Headroom learning surface through `AgyDispatchServer` transport side (not MCP). The dispatch server (`agy_dispatch.py`) would need a `--learn` POST endpoint exposed to agy's tool calls. This is transport-side work outside T9 scope. Deferred: **headroom-2i0**. |
| **ENABLE_TOOL_SEARCH** | **N/A** | `ENABLE_TOOL_SEARCH` is a claude-specific env var that activates the web-search tool in Claude Code. agy uses Google Search natively; no equivalent env plumbing needed. No ticket required. |
| **Retrieve MCP transport (url vs stdio)** | **STDIO (by design)** | agy 1.0.10 added `url`-type MCP entries. `AgyRetrieveServer` (`headroom/proxy/agy_retrieve.py`) is `AgyDispatchServer(plain_http=True)` — the same hypercorn lifecycle without the SNI TLS context or Host guard — so it answers plain HTTP/REST and does NOT implement MCP-over-HTTP (streamable HTTP). Registering it as a `url` entry would require adding an MCP-HTTP transport for zero added capability; the stdio child already works. Decision: stdio child stays; see ADR 0001 "Retrieve MCP transport". |
| **Cross-platform (Windows)** | **CODE SAFE; CI WIRED** | CA lifecycle and CONNECT terminator code is Windows-safe: `_assert_perms` is a no-op on non-POSIX; atomic bundle writes use `os.replace`; no POSIX-only crash path remains. The `agy-windows` CI job (`.github/workflows/ci.yml`) runs the agy CA/dispatch/terminator/retrieve/stats/registrar/wrap slice (`tests/test_agy_ca.py`, `test_agy_dispatch.py`, `test_agy_terminator.py`, `test_agy_retrieve.py`, `test_agy_stats.py`, `test_agy_registrar.py`, `test_proxy_google_cloudcode_route_aliases.py`, `test_wrap_agy.py`) on `windows-latest`, the only Windows coverage lane for this slice (the main shards run on Linux). Native-Windows E2E CI (`wrap-native-e2e.yml`, `install-native-e2e.yml`) remains excluded pending an upstream CRT issue — do not claim "Windows fully supported" until that native CI is green too. |
## Follow-up tickets
| Ticket | Feature | What's needed |
|--------|---------|---------------|
| **headroom-2i0** | Headroom MCP retrieve wiring — **DONE**; made **persistent + local-store-backed** in headroom-h76.6 (stable `build_headroom_spec()` spec, ledger-recorded, NOT reverted — so agy caches/exposes `headroom_retrieve` across sessions; the `AgyRetrieveServer` loopback listener remains an in-session HTTP fallback). MCP registration stays version-gated + smoke-verified; exposure-gated ccr downgrade in headroom-h76.5. Remaining: `--learn`; `--memory`. | Retrieve markers resolve via a second loopback `create_app()` sharing the process-global compression cache; see the matrix row above. `--learn`/`--memory` remain transport-side / no-equivalent-API work. |
| **headroom-30y.13** | Code-graph (`--code-graph`) — **DONE**, then **REALIGNED**: upstream repurposed the flag to the proxy-side watcher, so the agy-only MCP registration was removed. | `--code-graph` now only forwards to `_ensure_proxy(code_graph=...)`, identical to `wrap claude`; `unwrap agy` cleans a legacy cbm entry. |
| **headroom-30y.11** | Rust-proxy MITM parity — **RESOLVED N/A**. | The Rust proxy port (`crates/headroom-proxy`) carries **no `wrap` traffic** for any agent — every agent (incl. agy) runs through the Python proxy (`_start_proxy` → `python -m headroom.cli proxy`). agy MITM is **Python-only by design**; `wrap agy` hard-fails on a Rust backend. No silent drift (documented here + ADR 0001 alt-B). The Rust **core** (`headroom-core` smart_crusher + `auth_mode`) already has agy parity via PyO3. |

23
docs/content/docs/proxy.mdx
View file

@ -410,6 +410,24 @@ Google Cloud Code Assist / Antigravity compatibility endpoint used by Pi-style `
The proxy also accepts:
- `POST /v1/v1internal:streamGenerateContent`
**agy / Google Antigravity CLI:** `agy` cannot use this endpoint via a base-URL redirect
because it has no base-URL override knob. Headroom wraps `agy` via a selective TLS-MITM
transport instead:
- `HTTPS_PROXY` and `HTTP_PROXY` are set to the Headroom in-process CONNECT terminator
(loopback, bound to `127.0.0.1`).
- Three CA trust-bundle variables are set to a combined PEM bundle
(`SSL_CERT_FILE`, `CACERT_PATH`, `NODE_EXTRA_CA_CERTS`).
- Only `daily-cloudcode-pa.googleapis.com` (and `cloudcode-pa.googleapis.com`) are
TLS-terminated. All other CONNECT tunnels are byte-spliced unchanged.
- HTTP/2 and SSE are negotiated natively via an in-process hypercorn HTTPS server;
decrypted requests are routed to the existing `handle_google_cloudcode_stream` handler.
- The process-local CA is stored under `~/.headroom/ca/` and is **never** added to the
OS trust store.
Usage: `headroom wrap agy` (also `headroom unwrap agy`, `--no-intercept`).
See the [agy section in README.md](../../../README.md#using-headroom-with-agy) for full details.
### `POST /v1/compress`
Compression-only endpoint. Compresses messages and returns them without ever making a **completion request to an LLM provider** — no generation, no provider API key, no upstream chat call. Used by the TypeScript SDK, by LiteLLM's `headroom` guardrail, and by API gateways running Headroom as a sidecar.
@ -582,6 +600,11 @@ headroom wrap cursor
# Grok Build (updates ~/.grok/config.toml and starts the proxy)
headroom wrap grok-build
# Google Antigravity CLI (agy) — TLS-MITM transport, not a base-URL redirect
headroom wrap agy
headroom wrap agy --no-intercept # passthrough, no compression
headroom unwrap agy # revert GEMINI.md and MCP config changes
```
Cursor reads model endpoints from its settings UI, so `headroom wrap cursor`

12
headroom/cache/compression_store.py vendored
View file

@ -92,6 +92,16 @@ def _get_env_default_ttl_seconds() -> int:
return ttl_seconds
def default_ccr_hash(content: str) -> str:
"""SHA-256(content)[:24] -- the default CCR store key.
Single source of truth so the compression store's key and any exemption
recompute (e.g. the agy retrieve exemption in
``headroom.transforms.agy_fr_compressor``) cannot drift.
"""
return hashlib.sha256(content.encode()).hexdigest()[:24]
def format_retrieval_miss_detail(status: dict[str, Any]) -> str:
"""Return an operator-facing miss reason for CCR retrieval failures."""
default_ttl = status.get("default_ttl_seconds", DEFAULT_CCR_TTL_SECONDS)
@ -340,7 +350,7 @@ class CompressionStore:
# in-memory, so changing the hash function on upgrade has no
# persistence-side effect — the same content always hashes
# deterministically under whichever function is in use.
hash_key = hashlib.sha256(original.encode()).hexdigest()[:24]
hash_key = default_ccr_hash(original)
# Refuse to persist a bare CCR marker as an entry's "original"
# (#2694). A marker is a *pointer* to content, never content: an

48
headroom/ccr/mcp_server.py
View file

@ -75,6 +75,30 @@ COMPRESS_TOOL_NAME = "headroom_compress"
STATS_TOOL_NAME = "headroom_stats"
READ_TOOL_NAME = "headroom_read"
# Canonical schema for the retrieve tool. Single source of truth: the live
# ``list_tools()`` handler builds its ``Tool`` from these, and ``wrap agy``
# serialises them into agy's per-tool cache so the tool is exposed on the first
# run (see ``_setup_headroom_retrieve_mcp_agy``). Keeping both off one
# definition stops the primed cache from drifting from what the server offers.
CCR_RETRIEVE_TOOL_DESCRIPTION = (
"Retrieve original uncompressed content by hash. This is the ONLY "
"tool that expands Headroom compression markers — use it (not any "
"other retrieve/expand tool) whenever you see a marker containing "
"'hash=', including '[N items compressed... hash=abc123]' and "
"'[functionResponse compressed. Call headroom_retrieve to expand. "
"Retrieve more: hash=...]'. The hash is the value after 'hash='."
)
CCR_RETRIEVE_TOOL_INPUT_SCHEMA: dict[str, Any] = {
"type": "object",
"properties": {
"hash": {
"type": "string",
"description": "Hash key from compression (e.g., 'abc123' from hash=abc123)",
},
},
"required": ["hash"],
}
logger = logging.getLogger("headroom.ccr.mcp")
# Feature flag: enable headroom_read tool (file read caching via CCR)
@ -499,7 +523,11 @@ class HeadroomMCPServer:
try:
result = await self._retrieve_via_proxy(hash_key)
if "error" not in result:
result["source"] = "proxy"
# headroom-8tm WU-2b: `source` as the LEADING key (before the
# large `original_content`) so the agy FR compressor's
# content-based envelope exemption anchors survive truncation.
# Key-order only -- same keys/values.
result = {"source": "proxy", **result}
self._stats.record_retrieval(hash_key)
return result
except Exception:
@ -638,22 +666,8 @@ class HeadroomMCPServer:
),
Tool(
name=CCR_TOOL_NAME,
description=(
"Retrieve original uncompressed content by hash. "
"Use this when you need full details from previously compressed content. "
"The hash comes from headroom_compress results or from compression "
"markers like [N items compressed... hash=abc123]."
),
inputSchema={
"type": "object",
"properties": {
"hash": {
"type": "string",
"description": "Hash key from compression (e.g., 'abc123' from hash=abc123)",
},
},
"required": ["hash"],
},
description=CCR_RETRIEVE_TOOL_DESCRIPTION,
inputSchema=CCR_RETRIEVE_TOOL_INPUT_SCHEMA,
),
Tool(
name=STATS_TOOL_NAME,

15
headroom/ccr/tool_injection.py
View file

@ -22,6 +22,21 @@ from typing import Any, Protocol, runtime_checkable
CCR_TOOL_NAME = "headroom_retrieve"
def is_headroom_retrieve_name(name: object) -> bool:
"""True if a tool name is the headroom_retrieve tool.
Matches the bare name or an MCP-namespaced ``*__headroom_retrieve``
suffix (e.g. ``mcp__headroom__headroom_retrieve``). A single trailing
``retrieve`` fragment without the ``__`` boundary (e.g.
``xheadroom_retrieve``) does NOT match -- only an exact bare name or a
proper namespaced suffix does.
``name`` may come from untrusted request JSON; a non-str value (e.g.
int) would raise on ``.endswith``, so this guards with ``isinstance``.
"""
return isinstance(name, str) and (name == CCR_TOOL_NAME or name.endswith(f"__{CCR_TOOL_NAME}"))
@runtime_checkable
class _HashOwnershipStore(Protocol):
"""Structural type for verify_ownership()'s store dependency.

1163
headroom/cli/wrap.py
View file

File diff suppressed because it is too large Load diff

9
headroom/evals/session_probes.py
View file

@ -35,13 +35,14 @@ from pathlib import Path
from typing import Any
from headroom.learn.scanner import is_error_content
from headroom.parser import CCR_MARKER_ALTERNATION
DIMENSIONS = ("numerics", "artifacts", "errors")
# Mirrors the marker shapes matched by
# headroom.transforms.compression_units._CCR_MARKER_RE (kept local so the
# evals layer does not depend on a private transforms symbol).
_CCR_MARKER_RE = re.compile(r"Retrieve more: hash=|Retrieve original: hash=|<<ccr:[^>]+>>")
# Canonical CCR retrieval-marker alternation, shared with
# headroom.parser.CCR_RETRIEVAL_MARKER_RE and
# headroom.transforms.compression_units._CCR_MARKER_RE.
_CCR_MARKER_RE = re.compile(CCR_MARKER_ALTERNATION)
# A number with its immediate key context ("retry_limit: 3", "port=8787",
# JSON's '"latency_ms": 12'). Bare numbers are skipped: without context they

2
headroom/mcp_registry/__init__.py
View file

@ -13,6 +13,7 @@ without changing the calling code.
from __future__ import annotations
from .agy import AgyRegistrar
from .base import MCPRegistrar, RegisterResult, RegisterStatus, ServerSpec
from .claude import ClaudeConfigMutationError, ClaudeRegistrar
from .codex import CodexRegistrar
@ -31,6 +32,7 @@ from .server_json import build_server_json, render_server_json
__all__ = [
"DEFAULT_PROXY_URL",
"AgyRegistrar",
"CLAUDE_SERENA_CONTEXT",
"ClaudeConfigMutationError",
"ClaudeRegistrar",

265
headroom/mcp_registry/agy.py Normal file
View file

@ -0,0 +1,265 @@
"""Antigravity CLI (agy) MCP registrar.
agy 1.1.x reads MCP server configuration from the global, IDE-shared
``~/.gemini/config/mcp_config.json`` (migrated from the legacy
``~/.gemini/antigravity-cli/mcp_config.json``, which 1.1.x no longer reads —
google-antigravity/antigravity-cli#60), using the same JSON shape as Claude
Code's file path:
{"mcpServers": {"<name>": {"command": ..., "args": ..., "env": {...}}}}
There is no general-purpose CLI for editing this file, so we read/write the
JSON directly. We do NOT use marker blocks here (unlike codex.py) because the
JSON format does not admit inline comments; instead we operate on the
``mcpServers`` dict directly — adding a key to register and deleting it to
unregister — which is both safe and merge-friendly (preserves other user and
Antigravity-IDE entries untouched).
"""
from __future__ import annotations
import json
import logging
from pathlib import Path
from typing import Any
from .base import MCPRegistrar, RegisterResult, RegisterStatus, ServerSpec
logger = logging.getLogger(__name__)
#: Config file path relative to home, matching agy's own lookup.
#: agy 1.1.x MIGRATED the MCP-server read-path here from the legacy
#: ``.gemini/antigravity-cli/mcp_config.json`` (which 1.1.x no longer reads —
#: google-antigravity/antigravity-cli#60). This global config is SHARED with the
#: Antigravity IDE, so registration MUST stay merge-not-clobber.
_AGY_CONFIG_RELPATH = ".gemini/config/mcp_config.json"
#: agy's app-data directory (relative to home). Independent of the config file:
#: agy writes its per-tool cache to ``<appdata>/mcp/<server>/<tool>.json``
#: REGARDLESS of which config file declared the server (cli.log: appDataDir).
_AGY_APPDATA_RELPATH = ".gemini/antigravity-cli"
class AgyRegistrar(MCPRegistrar):
"""Register MCP servers with the Antigravity CLI (agy)."""
name = "agy"
display_name = "Antigravity CLI"
def __init__(self, *, home_dir: Path | None = None) -> None:
"""Allow ``home_dir`` override for testing (mirrors codex.py seam).
Pass ``home_dir`` in tests to redirect all file I/O to a tmp path so
the real ``~/.gemini`` is never touched.
"""
home = home_dir if home_dir is not None else Path.home()
self._config_file: Path = home / _AGY_CONFIG_RELPATH
self._appdata_dir: Path = home / _AGY_APPDATA_RELPATH
@property
def config_dir(self) -> Path:
"""Directory holding ``mcp_config.json`` (the read-path config)."""
return self._config_file.parent
@property
def cache_dir(self) -> Path:
"""agy's per-tool cache root: ``<appdata>/mcp``.
agy writes ``<appdata>/mcp/<server>/<tool>.json`` when it connects a
server — the presence of that file is the real exposure signal. This is
under the app-data dir, NOT the (migrated) config dir, so it is exposed
separately from ``config_dir``. Derived from the single ``home_dir`` seam.
"""
return self._appdata_dir / "mcp"
# ------------------------------------------------------------------
# MCPRegistrar interface
# ------------------------------------------------------------------
def detect(self) -> bool:
"""Return True if agy appears to be installed.
We key on agy's app-data directory (``~/.gemini/antigravity-cli``) — the
stable install marker — rather than the config dir, because the migrated
config dir (``~/.gemini/config``) may not exist until the first server is
written. A pre-existing config file is also accepted. We deliberately
do NOT shell out to ``shutil.which("agy")`` — the registrar is also used
in test environments where the CLI may not be on PATH.
"""
return self._appdata_dir.exists() or self._config_file.exists()
def get_server(self, server_name: str) -> ServerSpec | None:
"""Return the registered ServerSpec for ``server_name``, or ``None``."""
entry = _read_json(self._config_file).get("mcpServers", {}).get(server_name)
if not isinstance(entry, dict):
return None
return _entry_to_spec(server_name, entry)
def register_server(self, spec: ServerSpec, *, force: bool = False) -> RegisterResult:
"""Idempotently register an MCP server.
Semantics mirror claude.py's file-path path:
* Already present and matches → ALREADY.
* Already present, different, no ``force`` → MISMATCH (no clobber).
* Already present, different, ``force=True`` → overwrite → REGISTERED.
* Absent → write → REGISTERED.
In all cases, only ``spec.name`` is touched; all other ``mcpServers``
entries are preserved (merge-not-clobber).
"""
existing = self.get_server(spec.name)
if existing is not None:
if _specs_equivalent(existing, spec):
return RegisterResult(RegisterStatus.ALREADY, "matches current configuration")
if not force:
return RegisterResult(RegisterStatus.MISMATCH, _diff_specs(existing, spec))
# force=True: fall through and overwrite below.
return self._write_entry(spec)
def unregister_server(self, server_name: str) -> bool:
"""Remove ``server_name`` from the config; preserves all other entries.
Returns ``True`` on success, ``False`` if the server was absent or the
file could not be read/written.
"""
if not self._config_file.exists():
return False
try:
config = _read_json_for_write(self._config_file)
except (_MalformedConfigError, OSError) as exc:
logger.debug("agy: refusing to rewrite %s: %s", self._config_file, exc)
return False
servers: dict[str, Any] = config.get("mcpServers", {})
if server_name not in servers:
return False
del servers[server_name]
config["mcpServers"] = servers
try:
_write_json(self._config_file, config)
except OSError as exc:
logger.debug("agy: could not write %s: %s", self._config_file, exc)
return False
return True
# ------------------------------------------------------------------
# Private helpers
# ------------------------------------------------------------------
def _write_entry(self, spec: ServerSpec) -> RegisterResult:
try:
config = _read_json_for_write(self._config_file)
except (_MalformedConfigError, OSError) as exc:
return RegisterResult(
RegisterStatus.FAILED,
f"refusing to overwrite {self._config_file}: {exc}",
)
servers: dict[str, Any] = config.setdefault("mcpServers", {})
servers[spec.name] = _spec_to_entry(spec)
try:
_write_json(self._config_file, config)
except OSError as exc:
return RegisterResult(
RegisterStatus.FAILED, f"could not write {self._config_file}: {exc}"
)
return RegisterResult(RegisterStatus.REGISTERED, f"wrote {self._config_file}")
# ----------------------------------------------------------------------
# JSON helpers (private to this module; do NOT import from claude.py)
# ----------------------------------------------------------------------
class _MalformedConfigError(RuntimeError):
"""Raised when an existing config cannot be parsed before a full rewrite."""
def _read_json(path: Path) -> dict[str, Any]:
"""Read JSON file, returning empty dict if absent or unparseable.
Safe for READ-ONLY callers. Do NOT use before a full-file rewrite: an
unparseable file returns ``{}`` here, and writing that back destroys the
user's config. Use :func:`_read_json_for_write` on the write path instead.
"""
if not path.exists():
return {}
try:
with open(path, encoding="utf-8") as f:
data = json.load(f)
except (OSError, json.JSONDecodeError):
return {}
if not isinstance(data, dict):
return {}
return data
def _read_json_for_write(path: Path) -> dict[str, Any]:
"""Read a JSON object ahead of a full-file rewrite.
Returns ``{}`` only when the file is absent or empty (safe to create fresh).
When it has content that is not a JSON object, raises
:class:`_MalformedConfigError` so the caller aborts instead of overwriting an
unrelated user config — agy's ``mcp_config.json`` is shared with the
Antigravity IDE and holds the user's own servers alongside ours.
"""
if not path.exists():
return {}
raw = path.read_text(encoding="utf-8") # OSError propagates to the caller
if not raw.strip():
return {}
try:
data = json.loads(raw)
except json.JSONDecodeError as exc:
raise _MalformedConfigError(f"{path} is not valid JSON: {exc}") from exc
if not isinstance(data, dict):
raise _MalformedConfigError(f"{path} does not contain a JSON object")
return data
def _write_json(path: Path, data: dict[str, Any]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
json.dump(data, f, indent=2)
f.write("\n")
def _spec_to_entry(spec: ServerSpec) -> dict[str, Any]:
entry: dict[str, Any] = {"command": spec.command}
if spec.args:
entry["args"] = list(spec.args)
if spec.env:
entry["env"] = dict(spec.env)
return entry
def _entry_to_spec(name: str, entry: dict[str, Any]) -> ServerSpec:
return ServerSpec(
name=name,
command=str(entry.get("command", "")),
args=tuple(entry.get("args", ())),
env=dict(entry.get("env", {})),
)
def _specs_equivalent(a: ServerSpec, b: ServerSpec) -> bool:
return (
a.name == b.name
and a.command == b.command
and tuple(a.args) == tuple(b.args)
and dict(a.env) == dict(b.env)
)
def _diff_specs(existing: ServerSpec, requested: ServerSpec) -> str:
parts: list[str] = []
if existing.command != requested.command:
parts.append(f"command {existing.command!r} -> {requested.command!r}")
if tuple(existing.args) != tuple(requested.args):
parts.append(f"args {list(existing.args)} -> {list(requested.args)}")
if dict(existing.env) != dict(requested.env):
parts.append(f"env {dict(existing.env)} -> {dict(requested.env)}")
if not parts:
return "spec differs in unidentified field(s)"
return "; ".join(parts)

9
headroom/mcp_registry/install.py
View file

@ -6,6 +6,7 @@ from collections.abc import Iterable
from headroom.install.runtime import resolve_headroom_command
from .agy import AgyRegistrar
from .base import MCPRegistrar, RegisterResult, RegisterStatus, ServerSpec
from .claude import ClaudeRegistrar
from .codex import CodexRegistrar
@ -22,7 +23,13 @@ def get_all_registrars() -> list[MCPRegistrar]:
The list grows as we add adapters for Cursor, Continue, Cline, etc.
"""
return [ClaudeRegistrar(), CodexRegistrar(), GrokRegistrar(), OpencodeRegistrar()]
return [
ClaudeRegistrar(),
CodexRegistrar(),
AgyRegistrar(),
GrokRegistrar(),
OpencodeRegistrar(),
]
def build_headroom_spec(proxy_url: str = DEFAULT_PROXY_URL) -> ServerSpec:

10
headroom/parser.py
View file

@ -24,10 +24,12 @@ JSON_BLOCK_PATTERN = re.compile(r"\{[\s\S]{500,}\}")
# exit codes) and are not evidence of a re-read.
REREAD_MIN_TOKENS = 50
# Canonical CCR retrieval-marker shapes. Mirrors the alternation in
# transforms/compression_units._CCR_MARKER_RE; kept local because the parser
# is a base module and importing from transforms would create a cycle.
CCR_RETRIEVAL_MARKER_RE = re.compile(r"Retrieve more: hash=|Retrieve original: hash=|<<ccr:[^>]+>>")
# Canonical CCR retrieval-marker shapes. parser is a base module (content_router.py
# already imports from it), so this alternation is defined here and re-exported
# for transforms/compression_units.py and evals/session_probes.py to import,
# rather than kept as byte-identical local copies.
CCR_MARKER_ALTERNATION = r"Retrieve more: hash=|Retrieve original: hash=|<<ccr:[^>]+>>"
CCR_RETRIEVAL_MARKER_RE = re.compile(CCR_MARKER_ALTERNATION)
# Repeats this close (in message positions) to the previous serve are
# polling, not re-reads. Consecutive tool turns sit 2 apart (the

5
headroom/providers/agy/__init__.py Normal file
View file

@ -0,0 +1,5 @@
"""agy-specific provider helpers."""
from .runtime import build_agy_env
__all__ = ["build_agy_env"]

83
headroom/providers/agy/runtime.py Normal file
View file

@ -0,0 +1,83 @@
"""Runtime env builder for agy-specific MITM proxy wiring.
Pure data transform: given a terminator URL and a CA trust bundle path,
produce the child environment dict that routes agy through the CONNECT
terminator while trusting the minted CA bundle.
No side effects; no I/O; no subprocess.
"""
from __future__ import annotations
from pathlib import Path
def build_agy_env(
*,
terminator_url: str,
bundle_path: Path,
base_env: dict[str, str],
) -> dict[str, str]:
"""Return a new env dict suitable for launching agy through the MITM terminator.
Parameters
----------
terminator_url:
Full HTTP URL of the AgyCONNECTTerminator (e.g. ``http://127.0.0.1:<port>``).
bundle_path:
Path to the combined CA trust bundle produced by
``headroom.proxy.agy_ca.build_combined_bundle``. Set in all three
trust-bundle env vars so Python, Node.js, and curl all see it.
base_env:
Base environment (typically ``os.environ.copy()``). A fresh copy is
returned — ``base_env`` is never mutated.
Returns
-------
dict[str, str]
New environment dict with proxy and CA vars wired for agy.
Notes
-----
Corporate proxy chaining works without any extra plumbing here: this
function returns a COPY and never mutates ``base_env`` or
``os.environ``. The CONNECT terminator runs in the PARENT process and
therefore still reads the original corporate ``os.environ["HTTPS_PROXY"]``
when it chains non-allowlisted CONNECTs upstream
(see ``agy_terminator.py:_handle_blind_tunnel``). Only the CHILD agy
process receives ``HTTPS_PROXY=terminator_url`` so that all of its
traffic is routed into the terminator first.
"""
bundle_str = str(bundle_path)
env = dict(base_env) # copy — never mutate caller's dict
# Drop the wrapper's session-scoped savings redirection. Those vars belong to
# THIS process's in-proxy funnel (savings written to a temp dir that is
# deleted on exit, inbox-emit marker set); the agy child — and the
# `headroom mcp serve` grandchild it spawns, which is headroom code — must
# not inherit them and write into a sink that disappears.
for leaked in (
"HEADROOM_AGY_INBOX_EMIT",
"HEADROOM_SAVINGS_PATH",
"HEADROOM_SAVINGS_EVENTS_PATH",
"HEADROOM_OTEL_METRICS_ENABLED",
):
env.pop(leaked, None)
# Route all traffic through the CONNECT terminator.
env["HTTPS_PROXY"] = terminator_url
env["HTTP_PROXY"] = terminator_url
# Extend, never replace: on a corporate machine the inherited NO_PROXY names
# hosts that MUST bypass the proxy, and dropping them would tunnel them
# through the terminator.
inherited_no_proxy = (env.get("NO_PROXY") or env.get("no_proxy") or "").strip().strip(",")
env["NO_PROXY"] = (
f"127.0.0.1,localhost,{inherited_no_proxy}" if inherited_no_proxy else "127.0.0.1,localhost"
)
# Trust our minted CA bundle — blanking these would break MITM.
env["SSL_CERT_FILE"] = bundle_str
env["CACERT_PATH"] = bundle_str
env["NODE_EXTRA_CA_CERTS"] = bundle_str
return env

227
headroom/providers/agy/stats.py Normal file
View file

@ -0,0 +1,227 @@
"""Agy-session compression observability helpers.
Thread-safe, agy-scoped ONLY. No imports from gemini.py, transport, or
compression_store — those are imported lazily at call time.
Public surface
--------------
FailOpenWarnHandler logging.Handler that emits a one-time stderr notice on
the first Cloud-Code-Assist fail-open log record.
AgySessionStats Snapshot + delta + summary formatting; idempotent print.
Ref: headroom-30y.15
"""
from __future__ import annotations
import logging
import sys
import threading
from typing import Any
# The logger that gemini.py actually emits the fail-open warning on.
# CONFIRMED: gemini.py:25 is `logger = logging.getLogger("headroom.proxy")` and
# the fail-open warning at gemini.py:883 uses that logger. Python logging
# propagates child->parent (NOT parent->child), so a handler on the *child*
# "headroom.proxy.handlers.gemini" would NEVER receive these records — we must
# install on the actual emitting logger "headroom.proxy".
_GEMINI_LOGGER = "headroom.proxy"
# Substring that identifies the fail-open warning record (gemini.py:883). Used
# to filter out unrelated "headroom.proxy" warnings.
_FAIL_OPEN_SUBSTR = "Cloud Code Assist optimization failed"
class FailOpenWarnHandler(logging.Handler):
"""One-shot logging.Handler that prints a user-facing stderr notice on the
FIRST fail-open compression warning emitted by the gemini handler, then
counts all subsequent occurrences.
Install on the ``"headroom.proxy"`` logger (the logger gemini.py actually
emits the fail-open warning on) before launching agy. Remove in finally to
avoid leaking into the click process.
Thread-safe: the one-shot flag and counter use a single lock.
"""
def __init__(self) -> None:
super().__init__(level=logging.WARNING)
self._lock: threading.Lock = threading.Lock()
self._warned: bool = False
self._count: int = 0
# ------------------------------------------------------------------
# logging.Handler interface
# ------------------------------------------------------------------
def emit(self, record: logging.LogRecord) -> None:
if _FAIL_OPEN_SUBSTR not in record.getMessage():
return
with self._lock:
self._count += 1
if self._warned:
return
self._warned = True
# Print outside the lock to avoid holding it during I/O.
print(
"Headroom: compression failed for a request; forwarding it uncompressed"
" (fail-open). Further occurrences are summarized at exit.",
file=sys.stderr,
)
# ------------------------------------------------------------------
# Accessors (called from the main thread after agy exits)
# ------------------------------------------------------------------
@property
def fail_open_count(self) -> int:
"""Total number of fail-open log records observed (thread-safe)."""
with self._lock:
return self._count
def _get_compression_stats() -> dict[str, Any]:
"""Return get_compression_store().get_stats() — imported lazily so the
compression stack is not pulled in unless actually called."""
from headroom.cache.compression_store import get_compression_store
return get_compression_store().get_stats()
class AgySessionStats:
"""Snapshot start/end compression-store stats, format a one-line summary.
Usage::
stats = AgySessionStats() # call at session start (before agy)
stats.snapshot_start()
# ... agy runs ...
stats.print_summary(handler) # call in finally / SIGTERM handler
The summary is idempotent: ``print_summary`` prints exactly once regardless
of how many times it is called (safe for both the ``finally`` path and the
SIGTERM handler running close together).
"""
def __init__(self) -> None:
self._lock: threading.Lock = threading.Lock()
self._start: dict[str, Any] | None = None
self._printed: bool = False
def snapshot_start(self) -> None:
"""Capture the compression-store baseline before agy launches.
Best-effort: if the store is unavailable the snapshot is omitted and
``print_summary`` will emit a reduced message.
"""
try:
snap = _get_compression_stats()
except Exception: # noqa: BLE001
snap = None
with self._lock:
self._start = snap
def print_summary(self, handler: FailOpenWarnHandler | None = None) -> None:
"""Print a one-line session compression summary to stderr.
Idempotent: prints at most once per ``AgySessionStats`` instance.
Safe to call from both the ``finally`` block and the SIGTERM handler.
Args:
handler: The ``FailOpenWarnHandler`` installed for this session, or
``None`` if it was not installed (fail-open count omitted).
"""
with self._lock:
if self._printed:
return
self._printed = True
start = self._start
# Snapshot end outside the lock (I/O + potential lock in store).
try:
end = _get_compression_stats()
except Exception: # noqa: BLE001
end = None
fail_open = handler.fail_open_count if handler is not None else None
summary = _format_summary(start, end, fail_open_count=fail_open)
print(summary, file=sys.stderr)
def _format_summary(
start: dict[str, Any] | None,
end: dict[str, Any] | None,
*,
fail_open_count: int | None = None,
) -> str:
"""Format a session compression summary string.
Pure function for testability — no I/O side-effects.
Args:
start: ``get_stats()`` snapshot taken before the session.
end: ``get_stats()`` snapshot taken after the session.
fail_open_count: Number of fail-open warnings observed, or ``None``
when the handler was not installed.
Returns:
A single-line string suitable for printing to stderr.
"""
if start is None or end is None:
fail_suffix = (
f" Fail-open requests: {fail_open_count}" if fail_open_count is not None else ""
)
return f"Headroom agy session summary: compression stats unavailable.{fail_suffix}"
entries = max(0, end.get("entry_count", 0) - start.get("entry_count", 0))
orig = max(0, end.get("total_original_tokens", 0) - start.get("total_original_tokens", 0))
comp = max(0, end.get("total_compressed_tokens", 0) - start.get("total_compressed_tokens", 0))
# Report the share of the original that survived. "0.30x ratio" alone reads
# like a 30% expansion; "30% of original" cannot be misread.
if orig > 0:
ratio_str = f"{comp / orig:.0%} of original"
else:
ratio_str = "n/a (no compression)"
parts = [
f"Headroom agy session: {entries} entries compressed,",
f"{orig:,} → {comp:,} tokens ({ratio_str})",
]
if fail_open_count is not None:
parts.append(f"| {fail_open_count} fail-open request(s)")
return " ".join(parts)
def install_fail_open_handler() -> FailOpenWarnHandler:
"""Install a ``FailOpenWarnHandler`` on the gemini proxy logger.
Returns the installed handler so the caller can:
- read ``.fail_open_count`` after agy exits
- pass it to ``remove_fail_open_handler`` in finally
Safe to call multiple times (each call installs a fresh handler; the old
one is not removed — call ``remove_fail_open_handler`` explicitly).
"""
handler = FailOpenWarnHandler()
# Target "headroom.proxy" — the logger gemini.py:25 actually uses to emit the
# fail-open warning. emit() filters on _FAIL_OPEN_SUBSTR so unrelated
# headroom.proxy warnings are ignored.
logging.getLogger(_GEMINI_LOGGER).addHandler(handler)
return handler
def remove_fail_open_handler(handler: FailOpenWarnHandler | None) -> None:
"""Remove a previously-installed ``FailOpenWarnHandler`` (best-effort).
Called in the ``finally`` block of ``agy()`` to avoid leaking the handler
into the click process or subsequent agent invocations. Idempotent and
exception-safe. Accepts ``None`` to simplify callers that may not have
installed the handler (e.g. if an exception was raised before install).
"""
if handler is None:
return
try:
logging.getLogger(_GEMINI_LOGGER).removeHandler(handler)
except Exception: # noqa: BLE001
pass

7
headroom/providers/proxy_routes.py
View file

@ -38,6 +38,9 @@ from headroom.providers.openai_responses import (
from headroom.providers.proxy_targets import (
api_target as _api_target,
)
from headroom.providers.proxy_targets import (
cloudcode_host_base as _cloudcode_host_base,
)
from headroom.providers.proxy_targets import (
select_passthrough_base_url as _select_passthrough_base_url,
)
@ -537,10 +540,12 @@ def register_provider_routes(app: FastAPI, proxy: Any) -> None:
normalized_cloudcode_path = normalize_cloudcode_passthrough_path(path)
if normalized_cloudcode_path is not None:
normalize_request_path(request, normalized_cloudcode_path)
host = request.headers.get("host", "")
cloudcode_base = _cloudcode_host_base(host) or _api_target(proxy, "cloudcode")
return await proxy.handle_passthrough(
request,
_api_target(proxy, "cloudcode"),
cloudcode_base,
)
return await proxy.handle_passthrough(

19
headroom/providers/proxy_targets.py
View file

@ -14,6 +14,7 @@ from headroom.copilot_auth import (
from headroom.providers.codex import resolve_codex_routing
from headroom.providers.codex.endpoints import CHATGPT_BACKEND_API_URL
from headroom.providers.vertex import vertex_target_for_location as _vertex_target_for_location
from headroom.proxy.agy_terminator import DEFAULT_ALLOWLIST, normalize_host
from headroom.proxy.upstream_guard import is_safe_upstream_url
LEGACY_API_TARGET_ATTRS: dict[str, str] = {
@ -36,6 +37,22 @@ def vertex_target_for_location(proxy: Any, location: str) -> str:
return _vertex_target_for_location(api_target(proxy, "vertex"), location)
def cloudcode_host_base(host: str) -> str | None:
"""Passthrough base for an allowlisted Cloud Code host, else ``None``.
agy (Google Antigravity CLI) reaches the proxy via TLS-MITM that terminates
the WHOLE connection to the Cloud Code host it addressed, so control-plane
calls land on the catch-all rather than a recognized route. Those paths
exist only on the Cloud Code host itself; forward them back to it.
Membership in ``DEFAULT_ALLOWLIST`` — not a loose suffix match — is the trust
boundary: a forged Host such as ``evilcloudcode-pa.googleapis.com`` returns
``None`` (closing the SSRF) so the caller falls back to the configured
default.
"""
normalized = normalize_host(host)
return f"https://{normalized}" if normalized in DEFAULT_ALLOWLIST else None
logger = logging.getLogger("headroom.proxy")
@ -43,6 +60,8 @@ def select_passthrough_base_url(
proxy: Any, headers: Mapping[str, str], path: str | None = None
) -> str:
"""Resolve the upstream base URL for catch-all proxy passthrough requests."""
if base := cloudcode_host_base(headers.get("host", "")):
return base
routing = resolve_codex_routing(headers)
if routing.is_chatgpt_auth:
return CHATGPT_BACKEND_API_URL

593
headroom/proxy/agy_ca.py Normal file
View file

@ -0,0 +1,593 @@
"""Root CA lifecycle + combined trust bundle for the agy MITM transport.
Process-scoped: the CA is generated once and persisted under
``~/.headroom/ca/`` (or an injectable base dir for tests). The combined
bundle (system CAs + headroom root CA + filtered corporate CAs) is written
to ``~/.headroom/combined-ca-bundle.pem`` with strict permissions and is
intended for injection into the wrapped agy process via environment
variables (CACERT_PATH / SSL_CERT_FILE / NODE_EXTRA_CA_CERTS).
Security invariants (enforced by assertion):
- CA private key: 0600, parent dir: 0700.
- Combined bundle: 0600, parent dir: 0700.
- CA is NEVER written to any OS trust-store path.
- Only PEM objects with basicConstraints CA:TRUE are included from
corporate CA files (per-object parse-then-filter).
"""
from __future__ import annotations
import datetime
import logging
import os
import ssl
import stat
import sys
import tempfile
from collections.abc import Sequence
from pathlib import Path
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
from cryptography.x509 import Certificate
from cryptography.x509.oid import NameOID
logger = logging.getLogger("headroom.proxy.agy_ca")
# CA validity: 10 years; regeneration triggers when less than 30 days remain.
_CA_VALIDITY_DAYS = 3650
_REGEN_THRESHOLD_DAYS = 30
# Key size for the root CA.
_RSA_KEY_BITS = 4096
# Well-known OS trust store paths — CA must never be written here.
_OS_TRUST_PATHS: tuple[str, ...] = (
"/etc/ssl/certs",
"/etc/pki/ca-trust",
"/usr/local/share/ca-certificates",
"/etc/ca-certificates",
"/usr/share/ca-certificates",
"/System/Library/Keychains",
"/Library/Keychains",
)
# Candidate system CA bundle paths (ordered by prevalence).
_SYSTEM_BUNDLE_CANDIDATES: tuple[str, ...] = (
"/etc/ssl/certs/ca-certificates.crt", # Debian/Ubuntu/Alpine
"/etc/pki/tls/certs/ca-bundle.crt", # RHEL/CentOS/Fedora
"/etc/ssl/ca-bundle.pem", # openSUSE
"/usr/share/ssl/certs/ca-bundle.crt", # legacy RHEL
"/usr/local/etc/openssl/cert.pem", # macOS Homebrew OpenSSL
"/etc/ssl/cert.pem", # macOS system / BSDs
"/usr/local/share/certs/ca-root-nss.crt", # FreeBSD
"/etc/pki/tls/cacert.pem", # older RHEL
)
# Environment variables that may point at a corporate CA bundle.
_CORP_CA_ENV_VARS: tuple[str, ...] = ("SSL_CERT_FILE", "NODE_EXTRA_CA_CERTS")
# File names under the CA directory.
_CA_KEY_NAME = "ca.key"
_CA_CERT_NAME = "ca.crt"
_BUNDLE_NAME = "combined-ca-bundle.pem"
# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------
def _assert_perms(path: Path, expected_mode: int) -> None:
"""Raise PermissionError if *path* does not have exactly *expected_mode* bits.
No-op on non-POSIX platforms (Windows) where mode bits are not meaningful.
"""
if os.name != "posix":
return
actual = stat.S_IMODE(path.stat().st_mode)
if actual != expected_mode:
raise PermissionError(
f"Permission check failed for {path}: expected {oct(expected_mode)}, got {oct(actual)}"
)
def _secure_dir(path: Path) -> None:
"""Create *path* with 0700 if absent; enforce 0700 on return.
``parents=True`` only applies the mode to the leaf directory on some
platforms — intermediate parents get the umask-filtered mode. We
therefore chmod the leaf explicitly after mkdir so pre-existing or
newly-created paths are always corrected to 0700.
"""
path.mkdir(mode=0o700, parents=True, exist_ok=True)
os.chmod(path, 0o700)
_assert_perms(path, 0o700)
def _write_secure(path: Path, data: bytes) -> None:
"""Write *data* to *path* atomically with 0600; assert afterwards.
``mkstemp`` gives a fresh, exclusively-created name in the target directory,
opened 0600 and in binary mode (so Windows never translates ``\n`` ->
``\r\n`` and corrupts the PEM bytes). There is no world-readable window.
A fixed ``<stem>.tmp`` would collide two ways — ``ca.key`` and ``ca.crt``
both map to ``ca.tmp``, and two concurrent ``wrap agy`` runs share it — and
``O_CREAT`` without ``O_EXCL`` silently keeps an existing file's mode, so a
leftover 0644 temp would carry the private key.
"""
fd, tmp_name = tempfile.mkstemp(dir=str(path.parent), prefix=f".{path.name}.", suffix=".tmp")
tmp = Path(tmp_name)
try:
os.write(fd, data)
finally:
os.close(fd)
os.replace(tmp, path)
_assert_perms(path, 0o600)
def _not_in_os_trust(path: Path) -> None:
"""Raise RuntimeError if *path* resides under any known OS trust location."""
resolved = path.resolve()
for trust_path in _OS_TRUST_PATHS:
# Path-component comparison, not a string prefix: /etc/ssl/certs-mine is
# not inside /etc/ssl/certs.
if resolved == Path(trust_path) or resolved.is_relative_to(trust_path):
raise RuntimeError(
f"CA file {path} resolves to {resolved}, which is inside OS trust path {trust_path}"
)
def _now_utc() -> datetime.datetime:
return datetime.datetime.now(tz=datetime.timezone.utc)
# ---------------------------------------------------------------------------
# CA generation
# ---------------------------------------------------------------------------
def _generate_root_ca() -> tuple[RSAPrivateKey, Certificate]:
"""Generate a new RSA root CA key + self-signed certificate."""
key: RSAPrivateKey = rsa.generate_private_key(
public_exponent=65537,
key_size=_RSA_KEY_BITS,
)
now = _now_utc()
subject = issuer = x509.Name(
[
x509.NameAttribute(NameOID.COMMON_NAME, "Headroom Local CA"),
x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Headroom MITM"),
]
)
cert = (
x509.CertificateBuilder()
.subject_name(subject)
.issuer_name(issuer)
.public_key(key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(now)
.not_valid_after(now + datetime.timedelta(days=_CA_VALIDITY_DAYS))
.add_extension(
x509.BasicConstraints(ca=True, path_length=0),
critical=True,
)
.add_extension(
x509.KeyUsage(
digital_signature=True,
key_cert_sign=True,
crl_sign=True,
content_commitment=False,
key_encipherment=False,
data_encipherment=False,
key_agreement=False,
encipher_only=False,
decipher_only=False,
),
critical=True,
)
.add_extension(
x509.SubjectKeyIdentifier.from_public_key(key.public_key()),
critical=False,
)
.sign(key, hashes.SHA256())
)
return key, cert
def _is_ca_cert(cert: Certificate) -> bool:
"""Return True iff the certificate has basicConstraints CA:TRUE."""
try:
bc = cert.extensions.get_extension_for_class(x509.BasicConstraints)
return bool(bc.value.ca)
except x509.ExtensionNotFound:
return False
def _cert_near_expiry(cert: Certificate) -> bool:
"""Return True if the certificate expires within the regen threshold."""
threshold = _now_utc() + datetime.timedelta(days=_REGEN_THRESHOLD_DAYS)
return bool(cert.not_valid_after_utc <= threshold)
# ---------------------------------------------------------------------------
# Bundle helpers
# ---------------------------------------------------------------------------
def _detect_system_bundle() -> Path:
"""Return path to the system CA bundle; raise RuntimeError if not found."""
for candidate in _SYSTEM_BUNDLE_CANDIDATES:
p = Path(candidate)
if p.is_file() and p.stat().st_size > 0:
logger.debug("event=system_bundle_found path=%s", p)
return p
raise RuntimeError(
"No system CA bundle found. Searched: " + ", ".join(_SYSTEM_BUNDLE_CANDIDATES)
)
def _windows_trust_pem() -> bytes:
"""Collect CA:TRUE certs from the Windows system trust stores as PEM bytes.
Windows has no single on-disk CA bundle file, so ``_SYSTEM_BUNDLE_CANDIDATES``
never matches there. ``ssl.enum_certificates`` (Windows-only) enumerates the
ROOT and CA stores but returns *all* certs including leaf certs, so the
result is run through the same ``_parse_ca_certs_from_pem`` CA:TRUE filter
used for corporate bundles — never trust a non-CA cert as an anchor.
"""
blocks: list[bytes] = []
try:
for store in ("ROOT", "CA"):
for der, _enc, _trust in ssl.enum_certificates(store): # type: ignore[attr-defined,unused-ignore]
try:
blocks.append(ssl.DER_cert_to_PEM_cert(der).encode("ascii"))
except ValueError:
# A single malformed DER entry must not abort the whole store.
logger.debug("event=windows_der_skip reason=bad_der")
except OSError as exc:
# ssl.enum_certificates wraps the Win32 cert-store API; surface a clear
# cause instead of an opaque traceback at `wrap agy` launch.
raise RuntimeError("could not read the Windows system trust store") from exc
ca_pem = b"".join(_parse_ca_certs_from_pem(b"".join(blocks)))
if not ca_pem:
# An empty system-trust component would silently leave the combined
# bundle trusting only the headroom MITM root — fail loud, never ship a
# trust bundle with no system anchors.
raise RuntimeError(
"Windows system trust store yielded no CA anchors; refusing to build "
"a trust bundle with no system anchors"
)
return ca_pem
def _system_trust_pem() -> tuple[bytes, str]:
"""Return ``(system trust PEM bytes, source label)`` for this platform.
POSIX/macOS read the detected on-disk bundle; Windows enumerates the
system trust stores via stdlib ``ssl`` (no certifi dependency).
"""
# Prefer an on-disk bundle on every platform (so a corp-provided file or a
# test-injected candidate wins). Windows normally has no such file, so fall
# back to enumerating the system cert stores there.
try:
path = _detect_system_bundle()
except RuntimeError:
if sys.platform == "win32":
return _windows_trust_pem(), "windows-cert-store"
raise
return path.read_bytes(), str(path)
def _parse_ca_certs_from_pem(pem_data: bytes) -> list[bytes]:
"""Parse a multi-cert PEM file, returning PEM bytes for CA:TRUE certs only."""
results: list[bytes] = []
# Split on BEGIN CERTIFICATE boundaries; preserve header+body per cert.
parts = pem_data.split(b"-----BEGIN CERTIFICATE-----")
for part in parts[1:]: # skip leading empty fragment
pem_block = b"-----BEGIN CERTIFICATE-----" + part
# Trim trailing noise after END CERTIFICATE.
end_marker = b"-----END CERTIFICATE-----"
end_idx = pem_block.find(end_marker)
if end_idx == -1:
continue
pem_block = pem_block[: end_idx + len(end_marker)] + b"\n"
# cryptography parses lazily: load_pem_x509_certificate succeeds but
# accessing extensions (in _is_ca_cert) can still raise on a real-world
# cert with a non-strict-conformant field — e.g. some Windows ROOT-store
# certs have a UserNotice explicit_text that rust-asn1 rejects. A trust
# builder must skip such a cert, not crash, so the whole load+inspect is
# guarded.
try:
cert = x509.load_pem_x509_certificate(pem_block)
is_ca = _is_ca_cert(cert)
except Exception: # noqa: BLE001
logger.debug("event=pem_parse_skip reason=invalid_cert")
continue
if is_ca:
results.append(pem_block)
return results
def _collect_corporate_ca_pems(env_vars: Sequence[str] = _CORP_CA_ENV_VARS) -> list[bytes]:
"""
Collect CA-only PEM blocks from any pre-existing corporate CA env vars.
Reads SSL_CERT_FILE and NODE_EXTRA_CA_CERTS (if set and pointing at a
file), parses each PEM object, and retains only those with
basicConstraints CA:TRUE.
"""
ca_pems: list[bytes] = []
for var in env_vars:
path_str = os.environ.get(var)
if not path_str:
continue
p = Path(path_str)
if not p.is_file():
logger.warning("event=corp_ca_env_missing var=%s path=%r", var, path_str)
continue
data = p.read_bytes()
filtered = _parse_ca_certs_from_pem(data)
logger.info(
"event=corp_ca_loaded var=%s path=%s ca_count=%d",
var,
path_str,
len(filtered),
)
ca_pems.extend(filtered)
return ca_pems
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
def ensure_root_ca(
base_dir: Path | None = None,
) -> tuple[RSAPrivateKey, Certificate, Path, Path]:
"""Ensure the headroom root CA exists and is valid; regenerate if expired.
Parameters
----------
base_dir:
Root of the headroom state directory. Defaults to ``~/.headroom``.
Tests must pass a ``tmp_path``-derived value to avoid touching the
real home directory.
Returns
-------
(private_key, certificate, key_path, cert_path)
The in-memory key + cert objects and their on-disk paths.
"""
if base_dir is None:
base_dir = Path.home() / ".headroom"
_secure_dir(base_dir)
ca_dir = base_dir / "ca"
_secure_dir(ca_dir)
_not_in_os_trust(ca_dir)
key_path = ca_dir / _CA_KEY_NAME
cert_path = ca_dir / _CA_CERT_NAME
# --- load existing if present ---
if key_path.exists() and cert_path.exists():
_assert_perms(key_path, 0o600)
_assert_perms(cert_path, 0o600)
try:
existing_cert = x509.load_pem_x509_certificate(cert_path.read_bytes())
except Exception as exc:
logger.warning("event=ca_load_failed reason=%s; regenerating", exc)
existing_cert = None
if existing_cert is not None and not _cert_near_expiry(existing_cert):
try:
key_bytes = key_path.read_bytes()
existing_key = serialization.load_pem_private_key(key_bytes, password=None)
logger.info("event=ca_reused path=%s", cert_path)
return existing_key, existing_cert, key_path, cert_path # type: ignore[return-value]
except Exception as exc:
logger.warning("event=ca_key_load_failed reason=%s; regenerating", exc)
# Regenerate — delete stale artifacts.
logger.info("event=ca_regenerate reason=expired_or_corrupt path=%s", cert_path)
_delete_stale_artifacts(base_dir)
# --- generate fresh CA ---
key, cert = _generate_root_ca()
key_pem = key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.TraditionalOpenSSL,
encryption_algorithm=serialization.NoEncryption(),
)
cert_pem = cert.public_bytes(serialization.Encoding.PEM)
_write_secure(key_path, key_pem)
_write_secure(cert_path, cert_pem)
_assert_perms(ca_dir, 0o700)
_not_in_os_trust(key_path)
_not_in_os_trust(cert_path)
logger.info("event=ca_generated path=%s", cert_path)
return key, cert, key_path, cert_path
def _delete_stale_artifacts(base_dir: Path) -> None:
"""Remove old combined bundle and any leaf certs on CA regeneration."""
bundle = base_dir / _BUNDLE_NAME
if bundle.exists():
bundle.unlink()
logger.info("event=stale_bundle_deleted path=%s", bundle)
# Leaf certs would live under base_dir/leaves/ (T8). Delete the dir if present.
leaves_dir = base_dir / "leaves"
if leaves_dir.is_dir():
import shutil
shutil.rmtree(leaves_dir)
logger.info("event=stale_leaves_deleted path=%s", leaves_dir)
def build_combined_bundle(
base_dir: Path | None = None,
corp_env_vars: Sequence[str] = _CORP_CA_ENV_VARS,
) -> Path:
"""Build (or rebuild) the combined CA trust bundle.
Combines:
1. System CA bundle (detected cross-distro; fail-fast if absent).
2. Headroom root CA certificate.
3. Any pre-existing corporate CAs from env (CA:TRUE-only, per-object filter).
Writes to ``<base_dir>/combined-ca-bundle.pem`` with 0600 perms.
Parent dir is asserted 0700.
Parameters
----------
base_dir:
Headroom state directory. Defaults to ``~/.headroom``.
corp_env_vars:
Environment variable names to scan for corporate CA files.
Override in tests to inject fixture paths without touching the env.
Returns
-------
Path to the combined bundle.
"""
if base_dir is None:
base_dir = Path.home() / ".headroom"
_secure_dir(base_dir)
system_pem, system_source = _system_trust_pem()
_, ca_cert, _, ca_cert_path = ensure_root_ca(base_dir)
headroom_pem = ca_cert.public_bytes(serialization.Encoding.PEM)
corp_pems = _collect_corporate_ca_pems(corp_env_vars)
combined = system_pem
if not combined.endswith(b"\n"):
combined += b"\n"
combined += headroom_pem
for pem in corp_pems:
combined += pem
bundle_path = base_dir / _BUNDLE_NAME
_write_secure(bundle_path, combined)
_assert_perms(bundle_path, 0o600)
_assert_perms(base_dir, 0o700)
_not_in_os_trust(bundle_path)
logger.info(
"event=bundle_written path=%s system=%s corp_ca_count=%d",
bundle_path,
system_source,
len(corp_pems),
)
return bundle_path
# ---------------------------------------------------------------------------
# In-memory leaf cert/key loader
# ---------------------------------------------------------------------------
def load_cert_chain_in_memory(
ctx: ssl.SSLContext,
cert_pem: bytes,
key_pem: bytes,
) -> None:
"""Load *cert_pem* + *key_pem* into *ctx* without writing a persistent key file.
Leaf private keys are loaded from anonymous memory (memfd) on Linux and
never touch the filesystem; on platforms without memfd, a 0600 temp file
is written and unlinked immediately after load (perms asserted).
Primary path (Linux, ``os.memfd_create`` available):
An anonymous, unnamed in-kernel file descriptor is created via
``memfd_create``. The combined ``cert_pem + key_pem`` PEM is written
into it (looping on ``os.write`` to handle short-writes).
``load_cert_chain`` reads it through ``/proc/self/fd/{fd}``; the fd is
closed in a ``finally`` block *after* the load (the ``/proc`` path dies
the moment the fd is closed).
Fallback path (memfd absent or ``/proc`` unusable):
``tempfile.mkstemp`` creates a 0600 temp file. The combined PEM is
written in full (loop on ``os.write``). ``_assert_perms`` validates
the 0600 mode (fail-loud; no silent chmod since mkstemp already yields
0600). ``load_cert_chain`` is called; ``os.unlink`` removes the file
in a ``finally`` block even if ``load_cert_chain`` raises. Residual
risk: a hard crash (SIGKILL/OOM) during the brief load window could
orphan a 0600 temp until OS temp cleanup. This fallback only runs on
platforms without ``memfd_create``; Linux never writes the key to disk.
Parameters
----------
ctx:
Target ``ssl.SSLContext`` (must be server-side, ``PROTOCOL_TLS_SERVER``).
cert_pem:
Leaf certificate in PEM encoding.
key_pem:
Leaf private key in PEM encoding (unencrypted).
"""
combined = cert_pem + key_pem
if sys.platform == "linux" and hasattr(os, "memfd_create"):
fd = os.memfd_create("hr_leaf") # type: ignore[attr-defined]
try:
_write_all_fd(fd, combined)
ctx.load_cert_chain(f"/proc/self/fd/{fd}")
return
except (FileNotFoundError, PermissionError):
# /proc not mounted / inaccessible (some containers) — fall through
# to mkstemp. Caught narrowly ON PURPOSE: ssl.SSLError is a subclass
# of OSError, so a broad `except OSError` would swallow a malformed
# cert/key and silently disk-fall-back. Those propagate instead.
pass
finally:
os.close(fd)
_load_via_mkstemp(ctx, combined)
def _write_all_fd(fd: int, data: bytes) -> None:
"""Write all of *data* to *fd*, handling short-writes."""
view = memoryview(data)
written = 0
total = len(data)
while written < total:
n = os.write(fd, view[written:])
if n == 0:
raise OSError("os.write wrote 0 bytes; cannot persist leaf PEM")
written += n
def _load_via_mkstemp(ctx: ssl.SSLContext, combined: bytes) -> None:
"""Write *combined* to a 0600 mkstemp file, load it, then unlink.
Fallback for platforms without ``memfd_create`` (Windows, macOS). Unlike the
Linux memfd path, the leaf key is briefly on disk here. On POSIX it is 0600;
on Windows mode bits are not enforceable, so the protection is the per-user
temp dir + immediate unlink. See docs/adr/0001 "Leaf private-key posture".
"""
fd, path = tempfile.mkstemp(prefix="hr_leaf_", suffix=".pem")
try:
_write_all_fd(fd, combined)
os.close(fd)
fd = -1 # prevent double-close in finally
_assert_perms(Path(path), 0o600)
ctx.load_cert_chain(path)
finally:
if fd != -1:
try:
os.close(fd)
except OSError:
pass
try:
os.unlink(path)
except OSError:
pass

379
headroom/proxy/agy_dispatch.py Normal file
View file

@ -0,0 +1,379 @@
"""In-process hypercorn HTTPS dispatch server for agy MITM transport.
Serves the existing headroom FastAPI app on a loopback HTTPS port so that
the agy CONNECT terminator can byte-splice accepted client connections straight
to this server — no second upstream TLS dial, no logic duplication.
Architecture (ADR 0001 §"Dispatch via hypercorn"):
agy → CONNECT terminator (T8) → byte-splice → this server (TLS) → FastAPI app
↑ mints leaf per SNI via _LeafCache
Security invariants:
- Binds 127.0.0.1 only (loopback guard).
- Leaf private keys are loaded from anonymous memory (memfd) on Linux and
never touch the filesystem; on platforms without memfd, a 0600 temp file
is written and unlinked immediately after load (perms asserted).
- ALPN offers ["h2", "http/1.1"] matching the terminator leaf context.
Header handling: the Gemini handler strips the inbound ``accept-encoding``
header so the upstream returns a compressible (plain) body; agy UA and other
client headers are forwarded unchanged. The handler recompresses for the
upstream connection where applicable.
"""
from __future__ import annotations
import asyncio
import logging
import os
import socket
import ssl
from pathlib import Path
from typing import Any
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
from cryptography.x509 import Certificate
from headroom.proxy.agy_ca import ensure_root_ca, load_cert_chain_in_memory
from headroom.proxy.agy_terminator import DEFAULT_ALLOWLIST, _LeafCache, normalize_host
logger = logging.getLogger("headroom.proxy.agy_dispatch")
_BIND_HOST = "127.0.0.1"
_PLACEHOLDER_HOST = "headroom.internal"
# ---------------------------------------------------------------------------
# ASGI helpers
# ---------------------------------------------------------------------------
async def _send_421(send: Any) -> None:
"""Send a minimal HTTP 421 Misdirected Request response."""
body = b"Misdirected Request"
await send(
{
"type": "http.response.start",
"status": 421,
"headers": [
(b"content-type", b"text/plain"),
(b"content-length", str(len(body)).encode()),
],
}
)
await send({"type": "http.response.body", "body": body, "more_body": False})
def make_host_guard(app: Any, allowlist: frozenset[str], project: str | None = None) -> Any:
"""Wrap an ASGI *app* with a post-handshake Host/authority allowlist guard.
Mandatory defense-in-depth for the no-SNI / placeholder path (where the
TLS SNI guard may not fire). Hypercorn normalizes the HTTP/2 ``:authority``
pseudo-header into a ``host`` header, so reading ``host`` covers h2 and
http/1.1 uniformly. Module-level (not a closure) so it is unit-testable
with synthetic ASGI scopes.
When *project* is truthy, an ``x-headroom-project`` request header carrying
the launch-directory project label is injected (after the Host allowlist
check passes) so per-project savings attribute to the agy launch directory.
agy is a Go binary with no header knob, so this MUST happen at the MITM
boundary. Any client-forged ``x-headroom-project`` value is replaced (never
duplicated). ``project`` is computed once at launch and is already
RFC-3986 percent-encoded ASCII, safe for latin-1 encoding.
"""
async def _host_guard_app(scope: dict[str, Any], receive: Any, send: Any) -> None:
if scope.get("type") in ("http", "websocket"):
# Enforce exactly ONE Host header — multiple Host headers are a
# request-smuggling vector (guard validates one, backend may route
# on another). RFC 7230 §5.4 requires rejecting them.
host_values = [
value for name, value in scope.get("headers", ()) if name.lower() == b"host"
]
if len(host_values) != 1:
logger.warning("event=host_refused host_count=%d", len(host_values))
await _send_421(send)
return
host_str = host_values[0].decode("latin-1")
if not host_str:
logger.warning("event=host_refused host=%r", host_str)
await _send_421(send)
return
# Same normalization as the CONNECT target and the SNI guard, so no
# layer of the allowlist check can disagree with another.
normalized = normalize_host(host_str)
if normalized not in allowlist:
logger.warning("event=host_refused host=%s", host_str)
await _send_421(send)
return
if project:
# Replace any client-forged x-headroom-project value; never
# duplicate. Only touch http/websocket scopes.
scope["headers"] = [
(name, value)
for name, value in scope.get("headers", ())
if name.lower() != b"x-headroom-project"
] + [(b"x-headroom-project", project.encode("latin-1"))]
await app(scope, receive, send)
return _host_guard_app
# ---------------------------------------------------------------------------
# SNI-capable SSL context builder
# ---------------------------------------------------------------------------
def _build_sni_ssl_context(
leaf_cache: _LeafCache,
ca_key: RSAPrivateKey,
ca_cert: Certificate,
allowlist: frozenset[str],
) -> ssl.SSLContext:
"""Return a server SSLContext whose SNI callback mints leaf certs on demand.
The initial certfile/keyfile uses a wildcard placeholder cert so that
ssl.SSLContext accepts the load_cert_chain call; the SNI callback replaces
it per-connection before the handshake completes.
ALPN: ["h2", "http/1.1"] — required for HTTP/2 negotiation.
"""
# Mint a placeholder leaf for the initial load_cert_chain (SNI callback
# guards against it before the handshake completes — placeholder never
# served to real clients because SNI guard rejects non-allowlisted names).
init_cert_pem, init_key_pem = leaf_cache.get_or_mint(_PLACEHOLDER_HOST, ca_key, ca_cert)
ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
ctx.minimum_version = ssl.TLSVersion.TLSv1_2
ctx.set_alpn_protocols(["h2", "http/1.1"])
# Load the placeholder cert chain (required before SNI callback fires).
load_cert_chain_in_memory(ctx, init_cert_pem, init_key_pem)
def _sni_callback(
ssl_obj: ssl.SSLObject,
server_name: str | None,
ctx_in: ssl.SSLContext, # noqa: ARG001
) -> int | None:
"""Guard SNI then mint or reuse a leaf cert for *server_name* and swap it in-place."""
# Case-insensitive per RFC 6066; normalize once so the membership check
# AND the cache key match the allowlist, the Host guard and the CONNECT
# target.
host = normalize_host(server_name) if server_name is not None else None
if host is None or host not in allowlist:
logger.warning("event=sni_refused host=%s", server_name)
return ssl.ALERT_DESCRIPTION_UNRECOGNIZED_NAME
cert_pem, key_pem = leaf_cache.get_or_mint(host, ca_key, ca_cert)
new_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
new_ctx.minimum_version = ssl.TLSVersion.TLSv1_2
new_ctx.set_alpn_protocols(["h2", "http/1.1"])
load_cert_chain_in_memory(new_ctx, cert_pem, key_pem)
ssl_obj.context = new_ctx # type: ignore[assignment]
return None
ctx.set_servername_callback(_sni_callback) # type: ignore[arg-type]
return ctx
# ---------------------------------------------------------------------------
# AgyDispatchServer
# ---------------------------------------------------------------------------
class AgyDispatchServer:
"""In-process hypercorn HTTPS server serving the headroom FastAPI app.
Binds on loopback only; TLS via SNI callback (mints leaf per hostname
from the headroom root CA). Hypercorn handles h2/http1.1 + lifespan.
With ``plain_http=True`` the same plumbing serves the app over PLAIN HTTP:
no SSL context, no CA touched, no Host allowlist guard. That is the
retrieve listener (:class:`headroom.proxy.agy_retrieve.AgyRetrieveServer`),
which a stdio ``headroom mcp serve`` child must reach over loopback — it
cannot speak the Cloud-Code-SNI TLS the dispatch listener requires.
Usage::
server = AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert)
await server.start()
# server.address → ("127.0.0.1", <ephemeral-port>)
await server.stop()
Or as an async context manager::
async with AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) as srv:
host, port = srv.address
"""
def __init__(
self,
ca_key: RSAPrivateKey | None = None,
ca_cert: Certificate | None = None,
base_dir: Path | None = None,
port: int = 0,
allowlist: frozenset[str] | None = None,
project: str | None = None,
plain_http: bool = False,
) -> None:
self._plain_http = plain_http
self._ca_key_init = ca_key
self._ca_cert_init = ca_cert
self._base_dir = base_dir
self._port = port
self._allowlist: frozenset[str] = allowlist if allowlist is not None else DEFAULT_ALLOWLIST
self._project = project
self._server: asyncio.Server | None = None
self._lifespan_task: asyncio.Task[None] | None = None
self._lifespan: Any | None = None # hypercorn.asyncio.run.Lifespan
self._context: Any | None = None # hypercorn.asyncio.run.WorkerContext
self._app_wrapper: Any | None = None
self._config: Any | None = None
self._lifespan_state: dict[str, Any] = {}
self._leaf_cache: _LeafCache | None = None
async def start(self) -> None:
"""Start the hypercorn server; binds loopback (HTTPS, or plain HTTP) on a port."""
from hypercorn.asyncio import wrap_app
from hypercorn.asyncio.run import Lifespan, TCPServer, WorkerContext
from hypercorn.config import Config
ssl_ctx: ssl.SSLContext | None = None
if not self._plain_http:
# Resolve CA.
if self._ca_key_init is not None and self._ca_cert_init is not None:
ca_key = self._ca_key_init
ca_cert = self._ca_cert_init
else:
ca_key, ca_cert, _, _ = ensure_root_ca(base_dir=self._base_dir)
self._leaf_cache = _LeafCache(max_size=len(self._allowlist) + 1)
ssl_ctx = _build_sni_ssl_context(self._leaf_cache, ca_key, ca_cert, self._allowlist)
# Build minimal hypercorn Config (no certfile/keyfile — we supply ssl directly).
config = Config()
config.bind = [f"{_BIND_HOST}:{self._port}"]
config.accesslog = "-" # suppress hypercorn access log noise in tests
config.errorlog = "-"
config.loglevel = "WARNING"
self._config = config
# Import and build the FastAPI app.
from headroom.proxy.server import create_app
# Plain HTTP (retrieve listener): no Host allowlist guard — the client is
# a stdio child in the same trust boundary, addressing 127.0.0.1 directly.
app: Any = create_app()
if not self._plain_http:
app = make_host_guard(app, self._allowlist, self._project)
# wrap_app accepts the ASGI callable directly; ignore the narrow stub type.
app_wrapper = wrap_app(app, config.wsgi_max_body_size, mode="asgi") # type: ignore[arg-type]
self._app_wrapper = app_wrapper
# Run hypercorn lifespan (startup/shutdown events).
loop = asyncio.get_event_loop()
lifespan_state: dict[str, Any] = {}
self._lifespan_state = lifespan_state
lifespan = Lifespan(app_wrapper, config, loop, lifespan_state)
self._lifespan = lifespan
self._lifespan_task = loop.create_task(lifespan.handle_lifespan())
await lifespan.wait_for_startup()
if self._lifespan_task.done():
exc = self._lifespan_task.exception()
if exc is not None:
raise exc
worker_context = WorkerContext(max_requests=None)
self._context = worker_context
# Bind a plain TCP socket on loopback (wrapped with our SSL context
# below unless this is the plain-HTTP retrieve listener).
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
# SO_REUSEADDR means fast TIME_WAIT reuse on POSIX, but on Windows it
# lets a second process bind this same loopback port and intercept the
# decrypted MITM traffic. Restrict to POSIX; on Windows enforce
# exclusive use so a duplicate bind fails loudly.
if os.name == "posix":
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
elif hasattr(socket, "SO_EXCLUSIVEADDRUSE"):
sock.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
sock.bind((_BIND_HOST, self._port))
async def _connection_handler(
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
) -> None:
await TCPServer(
app_wrapper,
loop,
config,
worker_context,
lifespan_state,
reader,
writer,
)
# asyncio rejects ssl_handshake_timeout when ssl is None, so only pass
# it on the TLS path.
if ssl_ctx is None:
self._server = await asyncio.start_server(
_connection_handler,
sock=sock,
)
else:
self._server = await asyncio.start_server(
_connection_handler,
sock=sock,
ssl=ssl_ctx,
ssl_handshake_timeout=config.ssl_handshake_timeout,
)
addr = self._server.sockets[0].getsockname()
logger.info("event=%s_started address=%s:%d", self._event, addr[0], addr[1])
async def stop(self) -> None:
"""Gracefully shut down the server and hypercorn lifespan."""
if self._server is not None:
self._server.close()
await self._server.wait_closed()
self._server = None
if self._lifespan is not None:
try:
await self._lifespan.wait_for_shutdown()
except Exception: # noqa: BLE001
pass
self._lifespan = None
if self._lifespan_task is not None:
self._lifespan_task.cancel()
try:
await self._lifespan_task
except (asyncio.CancelledError, Exception): # noqa: BLE001
pass
self._lifespan_task = None
logger.info("event=%s_stopped", self._event)
@property
def _event(self) -> str:
"""Log event prefix: ``retrieve`` for the plain-HTTP listener."""
return "retrieve" if self._plain_http else "dispatch"
@property
def address(self) -> tuple[str, int]:
"""Return ``(host, port)`` the server is bound to. Requires :meth:`start`."""
if self._server is None:
raise RuntimeError(f"{type(self).__name__} not started")
sock = self._server.sockets[0]
host, port = sock.getsockname()[:2]
return host, port
async def __aenter__(self) -> AgyDispatchServer:
await self.start()
return self
async def __aexit__(self, *_: object) -> None:
await self.stop()

53
headroom/proxy/agy_retrieve.py Normal file
View file

@ -0,0 +1,53 @@
"""In-process hypercorn PLAIN-HTTP retrieve server for agy.
The proxy compresses tool_result payloads and emits ``[Retrieve more:
hash=…]`` markers. For agy those markers are produced on the decrypted
stream inside the HTTPS dispatch server (:mod:`headroom.proxy.agy_dispatch`).
To resolve a marker the agent runs the ``headroom mcp serve`` stdio child,
which calls the proxy's retrieve HTTP endpoint via ``HEADROOM_PROXY_URL``.
The dispatch server is HTTPS with a Cloud-Code-SNI leaf only, so a stdio
retrieve child cannot reach it over loopback. This module stands up a
SECOND loopback listener — PLAIN HTTP, no TLS — serving the same FastAPI
app on an ephemeral port for the session. The compression/marker cache is
a process-global singleton (:func:`headroom.cache.compression_store.get_compression_store`),
so this second ``create_app()`` shares the exact cache the dispatch server
populates: a marker minted on the HTTPS side resolves over plain HTTP here.
Why plain HTTP is safe: the listener binds ``127.0.0.1`` only, serves the
retrieve endpoints to a stdio child in the *same* trust boundary, and never
carries upstream credentials (it only reads the in-memory marker cache).
The hypercorn plumbing (lifespan, TCPServer, socket options, lifecycle) is
:class:`headroom.proxy.agy_dispatch.AgyDispatchServer`'s — this listener is
that same server in its ``plain_http`` mode: no SSL context, no CA touched,
no Host allowlist guard.
"""
from __future__ import annotations
from headroom.proxy.agy_dispatch import AgyDispatchServer
class AgyRetrieveServer(AgyDispatchServer):
"""PLAIN-HTTP loopback listener serving the headroom FastAPI app.
Serves the process-global compression cache via ``create_app()`` so
``GET /v1/retrieve/{hash}`` resolves markers the HTTPS dispatch server
populated.
Usage::
server = AgyRetrieveServer()
await server.start()
# server.address → ("127.0.0.1", <ephemeral-port>)
await server.stop()
Or as an async context manager::
async with AgyRetrieveServer() as srv:
host, port = srv.address
"""
def __init__(self, port: int = 0) -> None:
super().__init__(port=port, plain_http=True)

275
headroom/proxy/agy_savings_inbox.py Normal file
View file

@ -0,0 +1,275 @@
"""Cross-process savings inbox: agy -> shared proxy replay.
``agy`` runs as a **separate process** from the shared Headroom proxy, but its
per-request savings must show up on the shared dashboard, counted once, without
agy ever writing any shared durable state. This module is the bridge.
Mechanism (AT-LEAST-ONCE with best-effort dedup — *not* exactly-once):
* In the agy process, :func:`emit_event` drops one JSON file per request into a
canonical inbox directory (``workspace_dir()/savings.d``). Each file carries
the exact keyword arguments that :meth:`PrometheusMetrics.record_request`
(the single dashboard funnel) expects, plus a unique ``event_id``.
* In the shared proxy, :func:`drain_inbox` reads those files and replays each
event through its *own* ``record_request`` funnel — the one writer of shared
durable state (savings ledger, SavingsTracker, OTEL). agy itself redirects all
three of those to throwaway paths, so the proxy replay is the sole writer and
savings are counted exactly once on the dashboard.
Everything on the agy side is best-effort: emit never raises into the request
path, and drain never raises out into the proxy's lifespan / stats handler.
"""
from __future__ import annotations
import asyncio
import itertools
import json
import logging
import os
import random
import tempfile
from pathlib import Path
from typing import Any
from headroom.paths import workspace_dir
logger = logging.getLogger("headroom.proxy")
# Bump when the on-disk envelope shape changes incompatibly.
SCHEMA_VERSION = 1
# Env var (set only in the agy process) that turns on emit at the outcome hook.
AGY_INBOX_EMIT_ENV = "HEADROOM_AGY_INBOX_EMIT"
# Hard cap on pending event files; oldest are dropped (disclosed) past this.
MAX_INBOX = 5000
# Keep at most this many processed ids in the dedup file so it stays bounded.
MAX_PROCESSED_IDS = 20000
_INBOX_SUBDIR = "savings.d"
_PROCESSED_FILE = ".processed"
# Monotonic per-process sequence so two events from the same pid never collide.
_seq = itertools.count()
# Serialize drains so the periodic task and /stats-triggered drain never race.
_drain_lock = asyncio.Lock()
def inbox_dir() -> Path:
"""Return the canonical inbox directory, creating it on demand."""
path = workspace_dir() / _INBOX_SUBDIR
path.mkdir(parents=True, exist_ok=True)
return path
def agy_emit_enabled() -> bool:
"""True when the agy emit marker env var is set to ``"1"``."""
return os.environ.get(AGY_INBOX_EMIT_ENV, "").strip() == "1"
def _new_event_id() -> str:
"""Return a process-unique, collision-resistant event id."""
return f"{os.getpid()}-{next(_seq)}-{random.getrandbits(48):012x}"
def _json_safe(value: Any) -> Any:
"""Return ``value`` if it round-trips through JSON, else ``None``.
Non-scalar funnel args (``pipeline_timing``, ``waste_signals``) are dicts of
scalars and normally survive; anything that does not is dropped so a single
weird value can never make the whole envelope unwritable.
"""
try:
json.dumps(value)
return value
except (TypeError, ValueError):
return None
def _enforce_cap() -> None:
"""Drop the oldest event files if the inbox is at/over :data:`MAX_INBOX`."""
try:
files = sorted(
inbox_dir().glob("evt-*.json"),
key=lambda p: p.stat().st_mtime,
)
except OSError:
return
excess = len(files) - MAX_INBOX
if excess <= 0:
return
for stale in files[: excess + 1]:
try:
stale.unlink()
except OSError:
continue
logger.warning(
"agy savings inbox at cap (%d); dropped %d oldest event(s)",
MAX_INBOX,
excess + 1,
)
def emit_event(**funnel_kwargs: Any) -> None:
"""Atomically write one inbox event carrying ``record_request`` kwargs.
Best-effort: any failure is swallowed (logged at debug) so emit can never
break the request that triggered it.
"""
try:
directory = inbox_dir()
_enforce_cap()
safe_kwargs = {key: _json_safe(val) for key, val in funnel_kwargs.items()}
event_id = _new_event_id()
envelope = {
"v": SCHEMA_VERSION,
"event_id": event_id,
"kwargs": safe_kwargs,
}
fd, tmp_name = tempfile.mkstemp(dir=directory, prefix=".tmp-evt-", suffix=".json")
try:
with os.fdopen(fd, "w", encoding="utf-8") as fh:
json.dump(envelope, fh)
fh.flush()
os.fsync(fh.fileno())
os.replace(tmp_name, directory / f"evt-{event_id}.json")
except Exception:
try:
os.unlink(tmp_name)
except OSError:
pass
raise
except Exception as exc: # noqa: BLE001 - best-effort, never raise into caller
logger.debug("agy savings emit failed: %s", exc)
def _load_processed(path: Path) -> list[str]:
"""Return the processed-id list (order preserved), or ``[]`` if unreadable."""
try:
text = path.read_text(encoding="utf-8")
except OSError:
return []
ids: list[str] = []
for line in text.splitlines():
line = line.strip()
if line:
ids.append(line)
return ids
def _write_processed(path: Path, ids: list[str]) -> None:
"""Atomically persist the processed-id list, pruned to the newest N."""
pruned = ids[-MAX_PROCESSED_IDS:]
try:
fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=".tmp-proc-")
with os.fdopen(fd, "w", encoding="utf-8") as fh:
fh.write("\n".join(pruned))
if pruned:
fh.write("\n")
fh.flush()
os.fsync(fh.fileno())
os.replace(tmp_name, path)
except OSError as exc:
logger.debug("agy savings processed-file write failed: %s", exc)
async def drain_inbox(metrics: Any, *, max_events: int = 1000) -> int:
"""Replay pending inbox events through ``metrics.record_request``.
Delivery is at-least-once with best-effort dedup: an event is only unlinked
*after* it has been recorded (or found already-processed), so a crash
between record and unlink re-delivers it next drain — the ``.processed`` set
then suppresses the duplicate. Returns the number of events recorded.
Never raises: the whole body is defended so a drain error can never crash
the proxy lifespan loop or the stats handler.
"""
recorded = 0
async with _drain_lock:
try:
directory = inbox_dir()
processed_path = directory / _PROCESSED_FILE
processed_list = _load_processed(processed_path)
processed_set = set(processed_list)
try:
files = sorted(directory.glob("evt-*.json"))
except OSError:
return 0
dirty = False
for event_file in files[:max_events]:
try:
try:
raw = event_file.read_text(encoding="utf-8")
envelope = json.loads(raw)
except (OSError, ValueError):
# Malformed / unreadable: skip and remove, never fatal.
logger.debug("agy savings: dropping malformed %s", event_file.name)
_safe_unlink(event_file)
continue
event_id = envelope.get("event_id")
if not isinstance(event_id, str) or not event_id:
_safe_unlink(event_file)
continue
if event_id in processed_set:
# Crash-window duplicate: already recorded, just remove.
_safe_unlink(event_file)
continue
kwargs = envelope.get("kwargs")
if not isinstance(kwargs, dict):
_safe_unlink(event_file)
continue
await metrics.record_request(**kwargs)
recorded += 1
processed_set.add(event_id)
processed_list.append(event_id)
dirty = True
_safe_unlink(event_file)
except Exception as exc: # noqa: BLE001 - one bad event never aborts drain
logger.debug("agy savings: error replaying %s: %s", event_file.name, exc)
continue
if dirty:
_write_processed(processed_path, processed_list)
except Exception as exc: # noqa: BLE001 - drain never raises out
logger.debug("agy savings drain failed: %s", exc)
return recorded
def _safe_unlink(path: Path) -> None:
try:
path.unlink()
except OSError:
pass
__all__ = [
"AGY_INBOX_EMIT_ENV",
"SCHEMA_VERSION",
"MAX_INBOX",
"inbox_dir",
"agy_emit_enabled",
"emit_event",
"drain_inbox",
]

725
headroom/proxy/agy_terminator.py Normal file
View file

@ -0,0 +1,725 @@
"""Selective TLS-MITM forward-proxy listener for the agy MITM transport.
Binds to 127.0.0.1 ONLY. Accepts HTTP CONNECT:
- Allowlisted hosts: ACK the CONNECT and byte-splice the raw connection to the
in-process hypercorn HTTPS server at ``dispatch_port`` (AgyDispatchServer).
The hypercorn server owns TLS termination and ASGI routing.
- Non-allowlisted hosts: raw bidirectional byte-splice (blind tunnel).
If HTTPS_PROXY is set, forward CONNECT through that upstream proxy, deriving
Proxy-Authorization from its userinfo when present. An ``https://`` upstream
proxy is dialled over TLS (``ssl.create_default_context()``, default
certificate validation, SNI set to the proxy's own hostname, ALPN pinned to
``http/1.1``) instead of plaintext on :443.
NEVER chain to a loopback address (self-loop guard).
Security invariants:
- Leaf private keys are loaded from anonymous memory (memfd) on Linux and
never touch the filesystem; on platforms without memfd, a 0600 temp file
is written and unlinked immediately after load (perms asserted).
- Proxy-Authorization is never logged.
- Upstream proxy auth: when HTTPS_PROXY carries `user:pass@` userinfo, it is
percent-decoded and sent as HTTP Basic auth, only when the proxy scheme is
`http`/`https` (never to e.g. `socks5://`), and only the URL-derived
credential is used when the URL carries one (it takes precedence over an
inbound Proxy-Authorization header). This is sent in cleartext when the
upstream scheme is `http://` — same as curl, Go and requests do to a plain
HTTP proxy; the credential already lives in an env var every tool on the
box can read, and refusing to send it would break most corporate proxies.
- Listener bind address is 127.0.0.1, never 0.0.0.0.
"""
from __future__ import annotations
import asyncio
import base64
import datetime
import ipaddress
import logging
import os
import socket
import ssl
import urllib.parse
from pathlib import Path
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
from cryptography.x509 import Certificate
from cryptography.x509.oid import ExtendedKeyUsageOID, NameOID
from headroom.proxy.agy_ca import ensure_root_ca
logger = logging.getLogger("headroom.proxy.agy_terminator")
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
_LEAF_KEY_BITS = 2048
_LEAF_VALIDITY_HOURS = 72
_BIND_HOST = "127.0.0.1"
_CONNECT_TIMEOUT = 10.0
_SPLICE_BUF = 65536
DEFAULT_ALLOWLIST: frozenset[str] = frozenset(
{
"daily-cloudcode-pa.googleapis.com",
"cloudcode-pa.googleapis.com",
}
)
# ---------------------------------------------------------------------------
# Leaf certificate minting
# ---------------------------------------------------------------------------
def mint_leaf(
host: str,
ca_key: RSAPrivateKey,
ca_cert: Certificate,
) -> tuple[bytes, bytes]:
"""Mint a leaf TLS certificate for *host* signed by the root CA.
Parameters
----------
host:
Hostname for SAN=dNSName entry.
ca_key:
Root CA private key (in-memory, never written).
ca_cert:
Root CA certificate object.
Returns
-------
(cert_pem, key_pem)
Both as PEM bytes. Leaf private keys are loaded from anonymous memory
(memfd) on Linux and never touch the filesystem; on platforms without
memfd, a 0600 temp file is written and unlinked immediately after load
(perms asserted).
"""
leaf_key: RSAPrivateKey = rsa.generate_private_key(
public_exponent=65537,
key_size=_LEAF_KEY_BITS,
)
now = datetime.datetime.now(tz=datetime.timezone.utc)
not_after = now + datetime.timedelta(hours=_LEAF_VALIDITY_HOURS)
cert = (
x509.CertificateBuilder()
.subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, host)]))
.issuer_name(ca_cert.subject)
.public_key(leaf_key.public_key())
.serial_number(x509.random_serial_number())
.not_valid_before(now)
.not_valid_after(not_after)
.add_extension(
x509.SubjectAlternativeName([x509.DNSName(host)]),
critical=False,
)
.add_extension(
x509.ExtendedKeyUsage([ExtendedKeyUsageOID.SERVER_AUTH]),
critical=True,
)
.add_extension(
x509.BasicConstraints(ca=False, path_length=None),
critical=True,
)
.sign(ca_key, hashes.SHA256())
)
cert_pem = cert.public_bytes(serialization.Encoding.PEM)
key_pem = leaf_key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.TraditionalOpenSSL,
encryption_algorithm=serialization.NoEncryption(),
)
return cert_pem, key_pem
# ---------------------------------------------------------------------------
# Leaf cert cache
# ---------------------------------------------------------------------------
class _LeafCache:
"""Fixed-bound leaf cert cache keyed by hostname.
Bound to allowlist size (small dict). Entries are reused within
validity; expired entries are replaced in-place.
"""
def __init__(self, max_size: int) -> None:
self._max = max(max_size, 1)
# host -> (cert_pem, key_pem, not_after_utc)
self._cache: dict[str, tuple[bytes, bytes, datetime.datetime]] = {}
def get_or_mint(
self,
host: str,
ca_key: RSAPrivateKey,
ca_cert: Certificate,
) -> tuple[bytes, bytes]:
"""Return cached leaf or mint a fresh one."""
now = datetime.datetime.now(tz=datetime.timezone.utc)
if host in self._cache:
cert_pem, key_pem, not_after = self._cache[host]
if now < not_after - datetime.timedelta(minutes=5):
return cert_pem, key_pem
# Expired — re-mint in place.
del self._cache[host]
if len(self._cache) >= self._max:
# Evict oldest entry (FIFO; dict preserves insertion order in Python 3.7+).
oldest = next(iter(self._cache))
del self._cache[oldest]
cert_pem, key_pem = mint_leaf(host, ca_key, ca_cert)
# Parse just-minted cert to get its not_valid_after.
cert_obj = x509.load_pem_x509_certificate(cert_pem)
self._cache[host] = (cert_pem, key_pem, cert_obj.not_valid_after_utc)
logger.debug("event=leaf_minted host=%s", host)
return cert_pem, key_pem
# ---------------------------------------------------------------------------
# Loopback guard helper
# ---------------------------------------------------------------------------
def _is_loopback(host: str) -> bool:
"""Return True if *host* is a loopback address, IP-literal forms only.
Catches dotted-quad, IPv4-shorthand (``127.1``, decimal ``2130706433``),
the unspecified address (``0.0.0.0`` / ``0``, which Linux ``connect()``
treats as loopback), IPv6 ``::1``, IPv4-mapped IPv6 (``::ffff:127.0.0.1``,
normalised via ``ipv4_mapped`` so the result does not depend on the
interpreter version — see CPython gh-103365, fixed in 3.13), and
``localhost``/``localhost.`` case-insensitively. Does NOT resolve DNS: a
hostname that resolves to loopback (e.g. an attacker-controlled
``/etc/hosts`` entry) is not detected here and returns False.
"""
if host.lower() in ("localhost", "localhost."):
return True
try:
addr = ipaddress.ip_address(host)
except ValueError:
try:
packed = socket.inet_aton(host)
except (OSError, UnicodeError):
return False
addr = ipaddress.IPv4Address(packed)
if isinstance(addr, ipaddress.IPv6Address):
addr = addr.ipv4_mapped or addr
return bool(addr.is_loopback or addr.is_unspecified)
# ---------------------------------------------------------------------------
# Byte-splice helpers
# ---------------------------------------------------------------------------
async def _splice_half(
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
) -> None:
"""Forward bytes from reader to writer until EOF."""
try:
while True:
data = await reader.read(_SPLICE_BUF)
if not data:
break
writer.write(data)
await writer.drain()
except (ConnectionResetError, BrokenPipeError, asyncio.CancelledError):
pass
finally:
try:
writer.write_eof()
except Exception: # noqa: BLE001
pass
async def _blind_splice(
client_reader: asyncio.StreamReader,
client_writer: asyncio.StreamWriter,
target_reader: asyncio.StreamReader,
target_writer: asyncio.StreamWriter,
) -> None:
"""Bidirectional byte-splice until either side closes.
Waits until the FIRST half-stream closes (one side EOF'd / connection
dropped), then cancels the other. This avoids a hang when the target
closes after echoing but the client hasn't sent EOF yet.
"""
t1 = asyncio.create_task(_splice_half(client_reader, target_writer))
t2 = asyncio.create_task(_splice_half(target_reader, client_writer))
try:
done, pending = await asyncio.wait({t1, t2}, return_when=asyncio.FIRST_COMPLETED)
for task in pending:
task.cancel()
await asyncio.gather(*pending, return_exceptions=True)
except Exception: # noqa: BLE001
t1.cancel()
t2.cancel()
await asyncio.gather(t1, t2, return_exceptions=True)
finally:
for w in (client_writer, target_writer):
try:
w.close()
await w.wait_closed()
except Exception: # noqa: BLE001
pass
# ---------------------------------------------------------------------------
# CONNECT request parser
# ---------------------------------------------------------------------------
def normalize_host(value: str) -> str:
"""Return *value* as a bare, comparable hostname.
Strips an optional ``:port``, a trailing root dot, and case. Hostnames are
case-insensitive and ``example.com.`` names the same host as ``example.com``,
so every exact-match allowlist check in the agy path — CONNECT target, SNI,
Host header, passthrough base — must compare this form. Otherwise the layers
disagree: ``CloudCode-PA.googleapis.com`` skips TLS termination and passes
through uncompressed with no signal that anything was bypassed.
"""
host = value.strip()
# Strip a trailing ``:port`` only when it IS a port. ``example.com:abc`` is
# not a host with a port, so it stays whole and simply fails the allowlist;
# requiring exactly one colon leaves IPv6 literals such as ``::1`` alone.
if host.count(":") == 1:
left, _, right = host.rpartition(":")
if right.isdigit():
host = left
return host.rstrip(".").lower()
def _parse_connect(line: str) -> tuple[str, int]:
"""Parse 'CONNECT host:port HTTP/1.x' → (host, port). Raises ValueError."""
parts = line.strip().split()
if len(parts) < 2 or parts[0].upper() != "CONNECT":
raise ValueError(f"Not a CONNECT request: {line!r}")
hostport = parts[1]
if ":" not in hostport:
raise ValueError(f"Missing port in CONNECT target: {hostport!r}")
host, port_str = hostport.rsplit(":", 1)
return normalize_host(host), int(port_str)
# ---------------------------------------------------------------------------
# Upstream proxy (HTTPS_PROXY) tunnel
# ---------------------------------------------------------------------------
async def _connect_via_upstream_proxy(
proxy_host: str,
proxy_port: int,
target_host: str,
target_port: int,
proxy_auth: str | None,
ssl_context: ssl.SSLContext | None = None,
) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
"""Open a TCP connection through an upstream HTTP proxy using CONNECT.
*ssl_context* is non-None only for an ``https://`` upstream proxy: the
CONNECT dial itself is then wrapped in TLS to the proxy (SNI ==
``proxy_host``, the proxy's own name — the tunnelled payload carries the
target's TLS handshake and SNI separately, inside the tunnel).
"""
reader, writer = await asyncio.wait_for(
asyncio.open_connection(
proxy_host,
proxy_port,
ssl=ssl_context,
server_hostname=proxy_host if ssl_context is not None else None,
),
timeout=_CONNECT_TIMEOUT,
)
connect_line = (
f"CONNECT {target_host}:{target_port} HTTP/1.1\r\nHost: {target_host}:{target_port}\r\n"
)
if proxy_auth:
connect_line += f"Proxy-Authorization: {proxy_auth}\r\n"
connect_line += "\r\n"
writer.write(connect_line.encode())
await writer.drain()
# Read response — look for 200 Connection Established.
try:
response_line = await asyncio.wait_for(reader.readline(), timeout=_CONNECT_TIMEOUT)
if b"200" not in response_line:
raise OSError(f"Upstream proxy refused CONNECT: {response_line!r}")
# Drain remaining headers.
while True:
hdr = await asyncio.wait_for(reader.readline(), timeout=_CONNECT_TIMEOUT)
if hdr in (b"\r\n", b"\n", b""):
break
except (OSError, asyncio.TimeoutError):
writer.close()
raise
return reader, writer
# ---------------------------------------------------------------------------
# Main connection handler
# ---------------------------------------------------------------------------
async def _handle_connect(
client_reader: asyncio.StreamReader,
client_writer: asyncio.StreamWriter,
allowlist: frozenset[str],
dispatch_port: int,
self_port: int | None = None,
) -> None:
"""Handle one incoming TCP connection carrying an HTTP CONNECT request.
*self_port* is the terminator's own listening port, used to refuse a tunnel
that would loop back into this very listener.
"""
peer = client_writer.get_extra_info("peername", ("?", 0))
try:
first_line_bytes = await asyncio.wait_for(
client_reader.readline(), timeout=_CONNECT_TIMEOUT
)
except asyncio.TimeoutError:
logger.debug("event=connect_timeout peer=%s", peer)
client_writer.close()
return
first_line = first_line_bytes.decode("latin-1")
try:
target_host, target_port = _parse_connect(first_line)
except ValueError as exc:
logger.debug("event=parse_error peer=%s err=%s", peer, exc)
client_writer.write(b"HTTP/1.1 400 Bad Request\r\n\r\n")
await client_writer.drain()
client_writer.close()
return
# Drain remaining CONNECT request headers.
proxy_auth: str | None = None
while True:
try:
hdr_bytes = await asyncio.wait_for(client_reader.readline(), timeout=_CONNECT_TIMEOUT)
except asyncio.TimeoutError:
logger.debug("event=connect_header_timeout peer=%s", peer)
client_writer.close()
return
if hdr_bytes in (b"\r\n", b"\n", b""):
break
hdr = hdr_bytes.decode("latin-1")
if hdr.lower().startswith("proxy-authorization:"):
proxy_auth = hdr.split(":", 1)[1].strip()
logger.debug(
"event=connect_received peer=%s target=%s:%d allowlisted=%s",
peer,
target_host,
target_port,
target_host in allowlist,
)
if target_host in allowlist:
await _handle_mitm(client_reader, client_writer, dispatch_port)
else:
await _handle_blind_tunnel(
client_reader,
client_writer,
target_host,
target_port,
proxy_auth,
self_port=self_port,
)
async def _handle_mitm(
client_reader: asyncio.StreamReader,
client_writer: asyncio.StreamWriter,
dispatch_port: int,
) -> None:
"""Handle an allowlisted CONNECT: ACK it and byte-splice to hypercorn.
The raw connection is spliced to the loopback hypercorn HTTPS port
(AgyDispatchServer), which owns TLS termination, ALPN negotiation and
ASGI routing.
"""
# ACK the CONNECT so the client believes the tunnel is up.
client_writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
await client_writer.drain()
try:
dispatch_reader, dispatch_writer = await asyncio.wait_for(
asyncio.open_connection("127.0.0.1", dispatch_port),
timeout=_CONNECT_TIMEOUT,
)
except (OSError, asyncio.TimeoutError) as exc:
logger.error("event=dispatch_connect_failed port=%d err=%s", dispatch_port, exc)
try:
client_writer.close()
except Exception: # noqa: BLE001
pass
return
await _blind_splice(client_reader, client_writer, dispatch_reader, dispatch_writer)
async def _resolve_tunnel_target(host: str, port: int, self_port: int | None) -> str:
"""Return an address for *host* that is safe to tunnel to, else raise ValueError.
The terminator is an unauthenticated CONNECT proxy on loopback for the life of
an agy session, so anything running as the user can drive it. Two targets must
never be reachable through it:
* itself — ``CONNECT 127.0.0.1:<terminator_port>`` makes the terminator tunnel
into itself, burning two fds per nesting level until they run out;
* link-local — 169.254.0.0/16 carries the cloud instance-metadata service.
Other loopback ports are deliberately still reachable: a local process could
open them directly, so refusing them buys no security and would break plain
local tunnelling. The check runs on the *resolved* addresses, not the literal
(a name resolving to 127.0.0.1 is the same self-connect), and the vetted
address is what we connect to, so no second lookup can substitute another.
"""
loop = asyncio.get_running_loop()
infos = await loop.getaddrinfo(host, port, type=socket.SOCK_STREAM)
if not infos:
raise ValueError(f"no address for {host}")
for info in infos:
addr = ipaddress.ip_address(info[4][0])
if addr.is_link_local:
raise ValueError(f"{host} resolves to link-local {addr}")
if addr.is_loopback and self_port is not None and port == self_port:
raise ValueError("self-connect to the terminator's own port")
return str(ipaddress.ip_address(infos[0][4][0]))
def _upstream_proxy_auth(parsed: urllib.parse.ParseResult, inbound: str | None) -> str | None:
"""Resolve the Proxy-Authorization value to send to the upstream proxy.
Only pure string/URL logic — no socket I/O — so this is unit-testable
without spinning up a listener.
"""
if parsed.scheme not in ("http", "https"):
return None
if not parsed.username:
return inbound
password = urllib.parse.unquote(parsed.password or "")
userinfo = f"{urllib.parse.unquote(parsed.username)}:{password}".encode()
return "Basic " + base64.b64encode(userinfo).decode()
async def _handle_blind_tunnel(
client_reader: asyncio.StreamReader,
client_writer: asyncio.StreamWriter,
target_host: str,
target_port: int,
proxy_auth: str | None,
self_port: int | None = None,
) -> None:
"""Byte-splice tunnel for non-allowlisted targets.
When chaining through an upstream HTTPS_PROXY, Proxy-Authorization is
derived from that URL's userinfo (percent-decoded) and takes precedence
over any inbound Proxy-Authorization header: the URL is the operator's
configuration for this specific upstream proxy and is the only source
that can carry a working credential, since the child process is handed a
userinfo-free loopback URL and never sends a header of its own. The
inbound header remains a fallback for a caller that does supply one.
Trust boundary, chaining branch: the target (``target_host``) is NOT
vetted here — it is forwarded to the upstream proxy by name, which
re-resolves it in its own DNS view, so the upstream proxy's own egress
policy is the actual boundary, not anything checked in this process.
Only the proxy side is guarded (self-loop, scheme). The self-loop guard
(``_is_loopback``) covers IP-literal forms only; a DNS name that
resolves to loopback (e.g. a local ``/etc/hosts`` entry) is not detected
and is out of scope — see ``_is_loopback``'s docstring.
"""
upstream_proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy")
try:
if upstream_proxy:
parsed = urllib.parse.urlparse(upstream_proxy)
if parsed.scheme not in ("http", "https"):
raise ValueError(f"unsupported upstream proxy scheme: {parsed.scheme}")
proxy_host = parsed.hostname or ""
# Default per scheme: a scheme-less-port HTTPS_PROXY like
# "http://proxy.corp" speaks plain HTTP on :80, not :443.
proxy_port = parsed.port or (443 if parsed.scheme == "https" else 80)
# Self-loop guard: never chain through a loopback upstream proxy.
if _is_loopback(proxy_host):
# Log host:port only — never the full URL, which may embed
# user:pass@ credentials.
logger.warning(
"event=self_loop_blocked_proxy proxy=%s:%s",
proxy_host,
proxy_port,
)
client_writer.write(b"HTTP/1.1 403 Forbidden\r\n\r\n")
await client_writer.drain()
client_writer.close()
return
proxy_ssl_context: ssl.SSLContext | None = None
if parsed.scheme == "https":
proxy_ssl_context = ssl.create_default_context()
proxy_ssl_context.set_alpn_protocols(["http/1.1"])
target_reader, target_writer = await _connect_via_upstream_proxy(
proxy_host,
proxy_port,
target_host,
target_port,
_upstream_proxy_auth(parsed, proxy_auth),
proxy_ssl_context,
)
else:
target_addr = await asyncio.wait_for(
_resolve_tunnel_target(target_host, target_port, self_port),
timeout=_CONNECT_TIMEOUT,
)
target_reader, target_writer = await asyncio.wait_for(
asyncio.open_connection(target_addr, target_port),
timeout=_CONNECT_TIMEOUT,
)
except ValueError as exc:
logger.warning(
"event=tunnel_target_refused target=%s:%d reason=%s",
target_host,
target_port,
exc,
)
client_writer.write(b"HTTP/1.1 403 Forbidden\r\n\r\n")
await client_writer.drain()
client_writer.close()
return
except (OSError, asyncio.TimeoutError) as exc:
logger.debug(
"event=tunnel_connect_failed target=%s:%d err=%s",
target_host,
target_port,
exc,
)
client_writer.write(b"HTTP/1.1 502 Bad Gateway\r\n\r\n")
await client_writer.drain()
client_writer.close()
return
try:
client_writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n")
await client_writer.drain()
except OSError:
target_writer.close()
raise
await _blind_splice(client_reader, client_writer, target_reader, target_writer)
# ---------------------------------------------------------------------------
# Public API: Terminator server
# ---------------------------------------------------------------------------
class AgyCONNECTTerminator:
"""Asyncio forward-proxy listener implementing selective TLS-MITM.
Parameters
----------
dispatch_port:
Allowlisted CONNECT connections are ACK-ed and byte-spliced raw to
``127.0.0.1:<dispatch_port>`` (the in-process AgyDispatchServer).
allowlist:
Set of hostnames to TLS-terminate. Defaults to ``DEFAULT_ALLOWLIST``.
base_dir:
Headroom state directory (for CA; defaults to ~/.headroom).
Inject a ``tmp_path``-derived path in tests.
ca_key / ca_cert:
Pre-built CA key+cert. When provided, ``base_dir`` is not used for
CA loading. Intended for tests.
port:
Listener port. 0 = OS-assigned ephemeral (default; tests use this).
host:
Bind address. Hardcoded to ``127.0.0.1``; parameter exists only for
testing internal assertion — callers may not override to non-loopback.
"""
def __init__(
self,
dispatch_port: int,
allowlist: frozenset[str] | None = None,
base_dir: Path | None = None,
ca_key: RSAPrivateKey | None = None,
ca_cert: Certificate | None = None,
port: int = 0,
) -> None:
self._allowlist = allowlist if allowlist is not None else DEFAULT_ALLOWLIST
self._dispatch_port = dispatch_port
self._base_dir = base_dir
self._ca_key_init = ca_key
self._ca_cert_init = ca_cert
self._port = port
self._server: asyncio.Server | None = None
self._ca_key: RSAPrivateKey | None = None
self._ca_cert: Certificate | None = None
self._leaf_cache: _LeafCache | None = None
async def start(self) -> None:
"""Start the listener. Must be called before :meth:`address`."""
if self._ca_key_init is not None and self._ca_cert_init is not None:
self._ca_key = self._ca_key_init
self._ca_cert = self._ca_cert_init
else:
ca_key, ca_cert, _, _ = ensure_root_ca(base_dir=self._base_dir)
self._ca_key = ca_key
self._ca_cert = ca_cert
self._leaf_cache = _LeafCache(max_size=max(len(self._allowlist), 1))
self._server = await asyncio.start_server(
self._connection_handler,
host=_BIND_HOST,
port=self._port,
)
addr = self._server.sockets[0].getsockname()
logger.info("event=terminator_started address=%s:%d", addr[0], addr[1])
async def _connection_handler(
self,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
) -> None:
await _handle_connect(
reader,
writer,
self._allowlist,
self._dispatch_port,
self_port=self.address[1],
)
@property
def address(self) -> tuple[str, int]:
"""Return (host, port) the server is bound to. Requires :meth:`start`."""
if self._server is None:
raise RuntimeError("Terminator not started")
sock = self._server.sockets[0]
host, port = sock.getsockname()[:2]
return host, port
async def stop(self) -> None:
"""Stop the listener and wait for all connections to close."""
if self._server is not None:
self._server.close()
await self._server.wait_closed()
self._server = None
logger.info("event=terminator_stopped")
async def __aenter__(self) -> AgyCONNECTTerminator:
await self.start()
return self
async def __aexit__(self, *_: object) -> None:
await self.stop()

140
headroom/proxy/handlers/gemini.py
View file

@ -23,11 +23,36 @@ from headroom.proxy.helpers import COMPRESSION_TIMEOUT_SECONDS, extract_tags
from headroom.proxy.identity import resolve_memory_identity
from headroom.proxy.outcome import RequestOutcome
from headroom.proxy.token_counting import gemini_output_tokens
from headroom.transforms.agy_fr_compressor import ( # noqa: F401 (re-exported for existing import sites)
_FR_CCR_HASH_LEN,
_FR_CCR_MARKER_PREFIX,
_FR_CCR_MARKER_TEMPLATE,
_FR_MARKER_MIN_RATIO,
_RETRIEVE_HASH_RE,
_requested_agy_fr_mode,
_scan_hex_hashes,
compress_function_response_leaves,
)
logger = logging.getLogger("headroom.proxy")
DEFAULT_CLOUDCODE_API_URL = "https://cloudcode-pa.googleapis.com"
ANTIGRAVITY_DAILY_API_URL = "https://daily-cloudcode-pa.sandbox.googleapis.com"
ANTIGRAVITY_DAILY_API_URL = "https://daily-cloudcode-pa.googleapis.com"
def _resolve_agy_fr_mode() -> str:
"""Resolve the functionResponse compression mode for an agy run.
``HEADROOM_AGY_FR_MODE`` selects ``ccr`` (default) or ``lossless``;
``lossless`` must be requested explicitly to opt out of savings. When
``ccr`` is in effect but the CCR retrieve listener is not wired for this
run (``HEADROOM_AGY_RETRIEVE_WIRED`` != "1"), we must NOT ship
unrecoverable markers -- downgrade to ``lossless`` (byte-recoverable / no-op).
"""
mode = _requested_agy_fr_mode()
if mode == "ccr" and os.environ.get("HEADROOM_AGY_RETRIEVE_WIRED") != "1":
return "lossless"
return mode
def _usage_int(value: Any, default: int = 0) -> int:
@ -60,19 +85,63 @@ class GeminiHandlerMixin:
def _is_cloudcode_antigravity_request(
self, body: dict[str, Any], headers: dict[str, str]
) -> bool:
"""Detect Pi/OpenClaw antigravity requests routed via Cloud Code Assist."""
"""Detect Pi/OpenClaw and agy antigravity requests routed via Cloud Code Assist.
Detection paths (any one is sufficient):
- body requestType == "agent" (Pi/OpenClaw classic)
- body userAgent == "antigravity" (Pi/OpenClaw classic)
- HTTP User-Agent header starts with "antigravity/" (case-insensitive)
- body model is an agent-model name (e.g. "gemini-3-flash-agent")
- agy-shaped body: top-level model + project + request.contents present
"""
user_agent = headers.get("user-agent", "").lower()
body_user_agent = str(body.get("userAgent", "")).lower()
model = str(body.get("model", ""))
# Agent-model names carry "-agent" suffix (e.g. gemini-3-flash-agent)
is_agent_model = model.endswith("-agent")
# agy-shaped body confirmation: top-level project + request with contents.
# Only meaningful together with is_agent_model; the body shape alone is shared
# with regular Pi/OpenClaw traffic (CLOUDCODE_BODY has the same structure).
request_block = body.get("request", {})
is_agy_agent_body = is_agent_model and (
bool(body.get("project"))
and isinstance(request_block, dict)
and bool(request_block.get("contents"))
)
return (
body.get("requestType") == "agent"
or body_user_agent == "antigravity"
or user_agent.startswith("antigravity/")
or is_agy_agent_body
)
def _resolve_cloudcode_base_url(self, is_antigravity: bool) -> str:
"""Resolve upstream base URL for Pi Cloud Code Assist / Antigravity traffic."""
def _resolve_cloudcode_base_url(
self,
is_antigravity: bool,
original_host: str | None = None,
) -> str:
"""Resolve upstream base URL for Pi Cloud Code Assist / Antigravity traffic.
Resolution order (first match wins):
1. Antigravity path — env HEADROOM_ANTIGRAVITY_API_URL override (an explicit
operator escape hatch, honoured verbatim), else the host the client itself
chose, else the default backend.
2. Reverse-proxy path — ``CLOUDCODE_API_URL`` instance attr or DEFAULT_CLOUDCODE_API_URL.
``original_host`` carries the MITM CONNECT target (the allowlisted host agy
opened the tunnel to). Preserving it matters: the allowlist covers both
``cloudcode-pa`` and ``daily-cloudcode-pa``, so re-originating everything to
the default would send a client's request — and its bearer — to a backend it
never selected. Requests arriving via the reverse-proxy route have no CONNECT
host and fall through to the default.
"""
if is_antigravity:
return ANTIGRAVITY_DAILY_API_URL
override = os.environ.get("HEADROOM_ANTIGRAVITY_API_URL")
if override:
return override.rstrip("/")
from headroom.providers.proxy_targets import cloudcode_host_base
return cloudcode_host_base(original_host or "") or ANTIGRAVITY_DAILY_API_URL
return getattr(self, "CLOUDCODE_API_URL", DEFAULT_CLOUDCODE_API_URL).rstrip("/")
@staticmethod
@ -999,6 +1068,22 @@ class GeminiHandlerMixin:
},
)
def _compress_agy_function_responses(
self,
contents: list[dict],
mode: str,
tokenizer: Any,
store: Any,
) -> tuple[int, int, int]:
"""Uniformly compress functionResponse string leaves across ALL entries.
Thin delegate -- the algorithm now lives in
``headroom.transforms.agy_fr_compressor.compress_function_response_leaves``
(headroom-37g.36) so it is unit-testable standalone. See that
function's docstring for the full behavior contract.
"""
return compress_function_response_leaves(contents, mode, tokenizer, store)
async def handle_google_cloudcode_stream(
self,
request: Request,
@ -1128,6 +1213,33 @@ class GeminiHandlerMixin:
optimized_tokens = original_tokens
transforms_applied = []
# WU1 (headroom-37g.1): uniform deterministic recoverable compression of
# agy functionResponse leaves. Runs INDEPENDENT of the text-pipeline
# revert above so the per-turn tool-output bulk (which lives in preserved
# functionResponse parts the text compressor never sees) is compressed
# and counted even when the tiny residual text inflates and reverts.
fr_before = fr_after = fr_leaves = 0
if is_antigravity and _decision.should_compress and isinstance(contents, list):
try:
fr_mode = _resolve_agy_fr_mode()
fr_store = None
if fr_mode == "ccr":
from headroom.cache.compression_store import get_compression_store
fr_store = get_compression_store()
fr_before, fr_after, fr_leaves = self._compress_agy_function_responses(
contents, fr_mode, tokenizer, fr_store
)
if fr_leaves:
logger.info(
f"[{request_id}] agy functionResponse compression: "
f"mode={fr_mode} leaves={fr_leaves} "
f"tokens {fr_before}->{fr_after} retrieve_wired="
f"{os.environ.get('HEADROOM_AGY_RETRIEVE_WIRED') == '1'}"
)
except Exception as e:
logger.warning(f"[{request_id}] agy functionResponse compression failed: {e}")
if optimized_messages != messages:
optimized_contents, optimized_system = self._messages_to_gemini_contents(
optimized_messages
@ -1144,10 +1256,26 @@ class GeminiHandlerMixin:
request_payload["systemInstruction"] = optimized_system
elif "systemInstruction" in request_payload:
del request_payload["systemInstruction"]
elif fr_leaves:
# Text pipeline reverted (or produced no change) but functionResponse
# leaves were compressed in place. Ship the mutated contents as-is:
# original structure preserved, only FR string leaves replaced. Avoid
# the messages<->contents round-trip (which collapses multi-part text).
request_payload["contents"] = contents
# Fold the functionResponse leaf delta into the accounting so the saving
# ships and is recorded even when the text pipeline reverted. FR tokens
# are disjoint from the text-pipeline counts (messages excludes
# functionResponse), so this never double-counts the #819 waste path.
original_tokens += fr_before
optimized_tokens += fr_after
tokens_saved = original_tokens - optimized_tokens
optimization_latency = (time.time() - start_time) * 1000
base_url = self._resolve_cloudcode_base_url(is_antigravity)
# On the MITM path the Host header still carries the host agy CONNECTed to;
# keep the request on that backend instead of re-originating it.
base_url = self._resolve_cloudcode_base_url(
is_antigravity, original_host=request.headers.get("host")
)
stream_url = f"{base_url}/v1internal:streamGenerateContent"
if request.url.query:
stream_url = f"{stream_url}?{request.url.query}"

5
headroom/proxy/handlers/openai.py
View file

@ -49,6 +49,7 @@ import httpx
from headroom.agent_savings import proxy_pipeline_kwargs
from headroom.ccr.marker_resolution import resolve_markers_in_response
from headroom.ccr.tool_injection import is_headroom_retrieve_name
from headroom.config import unwrap_tool_call_name
from headroom.copilot_auth import (
apply_copilot_api_auth,
@ -2093,9 +2094,7 @@ class OpenAIHandlerMixin:
name = unwrap_tool_call_name(name, item.get("arguments") or item.get("input"))
if isinstance(name, str) and isinstance(call_id, str) and call_id:
function_name_by_call_id[call_id] = name
if isinstance(name, str) and (
name == "headroom_retrieve" or name.endswith("__headroom_retrieve")
):
if is_headroom_retrieve_name(name):
if isinstance(call_id, str) and call_id:
headroom_retrieve_call_ids.add(call_id)

22
headroom/proxy/handlers/streaming.py
View file

@ -149,6 +149,22 @@ class StreamingMixin:
int(cache_creation.get("ephemeral_1h_input_tokens", 0) or 0),
)
@staticmethod
def _gemini_usage_meta(data: dict) -> dict | None:
"""Return a Gemini chunk's ``usageMetadata`` dict, or ``None``.
Native Gemini puts it top-level; Cloud Code Assist (agy) wraps chunks in
a ``response`` envelope (``{"response": {"usageMetadata": {...}}}``),
mirroring the request-side wrap (gemini.py ``body["request"]``). A
non-dict/absent metadata yields ``None`` so callers skip cleanly — a
malformed upstream value can never reach ``.get()`` and crash the parser.
"""
meta = data.get("usageMetadata")
if not isinstance(meta, dict):
response = data.get("response")
meta = response.get("usageMetadata") if isinstance(response, dict) else None
return meta if isinstance(meta, dict) else None
def _parse_sse_usage(self, chunk: bytes, provider: str) -> dict[str, int] | None:
"""Parse usage information from SSE chunk.
@ -221,9 +237,7 @@ class StreamingMixin:
usage["cache_read_input_tokens"] = details.get("cached_tokens", 0)
elif provider == "gemini":
# Gemini sends usageMetadata in each streaming chunk
# Format: {"usageMetadata": {"promptTokenCount": N, "candidatesTokenCount": M}}
usage_meta = data.get("usageMetadata")
usage_meta = self._gemini_usage_meta(data)
if usage_meta:
usage["input_tokens"] = usage_meta.get("promptTokenCount", 0)
usage["output_tokens"] = gemini_output_tokens(usage_meta)
@ -341,7 +355,7 @@ class StreamingMixin:
)
elif provider == "gemini":
usage_meta = data.get("usageMetadata")
usage_meta = self._gemini_usage_meta(data)
if usage_meta:
usage_found["input_tokens"] = usage_meta.get("promptTokenCount", 0)
usage_found["output_tokens"] = gemini_output_tokens(usage_meta)

41
headroom/proxy/outcome.py
View file

@ -1,11 +1,10 @@
"""``RequestOutcome``: the canonical value type for "what happened during
one completed proxy request."
Per the P0 audit (``docs/superpowers/specs/P0-proxy-pipeline-audit.md``),
18 ``metrics.record_request`` call sites across four handler files
disagreed on argument shape — 9 of 18 omitted ``cached=``, 7 of 18
omitted ``attempted_input_tokens=``, only 4 sites emitted a structured
PERF log at all. This module is the structural fix: every handler
An earlier audit found that 18 ``metrics.record_request`` call sites across
four handler files disagreed on argument shape — 9 of 18 omitted ``cached=``,
7 of 18 omitted ``attempted_input_tokens=``, only 4 sites emitted a
structured PERF log at all. This module is the structural fix: every handler
converges on building a :class:`RequestOutcome` at end-of-request and
hands it to :func:`emit_request_outcome` (also exposed as
:meth:`HeadroomProxy._record_request_outcome`), which owns the four
@ -531,6 +530,38 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
savings_attribution=savings_breakdown,
)
# 1b. agy cross-process emit (best-effort, agy process only). When agy runs
# as a separate process from the shared proxy, this drops one inbox
# event carrying the SAME funnel kwargs; the shared proxy drains and
# replays it through its own record_request so the savings land on the
# shared dashboard, counted once. Gated by a marker env var the shared
# proxy never sets, so it never emits. Never raises into the request.
from headroom.proxy import agy_savings_inbox
if agy_savings_inbox.agy_emit_enabled():
try:
agy_savings_inbox.emit_event(
provider=outcome.provider,
model=outcome.model,
input_tokens=outcome.optimized_tokens,
output_tokens=outcome.output_tokens,
tokens_saved=outcome.tokens_saved,
latency_ms=outcome.total_latency_ms,
cached=outcome.cache_hit,
overhead_ms=outcome.overhead_ms,
ttfb_ms=outcome.ttfb_ms,
cache_read_tokens=outcome.cache_read_tokens,
cache_write_tokens=outcome.cache_write_tokens,
cache_write_5m_tokens=outcome.cache_write_5m_tokens,
cache_write_1h_tokens=outcome.cache_write_1h_tokens,
uncached_input_tokens=outcome.uncached_input_tokens,
attempted_input_tokens=outcome.attempted_input_tokens,
project=project,
client=outcome.client,
)
except Exception: # noqa: BLE001 - best-effort, never break the response
pass
# 2. Cost tracker (optional).
cost_tracker = getattr(handler, "cost_tracker", None)
if cost_tracker is not None:

58
headroom/proxy/server.py
View file

@ -2166,9 +2166,6 @@ class HeadroomProxy(
The real implementation lives in ``outcome.py`` as a free function so
test dummies and provider mixins can call it without inheriting from
``HeadroomProxy``.
See ``docs/superpowers/specs/P0-proxy-pipeline-audit.md`` for the
divergence catalog this funnel collapses.
"""
from headroom.proxy.outcome import emit_request_outcome
@ -2420,6 +2417,24 @@ async def _log_toin_stats_periodically(interval_seconds: int = 300) -> None:
logger.debug("Failed to log TOIN stats: %s", e)
async def _drain_agy_savings_periodically(metrics: Any, interval_seconds: int = 5) -> None:
"""Background task: drain the agy cross-process savings inbox on a timer.
Each agy process drops per-request savings events into a canonical inbox;
this replays them through the shared proxy's own ``record_request`` funnel so
they land on the dashboard, counted once. Best-effort — a drain error never
crashes the loop (``drain_inbox`` also never raises out on its own).
"""
from headroom.proxy import agy_savings_inbox
while True:
await asyncio.sleep(interval_seconds)
try:
await agy_savings_inbox.drain_inbox(metrics)
except Exception as e: # noqa: BLE001 - never let a drain error kill the loop
logger.debug("Failed to drain agy savings inbox: %s", e)
def _register_memory_components(proxy: HeadroomProxy, tracker: MemoryTracker) -> None:
"""Register all memory-tracked components with the tracker.
@ -2765,6 +2780,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
app.state.periodic_toin_stats_task = None
app.state.periodic_malloc_trim_task = None
agy_drain_task: asyncio.Task[None] | None = None
try:
try:
previous_handler = _install_loop_exception_handler()
@ -2774,6 +2790,22 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
app.state.periodic_toin_stats_task = asyncio.create_task(
_log_toin_stats_periodically()
)
# Periodically drain the agy cross-process savings inbox so
# agy sessions' savings surface on the shared dashboard.
# Tracked so it is cancelled on shutdown (no leaked task).
#
# A process that EMITS must never DRAIN: `wrap agy` builds
# create_app() in-process (twice — dispatch + retrieve) with its
# own savings paths redirected to a temp dir deleted at exit, so
# a drain loop here would consume events into a sink that is
# thrown away, and race the shared proxy for them.
from headroom.proxy.agy_savings_inbox import agy_emit_enabled
if not agy_emit_enabled():
agy_drain_task = asyncio.create_task(
_drain_agy_savings_periodically(proxy.metrics)
)
# Per-worker on purpose: allocator state is per-process, so
# every worker must trim its own zones (no beacon-owner gate).
if config.periodic_malloc_trim_enabled:
@ -2839,6 +2871,15 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
)
app.state.periodic_toin_stats_task = None
# Cancel the agy savings-drain loop so it does not leak past shutdown.
if agy_drain_task is not None:
agy_drain_task.cancel()
await _timed(
asyncio.gather(agy_drain_task, return_exceptions=True),
label="agy_drain_task.stop",
timeout=3.0,
)
periodic_malloc_trim_task = app.state.periodic_malloc_trim_task
if periodic_malloc_trim_task is not None:
periodic_malloc_trim_task.cancel()
@ -2847,7 +2888,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
label="periodic_malloc_trim.stop",
timeout=3.0,
)
app.state.periodic_malloc_trim_task = None
app.state.periodic_malloc_trim_task = None
if _cc_reconciler is not None:
await _timed(_cc_reconciler.stop(), label="cc_reconciler.stop", timeout=3.0)
@ -4581,6 +4622,15 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
only for loopback callers — the local dashboard. Network callers still
get the aggregate counters but never the per-request metadata.
"""
# Opportunistically drain the agy cross-process savings inbox so the
# dashboard reflects any pending agy events promptly. Best-effort.
try:
from headroom.proxy import agy_savings_inbox
await agy_savings_inbox.drain_inbox(proxy.metrics)
except Exception: # noqa: BLE001 - never let a drain error break /stats
pass
include_sensitive = _request_can_view_dashboard_metadata(
request,
trusted_dashboard_client_cidrs,

448
headroom/transforms/agy_fr_compressor.py Normal file
View file

@ -0,0 +1,448 @@
"""Deterministic, recoverable compression of agy functionResponse leaves.
Moved out of ``GeminiHandlerMixin`` (headroom-37g.36) so the compression
algorithm is unit-testable standalone, without booting the FastAPI app.
Pure move -- no behavior change; ``GeminiHandlerMixin._compress_agy_function_responses``
now delegates to ``compress_function_response_leaves`` below.
---------------------------------------------------------------------------
WU1 (headroom-37g.1): uniform deterministic recoverable compression of agy
functionResponse leaves.
agy's per-turn bulk lives in ``functionResponse`` parts. Those entries carry
non-text parts, so ``_gemini_contents_to_messages`` routes them into
``preserved_indices`` and ``_rebuild_gemini_contents`` restores them verbatim
-- the text compressor never sees them. Only tiny residual text is compressed,
it inflates, the revert guard fires, and tokens_saved collapses to 0 (PR
#1044: "704 -> 718, reverting").
We compress the large STRING leaves inside those parts with a DETERMINISTIC,
IDEMPOTENT, RECOVERABLE transform applied UNIFORMLY to every functionResponse
leaf (historical + tail). Because headroom is an in-flight MITM that never
rewrites agy's LOCAL history, agy re-sends the ORIGINAL bytes each turn; a
deterministic transform (same original -> identical bytes every turn) yields a
byte-stable compressed prefix that re-hits the Cloud Code Assist server-side
cache. Recoverability is mandatory: the model reads functionResponse back as
its own prior tool results, so lossy summaries would corrupt multi-turn
reasoning.
---------------------------------------------------------------------------
"""
from __future__ import annotations
import logging
import os
import re
from typing import Any
from headroom.cache.compression_store import default_ccr_hash
from headroom.ccr.tool_injection import is_headroom_retrieve_name
logger = logging.getLogger("headroom.proxy")
# Marker shipped in place of a compressed leaf (CCR mode). It carries fixed
# prose plus the 24-hex-char CCR hash (SHA-256(original)[:24], the
# compression_store default), which ``headroom_retrieve`` resolves back to the
# original bytes. Self-describing: it NAMES the ``headroom_retrieve`` tool and
# gives a one-line call-to-expand instruction, so a model that needs the
# compressed detail knows how to fetch it (a marker naming no tool led to 0
# retrieve calls in the WU4 live trial). All-ours single-hash form: the hash
# appears exactly once, in the trailing ``Retrieve more: hash=`` form that
# also matches the existing bracketed marker style / regex
# (parser.CCR_RETRIEVAL_MARKER_RE).
_FR_CCR_HASH_LEN = 24
_FR_CCR_MARKER_PREFIX = (
"[functionResponse compressed. Call headroom_retrieve to expand. Retrieve more: hash="
)
_FR_CCR_MARKER_TEMPLATE = _FR_CCR_MARKER_PREFIX + "{hash}]"
# Per-leaf floor DERIVED from marker overhead (not a magic 200). Replacing a
# leaf ships the marker in its place, so the net saving is
# ``leaf_tokens - marker_tokens``. Compressing is only worthwhile when that net
# saving exceeds the marker's OWN cost, i.e. ``leaf_tokens > 2 * marker_tokens``.
# We therefore set the floor to ``_FR_MARKER_MIN_RATIO`` times the marker's
# token cost, computed at runtime against the request's tokenizer.
_FR_MARKER_MIN_RATIO = 2
# WU2-A (headroom-37g.17): agy resends full history with ORIGINAL tool outputs
# every turn (it never rewrites local history to hold headroom's markers). The
# compressor above re-hashes and re-compresses the resent cold original on
# every subsequent turn, so a model that already retrieved a hash via
# ``headroom_retrieve`` is forced to re-retrieve it every turn (observed 236x
# thrash). agy has no call_id, so -- unlike the OpenAI/live_zone.rs path,
# which exempts by call_id (``live_zone.rs:2362-2384``) -- the exemption here
# keys on the retrieved HASH itself: any 24-hex-char token found in the args
# of a functionCall that references ``headroom_retrieve`` is treated as
# "already retrieved this turn" and its matching functionResponse leaf is
# left uncompressed.
_RETRIEVE_HASH_RE = re.compile(r"(?<![0-9a-f])[0-9a-f]{24}(?![0-9a-f])")
def _scan_hex_hashes(value: Any, hashes: set[str]) -> None:
"""Recursively collect 24-hex-char tokens from every STRING value in ``value``."""
if isinstance(value, dict):
for v in value.values():
_scan_hex_hashes(v, hashes)
elif isinstance(value, list):
for v in value:
_scan_hex_hashes(v, hashes)
elif isinstance(value, str):
hashes.update(_RETRIEVE_HASH_RE.findall(value.lower()))
# WU2-A follow-up (headroom-8tm): agy assigns the headroom_retrieve RESULT
# functionResponse a name that matches neither is_headroom_retrieve_name nor
# _args_mention_retrieve (its call args are just {"hash": ...}). So the
# name-based fr exemption AND the functionCall hash-collection both MISS agy's
# retrieve responses -- the resolved envelope re-compresses into a marker every
# turn (model re-retrieves it, L1) AND the ORIGINAL leaf keeps re-retrieving
# (the 236x, L2). We detect the envelope by CONTENT, name-independently, and use
# it to drive both exemptions. Envelope = json.dumps({"hash": <24hex>,
# "source": "local"|"proxy", "original_content": ...}, indent=2) (see
# ccr.mcp_server._retrieve_content); hash + source are LEADING value-bearing
# keys (a source read carries `"hash": hash_key` -- a variable, no 24-hex
# literal -- so it does NOT match), so detection survives a
# `saved to file://` original_content replacement.
_CCR_ENVELOPE_HASH_RE = re.compile(r'"hash"\s*:\s*"([0-9a-f]{24})"(?![0-9a-f])')
_CCR_ENVELOPE_SOURCE_RE = re.compile(r'"source"\s*:\s*"(?:local|proxy)"')
def _ccr_envelope_hash(value: Any) -> str | None:
"""Resolved hash if ``value`` is a headroom_retrieve result envelope, else None.
Name-independent detection of the ``ccr.mcp_server._retrieve_content``
envelope in either the dict form or the JSON-as-text form agy renders,
anchored on the two LEADING value-bearing keys (``hash`` 24-hex + ``source``
local|proxy).
"""
if isinstance(value, dict):
h = value.get("hash")
if (
isinstance(h, str)
and len(h) == _FR_CCR_HASH_LEN
and all(c in "0123456789abcdef" for c in h)
and value.get("source") in ("local", "proxy")
):
return h
return None
if isinstance(value, str):
m = _CCR_ENVELOPE_HASH_RE.search(value)
if m is not None and _CCR_ENVELOPE_SOURCE_RE.search(value) is not None:
return m.group(1)
return None
def _scan_envelope_hashes(value: Any, hashes: set[str]) -> None:
"""Collect resolved hashes from any headroom_retrieve envelope in ``value``.
L2 (headroom-8tm): the envelope carries the hash the model just retrieved;
adding it to ``retrieved_hashes`` exempts the ORIGINAL leaf agy resends
(via the existing exemption in ``_walk_fr_compress``).
"""
h = _ccr_envelope_hash(value)
if h is not None:
hashes.add(h)
return # envelope found; its original_content is resolved bytes, not another envelope
if isinstance(value, dict):
for v in value.values():
_scan_envelope_hashes(v, hashes)
elif isinstance(value, list):
for v in value:
_scan_envelope_hashes(v, hashes)
def _requested_agy_fr_mode() -> str:
"""Normalize the REQUESTED functionResponse mode from the environment.
``HEADROOM_AGY_FR_MODE`` selects ``lossless`` or ``ccr`` (default);
unset/invalid values fall back to ``ccr``. Single source of truth
shared by ``_resolve_agy_fr_mode`` (the downgrade decision) and the
wrap-agy downgrade warning (``headroom.cli.wrap._maybe_warn_agy_ccr_downgrade``)
so the two cannot drift.
"""
mode = (os.environ.get("HEADROOM_AGY_FR_MODE") or "ccr").strip().lower()
if mode not in ("ccr", "lossless"):
return "ccr"
return mode
def _fr_marker_tokens_and_floor(tokenizer: Any) -> tuple[int, int]:
"""Derive the CCR marker's token cost and the per-leaf compression floor.
A compressed leaf ships the marker in its place, so the net saving is
``leaf_tokens - marker_tokens``. We only compress when that saving
exceeds the marker's own cost (``_FR_MARKER_MIN_RATIO`` x marker).
Returns ``(marker_body_tokens, floor)``: both derive from ONE
``count_text`` call on the placeholder marker, so callers can thread
``marker_body_tokens`` down to every leaf instead of recomputing the
identical value per leaf (headroom-37g.35).
"""
marker_body_tokens = tokenizer.count_text(
_FR_CCR_MARKER_TEMPLATE.format(hash="0" * _FR_CCR_HASH_LEN)
)
floor = int(max(1, marker_body_tokens * _FR_MARKER_MIN_RATIO))
return marker_body_tokens, floor
def _compress_fr_leaf(
leaf: str,
mode: str,
store: Any,
tool_name: str | None,
*,
leaf_tokens: int,
marker_body_tokens: int,
hash_key: str,
) -> str:
"""Deterministically compress a single functionResponse string leaf.
``ccr``: cache the ORIGINAL and ship a hash marker. ``hash_key`` is
``default_ccr_hash(leaf)`` -- SHA-256(original)[:24] -- computed ONCE by
the caller (``_walk_fr_compress``, which also uses it for the retrieve
exemption check) and passed here as ``explicit_hash``. This is byte-for-
byte identical to the store's own implicit default (``compression_store.
store()`` falls back to ``default_ccr_hash(original)`` when no
``explicit_hash`` is given), so marker/store bytes are unchanged; we just
avoid a second identical hash + a second identical ``count_text(leaf)``
inside the store call. Idempotent: an already-compressed marker is
returned unchanged.
``lossless``: format-native reversible compaction (no marker).
"""
if mode == "ccr":
# Idempotency guard: never re-wrap our own marker.
if leaf.startswith(_FR_CCR_MARKER_PREFIX):
return leaf
store.store(
leaf,
_FR_CCR_MARKER_TEMPLATE,
original_tokens=leaf_tokens,
compressed_tokens=marker_body_tokens,
tool_name=tool_name,
explicit_hash=hash_key,
)
return _FR_CCR_MARKER_TEMPLATE.format(hash=hash_key)
# lossless: reversible, deterministic, self-verified smaller-or-unchanged.
from headroom.transforms.lossless_compaction import compact_lossless
return compact_lossless(leaf, "text")
def _args_mention_retrieve(value: Any) -> bool:
"""Bounded recursive scan for a ``"headroom_retrieve"`` substring.
Replaces the ``"headroom_retrieve" in json.dumps(args)`` substring test
with a direct walk of the ``args`` structure -- no serialization. Scans
the SAME surface ``json.dumps`` would have covered: dict KEYS (JSON
object keys are strings) and dict/list VALUES, recursively, short-
circuiting on first match. Keep the match set identical to the old
``json.dumps`` scan -- do not tighten it.
"""
if isinstance(value, dict):
for k, v in value.items():
if isinstance(k, str) and "headroom_retrieve" in k:
return True
if _args_mention_retrieve(v):
return True
return False
if isinstance(value, list):
return any(_args_mention_retrieve(v) for v in value)
if isinstance(value, str):
return "headroom_retrieve" in value
return False
def _collect_retrieved_hashes(contents: list[dict]) -> set[str]:
"""Collect CCR hashes the model already retrieved via ``headroom_retrieve``.
Scans every ``functionCall`` part across ALL of ``contents`` (any
entry, not just the tail -- agy resends the full history every turn)
for calls that reference ``headroom_retrieve`` (bare name, or the
generic MCP dispatch shape e.g. ``call_mcp_tool`` whose args mention
``headroom_retrieve``), then recursively pulls every 24-hex-char
token out of that call's ``args``. See the WU2-A comment above
``_RETRIEVE_HASH_RE`` for why this keys on the hash rather than a
call_id (agy has none).
"""
hashes: set[str] = set()
for content in contents:
if not isinstance(content, dict):
continue
parts = content.get("parts")
if not isinstance(parts, list):
continue
for part in parts:
if not isinstance(part, dict):
continue
fc = part.get("functionCall")
if isinstance(fc, dict):
name = fc.get("name", "")
args = fc.get("args") or {}
if is_headroom_retrieve_name(name) or _args_mention_retrieve(args):
_scan_hex_hashes(args, hashes)
# L2 (headroom-8tm): agy's opaque retrieve fr name defeats the
# functionCall-based collection above, so recover the resolved hash
# from the retrieve-result envelope the model already received --
# the ORIGINAL leaf agy resends then hits the exemption in
# ``_walk_fr_compress``.
fr = part.get("functionResponse")
if isinstance(fr, dict):
_scan_envelope_hashes(fr.get("response"), hashes)
return hashes
def _walk_fr_compress(
value: Any,
mode: str,
tokenizer: Any,
store: Any,
floor: int,
marker_body_tokens: int,
tool_name: str | None,
stats: dict[str, int],
retrieved_hashes: set[str],
) -> Any:
"""Recurse dict/list; compress every string leaf >= ``floor`` in place.
Non-string scalars and sub-floor leaves are skipped. A leaf whose
default CCR hash is in ``retrieved_hashes`` is exempt (WU2-A: the
model already retrieved it this turn; re-compressing it would force
an endless re-retrieve loop). Mutates containers in place and
returns ``value`` for convenient reassignment.
"""
if isinstance(value, dict):
if _ccr_envelope_hash(value) is not None:
stats["fr_envelope_exempt"] = stats.get("fr_envelope_exempt", 0) + 1
return value # L1: headroom_retrieve envelope -- never re-compress (headroom-8tm)
for k, v in value.items():
value[k] = _walk_fr_compress(
v,
mode,
tokenizer,
store,
floor,
marker_body_tokens,
tool_name,
stats,
retrieved_hashes,
)
return value
if isinstance(value, list):
for i, v in enumerate(value):
value[i] = _walk_fr_compress(
v,
mode,
tokenizer,
store,
floor,
marker_body_tokens,
tool_name,
stats,
retrieved_hashes,
)
return value
if isinstance(value, str):
try:
leaf_tokens = tokenizer.count_text(value)
if leaf_tokens < floor:
return value
if _ccr_envelope_hash(value) is not None:
stats["fr_envelope_exempt"] = stats.get("fr_envelope_exempt", 0) + 1
return value # L1: retrieve envelope as text -- never re-compress (headroom-8tm)
hash_key = default_ccr_hash(value)
if hash_key in retrieved_hashes:
return value # exempt: model already retrieved this hash (live_zone.rs parity)
new_leaf = _compress_fr_leaf(
value,
mode,
store,
tool_name,
leaf_tokens=leaf_tokens,
marker_body_tokens=marker_body_tokens,
hash_key=hash_key,
)
except Exception:
# Broad by design: one malformed leaf must not abort the whole
# walk and strand earlier leaves half-compressed in the shared
# `contents` object. Reachable from untrusted tool output -- e.g.
# a lone UTF-16 surrogate makes default_ccr_hash's str.encode()
# raise UnicodeEncodeError. Leave this one leaf verbatim, keep going.
logger.warning(
"agy FR: leaving one functionResponse leaf uncompressed (failed to hash/compress)",
exc_info=True,
)
return value
if new_leaf != value:
new_tokens = tokenizer.count_text(new_leaf)
# Guard: only accept an actual reduction (lossless may no-op).
if new_tokens < leaf_tokens:
stats["before"] += leaf_tokens
stats["after"] += new_tokens
stats["leaves"] += 1
return new_leaf
return value
# Non-string scalar (int/float/bool/None): skipped, JSON shape preserved.
return value
def compress_function_response_leaves(
contents: list[dict],
mode: str,
tokenizer: Any,
store: Any,
) -> tuple[int, int, int]:
"""Uniformly compress functionResponse string leaves across ALL entries.
Walks every ``contents[]`` entry (historical + tail), every ``parts[]``
entry, and every ``functionResponse`` part (an entry may carry several),
recursing into the ``response`` value to compress its large string leaves
in place. ``functionCall`` parts are never touched; JSON shape and
functionCall/functionResponse pairing are preserved.
EXEMPTION: a functionResponse named ``headroom_retrieve`` (bare or
MCP-namespaced, see ``is_headroom_retrieve_name``) is left untouched.
That tool's own output is the just-resolved ORIGINAL of a marker the
model expanded; re-compressing it back into the same marker is a
self-defeating loop (the OpenAI path already exempts this -- see
``headroom_retrieve_call_ids`` in ``live_zone.rs``).
Returns ``(fr_tokens_before, fr_tokens_after, leaves_compressed)`` over
the leaves that were actually compressed.
"""
marker_body_tokens, floor = _fr_marker_tokens_and_floor(tokenizer)
stats: dict[str, int] = {"before": 0, "after": 0, "leaves": 0, "fr_envelope_exempt": 0}
retrieved_hashes = _collect_retrieved_hashes(contents)
for content in contents:
if not isinstance(content, dict):
continue
parts = content.get("parts")
if not isinstance(parts, list):
continue
for part in parts:
if not isinstance(part, dict):
continue
fr = part.get("functionResponse")
if not isinstance(fr, dict):
continue
response = fr.get("response")
if response is None:
continue
if is_headroom_retrieve_name(fr.get("name")):
continue
fr["response"] = _walk_fr_compress(
response,
mode,
tokenizer,
store,
floor,
marker_body_tokens,
fr.get("name"),
stats,
retrieved_hashes,
)
if stats["fr_envelope_exempt"]:
logger.info(
"agy FR: exempted %d headroom_retrieve envelope leaf(s) from re-compression",
stats["fr_envelope_exempt"],
)
return stats["before"], stats["after"], stats["leaves"]

5
headroom/transforms/compression_units.py
View file

@ -13,6 +13,7 @@ from collections.abc import Iterable
from dataclasses import dataclass, field, replace
from typing import Protocol
from ..parser import CCR_MARKER_ALTERNATION
from .content_router import (
CompressionStrategy,
ContentRouter,
@ -108,9 +109,7 @@ class RoutedCompressionUnit:
slot: object
_CCR_MARKER_RE = re.compile(
r"(?m)^.*(?:Retrieve more: hash=|Retrieve original: hash=|<<ccr:[^>]+>>).*$"
)
_CCR_MARKER_RE = re.compile(rf"(?m)^.*(?:{CCR_MARKER_ALTERNATION}).*$")
_LOSSY_UNMARKED_STRATEGIES = {
CompressionStrategy.KOMPRESS.value,

3
llms.txt
View file

@ -21,7 +21,8 @@ The canonical, always-current documentation index lives at the docs site below.
- TypeScript / Node: `npm install headroom-ai` (or `pnpm add headroom-ai`, `bun add headroom-ai`)
- Docker: `docker run -p 8787:8787 ghcr.io/headroomlabs-ai/headroom:latest`
- Run the proxy: `headroom proxy --port 8787` then point any client at `http://127.0.0.1:8787`
- Wrap an agent in one command: `headroom wrap claude` (also: `codex`, `copilot`, `cursor`, `aider`, `opencode`, `cline`, `continue`, `goose`, `openhands`, `openclaw`, `vibe`, `omp`)
- Wrap an agent in one command: `headroom wrap claude` (also: `codex`, `copilot`, `cursor`, `aider`, `opencode`, `cline`, `continue`, `goose`, `openhands`, `openclaw`, `vibe`, `omp`, `agy`)
- Wrap agy (Google Antigravity CLI) with TLS-MITM transport: `headroom wrap agy` (pass `--no-intercept` to skip interception; `headroom unwrap agy` reverts GEMINI.md and MCP config)
## Entry points

4
pyproject.toml
View file

@ -92,6 +92,8 @@ proxy = [
"transformers>=5.5.0,<6.0", # Tokenizer only (for Kompress)
"watchdog>=4.0.0", # File watcher for live code graph reindexing (--code-graph)
"sqlite-vec>=0.1.6", # Vector index for memory (--memory). Lightweight, no torch.
"cryptography>=42.0.0", # Root CA + leaf minting for agy TLS-MITM wrap (headroom wrap agy)
"hypercorn>=0.16", # In-process HTTPS dispatch server for agy MITM transport
]
# Production ASGI/WSGI server — Unix-only (gunicorn does not support Windows).
# Kept separate from [proxy] so that dev, CI, and Windows users are not forced
@ -296,6 +298,8 @@ dev = [
"sentence-transformers>=2.2.0,<6.0",
"numpy>=1.24.0",
"openpyxl>=3.1.0", # exercises spreadsheet_ingest (.xlsx) in the test suite
"cryptography>=42.0.0", # agy TLS-MITM CA/leaf minting (tests: test_agy_ca/dispatch/terminator)
"hypercorn>=0.16", # agy in-process HTTPS dispatch server (tests: test_agy_dispatch)
"respx>=0.20.0", # HTTP mock transport for passthrough handler tests
]
# All optional dependencies (everything you need).

47
tests/conftest.py
View file

@ -85,6 +85,53 @@ def _isolate_mcp_ledger(monkeypatch, tmp_path_factory):
monkeypatch.setattr(ledger, "ledger_path", lambda: ledger_file)
# Any test that builds the proxy app (`create_app()`) runs its lifespan startup,
# which calls `registry.start_all()` and starts the subscription tracker. That
# tracker polls `https://api.anthropic.com/api/oauth/usage` for real — observed
# on every local run of tests/test_agy_dispatch.py as
# `httpx - INFO - HTTP Request: GET https://api.anthropic.com/api/oauth/usage`.
# Same class of bug as the beacon and MCP-ledger fixtures above: a unit test
# reaching a live external service.
#
# There is no env knob that reaches this from a test: `ProxyConfig.
# subscription_tracking_enabled` defaults to True and HEADROOM_NO_SUBSCRIPTION_
# TRACKING is wired only through the `headroom proxy` CLI, so `create_app()`
# always starts the poller. Force `enabled=False` at the tracker factory
# instead, which is the single seam every caller funnels through.
#
# Beyond hygiene this is a determinism fix: the poll made first-request latency
# depend on live network egress, which timed out the 10s guard in
# test_dispatch_server_tls_and_route on the Windows CI lane (no offline env
# there, unlike the Linux shards' HF_HUB_OFFLINE) while passing on Linux.
@pytest.fixture(autouse=True)
def _disable_subscription_polling(monkeypatch):
# Same guard as the sibling fixtures: the macos/windows-native-wrapper CI
# jobs install only pytest and drive the installer shell scripts via
# subprocess, so headroom isn't importable and there is no tracker to
# disable.
try:
from headroom.subscription import tracker as _tracker
except ModuleNotFoundError:
return
real_configure = _tracker.configure_subscription_tracker
def _configure_disabled(*args, **kwargs):
kwargs["enabled"] = False
return real_configure(*args, **kwargs)
monkeypatch.setattr(_tracker, "configure_subscription_tracker", _configure_disabled)
# server.py imports the symbol directly (`from ... import
# configure_subscription_tracker`), so patching only the defining module
# would leave that already-bound reference pointing at the real one.
try:
from headroom.proxy import server as _server
except ModuleNotFoundError:
return
if hasattr(_server, "configure_subscription_tracker"):
monkeypatch.setattr(_server, "configure_subscription_tracker", _configure_disabled)
# The Copilot "routed to Copilot" flag is a module-global ContextVar that
# build_copilot_upstream_url() sets as a side effect. Unit tests that call that
# builder directly (or otherwise run in the shared root context) would leave it

1299
tests/test_agy_ca.py Normal file
View file

File diff suppressed because it is too large Load diff

238
tests/test_agy_ccr_downgrade_warning.py Normal file
View file

@ -0,0 +1,238 @@
"""Tests for the loud ccr->lossless downgrade warning in ``headroom wrap agy``.
Originally written test-first (red) before ``_maybe_warn_agy_ccr_downgrade``
existed in ``headroom/cli/wrap.py``; the implementation has since landed.
Scope (headroom-svf; ccr-default per headroom-37g.32): when ``headroom wrap
agy`` runs in ``ccr`` mode (now the default -- unset/invalid resolve to ccr)
but the retrieve MCP could NOT be wired for the run, the Cloud Code Assist handler
(``headroom.proxy.handlers.gemini._resolve_agy_fr_mode``) silently downgrades
functionResponse compression to ``lossless`` (a no-op), so tool-output
savings collapse to ~0 with no user-visible warning. This must become loud
and actionable, with best-effort cause detection:
* ``mcp`` not importable in *this* (parent) interpreter -> ADVISORY hint to
install ``headroom-ai[proxy]`` (the agy child is resolved via
``shutil.which("headroom")`` and need NOT share this venv, hence a
likely-cause hint, not certainty).
* ``mcp`` importable -> the retrieve MCP failed to register or complete its
handshake; point at the ``MCP retrieve tool:`` failure line already printed
to the console (the agy path runs in-process servers and writes no
``proxy.log``).
The warning fires whenever the resolved mode is ccr (explicit, OR the
unset/invalid default) AND the retrieve MCP did not wire. It must stay silent
when retrieve DID wire, or when ``lossless`` was requested explicitly (no
downgrade occurred).
"""
from __future__ import annotations
from types import SimpleNamespace
from typing import Any
import pytest
import headroom.cli.wrap as wrap_mod
from headroom.cli.wrap import _maybe_warn_agy_ccr_downgrade
class TestMaybeWarnAgyCcrDowngrade:
# ------------------------------------------------------------------
# Gating: fires only for ccr + not-wired.
# ------------------------------------------------------------------
def test_silent_when_retrieve_registered(
self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
monkeypatch.delenv("HEADROOM_AGY_FR_MODE", raising=False)
_maybe_warn_agy_ccr_downgrade(retrieve_wired=True)
out = capsys.readouterr().out
assert out == ""
def test_silent_when_lossless_requested(
self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "lossless")
_maybe_warn_agy_ccr_downgrade(retrieve_wired=False)
out = capsys.readouterr().out
assert out == ""
def test_warns_when_unset_defaults_to_ccr_and_not_registered(
self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
# ccr is now the default (WU-CCRDEFAULT): unset + not-wired downgrades,
# so the warning must fire (previously silent when lossless was default).
monkeypatch.delenv("HEADROOM_AGY_FR_MODE", raising=False)
_maybe_warn_agy_ccr_downgrade(retrieve_wired=False)
out = capsys.readouterr().out
assert "DISABLED" in out
def test_warns_when_explicit_ccr_and_not_registered(
self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "ccr")
_maybe_warn_agy_ccr_downgrade(retrieve_wired=False)
out = capsys.readouterr().out
assert "DISABLED" in out
def test_invalid_mode_value_treated_as_ccr_default(
self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
# Mirrors _requested_agy_fr_mode's fallback-to-ccr for garbage values:
# invalid -> ccr default -> not-wired -> warns.
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "bogus")
_maybe_warn_agy_ccr_downgrade(retrieve_wired=False)
out = capsys.readouterr().out
assert "DISABLED" in out
# ------------------------------------------------------------------
# Cause detection: in-parent `mcp` importability drives the branch.
# The fakes are NAME-SENSITIVE (keyed on the probed module name), so the
# branch tests also verify the probe asks about "mcp" specifically.
# ------------------------------------------------------------------
def test_mcp_missing_branch_recommends_proxy_extra(
self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
# ccr must be requested explicitly now that lossless is the default.
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "ccr")
# False ONLY for "mcp": probing any other name would flip the branch.
monkeypatch.setattr("headroom.cli.wrap._module_available", lambda name: name != "mcp")
_maybe_warn_agy_ccr_downgrade(retrieve_wired=False)
out = capsys.readouterr().out
assert "headroom-ai[proxy]" in out
assert "pip install mcp" in out
# Advisory caveat: parent-mcp-present/absent doesn't guarantee child state.
assert "ADVISORY" in out or "likely cause" in out
assert "proxy.log" not in out
def test_mcp_present_branch_points_at_console_failure_line(
self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
# ccr must be requested explicitly now that lossless is the default.
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "ccr")
# True ONLY for "mcp": probing any other name would flip the branch.
monkeypatch.setattr("headroom.cli.wrap._module_available", lambda name: name == "mcp")
_maybe_warn_agy_ccr_downgrade(retrieve_wired=False)
out = capsys.readouterr().out
# The agy path runs in-process servers and writes NO proxy.log; the
# handshake failure detail is the "MCP retrieve tool:" console line.
assert "MCP retrieve tool:" in out
# Cause text broadened for the exposure gate: handshake-OK-but-uncached
# is now a distinct downgrade reason alongside register/handshake failure.
assert "did not register/handshake" in out
assert "exposed it as a callable tool" in out
assert "proxy.log" not in out
assert "headroom-ai[proxy]" not in out
class TestAgyCallSiteWiring:
"""Prove ``agy()`` actually invokes the warning on the downgrade path.
All helper tests above exercise ``_maybe_warn_agy_ccr_downgrade`` in
isolation; without this test, deleting the call site inside ``agy()``
would leave the suite green — the exact silent-downgrade regression this
feature exists to prevent. Behavioral spy: the call site raising through
the spy aborts ``agy()`` before it would exec the agy binary.
"""
def test_agy_invokes_downgrade_warning_when_retrieve_not_wired(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
for key in (
"HEADROOM_AGY_FR_MODE",
"HEADROOM_AGY_RETRIEVE_WIRED",
"HEADROOM_BACKEND",
"HEADROOM_SAVINGS_PATH",
"HEADROOM_SAVINGS_EVENTS_PATH",
"HEADROOM_OTEL_METRICS_ENABLED",
"HEADROOM_AGY_INBOX_EMIT",
):
monkeypatch.delenv(key, raising=False)
# -- Binary resolution: agy "installed", rtk absent. ---------------
monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/agy" if name == "agy" else None)
# -- CA / child-env plumbing: no real crypto, no real env build. ----
monkeypatch.setattr(
"headroom.proxy.agy_ca.ensure_root_ca", lambda: (None, None, None, None)
)
monkeypatch.setattr("headroom.proxy.agy_ca.build_combined_bundle", lambda: "/dev/null")
monkeypatch.setattr("headroom.providers.agy.build_agy_env", lambda **kwargs: {})
# -- Session stats / fail-open observability: inert fakes. ----------
class _FakeStats:
def snapshot_start(self) -> None:
pass
def print_summary(self, handler: Any) -> None:
pass
monkeypatch.setattr("headroom.providers.agy.stats.AgySessionStats", _FakeStats)
monkeypatch.setattr("headroom.providers.agy.stats.install_fail_open_handler", lambda: None)
monkeypatch.setattr(
"headroom.providers.agy.stats.remove_fail_open_handler", lambda handler: None
)
# -- MCP registrar + tooling setup: inert fakes. --------------------
class _FakeRegistrar:
name = "agy"
def register_server(self, spec: Any, force: bool = False) -> Any:
raise AssertionError("register_server must not be reached in this test")
def unregister_server(self, name: str) -> bool:
return False
monkeypatch.setattr("headroom.mcp_registry.agy.AgyRegistrar", _FakeRegistrar)
monkeypatch.setattr(
"headroom.cli.wrap._disable_tokensave_mcp", lambda *args, **kwargs: None
)
monkeypatch.setattr("headroom.cli.wrap._disable_serena_mcp", lambda *args, **kwargs: None)
# -- In-process servers: fake handle with a live retrieve port. -----
fake_servers = SimpleNamespace(
terminator=SimpleNamespace(address=("127.0.0.1", 1)), retrieve_port=12345
)
monkeypatch.setattr(
"headroom.cli.wrap._start_agy_servers", lambda *args, **kwargs: fake_servers
)
monkeypatch.setattr("headroom.cli.wrap._stop_agy_servers", lambda servers: None)
# -- Downgrade scenario: retrieve MCP does not wire. -----------------
monkeypatch.setattr(
"headroom.cli.wrap._setup_headroom_retrieve_mcp_agy",
lambda *args, **kwargs: False,
)
# -- Spy: record the call, abort agy() before it would exec agy. -----
calls: list[bool] = []
def _spy(retrieve_registered: bool) -> None:
calls.append(retrieve_registered)
raise SystemExit(0)
monkeypatch.setattr("headroom.cli.wrap._maybe_warn_agy_ccr_downgrade", _spy)
# -- Guard: if the call site is ever removed, never exec a binary. ---
monkeypatch.setattr("subprocess.run", lambda *args, **kwargs: SimpleNamespace(returncode=0))
# _register_proxy_client writes a durable marker under workspace_dir()
# (~/.headroom/proxy_clients/<port>/); stub it so this test never touches
# the real client registry (conftest provides no HOME isolation).
monkeypatch.setattr("headroom.cli.wrap._register_proxy_client", lambda *a, **k: None)
with pytest.raises(SystemExit):
wrap_mod.agy.callback(
port=8899,
no_proxy=True,
no_intercept=False,
backend=None,
no_mcp=False,
no_serena=True,
no_tokensave=True,
code_graph=False,
agy_args=(),
)
assert calls == [False]

206
tests/test_agy_ccr_retrieve_roundtrip.py Normal file
View file

@ -0,0 +1,206 @@
"""headroom-vb5: deterministic proof that the CCR retrieve MECHANISM
resolves a ccr-compressed functionResponse marker's hash back to the
byte-identical original leaf.
Scope: this module proves the RECOVERY MECHANISM only --
1. WU1's ccr compressor (``HeadroomProxy._compress_agy_function_responses``)
really does replace a functionResponse leaf with a self-describing
``headroom_retrieve`` marker (single-hash, ``Retrieve more: hash=...``
form), and the original bytes are gone from the shipped payload.
2. ``CompressionStore.retrieve(hash)`` resolves that hash back to the
byte-identical original via TWO independent paths that mirror the real
``headroom mcp serve`` child:
a) PRIMARY -- two ``SQLiteBackend`` handles opened on the SAME on-disk
ccr_store.db file (the child's local-resolution path, since proxy and
child share one sqlite file).
b) SECONDARY -- the ``POST /v1/retrieve`` HTTP endpoint that
``_retrieve_via_proxy`` (headroom/ccr/mcp_server.py) falls back to when
local resolution misses (memory backend / workspace mismatch / sqlite
init failure).
3. A bogus hash never produces a false recovery.
OUT OF SCOPE (explicitly, so the finding is owned and not dropped): whether
the MODEL actually chooses to *emit* a ``headroom_retrieve`` tool call when
it sees a marker in context ("0 retrieve calls" in the live trial) is MODEL
BEHAVIOR, not a mechanism defect. That is owned by headroom-y4q.
Fully deterministic and hermetic: no live agy, no Cloud Code Assist network,
no ``:8787`` proxy. Only ``fastapi.testclient.TestClient`` (in-process ASGI)
and in-process ``CompressionStore``/``SQLiteBackend`` handles. An autouse
fixture isolates the process-global CCR store to a per-test tmp SQLite file
so no test ever touches the real ``~/.headroom`` workspace db.
"""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
import pytest
from fastapi.testclient import TestClient
from headroom.cache.backends import InMemoryBackend, SQLiteBackend
from headroom.cache.compression_store import (
CompressionStore,
get_compression_store,
reset_compression_store,
)
from headroom.proxy.handlers.gemini import _FR_CCR_MARKER_PREFIX
from headroom.proxy.server import ProxyConfig, create_app
from headroom.tokenizers import get_tokenizer
_MODEL = "gemini-3-flash-agent"
# A unique needle so we can prove it is (a) absent from the shipped marker
# bytes and (b) present, byte-identically, in whatever retrieve() returns.
NEEDLE = "UNIQUE-NEEDLE-c9f3a7d1-92be-4e6a-8c31-roundtrip-marker"
# Large, single-line leaf (no repeated lines, so lossless compaction would be
# a no-op) well above the marker-derived compression floor -- mirrors the
# BIG_LEAF fixture in tests/test_agy_functionresponse_compression.py.
BIG_LEAF = "search result row alpha beta gamma delta epsilon zeta eta " * 40 + NEEDLE
def _fr_entry(role: str = "user", leaf: Any = BIG_LEAF, name: str = "search") -> dict:
"""Build a historical (non-tail) functionResponse contents[] entry."""
return {
"role": role,
"parts": [{"functionResponse": {"name": name, "response": {"output": leaf}}}],
}
def _fr_leaf(contents: list, entry: int = 0, part: int = 0, key: str = "output") -> Any:
return contents[entry]["parts"][part]["functionResponse"]["response"][key]
def _hash_of(marker: str) -> str:
"""Extract the hash from a ``headroom_retrieve`` marker.
Mirrors ``_hash_of`` in test_agy_functionresponse_compression.py:
single-hash marker, in the ``Retrieve more: hash=<h>]`` form that
``parser.CCR_RETRIEVAL_MARKER_RE`` keys on.
"""
assert marker.startswith(_FR_CCR_MARKER_PREFIX), marker
return marker.split("hash=", 1)[1].rstrip("]")
@pytest.fixture(autouse=True)
def _isolate_global_compression_store(
tmp_path_factory: pytest.TempPathFactory, monkeypatch: pytest.MonkeyPatch
) -> Any:
"""Isolate the process-global CCR store for every test in this module.
``create_app()`` lazily calls ``get_compression_store()`` at startup
(memory-tracker registration), which would otherwise bind the global
singleton to the real ``workspace_dir()/ccr_store.db``. Point it at a
per-test tmp file instead and reset the singleton around the test so no
test touches real on-disk state or leaks into another test.
"""
db_path = tmp_path_factory.mktemp("ccr") / "global_ccr.db"
monkeypatch.setenv("HEADROOM_CCR_SQLITE_PATH", str(db_path))
monkeypatch.delenv("HEADROOM_CCR_BACKEND", raising=False)
reset_compression_store()
yield
reset_compression_store()
@pytest.fixture
def proxy() -> Any:
"""A HeadroomProxy instance exposing ``_compress_agy_function_responses``."""
with TestClient(create_app(ProxyConfig(optimize=True))) as client:
yield client.app.state.proxy # type: ignore[attr-defined]
@pytest.fixture
def tok() -> Any:
return get_tokenizer(_MODEL)
def test_marker_replaces_needle_in_shipped_bytes(proxy: Any, tok: Any) -> None:
"""WU1's ccr compressor ships a marker in place of the leaf; the NEEDLE
must NOT be literally present anywhere in the shipped contents[] bytes."""
store = CompressionStore(backend=InMemoryBackend())
contents = [_fr_entry()]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, store)
assert leaves == 1
assert before > after
marker = _fr_leaf(contents)
assert marker.startswith(_FR_CCR_MARKER_PREFIX)
assert marker != BIG_LEAF
shipped_bytes = json.dumps(contents)
assert NEEDLE not in shipped_bytes
def test_shared_sqlite_two_handles_resolve_byte_identical(
proxy: Any, tok: Any, tmp_path: Path
) -> None:
"""PRIMARY: two independent SQLiteBackend handles opened on the SAME
ccr_store.db file -- the real local-resolution path the ``headroom mcp
serve`` child uses, since proxy and child open the same sqlite file in
one interpreter. NOT a claim of OS-level cross-process resolution."""
db_path = tmp_path / "shared_ccr.db"
store_a = CompressionStore(backend=SQLiteBackend(db_path))
contents = [_fr_entry()]
proxy._compress_agy_function_responses(contents, "ccr", tok, store_a)
marker = _fr_leaf(contents)
hash_key = _hash_of(marker)
# SECOND handle, independently opened on the SAME sqlite file.
store_b = CompressionStore(backend=SQLiteBackend(db_path))
entry = store_b.retrieve(hash_key)
assert entry is not None
assert entry.original_content == BIG_LEAF # byte-identical
assert NEEDLE in entry.original_content
def test_http_retrieve_fallback_byte_identical(proxy: Any, tok: Any) -> None:
"""SECONDARY: HTTP fallback via POST /v1/retrieve -- the path
``_retrieve_via_proxy`` (headroom/ccr/mcp_server.py) uses when local
resolution misses. We force an empty LOCAL store (a fresh in-memory
handle that never saw this hash, e.g. memory backend / workspace
mismatch) and resolve via the proxy's HTTP endpoint instead, which is
backed by the same process-global store the compressor wrote to."""
contents = [_fr_entry()]
proxy._compress_agy_function_responses(contents, "ccr", tok, get_compression_store())
marker = _fr_leaf(contents)
hash_key = _hash_of(marker)
# Local store miss: a fresh, unrelated in-memory store never populated
# with this hash. This is the condition that forces the HTTP fallback.
empty_local_store = CompressionStore(backend=InMemoryBackend())
assert empty_local_store.retrieve(hash_key) is None
app = create_app(ProxyConfig(optimize=True))
with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client:
response = client.post("/v1/retrieve", json={"hash": hash_key})
assert response.status_code == 200
body = response.json()
assert body["original_content"] == BIG_LEAF # byte-identical
assert NEEDLE in body["original_content"]
def test_bogus_hash_returns_no_recovery(proxy: Any, tok: Any) -> None:
"""A bogus hash (well-formed hex, never stored) must not resolve --
neither locally nor via the HTTP endpoint. No false recovery."""
# Populate the store with something so the store is non-empty, then ask
# for a hash that was never returned by store().
contents = [_fr_entry()]
proxy._compress_agy_function_responses(contents, "ccr", tok, get_compression_store())
real_hash = _hash_of(_fr_leaf(contents))
bogus_hash = "0" * len(real_hash)
assert bogus_hash != real_hash
assert get_compression_store().retrieve(bogus_hash) is None
app = create_app(ProxyConfig(optimize=True))
with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client:
response = client.post("/v1/retrieve", json={"hash": bogus_hash})
assert response.status_code == 404

1444
tests/test_agy_dispatch.py Normal file
View file

File diff suppressed because it is too large Load diff

101
tests/test_agy_fr_compressor_unit.py Normal file
View file

@ -0,0 +1,101 @@
"""headroom-37g.36: standalone unit coverage for the moved agy FR compressor.
Proves the algorithm is unit-testable directly from
``headroom.transforms.agy_fr_compressor`` without booting the FastAPI app
(no ``create_app`` / ``TestClient``) -- the altitude payoff of the pure move
out of ``GeminiHandlerMixin``.
"""
from __future__ import annotations
from typing import Any
import pytest
from headroom.cache.compression_store import get_compression_store, reset_compression_store
from headroom.tokenizers import get_tokenizer
from headroom.transforms.agy_fr_compressor import (
_FR_CCR_MARKER_PREFIX,
compress_function_response_leaves,
)
_MODEL = "gemini-3-flash-agent"
# Large, single-line, non-repeating-line leaf: well above the marker-derived
# floor (~2x a ~20-token marker), so it is compressed.
_COMPRESSIBLE_LEAF = "search result row alpha beta gamma delta epsilon zeta eta " * 40
# Tiny leaf: below the marker-derived floor, so it must be left untouched.
_SUB_FLOOR_LEAF = "ok"
@pytest.fixture
def tok() -> Any:
return get_tokenizer(_MODEL)
@pytest.fixture
def store(monkeypatch: pytest.MonkeyPatch) -> Any:
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
reset_compression_store()
s = get_compression_store()
yield s
reset_compression_store()
def _contents() -> list[dict]:
return [
{
"role": "user",
"parts": [
{
"functionResponse": {
"name": "search",
"response": {"output": _COMPRESSIBLE_LEAF},
}
}
],
},
{
"role": "user",
"parts": [
{
"functionResponse": {
"name": "search",
"response": {"output": _SUB_FLOOR_LEAF},
}
}
],
},
{
"role": "user",
"parts": [
{
"functionResponse": {
# Exempt: headroom_retrieve's own output is never
# re-compressed (would self-defeating-loop).
"name": "headroom_retrieve",
"response": {"output": _COMPRESSIBLE_LEAF},
}
}
],
},
]
def test_compress_function_response_leaves_standalone(tok: Any, store: Any) -> None:
contents = _contents()
before, after, leaves = compress_function_response_leaves(contents, "ccr", tok, store)
# Only the one compressible leaf is counted.
assert leaves == 1
assert before > after > 0
compressible = contents[0]["parts"][0]["functionResponse"]["response"]["output"]
sub_floor = contents[1]["parts"][0]["functionResponse"]["response"]["output"]
exempt = contents[2]["parts"][0]["functionResponse"]["response"]["output"]
assert compressible.startswith(_FR_CCR_MARKER_PREFIX)
assert sub_floor == _SUB_FLOOR_LEAF
assert exempt == _COMPRESSIBLE_LEAF

37
tests/test_agy_fr_mode_default.py Normal file
View file

@ -0,0 +1,37 @@
"""Tests for the ``ccr``-default of ``_requested_agy_fr_mode``.
Scope (headroom-37g.32, WU-CCRDEFAULT): ``HEADROOM_AGY_FR_MODE`` defaults to
``ccr`` -- both when unset and when set to an invalid value -- so agy users get
tool-output savings by default (WU2 retrieve-exemption converges voluntary
retrieval). ``lossless`` remains available but must be requested explicitly.
The unrecoverable-marker safety net is preserved downstream: ``_resolve_agy_fr_mode``
still downgrades ccr->lossless when the retrieve MCP is not wired.
"""
from __future__ import annotations
import pytest
from headroom.proxy.handlers.gemini import _requested_agy_fr_mode
class TestRequestedAgyFrModeDefault:
def test_unset_defaults_to_ccr(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("HEADROOM_AGY_FR_MODE", raising=False)
assert _requested_agy_fr_mode() == "ccr"
def test_invalid_value_falls_back_to_ccr(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "xyz")
assert _requested_agy_fr_mode() == "ccr"
def test_explicit_ccr_is_honored(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "ccr")
assert _requested_agy_fr_mode() == "ccr"
def test_explicit_lossless_is_honored(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "lossless")
assert _requested_agy_fr_mode() == "lossless"
def test_normalizes_case_and_whitespace(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", " CCR ")
assert _requested_agy_fr_mode() == "ccr"

217
tests/test_agy_fr_perf_37g35.py Normal file
View file

@ -0,0 +1,217 @@
"""headroom-37g.35: instrumentation proving the per-leaf double-work is gone.
Byte-parity alone does not prove the redundant work was eliminated -- a
regression that silently reintroduces a duplicate ``count_text`` or
``default_ccr_hash`` call would still pass every existing behavioral test.
These tests spy on the real call counts instead.
Item A: ``_compress_fr_leaf`` used to recompute (a) the marker's own token
cost (identical to the per-request floor calculation) and (b) the leaf's
own token count (identical to what ``_walk_fr_compress`` already computed),
and ``store.store()`` used to re-derive SHA-256(leaf)[:24] internally even
though ``_walk_fr_compress`` already computed it for the retrieve-hash
exemption check. All three are now computed exactly once and threaded
through.
Item D: ``_collect_retrieved_hashes`` used to serialize a functionCall's
``args`` via ``json.dumps`` just to substring-search it for
``"headroom_retrieve"``. That is replaced with a direct recursive scan
(``_args_mention_retrieve``) -- ``json.dumps`` must not be called at all.
"""
from __future__ import annotations
from typing import Any
from unittest.mock import patch
import pytest
from headroom.cache.compression_store import (
default_ccr_hash,
get_compression_store,
reset_compression_store,
)
from headroom.tokenizers import get_tokenizer
from headroom.transforms import agy_fr_compressor
from headroom.transforms.agy_fr_compressor import (
_FR_CCR_MARKER_TEMPLATE,
_collect_retrieved_hashes,
compress_function_response_leaves,
)
_MODEL = "gemini-3-flash-agent"
# Large, single-line, non-repeating-line leaf: well above the marker-derived
# floor (~2x a ~20-token marker), so it is compressed exactly once.
_COMPRESSIBLE_LEAF = "search result row alpha beta gamma delta epsilon zeta eta " * 40
_MARKER_PLACEHOLDER = _FR_CCR_MARKER_TEMPLATE.format(hash="0" * 24)
@pytest.fixture
def tok() -> Any:
return get_tokenizer(_MODEL)
@pytest.fixture
def store(monkeypatch: pytest.MonkeyPatch) -> Any:
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
reset_compression_store()
s = get_compression_store()
yield s
reset_compression_store()
def _contents_one_compressible_leaf() -> list[dict]:
return [
{
"role": "user",
"parts": [
{
"functionResponse": {
"name": "search",
"response": {"output": _COMPRESSIBLE_LEAF},
}
}
],
}
]
# ---------------------------------------------------------------------------
# Item A: count_text call count.
# ---------------------------------------------------------------------------
def test_count_text_called_exactly_once_per_leaf_plus_one_per_request(tok: Any, store: Any) -> None:
"""Pin the exact ``count_text`` call count for ONE compressed leaf.
Expected calls (3 total, not the pre-fix 5):
1. ``_fr_marker_tokens_and_floor`` -- ONE per-request call on the
placeholder marker (``hash="0"*24``), shared for both the floor
and ``marker_body_tokens`` (no longer recomputed inside
``_compress_fr_leaf``).
2. ``_walk_fr_compress`` -- ONE call on the original leaf text
(``leaf_tokens``), no longer recomputed a second time as
``original_tokens`` inside ``store.store()``.
3. ``_walk_fr_compress`` -- ONE call on the resulting marker text
(real hash digits, not the placeholder) to compute the actual
``stats["after"]`` token count. This one is NOT eliminated by
37g.35 -- it counts a different string than call 1 (real hash
vs. placeholder hash) and is required for accurate before/after
stats.
"""
with patch.object(tok, "count_text", wraps=tok.count_text) as spy:
before, after, leaves = compress_function_response_leaves(
_contents_one_compressible_leaf(), "ccr", tok, store
)
assert leaves == 1
assert before > after > 0
calls = [call.args[0] for call in spy.call_args_list]
assert len(calls) == 3, f"expected exactly 3 count_text calls, got {len(calls)}: {calls}"
# The placeholder marker is counted exactly once (the old code counted
# it a second time inside _compress_fr_leaf -- that recompute is gone).
assert calls.count(_MARKER_PLACEHOLDER) == 1
# The original leaf text is counted exactly once (the old code counted
# it a second time as store.store()'s `original_tokens` arg -- gone).
assert calls.count(_COMPRESSIBLE_LEAF) == 1
# The remaining call is the actual (real-hash) marker text, required
# for stats and NOT part of the eliminated double-work.
remaining = [c for c in calls if c not in (_MARKER_PLACEHOLDER, _COMPRESSIBLE_LEAF)]
assert len(remaining) == 1
assert remaining[0].startswith("[functionResponse compressed.")
# ---------------------------------------------------------------------------
# Item A: default_ccr_hash call count.
# ---------------------------------------------------------------------------
def test_default_ccr_hash_called_exactly_once_per_leaf(tok: Any, store: Any) -> None:
"""``default_ccr_hash`` must be called exactly once for one compressed
leaf: once in ``_walk_fr_compress`` for the exemption check, and that
SAME value is threaded into ``store.store(..., explicit_hash=...)``.
Scope note: this spy patches the compressor module's ``default_ccr_hash``
binding, so it proves the compressor computes the hash once (not twice).
That the store then SKIPS its own internal recompute when ``explicit_hash``
is passed is a separate fact, verified by inspection of
``compression_store.store`` (the ``explicit_hash is not None`` branch skips
``default_ccr_hash(original)``), not asserted by this spy.
"""
with patch.object(agy_fr_compressor, "default_ccr_hash", wraps=default_ccr_hash) as hash_spy:
before, after, leaves = compress_function_response_leaves(
_contents_one_compressible_leaf(), "ccr", tok, store
)
assert leaves == 1
assert hash_spy.call_count == 1
assert hash_spy.call_args.args[0] == _COMPRESSIBLE_LEAF
# ---------------------------------------------------------------------------
# Item A: explicit_hash produces byte-identical store keys / markers.
# ---------------------------------------------------------------------------
def test_explicit_hash_matches_implicit_default_hash(tok: Any, store: Any) -> None:
"""The threaded ``explicit_hash`` must yield the SAME store key as the
old implicit default (``default_ccr_hash(original)``) -- otherwise
``/v1/retrieve/{hash}`` would 404 for previously-cached content."""
contents = _contents_one_compressible_leaf()
compress_function_response_leaves(contents, "ccr", tok, store)
marker = contents[0]["parts"][0]["functionResponse"]["response"]["output"]
expected_hash = default_ccr_hash(_COMPRESSIBLE_LEAF)
assert marker.endswith(f"hash={expected_hash}]")
entry = store.retrieve(expected_hash)
assert entry is not None
assert entry.original_content == _COMPRESSIBLE_LEAF
# ---------------------------------------------------------------------------
# Item D: json.dumps must not be called by _collect_retrieved_hashes.
# ---------------------------------------------------------------------------
def _mcp_retrieve_call_entry(hash_value: str) -> dict:
"""Generic MCP dispatch shape: a ``call_mcp_tool`` functionCall whose
args reference ``headroom_retrieve`` (as a VALUE, not a key) and carry
the target hash -- the case the old ``json.dumps(args)`` substring scan
was covering."""
return {
"role": "model",
"parts": [
{
"functionCall": {
"name": "call_mcp_tool",
"args": {"tool": "headroom_retrieve", "arguments": {"hash": hash_value}},
}
}
],
}
def test_collect_retrieved_hashes_never_calls_json_dumps() -> None:
h = default_ccr_hash(_COMPRESSIBLE_LEAF)
contents = [_mcp_retrieve_call_entry(h)]
with patch("json.dumps", side_effect=AssertionError("json.dumps must not be called")) as spy:
hashes = _collect_retrieved_hashes(contents)
assert spy.call_count == 0
assert h in hashes
def test_collect_retrieved_hashes_still_finds_bare_retrieve_call() -> None:
"""Sanity check the replacement scan is not merely absent-of-crash --
it must still find hashes via the bare ``headroom_retrieve`` name path
(which never touched ``json.dumps`` even before this change)."""
h = default_ccr_hash(_COMPRESSIBLE_LEAF)
contents = [
{
"role": "model",
"parts": [{"functionCall": {"name": "headroom_retrieve", "args": {"hash": h}}}],
}
]
with patch("json.dumps", side_effect=AssertionError("json.dumps must not be called")):
hashes = _collect_retrieved_hashes(contents)
assert h in hashes

250
tests/test_agy_fr_retrieve_envelope_exempt.py Normal file
View file

@ -0,0 +1,250 @@
"""headroom-8tm: the headroom_retrieve result envelope must NOT be re-compressed.
On agy the retrieve-result functionResponse carries a name that matches neither
``is_headroom_retrieve_name`` nor ``_args_mention_retrieve``, so BOTH the
name-based fr exemption and the functionCall hash-collection miss it. Left
unfixed, the resolved envelope re-compresses into a marker every turn (the model
re-retrieves it, L1) and the ORIGINAL resent leaf keeps re-retrieving (the 236x,
L2). These tests pin the name-INDEPENDENT, content-based exemption.
Fixtures mirror the real ``ccr.mcp_server._retrieve_content`` serialization
(``json.dumps(result, indent=2)``) plus the agy ``Created At:/Completed At:``
text wrapper observed in the fry run3 store.
"""
from __future__ import annotations
import json
from typing import Any
import pytest
import headroom.transforms.agy_fr_compressor as mod
from headroom.cache.compression_store import (
default_ccr_hash,
get_compression_store,
reset_compression_store,
)
from headroom.tokenizers import get_tokenizer
from headroom.transforms.agy_fr_compressor import (
_FR_CCR_MARKER_PREFIX,
_ccr_envelope_hash,
_collect_retrieved_hashes,
compress_function_response_leaves,
)
_MODEL = "gemini-3-flash-agent"
# Big, non-repeating original content -- well above the compression floor, so if
# it were NOT exempt it would compress to a marker.
_ORIGINAL = "watermark CRIMSON-WALRUS " + ("archive row alpha beta gamma delta epsilon " * 90)
@pytest.fixture
def tok() -> Any:
return get_tokenizer(_MODEL)
@pytest.fixture
def store(monkeypatch: pytest.MonkeyPatch) -> Any:
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
reset_compression_store()
s = get_compression_store()
yield s
reset_compression_store()
def _local_envelope_dict(hash_key: str, original: str = _ORIGINAL) -> dict:
return {
"hash": hash_key,
"source": "local",
"original_content": original,
"original_item_count": 10,
"compressed_item_count": 1,
"retrieval_count": 1,
}
def _proxy_envelope_dict(hash_key: str, original: str = _ORIGINAL) -> dict:
# WU-2b: source is the leading key on the proxy path too.
return {
"source": "proxy",
"hash": hash_key,
"original_content": original,
"original_tokens": 500,
"original_item_count": 10,
"compressed_item_count": 1,
"tool_name": "headroom_retrieve",
"retrieval_count": 1,
}
def _agy_text(envelope: dict) -> str:
"""The envelope as agy renders it: a timestamped text wrapper + indent=2 JSON."""
body = json.dumps(envelope, indent=2)
return f"Created At: 2026-07-11T15:59:27+02:00\nCompleted At: 2026-07-11T15:59:28+02:00\n{body}"
# --- Detector -------------------------------------------------------------
class TestDetector:
def test_local_dict(self) -> None:
h = default_ccr_hash(_ORIGINAL)
assert _ccr_envelope_hash(_local_envelope_dict(h)) == h
def test_proxy_dict(self) -> None:
h = default_ccr_hash(_ORIGINAL)
assert _ccr_envelope_hash(_proxy_envelope_dict(h)) == h
def test_local_text(self) -> None:
h = default_ccr_hash(_ORIGINAL)
assert _ccr_envelope_hash(_agy_text(_local_envelope_dict(h))) == h
def test_proxy_text(self) -> None:
h = default_ccr_hash(_ORIGINAL)
assert _ccr_envelope_hash(_agy_text(_proxy_envelope_dict(h))) == h
def test_file_pointer_variant_still_detected(self) -> None:
# agy replaced a large original_content with a saved-to-file pointer;
# the LEADING hash+source keys survive.
h = "a" * 24
env = {
"hash": h,
"source": "local",
"original_content": "The output was large and was saved to: file:///tmp/x",
}
assert _ccr_envelope_hash(_agy_text(env)) == h
def test_source_read_is_not_an_envelope(self) -> None:
# A leaf that READS headroom's own source: key NAMES present, but the
# hash value is a variable (`hash_key`), not a 24-hex literal.
leaf = (
'return {\n "hash": hash_key,\n "source": "local",\n "original_content": entry.x,\n}'
)
assert _ccr_envelope_hash(leaf) is None
def test_uppercase_hash_rejected(self) -> None:
env = _local_envelope_dict("A" * 24)
assert _ccr_envelope_hash(env) is None
assert _ccr_envelope_hash(_agy_text(env)) is None
def test_forty_hex_rejected(self) -> None:
# A 40-hex git sha must not match the {24} anchor (hex boundary).
leaf = '{\n "hash": "' + ("d" * 40) + '",\n "source": "local"\n}'
assert _ccr_envelope_hash(leaf) is None
def test_plain_text_is_none(self) -> None:
assert _ccr_envelope_hash("just some large file content\n" * 50) is None
# --- L1: envelope leaf never compressed -----------------------------------
def _fr_contents(name: str, leaf: Any) -> list[dict]:
return [
{
"role": "user",
"parts": [{"functionResponse": {"name": name, "response": {"output": leaf}}}],
}
]
class TestL1Exempt:
@pytest.mark.parametrize(
"name", ["headroom.headroom_retrieve", "headroom", "call_mcp_tool", None]
)
def test_text_envelope_not_compressed_regardless_of_name(
self, tok: Any, store: Any, name: str, monkeypatch: pytest.MonkeyPatch
) -> None:
h = default_ccr_hash(_ORIGINAL)
env_text = _agy_text(_local_envelope_dict(h))
contents = _fr_contents(name, env_text)
calls: list = []
orig = mod._compress_fr_leaf
monkeypatch.setattr(
mod, "_compress_fr_leaf", lambda *a, **k: (calls.append(1), orig(*a, **k))[1]
)
before, after, leaves = compress_function_response_leaves(contents, "ccr", tok, store)
out = contents[0]["parts"][0]["functionResponse"]["response"]["output"]
assert out == env_text # verbatim -- not a marker
assert not out.startswith(_FR_CCR_MARKER_PREFIX)
assert leaves == 0
assert calls == [] # _compress_fr_leaf never called for the envelope
def test_proxy_and_file_pointer_variants_exempt(self, tok: Any, store: Any) -> None:
h = default_ccr_hash(_ORIGINAL)
for leaf in (
_agy_text(_proxy_envelope_dict(h)),
_agy_text(
{"hash": "b" * 24, "source": "local", "original_content": "saved to: file:///tmp/y"}
),
):
contents = _fr_contents("headroom.headroom_retrieve", leaf)
compress_function_response_leaves(contents, "ccr", tok, store)
assert contents[0]["parts"][0]["functionResponse"]["response"]["output"] == leaf
def test_dict_envelope_response_exempt(self, tok: Any, store: Any) -> None:
# response IS the envelope dict (structured, not text-rendered).
h = default_ccr_hash(_ORIGINAL)
env = _local_envelope_dict(h)
contents = [
{"role": "user", "parts": [{"functionResponse": {"name": "headroom", "response": env}}]}
]
compress_function_response_leaves(contents, "ccr", tok, store)
# original_content left verbatim (dict exempt as a whole)
assert (
contents[0]["parts"][0]["functionResponse"]["response"]["original_content"] == _ORIGINAL
)
# --- L2: envelope hash exempts the ORIGINAL resent leaf --------------------
class TestL2Exempt:
def test_envelope_hash_collected(self) -> None:
h = default_ccr_hash(_ORIGINAL)
contents = _fr_contents("headroom.headroom_retrieve", _agy_text(_local_envelope_dict(h)))
assert h in _collect_retrieved_hashes(contents)
def test_original_leaf_exempt_when_envelope_present(self, tok: Any, store: Any) -> None:
h = default_ccr_hash(_ORIGINAL)
contents = [
{
"role": "user",
"parts": [
{
"functionResponse": {
"name": "headroom.headroom_retrieve",
"response": {"output": _agy_text(_local_envelope_dict(h))},
}
},
],
},
{
"role": "user",
"parts": [
{"functionResponse": {"name": "read_file", "response": {"output": _ORIGINAL}}},
],
},
]
compress_function_response_leaves(contents, "ccr", tok, store)
# the resent ORIGINAL leaf is exempt (its hash is in retrieved_hashes)
assert contents[1]["parts"][0]["functionResponse"]["response"]["output"] == _ORIGINAL
# --- Negatives: normal / false-positive leaves STILL compress -------------
class TestStillCompresses:
def test_normal_large_leaf_compresses(self, tok: Any, store: Any) -> None:
contents = _fr_contents("read_file", _ORIGINAL)
before, after, leaves = compress_function_response_leaves(contents, "ccr", tok, store)
out = contents[0]["parts"][0]["functionResponse"]["response"]["output"]
assert leaves == 1
assert out.startswith(_FR_CCR_MARKER_PREFIX)
def test_source_read_leaf_still_compresses(self, tok: Any, store: Any) -> None:
# Large leaf mentioning the key NAMES but no 24-hex hash value.
src = (
'def build():\n return {\n "hash": hash_key,\n "source": "local",\n'
' "original_content": entry.original_content,\n }\n'
) * 40
contents = _fr_contents("read_file", src)
before, after, leaves = compress_function_response_leaves(contents, "ccr", tok, store)
out = contents[0]["parts"][0]["functionResponse"]["response"]["output"]
assert leaves == 1
assert out.startswith(_FR_CCR_MARKER_PREFIX)

441
tests/test_agy_functionresponse_compression.py Normal file
View file

@ -0,0 +1,441 @@
"""WU1 (headroom-37g.1): uniform deterministic recoverable compression of agy
functionResponse leaves in ``handle_google_cloudcode_stream``.
Scope (proves WU1's DoD):
- determinism / idempotency: same leaf -> identical bytes; f(f(x)) == f(x).
- CCR byte-recovery: retrieve(hash) returns the ORIGINAL leaf bytes.
- delivery-on-revert: text pipeline reverts, but a large functionResponse leaf
still ships compressed in ``request_payload["contents"]`` and tokens_saved > 0.
- uniform: a historical (non-tail) functionResponse entry is ALSO compressed.
- pairing/shape preserved; functionCall untouched; non-string leaf skipped;
multi-functionResponse-part entry all compressed.
- retrieve-gating: mode=ccr without HEADROOM_AGY_RETRIEVE_WIRED -> no
unrecoverable marker (lossless / no-op).
All upstream/network calls are stubbed via monkeypatch on
``HeadroomProxy._stream_response`` / ``openai_pipeline.apply``. Never contacts
the real 8787 proxy or any network destination.
"""
from __future__ import annotations
import copy
from typing import Any
import pytest
from fastapi.responses import StreamingResponse
from fastapi.testclient import TestClient
from headroom.cache.compression_store import (
get_compression_store,
reset_compression_store,
)
from headroom.ccr.tool_injection import is_headroom_retrieve_name
from headroom.parser import CCR_RETRIEVAL_MARKER_RE
from headroom.proxy.handlers.gemini import (
_FR_CCR_MARKER_PREFIX,
_resolve_agy_fr_mode,
)
from headroom.proxy.server import HeadroomProxy, ProxyConfig, create_app
from headroom.tokenizers import get_tokenizer
# A large, single-line string leaf: no repeated lines (so lossless is a no-op),
# well above the marker-derived floor (~2x a ~20-token marker).
BIG_LEAF = "search result row alpha beta gamma delta epsilon zeta eta " * 40
# Enough text to make CompressionDecision.should_compress True (mirrors the
# shared fixture in test_proxy_agy_compression.py).
_REPEAT_UNIT = "The quick brown fox jumps over the lazy dog. " * 60
_MODEL = "gemini-3-flash-agent"
_SSE_PAYLOAD = (
b'data: {"candidates":[{"content":{"parts":[{"text":"ok"}]}}]}\r\n\r\ndata: [DONE]\r\n\r\n'
)
def _make_sse() -> StreamingResponse:
async def _body() -> Any:
yield _SSE_PAYLOAD
return StreamingResponse(_body(), status_code=200, media_type="text/event-stream")
def _hash_of(marker: str) -> str:
assert marker.startswith(_FR_CCR_MARKER_PREFIX), marker
# split: single-hash marker now, in the ``Retrieve more: hash=<h>]`` form
# the parser regex keys on.
return marker.split("hash=", 1)[1].rstrip("]")
def _fr_leaf(contents: list, entry: int, part: int = 0, key: str = "output") -> Any:
return contents[entry]["parts"][part]["functionResponse"]["response"][key]
class _FakeResult:
"""Stand-in for the compression pipeline result."""
def __init__(self, messages: Any, tokens_before: int, tokens_after: int) -> None:
self.messages = messages
self.tokens_before = tokens_before
self.tokens_after = tokens_after
self.transforms_applied: list[str] = ["noop"]
@pytest.fixture
def proxy() -> Any:
with TestClient(create_app(ProxyConfig(optimize=True))) as client:
yield client.app.state.proxy # type: ignore[attr-defined]
@pytest.fixture
def tok() -> Any:
return get_tokenizer(_MODEL)
@pytest.fixture
def ccr_store(monkeypatch: pytest.MonkeyPatch) -> Any:
# Force an in-memory backend and a clean store for hermetic byte-recovery.
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
reset_compression_store()
store = get_compression_store()
yield store
reset_compression_store()
def _fr_entry(role: str = "user", leaf: Any = BIG_LEAF, name: str = "search") -> dict:
return {
"role": role,
"parts": [{"functionResponse": {"name": name, "response": {"output": leaf}}}],
}
# ---------------------------------------------------------------------------
# Determinism + idempotency
# ---------------------------------------------------------------------------
def test_ccr_deterministic_and_idempotent(proxy: Any, tok: Any, ccr_store: Any) -> None:
c1 = [_fr_entry()]
c2 = [_fr_entry()]
b1, a1, l1 = proxy._compress_agy_function_responses(c1, "ccr", tok, ccr_store)
b2, a2, l2 = proxy._compress_agy_function_responses(c2, "ccr", tok, ccr_store)
assert l1 == 1 and l2 == 1
# Deterministic: identical original -> identical marker bytes.
assert _fr_leaf(c1, 0) == _fr_leaf(c2, 0)
assert (b1, a1) == (b2, a2)
# Idempotent: f(f(x)) == f(x) -- re-running is a no-op, bytes stable.
stable = _fr_leaf(c1, 0)
b3, a3, l3 = proxy._compress_agy_function_responses(c1, "ccr", tok, ccr_store)
assert l3 == 0
assert _fr_leaf(c1, 0) == stable
# ---------------------------------------------------------------------------
# CCR byte-recovery
# ---------------------------------------------------------------------------
def test_ccr_byte_recovery(proxy: Any, tok: Any, ccr_store: Any) -> None:
contents = [_fr_entry()]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 1 and before > after > 0
marker = _fr_leaf(contents, 0)
entry = ccr_store.retrieve(_hash_of(marker))
assert entry is not None
assert entry.original_content == BIG_LEAF
# ---------------------------------------------------------------------------
# Self-describing marker: names the retrieval tool so a model that needs the
# compressed detail knows how to expand it (WU4 observed 0 retrieve calls
# because the old marker named no tool).
# ---------------------------------------------------------------------------
def test_marker_names_headroom_retrieve_tool(proxy: Any, tok: Any, ccr_store: Any) -> None:
c1 = [_fr_entry()]
c2 = [_fr_entry()]
proxy._compress_agy_function_responses(c1, "ccr", tok, ccr_store)
proxy._compress_agy_function_responses(c2, "ccr", tok, ccr_store)
marker = _fr_leaf(c1, 0)
# Names the tool + gives a one-line call-to-expand instruction.
assert "headroom_retrieve" in marker
# Store-lookup path substring intact: parser.CCR_RETRIEVAL_MARKER_RE and
# the CCR retrieval path both key on this exact substring.
assert "Retrieve more: hash=" in marker
assert CCR_RETRIEVAL_MARKER_RE.search(marker) is not None
# Deterministic: identical original -> identical marker bytes, twice over.
assert marker == _fr_leaf(c2, 0)
# Round-trips via the store using the trailing canonical hash.
entry = ccr_store.retrieve(_hash_of(marker))
assert entry is not None
assert entry.original_content == BIG_LEAF
# ---------------------------------------------------------------------------
# Floor + non-string leaves
# ---------------------------------------------------------------------------
def test_sub_floor_and_non_string_leaves_skipped(proxy: Any, tok: Any, ccr_store: Any) -> None:
contents = [
{
"role": "user",
"parts": [
{
"functionResponse": {
"name": "search",
"response": {
"small": "hi there", # below floor
"count": 12345, # non-string scalar
"ok": True, # non-string scalar
"nothing": None, # non-string scalar
"big": BIG_LEAF, # compressed
},
}
}
],
}
]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 1
resp = contents[0]["parts"][0]["functionResponse"]["response"]
assert resp["small"] == "hi there"
assert resp["count"] == 12345
assert resp["ok"] is True
assert resp["nothing"] is None
assert resp["big"].startswith(_FR_CCR_MARKER_PREFIX)
# ---------------------------------------------------------------------------
# Pairing / shape: functionCall untouched
# ---------------------------------------------------------------------------
def test_functioncall_untouched_pairing_preserved(proxy: Any, tok: Any, ccr_store: Any) -> None:
contents = [
{
"role": "model",
"parts": [{"functionCall": {"name": "search", "args": {"query": BIG_LEAF}}}],
},
_fr_entry(role="user"),
]
original_call = copy.deepcopy(contents[0])
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 1
# functionCall entry byte-identical (never touched).
assert contents[0] == original_call
# functionResponse leaf compressed.
assert _fr_leaf(contents, 1).startswith(_FR_CCR_MARKER_PREFIX)
# ---------------------------------------------------------------------------
# Multiple functionResponse parts in one entry: all compressed
# ---------------------------------------------------------------------------
def test_multi_functionresponse_parts_all_compressed(proxy: Any, tok: Any, ccr_store: Any) -> None:
contents = [
{
"role": "user",
"parts": [
{"functionResponse": {"name": "a", "response": {"output": BIG_LEAF}}},
{"functionResponse": {"name": "b", "response": {"output": BIG_LEAF + "!"}}},
],
}
]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 2
for part in contents[0]["parts"]:
assert part["functionResponse"]["response"]["output"].startswith(_FR_CCR_MARKER_PREFIX)
# ---------------------------------------------------------------------------
# Retrieve-gating: ccr without wired -> lossless (no unrecoverable marker)
# ---------------------------------------------------------------------------
def test_mode_downgrades_to_lossless_when_retrieve_not_wired(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("HEADROOM_AGY_RETRIEVE_WIRED", raising=False)
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "ccr")
assert _resolve_agy_fr_mode() == "lossless"
monkeypatch.setenv("HEADROOM_AGY_RETRIEVE_WIRED", "1")
assert _resolve_agy_fr_mode() == "ccr"
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "lossless")
assert _resolve_agy_fr_mode() == "lossless"
def test_lossless_mode_emits_no_unrecoverable_marker(proxy: Any, tok: Any) -> None:
contents = [_fr_entry()]
# lossless never needs a store; a single-line leaf is a no-op.
before, after, leaves = proxy._compress_agy_function_responses(contents, "lossless", tok, None)
leaf = _fr_leaf(contents, 0)
assert "Retrieve more: hash=" not in leaf
assert leaf == BIG_LEAF # unchanged no-op, still recoverable (byte-identical)
# ---------------------------------------------------------------------------
# Uniform: historical (non-tail) functionResponse entry also compressed
# ---------------------------------------------------------------------------
def test_uniform_historical_and_tail_compressed(proxy: Any, tok: Any, ccr_store: Any) -> None:
contents = [
_fr_entry(role="user", name="hist"), # historical (index 0)
{"role": "model", "parts": [{"text": "some reasoning"}]},
_fr_entry(role="user", name="tail"), # tail (index 2)
]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 2
assert _fr_leaf(contents, 0).startswith(_FR_CCR_MARKER_PREFIX)
assert _fr_leaf(contents, 2).startswith(_FR_CCR_MARKER_PREFIX)
# ---------------------------------------------------------------------------
# Anti-self-defeating-loop: headroom_retrieve's OWN output must never be
# re-compressed back into the marker it was expanded from (headroom-37g.8).
# ---------------------------------------------------------------------------
def test_is_headroom_retrieve_name_matching() -> None:
# Bare name and MCP-prefixed / custom-prefixed variants match.
assert is_headroom_retrieve_name("headroom_retrieve") is True
assert is_headroom_retrieve_name("mcp__headroom__headroom_retrieve") is True
assert is_headroom_retrieve_name("custom__headroom_retrieve") is True
# Unrelated / near-miss names must NOT match.
assert is_headroom_retrieve_name("read_file") is False
assert is_headroom_retrieve_name("my_headroom_retrieve_helper") is False
# No double-underscore boundary -- must NOT match (single "x" prefix).
assert is_headroom_retrieve_name("xheadroom_retrieve") is False
assert is_headroom_retrieve_name(None) is False
assert is_headroom_retrieve_name("") is False
# untrusted JSON: non-str name must return False, never raise
assert is_headroom_retrieve_name(123) is False
assert is_headroom_retrieve_name({"headroom_retrieve": 1}) is False
assert is_headroom_retrieve_name(["headroom_retrieve"]) is False
def test_headroom_retrieve_output_exempted_from_recompression(
proxy: Any, tok: Any, ccr_store: Any
) -> None:
contents = [
{
"role": "user",
"parts": [
{
"functionResponse": {
"name": "mcp__headroom__headroom_retrieve",
"response": {"output": BIG_LEAF},
}
},
{
"functionResponse": {
"name": "read_file",
"response": {"output": BIG_LEAF + "!"},
}
},
],
}
]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
# Only the normal tool's leaf is compressed; the retrieve leaf is skipped.
assert leaves == 1
retrieve_leaf = _fr_leaf(contents, 0, part=0)
normal_leaf = _fr_leaf(contents, 0, part=1)
assert retrieve_leaf == BIG_LEAF
assert not retrieve_leaf.startswith(_FR_CCR_MARKER_PREFIX)
assert normal_leaf.startswith(_FR_CCR_MARKER_PREFIX)
def test_headroom_retrieve_bare_and_suffixed_names_both_exempted(
proxy: Any, tok: Any, ccr_store: Any
) -> None:
contents = [
{
"role": "user",
"parts": [
{
"functionResponse": {
"name": "headroom_retrieve",
"response": {"output": BIG_LEAF},
}
},
{
"functionResponse": {
"name": "toolgroup__headroom_retrieve",
"response": {"output": BIG_LEAF + "!"},
}
},
],
}
]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 0
assert _fr_leaf(contents, 0, part=0) == BIG_LEAF
assert _fr_leaf(contents, 0, part=1) == BIG_LEAF + "!"
# ---------------------------------------------------------------------------
# Integration: delivery + accounting even when the text pipeline REVERTS
# ---------------------------------------------------------------------------
def test_delivery_on_text_revert_ships_and_counts(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "ccr")
monkeypatch.setenv("HEADROOM_AGY_RETRIEVE_WIRED", "1")
reset_compression_store()
captured: dict[str, Any] = {}
async def _fake_stream(
proxy_self: Any,
url: str,
headers: dict,
body: dict,
provider: str,
model: str,
request_id: Any,
original_tokens: int,
optimized_tokens: int,
tokens_saved: int,
*args: Any,
**kwargs: Any,
) -> StreamingResponse:
captured["body"] = body
captured["tokens_saved"] = tokens_saved
return _make_sse()
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
# Force the text pipeline to REVERT: report inflation (after > before).
def _inflating_apply(**kw: Any) -> _FakeResult:
return _FakeResult(
messages=[{"role": "user", "content": "x"}], tokens_before=5, tokens_after=99999
)
body = {
"model": _MODEL,
"request": {
"contents": [
{"role": "user", "parts": [{"text": _REPEAT_UNIT}]},
_fr_entry(role="user"),
]
},
}
with TestClient(create_app(ProxyConfig(optimize=True))) as client:
proxy = client.app.state.proxy # type: ignore[attr-defined]
proxy.openai_pipeline.apply = _inflating_apply # type: ignore[method-assign]
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": "antigravity/1.0.5"},
json=body,
)
assert response.status_code == 200
contents = captured["body"]["request"]["contents"]
marker = _fr_leaf(contents, 1)
# Compressed leaf SHIPPED despite the text-pipeline revert.
assert marker.startswith(_FR_CCR_MARKER_PREFIX)
# Text entry preserved (revert kept original text).
assert contents[0]["parts"][0]["text"] == _REPEAT_UNIT
# Saving counted even though the text pipeline reverted (704->718 case).
assert captured["tokens_saved"] > 0
# Byte-recoverable.
entry = get_compression_store().retrieve(_hash_of(marker))
assert entry is not None and entry.original_content == BIG_LEAF
reset_compression_store()

394
tests/test_agy_functionresponse_compression_matrix.py Normal file
View file

@ -0,0 +1,394 @@
"""WU2 (headroom-37g.2): matrix coverage for WU1's agy functionResponse leaf
compression that is NOT already covered by
``tests/test_agy_functionresponse_compression.py``.
Five gaps closed here:
1. Cross-turn cache stability: the SAME functionResponse entry, re-sent by agy
unchanged across two turns (once as tail, once as history), must produce
BYTE-IDENTICAL outbound compressed leaves -- this is what keeps Cloud Code
Assist's server-side cached prefix stable.
2. Mixed functionResponse+text entry: the functionResponse leaf compresses and
the outbound entry is the MUTATED one, not the pristine original; the
co-located text is untouched.
3. No-double-count: ``tokens_saved`` reflects the functionResponse delta only
-- the #819 ``waste_messages`` telemetry path never contributes to it.
4. #819 waste-signal non-regression: the ``include_function_responses=True``
conversion feeding ``TransformPipeline.apply(waste_messages=...)`` still
fires correctly alongside WU1's functionResponse compression.
5. Non-antigravity non-regression: a plain (non-antigravity) Gemini
cloudcode/generateContent request is untouched by the functionResponse-leaf
pass.
All upstream/network calls are stubbed via monkeypatch on
``HeadroomProxy._stream_response`` / ``openai_pipeline.apply``. Never contacts
the real 8787 proxy or any network destination.
"""
from __future__ import annotations
import copy
import json
from typing import Any
import pytest
from fastapi.responses import StreamingResponse
from fastapi.testclient import TestClient
from headroom.cache.compression_store import (
get_compression_store,
reset_compression_store,
)
from headroom.proxy.handlers.gemini import _FR_CCR_MARKER_PREFIX
from headroom.proxy.server import HeadroomProxy, ProxyConfig, create_app
from headroom.tokenizers import get_tokenizer
# Mirrors tests/test_agy_functionresponse_compression.py -- large, single-line
# (no repeated lines, so lossless is a no-op), well above the marker floor.
BIG_LEAF = "search result row alpha beta gamma delta epsilon zeta eta " * 40
# Enough text to make CompressionDecision.should_compress True.
_REPEAT_UNIT = "The quick brown fox jumps over the lazy dog. " * 60
_MODEL = "gemini-3-flash-agent"
_SSE_PAYLOAD = (
b'data: {"candidates":[{"content":{"parts":[{"text":"ok"}]}}]}\r\n\r\ndata: [DONE]\r\n\r\n'
)
def _make_sse() -> StreamingResponse:
async def _body() -> Any:
yield _SSE_PAYLOAD
return StreamingResponse(_body(), status_code=200, media_type="text/event-stream")
def _fr_leaf(contents: list, entry: int, part: int = 0, key: str = "output") -> Any:
return contents[entry]["parts"][part]["functionResponse"]["response"][key]
def _fr_entry(role: str = "user", leaf: Any = BIG_LEAF, name: str = "search") -> dict:
return {
"role": role,
"parts": [{"functionResponse": {"name": name, "response": {"output": leaf}}}],
}
class _FakeResult:
"""Stand-in for the compression pipeline result."""
def __init__(self, messages: Any, tokens_before: int, tokens_after: int) -> None:
self.messages = messages
self.tokens_before = tokens_before
self.tokens_after = tokens_after
self.transforms_applied: list[str] = ["noop"]
@pytest.fixture
def proxy() -> Any:
with TestClient(create_app(ProxyConfig(optimize=True))) as client:
yield client.app.state.proxy # type: ignore[attr-defined]
@pytest.fixture
def tok() -> Any:
return get_tokenizer(_MODEL)
@pytest.fixture
def ccr_store(monkeypatch: pytest.MonkeyPatch) -> Any:
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
reset_compression_store()
store = get_compression_store()
yield store
reset_compression_store()
# ---------------------------------------------------------------------------
# Gap 1: cross-turn cache stability
# ---------------------------------------------------------------------------
def test_cross_turn_cache_bytes_identical_for_resent_entry(
proxy: Any, tok: Any, ccr_store: Any
) -> None:
# Turn N: the functionResponse entry is the newest turn (tail).
turn_n = [
{"role": "user", "parts": [{"text": "call search"}]},
_fr_entry(role="user", name="search"),
]
proxy._compress_agy_function_responses(turn_n, "ccr", tok, ccr_store)
marker_turn_n = _fr_leaf(turn_n, 1)
assert marker_turn_n.startswith(_FR_CCR_MARKER_PREFIX)
# Turn N+1: agy re-sends the SAME entry with its ORIGINAL bytes, now
# historical, plus a new tail turn.
turn_n1 = [
{"role": "user", "parts": [{"text": "call search"}]},
_fr_entry(role="user", name="search"), # identical original leaf, resent
{"role": "model", "parts": [{"text": "reasoning about the result"}]},
_fr_entry(role="user", name="search2", leaf=BIG_LEAF + "!"), # new tail
]
proxy._compress_agy_function_responses(turn_n1, "ccr", tok, ccr_store)
marker_turn_n1 = _fr_leaf(turn_n1, 1)
assert marker_turn_n1.startswith(_FR_CCR_MARKER_PREFIX)
# BYTE-IDENTICAL across turns: the anti-cache-bust guarantee.
assert marker_turn_n1 == marker_turn_n
# Sanity: the new tail entry got its OWN, different marker.
assert _fr_leaf(turn_n1, 3) != marker_turn_n
# ---------------------------------------------------------------------------
# Gap 2: mixed functionResponse + text entry
# ---------------------------------------------------------------------------
def test_mixed_text_and_functionresponse_entry_leaf_compressed_and_entry_mutated(
proxy: Any, tok: Any, ccr_store: Any
) -> None:
contents = [
{
"role": "user",
"parts": [
{"text": "here is context"},
{"functionResponse": {"name": "search", "response": {"output": BIG_LEAF}}},
],
}
]
original = copy.deepcopy(contents[0])
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 1
assert before > after > 0
# functionResponse leaf compressed.
assert _fr_leaf(contents, 0, part=1).startswith(_FR_CCR_MARKER_PREFIX)
# Co-located text part untouched.
assert contents[0]["parts"][0]["text"] == "here is context"
# The outbound entry as a whole is the MUTATED one, not the pristine original.
assert contents[0] != original
assert contents[0]["parts"][1] != original["parts"][1]
# ---------------------------------------------------------------------------
# Gap 3: no-double-count against the #819 waste_messages path
# ---------------------------------------------------------------------------
def test_tokens_saved_reflects_fr_delta_only_no_double_count(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "ccr")
monkeypatch.setenv("HEADROOM_AGY_RETRIEVE_WIRED", "1")
reset_compression_store()
captured: dict[str, Any] = {}
async def _fake_stream(
proxy_self: Any,
url: str,
headers: dict,
body: dict,
provider: str,
model: str,
request_id: Any,
original_tokens: int,
optimized_tokens: int,
tokens_saved: int,
*args: Any,
**kwargs: Any,
) -> StreamingResponse:
captured["tokens_saved"] = tokens_saved
return _make_sse()
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
def _noop_apply(*, messages: Any, **kw: Any) -> _FakeResult:
# Identical messages object -> text-side delta is exactly zero,
# isolating the accounting to the functionResponse leaf pass.
return _FakeResult(messages=messages, tokens_before=0, tokens_after=0)
contents = [
{"role": "user", "parts": [{"text": "call search"}]},
_fr_entry(role="user", leaf=BIG_LEAF),
]
body = {"model": _MODEL, "request": {"contents": copy.deepcopy(contents)}}
with TestClient(create_app(ProxyConfig(optimize=True))) as client:
proxy = client.app.state.proxy # type: ignore[attr-defined]
proxy.openai_pipeline.apply = _noop_apply # type: ignore[method-assign]
# Expected FR delta computed directly via the production method against
# an INDEPENDENT copy (same tokenizer/store) -- proves the shipped
# tokens_saved is exactly the functionResponse delta and nothing more
# (the #819 waste_messages telemetry path contributes zero tokens).
tok_ = get_tokenizer(_MODEL)
store = get_compression_store()
expected_contents = copy.deepcopy(contents)
exp_before, exp_after, exp_leaves = proxy._compress_agy_function_responses(
expected_contents, "ccr", tok_, store
)
assert exp_leaves == 1
assert exp_before > exp_after > 0
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": "antigravity/1.0.5"},
json=body,
)
assert response.status_code == 200
assert captured["tokens_saved"] == exp_before - exp_after
reset_compression_store()
# ---------------------------------------------------------------------------
# Gap 4: #819 waste-signal non-regression
# ---------------------------------------------------------------------------
def test_waste_signal_detection_path_intact_with_fr_compression(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "ccr")
monkeypatch.setenv("HEADROOM_AGY_RETRIEVE_WIRED", "1")
reset_compression_store()
captured: dict[str, Any] = {}
async def _fake_stream(
proxy_self: Any,
url: str,
headers: dict,
body: dict,
provider: str,
model: str,
request_id: Any,
original_tokens: int,
optimized_tokens: int,
tokens_saved: int,
*args: Any,
**kwargs: Any,
) -> StreamingResponse:
captured["body"] = body
return _make_sse()
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
def _spy_apply(*, messages: Any, waste_messages: Any = None, **kw: Any) -> _FakeResult:
captured["waste_messages"] = waste_messages
return _FakeResult(messages=messages, tokens_before=0, tokens_after=0)
# A payload with BOTH a large string leaf (WU1's compression target) and a
# bulky array (drives #819 json-bloat waste-signal detection).
tool_payload = {
"output": BIG_LEAF,
"rows": [{"id": i, "name": f"item_{i}"} for i in range(50)],
}
body = {
"model": _MODEL,
"request": {
"contents": [
{"role": "user", "parts": [{"text": _REPEAT_UNIT}]},
{
"role": "user",
"parts": [
{"functionResponse": {"name": "fetch_data", "response": tool_payload}}
],
},
]
},
}
with TestClient(create_app(ProxyConfig(optimize=True))) as client:
proxy = client.app.state.proxy # type: ignore[attr-defined]
proxy.openai_pipeline.apply = _spy_apply # type: ignore[method-assign]
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": "antigravity/1.0.5"},
json=body,
)
assert response.status_code == 200
# #819: the tool-output payload reached waste-signal detection via the
# include_function_responses=True conversion, unaffected by WU1.
waste_msgs = captured.get("waste_messages")
assert waste_msgs is not None
tool_msgs = [m for m in waste_msgs if m.get("role") == "tool"]
assert len(tool_msgs) == 1
assert json.loads(tool_msgs[0]["content"]) == tool_payload
# WU1's functionResponse compression also ran on the same request: the
# large "output" leaf compressed; the small "rows" array untouched.
contents = captured["body"]["request"]["contents"]
fr_response = contents[1]["parts"][0]["functionResponse"]["response"]
assert fr_response["output"].startswith(_FR_CCR_MARKER_PREFIX)
assert fr_response["rows"] == tool_payload["rows"]
reset_compression_store()
# ---------------------------------------------------------------------------
# Gap 5: non-antigravity non-regression
# ---------------------------------------------------------------------------
def test_non_antigravity_request_functionresponse_passthrough_unchanged(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
monkeypatch.setenv("HEADROOM_AGY_FR_MODE", "ccr")
monkeypatch.setenv("HEADROOM_AGY_RETRIEVE_WIRED", "1")
reset_compression_store()
captured: dict[str, Any] = {}
async def _fake_stream(
proxy_self: Any,
url: str,
headers: dict,
body: dict,
provider: str,
model: str,
request_id: Any,
original_tokens: int,
optimized_tokens: int,
tokens_saved: int,
*args: Any,
**kwargs: Any,
) -> StreamingResponse:
captured["body"] = body
return _make_sse()
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
def _noop_apply(*, messages: Any, **kw: Any) -> _FakeResult:
return _FakeResult(messages=messages, tokens_before=0, tokens_after=0)
body = {
# NOTE: model does not end in "-agent" and no antigravity User-Agent /
# userAgent / requestType / project field is present -> is_antigravity
# resolves to False (see _is_cloudcode_antigravity_request).
"model": "gemini-3-pro",
"request": {
"contents": [
{"role": "user", "parts": [{"text": _REPEAT_UNIT}]},
_fr_entry(role="user"),
]
},
}
with TestClient(create_app(ProxyConfig(optimize=True))) as client:
proxy = client.app.state.proxy # type: ignore[attr-defined]
proxy.openai_pipeline.apply = _noop_apply # type: ignore[method-assign]
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
# Deliberately NOT antigravity: default TestClient User-Agent.
json=body,
)
assert response.status_code == 200
contents = captured["body"]["request"]["contents"]
fr_leaf_value = _fr_leaf(contents, 1)
# Untouched: byte-identical to the original, no marker, no compaction.
assert fr_leaf_value == BIG_LEAF
assert _FR_CCR_MARKER_PREFIX not in fr_leaf_value
reset_compression_store()

136
tests/test_agy_print_mode_version_gate.py Normal file
View file

@ -0,0 +1,136 @@
"""agy print-mode MCP version gate (headroom-37g.37).
Older/unknown agy binaries hang on ANY mcpServers entry in --print mode.
`headroom wrap agy` therefore gates print-mode MCP wiring on a runtime
`agy --version` preflight: enable only when >= 1.0.16, otherwise SUPPRESS and
actively PURGE any persisted entries so a prior interactive run can't leave a
config that still hangs. Interactive mode is never gated (the hang is
print-mode-only). No real agy binary is invoked here — everything is mocked.
"""
import subprocess
from unittest import mock
from headroom.cli import wrap
from headroom.cli.wrap import (
_AGY_PRINT_MODE_MCP_MIN_VERSION,
_agy_print_mode_mcp_allowed,
_detect_agy_version,
_purge_agy_mcp_entries,
)
def _run_result(returncode: int, stdout: str) -> mock.Mock:
return mock.Mock(returncode=returncode, stdout=stdout)
# -- _detect_agy_version (parse + never-raises) -----------------------------
def test_detect_version_parses_bare_line():
with mock.patch("subprocess.run", return_value=_run_result(0, "1.0.16\n")):
assert _detect_agy_version("/usr/bin/agy") == (1, 0, 16)
def test_detect_version_trims_whitespace_and_banner():
with mock.patch("subprocess.run", return_value=_run_result(0, " agy version 1.0.16 \n")):
assert _detect_agy_version("/usr/bin/agy") == (1, 0, 16)
def test_detect_version_takes_last_match_over_wrapper_banner():
# A wrapper prints its own version first, then the real agy version.
with mock.patch("subprocess.run", return_value=_run_result(0, "wrapper 9.9.9\nagy 1.0.16\n")):
assert _detect_agy_version("/usr/bin/agy") == (1, 0, 16)
def test_detect_version_none_on_garbage():
with mock.patch("subprocess.run", return_value=_run_result(0, "no version here")):
assert _detect_agy_version("/usr/bin/agy") is None
def test_detect_version_none_on_empty():
with mock.patch("subprocess.run", return_value=_run_result(0, "")):
assert _detect_agy_version("/usr/bin/agy") is None
def test_detect_version_none_on_nonzero_exit():
with mock.patch("subprocess.run", return_value=_run_result(2, "1.0.16")):
assert _detect_agy_version("/usr/bin/agy") is None
def test_detect_version_none_on_timeout():
with mock.patch("subprocess.run", side_effect=subprocess.TimeoutExpired("agy", 1.0)):
assert _detect_agy_version("/usr/bin/agy") is None
def test_detect_version_none_on_oserror():
with mock.patch("subprocess.run", side_effect=OSError("boom")):
assert _detect_agy_version("/usr/bin/agy") is None
def test_detect_version_none_when_bin_missing():
assert _detect_agy_version(None) is None
assert _detect_agy_version("") is None
def test_detect_version_uses_short_timeout_and_devnull_stderr():
with mock.patch("subprocess.run", return_value=_run_result(0, "1.0.16")) as run:
_detect_agy_version("/usr/bin/agy")
_args, kwargs = run.call_args
assert kwargs["timeout"] == 1.0
assert kwargs["stderr"] is subprocess.DEVNULL
# -- _agy_print_mode_mcp_allowed (the gate) ---------------------------------
def test_gate_allows_print_mode_when_known_good():
with mock.patch.object(
wrap, "_detect_agy_version", return_value=_AGY_PRINT_MODE_MCP_MIN_VERSION
):
assert _agy_print_mode_mcp_allowed(("--print", "hi"), "/usr/bin/agy") is True
def test_gate_allows_print_mode_when_newer():
with mock.patch.object(wrap, "_detect_agy_version", return_value=(1, 1, 0)):
assert _agy_print_mode_mcp_allowed(("-p", "hi"), "/usr/bin/agy") is True
def test_gate_suppresses_print_mode_when_older():
with mock.patch.object(wrap, "_detect_agy_version", return_value=(1, 0, 15)):
assert _agy_print_mode_mcp_allowed(("--print", "hi"), "/usr/bin/agy") is False
def test_gate_suppresses_print_mode_when_unknown():
with mock.patch.object(wrap, "_detect_agy_version", return_value=None):
assert _agy_print_mode_mcp_allowed(("--prompt", "hi"), "/usr/bin/agy") is False
def test_gate_allows_interactive_without_version_check():
# Interactive mode (no print flag) is always allowed and must NOT even
# spend a version-detection subprocess.
with mock.patch.object(wrap, "_detect_agy_version") as detect:
assert _agy_print_mode_mcp_allowed((), "/usr/bin/agy") is True
assert _agy_print_mode_mcp_allowed(("--model", "x"), "/usr/bin/agy") is True
detect.assert_not_called()
# -- _purge_agy_mcp_entries (load-bearing: all persisted types removed) ------
def test_purge_targets_all_four_entry_types():
registrar = mock.Mock()
with (
mock.patch.object(wrap, "_disable_tokensave_mcp") as dis_tok,
mock.patch.object(wrap, "_disable_serena_mcp") as dis_ser,
):
_purge_agy_mcp_entries(registrar)
# tokensave + serena removed via the LEDGER-AWARE helpers
# (so a user-owned entry is never clobbered).
dis_tok.assert_called_once_with(registrar)
dis_ser.assert_called_once()
# code-graph + retrieve removed via raw unregister (Headroom-owned names).
unregistered = {c.args[0] for c in registrar.unregister_server.call_args_list}
assert wrap._CBM_MCP_SERVER_NAME in unregistered
assert "headroom" in unregistered

134
tests/test_agy_provider_env.py Normal file
View file

@ -0,0 +1,134 @@
"""Tests for headroom.providers.agy env builder.
TDD: written before implementation — all tests should fail on first run.
"""
from __future__ import annotations
from pathlib import Path
from headroom.providers.agy.runtime import build_agy_env
class TestBuildAgyEnv:
"""Pure-function tests for build_agy_env."""
def test_sets_https_and_http_proxy_to_terminator(self, tmp_path: Path) -> None:
bundle = tmp_path / "bundle.pem"
bundle.touch()
env = build_agy_env(
terminator_url="http://127.0.0.1:54321",
bundle_path=bundle,
base_env={},
)
assert env["HTTPS_PROXY"] == "http://127.0.0.1:54321"
assert env["HTTP_PROXY"] == "http://127.0.0.1:54321"
def test_sets_no_proxy_loopback(self, tmp_path: Path) -> None:
bundle = tmp_path / "bundle.pem"
bundle.touch()
env = build_agy_env(
terminator_url="http://127.0.0.1:54321",
bundle_path=bundle,
base_env={},
)
assert env["NO_PROXY"] == "127.0.0.1,localhost"
def test_preserves_inherited_no_proxy_entries(self, tmp_path: Path) -> None:
"""A corporate NO_PROXY names hosts that must bypass the proxy.
Replacing it would tunnel them through the terminator; the loopback
entries are prepended to what the user already had.
"""
bundle = tmp_path / "bundle.pem"
bundle.touch()
env = build_agy_env(
terminator_url="http://127.0.0.1:54321",
bundle_path=bundle,
base_env={"NO_PROXY": "internal.corp,10.0.0.0/8"},
)
assert env["NO_PROXY"] == "127.0.0.1,localhost,internal.corp,10.0.0.0/8"
def test_session_scoped_savings_vars_are_not_inherited(self, tmp_path: Path) -> None:
"""The wrapper redirects its OWN funnel to a temp dir deleted at exit.
The agy child (and the `headroom mcp serve` grandchild it spawns) must
not inherit that redirection, or it writes savings into a sink that
disappears and marks itself as an inbox emitter.
"""
bundle = tmp_path / "bundle.pem"
bundle.touch()
env = build_agy_env(
terminator_url="http://127.0.0.1:54321",
bundle_path=bundle,
base_env={
"HEADROOM_AGY_INBOX_EMIT": "1",
"HEADROOM_SAVINGS_PATH": "/tmp/gone/proxy_savings.json",
"HEADROOM_SAVINGS_EVENTS_PATH": "/tmp/gone/savings_events.jsonl",
"HEADROOM_OTEL_METRICS_ENABLED": "0",
"PATH": "/usr/bin",
},
)
assert "HEADROOM_AGY_INBOX_EMIT" not in env
assert "HEADROOM_SAVINGS_PATH" not in env
assert "HEADROOM_SAVINGS_EVENTS_PATH" not in env
assert "HEADROOM_OTEL_METRICS_ENABLED" not in env
assert env["PATH"] == "/usr/bin"
def test_sets_all_three_ca_vars_to_bundle(self, tmp_path: Path) -> None:
bundle = tmp_path / "bundle.pem"
bundle.touch()
env = build_agy_env(
terminator_url="http://127.0.0.1:54321",
bundle_path=bundle,
base_env={},
)
assert env["SSL_CERT_FILE"] == str(bundle)
assert env["CACERT_PATH"] == str(bundle)
assert env["NODE_EXTRA_CA_CERTS"] == str(bundle)
def test_corp_proxy_not_leaked_into_child_and_base_env_unmutated(self, tmp_path: Path) -> None:
"""A pre-existing corporate HTTPS_PROXY must NOT leak into the child agy
env as its proxy (the child must talk to the terminator), and build_agy_env
must NOT mutate base_env — so the terminator, running in the PARENT process,
still reads the original corporate os.environ["HTTPS_PROXY"] for chaining
non-allowlisted CONNECTs."""
bundle = tmp_path / "bundle.pem"
bundle.touch()
upstream = "http://corp-proxy.internal:3128"
base_env = {"HTTPS_PROXY": upstream}
env = build_agy_env(
terminator_url="http://127.0.0.1:54321",
bundle_path=bundle,
base_env=base_env,
)
# Child agy talks to the terminator, NOT the corp proxy directly.
assert env["HTTPS_PROXY"] == "http://127.0.0.1:54321"
# Corp proxy is preserved in the caller's env (parent keeps it for the
# terminator's blind-tunnel chaining); build_agy_env never clobbers it.
assert base_env["HTTPS_PROXY"] == upstream
# No dead chaining var is fabricated in the child env.
assert "HEADROOM_UPSTREAM_HTTPS_PROXY" not in env
def test_base_env_merged_into_result(self, tmp_path: Path) -> None:
"""Other base_env keys must be present in the returned dict."""
bundle = tmp_path / "bundle.pem"
bundle.touch()
env = build_agy_env(
terminator_url="http://127.0.0.1:54321",
bundle_path=bundle,
base_env={"MY_CUSTOM_KEY": "my_value"},
)
assert env["MY_CUSTOM_KEY"] == "my_value"
def test_returns_new_dict_does_not_mutate_base_env(self, tmp_path: Path) -> None:
bundle = tmp_path / "bundle.pem"
bundle.touch()
base = {"SOME_KEY": "val"}
result = build_agy_env(
terminator_url="http://127.0.0.1:54321",
bundle_path=bundle,
base_env=base,
)
assert result is not base
assert base == {"SOME_KEY": "val"}

305
tests/test_agy_registrar.py Normal file
View file

@ -0,0 +1,305 @@
"""Tests for headroom.mcp_registry.agy.AgyRegistrar.
All tests use a tmp_path home_dir seam so the real ~/.gemini is never touched.
"""
from __future__ import annotations
import json
from pathlib import Path
from headroom.mcp_registry.agy import AgyRegistrar
from headroom.mcp_registry.base import RegisterStatus, ServerSpec
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
_SPEC = ServerSpec(
name="headroom",
command="headroom",
args=("mcp", "serve"),
env={"HEADROOM_PROXY_URL": "http://127.0.0.1:9999"},
)
_OTHER_SPEC = ServerSpec(
name="other-server",
command="/usr/bin/other",
args=("--flag",),
env={},
)
def _make_reg(tmp_path: Path) -> AgyRegistrar:
return AgyRegistrar(home_dir=tmp_path)
def _config_path(tmp_path: Path) -> Path:
# agy 1.1.x read-path (migrated from .gemini/antigravity-cli/mcp_config.json).
return tmp_path / ".gemini" / "config" / "mcp_config.json"
def _write_config(tmp_path: Path, data: dict) -> None:
p = _config_path(tmp_path)
p.parent.mkdir(parents=True, exist_ok=True)
p.write_text(json.dumps(data, indent=2) + "\n")
def _read_config(tmp_path: Path) -> dict:
p = _config_path(tmp_path)
return json.loads(p.read_text())
# ---------------------------------------------------------------------------
# detect
# ---------------------------------------------------------------------------
class TestDetect:
def test_returns_false_when_dir_absent(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
assert reg.detect() is False
def test_returns_true_when_appdata_dir_exists(self, tmp_path: Path) -> None:
# detect() keys on agy's app-data dir (the stable install marker), not
# the migrated config dir (which may not exist until the first write).
(tmp_path / ".gemini" / "antigravity-cli").mkdir(parents=True, exist_ok=True)
reg = _make_reg(tmp_path)
assert reg.detect() is True
def test_returns_true_when_config_file_exists(self, tmp_path: Path) -> None:
_write_config(tmp_path, {"mcpServers": {}})
reg = _make_reg(tmp_path)
assert reg.detect() is True
# ---------------------------------------------------------------------------
# get_server
# ---------------------------------------------------------------------------
class TestGetServer:
def test_returns_none_when_file_absent(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
assert reg.get_server("headroom") is None
def test_returns_none_when_server_absent(self, tmp_path: Path) -> None:
_write_config(tmp_path, {"mcpServers": {}})
reg = _make_reg(tmp_path)
assert reg.get_server("headroom") is None
def test_returns_spec_when_present(self, tmp_path: Path) -> None:
_write_config(
tmp_path,
{
"mcpServers": {
"headroom": {
"command": "headroom",
"args": ["mcp", "serve"],
"env": {"HEADROOM_PROXY_URL": "http://127.0.0.1:9999"},
}
}
},
)
reg = _make_reg(tmp_path)
spec = reg.get_server("headroom")
assert spec is not None
assert spec.name == "headroom"
assert spec.command == "headroom"
assert tuple(spec.args) == ("mcp", "serve")
assert spec.env == {"HEADROOM_PROXY_URL": "http://127.0.0.1:9999"}
def test_malformed_json_returns_none(self, tmp_path: Path) -> None:
p = _config_path(tmp_path)
p.parent.mkdir(parents=True, exist_ok=True)
p.write_text("{not valid json")
reg = _make_reg(tmp_path)
assert reg.get_server("headroom") is None
# ---------------------------------------------------------------------------
# register_server — REGISTERED
# ---------------------------------------------------------------------------
class TestRegisterServer:
def test_malformed_config_is_not_overwritten(self, tmp_path: Path) -> None:
"""A config we cannot parse must abort the write, not get replaced.
``mcp_config.json`` is shared with the Antigravity IDE and holds the
user's own servers; treating an unreadable file as ``{}`` and writing our
single entry back would delete all of them.
"""
p = _config_path(tmp_path)
p.parent.mkdir(parents=True, exist_ok=True)
original = '{"mcpServers": {"user-server": {"command": "x"}}, TRUNCATED'
p.write_text(original)
result = _make_reg(tmp_path).register_server(_SPEC)
assert result.status == RegisterStatus.FAILED
assert p.read_text() == original
def test_unregister_leaves_malformed_config_untouched(self, tmp_path: Path) -> None:
p = _config_path(tmp_path)
p.parent.mkdir(parents=True, exist_ok=True)
original = "{not valid json"
p.write_text(original)
assert _make_reg(tmp_path).unregister_server("headroom") is False
assert p.read_text() == original
def test_registers_new_server(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
result = reg.register_server(_SPEC)
assert result.status == RegisterStatus.REGISTERED
config = _read_config(tmp_path)
assert "headroom" in config["mcpServers"]
entry = config["mcpServers"]["headroom"]
assert entry["command"] == "headroom"
assert entry["args"] == ["mcp", "serve"]
def test_registers_creates_parent_dirs(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
assert not _config_path(tmp_path).parent.exists()
reg.register_server(_SPEC)
assert _config_path(tmp_path).exists()
# ALREADY
def test_already_when_spec_matches(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
reg.register_server(_SPEC)
result = reg.register_server(_SPEC)
assert result.status == RegisterStatus.ALREADY
# MISMATCH without force
def test_mismatch_when_command_differs_no_force(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
reg.register_server(_SPEC)
different = ServerSpec(
name="headroom",
command="/different/path/headroom",
args=("mcp", "serve"),
env={"HEADROOM_PROXY_URL": "http://127.0.0.1:9999"},
)
result = reg.register_server(different)
assert result.status == RegisterStatus.MISMATCH
# Config must be unchanged
config = _read_config(tmp_path)
assert config["mcpServers"]["headroom"]["command"] == "headroom"
def test_mismatch_when_env_differs_no_force(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
reg.register_server(_SPEC)
different = ServerSpec(
name="headroom",
command="headroom",
args=("mcp", "serve"),
env={"HEADROOM_PROXY_URL": "http://127.0.0.1:1111"},
)
result = reg.register_server(different)
assert result.status == RegisterStatus.MISMATCH
# force overwrite
def test_force_overwrites_existing(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
reg.register_server(_SPEC)
updated = ServerSpec(
name="headroom",
command="/new/headroom",
args=("mcp", "serve"),
env={"HEADROOM_PROXY_URL": "http://127.0.0.1:9999"},
)
result = reg.register_server(updated, force=True)
assert result.status == RegisterStatus.REGISTERED
config = _read_config(tmp_path)
assert config["mcpServers"]["headroom"]["command"] == "/new/headroom"
# MERGE: other entries preserved
def test_merge_preserves_other_user_servers(self, tmp_path: Path) -> None:
# Pre-populate with a user-managed server.
_write_config(
tmp_path,
{
"mcpServers": {
"user-server": {
"command": "/usr/bin/user-mcp",
"args": ["--some-flag"],
}
}
},
)
reg = _make_reg(tmp_path)
result = reg.register_server(_SPEC)
assert result.status == RegisterStatus.REGISTERED
config = _read_config(tmp_path)
# Both entries must exist.
assert "user-server" in config["mcpServers"]
assert "headroom" in config["mcpServers"]
# User entry untouched.
assert config["mcpServers"]["user-server"]["command"] == "/usr/bin/user-mcp"
def test_missing_file_treated_as_empty(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
result = reg.register_server(_SPEC)
assert result.status == RegisterStatus.REGISTERED
def test_registers_spec_without_env(self, tmp_path: Path) -> None:
spec = ServerSpec(name="minimal", command="headroom", args=())
reg = _make_reg(tmp_path)
result = reg.register_server(spec)
assert result.status == RegisterStatus.REGISTERED
config = _read_config(tmp_path)
entry = config["mcpServers"]["minimal"]
# env key should not be present when empty.
assert "env" not in entry
# ---------------------------------------------------------------------------
# unregister_server
# ---------------------------------------------------------------------------
class TestUnregisterServer:
def test_returns_false_when_file_absent(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
assert reg.unregister_server("headroom") is False
def test_returns_false_when_server_absent(self, tmp_path: Path) -> None:
_write_config(tmp_path, {"mcpServers": {}})
reg = _make_reg(tmp_path)
assert reg.unregister_server("headroom") is False
def test_removes_named_server(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
reg.register_server(_SPEC)
removed = reg.unregister_server("headroom")
assert removed is True
config = _read_config(tmp_path)
assert "headroom" not in config["mcpServers"]
def test_unregister_only_removes_named_server(self, tmp_path: Path) -> None:
"""Unregistering 'headroom' MUST NOT remove other user entries."""
_write_config(
tmp_path,
{
"mcpServers": {
"user-server": {"command": "/bin/user-mcp"},
"headroom": {
"command": "headroom",
"args": ["mcp", "serve"],
},
}
},
)
reg = _make_reg(tmp_path)
reg.unregister_server("headroom")
config = _read_config(tmp_path)
assert "user-server" in config["mcpServers"]
assert "headroom" not in config["mcpServers"]
def test_idempotent_double_unregister(self, tmp_path: Path) -> None:
reg = _make_reg(tmp_path)
reg.register_server(_SPEC)
assert reg.unregister_server("headroom") is True
assert reg.unregister_server("headroom") is False

422
tests/test_agy_retrieve.py Normal file
View file

@ -0,0 +1,422 @@
"""Tests for headroom.proxy.agy_retrieve.AgyRetrieveServer.
The retrieve server is a PLAIN-HTTP loopback listener that serves the same
FastAPI app (``create_app()``) as the HTTPS dispatch server — it *is*
``AgyDispatchServer(plain_http=True)``, so the hypercorn plumbing under test
here lives in :mod:`headroom.proxy.agy_dispatch`. Its load-bearing property:
it shares the *process-global* compression store, so a marker stored on the
dispatch side resolves via ``GET /v1/retrieve/{hash}`` on this side.
All tests use ephemeral loopback ports; no TLS, no real network, no
``~/.headroom`` mutation beyond the in-memory process-global store (which is
reset around each test).
"""
from __future__ import annotations
import asyncio
import socket
import httpx
import pytest
from headroom.cache.compression_store import (
get_compression_store,
reset_compression_store,
)
from headroom.proxy import agy_dispatch
from headroom.proxy.agy_retrieve import AgyRetrieveServer
@pytest.fixture(autouse=True)
def _clean_compression_store():
"""Isolate the process-global compression store around each test."""
reset_compression_store()
yield
reset_compression_store()
async def test_retrieve_server_starts_on_loopback_plain_http() -> None:
"""Server binds loopback and answers plain HTTP (no TLS handshake)."""
srv = AgyRetrieveServer(port=0)
await srv.start()
try:
host, port = srv.address
assert host == "127.0.0.1"
assert isinstance(port, int) and port > 0
# Plain HTTP (http://) must succeed — proving there is NO TLS layer.
async with httpx.AsyncClient() as client:
resp = await client.get(f"http://127.0.0.1:{port}/v1/retrieve/stats")
assert resp.status_code == 200
finally:
await srv.stop()
async def test_get_retrieve_returns_store_populated_content() -> None:
"""LOAD-BEARING: a hash stored via the process-global store resolves over
plain HTTP from a SECOND create_app() — proving the cache is shared.
This is exactly the dispatch-populates / retrieve-resolves contract: the
HTTPS dispatch server stores markers into the same process-global singleton
that this plain-HTTP listener serves.
"""
original = '{"rows": [{"id": 1, "name": "alpha"}, {"id": 2, "name": "beta"}]}'
compressed = '{"rows": "[Retrieve more]"}'
# Populate the process-global store DIRECTLY (as the dispatch side would,
# via the same get_compression_store() singleton) — the server is a
# *separate* create_app() instance and must still see this entry.
store = get_compression_store()
hash_key = store.store(
original=original,
compressed=compressed,
original_tokens=42,
compressed_tokens=7,
tool_name="search_api",
)
srv = AgyRetrieveServer(port=0)
await srv.start()
try:
_, port = srv.address
async with httpx.AsyncClient() as client:
resp = await client.get(f"http://127.0.0.1:{port}/v1/retrieve/{hash_key}")
assert resp.status_code == 200
body = resp.json()
assert body["hash"] == hash_key
assert body["original_content"] == original
assert body["tool_name"] == "search_api"
finally:
await srv.stop()
async def test_get_unknown_hash_returns_404() -> None:
"""An unknown marker hash returns 404 (not a 500/hang)."""
srv = AgyRetrieveServer(port=0)
await srv.start()
try:
_, port = srv.address
async with httpx.AsyncClient() as client:
resp = await client.get(f"http://127.0.0.1:{port}/v1/retrieve/deadbeefdeadbeefdeadbeef")
assert resp.status_code == 404
finally:
await srv.stop()
async def test_retrieve_server_binds_loopback_only() -> None:
"""The listener socket family/host must be loopback (127.0.0.1)."""
srv = AgyRetrieveServer(port=0)
await srv.start()
try:
host, _ = srv.address
assert host == "127.0.0.1"
finally:
await srv.stop()
async def test_retrieve_server_clean_start_stop_no_leaked_server_tasks() -> None:
"""start()/stop() leaves no server-owned tasks (lifespan / connection).
The only surviving task may be the FastAPI app's *periodic* TOIN-stats
background task — an app-level concern that the production
``_start_agy_servers`` reaps at loop teardown (it cancels all pending tasks
in its ``finally`` before ``loop.close()``). This test mirrors that final
sweep and asserts every leftover is cancellable (i.e. no task wedges the
shutdown), and that the server's OWN lifespan task is gone.
"""
loop = asyncio.get_running_loop()
before = {t for t in asyncio.all_tasks(loop) if not t.done()}
srv = AgyRetrieveServer(port=0)
await srv.start()
await srv.stop()
assert srv._lifespan_task is None, "stop() must clear the lifespan task"
await asyncio.sleep(0)
after = {t for t in asyncio.all_tasks(loop) if not t.done()}
leaked = after - before
# Any leftover must be ONLY the app-level periodic stats task; no hypercorn
# connection / lifespan task may survive stop().
offending = [t for t in leaked if "_log_toin_stats_periodically" not in repr(t.get_coro())]
assert not offending, f"retrieve server leaked server-owned tasks: {offending}"
# Model the production loop-teardown sweep: every leftover cancels cleanly.
for task in leaked:
task.cancel()
if leaked:
await asyncio.gather(*leaked, return_exceptions=True)
async def test_retrieve_server_stop_idempotent() -> None:
"""stop() after stop() does not raise."""
srv = AgyRetrieveServer(port=0)
await srv.start()
await srv.stop()
await srv.stop() # idempotent
def test_retrieve_server_address_raises_before_start() -> None:
"""address property raises RuntimeError before start()."""
srv = AgyRetrieveServer(port=0)
with pytest.raises(RuntimeError):
_ = srv.address
async def test_start_raises_when_lifespan_startup_fails(monkeypatch: pytest.MonkeyPatch) -> None:
"""If the hypercorn lifespan task fails during startup, start() surfaces
that exception (instead of silently continuing on to bind a socket)."""
import hypercorn.asyncio.run as hypercorn_run
class _FailingLifespan:
def __init__(self, *_args: object, **_kwargs: object) -> None:
pass
async def handle_lifespan(self) -> None:
raise RuntimeError("lifespan startup boom")
async def wait_for_startup(self) -> None:
# Yield control so the handle_lifespan task (already scheduled by
# loop.create_task) runs to completion — synchronously raising —
# before this coroutine resumes and returns.
await asyncio.sleep(0)
monkeypatch.setattr(hypercorn_run, "Lifespan", _FailingLifespan)
srv = AgyRetrieveServer(port=0)
with pytest.raises(RuntimeError, match="lifespan startup boom"):
await srv.start()
# The failure must be surfaced before any socket gets bound.
assert srv._server is None
async def test_start_continues_when_lifespan_task_completes_without_exception(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""If the lifespan task is already ``done()`` by the time
``wait_for_startup()`` returns, but *without* an exception, start() must
NOT raise — it continues on to bind the socket normally."""
import hypercorn.asyncio.run as hypercorn_run
class _InstantSucceedingLifespan:
def __init__(self, *_args: object, **_kwargs: object) -> None:
pass
async def handle_lifespan(self) -> None:
return None
async def wait_for_startup(self) -> None:
# Yield so the handle_lifespan task (scheduled by
# loop.create_task) runs to completion synchronously — with no
# exception — before this coroutine resumes and returns.
await asyncio.sleep(0)
monkeypatch.setattr(hypercorn_run, "Lifespan", _InstantSucceedingLifespan)
srv = AgyRetrieveServer(port=0)
await srv.start() # must NOT raise: task is done(), but exception() is None
try:
assert srv._lifespan_task is not None
assert srv._lifespan_task.done()
assert srv._lifespan_task.exception() is None
host, port = srv.address
assert host == "127.0.0.1"
assert isinstance(port, int) and port > 0
finally:
await srv.stop()
async def test_start_uses_so_exclusiveaddruse_on_non_posix(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""On a non-POSIX ``os.name`` (e.g. Windows), the listener applies the
exclusive-address-use socket option instead of SO_REUSEADDR — per the
module docstring, plain SO_REUSEADDR on Windows would let a second
process bind the same loopback port and intercept decrypted retrieve
traffic."""
class _OsNameShim:
"""Proxies the real ``os`` module except for ``.name``.
We can't monkeypatch the real ``os.name`` attribute directly: pathlib
(used transitively by ``create_app()``/hypercorn ``Config()`` during
``start()``) also reads ``os.name`` to pick ``WindowsPath`` vs.
``PosixPath`` and would break. Instead we rebind agy_dispatch's own
module-level ``os`` reference (the shared plumbing this listener runs
on) to this shim, leaving the real ``os`` module untouched.
"""
def __init__(self, real_os: object, forced_name: str) -> None:
self._real_os = real_os
self.name = forced_name
def __getattr__(self, item: str) -> object:
return getattr(self._real_os, item)
monkeypatch.setattr(agy_dispatch, "os", _OsNameShim(agy_dispatch.os, "nt"))
# Real SO_EXCLUSIVEADDRUSE only exists on Windows; alias it to
# SO_REUSEADDR's numeric value so the real setsockopt() syscall below
# succeeds on this (POSIX) test host.
monkeypatch.setattr(
agy_dispatch.socket, "SO_EXCLUSIVEADDRUSE", socket.SO_REUSEADDR, raising=False
)
setsockopt_calls: list[tuple[socket.socket, int, int, int]] = []
real_setsockopt = socket.socket.setsockopt
def _spy_setsockopt(
self: socket.socket, level: int, optname: int, value: int, *a: object, **kw: object
) -> None:
setsockopt_calls.append((self, level, optname, value))
real_setsockopt(self, level, optname, value, *a, **kw)
monkeypatch.setattr(socket.socket, "setsockopt", _spy_setsockopt)
class _FakeStartedServer:
"""Stand-in for the object asyncio.start_server() returns, so the
forced (non-posix) code window doesn't have to drive real
asyncio loop-internal connection machinery."""
def __init__(self, sock: socket.socket) -> None:
self.sockets = [sock]
def close(self) -> None:
self.sockets[0].close()
async def wait_closed(self) -> None:
return None
async def _fake_start_server(
_handler: object, sock: socket.socket | None = None, **_kw: object
) -> _FakeStartedServer:
assert sock is not None
return _FakeStartedServer(sock)
monkeypatch.setattr(agy_dispatch.asyncio, "start_server", _fake_start_server)
srv = AgyRetrieveServer(port=0)
await srv.start()
try:
listener = srv._server.sockets[0] # type: ignore[union-attr]
calls_on_listener = [c for c in setsockopt_calls if c[0] is listener]
# Exactly one setsockopt call was made on our listener socket, and it
# went through the (non-posix) elif branch — the `if os.name ==
# "posix"` branch never ran because we patched os.name to "nt".
assert calls_on_listener == [(listener, socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)]
# The applied option is actually in effect on the real socket. Read
# back a *non-zero* value rather than exactly 1: on macOS getsockopt()
# reports SO_REUSEADDR's internal bitmask (4) while Linux echoes the 1
# we set — both mean "enabled".
assert listener.getsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR) != 0
finally:
await srv.stop()
async def test_start_skips_sockopt_when_neither_posix_nor_exclusiveaddruse(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""When ``os.name`` isn't "posix" AND the platform lacks
SO_EXCLUSIVEADDRUSE, neither socket-opt branch applies — the listener is
still bound successfully, with no setsockopt call made at all."""
class _OsNameShim:
# See test_start_uses_so_exclusiveaddruse_on_non_posix for rationale:
# we rebind agy_dispatch's own module-level `os` reference rather
# than mutating the real `os` module (which pathlib etc. also read).
def __init__(self, real_os: object, forced_name: str) -> None:
self._real_os = real_os
self.name = forced_name
def __getattr__(self, item: str) -> object:
return getattr(self._real_os, item)
monkeypatch.setattr(agy_dispatch, "os", _OsNameShim(agy_dispatch.os, "nt"))
monkeypatch.delattr(socket, "SO_EXCLUSIVEADDRUSE", raising=False)
setsockopt_calls: list[tuple[socket.socket, int, int, int]] = []
real_setsockopt = socket.socket.setsockopt
def _spy_setsockopt(
self: socket.socket, level: int, optname: int, value: int, *a: object, **kw: object
) -> None:
setsockopt_calls.append((self, level, optname, value))
real_setsockopt(self, level, optname, value, *a, **kw)
monkeypatch.setattr(socket.socket, "setsockopt", _spy_setsockopt)
srv = AgyRetrieveServer(port=0)
await srv.start()
try:
listener = srv._server.sockets[0] # type: ignore[union-attr]
assert [c for c in setsockopt_calls if c[0] is listener] == []
host, port = srv.address
assert host == "127.0.0.1"
assert isinstance(port, int) and port > 0
finally:
await srv.stop()
async def test_stop_swallows_lifespan_shutdown_exception() -> None:
"""stop() must not propagate an exception raised by
``lifespan.wait_for_shutdown()`` — it logs/ignores it and still tears
down the rest of the server cleanly."""
srv = AgyRetrieveServer(port=0)
await srv.start()
async def _boom() -> None:
raise RuntimeError("shutdown boom")
assert srv._lifespan is not None
srv._lifespan.wait_for_shutdown = _boom # type: ignore[method-assign]
await srv.stop() # must not raise despite wait_for_shutdown() failing
assert srv._lifespan is None
assert srv._server is None
async def test_stop_swallows_lifespan_task_cancel_exception() -> None:
"""stop() must not propagate an exception raised while awaiting the
(just-cancelled) lifespan task — it cancels, swallows, and clears the
reference regardless."""
srv = AgyRetrieveServer(port=0)
await srv.start()
class _FakeCancelTask:
def __init__(self) -> None:
self.cancel_called = False
def cancel(self) -> None:
self.cancel_called = True
def __await__(self) -> object:
raise RuntimeError("await-after-cancel boom")
fake_task = _FakeCancelTask()
srv._lifespan_task = fake_task # type: ignore[assignment]
await srv.stop() # must not raise despite awaiting the fake task failing
assert fake_task.cancel_called is True
assert srv._lifespan_task is None
async def test_async_context_manager_starts_and_stops() -> None:
"""Used as ``async with``, the server starts on __aenter__ and stops on
__aexit__."""
async with AgyRetrieveServer(port=0) as srv:
assert isinstance(srv, AgyRetrieveServer)
host, port = srv.address
assert host == "127.0.0.1"
assert isinstance(port, int) and port > 0
async with httpx.AsyncClient() as client:
resp = await client.get(f"http://127.0.0.1:{port}/v1/retrieve/stats")
assert resp.status_code == 200
# __aexit__ ran stop(): lifespan task cleared, socket torn down.
assert srv._lifespan_task is None
with pytest.raises(RuntimeError):
_ = srv.address

238
tests/test_agy_retrieve_exemption.py Normal file
View file

@ -0,0 +1,238 @@
"""WU2-A (headroom-37g.17): agy cold-original retrieve-hash exemption.
agy resends full history with ORIGINAL tool outputs every turn (it never
rewrites local history to hold headroom's markers). The agy FR compressor
(``_compress_agy_function_responses``) therefore re-compresses the resent
cold original into the SAME marker every turn -- but a model that already
retrieved that hash via ``headroom_retrieve`` this turn should not be forced
to re-retrieve it again (observed 236x thrash without this exemption).
At parity with the Rust path (``live_zone.rs:2362-2384``, which exempts by
call_id), this exemption keys on the retrieved HASH itself -- agy has no
call_id: a functionResponse string leaf is exempt from compression iff a
``headroom_retrieve`` call for its default CCR hash appears ANYWHERE in the
same request's ``contents`` (any entry, historical or tail).
Scope:
1. cold-original leaf exempt when call_mcp_tool-shaped functionCall args
reference headroom_retrieve + the leaf's hash.
2. same, via a bare ``headroom_retrieve`` functionCall.
3. ordering-independent: retrieve call in a LATER contents[] entry than the
leaf still exempts it (proves the pre-scan runs before compression).
4. case-insensitive: an uppercased hash in args still exempts.
5. no over-exemption: a leaf whose hash is NOT retrieved is compressed.
6. convergence: running the compressor twice leaves the exempt leaf stable.
7. multi-leaf: only the leaf matching a retrieved hash is exempt; sibling
leaves are still compressed.
8. ``default_ccr_hash`` is the single source of truth shared with
``CompressionStore.store``'s default hash, and matches ``_FR_CCR_HASH_LEN``.
"""
from __future__ import annotations
from typing import Any
import pytest
from fastapi.testclient import TestClient
from headroom.cache.compression_store import (
default_ccr_hash,
get_compression_store,
reset_compression_store,
)
from headroom.proxy.handlers.gemini import _FR_CCR_HASH_LEN, _FR_CCR_MARKER_PREFIX
from headroom.proxy.server import ProxyConfig, create_app
from headroom.tokenizers import get_tokenizer
_MODEL = "gemini-3-flash-agent"
# Distinct large, single-line strings (no repeated lines, so lossless would be
# a no-op) -- well above the marker-derived compression floor.
_LEAF_A = "search result row alpha beta gamma delta epsilon zeta eta " * 40
_LEAF_B = "search result row omega psi chi phi upsilon tau sigma rho " * 40
_LEAF_C = "search result row kappa iota theta eta zeta epsilon delta " * 40
@pytest.fixture
def proxy() -> Any:
with TestClient(create_app(ProxyConfig(optimize=True))) as client:
yield client.app.state.proxy # type: ignore[attr-defined]
@pytest.fixture
def tok() -> Any:
return get_tokenizer(_MODEL)
@pytest.fixture
def ccr_store(monkeypatch: pytest.MonkeyPatch) -> Any:
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
reset_compression_store()
store = get_compression_store()
yield store
reset_compression_store()
def _fr_entry(leaf: Any, role: str = "user", name: str = "search") -> dict:
return {
"role": role,
"parts": [{"functionResponse": {"name": name, "response": {"output": leaf}}}],
}
def _fr_leaf(contents: list, entry: int, part: int = 0, key: str = "output") -> Any:
return contents[entry]["parts"][part]["functionResponse"]["response"][key]
def _mcp_retrieve_call_entry(hash_value: str, role: str = "model") -> dict:
"""Generic MCP dispatch shape: a ``call_mcp_tool`` functionCall whose args
reference ``headroom_retrieve`` and carry the target hash."""
return {
"role": role,
"parts": [
{
"functionCall": {
"name": "call_mcp_tool",
"args": {"tool": "headroom_retrieve", "arguments": {"hash": hash_value}},
}
}
],
}
def _bare_retrieve_call_entry(hash_value: str, role: str = "model") -> dict:
return {
"role": role,
"parts": [{"functionCall": {"name": "headroom_retrieve", "args": {"hash": hash_value}}}],
}
# ---------------------------------------------------------------------------
# 1. call_mcp_tool-shaped retrieve call exempts the matching cold leaf.
# ---------------------------------------------------------------------------
def test_mcp_dispatch_retrieve_call_exempts_matching_leaf(
proxy: Any, tok: Any, ccr_store: Any
) -> None:
h = default_ccr_hash(_LEAF_A)
contents = [_mcp_retrieve_call_entry(h), _fr_entry(_LEAF_A)]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 0
assert _fr_leaf(contents, 1) == _LEAF_A
assert not _fr_leaf(contents, 1).startswith(_FR_CCR_MARKER_PREFIX)
# ---------------------------------------------------------------------------
# 2. bare headroom_retrieve functionCall exempts the matching cold leaf.
# ---------------------------------------------------------------------------
def test_bare_retrieve_call_exempts_matching_leaf(proxy: Any, tok: Any, ccr_store: Any) -> None:
h = default_ccr_hash(_LEAF_A)
contents = [_bare_retrieve_call_entry(h), _fr_entry(_LEAF_A)]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 0
assert _fr_leaf(contents, 1) == _LEAF_A
# ---------------------------------------------------------------------------
# 3. Ordering: retrieve call AFTER the leaf still exempts it (pre-scan).
# ---------------------------------------------------------------------------
def test_retrieve_call_after_leaf_still_exempts(proxy: Any, tok: Any, ccr_store: Any) -> None:
h = default_ccr_hash(_LEAF_A)
contents = [_fr_entry(_LEAF_A), _mcp_retrieve_call_entry(h)]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 0
assert _fr_leaf(contents, 0) == _LEAF_A
# ---------------------------------------------------------------------------
# 4. Case-insensitivity: an uppercased hash in args still exempts.
# ---------------------------------------------------------------------------
def test_uppercased_hash_in_args_still_exempts(proxy: Any, tok: Any, ccr_store: Any) -> None:
h = default_ccr_hash(_LEAF_A).upper()
contents = [_mcp_retrieve_call_entry(h), _fr_entry(_LEAF_A)]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 0
assert _fr_leaf(contents, 1) == _LEAF_A
# ---------------------------------------------------------------------------
# 5. No over-exemption: a leaf whose hash is NOT retrieved is compressed.
# ---------------------------------------------------------------------------
def test_unrelated_retrieve_hash_does_not_exempt(proxy: Any, tok: Any, ccr_store: Any) -> None:
unrelated_hash = default_ccr_hash("something else entirely")
contents = [_mcp_retrieve_call_entry(unrelated_hash), _fr_entry(_LEAF_A)]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
assert leaves == 1
assert _fr_leaf(contents, 1).startswith(_FR_CCR_MARKER_PREFIX)
# ---------------------------------------------------------------------------
# 6. Convergence: f(f(x)) == f(x) for the exempt leaf.
# ---------------------------------------------------------------------------
def test_convergence_exempt_leaf_stable_across_runs(proxy: Any, tok: Any, ccr_store: Any) -> None:
h = default_ccr_hash(_LEAF_A)
contents = [_mcp_retrieve_call_entry(h), _fr_entry(_LEAF_A)]
proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
first = _fr_leaf(contents, 1)
assert first == _LEAF_A
proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
second = _fr_leaf(contents, 1)
assert second == first == _LEAF_A
# ---------------------------------------------------------------------------
# 7. Multi-leaf: only the leaf matching a retrieved hash is exempt.
# ---------------------------------------------------------------------------
def test_multi_leaf_only_matching_hash_exempt(proxy: Any, tok: Any, ccr_store: Any) -> None:
h_b = default_ccr_hash(_LEAF_B)
contents = [
_mcp_retrieve_call_entry(h_b),
{
"role": "user",
"parts": [
{
"functionResponse": {
"name": "search",
"response": {
"leaf_a": _LEAF_A,
"nested": {"leaf_b": _LEAF_B, "list": [_LEAF_C]},
},
}
}
],
},
]
before, after, leaves = proxy._compress_agy_function_responses(contents, "ccr", tok, ccr_store)
response = contents[1]["parts"][0]["functionResponse"]["response"]
assert leaves == 2 # leaf_a and leaf_c compressed; leaf_b exempt
assert response["leaf_a"].startswith(_FR_CCR_MARKER_PREFIX)
assert response["nested"]["list"][0].startswith(_FR_CCR_MARKER_PREFIX)
assert response["nested"]["leaf_b"] == _LEAF_B
# ---------------------------------------------------------------------------
# 8. default_ccr_hash is the single source of truth shared with the store.
# ---------------------------------------------------------------------------
def test_default_ccr_hash_matches_store_and_marker_length(ccr_store: Any) -> None:
store_hash = ccr_store.store(
_LEAF_A,
"compressed-placeholder",
original_tokens=1,
compressed_tokens=1,
tool_name="x",
)
assert default_ccr_hash(_LEAF_A) == store_hash
assert len(default_ccr_hash(_LEAF_A)) == _FR_CCR_HASH_LEN

212
tests/test_agy_retrieve_exposure_gate.py Normal file
View file

@ -0,0 +1,212 @@
"""Exposure gate for agy ``headroom_retrieve`` (headroom-h76.5).
The wrap↔child MCP ``initialize`` handshake proves only that wrap can spawn the
retrieve child; it does NOT prove agy will surface the tool. agy exposes tools
only from its persistent per-tool cache (``<appdata>/mcp/<server>/<tool>.json``),
so a registered-then-reverted entry is rejected at call time as
"Unknown tool: headroom_retrieve". These tests pin the exposure signal that
gates ``HEADROOM_AGY_RETRIEVE_WIRED`` — the flag that keeps ccr compression on —
so unrecoverable markers never ship on a false-positive handshake.
"""
from __future__ import annotations
import os
from pathlib import Path
from types import SimpleNamespace
from typing import Any
import pytest
from headroom.ccr.mcp_server import CCR_TOOL_NAME
from headroom.cli.wrap import (
_agy_exposes_retrieve_tool,
_ccr_backend_is_cross_process,
)
from headroom.mcp_registry import build_headroom_spec
from headroom.mcp_registry.agy import AgyRegistrar
def _registrar(tmp_path: Path) -> AgyRegistrar:
return AgyRegistrar(home_dir=tmp_path)
def _write_tool_cache(reg: AgyRegistrar, tool: str = CCR_TOOL_NAME) -> None:
"""Simulate agy caching a discovered tool for the headroom server.
Cache lives under agy's app-data dir (``cache_dir``), decoupled from the
(migrated) config dir.
"""
cache = reg.cache_dir / "headroom" / f"{tool}.json"
cache.parent.mkdir(parents=True, exist_ok=True)
cache.write_text(f'{{"name": "{tool}"}}')
class TestBackendCrossProcess:
@pytest.mark.parametrize(
"value,expected",
[
(None, True), # unset → default sqlite → shared
("", True),
("sqlite", True),
("redis", True), # external shared store
("memory", False), # per-process dict → child sees empty store
("MEMORY", False), # case-insensitive
(" memory ", False), # whitespace-insensitive
],
)
def test_only_memory_is_process_local(
self, monkeypatch: pytest.MonkeyPatch, value: str | None, expected: bool
) -> None:
if value is None:
monkeypatch.delenv("HEADROOM_CCR_BACKEND", raising=False)
else:
monkeypatch.setenv("HEADROOM_CCR_BACKEND", value)
assert _ccr_backend_is_cross_process() is expected
class TestExposureSignal:
"""All three conjuncts required: live config entry + tool cache + shared backend."""
def test_all_present_is_exposed(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("HEADROOM_CCR_BACKEND", raising=False)
reg = _registrar(tmp_path)
reg.register_server(build_headroom_spec(), force=True)
_write_tool_cache(reg)
assert _agy_exposes_retrieve_tool(reg) is True
def test_missing_config_entry_not_exposed(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("HEADROOM_CCR_BACKEND", raising=False)
reg = _registrar(tmp_path)
# cache present, but NO mcp_config "headroom" entry (reverted per-run entry)
_write_tool_cache(reg)
assert _agy_exposes_retrieve_tool(reg) is False
def test_missing_tool_cache_not_exposed(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("HEADROOM_CCR_BACKEND", raising=False)
reg = _registrar(tmp_path)
reg.register_server(build_headroom_spec(), force=True)
# config entry present, but agy has NOT cached the tool (never discovered)
assert _agy_exposes_retrieve_tool(reg) is False
def test_wrong_tool_cached_not_exposed(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("HEADROOM_CCR_BACKEND", raising=False)
reg = _registrar(tmp_path)
reg.register_server(build_headroom_spec(), force=True)
# a different tool cached under headroom/ must not count
_write_tool_cache(reg, tool="something_else")
assert _agy_exposes_retrieve_tool(reg) is False
def test_memory_backend_forces_unverified(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("HEADROOM_CCR_BACKEND", "memory")
reg = _registrar(tmp_path)
reg.register_server(build_headroom_spec(), force=True)
_write_tool_cache(reg)
# config + cache present, but a per-process store can't resolve hashes
assert _agy_exposes_retrieve_tool(reg) is False
class TestWiredGate:
"""The ``agy()`` call site sets WIRED only on positive exposure.
Behavioral spy on the exposure probe: the boolean passed to the downgrade
warning IS the gated signal that drives ``HEADROOM_AGY_RETRIEVE_WIRED``, so
asserting on it proves the gate without exec'ing agy.
"""
@pytest.mark.parametrize("exposed", [True, False])
def test_wired_follows_exposure(self, monkeypatch: pytest.MonkeyPatch, exposed: bool) -> None:
import headroom.cli.wrap as wrap_mod
for key in (
"HEADROOM_AGY_FR_MODE",
"HEADROOM_AGY_RETRIEVE_WIRED",
"HEADROOM_BACKEND",
"HEADROOM_SAVINGS_PATH",
"HEADROOM_SAVINGS_EVENTS_PATH",
"HEADROOM_OTEL_METRICS_ENABLED",
"HEADROOM_AGY_INBOX_EMIT",
):
monkeypatch.delenv(key, raising=False)
monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/agy" if name == "agy" else None)
monkeypatch.setattr(
"headroom.proxy.agy_ca.ensure_root_ca", lambda: (None, None, None, None)
)
monkeypatch.setattr("headroom.proxy.agy_ca.build_combined_bundle", lambda: "/dev/null")
monkeypatch.setattr("headroom.providers.agy.build_agy_env", lambda **kwargs: {})
class _FakeStats:
def snapshot_start(self) -> None:
pass
def print_summary(self, handler: Any) -> None:
pass
monkeypatch.setattr("headroom.providers.agy.stats.AgySessionStats", _FakeStats)
monkeypatch.setattr("headroom.providers.agy.stats.install_fail_open_handler", lambda: None)
monkeypatch.setattr(
"headroom.providers.agy.stats.remove_fail_open_handler",
lambda handler: None,
)
class _FakeRegistrar:
name = "agy"
def unregister_server(self, name: str) -> bool:
return False
monkeypatch.setattr("headroom.mcp_registry.agy.AgyRegistrar", _FakeRegistrar)
monkeypatch.setattr("headroom.cli.wrap._disable_tokensave_mcp", lambda *a, **k: None)
monkeypatch.setattr("headroom.cli.wrap._disable_serena_mcp", lambda *a, **k: None)
fake_servers = SimpleNamespace(
terminator=SimpleNamespace(address=("127.0.0.1", 1)), retrieve_port=12345
)
monkeypatch.setattr("headroom.cli.wrap._start_agy_servers", lambda *a, **k: fake_servers)
monkeypatch.setattr("headroom.cli.wrap._stop_agy_servers", lambda servers: None)
# Handshake succeeds (registered) — exposure alone decides WIRED.
monkeypatch.setattr(
"headroom.cli.wrap._setup_headroom_retrieve_mcp_agy",
lambda *a, **k: True,
)
monkeypatch.setattr(
"headroom.cli.wrap._agy_exposes_retrieve_tool", lambda registrar: exposed
)
monkeypatch.setattr("headroom.cli.wrap._register_proxy_client", lambda *a, **k: None)
seen: list[bool] = []
def _spy(retrieve_wired: bool) -> None:
seen.append(retrieve_wired)
raise SystemExit(0)
monkeypatch.setattr("headroom.cli.wrap._maybe_warn_agy_ccr_downgrade", _spy)
monkeypatch.setattr("subprocess.run", lambda *a, **k: SimpleNamespace(returncode=0))
with pytest.raises(SystemExit):
wrap_mod.agy.callback(
port=8899,
no_proxy=True,
no_intercept=False,
backend=None,
no_mcp=False,
no_serena=True,
no_tokensave=True,
code_graph=False,
agy_args=(),
)
assert seen == [exposed]
if exposed:
assert os.environ.get("HEADROOM_AGY_RETRIEVE_WIRED") == "1"
else:
assert "HEADROOM_AGY_RETRIEVE_WIRED" not in os.environ

183
tests/test_agy_retrieve_persistent.py Normal file
View file

@ -0,0 +1,183 @@
"""Persistent, local-store-backed headroom_retrieve registration (headroom-h76.6).
The retrieve MCP is registered with agy PERSISTENTLY and recorded in the install
ledger (mirroring codebase-memory-mcp / Serena) so agy caches and exposes the
tool across sessions. These tests pin: a stable port-independent spec, ledger
recording on REGISTERED and ALREADY, ledger-cleared handshake failure, and a
ledger-gated cooperative uninstall.
"""
from __future__ import annotations
from pathlib import Path
import pytest
from headroom.ccr.mcp_server import CCR_TOOL_NAME
from headroom.cli.wrap import (
_remove_headroom_installed_retrieve_mcp,
_setup_headroom_retrieve_mcp_agy,
)
from headroom.mcp_registry import build_headroom_spec
from headroom.mcp_registry.agy import AgyRegistrar
from headroom.mcp_registry.ledger import headroom_installed_matching
@pytest.fixture(autouse=True)
def _isolated_ledger(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
"""Redirect the install ledger to a tmp file (no global state touched)."""
ledger_file = tmp_path / "install_ledger.json"
monkeypatch.setattr("headroom.mcp_registry.ledger.ledger_path", lambda: ledger_file)
return ledger_file
def _reg(tmp_path: Path) -> AgyRegistrar:
return AgyRegistrar(home_dir=tmp_path / "home")
def _ledgered(reg: AgyRegistrar) -> bool:
return headroom_installed_matching(reg.name, reg.get_server("headroom"))
class TestSpecShape:
def test_stable_port_independent_local_store_spec(self) -> None:
spec = build_headroom_spec()
assert spec.name == "headroom"
# No ephemeral proxy URL -> child resolves from the on-disk store.
assert dict(spec.env) == {}
assert tuple(spec.args[-2:]) == ("mcp", "serve")
class TestPersistentRegistration:
def test_registers_and_records_ledger(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr("headroom.cli.wrap._smoke_verify_mcp_handshake", lambda *a, **k: True)
reg = _reg(tmp_path)
assert _setup_headroom_retrieve_mcp_agy(reg) is True
assert reg.get_server("headroom") is not None
assert _ledgered(reg) is True # persistent: recorded, not reverted
def test_idempotent_already_still_recorded(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr("headroom.cli.wrap._smoke_verify_mcp_handshake", lambda *a, **k: True)
reg = _reg(tmp_path)
assert _setup_headroom_retrieve_mcp_agy(reg) is True
# Second run hits ALREADY; record_install upserts, ledger stays valid.
assert _setup_headroom_retrieve_mcp_agy(reg) is True
assert _ledgered(reg) is True
def test_reclaims_ledger_after_loss(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr("headroom.cli.wrap._smoke_verify_mcp_handshake", lambda *a, **k: True)
reg = _reg(tmp_path)
# Pre-existing matching entry with NO ledger record (e.g. after the
# old-agy print-mode purge cleared it) — ALREADY must re-record.
reg.register_server(build_headroom_spec(), force=True)
assert _ledgered(reg) is False
assert _setup_headroom_retrieve_mcp_agy(reg) is True
assert _ledgered(reg) is True
def test_handshake_failure_removes_entry_and_clears_ledger(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
reg = _reg(tmp_path)
# First: succeed to seed a ledger record + entry.
monkeypatch.setattr("headroom.cli.wrap._smoke_verify_mcp_handshake", lambda *a, **k: True)
assert _setup_headroom_retrieve_mcp_agy(reg) is True
assert _ledgered(reg) is True
# Now a broken child: entry removed AND ledger cleared (no dead pointer,
# no stale ownership claim).
monkeypatch.setattr("headroom.cli.wrap._smoke_verify_mcp_handshake", lambda *a, **k: False)
assert _setup_headroom_retrieve_mcp_agy(reg) is False
assert reg.get_server("headroom") is None
assert _ledgered(reg) is False
class TestLedgerGatedUninstall:
def test_removes_ledgered_entry(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("headroom.cli.wrap._smoke_verify_mcp_handshake", lambda *a, **k: True)
reg = _reg(tmp_path)
_setup_headroom_retrieve_mcp_agy(reg)
assert _remove_headroom_installed_retrieve_mcp(reg) == "removed"
assert reg.get_server("headroom") is None
assert _ledgered(reg) is False
def test_leaves_non_ledgered_entry(self, tmp_path: Path) -> None:
reg = _reg(tmp_path)
# A user/fleet-managed "headroom" entry NOT recorded by wrap agy.
reg.register_server(build_headroom_spec(), force=True)
assert _remove_headroom_installed_retrieve_mcp(reg) == "not_headroom_owned"
assert reg.get_server("headroom") is not None # left untouched
def test_absent_entry_is_not_owned(self, tmp_path: Path) -> None:
reg = _reg(tmp_path)
assert _remove_headroom_installed_retrieve_mcp(reg) == "not_headroom_owned"
class TestMarkerToolAlignment:
def test_marker_names_exact_tool(self) -> None:
from headroom.transforms.agy_fr_compressor import _FR_CCR_MARKER_PREFIX
assert CCR_TOOL_NAME in _FR_CCR_MARKER_PREFIX
def test_tool_description_claims_headroom_markers(self) -> None:
import inspect
from headroom.ccr import mcp_server
src = inspect.getsource(mcp_server)
# Description must steer the model to this tool for headroom markers,
# disambiguating from any other expand/retrieve tool in the session.
assert "ONLY" in src and "Headroom compression markers" in src
assert "functionResponse compressed. Call headroom_retrieve" in src
class TestFirstRunToolCachePriming:
"""wrap agy pre-seeds agy's per-tool cache so retrieve is exposed run 1."""
def _cache_file(self, reg: AgyRegistrar) -> Path:
return reg.cache_dir / "headroom" / f"{CCR_TOOL_NAME}.json"
def test_setup_primes_retrieve_tool_cache(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
import json
from headroom.ccr.mcp_server import (
CCR_RETRIEVE_TOOL_DESCRIPTION,
CCR_RETRIEVE_TOOL_INPUT_SCHEMA,
)
monkeypatch.setattr("headroom.cli.wrap._smoke_verify_mcp_handshake", lambda *a, **k: True)
reg = _reg(tmp_path)
cache = self._cache_file(reg)
assert not cache.exists() # clean install: no cache yet
assert _setup_headroom_retrieve_mcp_agy(reg) is True
assert cache.is_file() # primed on the first setup, not after a relaunch
payload = json.loads(cache.read_text(encoding="utf-8"))
# Schema must mirror the live list_tools() entry (single source), with
# MCP inputSchema serialised under agy's ``parameters`` key.
assert payload == {
"name": CCR_TOOL_NAME,
"description": CCR_RETRIEVE_TOOL_DESCRIPTION,
"parameters": CCR_RETRIEVE_TOOL_INPUT_SCHEMA,
}
def test_priming_does_not_clobber_existing_cache(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr("headroom.cli.wrap._smoke_verify_mcp_handshake", lambda *a, **k: True)
reg = _reg(tmp_path)
cache = self._cache_file(reg)
cache.parent.mkdir(parents=True, exist_ok=True)
cache.write_text('{"name": "headroom_retrieve", "stale": true}', encoding="utf-8")
assert _setup_headroom_retrieve_mcp_agy(reg) is True
# An agy-written cache is authoritative; priming must not overwrite it.
assert cache.read_text(encoding="utf-8") == '{"name": "headroom_retrieve", "stale": true}'

178
tests/test_agy_savings_inbox.py Normal file
View file

@ -0,0 +1,178 @@
"""Isolated tests for the agy cross-process savings inbox.
These exercise :mod:`headroom.proxy.agy_savings_inbox` in a temp HOME so no
shared state is touched. The proxy funnel is replaced by a fake object whose
async ``record_request`` just records the kwargs it was called with.
"""
from __future__ import annotations
import json
from pathlib import Path
import pytest
from headroom.proxy import agy_savings_inbox
class FakeMetrics:
"""Stand-in for PrometheusMetrics: async record_request captures kwargs."""
def __init__(self) -> None:
self.calls: list[dict] = []
async def record_request(self, **kwargs) -> None:
self.calls.append(kwargs)
@pytest.fixture
def isolated_home(tmp_path, monkeypatch):
"""Point the workspace (and thus the inbox) at a throwaway dir."""
home = tmp_path / "home"
home.mkdir()
monkeypatch.setenv("HOME", str(home))
monkeypatch.setenv("HEADROOM_WORKSPACE_DIR", str(home / ".headroom"))
return home
_FUNNEL_KWARGS = {
"provider": "anthropic",
"model": "claude-sonnet",
"input_tokens": 1200,
"output_tokens": 340,
"tokens_saved": 800,
"latency_ms": 42.5,
"cached": False,
"overhead_ms": 3.0,
"ttfb_ms": 10.0,
"cache_read_tokens": 0,
"cache_write_tokens": 0,
"cache_write_5m_tokens": 0,
"cache_write_1h_tokens": 0,
"uncached_input_tokens": 1200,
"attempted_input_tokens": 2000,
"project": "myproj",
"client": "agy",
}
def _evt_files() -> list[Path]:
return sorted(agy_savings_inbox.inbox_dir().glob("evt-*.json"))
def test_emit_event_writes_one_roundtrippable_file(isolated_home):
agy_savings_inbox.emit_event(**_FUNNEL_KWARGS)
files = _evt_files()
assert len(files) == 1
envelope = json.loads(files[0].read_text())
assert envelope["v"] == agy_savings_inbox.SCHEMA_VERSION
assert isinstance(envelope["event_id"], str) and envelope["event_id"]
kwargs = envelope["kwargs"]
assert kwargs == _FUNNEL_KWARGS
# Sanity: the funnel-only fields are present...
assert kwargs["output_tokens"] == 340
assert kwargs["latency_ms"] == 42.5
# ...and SavingsTracker-only fields never leak in.
assert "total_input_tokens" not in kwargs
assert "total_input_cost_usd" not in kwargs
assert "timestamp" not in kwargs
@pytest.mark.asyncio
async def test_drain_replays_each_event_once_then_deletes(isolated_home):
agy_savings_inbox.emit_event(**_FUNNEL_KWARGS)
fake = FakeMetrics()
recorded = await agy_savings_inbox.drain_inbox(fake)
assert recorded == 1
assert fake.calls == [_FUNNEL_KWARGS]
assert _evt_files() == []
@pytest.mark.asyncio
async def test_redrain_dedups_and_survives_crash_window(isolated_home):
# Normal completed drain: event recorded, file gone, id in .processed.
agy_savings_inbox.emit_event(**_FUNNEL_KWARGS)
fake = FakeMetrics()
assert await agy_savings_inbox.drain_inbox(fake) == 1
assert len(fake.calls) == 1
# Re-draining does NOT re-record the same event id.
assert await agy_savings_inbox.drain_inbox(fake) == 0
assert len(fake.calls) == 1
# Crash window: an event whose id is already in .processed but whose evt
# file still exists (recorded, crashed before unlink) must be unlinked
# WITHOUT being recorded again.
inbox = agy_savings_inbox.inbox_dir()
seen_id = "crash-1"
(inbox / ".processed").write_text(seen_id + "\n")
(inbox / f"evt-{seen_id}.json").write_text(
json.dumps({"v": 1, "event_id": seen_id, "kwargs": _FUNNEL_KWARGS})
)
assert await agy_savings_inbox.drain_inbox(fake) == 0
assert len(fake.calls) == 1
assert not (inbox / f"evt-{seen_id}.json").exists()
@pytest.mark.asyncio
async def test_two_events_from_two_pids_recorded_once_each(isolated_home):
inbox = agy_savings_inbox.inbox_dir()
for pid in (111, 222):
eid = f"{pid}-0-abc"
(inbox / f"evt-{eid}.json").write_text(
json.dumps({"v": 1, "event_id": eid, "kwargs": _FUNNEL_KWARGS})
)
fake = FakeMetrics()
recorded = await agy_savings_inbox.drain_inbox(fake)
assert recorded == 2
assert len(fake.calls) == 2
assert _evt_files() == []
@pytest.mark.asyncio
async def test_malformed_event_skipped_not_fatal(isolated_home):
inbox = agy_savings_inbox.inbox_dir()
# Bad JSON file.
bad = inbox / "evt-bad.json"
bad.write_text("{ this is not json")
# A good event alongside it.
good_id = "999-0-def"
(inbox / f"evt-{good_id}.json").write_text(
json.dumps({"v": 1, "event_id": good_id, "kwargs": _FUNNEL_KWARGS})
)
fake = FakeMetrics()
recorded = await agy_savings_inbox.drain_inbox(fake)
# The malformed file is dropped; the good one still recorded.
assert recorded == 1
assert len(fake.calls) == 1
assert not bad.exists()
assert _evt_files() == []
@pytest.mark.asyncio
async def test_empty_inbox_returns_zero(isolated_home):
fake = FakeMetrics()
assert await agy_savings_inbox.drain_inbox(fake) == 0
assert fake.calls == []
def test_agy_emit_enabled_reflects_env(monkeypatch):
monkeypatch.delenv(agy_savings_inbox.AGY_INBOX_EMIT_ENV, raising=False)
assert agy_savings_inbox.agy_emit_enabled() is False
monkeypatch.setenv(agy_savings_inbox.AGY_INBOX_EMIT_ENV, "1")
assert agy_savings_inbox.agy_emit_enabled() is True
monkeypatch.setenv(agy_savings_inbox.AGY_INBOX_EMIT_ENV, "0")
assert agy_savings_inbox.agy_emit_enabled() is False

238
tests/test_agy_savings_integration.py Normal file
View file

@ -0,0 +1,238 @@
"""WU3 integration test: agy inbox event -> proxy drain -> real dashboard surfaces.
Proves the end-to-end replay path claimed by headroom-4l8: an event emitted by
agy, when drained by the shared proxy, moves the SAME in-memory metrics the
dashboard renders — the token-savings counter (``tokens_saved_total``, the source
of the dashboard token hero) AND the per-project SavingsTracker rows — and does
so exactly once across repeated drains (at-least-once + dedup).
Isolated: constructs a real ``PrometheusMetrics`` + ``SavingsTracker`` in-process,
no network, no live proxy, HOME pinned to a tmp dir. Never runs the broad suite.
"""
from __future__ import annotations
from pathlib import Path
import pytest
from headroom.proxy import agy_savings_inbox
from headroom.proxy.prometheus_metrics import PrometheusMetrics
from headroom.proxy.savings_tracker import SavingsTracker
def _verbose_log_blob() -> str:
"""A large, verbose pytest-style log — the shape Headroom's log compressor
collapses dramatically (dedupes repetitive PASS/INFO lines, keeps errors)."""
lines = ["============================= test session starts ============================="]
for i in range(600):
lines.append(f"tests/test_module_{i % 12}.py::test_case_{i} PASSED [{i % 100}%]")
lines.append(f"2026-07-05 18:00:{i % 60:02d},123 INFO worker.pool handled request id={i}")
lines += [
"tests/test_x.py::test_broken FAILED",
"E AssertionError: expected 3 got 4",
"======================== 1 failed, 1200 passed in 5.2s =========================",
]
return "Here is the failing test log. Find the root cause:\n\n" + "\n".join(lines)
def _event(project: str, *, tokens_saved: int, input_tokens: int) -> dict:
"""A minimal-but-complete funnel-kwargs payload for one agy request."""
return {
"provider": "anthropic",
"model": "claude-sonnet",
"input_tokens": input_tokens,
"output_tokens": 100,
"tokens_saved": tokens_saved,
"latency_ms": 25.0,
"cached": False,
"overhead_ms": 1.0,
"ttfb_ms": 5.0,
"cache_read_tokens": 0,
"cache_write_tokens": 0,
"cache_write_5m_tokens": 0,
"cache_write_1h_tokens": 0,
"uncached_input_tokens": input_tokens,
"attempted_input_tokens": input_tokens + tokens_saved,
"project": project,
"client": "agy",
}
@pytest.fixture
def isolated_home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
monkeypatch.setenv("HOME", str(tmp_path))
# Belt-and-suspenders: pin every savings sink under tmp so nothing global is touched.
monkeypatch.setenv("HEADROOM_SAVINGS_PATH", str(tmp_path / "proxy_savings.json"))
monkeypatch.setenv("HEADROOM_SAVINGS_EVENTS_PATH", str(tmp_path / "savings_events.jsonl"))
monkeypatch.setenv("HEADROOM_OTEL_METRICS_ENABLED", "0")
return tmp_path
@pytest.mark.parametrize(
("emit_marker", "expect_drain"),
[("1", False), ("", True)],
)
def test_emitting_process_never_drains(
isolated_home: Path,
monkeypatch: pytest.MonkeyPatch,
emit_marker: str,
expect_drain: bool,
) -> None:
"""A process that emits inbox events must not also drain them.
``wrap agy`` builds ``create_app()`` in-process (dispatch + retrieve) with its
savings paths redirected to a temp dir that is deleted at exit. A drain loop
there would consume events into that throwaway sink — and race the shared
proxy, which is the only process allowed to replay them.
"""
from fastapi.testclient import TestClient
from headroom.proxy import server as server_mod
from headroom.proxy.server import ProxyConfig
monkeypatch.setenv("HEADROOM_AGY_INBOX_EMIT", emit_marker)
started = False
async def _record(metrics: object, interval_seconds: int = 5) -> None:
nonlocal started
started = True
monkeypatch.setattr(server_mod, "_drain_agy_savings_periodically", _record)
with TestClient(server_mod.create_app(ProxyConfig(optimize=False))):
pass
assert started is expect_drain
async def test_drain_moves_token_hero_and_per_project(isolated_home: Path) -> None:
tracker = SavingsTracker(path=str(isolated_home / "proxy_savings.json"))
metrics = PrometheusMetrics(savings_tracker=tracker)
# Two agy requests in two projects land in the inbox.
agy_savings_inbox.emit_event(**_event("proj-a", tokens_saved=800, input_tokens=1200))
agy_savings_inbox.emit_event(**_event("proj-b", tokens_saved=300, input_tokens=500))
recorded = await agy_savings_inbox.drain_inbox(metrics)
assert recorded == 2
# Token hero source: the dashboard reads m.tokens_saved_total (server.py:2685).
assert metrics.tokens_saved_total == 1100
# Request-count fidelity: both requests are reflected, not just the savings.
assert metrics.requests_total == 2
# Per-project section: the dashboard reads savings_tracker.stats_preview()["projects"].
projects = metrics.savings_tracker.stats_preview()["projects"]
assert "proj-a" in projects and "proj-b" in projects
assert projects["proj-a"]["tokens_saved"] == 800
assert projects["proj-b"]["tokens_saved"] == 300
# Inbox drained empty.
assert not list(agy_savings_inbox.inbox_dir().glob("evt-*.json"))
async def test_redrain_does_not_double_count(isolated_home: Path) -> None:
tracker = SavingsTracker(path=str(isolated_home / "proxy_savings.json"))
metrics = PrometheusMetrics(savings_tracker=tracker)
agy_savings_inbox.emit_event(**_event("proj-a", tokens_saved=800, input_tokens=1200))
assert await agy_savings_inbox.drain_inbox(metrics) == 1
# A second drain with nothing new must not re-apply the event.
assert await agy_savings_inbox.drain_inbox(metrics) == 0
assert metrics.tokens_saved_total == 800
assert metrics.requests_total == 1
def test_real_agy_path_compression_reaches_dashboard(
isolated_home: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""End-to-end: a large agy (antigravity UA) cloudcode request is REALLY
compressed by the production pipeline, and the genuine savings delta reaches
the dashboard's metrics via the WU2 inbox. Asserts a real reduction (not a
magic number) so it is robust across compressor tuning.
"""
from fastapi.responses import StreamingResponse
from starlette.testclient import TestClient
from headroom.proxy.server import HeadroomProxy, ProxyConfig, create_app
monkeypatch.setenv("HEADROOM_AGY_INBOX_EMIT", "1")
async def _fake_stream(proxy_self, url, headers, body, *a, **k): # type: ignore[no-untyped-def]
async def _b():
yield b'data: {"candidates":[{"content":{"parts":[{"text":"ok"}]}}]}\r\n\r\ndata: [DONE]\r\n\r\n'
return StreamingResponse(_b(), status_code=200, media_type="text/event-stream")
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
body = {
"project": "agy-proof",
"model": "gemini-3-flash-agent",
"request": {"contents": [{"role": "user", "parts": [{"text": _verbose_log_blob()}]}]},
}
deltas: list[tuple[int, int]] = []
with TestClient(
create_app(
ProxyConfig(
optimize=True,
compress_user_messages=True,
protect_recent=0,
min_tokens_to_crush=100,
)
)
) as client:
proxy: HeadroomProxy = client.app.state.proxy # type: ignore[attr-defined]
real_apply = proxy.openai_pipeline.apply
def _spy(*a, **k): # type: ignore[no-untyped-def]
r = real_apply(*a, **k)
deltas.append((r.tokens_before, r.tokens_after))
return r
proxy.openai_pipeline.apply = _spy # type: ignore[method-assign]
resp = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": "antigravity/1.0.5"},
json=body,
)
assert resp.status_code == 200
assert deltas, "compression pipeline was not invoked on the agy path"
before, after = deltas[0]
saved = before - after
# Real, substantial compression happened on agy-shaped traffic.
assert saved > 0 and after < before, f"expected real compression, got {before}->{after}"
# The genuine delta flows through the WU2 inbox to the dashboard metrics.
agy_savings_inbox.emit_event(
provider="google",
model="gemini-3-flash-agent",
input_tokens=after,
output_tokens=5,
tokens_saved=saved,
latency_ms=30.0,
cached=False,
overhead_ms=0.0,
ttfb_ms=0.0,
cache_read_tokens=0,
cache_write_tokens=0,
cache_write_5m_tokens=0,
cache_write_1h_tokens=0,
uncached_input_tokens=after,
attempted_input_tokens=before,
project="agy-proof",
client="agy",
)
metrics = PrometheusMetrics(
savings_tracker=SavingsTracker(path=str(isolated_home / "dash.json"))
)
import asyncio
assert asyncio.run(agy_savings_inbox.drain_inbox(metrics)) == 1
assert metrics.tokens_saved_total == saved
assert metrics.savings_tracker.stats_preview()["projects"]["agy-proof"]["tokens_saved"] == saved

93
tests/test_agy_sse_usage_envelope.py Normal file
View file

@ -0,0 +1,93 @@
"""agy / Cloud Code Assist SSE usage: unwrap the response envelope (headroom-sit).
Cloud Code Assist wraps streaming chunks in a ``response`` envelope
(``{"response": {"usageMetadata": {...}}}``), mirroring the request-side wrap
(gemini.py ``body.get("request")``). Both gemini SSE usage parsers read
``usageMetadata`` at the top level only, so agy's ``candidatesTokenCount``
(output tokens) never parsed and every turn fell back to a bytes//40 estimate
(PR #1044 symptom (b): "Could not parse output_tokens from SSE, estimating ...").
Native-Gemini (top-level ``usageMetadata``) must keep working.
"""
import json
from headroom.proxy.server import HeadroomProxy
def _proxy() -> HeadroomProxy:
# The gemini branch of both parsers is pure JSON parsing — no proxy
# dependencies are touched, so a bare instance is sufficient.
return object.__new__(HeadroomProxy)
def _sse(payload: dict) -> bytes:
return f"data: {json.dumps(payload)}\n\n".encode()
# -- _parse_sse_usage_from_buffer (buffered path, primary agy streaming path) --
def test_buffer_gemini_unwraps_cloudcode_response_envelope():
chunk = {
"response": {
"usageMetadata": {
"promptTokenCount": 1234,
"candidatesTokenCount": 567,
"cachedContentTokenCount": 89,
}
}
}
state = {"sse_buffer": bytearray(_sse(chunk))}
usage = _proxy()._parse_sse_usage_from_buffer(state, "gemini")
assert usage == {
"input_tokens": 1234,
"output_tokens": 567,
"cache_read_input_tokens": 89,
}
def test_buffer_gemini_native_top_level_still_parses():
chunk = {"usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 20}}
state = {"sse_buffer": bytearray(_sse(chunk))}
usage = _proxy()._parse_sse_usage_from_buffer(state, "gemini")
assert usage["input_tokens"] == 10
assert usage["output_tokens"] == 20
# -- _parse_sse_usage (raw-chunk path) --
def test_chunk_gemini_unwraps_cloudcode_response_envelope():
chunk = {"response": {"usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 7}}}
usage = _proxy()._parse_sse_usage(_sse(chunk), "gemini")
assert usage is not None
assert usage["output_tokens"] == 7
assert usage["input_tokens"] == 5
def test_chunk_gemini_native_top_level_still_parses():
chunk = {"usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 4}}
usage = _proxy()._parse_sse_usage(_sse(chunk), "gemini")
assert usage is not None
assert usage["output_tokens"] == 4
# -- _gemini_usage_meta helper: guard malformed upstream (never crash) --
def test_gemini_usage_meta_guards_non_dict_metadata():
m = HeadroomProxy._gemini_usage_meta
# A truthy non-dict usageMetadata must NOT reach .get() and crash the parser.
assert m({"usageMetadata": "garbage"}) is None
assert m({"usageMetadata": [1, 2]}) is None
assert m({"usageMetadata": 42}) is None
assert m({"response": {"usageMetadata": "x"}}) is None
assert m({"response": "notadict"}) is None
assert m({}) is None
# Well-formed top-level and enveloped both resolve to the inner dict.
assert m({"usageMetadata": {"candidatesTokenCount": 7}}) == {"candidatesTokenCount": 7}
assert m({"response": {"usageMetadata": {"candidatesTokenCount": 7}}}) == {
"candidatesTokenCount": 7
}
# Empty-but-present dict passes through; callers skip it on falsy (no zero-overwrite).
assert m({"usageMetadata": {}}) == {}

439
tests/test_agy_stats.py Normal file
View file

@ -0,0 +1,439 @@
"""Unit tests for headroom.providers.agy.stats.
Tests are headless and isolated: no live agy, no network, no port :8787.
Ref: headroom-30y.15
"""
from __future__ import annotations
import logging
import threading
from typing import Any
from unittest.mock import MagicMock, patch
import pytest
from headroom.providers.agy.stats import (
_FAIL_OPEN_SUBSTR,
_GEMINI_LOGGER,
AgySessionStats,
FailOpenWarnHandler,
_format_summary,
_get_compression_stats,
install_fail_open_handler,
remove_fail_open_handler,
)
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _emit_fail_open(handler: FailOpenWarnHandler) -> None:
"""Emit a synthetic fail-open log record directly into *handler*."""
record = logging.LogRecord(
name=_GEMINI_LOGGER,
level=logging.WARNING,
pathname="",
lineno=883,
msg=f"[req-1] {_FAIL_OPEN_SUBSTR}: some error",
args=(),
exc_info=None,
)
handler.emit(record)
# ---------------------------------------------------------------------------
# FailOpenWarnHandler — one-shot notice
# ---------------------------------------------------------------------------
class TestFailOpenWarnHandler:
"""FailOpenWarnHandler emits exactly ONE user notice regardless of fire count."""
def test_emits_one_notice_on_first_fail_open(self, capsys: pytest.CaptureFixture[str]) -> None:
handler = FailOpenWarnHandler()
_emit_fail_open(handler)
captured = capsys.readouterr()
assert "Headroom: compression failed" in captured.err
assert "fail-open" in captured.err
def test_does_not_emit_second_notice(self, capsys: pytest.CaptureFixture[str]) -> None:
handler = FailOpenWarnHandler()
_emit_fail_open(handler)
capsys.readouterr() # drain first notice
_emit_fail_open(handler)
_emit_fail_open(handler)
captured = capsys.readouterr()
assert captured.err == "", "no second notice must be printed"
def test_counts_all_occurrences(self) -> None:
handler = FailOpenWarnHandler()
for _ in range(5):
_emit_fail_open(handler)
assert handler.fail_open_count == 5
def test_ignores_unrelated_warning(self, capsys: pytest.CaptureFixture[str]) -> None:
handler = FailOpenWarnHandler()
record = logging.LogRecord(
name=_GEMINI_LOGGER,
level=logging.WARNING,
pathname="",
lineno=1,
msg="Some unrelated warning",
args=(),
exc_info=None,
)
handler.emit(record)
captured = capsys.readouterr()
assert captured.err == ""
assert handler.fail_open_count == 0
def test_thread_safe_one_shot(self, capsys: pytest.CaptureFixture[str]) -> None:
"""Concurrent emit()s from multiple threads must produce exactly ONE notice."""
handler = FailOpenWarnHandler()
barrier = threading.Barrier(10)
def _fire() -> None:
barrier.wait()
# Capture via a temp stderr replacement per thread is unreliable;
# instead count the _warned flag transitions by checking stderr
# using capsys after all threads complete.
_emit_fail_open(handler)
threads = [threading.Thread(target=_fire) for _ in range(10)]
for t in threads:
t.start()
for t in threads:
t.join()
# count = 10, one-shot flag set exactly once
assert handler.fail_open_count == 10
# The one-shot flag must be True
assert handler._warned is True # noqa: SLF001
# ---------------------------------------------------------------------------
# install / remove lifecycle
# ---------------------------------------------------------------------------
class TestInstallRemoveHandler:
"""install_fail_open_handler adds; remove_fail_open_handler removes — no leak."""
def test_install_adds_handler_to_logger(self) -> None:
logger = logging.getLogger(_GEMINI_LOGGER)
before = list(logger.handlers)
handler = install_fail_open_handler()
try:
assert handler in logger.handlers
finally:
remove_fail_open_handler(handler)
assert list(logger.handlers) == before
def test_remove_is_idempotent(self) -> None:
handler = install_fail_open_handler()
remove_fail_open_handler(handler)
remove_fail_open_handler(handler) # must not raise
def test_handler_receives_real_log_record(self, capsys: pytest.CaptureFixture[str]) -> None:
logger = logging.getLogger(_GEMINI_LOGGER)
logger.setLevel(logging.WARNING)
handler = install_fail_open_handler()
try:
logger.warning(f"[req] {_FAIL_OPEN_SUBSTR}: boom")
finally:
remove_fail_open_handler(handler)
captured = capsys.readouterr()
assert "Headroom: compression failed" in captured.err
assert handler.fail_open_count == 1
def test_no_handler_leaks_after_remove(self) -> None:
logger = logging.getLogger(_GEMINI_LOGGER)
original_handlers = list(logger.handlers)
h = install_fail_open_handler()
remove_fail_open_handler(h)
assert logger.handlers == original_handlers
def test_remove_none_handler_returns_without_error(self) -> None:
"""Explicit None early-return (stats.py:221-222): no-op, no exception."""
logger = logging.getLogger(_GEMINI_LOGGER)
before = list(logger.handlers)
result = remove_fail_open_handler(None)
assert result is None
assert list(logger.handlers) == before
def test_remove_handler_swallows_removehandler_exception(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""removeHandler raising (stats.py:225-226) is swallowed, not propagated."""
handler = install_fail_open_handler()
logger = logging.getLogger(_GEMINI_LOGGER)
calls: list[logging.Handler] = []
def _raise(_h: logging.Handler) -> None:
calls.append(_h)
raise RuntimeError("boom")
monkeypatch.setattr(logger, "removeHandler", _raise)
try:
remove_fail_open_handler(handler) # must not raise despite removeHandler blowing up
# Prove the raising removeHandler was ACTUALLY invoked — otherwise the
# swallow branch (stats.py:225-226) is untested (a mutant that never
# calls removeHandler would leave calls == [] and fail here).
assert calls == [handler]
# And the handler is still attached (our stub raised before real removal).
assert handler in logger.handlers
finally:
monkeypatch.undo()
logger.removeHandler(handler) # actually detach; avoid cross-test leakage
assert handler not in logger.handlers
# ---------------------------------------------------------------------------
# Falsification: emit on the ACTUAL production logger ("headroom.proxy").
# This hardcodes the production logger name (gemini.py:25) — it does NOT use
# _GEMINI_LOGGER — so it MUST fail if the install target ever regresses to the
# child "headroom.proxy.handlers.gemini" (parent->child does not propagate).
# ---------------------------------------------------------------------------
class TestFailOpenOnProductionLogger:
"""Records emitted on "headroom.proxy" (gemini.py's logger) must be caught."""
# The exact logger gemini.py:25 uses. Hardcoded on purpose — independent of
# the stats module's _GEMINI_LOGGER constant so a regression is detectable.
PROD_LOGGER = "headroom.proxy"
def test_first_fail_open_on_prod_logger_emits_one_notice(
self, capsys: pytest.CaptureFixture[str]
) -> None:
prod = logging.getLogger(self.PROD_LOGGER)
prod.setLevel(logging.WARNING)
handler = install_fail_open_handler()
try:
prod.warning("[req-1] Cloud Code Assist optimization failed: boom")
first = capsys.readouterr()
assert "Headroom: compression failed" in first.err
assert handler.fail_open_count == 1
# Second such record on the real logger -> still ONE notice, count==2.
prod.warning("[req-2] Cloud Code Assist optimization failed: boom2")
second = capsys.readouterr()
assert second.err == "", "no second user notice must be printed"
assert handler.fail_open_count == 2
finally:
remove_fail_open_handler(handler)
def test_unrelated_warning_on_prod_logger_no_notice(
self, capsys: pytest.CaptureFixture[str]
) -> None:
prod = logging.getLogger(self.PROD_LOGGER)
prod.setLevel(logging.WARNING)
handler = install_fail_open_handler()
try:
prod.warning("[req] some unrelated headroom.proxy warning")
captured = capsys.readouterr()
assert captured.err == ""
assert handler.fail_open_count == 0
finally:
remove_fail_open_handler(handler)
# ---------------------------------------------------------------------------
# _format_summary — pure function
# ---------------------------------------------------------------------------
class TestFormatSummary:
def _make_stats(
self,
entry_count: int = 0,
orig: int = 0,
comp: int = 0,
) -> dict[str, Any]:
return {
"entry_count": entry_count,
"total_original_tokens": orig,
"total_compressed_tokens": comp,
}
def test_with_compression_data(self) -> None:
start = self._make_stats(entry_count=0, orig=0, comp=0)
end = self._make_stats(entry_count=3, orig=1000, comp=400)
summary = _format_summary(start, end)
assert "3 entries compressed" in summary
assert "1,000" in summary
assert "400" in summary
# Share of the original that survived — unambiguous in a way "0.40x" is not.
assert "40% of original" in summary
def test_divide_by_zero_guard_no_compression(self) -> None:
start = self._make_stats()
end = self._make_stats()
summary = _format_summary(start, end)
assert "n/a" in summary or "no compression" in summary
def test_fail_open_count_included_when_provided(self) -> None:
start = self._make_stats()
end = self._make_stats(entry_count=1, orig=500, comp=200)
summary = _format_summary(start, end, fail_open_count=3)
assert "3 fail-open" in summary
def test_fail_open_count_omitted_when_none(self) -> None:
start = self._make_stats()
end = self._make_stats(entry_count=1, orig=500, comp=200)
summary = _format_summary(start, end, fail_open_count=None)
assert "fail-open" not in summary
def test_none_start_returns_unavailable(self) -> None:
summary = _format_summary(None, self._make_stats())
assert "unavailable" in summary
def test_none_end_returns_unavailable(self) -> None:
summary = _format_summary(self._make_stats(), None)
assert "unavailable" in summary
def test_delta_is_correct_over_preexisting_entries(self) -> None:
"""Delta must subtract the baseline, not report absolute store totals."""
start = self._make_stats(entry_count=10, orig=5000, comp=2000)
end = self._make_stats(entry_count=13, orig=6500, comp=2800)
summary = _format_summary(start, end)
# 3 new entries, 1500 orig, 800 comp
assert "3 entries" in summary
assert "1,500" in summary
assert "800" in summary
def test_negative_delta_clamped_to_zero(self) -> None:
"""Entries can be evicted between snapshots; clamp negatives to 0."""
start = self._make_stats(entry_count=10, orig=5000, comp=2000)
end = self._make_stats(entry_count=8, orig=4800, comp=1900)
# Should not raise or produce negative numbers
summary = _format_summary(start, end)
assert "0 entries" in summary
# ---------------------------------------------------------------------------
# AgySessionStats — idempotent print_summary
# ---------------------------------------------------------------------------
class TestAgySessionStats:
"""print_summary is idempotent: prints exactly once."""
def _make_stats_patch(self, stats_list: list[dict[str, Any]]):
"""Patch _get_compression_stats to return successive values from stats_list."""
call_count = [0]
def _fake() -> dict[str, Any]:
idx = min(call_count[0], len(stats_list) - 1)
call_count[0] += 1
return stats_list[idx]
return patch("headroom.providers.agy.stats._get_compression_stats", side_effect=_fake)
def test_print_summary_outputs_once(self, capsys: pytest.CaptureFixture[str]) -> None:
start_snap = {"entry_count": 0, "total_original_tokens": 0, "total_compressed_tokens": 0}
end_snap = {"entry_count": 2, "total_original_tokens": 800, "total_compressed_tokens": 320}
with self._make_stats_patch([start_snap, end_snap]):
stats = AgySessionStats()
stats.snapshot_start()
stats.print_summary()
stats.print_summary() # second call must NOT print
captured = capsys.readouterr()
assert captured.err.count("Headroom agy session") == 1
def test_print_summary_idempotent_across_threads(
self, capsys: pytest.CaptureFixture[str]
) -> None:
start_snap = {"entry_count": 0, "total_original_tokens": 0, "total_compressed_tokens": 0}
end_snap = {"entry_count": 1, "total_original_tokens": 400, "total_compressed_tokens": 160}
with self._make_stats_patch([start_snap, end_snap, end_snap, end_snap]):
stats = AgySessionStats()
stats.snapshot_start()
barrier = threading.Barrier(4)
def _call_print() -> None:
barrier.wait()
# Redirect stderr per-thread is tricky; just count
stats.print_summary()
threads = [threading.Thread(target=_call_print) for _ in range(4)]
for t in threads:
t.start()
for t in threads:
t.join()
captured = capsys.readouterr()
assert captured.err.count("Headroom agy session") == 1
def test_snapshot_start_graceful_on_import_error(
self, capsys: pytest.CaptureFixture[str]
) -> None:
"""If compression_store is unavailable, snapshot_start must not raise."""
with patch(
"headroom.providers.agy.stats._get_compression_stats",
side_effect=ImportError("no compression_store"),
):
stats = AgySessionStats()
stats.snapshot_start() # must not raise
stats.print_summary()
captured = capsys.readouterr()
assert "unavailable" in captured.err
def test_print_summary_includes_fail_open_count(
self, capsys: pytest.CaptureFixture[str]
) -> None:
start_snap = {"entry_count": 0, "total_original_tokens": 0, "total_compressed_tokens": 0}
end_snap = {"entry_count": 1, "total_original_tokens": 200, "total_compressed_tokens": 80}
handler = FailOpenWarnHandler()
_emit_fail_open(handler)
_emit_fail_open(handler)
with self._make_stats_patch([start_snap, end_snap]):
stats = AgySessionStats()
stats.snapshot_start()
stats.print_summary(handler=handler)
captured = capsys.readouterr()
assert "2 fail-open" in captured.err
# ---------------------------------------------------------------------------
# _get_compression_stats — lazy-import delegation to the compression store
# ---------------------------------------------------------------------------
class TestGetCompressionStats:
"""_get_compression_stats (stats.py:82-87) delegates to
get_compression_store().get_stats(), importing the store lazily at call
time (the import happens inside the function body, not at module load)."""
def test_delegates_to_compression_store_get_stats(self) -> None:
fake_stats: dict[str, Any] = {
"entry_count": 7,
"max_entries": 1000,
"total_original_tokens": 1234,
"total_compressed_tokens": 567,
}
fake_store = MagicMock()
fake_store.get_stats.return_value = fake_stats
with patch(
"headroom.cache.compression_store.get_compression_store",
return_value=fake_store,
) as mock_get_store:
result = _get_compression_stats()
mock_get_store.assert_called_once_with()
fake_store.get_stats.assert_called_once_with()
assert result == fake_stats
assert result["entry_count"] == 7

1624
tests/test_agy_terminator.py Normal file
View file

File diff suppressed because it is too large Load diff

40
tests/test_cli/test_wrap_agy_port_alias.py Normal file
View file

@ -0,0 +1,40 @@
"""Regression for headroom-r9k: `wrap agy -p X` must route `-p` to agy (print
mode), not be swallowed as the proxy ``--port``.
508.1 (b5814ffa) added ``@click.option("--port", "-p", ...)`` to the agy
subcommand, but agy's own ``-p`` is ``--print``. Click consumed ``-p`` as
``--port`` before it reached ``agy_args`` -> ``wrap agy -p PROMPT`` failed with
"Invalid value for --port/-p". Fix: the agy ``--port`` option no longer carries
the ``-p`` short alias (long ``--port`` only), so ``-p`` flows through
``ignore_unknown_options`` into ``agy_args`` and ``_agy_print_mode`` recognizes it.
These are parse-level tests via ``make_context`` -- it parses args WITHOUT
invoking the command callback, so no proxy is started.
"""
from __future__ import annotations
from headroom.cli.wrap import _agy_print_mode, agy
def _port_option():
return next(p for p in agy.params if getattr(p, "name", None) == "port")
class TestAgyPortAliasNoShadow:
def test_dash_p_routes_to_agy_print_not_port(self) -> None:
ctx = agy.make_context("agy", ["-p", "hello"])
assert ctx.params["port"] == 8787 # -p did NOT set the proxy port
assert ctx.params["agy_args"] == ("-p", "hello")
# Ticket-mandated: proves routing to agy PRINT MODE, not just presence.
assert _agy_print_mode(ctx.params["agy_args"]) is True
def test_long_port_still_sets_port(self) -> None:
ctx = agy.make_context("agy", ["--port", "9000", "foo"])
assert ctx.params["port"] == 9000
assert ctx.params["agy_args"] == ("foo",)
def test_port_option_has_no_dash_p_alias(self) -> None:
opt = _port_option()
assert opt.opts == ["--port"]
assert opt.secondary_opts == []

69
tests/test_provider_proxy_routes.py
View file

@ -456,6 +456,75 @@ def test_proxy_route_helpers_prefer_legacy_targets_and_gemini_passthrough() -> N
assert proxy_routes._select_passthrough_base_url(proxy, {}) == "https://legacy.anthropic.test"
def test_cloudcode_host_base_allowlists_exact_hosts_only() -> None:
proxy_targets = importlib.import_module("headroom.providers.proxy_targets")
# Every allowlisted host maps to its own https URL (guards against an
# all-None regression where the helper rejects legitimate hosts too).
assert proxy_targets.DEFAULT_ALLOWLIST, "allowlist must be non-empty"
for host in proxy_targets.DEFAULT_ALLOWLIST:
assert proxy_targets.cloudcode_host_base(host) == f"https://{host}"
# SSRF guard: a suffix-collision host that the old endswith() check accepted
# is now rejected, as are empty / unrelated hosts.
assert proxy_targets.cloudcode_host_base("evilcloudcode-pa.googleapis.com") is None
assert proxy_targets.cloudcode_host_base("cloudcode-pa.googleapis.com.evil.test") is None
assert proxy_targets.cloudcode_host_base("") is None
# The catch-all base selection honours the allowlist forward, so an
# allowlisted Host on an unrecognised path routes back to that host
# instead of falling through to the provider-heuristic default.
class _Runtime:
@staticmethod
def model_metadata_provider(_headers: dict[str, str]) -> str:
return "anthropic"
class _Proxy:
provider_runtime = _Runtime()
ANTHROPIC_API_URL = "https://legacy.anthropic.test"
allowlisted = next(iter(proxy_targets.DEFAULT_ALLOWLIST))
assert (
proxy_targets.select_passthrough_base_url(_Proxy(), {"host": allowlisted})
== f"https://{allowlisted}"
)
def test_select_passthrough_rejects_forged_cloudcode_host() -> None:
proxy_routes = importlib.import_module("headroom.providers.proxy_routes")
proxy = type(
"Proxy",
(),
{
"ANTHROPIC_API_URL": "https://legacy.anthropic.test",
"GEMINI_API_URL": "https://legacy.gemini.test",
"provider_runtime": type(
"Runtime",
(),
{
"api_target": staticmethod(lambda provider: f"https://runtime.{provider}.test"),
"model_metadata_provider": staticmethod(lambda headers: "anthropic"),
},
)(),
},
)()
# Positive: an allowlisted host is forwarded back to itself via the same path.
assert (
proxy_routes._select_passthrough_base_url(
proxy, {"host": "daily-cloudcode-pa.googleapis.com"}
)
== "https://daily-cloudcode-pa.googleapis.com"
)
# SSRF fallthrough: a forged suffix-collision host is NOT echoed back; it
# falls through to the configured default selection instead.
forged = "evilcloudcode-pa.googleapis.com"
base = proxy_routes._select_passthrough_base_url(proxy, {"host": forged})
assert base != f"https://{forged}"
assert base == "https://legacy.anthropic.test"
def test_provider_specific_routes_delegate_to_expected_proxy_handlers(monkeypatch) -> None:
delegated: list[tuple[str, str, tuple[str, ...]]] = []

607
tests/test_proxy_agy_compression.py Normal file
View file

@ -0,0 +1,607 @@
"""Tests for agy/antigravity path in handle_google_cloudcode_stream.
Scope: compression behaviour, routing, stealth, SSE pass-through,
accept-encoding stripping, single-upstream-origination, gzip request body,
auth-redaction on the default log path, and fail-open observability.
All tests use TestClient(create_app(…)) — in-process, no real port bind.
ALL upstream/network calls are stubbed via monkeypatch on HeadroomProxy._stream_response
or HeadroomProxy.openai_pipeline (the compression pipeline).
Never contacts 8787 or any real network destination.
"""
from __future__ import annotations
import gzip
import json
import logging
from typing import Any
from urllib.parse import urlparse
import pytest
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.testclient import TestClient
from headroom.proxy.server import HeadroomProxy, ProxyConfig, create_app
# ---------------------------------------------------------------------------
# Shared fixture body — large enough that CompressionDecision.should_compress
# is True when optimize=True (default). Repeated text triggers compression.
# ---------------------------------------------------------------------------
_REPEAT_UNIT = "The quick brown fox jumps over the lazy dog. " * 60 # ~2 700 chars
_LARGE_AGY_BODY: dict[str, Any] = {
"project": "test-project-123",
"model": "gemini-3-flash-agent",
"request": {
"contents": [
{
"role": "user",
"parts": [{"text": _REPEAT_UNIT}],
}
]
},
}
# Minimal SSE payload the handler's _stream_response would return.
_SSE_PAYLOAD = (
b'data: {"candidates":[{"content":{"parts":[{"text":"ok"}]}}]}\r\n\r\ndata: [DONE]\r\n\r\n'
)
# ---------------------------------------------------------------------------
# Helper: build a minimal SSE StreamingResponse suitable for the stub
# ---------------------------------------------------------------------------
def _make_sse_streaming_response() -> StreamingResponse:
async def _body(): # type: ignore[return]
yield _SSE_PAYLOAD
return StreamingResponse(_body(), status_code=200, media_type="text/event-stream")
# ---------------------------------------------------------------------------
# 1. COMPRESSION DELTA — optimization runs on the cloudcode/antigravity path
# ---------------------------------------------------------------------------
def test_compression_delta_on_antigravity_path(monkeypatch: pytest.MonkeyPatch) -> None:
"""A sufficiently large/redundant body triggers the compression code path.
We spy on openai_pipeline.apply to confirm it is called at least once,
confirming the CloudCode/antigravity handler enters the compression branch
when should_compress=True. The spy wraps the real apply so the return
value is genuine (no fake result required).
Note: when monkeypatching an *instance* attribute, the function receives
no implicit self — use *args/**kwargs to capture the call faithfully.
"""
call_log: list[dict[str, Any]] = []
async def _fake_stream(
proxy_self: Any, url: str, headers: dict, body: dict, *args: Any, **kwargs: Any
) -> StreamingResponse:
return _make_sse_streaming_response()
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
with TestClient(create_app(ProxyConfig(optimize=True))) as client:
proxy: HeadroomProxy = client.app.state.proxy # type: ignore[attr-defined]
real_apply = proxy.openai_pipeline.apply
# Instance-level patch: function is called without implicit self.
def _spy_apply(*args: Any, **kwargs: Any) -> Any:
result = real_apply(*args, **kwargs)
call_log.append(
{
"tokens_before": result.tokens_before,
"tokens_after": result.tokens_after,
"transforms": result.transforms_applied,
}
)
return result
proxy.openai_pipeline.apply = _spy_apply # type: ignore[method-assign]
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": "antigravity/1.0.5"},
json=_LARGE_AGY_BODY,
)
assert response.status_code == 200
# Pipeline must have been called at least once.
assert len(call_log) >= 1, "openai_pipeline.apply was never called — compression path not taken"
# ---------------------------------------------------------------------------
# 2. CORRECT HOST — antigravity traffic routes to ANTIGRAVITY_DAILY_API_URL
# ---------------------------------------------------------------------------
def test_antigravity_routes_to_daily_endpoint(monkeypatch: pytest.MonkeyPatch) -> None:
"""antigravity UA → https://daily-cloudcode-pa.googleapis.com target URL."""
captured: list[str] = []
async def _fake_stream(proxy_self: Any, url: str, *args: Any, **kwargs: Any) -> JSONResponse:
captured.append(url)
return JSONResponse({"url": url})
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": "antigravity/1.0.5"},
json=_LARGE_AGY_BODY,
)
assert response.status_code == 200
assert len(captured) == 1
# Parse and compare scheme+host exactly (not a URL-prefix substring check).
parsed = urlparse(captured[0])
assert (parsed.scheme, parsed.netloc) == ("https", "daily-cloudcode-pa.googleapis.com"), (
f"Expected daily endpoint, got: {captured[0]}"
)
# ---------------------------------------------------------------------------
# 3. SSE PRESERVED — response Content-Type text/event-stream passes through
# ---------------------------------------------------------------------------
def test_sse_response_content_type_preserved(monkeypatch: pytest.MonkeyPatch) -> None:
"""StreamingResponse with text/event-stream is forwarded unchanged."""
async def _fake_stream(
proxy_self: Any, url: str, *args: Any, **kwargs: Any
) -> StreamingResponse:
return _make_sse_streaming_response()
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": "antigravity/1.0.5"},
json=_LARGE_AGY_BODY,
)
assert response.status_code == 200
assert "text/event-stream" in response.headers.get("content-type", ""), (
f"Expected text/event-stream content-type, got: {response.headers.get('content-type')}"
)
# ---------------------------------------------------------------------------
# 4a. STEALTH — no x-headroom-* headers reach upstream
# ---------------------------------------------------------------------------
def test_stealth_no_x_headroom_headers_upstream(monkeypatch: pytest.MonkeyPatch) -> None:
"""x-headroom-* headers are stripped before the upstream call (gemini.py:826)."""
captured_headers: dict[str, str] = {}
async def _fake_stream(
proxy_self: Any, url: str, headers: dict, *args: Any, **kwargs: Any
) -> JSONResponse:
captured_headers.update(headers)
return JSONResponse({"ok": True})
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={
"User-Agent": "antigravity/1.0.5",
"x-headroom-bypass": "true",
"x-headroom-user-id": "tester",
"x-headroom-mode": "passthrough",
},
json=_LARGE_AGY_BODY,
)
assert response.status_code == 200
x_headroom_keys = [k for k in captured_headers if k.lower().startswith("x-headroom-")]
assert x_headroom_keys == [], f"x-headroom-* headers leaked to upstream: {x_headroom_keys}"
# ---------------------------------------------------------------------------
# 4b. STEALTH — agy User-Agent is passed through unchanged
# ---------------------------------------------------------------------------
def test_stealth_agy_user_agent_unchanged(monkeypatch: pytest.MonkeyPatch) -> None:
"""The agy UA is not rewritten by the handler."""
captured_headers: dict[str, str] = {}
_AGY_UA = "antigravity/1.0.5 linux/x86_64"
async def _fake_stream(
proxy_self: Any, url: str, headers: dict, *args: Any, **kwargs: Any
) -> JSONResponse:
captured_headers.update(headers)
return JSONResponse({"ok": True})
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": _AGY_UA},
json=_LARGE_AGY_BODY,
)
assert response.status_code == 200
sent_ua = captured_headers.get("user-agent", "")
assert sent_ua == _AGY_UA, f"UA was rewritten: expected {_AGY_UA!r}, got {sent_ua!r}"
# ---------------------------------------------------------------------------
# 5. ACCEPT-ENCODING STRIPPED — handler removes it before upstream (gemini.py:817)
# ---------------------------------------------------------------------------
def test_accept_encoding_stripped_before_upstream(monkeypatch: pytest.MonkeyPatch) -> None:
"""Handler pops accept-encoding from headers before calling _stream_response."""
captured_headers: dict[str, str] = {}
async def _fake_stream(
proxy_self: Any, url: str, headers: dict, *args: Any, **kwargs: Any
) -> JSONResponse:
captured_headers.update(headers)
return JSONResponse({"ok": True})
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={
"User-Agent": "antigravity/1.0.5",
"Accept-Encoding": "gzip, deflate, br",
},
json=_LARGE_AGY_BODY,
)
assert response.status_code == 200
assert "accept-encoding" not in {k.lower() for k in captured_headers}, (
f"accept-encoding reached upstream: {captured_headers}"
)
# ---------------------------------------------------------------------------
# 6. SINGLE-UPSTREAM-ORIGINATION — _stream_response called exactly once
# ---------------------------------------------------------------------------
def test_single_upstream_origination(monkeypatch: pytest.MonkeyPatch) -> None:
"""_stream_response is called EXACTLY once per request (no duplicate origination)."""
call_count = 0
async def _fake_stream(proxy_self: Any, url: str, *args: Any, **kwargs: Any) -> JSONResponse:
nonlocal call_count
call_count += 1
return JSONResponse({"ok": True})
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": "antigravity/1.0.5"},
json=_LARGE_AGY_BODY,
)
assert response.status_code == 200
assert call_count == 1, f"_stream_response called {call_count} times (expected exactly 1)"
# ---------------------------------------------------------------------------
# 7. GZIP REQUEST BODY — handler decodes gzip-encoded JSON body correctly
# ---------------------------------------------------------------------------
def test_gzip_request_body_decoded_correctly(monkeypatch: pytest.MonkeyPatch) -> None:
"""If the client sends a gzip-encoded request body, _read_request_json decompresses it.
_read_request_body_bytes (helpers.py:2689) handles Content-Encoding: gzip.
We confirm that handle_google_cloudcode_stream successfully parses the body
(returns 200, not 400) and forwards the correct model to _stream_response.
"""
captured_body: dict[str, Any] = {}
async def _fake_stream(
proxy_self: Any, url: str, headers: dict, body: dict, *args: Any, **kwargs: Any
) -> JSONResponse:
captured_body.update(body)
return JSONResponse({"ok": True})
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
raw_json = json.dumps(_LARGE_AGY_BODY).encode()
compressed = gzip.compress(raw_json)
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={
"User-Agent": "antigravity/1.0.5",
"Content-Encoding": "gzip",
"Content-Type": "application/json",
},
content=compressed,
)
assert response.status_code == 200, (
f"Expected 200 for gzip body, got {response.status_code}: {response.text}"
)
assert captured_body.get("model") == "gemini-3-flash-agent", (
f"Body not correctly decoded: model={captured_body.get('model')!r}"
)
# ---------------------------------------------------------------------------
# 8. AUTH REDACTION on default log path — Bearer + x-goog-api-key must NOT
# appear in plaintext in default-level logs (caplog).
# ---------------------------------------------------------------------------
def test_auth_not_leaked_in_default_logs(
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
) -> None:
"""Authorization and x-goog-api-key values must not appear in headroom logs.
The handler uses log_outbound_headers (gemini.py:828-832) which only logs
stripped_count, never header values. This asserts that the default log path
does not leak secrets for the cloudcode/antigravity handler.
"""
SECRET_BEARER = "supersecret-bearer-token-xyz789"
SECRET_API_KEY = "AIzaSyFakeSecret1234567890"
async def _fake_stream(proxy_self: Any, url: str, *args: Any, **kwargs: Any) -> JSONResponse:
return JSONResponse({"ok": True})
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
with caplog.at_level(logging.DEBUG, logger="headroom"):
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={
"User-Agent": "antigravity/1.0.5",
"Authorization": f"Bearer {SECRET_BEARER}",
"x-goog-api-key": SECRET_API_KEY,
},
json=_LARGE_AGY_BODY,
)
assert response.status_code == 200
log_text = "\n".join(r.getMessage() for r in caplog.records)
assert SECRET_BEARER not in log_text, "Bearer token leaked into headroom logs"
assert SECRET_API_KEY not in log_text, "x-goog-api-key leaked into headroom logs"
# ---------------------------------------------------------------------------
# 9. FAIL-OPEN OBSERVABILITY — pipeline raises → original bytes forwarded,
# warning logged.
# ---------------------------------------------------------------------------
def test_fail_open_on_compression_pipeline_exception(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""If openai_pipeline.apply raises, handler falls through (fail-open) with
original messages and emits a warning log. _stream_response is still called
exactly once (original body forwarded, not dropped).
The production code at gemini.py:882-883:
except Exception as e:
logger.warning(f"[{request_id}] Cloud Code Assist optimization failed: {e}")
ensures the outer _stream_response call still proceeds with original messages.
We capture the warning via a direct logging.Handler installed on the
headroom.proxy logger to avoid scope-ordering issues between TestClient's
event-loop dispatch and pytest caplog's propagation-reset fixture.
"""
call_count = 0
upstream_body_received: dict[str, Any] = {}
async def _fake_stream(
proxy_self: Any, url: str, headers: dict, body: dict, *args: Any, **kwargs: Any
) -> JSONResponse:
nonlocal call_count
call_count += 1
upstream_body_received.update(body)
return JSONResponse({"ok": True})
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
def _exploding_apply(*_args: Any, **_kw: Any) -> None:
raise RuntimeError("Simulated compression pipeline failure")
# Direct handler on headroom.proxy so we capture regardless of propagation state.
warning_messages: list[str] = []
class _CapturingHandler(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
if record.levelno >= logging.WARNING:
warning_messages.append(record.getMessage())
proxy_logger = logging.getLogger("headroom.proxy")
cap_handler = _CapturingHandler()
proxy_logger.addHandler(cap_handler)
# Pin the emit logger's own level so WARNING records are enabled regardless
# of any ancestor level another test left raised (isEnabledFor walks parents).
prev_level = proxy_logger.level
proxy_logger.setLevel(logging.WARNING)
try:
with TestClient(create_app(ProxyConfig(optimize=True))) as client:
proxy: HeadroomProxy = client.app.state.proxy # type: ignore[attr-defined]
# Direct instance assignment (not monkeypatch.setattr) so the function
# is stored exactly as given — no implicit self when called.
proxy.openai_pipeline.apply = _exploding_apply # type: ignore[method-assign]
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": "antigravity/1.0.5"},
json=_LARGE_AGY_BODY,
)
finally:
proxy_logger.removeHandler(cap_handler)
proxy_logger.setLevel(prev_level)
# Fail-open: must not 500/502; upstream call must proceed.
assert response.status_code == 200, (
f"Expected fail-open 200, got {response.status_code}: {response.text}"
)
# Upstream called exactly once.
assert call_count == 1, f"_stream_response called {call_count} times (expected 1)"
# Original body forwarded (model unchanged).
assert upstream_body_received.get("model") == "gemini-3-flash-agent", (
f"Body not forwarded correctly: {upstream_body_received.get('model')!r}"
)
# Warning was emitted on the headroom.proxy logger.
assert any(
"optimization failed" in msg.lower() or "cloud code assist" in msg.lower()
for msg in warning_messages
), "Expected a warning about compression failure. Got: " + "\n".join(warning_messages)
# ---------------------------------------------------------------------------
# 10. FAIL-OPEN BODY IDENTITY — compression raises → original body forwarded
# byte-for-byte (no mutation, no gzip, no truncation).
# ---------------------------------------------------------------------------
def test_fail_open_compression_degrades_open(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Compression pipeline raises → request completes successfully (degrade open).
Complementary to test_fail_open_on_compression_pipeline_exception which
verifies status-200 + exactly-one upstream call + warning log. This test
pins the BODY IDENTITY guarantee: the body forwarded upstream when
compression explodes is identical to the original request body — no
tokens modified, no gzip wrapping, no partial writes.
Also verifies the agy session is not crashed: a second request in the
same session after the first fail-open also completes with status 200.
"""
received_bodies: list[dict[str, Any]] = []
async def _fake_stream(
proxy_self: Any, url: str, headers: dict, body: dict, *args: Any, **kwargs: Any
) -> JSONResponse:
received_bodies.append(dict(body))
return JSONResponse({"ok": True})
monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream)
def _exploding_apply(*_args: Any, **_kw: Any) -> None:
raise RuntimeError("Simulated compression pipeline failure — body identity check")
warning_messages: list[str] = []
class _CapturingHandler(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
if record.levelno >= logging.WARNING:
warning_messages.append(record.getMessage())
proxy_logger = logging.getLogger("headroom.proxy")
cap_handler = _CapturingHandler()
proxy_logger.addHandler(cap_handler)
# Pin the emit logger's own level so WARNING records are enabled regardless
# of any ancestor level another test left raised (isEnabledFor walks parents).
prev_level = proxy_logger.level
proxy_logger.setLevel(logging.WARNING)
try:
with TestClient(create_app(ProxyConfig(optimize=True))) as client:
proxy: HeadroomProxy = client.app.state.proxy # type: ignore[attr-defined]
proxy.openai_pipeline.apply = _exploding_apply # type: ignore[method-assign]
# First request — fails compression, must degrade open.
response1 = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": "antigravity/1.0.5"},
json=_LARGE_AGY_BODY,
)
# Second request in same session — session must still be alive.
response2 = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"User-Agent": "antigravity/1.0.5"},
json=_LARGE_AGY_BODY,
)
finally:
proxy_logger.removeHandler(cap_handler)
proxy_logger.setLevel(prev_level)
# Both requests must degrade open (200).
assert response1.status_code == 200, (
f"First fail-open request must return 200, got {response1.status_code}: {response1.text}"
)
assert response2.status_code == 200, (
f"Second request proves session not crashed; got {response2.status_code}: {response2.text}"
)
# Both requests must have reached upstream — session not aborted.
assert len(received_bodies) == 2, (
f"Expected 2 upstream calls (one per request); got {len(received_bodies)}"
)
# Body identity: every upstream call received the original (uncompressed) body.
for i, body in enumerate(received_bodies):
assert body.get("model") == _LARGE_AGY_BODY["model"], (
f"Request {i + 1}: model field mutated — body identity broken: {body.get('model')!r}"
)
assert body.get("project") == _LARGE_AGY_BODY["project"], (
f"Request {i + 1}: project field mutated — body identity broken: {body.get('project')!r}"
)
contents = body.get("request", {}).get("contents", [])
assert len(contents) == 1, (
f"Request {i + 1}: contents list mutated — expected 1 item, got {len(contents)}"
)
text = contents[0].get("parts", [{}])[0].get("text", "")
assert text == _REPEAT_UNIT, (
f"Request {i + 1}: text body mutated or truncated — body identity broken"
)
# Fail-open is observable: at least one warning logged per fail.
assert len(warning_messages) >= 2, (
f"Expected at least 2 warnings (one per fail-open); got {len(warning_messages)}: "
+ "\n".join(warning_messages)
)
for msg in warning_messages:
assert "optimization failed" in msg.lower() or "cloud code assist" in msg.lower(), (
f"Warning message does not mention compression failure: {msg!r}"
)
# ---------------------------------------------------------------------------
# CROSS-AGENT REGRESSION: aider wrap-env byte-identity
#
# test_cli/test_wrap_aider.py already covers the aider env builder:
# test_wrap_aider_sets_provider_envs asserts OPENAI_API_BASE + ANTHROPIC_BASE_URL
# and agent_type == "aider".
#
# Recorded as: covered: tests/test_cli/test_wrap_aider.py::test_wrap_aider_sets_provider_envs
# ---------------------------------------------------------------------------

201
tests/test_proxy_google_cloudcode_route_aliases.py
View file

@ -1,3 +1,4 @@
import pytest
from fastapi.responses import JSONResponse
from fastapi.testclient import TestClient
@ -64,7 +65,7 @@ def test_antigravity_cloudcode_route_uses_daily_endpoint(monkeypatch):
assert response.status_code == 200
assert response.json() == {
"url": "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:streamGenerateContent?alt=sse",
"url": "https://daily-cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse",
"provider": "gemini",
"model": "claude-sonnet-4-6",
}
@ -135,7 +136,7 @@ def test_antigravity_header_detection_is_case_insensitive(monkeypatch):
assert response.status_code == 200
assert response.json() == {
"url": "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:streamGenerateContent?alt=sse",
"url": "https://daily-cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse",
"provider": "gemini",
"model": "claude-opus-4-6-thinking",
}
@ -158,7 +159,7 @@ def test_antigravity_route_does_not_cross_route_to_cloudcode_override(monkeypatc
assert response.status_code == 200
assert response.json() == {
"url": "https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:streamGenerateContent?alt=sse",
"url": "https://daily-cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse",
"provider": "gemini",
"model": "claude-sonnet-4-6",
}
@ -196,3 +197,197 @@ def test_cloudcode_override_does_not_leak_between_app_instances(monkeypatch):
second.json()["url"]
== "https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"
)
# ---------------------------------------------------------------------------
# T4: agy agent-model + body detection; env override; Pi/OpenClaw non-regression
# ---------------------------------------------------------------------------
AGY_AGENT_BODY = {
"project": "my-gcp-project",
"model": "gemini-3-flash-agent",
"request": {
"contents": [
{
"role": "user",
"parts": [{"text": "Hello from agy"}],
}
]
},
}
def test_agy_agent_model_body_routes_to_daily_endpoint(monkeypatch):
"""agy traffic with agent-model name + project + request.contents hits non-sandbox daily host."""
async def fake_stream(self, url, _headers, _body, provider, model, *_args, **_kwargs): # type: ignore[no-untyped-def]
return JSONResponse({"url": url, "provider": provider, "model": model})
monkeypatch.setattr(HeadroomProxy, "_stream_response", fake_stream)
# No antigravity/ UA header on purpose: detection must rest SOLELY on the
# agy body shape (agent-model name + project + request.contents), so this
# test fails if the body-shape detection branch is removed.
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
json=AGY_AGENT_BODY,
)
assert response.status_code == 200
assert response.json() == {
"url": "https://daily-cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse",
"provider": "gemini",
"model": "gemini-3-flash-agent",
}
@pytest.mark.parametrize(
("connect_host", "expected_host"),
[
("cloudcode-pa.googleapis.com", "cloudcode-pa.googleapis.com"),
("daily-cloudcode-pa.googleapis.com", "daily-cloudcode-pa.googleapis.com"),
# Not an allowlisted Cloud Code host -> falls back to the default backend
# rather than letting a forged Host steer the upstream (SSRF).
("evil-cloudcode-pa.googleapis.com", "daily-cloudcode-pa.googleapis.com"),
],
)
def test_mitm_request_stays_on_the_host_the_client_connected_to(
monkeypatch, connect_host, expected_host
):
"""A MITM'd request must be re-originated to the host agy CONNECTed to.
The terminator allowlist covers both cloudcode-pa and daily-cloudcode-pa, so
resolving every antigravity request to one default would send the client's
request — and its bearer — to a backend it never selected.
"""
async def fake_stream(self, url, _headers, _body, provider, model, *_args, **_kwargs): # type: ignore[no-untyped-def]
return JSONResponse({"url": url, "provider": provider, "model": model})
monkeypatch.setattr(HeadroomProxy, "_stream_response", fake_stream)
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
headers={"host": connect_host},
json=ANTIGRAVITY_BODY,
)
assert response.status_code == 200
assert (
response.json()["url"]
== f"https://{expected_host}/v1internal:streamGenerateContent?alt=sse"
)
def test_headroom_antigravity_api_url_env_override(monkeypatch):
"""HEADROOM_ANTIGRAVITY_API_URL env var overrides the corrected default for antigravity traffic."""
async def fake_stream(self, url, _headers, _body, provider, model, *_args, **_kwargs): # type: ignore[no-untyped-def]
return JSONResponse({"url": url, "provider": provider, "model": model})
monkeypatch.setattr(HeadroomProxy, "_stream_response", fake_stream)
monkeypatch.setenv("HEADROOM_ANTIGRAVITY_API_URL", "https://my-custom-agy.example.com")
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
json=ANTIGRAVITY_BODY,
)
assert response.status_code == 200
assert (
response.json()["url"]
== "https://my-custom-agy.example.com/v1internal:streamGenerateContent?alt=sse"
)
def test_pi_openclaw_requesttype_agent_still_detected(monkeypatch):
"""Pi/OpenClaw requestType=='agent' detection is not broken by new agy checks."""
async def fake_stream(self, url, _headers, _body, provider, model, *_args, **_kwargs): # type: ignore[no-untyped-def]
return JSONResponse({"url": url, "provider": provider, "model": model})
monkeypatch.setattr(HeadroomProxy, "_stream_response", fake_stream)
pi_body = {
"project": "pi-project",
"model": "gemini-1.5-pro",
"requestType": "agent",
"userAgent": "pi-coding-agent",
"request": {"contents": [{"role": "user", "parts": [{"text": "ping"}]}]},
}
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:streamGenerateContent",
params={"alt": "sse"},
json=pi_body,
)
assert response.status_code == 200
assert (
response.json()["url"]
== "https://daily-cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse"
)
def test_agy_control_plane_passthrough_routes_to_cloudcode_host(monkeypatch):
"""agy's non-streamGenerateContent control-plane calls (loadCodeAssist,
setUserSettings, …) reach the catch-all and MUST be proxied back to the
Cloud Code host agy addressed — not the generic Gemini endpoint that the
x-goog-api-key header would otherwise select. Without this the MITM dispatch
404s agy's onboarding and agy never issues a generateContent call."""
async def fake_passthrough(self, request, base_url, *args, **kwargs): # type: ignore[no-untyped-def]
return JSONResponse({"base_url": base_url, "path": request.url.path})
monkeypatch.setattr(HeadroomProxy, "handle_passthrough", fake_passthrough)
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:loadCodeAssist",
headers={
"host": "daily-cloudcode-pa.googleapis.com",
# agy sends x-goog-api-key; this previously forced the generic
# Gemini host. The Cloud Code host check must win over it.
"x-goog-api-key": "test-key",
},
json={"metadata": {"pluginType": "GEMINI"}},
)
assert response.status_code == 200
body = response.json()
assert body["path"] == "/v1internal:loadCodeAssist"
assert body["base_url"] == "https://daily-cloudcode-pa.googleapis.com"
def test_agy_control_plane_passthrough_rejects_non_allowlisted_host(monkeypatch):
"""The v1internal control-plane branch forwards to the incoming Host only
when it is an allowlisted Cloud Code host. A look-alike host (e.g. a suffix
match like ``evilcloudcode-pa.googleapis.com``) must NOT be used as the
upstream — it falls back to the static cloudcode target, so a forged Host
header cannot steer the MITM passthrough to an attacker-controlled origin."""
async def fake_passthrough(self, request, base_url, *args, **kwargs): # type: ignore[no-untyped-def]
return JSONResponse({"base_url": base_url, "path": request.url.path})
monkeypatch.setattr(HeadroomProxy, "handle_passthrough", fake_passthrough)
with TestClient(create_app(ProxyConfig(optimize=False))) as client:
response = client.post(
"/v1internal:loadCodeAssist",
headers={
"host": "evilcloudcode-pa.googleapis.com",
"x-goog-api-key": "test-key",
},
json={"metadata": {"pluginType": "GEMINI"}},
)
assert response.status_code == 200
body = response.json()
assert body["base_url"] != "https://evilcloudcode-pa.googleapis.com"

1541
tests/test_wrap_agy.py Normal file
View file

File diff suppressed because it is too large Load diff

167
tests/test_wrap_agy_proxy_wiring.py Normal file
View file

@ -0,0 +1,167 @@
"""headroom-508.1: `wrap agy` ensures the shared 8787 proxy (drain -> dashboard).
These tests are FULLY ISOLATED from any real proxy: `_ensure_proxy`,
`_register_proxy_client`, `_make_cleanup`, `ensure_root_ca`,
`build_combined_bundle`, and `shutil.which` are all patched, and `agy()` is
short-circuited at the patched `_ensure_proxy` (before any MITM server or the
real agy launch). A throwaway ``--port`` is passed as a second safeguard so no
code path can contact port 8787. Nothing here starts, probes, or tears down a
real proxy.
Blocker regression guards (from the plan-review gate):
- BLOCKER 1: `_ensure_proxy` must run BEFORE `agy()` poisons `os.environ`
(HEADROOM_AGY_INBOX_EMIT / HEADROOM_SAVINGS_PATH / HEADROOM_SAVINGS_EVENTS_PATH
/ HEADROOM_OTEL_METRICS_ENABLED), else a freshly-spawned shared proxy inherits
those via `os.environ.copy()` and corrupts shared state.
- BLOCKER 2: teardown (`cleanup`) is wired into agy's `finally`. (The refcounted
correctness of `_make_cleanup`/`_register_proxy_client` is agent-agnostic and
covered in tests/test_cli/test_wrap_helpers.py; agy reuses them unchanged.)
"""
from __future__ import annotations
from typing import Any
import pytest
from click.testing import CliRunner
_POISON_VARS = (
"HEADROOM_AGY_INBOX_EMIT",
"HEADROOM_SAVINGS_PATH",
"HEADROOM_SAVINGS_EVENTS_PATH",
"HEADROOM_OTEL_METRICS_ENABLED",
)
_THROWAWAY_PORT = "59123" # never 8787; also unreached because _ensure_proxy is faked
def _get_main() -> Any:
from headroom.cli import main
return main
class _StopBeforeLaunch(SystemExit):
"""Raised by the fake _ensure_proxy to short-circuit agy() cleanly."""
def _isolate(monkeypatch: pytest.MonkeyPatch, record: dict) -> None:
"""Patch every collaborator agy() reaches up to and including _ensure_proxy,
so the command never touches a real proxy/port or the real ~/.headroom."""
import os
# agy binary present -> agy() does not bail early
monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/agy" if name == "agy" else None)
# CA setup (lazily imported from headroom.proxy.agy_ca inside agy())
monkeypatch.setattr(
"headroom.proxy.agy_ca.ensure_root_ca",
lambda: (b"kkey", b"ccert", "/tmp/k.pem", "/tmp/c.pem"),
)
monkeypatch.setattr("headroom.proxy.agy_ca.build_combined_bundle", lambda: "/tmp/bundle.pem")
# proxy lifecycle -> no real markers, no real proxy
monkeypatch.setattr("headroom.cli.wrap._register_proxy_client", lambda port: None)
monkeypatch.setattr(
"headroom.cli.wrap._make_cleanup",
lambda holder, port: record.setdefault("cleanup", _RecordingCleanup()),
)
def _fake_ensure_proxy(port: int, no_proxy: bool, **kwargs: Any) -> None:
# Snapshot env at call time: agy() must not have set the poison vars YET.
record["env_at_ensure"] = dict(os.environ)
record["ensure_args"] = {"port": port, "no_proxy": no_proxy, "kwargs": kwargs}
raise _StopBeforeLaunch(0)
monkeypatch.setattr("headroom.cli.wrap._ensure_proxy", _fake_ensure_proxy)
# Clean slate so any poison var in the snapshot can only come from agy().
for var in _POISON_VARS:
monkeypatch.delenv(var, raising=False)
def _invoke_agy(*extra_args: str) -> Any:
"""Run `wrap agy` under the isolation harness and assert it got there cleanly.
``_fake_ensure_proxy`` short-circuits with ``_StopBeforeLaunch(0)``, so a
non-zero exit or any other exception means the command died somewhere else —
in which case the recorded assertions below would be checking a run that
never happened.
"""
result = CliRunner().invoke(
_get_main(), ["wrap", "agy", "--port", _THROWAWAY_PORT, *extra_args]
)
assert result.exit_code == 0, (
f"wrap agy exited {result.exit_code} before the assertion point: "
f"{result.exception!r}\n{result.output}"
)
if result.exception is not None:
assert isinstance(result.exception, SystemExit), result.exception
return result
class _RecordingCleanup:
def __init__(self) -> None:
self.calls = 0
def __call__(self, *_a: Any, **_k: Any) -> None:
self.calls += 1
class TestAgyEnsuresSharedProxy:
def test_ensure_proxy_runs_before_env_poisoning(self, monkeypatch: pytest.MonkeyPatch) -> None:
record: dict = {}
_isolate(monkeypatch, record)
_invoke_agy()
assert "env_at_ensure" in record, "agy() never reached _ensure_proxy"
leaked = [v for v in _POISON_VARS if v in record["env_at_ensure"]]
assert leaked == [], (
f"env poisoned before _ensure_proxy (would corrupt shared proxy): {leaked}"
)
def test_ensure_proxy_called_with_agy_agent_type(self, monkeypatch: pytest.MonkeyPatch) -> None:
record: dict = {}
_isolate(monkeypatch, record)
_invoke_agy()
args = record.get("ensure_args", {})
assert args.get("port") == int(_THROWAWAY_PORT)
assert args.get("kwargs", {}).get("agent_type") == "agy"
assert args.get("no_proxy") is False
def test_no_proxy_flag_is_passed_through(self, monkeypatch: pytest.MonkeyPatch) -> None:
record: dict = {}
_isolate(monkeypatch, record)
_invoke_agy("--no-proxy")
assert record.get("ensure_args", {}).get("no_proxy") is True
def test_code_graph_flag_forwards_to_proxy_watcher(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
"""--code-graph drives the proxy-side watcher, not an agy MCP entry.
Upstream repurposed ``--code-graph``: every wrap subcommand forwards it
to ``_ensure_proxy``, which starts the live reindex watcher. agy must
follow that contract instead of registering codebase-memory-mcp itself.
"""
record: dict = {}
_isolate(monkeypatch, record)
_invoke_agy("--code-graph")
assert record.get("ensure_args", {}).get("kwargs", {}).get("code_graph") is True
def test_code_graph_defaults_off(self, monkeypatch: pytest.MonkeyPatch) -> None:
record: dict = {}
_isolate(monkeypatch, record)
_invoke_agy()
assert record.get("ensure_args", {}).get("kwargs", {}).get("code_graph") is False
def test_cleanup_runs_on_teardown(self, monkeypatch: pytest.MonkeyPatch) -> None:
record: dict = {}
_isolate(monkeypatch, record)
_invoke_agy()
cleanup = record.get("cleanup")
assert cleanup is not None, "_make_cleanup was never built"
# agy()'s finally must invoke cleanup() even though it short-circuited.
assert cleanup.calls >= 1, "cleanup() not called on teardown (proxy would leak)"

265
uv.lock generated
View file

@ -280,12 +280,12 @@ name = "any-llm-sdk"
version = "1.12.1"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "httpx", marker = "python_full_version >= '3.11'" },
{ name = "openai", marker = "python_full_version >= '3.11'" },
{ name = "openresponses-types", marker = "python_full_version >= '3.11'" },
{ name = "pydantic", marker = "python_full_version >= '3.11'" },
{ name = "rich", marker = "python_full_version >= '3.11'" },
{ name = "typing-extensions", marker = "python_full_version >= '3.11'" },
{ name = "httpx" },
{ name = "openai" },
{ name = "openresponses-types" },
{ name = "pydantic" },
{ name = "rich" },
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/21/d5/21ef27d031b72f0054ba0015a66cc1f10cda9410e4a3acb6e795d848ad0d/any_llm_sdk-1.12.1.tar.gz", hash = "sha256:76e043fcaa56fccfb375a908511869dc7dbf393dd97a82d06bb764d8201724e8", size = 152078, upload-time = "2026-03-18T13:13:01.735Z" }
wheels = [
@ -811,7 +811,7 @@ name = "coloredlogs"
version = "15.0.1"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "humanfriendly", marker = "python_full_version < '3.11'" },
{ name = "humanfriendly" },
]
sdist = { url = "https://files.pythonhosted.org/packages/cc/c7/eed8f27100517e8c0e6b923d5f0845d0cb99763da6fdee00478f91db7325/coloredlogs-15.0.1.tar.gz", hash = "sha256:7c991aa71a4577af2f82600d8f8f3a89f936baeaf9b50a9c197da014e5bf16b0", size = 278520, upload-time = "2021-06-11T10:22:45.202Z" }
wheels = [
@ -823,7 +823,7 @@ name = "colorlog"
version = "6.10.1"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "colorama", marker = "python_full_version >= '3.13' and sys_platform == 'win32'" },
{ name = "colorama", marker = "sys_platform == 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/a2/61/f083b5ac52e505dfc1c624eafbf8c7589a0d7f32daa398d2e7590efa5fda/colorlog-6.10.1.tar.gz", hash = "sha256:eb4ae5cb65fe7fec7773c2306061a8e63e02efc2c72eba9d27b0fa23c94f1321", size = 17162, upload-time = "2025-10-16T16:14:11.978Z" }
wheels = [
@ -1045,7 +1045,7 @@ name = "cuda-bindings"
version = "13.3.1"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "cuda-pathfinder", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
{ name = "cuda-pathfinder" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/a9/21/8464d133752951c154feafb3b65c297e7d80f301183d220bec4c830f1441/cuda_bindings-13.3.1-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:120fcc53d57903df529c3486962c56528cba5b7d6c57c99537320ed9922c8b86", size = 6073403, upload-time = "2026-05-29T23:11:36.22Z" },
@ -1080,43 +1080,43 @@ wheels = [
[package.optional-dependencies]
cublas = [
{ name = "nvidia-cublas", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
{ name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
cudart = [
{ name = "nvidia-cuda-runtime", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
cufft = [
{ name = "nvidia-cufft", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-cufft", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
{ name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
cufile = [
{ name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-cufile", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
cupti = [
{ name = "nvidia-cuda-cupti", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
curand = [
{ name = "nvidia-curand", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-curand", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
cusolver = [
{ name = "nvidia-cublas", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-cusolver", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-cublas", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
{ name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
{ name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
{ name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
cusparse = [
{ name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
{ name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
nvjitlink = [
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
nvrtc = [
{ name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
nvtx = [
{ name = "nvidia-nvtx", marker = "(python_full_version < '3.11' and platform_machine == 'AMD64' and sys_platform == 'win32') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" },
{ name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' or platform_machine == 'x86_64'" },
]
[[package]]
@ -1220,7 +1220,7 @@ name = "exceptiongroup"
version = "1.3.1"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "typing-extensions", marker = "python_full_version < '3.11'" },
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" }
wheels = [
@ -1637,7 +1637,7 @@ name = "gunicorn"
version = "26.0.0"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "packaging", marker = "python_full_version < '3.11' or sys_platform != 'win32'" },
{ name = "packaging" },
]
sdist = { url = "https://files.pythonhosted.org/packages/6d/b7/a4a3f632f823e432ce6bc65f62961b7980c898c77f075a2f7118cb3846fe/gunicorn-26.0.0.tar.gz", hash = "sha256:ca9346f85e3a4aeeb64d491045c16b9a35647abd37ea15efe53080eb8b090baf", size = 727286, upload-time = "2026-05-05T06:38:25.529Z" }
wheels = [
@ -1668,7 +1668,7 @@ wheels = [
[[package]]
name = "headroom-ai"
version = "0.36.0"
version = "0.36.5"
source = { editable = "." }
dependencies = [
{ name = "ast-grep-cli" },
@ -1689,11 +1689,13 @@ agno = [
]
all = [
{ name = "anthropic" },
{ name = "cryptography" },
{ name = "datasets" },
{ name = "fastapi" },
{ name = "fastembed" },
{ name = "httpx", extra = ["http2"] },
{ name = "huggingface-hub" },
{ name = "hypercorn" },
{ name = "jinja2" },
{ name = "magika" },
{ name = "mcp" },
@ -1745,9 +1747,11 @@ crewai = [
]
dev = [
{ name = "anthropic" },
{ name = "cryptography" },
{ name = "fastapi" },
{ name = "hnswlib" },
{ name = "httpx", extra = ["http2"] },
{ name = "hypercorn" },
{ name = "langchain-ollama" },
{ name = "litellm", marker = "python_full_version < '3.14'" },
{ name = "mypy" },
@ -1818,8 +1822,10 @@ otel = [
{ name = "opentelemetry-sdk" },
]
proxy = [
{ name = "cryptography" },
{ name = "fastapi" },
{ name = "httpx", extra = ["http2"] },
{ name = "hypercorn" },
{ name = "magika" },
{ name = "mcp" },
{ name = "onnxruntime", version = "1.23.2", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version < '3.11'" },
@ -1834,9 +1840,11 @@ proxy = [
{ name = "zstandard" },
]
proxy-prod = [
{ name = "cryptography" },
{ name = "fastapi" },
{ name = "gunicorn", marker = "sys_platform != 'win32'" },
{ name = "httpx", extra = ["http2"] },
{ name = "hypercorn" },
{ name = "magika" },
{ name = "mcp" },
{ name = "onnxruntime", version = "1.23.2", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version < '3.11'" },
@ -1863,9 +1871,11 @@ reports = [
{ name = "jinja2" },
]
sandbox = [
{ name = "cryptography" },
{ name = "fastapi" },
{ name = "fastembed" },
{ name = "httpx", extra = ["http2"] },
{ name = "hypercorn" },
{ name = "jinja2" },
{ name = "magika" },
{ name = "mcp" },
@ -1928,6 +1938,8 @@ requires-dist = [
{ name = "botocore", extras = ["crt"], marker = "extra == 'bedrock'", specifier = ">=1.41.0" },
{ name = "click", specifier = ">=8.3.3" },
{ name = "crewai", marker = "extra == 'crewai'", specifier = ">=1.0" },
{ name = "cryptography", marker = "extra == 'dev'", specifier = ">=42.0.0" },
{ name = "cryptography", marker = "extra == 'proxy'", specifier = ">=42.0.0" },
{ name = "datasets", marker = "extra == 'evals'", specifier = ">=5.0.1" },
{ name = "datasets", marker = "extra == 'voice-train'", specifier = ">=5.0.1" },
{ name = "fastapi", marker = "extra == 'dev'", specifier = ">=0.100.0" },
@ -1944,6 +1956,8 @@ requires-dist = [
{ name = "httpx", extras = ["http2"], marker = "extra == 'dev'", specifier = ">=0.24.0" },
{ name = "httpx", extras = ["http2"], marker = "extra == 'proxy'", specifier = ">=0.24.0" },
{ name = "huggingface-hub", marker = "extra == 'ml'", specifier = ">=1.5.0,<2.0" },
{ name = "hypercorn", marker = "extra == 'dev'", specifier = ">=0.16" },
{ name = "hypercorn", marker = "extra == 'proxy'", specifier = ">=0.16" },
{ name = "jinja2", marker = "extra == 'reports'", specifier = ">=3.0.0" },
{ name = "langchain-core", marker = "extra == 'langchain'", specifier = ">=1.3.3,<4.0" },
{ name = "langchain-ollama", marker = "extra == 'dev'", specifier = ">=0.2.0" },
@ -2207,13 +2221,32 @@ name = "humanfriendly"
version = "10.0"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "pyreadline3", marker = "python_full_version < '3.11' and sys_platform == 'win32'" },
{ name = "pyreadline3", marker = "sys_platform == 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/cc/3f/2c29224acb2e2df4d2046e4c73ee2662023c58ff5b113c4c1adac0886c43/humanfriendly-10.0.tar.gz", hash = "sha256:6b0b831ce8f15f7300721aa49829fc4e83921a9a301cc7f606be6686a2288ddc", size = 360702, upload-time = "2021-09-17T21:40:43.31Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/f0/0f/310fb31e39e2d734ccaa2c0fb981ee41f7bd5056ce9bc29b2248bd569169/humanfriendly-10.0-py2.py3-none-any.whl", hash = "sha256:1697e1a8a8f550fd43c2865cd84542fc175a61dcb779b6fee18cf6b6ccba1477", size = 86794, upload-time = "2021-09-17T21:40:39.897Z" },
]
[[package]]
name = "hypercorn"
version = "0.18.0"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "exceptiongroup", marker = "python_full_version < '3.11'" },
{ name = "h11" },
{ name = "h2" },
{ name = "priority" },
{ name = "taskgroup", marker = "python_full_version < '3.11'" },
{ name = "tomli", marker = "python_full_version < '3.11'" },
{ name = "typing-extensions", marker = "python_full_version < '3.11'" },
{ name = "wsproto" },
]
sdist = { url = "https://files.pythonhosted.org/packages/44/01/39f41a014b83dd5c795217362f2ca9071cf243e6a75bdcd6cd5b944658cc/hypercorn-0.18.0.tar.gz", hash = "sha256:d63267548939c46b0247dc8e5b45a9947590e35e64ee73a23c074aa3cf88e9da", size = 68420, upload-time = "2025-11-08T13:54:04.78Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/93/35/850277d1b17b206bd10874c8a9a3f52e059452fb49bb0d22cbb908f6038b/hypercorn-0.18.0-py3-none-any.whl", hash = "sha256:225e268f2c1c2f28f6d8f6db8f40cb8c992963610c5725e13ccfcddccb24b1cd", size = 61640, upload-time = "2025-11-08T13:54:03.202Z" },
]
[[package]]
name = "hyperframe"
version = "6.1.0"
@ -2690,18 +2723,18 @@ name = "litellm"
version = "1.88.1"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "aiohttp", marker = "python_full_version < '3.14'" },
{ name = "click", marker = "python_full_version < '3.14'" },
{ name = "fastuuid", marker = "python_full_version < '3.14'" },
{ name = "httpx", marker = "python_full_version < '3.14'" },
{ name = "importlib-metadata", marker = "python_full_version < '3.14'" },
{ name = "jinja2", marker = "python_full_version < '3.14'" },
{ name = "jsonschema", marker = "python_full_version < '3.14'" },
{ name = "openai", marker = "python_full_version < '3.14'" },
{ name = "pydantic", marker = "python_full_version < '3.14'" },
{ name = "python-dotenv", marker = "python_full_version < '3.14'" },
{ name = "tiktoken", marker = "python_full_version < '3.14'" },
{ name = "tokenizers", marker = "python_full_version < '3.14'" },
{ name = "aiohttp" },
{ name = "click" },
{ name = "fastuuid" },
{ name = "httpx" },
{ name = "importlib-metadata" },
{ name = "jinja2" },
{ name = "jsonschema" },
{ name = "openai" },
{ name = "pydantic" },
{ name = "python-dotenv" },
{ name = "tiktoken" },
{ name = "tokenizers" },
]
sdist = { url = "https://files.pythonhosted.org/packages/16/ea/f99ececb7f22703fe120f1d8be9ffb749ec9453fbbbbbebc0d6a6b4d7864/litellm-1.88.1.tar.gz", hash = "sha256:89c6b74cc7912d6365793006ff951c0450fe847625008dfe49de8a7dc4529aa5", size = 13885969, upload-time = "2026-06-09T01:06:25.192Z" }
wheels = [
@ -3583,7 +3616,7 @@ name = "nvidia-cublas"
version = "13.1.1.3"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
{ name = "nvidia-cuda-nvrtc" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" },
@ -3622,7 +3655,7 @@ name = "nvidia-cudnn-cu13"
version = "9.20.0.48"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
{ name = "nvidia-cublas" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" },
@ -3634,7 +3667,7 @@ name = "nvidia-cufft"
version = "12.0.0.61"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
{ name = "nvidia-nvjitlink" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" },
@ -3664,9 +3697,9 @@ name = "nvidia-cusolver"
version = "12.0.4.66"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
{ name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
{ name = "nvidia-cublas" },
{ name = "nvidia-cusparse" },
{ name = "nvidia-nvjitlink" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" },
@ -3678,7 +3711,7 @@ name = "nvidia-cusparse"
version = "12.6.3.3"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" },
{ name = "nvidia-nvjitlink" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" },
@ -3757,8 +3790,8 @@ name = "omegaconf"
version = "2.3.0"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "antlr4-python3-runtime", marker = "python_full_version >= '3.13'" },
{ name = "pyyaml", marker = "python_full_version >= '3.13'" },
{ name = "antlr4-python3-runtime" },
{ name = "pyyaml" },
]
sdist = { url = "https://files.pythonhosted.org/packages/09/48/6388f1bb9da707110532cb70ec4d2822858ddfb44f1cdf1233c20a80ea4b/omegaconf-2.3.0.tar.gz", hash = "sha256:d5d4b6d29955cc50ad50c46dc269bcd92c6e00f5f90d23ab5fee7bfca4ba4cc7", size = 3298120, upload-time = "2022-12-08T20:59:22.753Z" }
wheels = [
@ -3773,12 +3806,12 @@ resolution-markers = [
"python_full_version < '3.11'",
]
dependencies = [
{ name = "coloredlogs", marker = "python_full_version < '3.11'" },
{ name = "flatbuffers", marker = "python_full_version < '3.11'" },
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version < '3.11'" },
{ name = "packaging", marker = "python_full_version < '3.11'" },
{ name = "protobuf", marker = "python_full_version < '3.11'" },
{ name = "sympy", marker = "python_full_version < '3.11'" },
{ name = "coloredlogs" },
{ name = "flatbuffers" },
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple/" } },
{ name = "packaging" },
{ name = "protobuf" },
{ name = "sympy" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/35/d6/311b1afea060015b56c742f3531168c1644650767f27ef40062569960587/onnxruntime-1.23.2-cp310-cp310-macosx_13_0_arm64.whl", hash = "sha256:a7730122afe186a784660f6ec5807138bf9d792fa1df76556b27307ea9ebcbe3", size = 17195934, upload-time = "2025-10-27T23:06:14.143Z" },
@ -3824,10 +3857,10 @@ resolution-markers = [
"python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'",
]
dependencies = [
{ name = "flatbuffers", marker = "python_full_version >= '3.11'" },
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version >= '3.11'" },
{ name = "packaging", marker = "python_full_version >= '3.11'" },
{ name = "protobuf", marker = "python_full_version >= '3.11'" },
{ name = "flatbuffers" },
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple/" } },
{ name = "packaging" },
{ name = "protobuf" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/d4/81/29a9eb470994a75eb7b3ccf32be314d7c66675a00ac7b50294816cc2db27/onnxruntime-1.26.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:ee1109ef4ef27cad90e823399e61e03b3c6c7bfe0fb820b4baf3678c15be8b3c", size = 18005108, upload-time = "2026-05-08T19:08:11.728Z" },
@ -3911,7 +3944,7 @@ name = "openresponses-types"
version = "2.3.0.post1"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "pydantic", marker = "python_full_version >= '3.11'" },
{ name = "pydantic" },
]
sdist = { url = "https://files.pythonhosted.org/packages/d9/26/b612c3215f5599714fa94d63eb5ee59b4eb66dbdeeaf86bb4d848359484d/openresponses_types-2.3.0.post1.tar.gz", hash = "sha256:11b8896d3621d2ac2439f6ff106f34ddcb1bbd517c317a6c852a9df2e98a0753", size = 19254, upload-time = "2026-01-22T20:02:03.933Z" }
wheels = [
@ -4154,10 +4187,10 @@ resolution-markers = [
"python_full_version < '3.11'",
]
dependencies = [
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version < '3.11'" },
{ name = "python-dateutil", marker = "python_full_version < '3.11'" },
{ name = "pytz", marker = "python_full_version < '3.11'" },
{ name = "tzdata", marker = "python_full_version < '3.11'" },
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple/" } },
{ name = "python-dateutil" },
{ name = "pytz" },
{ name = "tzdata" },
]
sdist = { url = "https://files.pythonhosted.org/packages/33/01/d40b85317f86cf08d853a4f495195c73815fdf205eef3993821720274518/pandas-2.3.3.tar.gz", hash = "sha256:e05e1af93b977f7eafa636d043f9f94c7ee3ac81af99c13508215942e64c993b", size = 4495223, upload-time = "2025-09-29T23:34:51.853Z" }
wheels = [
@ -4229,9 +4262,9 @@ resolution-markers = [
"python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'",
]
dependencies = [
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version >= '3.11'" },
{ name = "python-dateutil", marker = "python_full_version >= '3.11'" },
{ name = "tzdata", marker = "(python_full_version >= '3.11' and sys_platform == 'emscripten') or (python_full_version >= '3.11' and sys_platform == 'win32')" },
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple/" } },
{ name = "python-dateutil" },
{ name = "tzdata", marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/de/da/b1dc0481ab8d55d0f46e343cfe67d4551a0e14fcee52bd38ca1bd73258d8/pandas-3.0.0.tar.gz", hash = "sha256:0facf7e87d38f721f0af46fe70d97373a37701b1c09f7ed7aeeb292ade5c050f", size = 4633005, upload-time = "2026-01-21T15:52:04.726Z" }
wheels = [
@ -4476,6 +4509,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/5d/19/fd3ef348460c80af7bb4669ea7926651d1f95c23ff2df18b9d24bab4f3fa/pre_commit-4.5.1-py2.py3-none-any.whl", hash = "sha256:3b3afd891e97337708c1674210f8eba659b52a38ea5f822ff142d10786221f77", size = 226437, upload-time = "2025-12-16T21:14:32.409Z" },
]
[[package]]
name = "priority"
version = "2.0.0"
source = { registry = "https://pypi.org/simple/" }
sdist = { url = "https://files.pythonhosted.org/packages/f5/3c/eb7c35f4dcede96fca1842dac5f4f5d15511aa4b52f3a961219e68ae9204/priority-2.0.0.tar.gz", hash = "sha256:c965d54f1b8d0d0b19479db3924c7c36cf672dbf2aec92d43fbdaf4492ba18c0", size = 24792, upload-time = "2021-06-27T10:15:05.487Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/5e/5f/82c8074f7e84978129347c2c6ec8b6c59f3584ff1a20bc3c940a3e061790/priority-2.0.0-py3-none-any.whl", hash = "sha256:6f8eefce5f3ad59baf2c080a664037bb4725cd0a790d53d59ab4059288faf6aa", size = 8946, upload-time = "2021-06-27T10:15:03.856Z" },
]
[[package]]
name = "propcache"
version = "0.4.1"
@ -5393,17 +5435,17 @@ name = "rapidocr"
version = "3.8.1"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "colorlog", marker = "python_full_version >= '3.13'" },
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version >= '3.13'" },
{ name = "omegaconf", marker = "python_full_version >= '3.13'" },
{ name = "opencv-python", marker = "python_full_version >= '3.13'" },
{ name = "pillow", marker = "python_full_version >= '3.13'" },
{ name = "pyclipper", marker = "python_full_version >= '3.13'" },
{ name = "pyyaml", marker = "python_full_version >= '3.13'" },
{ name = "requests", marker = "python_full_version >= '3.13'" },
{ name = "shapely", marker = "python_full_version >= '3.13'" },
{ name = "six", marker = "python_full_version >= '3.13'" },
{ name = "tqdm", marker = "python_full_version >= '3.13'" },
{ name = "colorlog" },
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple/" } },
{ name = "omegaconf" },
{ name = "opencv-python" },
{ name = "pillow" },
{ name = "pyclipper" },
{ name = "pyyaml" },
{ name = "requests" },
{ name = "shapely" },
{ name = "six" },
{ name = "tqdm" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/ea/4a/fa521d947f0fc7bb304bf11bec4cb66266bd81494588b4cb48dc01001719/rapidocr-3.8.1-py3-none-any.whl", hash = "sha256:650044b1fbce9e6bae5cae462dcf8be754cde11e2f23fc51f65dcc08deae2c46", size = 15080319, upload-time = "2026-04-11T07:13:22.56Z" },
@ -5414,17 +5456,17 @@ name = "rapidocr-onnxruntime"
version = "1.4.4"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version < '3.11'" },
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
{ name = "onnxruntime", version = "1.23.2", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version < '3.11'" },
{ name = "onnxruntime", version = "1.23.2", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" },
{ name = "onnxruntime", version = "1.26.0", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" },
{ name = "opencv-python", marker = "python_full_version < '3.13'" },
{ name = "pillow", marker = "python_full_version < '3.13'" },
{ name = "pyclipper", marker = "python_full_version < '3.13'" },
{ name = "pyyaml", marker = "python_full_version < '3.13'" },
{ name = "shapely", marker = "python_full_version < '3.13'" },
{ name = "six", marker = "python_full_version < '3.13'" },
{ name = "tqdm", marker = "python_full_version < '3.13'" },
{ name = "opencv-python" },
{ name = "pillow" },
{ name = "pyclipper" },
{ name = "pyyaml" },
{ name = "shapely" },
{ name = "six" },
{ name = "tqdm" },
]
wheels = [
{ url = "https://files.pythonhosted.org/packages/ba/12/1e5497183bdbe782dbb91bad1d0d2297dba4d2831b2652657f7517bfc6df/rapidocr_onnxruntime-1.4.4-py3-none-any.whl", hash = "sha256:971d7d5f223a7a808662229df1ef69893809d8457d834e6373d3854bc1782cbf", size = 14915192, upload-time = "2025-01-17T01:48:25.104Z" },
@ -5821,10 +5863,10 @@ resolution-markers = [
"python_full_version < '3.11'",
]
dependencies = [
{ name = "joblib", marker = "python_full_version < '3.11'" },
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version < '3.11'" },
{ name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version < '3.11'" },
{ name = "threadpoolctl", marker = "python_full_version < '3.11'" },
{ name = "joblib" },
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple/" } },
{ name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple/" } },
{ name = "threadpoolctl" },
]
sdist = { url = "https://files.pythonhosted.org/packages/98/c2/a7855e41c9d285dfe86dc50b250978105dce513d6e459ea66a6aeb0e1e0c/scikit_learn-1.7.2.tar.gz", hash = "sha256:20e9e49ecd130598f1ca38a1d85090e1a600147b9c02fa6f15d69cb53d968fda", size = 7193136, upload-time = "2025-09-09T08:21:29.075Z" }
wheels = [
@ -5879,10 +5921,10 @@ resolution-markers = [
"python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'",
]
dependencies = [
{ name = "joblib", marker = "python_full_version >= '3.11'" },
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version >= '3.11'" },
{ name = "scipy", version = "1.17.0", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version >= '3.11'" },
{ name = "threadpoolctl", marker = "python_full_version >= '3.11'" },
{ name = "joblib" },
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple/" } },
{ name = "scipy", version = "1.17.0", source = { registry = "https://pypi.org/simple/" } },
{ name = "threadpoolctl" },
]
sdist = { url = "https://files.pythonhosted.org/packages/0e/d4/40988bf3b8e34feec1d0e6a051446b1f66225f8529b9309becaeef62b6c4/scikit_learn-1.8.0.tar.gz", hash = "sha256:9bccbb3b40e3de10351f8f5068e105d0f4083b1a65fa07b6634fbc401a6287fd", size = 7335585, upload-time = "2025-12-10T07:08:53.618Z" }
wheels = [
@ -5932,7 +5974,7 @@ resolution-markers = [
"python_full_version < '3.11'",
]
dependencies = [
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version < '3.11'" },
{ name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple/" } },
]
sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" }
wheels = [
@ -6002,7 +6044,7 @@ resolution-markers = [
"python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'",
]
dependencies = [
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple/" }, marker = "python_full_version >= '3.11'" },
{ name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple/" } },
]
sdist = { url = "https://files.pythonhosted.org/packages/56/3e/9cca699f3486ce6bc12ff46dc2031f1ec8eb9ccc9a320fdaf925f1417426/scipy-1.17.0.tar.gz", hash = "sha256:2591060c8e648d8b96439e111ac41fd8342fdeff1876be2e19dea3fe8930454e", size = 30396830, upload-time = "2026-01-10T21:34:23.009Z" }
wheels = [
@ -6399,6 +6441,19 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" },
]
[[package]]
name = "taskgroup"
version = "0.2.2"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "exceptiongroup" },
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/f0/8d/e218e0160cc1b692e6e0e5ba34e8865dbb171efeb5fc9a704544b3020605/taskgroup-0.2.2.tar.gz", hash = "sha256:078483ac3e78f2e3f973e2edbf6941374fbea81b9c5d0a96f51d297717f4752d", size = 11504, upload-time = "2025-01-03T09:24:13.761Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/d1/b1/74babcc824a57904e919f3af16d86c08b524c0691504baf038ef2d7f655c/taskgroup-0.2.2-py2.py3-none-any.whl", hash = "sha256:e2c53121609f4ae97303e9ea1524304b4de6faf9eb2c9280c7f87976479a52fb", size = 14237, upload-time = "2025-01-03T09:24:11.41Z" },
]
[[package]]
name = "tenacity"
version = "9.1.2"
@ -7333,6 +7388,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/1f/f6/a933bd70f98e9cf3e08167fc5cd7aaaca49147e48411c0bd5ae701bb2194/wrapt-1.17.3-py3-none-any.whl", hash = "sha256:7171ae35d2c33d326ac19dd8facb1e82e5fd04ef8c6c0e394d7af55a55051c22", size = 23591, upload-time = "2025-08-12T05:53:20.674Z" },
]
[[package]]
name = "wsproto"
version = "1.3.2"
source = { registry = "https://pypi.org/simple/" }
dependencies = [
{ name = "h11" },
]
sdist = { url = "https://files.pythonhosted.org/packages/c7/79/12135bdf8b9c9367b8701c2c19a14c913c120b882d50b014ca0d38083c2c/wsproto-1.3.2.tar.gz", hash = "sha256:b86885dcf294e15204919950f666e06ffc6c7c114ca900b060d6e16293528294", size = 50116, upload-time = "2025-11-20T18:18:01.871Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/a4/f5/10b68b7b1544245097b2a1b8238f66f2fc6dcaeb24ba5d917f52bd2eed4f/wsproto-1.3.2-py3-none-any.whl", hash = "sha256:61eea322cdf56e8cc904bd3ad7573359a242ba65688716b0710a5eb12beab584", size = 24405, upload-time = "2025-11-20T18:18:00.454Z" },
]
[[package]]
name = "xlrd"
version = "2.0.2"