diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 431860278..16c481585 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -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' diff --git a/.gitignore b/.gitignore index 8a477f307..ace6f561b 100644 --- a/.gitignore +++ b/.gitignore @@ -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 diff --git a/README.md b/README.md index 0b923e940..2ed29d116 100644 --- a/README.md +++ b/README.md @@ -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 `` / +``; 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…** diff --git a/crates/headroom-core/tests/auth_mode.rs b/crates/headroom-core/tests/auth_mode.rs index f0de13301..e0bdce887 100644 --- a/crates/headroom-core/tests/auth_mode.rs +++ b/crates/headroom-core/tests/auth_mode.rs @@ -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/`. + // 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 diff --git a/docs/adr/0001-agy-mitm-transport.md b/docs/adr/0001-agy-mitm-transport.md new file mode 100644 index 000000000..7e6faea8d --- /dev/null +++ b/docs/adr/0001-agy-mitm-transport.md @@ -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:` 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:` 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. diff --git a/docs/agy-parity-matrix.md b/docs/agy-parity-matrix.md new file mode 100644 index 000000000..641cb6575 --- /dev/null +++ b/docs/agy-parity-matrix.md @@ -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 `` 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. | diff --git a/docs/content/docs/proxy.mdx b/docs/content/docs/proxy.mdx index 3390a086d..936cb5e2b 100644 --- a/docs/content/docs/proxy.mdx +++ b/docs/content/docs/proxy.mdx @@ -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` diff --git a/headroom/cache/compression_store.py b/headroom/cache/compression_store.py index bbda697cd..5844827ab 100644 --- a/headroom/cache/compression_store.py +++ b/headroom/cache/compression_store.py @@ -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 diff --git a/headroom/ccr/mcp_server.py b/headroom/ccr/mcp_server.py index 8294c0e21..7e74085b3 100644 --- a/headroom/ccr/mcp_server.py +++ b/headroom/ccr/mcp_server.py @@ -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, diff --git a/headroom/ccr/tool_injection.py b/headroom/ccr/tool_injection.py index 86f6cac3f..6a90ae3c0 100644 --- a/headroom/ccr/tool_injection.py +++ b/headroom/ccr/tool_injection.py @@ -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. diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 966d150ac..6a42b2dc4 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -19,6 +19,7 @@ Usage: from __future__ import annotations +import asyncio import errno import importlib.util import io @@ -30,6 +31,8 @@ import signal import socket import subprocess import sys +import tempfile +import threading import time import urllib.parse from collections.abc import Callable, Mapping @@ -38,7 +41,7 @@ from functools import wraps from pathlib import Path from typing import Any, NamedTuple, cast -from headroom._subprocess import pid_alive, run +from headroom._subprocess import Popen, pid_alive, run # Fix Windows cp1252 encoding — box-drawing characters require UTF-8 if sys.platform == "win32" and hasattr(sys.stdout, "buffer"): @@ -1024,6 +1027,379 @@ _code_memory_option = click.option( # hooks are removed separately, by # headroom.context_tool_cleanup.purge_context_tool_artifacts.) _HEADROOM_HOOK_MARKERS = ("headroom-init-claude",) +#: agy flags that put it into single-shot, non-interactive output mode. +#: In this mode agy hangs indefinitely whenever ANY mcpServers entry is present +#: in its MCP config (post-migration ~/.gemini/config/mcp_config.json; verified +#: live on the legacy path: every server tried — serena, and even a nonexistent +#: command — hangs; empty mcpServers answers in seconds). So Headroom must NOT +#: activate any MCP server for print-mode invocations. +_AGY_PRINT_FLAGS = ("--print", "-p", "--prompt") + +#: Minimum agy version known to no longer hang on a registered MCP server in +#: print mode (re-verified 2026-07-05: serena and the headroom retrieve server +#: both answer in ~4s on 1.0.16). Older or unparseable/unknown versions are +#: treated as unsafe — see ``_agy_print_mode_mcp_allowed``. +_AGY_PRINT_MODE_MCP_MIN_VERSION = (1, 0, 16) + + +_PROXY_URL_REDACTED_PLACEHOLDER = "" + + +def redact_proxy_url(url: str) -> str: + """Render ``scheme://host:port`` for a corporate proxy URL, dropping userinfo. + + Invariant: userinfo (``user:pass@``) is never rendered, and no failure + path echoes the raw input — any parse error, a ``ValueError`` from an + invalid ``.port``, or a falsy ``.hostname`` (e.g. a schemeless URL where + urlparse puts the credentials in ``.path``/``.scheme`` instead) returns a + fixed placeholder. Non-printable characters are stripped from the + rendered scheme/host so control bytes (e.g. ESC) can never reach output. + """ + try: + parsed = urllib.parse.urlparse(url) + host = parsed.hostname + if not host: + return _PROXY_URL_REDACTED_PLACEHOLDER + port = parsed.port or (443 if parsed.scheme == "https" else 80) + except ValueError: + return _PROXY_URL_REDACTED_PLACEHOLDER + + scheme = "".join(ch for ch in parsed.scheme if ch.isprintable()) + host = "".join(ch for ch in host if ch.isprintable()) + if ":" in host: + host = f"[{host}]" + return f"{scheme}://{host}:{port}" + + +def _agy_print_mode(agy_args: tuple[str, ...] | list[str]) -> bool: + """Return True if agy is being launched in non-interactive print mode. + + agy treats ``--print`` / ``-p`` / ``--prompt`` as "run one prompt and exit". + Older/unknown agy versions hang in this mode whenever a registered MCP + server is present, so callers use this to gate MCP wiring behind an agy + version preflight (see ``_agy_print_mode_mcp_allowed``); interactive mode + is never suppressed. + + Matches both space-separated forms (``--print hi``) and ``=``-joined forms + (``--print=hi``, ``--prompt=hi``, ``-p=hi``) — all live-verified as valid + agy print invocations (2026-06-16). The attached short form ``-pVALUE`` is + *not* matched because agy rejects it (``flags provided but not defined``, + exit 2) — it never reaches MCP init, so it cannot trigger the hang. + """ + return any(arg.split("=", 1)[0] in _AGY_PRINT_FLAGS for arg in agy_args) + + +def _detect_agy_version(agy_bin: str | None) -> tuple[int, ...] | None: + """Best-effort detect the installed agy binary's version. + + Runs `` --version`` with a 1s timeout and parses the LAST + ``\\d+(\\.\\d+)+`` token found in stdout (defends against a wrapper + script printing its own version banner before delegating to the real + binary). Returns ``None`` whenever the version cannot be established — + no binary, non-zero exit, no parseable version token, or a slow/hung + process (killed after the timeout) — so callers can treat "unknown" the + same as "known old" (safe-by-default). Never raises. + """ + try: + if not agy_bin: + return None + result = run( + [agy_bin, "--version"], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + timeout=1.0, + ) + if result.returncode != 0: + return None + matches = re.findall(r"\d+(?:\.\d+)+", result.stdout or "") + if not matches: + return None + return tuple(int(part) for part in matches[-1].split(".")) + except Exception: + return None + + +def _agy_print_mode_mcp_allowed(agy_args: tuple[str, ...] | list[str], agy_bin: str | None) -> bool: + """Gate print-mode MCP wiring on a known-good agy version. + + Interactive mode is always allowed — the hang is print-mode-only, so no + version check is performed. In print mode, MCP wiring is allowed only + when the detected agy version is >= ``_AGY_PRINT_MODE_MCP_MIN_VERSION``; + an unparseable/unknown version is treated as too old (safe-by-default). + """ + if not _agy_print_mode(agy_args): + return True + version = _detect_agy_version(agy_bin) + return version is not None and version >= _AGY_PRINT_MODE_MCP_MIN_VERSION + + +def _smoke_verify_mcp_handshake( + command: str, args: list[str], env: dict[str, str], *, timeout: float = 8.0 +) -> bool: + """Spawn an stdio MCP server and assert it answers an ``initialize`` request. + + Sends a minimal JSON-RPC ``initialize`` over stdin and waits up to + ``timeout`` seconds for a JSON-RPC response on stdout. Returns True iff a + well-formed response object (``"jsonrpc"`` + matching ``"id"``) is seen. + The process is always terminated before returning. This is a *guard*: a + passing handshake does not by itself prove agy will be happy (an unrelated + agy print-mode bug hangs on any MCP), but a *failing* handshake proves the + entry is broken and must not be persisted. + """ + request = ( + json.dumps( + { + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "headroom-smoke", "version": "1"}, + }, + } + ) + + "\n" + ) + full_env = {**os.environ, **env} + proc: subprocess.Popen[str] | None = None + try: + proc = Popen( + [command, *args], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + env=full_env, + ) + try: + stdout, _ = proc.communicate(input=request, timeout=timeout) + except subprocess.TimeoutExpired: + return False + for line in (stdout or "").splitlines(): + line = line.strip() + if not line: + continue + try: + payload = json.loads(line) + except json.JSONDecodeError: + continue + if ( + isinstance(payload, dict) + and payload.get("jsonrpc") == "2.0" + and payload.get("id") == 1 + ): + return True + return False + except (OSError, ValueError): + return False + finally: + if proc is not None and proc.poll() is None: + proc.kill() + try: + proc.wait(timeout=2) + except subprocess.TimeoutExpired: + pass + + +def _setup_headroom_retrieve_mcp_agy(registrar: Any, *, verbose: bool = False) -> bool: + """Register the headroom retrieve MCP with agy PERSISTENTLY (mirrors CBM). + + The retrieve tool is an ``headroom mcp serve`` stdio child that resolves + ``[Retrieve more: hash=…]`` markers. It resolves them from the shared + on-disk CCR store (``ccr_store.db``) FIRST — see + ``ccr.mcp_server._retrieve_content`` — so it needs no live proxy and no + per-run ephemeral port; the spec is stable and port-independent + (``build_headroom_spec()`` with the default URL yields ``env={}``). + + agy only surfaces tools from servers in its persistent per-tool cache, so + the entry is registered persistently and RECORDED in the install ledger + (like codebase-memory-mcp / Serena), NOT reverted on teardown. That is what + lets agy discover, cache, and expose ``headroom_retrieve`` across sessions — + the exposure the h76.5 gate then checks before keeping ccr compression on. + + Returns True iff the entry is registered AND survives the smoke handshake. + """ + from headroom.mcp_registry import build_headroom_spec + from headroom.mcp_registry.base import RegisterStatus + from headroom.mcp_registry.ledger import clear_install, record_install + + spec = build_headroom_spec() + result = registrar.register_server(spec, force=True) + if result.status not in (RegisterStatus.REGISTERED, RegisterStatus.ALREADY): + click.echo( + f" MCP retrieve tool: could not register headroom MCP — skipping ({result.detail})." + ) + return False + + if _smoke_verify_mcp_handshake(spec.command, list(spec.args), dict(spec.env)): + # Record on BOTH REGISTERED and ALREADY: a matching on-disk entry whose + # ledger record was lost (e.g. cleared by the old-agy print-mode purge) + # must be re-claimed as Headroom-owned so ledger-gated uninstall works. + # record_install upserts on spec.name, so this never double-counts. + record_install(registrar.name, spec) + _prime_agy_retrieve_tool_cache(registrar) + if verbose: + click.echo( + " MCP retrieve tool: headroom MCP registered persistently " + "(local-store resolution) and handshake-verified." + ) + else: + click.echo(" MCP retrieve tool: headroom MCP wired (persistent, handshake verified).") + return True + + # Handshake failed: remove the entry AND clear any ledger record so a broken + # pointer can never persist or masquerade as Headroom-owned. + registrar.unregister_server("headroom") + clear_install(registrar.name, "headroom") + click.echo( + " MCP retrieve tool: headroom MCP failed handshake — entry removed (agy left transport-only)." + ) + return False + + +def _prime_agy_retrieve_tool_cache(registrar: Any) -> None: + """Pre-write agy's per-tool cache for ``headroom_retrieve`` on first run. + + agy only surfaces a server's tools from its persistent per-tool cache + (``//.json``), which it otherwise writes only + *during* a session. So on a clean install the h76.5 exposure gate + (``_agy_exposes_retrieve_tool``) sees no cache file, withholds ``WIRED``, + and ccr downgrades to lossless for that first run — the "launch, exit, + re-launch" tax. Seeding the file here (verified: agy reads it at startup and + exposes the tool immediately) removes it. The schema mirrors the live + ``list_tools()`` entry via the shared ``CCR_RETRIEVE_TOOL_*`` constants, so + the primed cache cannot drift from what the server actually offers; ``agy`` + serialises MCP ``inputSchema`` under the ``parameters`` key. Existing caches + are left untouched, and any write error degrades to the pre-fix behavior. + """ + from headroom.ccr.mcp_server import ( + CCR_RETRIEVE_TOOL_DESCRIPTION, + CCR_RETRIEVE_TOOL_INPUT_SCHEMA, + CCR_TOOL_NAME, + ) + + tool_cache = registrar.cache_dir / "headroom" / f"{CCR_TOOL_NAME}.json" + if tool_cache.is_file(): + return + try: + tool_cache.parent.mkdir(parents=True, exist_ok=True) + tool_cache.write_text( + json.dumps( + { + "name": CCR_TOOL_NAME, + "description": CCR_RETRIEVE_TOOL_DESCRIPTION, + "parameters": CCR_RETRIEVE_TOOL_INPUT_SCHEMA, + } + ), + encoding="utf-8", + ) + except OSError: + # A cache-write failure must never abort setup; the exposure gate simply + # falls back to the pre-fix downgrade-to-lossless for the first run. + pass + + +def _ccr_backend_is_cross_process() -> bool: + """True unless the CCR store backend is process-local (``memory``). + + The agy-spawned ``headroom mcp serve`` child resolves markers against the + CCR store. With ``HEADROOM_CCR_BACKEND=memory`` that store is a per-process + dict (compression_store.py ``_create_default_ccr_backend``), so the child + sees an empty store and cannot resolve the proxy's hashes — markers would + ship unrecoverable. Every other backend (default sqlite, redis, custom + entry points) is shared across processes. This is a PRODUCT guard on + ``WIRED``, not a test-only check. + """ + return (os.environ.get("HEADROOM_CCR_BACKEND") or "").strip().lower() != "memory" + + +def _agy_exposes_retrieve_tool(registrar: Any) -> bool: + """True iff agy will actually expose ``headroom_retrieve`` as a callable tool. + + A successful wrap↔child ``initialize`` handshake is necessary but NOT + sufficient: agy only surfaces tools from servers in its persistent per-tool + cache (``/mcp//.json``, written *during* a session), + so an entry that is registered-then-reverted every run never enters that + cache and agy rejects the call with "Unknown tool: headroom_retrieve". + + Positive exposure requires ALL of: + 1. a live ``headroom`` entry in ``mcp_config.json`` (registrar.get_server), + 2. an agy-written tool-cache file for ``headroom_retrieve``, and + 3. a cross-process CCR backend (so the child can resolve hashes). + + Anything else = UNVERIFIED → caller withholds ``WIRED`` → ccr downgrades to + lossless (fail-safe; never ships unrecoverable compression). The cache is a + *previous* session's artifact, so pairing it with the live config entry + avoids a stale-cache false positive. + """ + from headroom.ccr.mcp_server import CCR_TOOL_NAME + + if registrar.get_server("headroom") is None: + return False + # Cache lives under agy's app-data dir (cache_dir), NOT the migrated config + # dir — agy writes /mcp//.json regardless of which + # config file declared the server. + tool_cache = registrar.cache_dir / "headroom" / f"{CCR_TOOL_NAME}.json" + if not tool_cache.is_file(): + return False + return _ccr_backend_is_cross_process() + + +def _maybe_warn_agy_ccr_downgrade(retrieve_wired: bool) -> None: + """Loudly warn when ccr mode silently downgraded to lossless this run. + + Fires iff ``headroom.proxy.handlers.gemini._resolve_agy_fr_mode`` would + downgrade: both read the requested mode from the shared + ``_requested_agy_fr_mode`` helper (single source of truth). ccr is the + default mode, and it is the only mode that ships recoverable + functionResponse compression, so it requires the retrieve MCP + to resolve ``[Retrieve more: hash=…]`` markers. When the retrieve MCP did not wire + for this run, that handler falls back to ``lossless`` -- a byte-recoverable + no-op -- so tool-output savings collapse to ~0 with no other signal to the + user. Stays silent when retrieve DID wire, or when ``lossless`` was + requested explicitly (no downgrade occurred). + + Cause detection is ADVISORY only: this probes ``mcp`` importability in + THIS (parent) interpreter, but the agy child is launched via + ``resolve_headroom_command()`` (``shutil.which("headroom")``), which need + not share this venv. A false negative here (mcp present in the parent, + absent in the child) still degrades gracefully to the generic + handshake-failure branch. + + ``retrieve_wired`` is the EXPOSURE-gated signal (handshake AND agy actually + caching the tool), not the bare handshake result -- so this also fires when + the child registers/handshakes fine but agy has not yet exposed the tool. + """ + from headroom.proxy.handlers.gemini import _requested_agy_fr_mode + + if _requested_agy_fr_mode() != "ccr" or retrieve_wired: + return + + if _module_available("mcp"): + cause = ( + "the retrieve MCP did not register/handshake, or agy has not yet " + "exposed it as a callable tool (see the 'MCP retrieve tool:' line " + "above)" + ) + remedy = "Fix the failure shown on that line, then re-run `headroom wrap agy`." + else: + cause = ( + "mcp is not importable in this interpreter (ADVISORY: likely cause -- " + "the agy child is resolved via `headroom` on PATH and may run in a " + "different environment than this one)" + ) + remedy = "Install with: pip install 'headroom-ai[proxy]' (or: pip install mcp)" + + click.echo() + click.echo(" ⚠️ WARNING: agy compression savings are DISABLED this run.") + click.echo(" ⚠️ ccr mode requires the retrieve MCP; it did not wire, so") + click.echo(" ⚠️ functionResponse compression fell back to lossless (saves ~0 on tool output).") + click.echo(f" ⚠️ Cause: {cause}.") + click.echo(f" ⚠️ Fix: {remedy}") + click.echo() + # Env vars Headroom's init/wrap inject into Claude settings.json; unwrap removes # them. ENABLE_TOOL_SEARCH keeps Claude Code's tool deferral on behind the proxy @@ -2070,8 +2446,8 @@ def _serena_instruction_file(registrar: Any) -> Path: def _inject_serena_instructions(file_path: Path, verbose: bool = False) -> bool: """Steer the agent toward Serena's symbol tools over whole-file reads. - Opt-in (off by default): mirrors :func:`_inject_rtk_instructions` and - early-returns unless ``--serena-instructions`` / ``HEADROOM_SERENA_INSTRUCTIONS`` + Opt-in (off by default): early-returns unless + ``--serena-instructions`` / ``HEADROOM_SERENA_INSTRUCTIONS`` is set, so the user's hint file is left untouched by default. Idempotent — skips if the marker is already present. Appends to an existing @@ -2408,6 +2784,24 @@ def _remove_headroom_installed_serena_mcp(registrar: Any) -> str: return "failed" +def _remove_headroom_installed_retrieve_mcp(registrar: Any) -> str: + """Remove the headroom retrieve MCP only if the ledger proves Headroom installed it. + + Mirrors ``_remove_headroom_installed_serena_mcp``: the retrieve entry is now a + persistent, ledger-recorded server, so cooperative uninstall is ledger-gated + (never clobber a user's own "headroom" entry) and clears the ledger record. + """ + from headroom.mcp_registry.ledger import clear_install, headroom_installed_matching + + current = registrar.get_server("headroom") + if not headroom_installed_matching(registrar.name, current): + return "not_headroom_owned" + if registrar.unregister_server("headroom"): + clear_install(registrar.name, "headroom") + return "removed" + return "failed" + + def _disable_serena_mcp( registrar: Any, *, verbose: bool = False, reason: str = "--no-serena" ) -> None: @@ -2531,11 +2925,37 @@ def _setup_coding_compressor(registrar: Any, *, serena_context: str, **kwargs: A _CBM_MCP_SERVER_NAME = "codebase-memory-mcp" +def _purge_agy_mcp_entries(registrar: Any) -> None: + """Actively remove all agy MCP entries that could hang a print-mode run. + + Used when the print-mode agy-version preflight fails (older or unknown + agy): merely skipping *new* registration is not enough, since a prior + interactive ``headroom wrap agy`` run may have already persisted entries + in mcp_config.json. Every call here is idempotent -- a no-op when the + entry is already absent -- so this is safe to call unconditionally. + """ + from headroom.mcp_registry.ledger import clear_install + + _disable_tokensave_mcp(registrar) + _disable_serena_mcp(registrar, reason="agy print-mode MCP preflight failed") + registrar.unregister_server(_CBM_MCP_SERVER_NAME) + # headroom retrieve is now a ledger-recorded PERSISTENT entry; old agy hangs + # in print mode on ANY MCP entry, so purge it AND clear its ledger record so + # the persistent-skip on the next compatible-agy run does not treat the now + # absent entry as still-installed. Re-registration happens on that next wrap. + registrar.unregister_server("headroom") + clear_install(registrar.name, "headroom") + + # Memory MCP markers _MEMORY_MCP_MARKER = "# --- Headroom memory MCP (auto-injected) ---" _MEMORY_MCP_END = "# --- end Headroom memory ---" _MEMORY_AGENTS_MARKER = "" +# agy / GEMINI.md instruction-block markers +_AGY_GEMINI_BLOCK_START = "" +_AGY_GEMINI_BLOCK_END = "" + # Codex config injection markers _CODEX_TOP_LEVEL_MARKER = "# --- Headroom proxy (auto-injected by headroom wrap codex) ---" _CODEX_END_MARKER = "# --- end Headroom ---" @@ -3369,6 +3789,35 @@ def _run_proxy_only_watcher( cleanup() +def _remove_gemini_md_block(gemini_md: Path, verbose: bool = False) -> bool: + """Remove the Headroom-marked block from GEMINI.md (idempotent). + + Only removes the delimited block; all user content outside the markers is + preserved. Returns ``True`` if a block was found and removed. + """ + if not gemini_md.exists(): + return False + existing = gemini_md.read_text(encoding="utf-8") + if _AGY_GEMINI_BLOCK_START not in existing or _AGY_GEMINI_BLOCK_END not in existing: + return False + start = existing.index(_AGY_GEMINI_BLOCK_START) + end = existing.index(_AGY_GEMINI_BLOCK_END) + len(_AGY_GEMINI_BLOCK_END) + before = existing[:start].rstrip("\n") + after = existing[end:].lstrip("\n") + if before and after: + new_text = before + "\n\n" + after + elif before: + new_text = before + "\n" + elif after: + new_text = after + else: + new_text = "" + gemini_md.write_text(new_text, encoding="utf-8") + if verbose: + click.echo(f" headroom block removed from {gemini_md}") + return True + + def _inject_memory_mcp_config(user_id: str) -> None: """Register headroom memory as an MCP server in Codex's config.toml. @@ -8154,6 +8603,714 @@ def unwrap_codex(port: int, no_stop_proxy: bool) -> None: click.echo() +# ============================================================================= +# agy MITM lifecycle helpers +# ============================================================================= + + +class _AgyServers: + """Handle to the running terminator + dispatch pair. + + Holds the async event-loop thread and exposes a synchronous ``stop()`` + that schedules cleanup on that loop and joins the thread. + """ + + def __init__( + self, + terminator: Any, + dispatch: Any, + loop: asyncio.AbstractEventLoop, + thread: threading.Thread, + stop_flag: asyncio.Event, + retrieve: Any | None = None, + retrieve_port: int | None = None, + ) -> None: + self.terminator = terminator + self.dispatch = dispatch + # Plain-HTTP loopback retrieve listener (interactive mode only). ``None`` + # in print mode, where no MCP server may run (agy hangs otherwise). + self.retrieve = retrieve + self.retrieve_port = retrieve_port + self._loop = loop + self._thread = thread + self._stop_flag = stop_flag + self._lock = threading.Lock() + self._stopped = False + + def stop(self) -> None: + """Best-effort graceful shutdown (idempotent).""" + with self._lock: + if self._stopped: + return + self._stopped = True + # Wake the event loop so it can stop() the servers and exit. + self._loop.call_soon_threadsafe(self._stop_flag.set) + self._thread.join(timeout=10) + + +def _start_agy_servers( + ca_key: Any, + ca_cert: Any, + base_dir: Path | None = None, + *, + start_retrieve: bool = False, + project: str | None = None, +) -> _AgyServers: + """Start AgyCONNECTTerminator + AgyDispatchServer on a dedicated thread. + + Both servers bind loopback ephemeral ports (port=0). Readiness is + signalled via a threading.Event; startup errors raise RuntimeError fast. + + When ``start_retrieve`` is True an additional PLAIN-HTTP loopback + :class:`AgyRetrieveServer` is started on the same loop; its port is exposed + via ``.retrieve_port`` so the headroom retrieve MCP can point at it. The + caller now passes ``start_retrieve=True`` in every mode: the listener is a + harmless idle loopback socket, and whether the retrieve MCP *entry* is + registered is decided separately by the print-mode version gate + (``_agy_print_mode_mcp_allowed``, headroom-37g.37). The retrieve server + shares the process-global compression cache the dispatch server populates, + so ``[Retrieve more: hash=…]`` markers resolve. + + Returns an _AgyServers handle with ``.terminator`` and ``.dispatch`` + already started, and a ``.stop()`` method for clean shutdown. + """ + from headroom.proxy.agy_dispatch import AgyDispatchServer + from headroom.proxy.agy_retrieve import AgyRetrieveServer + from headroom.proxy.agy_terminator import DEFAULT_ALLOWLIST, AgyCONNECTTerminator + + allowlist = DEFAULT_ALLOWLIST + + ready_event: threading.Event = threading.Event() + error_holder: list[Exception] = [] + result_holder: list[_AgyServers] = [] + + def _run_loop() -> None: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + stop_flag = asyncio.Event() + + async def _main() -> None: + dispatch = AgyDispatchServer( + ca_key=ca_key, + ca_cert=ca_cert, + base_dir=base_dir, + port=0, + allowlist=allowlist, + project=project, + ) + await dispatch.start() + _, dispatch_port = dispatch.address + + terminator = AgyCONNECTTerminator( + ca_key=ca_key, + ca_cert=ca_cert, + base_dir=base_dir, + port=0, + dispatch_port=dispatch_port, + allowlist=allowlist, + ) + await terminator.start() + + retrieve: AgyRetrieveServer | None = None + retrieve_port: int | None = None + if start_retrieve: + retrieve = AgyRetrieveServer(port=0) + await retrieve.start() + _, retrieve_port = retrieve.address + + servers = _AgyServers( + terminator=terminator, + dispatch=dispatch, + loop=loop, + thread=current_thread, + stop_flag=stop_flag, + retrieve=retrieve, + retrieve_port=retrieve_port, + ) + result_holder.append(servers) + ready_event.set() + + # Keep event loop alive until stop_flag is set. + await stop_flag.wait() + + # Graceful shutdown. + await terminator.stop() + await dispatch.stop() + if retrieve is not None: + await retrieve.stop() + + try: + loop.run_until_complete(_main()) + except Exception as exc: # noqa: BLE001 + error_holder.append(exc) + ready_event.set() + finally: + # Cancel and drain any stragglers (hypercorn per-connection tasks, + # the app lifespan's periodic stats task) before closing the loop — + # otherwise loop.close() with pending tasks spews "Task was destroyed + # but it is pending" / "Event loop is closed" on every agy exit. + try: + pending = [t for t in asyncio.all_tasks(loop) if not t.done()] + for task in pending: + task.cancel() + if pending: + loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + except Exception: # noqa: BLE001 + pass + loop.close() + + current_thread = threading.Thread(target=_run_loop, daemon=True, name="headroom-agy-mitm") + current_thread.start() + ready_event.wait(timeout=15) + + if error_holder: + raise RuntimeError(f"agy MITM server startup failed: {error_holder[0]}") from error_holder[ + 0 + ] + if not result_holder: + raise RuntimeError("agy MITM servers did not start within 15 seconds") + + return result_holder[0] + + +def _stop_agy_servers(servers: _AgyServers | None) -> None: + """Best-effort stop of agy servers (called from finally block). + + Accepts the _AgyServers handle returned by _start_agy_servers. + Idempotent and None-safe so it can run from both the normal exit path + and the SIGTERM handler without double-teardown errors. + """ + if servers is None: + return + try: + servers.stop() + except Exception: # noqa: BLE001 + pass + + +# ============================================================================= +# wrap agy +# ============================================================================= + + +@wrap.command(context_settings={"ignore_unknown_options": True}) +@click.option( + # NOTE: no "-p" short alias here (unlike sibling wrap subcommands): agy's + # own CLI uses -p for --print, so a -p alias on --port would swallow the + # user's prompt as the proxy port (headroom-r9k). Long --port only. + "--port", + default=8787, + type=click.IntRange(1, 65535), + help="Proxy port (default: 8787)", +) +@click.option( + "--no-intercept", + is_flag=True, + help=( + "Passthrough / escape hatch: launch agy unchanged, with no TLS interception. " + "agy traffic is NOT compressed or inspected by Headroom. Use this to verify " + "issues are caused by the MITM transport, or to opt out entirely. " + "Run 'headroom unwrap agy' to revert any persistent changes." + ), +) +@click.option( + "--backend", + default=None, + help="API backend for the proxy (env: HEADROOM_BACKEND). NOTE: only Python backend is supported for agy.", +) +@click.option("--no-mcp", is_flag=True, help="Skip headroom MCP server registration") +@click.option("--no-serena", is_flag=True, help="Skip Serena MCP server registration") +@click.option( + "--no-tokensave", + is_flag=True, + hidden=True, + help="Deprecated and ignored: tokensave was retired; Serena is the default code memory.", +) +@click.option( + "--code-graph", + is_flag=True, + default=False, + help="Enable code graph indexing via codebase-memory-mcp (optional)", +) +@click.option("--no-proxy", is_flag=True, help="Skip proxy startup (use existing proxy)") +@click.argument("agy_args", nargs=-1, type=click.UNPROCESSED) +def agy( + port: int, + no_intercept: bool, + backend: str | None, + no_mcp: bool, + no_serena: bool, + no_tokensave: bool, + code_graph: bool, + no_proxy: bool, + agy_args: tuple, +) -> None: + """Launch agy through Headroom's selective TLS-MITM transport. + + \b + agy has no base-URL override knob, so Headroom intercepts its traffic via + an in-process HTTP CONNECT terminator that TLS-terminates only the Cloud + Code backend hosts in the terminator allowlist (daily-cloudcode-pa and + cloudcode-pa googleapis.com). All other connections are byte-spliced + unchanged (and chained through any pre-existing corporate HTTPS_PROXY). + + \b + The process-local CA (headroom.proxy.agy_ca) is used to mint leaf + certificates for the intercepted host. It is NEVER added to the OS trust + store; it lives only in the child process environment. + + \b + Use --no-intercept to launch agy with no interception (passthrough mode). + Run 'headroom unwrap agy' to undo any persistent configuration changes. + + \b + Examples: + headroom wrap agy # Start with MITM transport + headroom wrap agy -- --help # Pass args to agy + headroom wrap agy --no-intercept # Passthrough / escape hatch + """ + from headroom.mcp_registry.agy import AgyRegistrar + + # Resolve binary first — fast exit if not installed. + agy_bin = shutil.which("agy") + if not agy_bin: + raise click.ClickException( + "'agy' not found in PATH. " + "Install agy: https://github.com/google/agy (or via your package manager)" + ) + + # Rust backend is Python-only for agy (T11 deferred). + effective_backend = backend or os.environ.get("HEADROOM_BACKEND") + if effective_backend == "rust": + click.echo( + "Error: agy MITM transport is Python-only. " + "Rust backend support is deferred (T11). " + "Use the Python backend (omit --backend rust / unset HEADROOM_BACKEND)." + ) + raise SystemExit(1) + + if no_intercept: + # Passthrough: launch agy unchanged, zero modification to its env. + click.echo() + click.echo(" ╔═══════════════════════════════════════════════╗") + click.echo(" ║ HEADROOM WRAP: AGY ║") + click.echo(" ╚═══════════════════════════════════════════════╝") + click.echo() + click.echo( + " Mode: --no-intercept (passthrough). Headroom does NOT intercept agy traffic." + ) + click.echo() + result = subprocess.run([agy_bin, *agy_args]) + raise SystemExit(result.returncode) + + # ----------------------------------------------------------------------- + # MITM path + # ----------------------------------------------------------------------- + # Quiet litellm's "Provider List: https://..." banner, which it prints to + # stderr on every cost lookup for models it doesn't know (agy's Cloud Code + # model ids). Set the global flags once, before the dispatch handles any + # request. Assign through an Any alias (the flags aren't in litellm's stubs). + try: + import litellm + + _litellm: Any = litellm + _litellm.suppress_debug_info = True + _litellm.set_verbose = False + except Exception: # noqa: BLE001 - best-effort noise suppression + pass + + from headroom.providers.agy import build_agy_env + from headroom.proxy.agy_ca import build_combined_bundle, ensure_root_ca + from headroom.proxy.agy_terminator import DEFAULT_ALLOWLIST + + allowlist = DEFAULT_ALLOWLIST + + ca_key, ca_cert, _key_path, _cert_path = ensure_root_ca() + bundle_path = build_combined_bundle() + + # Capture the corporate HTTPS_PROXY (if any) BEFORE building the child env, + # for transparency only. Chaining itself needs no plumbing: build_agy_env + # returns a copy and never mutates os.environ, so the terminator (running in + # THIS parent process) still reads the original corporate + # os.environ["HTTPS_PROXY"] for non-allowlisted CONNECT chaining. + corp_proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy") + + # ------------------------------------------------------------------ + # Observability: fail-open warning + session compression summary. + # Ref: headroom-30y.15 + # ------------------------------------------------------------------ + from headroom.providers.agy.stats import ( + AgySessionStats, + FailOpenWarnHandler, + install_fail_open_handler, + remove_fail_open_handler, + ) + + session_stats = AgySessionStats() + fail_open_handler: FailOpenWarnHandler | None = None + + servers: _AgyServers | None = None + old_sigint: Any = None + old_sigterm: Any = None + retrieve_registered = False + + # Shared Headroom proxy (default :8787). Its savings-inbox DRAIN loop + # (_drain_agy_savings_periodically) is what turns the savings.d events + # emitted below into the dashboard's $/token hero number and this + # project's row. Set up the same way every other wrap subcommand does + # (_make_cleanup / _register_proxy_client) but WITHOUT a second + # signal.signal(SIGTERM, ...): agy installs its own SIGTERM handler + # (_agy_sigterm below), which calls cleanup() itself, and the `finally` + # block below also calls cleanup() — together they cover both the + # signal-exit and normal-exit paths without clobbering agy's handler. + proxy_holder: list[subprocess.Popen | None] = [None] + cleanup = _make_cleanup(proxy_holder, port) + _register_proxy_client(port) + + # Cross-process savings: redirect THIS process's in-proxy funnel writes to a + # throwaway dir and turn on the inbox emit marker. agy runs its dispatch app + # in this process, so the funnel's durable writes (savings ledger, + # SavingsTracker, OTEL) must go nowhere durable — the shared proxy replays + # each emitted inbox event through its OWN funnel and is the sole writer of + # shared state. The tmp dir + env vars live only for this agy session. + agy_savings_tmp: str | None = None + try: + # MUST run before the os.environ mutations below: when no proxy is + # already running, _ensure_proxy -> _start_proxy snapshots + # os.environ.copy() to launch the shared proxy subprocess. If this ran + # after HEADROOM_AGY_INBOX_EMIT / HEADROOM_SAVINGS_PATH / + # HEADROOM_SAVINGS_EVENTS_PATH / HEADROOM_OTEL_METRICS_ENABLED were + # set, the shared durable proxy would inherit them: double-count its + # own traffic, redirect its durable savings ledger into this session's + # throwaway tmp dir (deleted on agy exit), and disable OTEL for every + # client sharing the proxy. agy uses its own MITM env (build_agy_env + # below) rather than a base-URL redirect, so unlike the other wrap + # subcommands we do NOT call _push_runtime_env here. + # _ensure_proxy returns (proxy, actual_port); agy addresses the shared + # proxy by the requested `port` throughout (_make_cleanup / + # _register_proxy_client above), so the bound port is unused here — but + # proxy_holder[0] MUST be the Popen, not the tuple, for cleanup to reap + # an agy-started proxy. + proxy_holder[0], _actual_port = _ensure_proxy( + port, no_proxy, agent_type="agy", code_graph=code_graph + ) + + agy_savings_tmp = tempfile.mkdtemp(prefix="headroom-agy-savings-") + os.environ["HEADROOM_SAVINGS_PATH"] = str(Path(agy_savings_tmp) / "proxy_savings.json") + os.environ["HEADROOM_SAVINGS_EVENTS_PATH"] = str( + Path(agy_savings_tmp) / "savings_events.jsonl" + ) + os.environ["HEADROOM_OTEL_METRICS_ENABLED"] = "0" + os.environ["HEADROOM_AGY_INBOX_EMIT"] = "1" + # Snapshot compression-store baseline and install the fail-open warning + # handler BEFORE the dispatch thread starts so we catch every event. + session_stats.snapshot_start() + fail_open_handler = install_fail_open_handler() + + agy_project = _project_name_from_cwd() + servers = _start_agy_servers(ca_key, ca_cert, start_retrieve=True, project=agy_project) + term_host, term_port = servers.terminator.address + terminator_url = f"http://{term_host}:{term_port}" + + env = build_agy_env( + terminator_url=terminator_url, + bundle_path=bundle_path, + base_env=os.environ.copy(), + ) + + env_vars_display = [ + f"HTTPS_PROXY={terminator_url} (agy CONNECT terminator)", + f"HTTP_PROXY={terminator_url}", + "NO_PROXY=127.0.0.1,localhost", + f"SSL_CERT_FILE={bundle_path}", + f"CACERT_PATH={bundle_path}", + f"NODE_EXTRA_CA_CERTS={bundle_path}", + ] + if corp_proxy: + env_vars_display.append( + f"chaining non-allowlisted CONNECTs via {redact_proxy_url(corp_proxy)}" + ) + + click.echo() + click.echo(" ╔═══════════════════════════════════════════════╗") + click.echo(" ║ HEADROOM WRAP: AGY ║") + click.echo(" ╚═══════════════════════════════════════════════╝") + click.echo() + click.echo(" ┌─ TLS INTERCEPTION DISCLOSURE ──────────────────") + for _intercepted_host in sorted(allowlist): + click.echo(f" │ Headroom terminates TLS for: {_intercepted_host}") + click.echo(" │ A process-local CA mints leaf certificates for those hosts.") + click.echo(" │ This CA is NEVER added to the OS trust store.") + click.echo(" │ Compression and context injection are applied on the decrypted stream.") + click.echo(" │ Leaf private keys: held in anonymous process memory (memfd) on Linux;") + click.echo(" │ on other platforms a 0600 temp file is written and unlinked immediately") + click.echo(" │ after load (permissions asserted). Keys are never trust-stored.") + click.echo(" │") + click.echo(" │ To opt out of interception: headroom wrap agy --no-intercept") + click.echo(" │ To revert all changes: headroom unwrap agy") + click.echo(" └────────────────────────────────────────────────") + click.echo() + click.echo(" Launching agy (traffic routed through Headroom MITM transport)...") + for var in env_vars_display: + click.echo(f" {var}") + if agy_args: + click.echo(f" Extra args: {' '.join(agy_args)}") + _print_telemetry_notice() + click.echo() + + # ------------------------------------------------------------------ + # MCP tooling wiring is gated on a runtime agy-version preflight: older + # or unknown agy binaries hang on ANY mcpServers entry when launched in + # print mode. Interactive mode is unaffected and is always wired. + # ------------------------------------------------------------------ + mcp_allowed = _agy_print_mode_mcp_allowed(agy_args, agy_bin) + if mcp_allowed: + # ------------------------------------------------------------------ + # MCP tooling is wired identically in print and interactive mode. agy + # 1.0.16 no longer hangs on MCP servers in --print mode (re-verified + # 2026-07-05: serena and the headroom retrieve server both answer in + # ~4s), so agy gets first-class MCP parity in every mode, like any + # other client. + # ------------------------------------------------------------------ + + # ------------------------------------------------------------------ + # Code-memory MCP — Serena is the active engine on current main. + # Retire any Headroom-installed tokensave entry left by older builds, + # then register Serena unless explicitly disabled. Wired in all modes + # once the agy print-mode version gate allows MCP. + # ------------------------------------------------------------------ + _disable_tokensave_mcp(AgyRegistrar(), verbose=False) + if not no_serena: + _setup_serena_mcp( + AgyRegistrar(), context="ide-assistant", verbose=False, force=True + ) + else: + _disable_serena_mcp( + AgyRegistrar(), + verbose=False, + reason="--no-serena", + ) + + # ------------------------------------------------------------------ + # Headroom retrieve MCP. The retrieve tool is an ``headroom mcp serve`` + # stdio child that resolves ``[Retrieve more: hash=…]`` markers from the + # shared on-disk CCR store. It is registered PERSISTENTLY and recorded + # in the install ledger (like Serena) so agy can cache and expose it + # across sessions — it is NOT reverted on teardown. + # Wired in all print-mode-capable agy versions. + # ------------------------------------------------------------------ + if no_mcp: + # Parity with `wrap claude` / `wrap opencode`: --no-mcp skips + # registration entirely. Compression markers then have no tool + # that can resolve them, so the handler must not ship any (the + # HEADROOM_AGY_RETRIEVE_WIRED gate below stays unset). + retrieve_registered = False + click.echo(" Skipping MCP retrieve tool (--no-mcp)") + elif servers is not None and servers.retrieve_port is not None: + retrieve_registered = _setup_headroom_retrieve_mcp_agy( + AgyRegistrar(), verbose=False + ) + else: + # No in-process servers this run. Leave a ledger-recorded + # PERSISTENT headroom entry in place (it resolves from the on-disk + # store, no live port required); only purge a stale NON-ledgered + # entry left by a pre-persistent SIGKILLed session (its ephemeral + # proxy URL is dead). Idempotent — no-op when absent. + from headroom.mcp_registry.ledger import headroom_installed_matching + + _reg = AgyRegistrar() + if not headroom_installed_matching(_reg.name, _reg.get_server("headroom")): + _reg.unregister_server("headroom") + + else: + # Print-mode MCP preflight failed: agy is older than + # _AGY_PRINT_MODE_MCP_MIN_VERSION, or its version could not be + # detected. Actively PURGE any MCP entries a prior interactive run + # may have persisted in mcp_config.json -- merely skipping + # registration is not enough, since a stale entry from an earlier + # run would still hang this print-mode invocation. All calls below + # are idempotent (no-op when the entry is already absent). The + # retrieve LISTENER started above still runs (harmless idle loopback) + # -- only MCP *registration* is suppressed here. + _purge_agy_mcp_entries(AgyRegistrar()) + _detected_version = _detect_agy_version(agy_bin) + _detected_str = ( + ".".join(str(part) for part in _detected_version) + if _detected_version is not None + else "unknown" + ) + click.echo( + f" MCP tooling: suppressed (detected agy version {_detected_str}; " + "print-mode MCP requires agy >= " + f"{'.'.join(str(p) for p in _AGY_PRINT_MODE_MCP_MIN_VERSION)}). " + "agy still runs transport-only.", + err=True, + ) + + # WU1 (headroom-37g.1): tell the in-process Cloud Code Assist handler + # whether the CCR retrieve listener is wired for this run. The handler + # ships recoverable functionResponse hash markers only when retrieval can + # resolve them; otherwise it falls back to lossless. The dispatch app + # runs in THIS process, so the signal must live in os.environ (mirrors + # HEADROOM_AGY_INBOX_EMIT above); also mirror it into the child env. + # HEADROOM_AGY_FR_MODE is already inherited via os.environ.copy() above. + # + # WIRED requires POSITIVE agy exposure, not just a successful handshake: + # the handshake proves wrap can spawn the child, but agy only surfaces + # tools it has cached, so a registered-then-reverted entry is rejected as + # "Unknown tool: headroom_retrieve". Gate WIRED on the exposure signal so + # ccr never ships unrecoverable markers on a false-positive handshake. + retrieve_exposed = retrieve_registered and _agy_exposes_retrieve_tool(AgyRegistrar()) + if retrieve_exposed: + os.environ["HEADROOM_AGY_RETRIEVE_WIRED"] = "1" + env["HEADROOM_AGY_RETRIEVE_WIRED"] = "1" + else: + os.environ.pop("HEADROOM_AGY_RETRIEVE_WIRED", None) + env.pop("HEADROOM_AGY_RETRIEVE_WIRED", None) + + _maybe_warn_agy_ccr_downgrade(retrieve_exposed) + + # ------------------------------------------------------------------ + # Install signal handlers so the terminator/dispatch are always torn + # down on SIGINT/SIGTERM (mirrors _launch_tool's signal-safe teardown + # without registering agy as a proxy client). SIGINT is ignored here + # so agy itself owns Ctrl-C; SIGTERM stops our servers then exits via + # SystemExit(143) so the finally below also runs. + def _agy_sigterm(_signum: int | None = None, _frame: Any = None) -> None: + # retrieve + code_graph are persistent ledger-recorded entries (like + # Serena), NOT reverted on exit — they resolve from the on-disk store + # and must survive so agy can cache/expose them next session. + _stop_agy_servers(servers) + cleanup() + # Flush compression summary on kill (idempotent — won't double-print + # if the finally below also runs). Ref: headroom-30y.15 + session_stats.print_summary(fail_open_handler) + remove_fail_open_handler(fail_open_handler) + raise SystemExit(143) + + old_sigint = signal.signal(signal.SIGINT, _ignore_child_sigint) + old_sigterm = signal.signal(signal.SIGTERM, _agy_sigterm) + + result = subprocess.run([agy_bin, *agy_args], env=env) + raise SystemExit(result.returncode) + + except SystemExit: + raise + except Exception as e: + # Walk the exception chain to surface a specific port-in-use message. + cause: BaseException | None = e + _port_in_use = False + while cause is not None: + if isinstance(cause, OSError) and cause.errno == errno.EADDRINUSE: + _port_in_use = True + break + cause = cause.__cause__ or cause.__context__ + if _port_in_use: + click.echo( + f"Error: a required proxy port is already in use ({cause}). " + "Stop the conflicting process and retry.", + err=True, + ) + else: + click.echo(f"Error: agy MITM transport failed to start: {e}", err=True) + raise SystemExit(1) from e + finally: + # The headroom retrieve entry is PERSISTENT (ledger-recorded, resolves + # from the on-disk store) — like Serena/CBM it is intentionally NOT + # reverted here so agy can cache and expose it on the next session. + # Restore prior signal handlers so they don't leak into the click process. + if old_sigint is not None: + signal.signal(signal.SIGINT, old_sigint) + if old_sigterm is not None: + signal.signal(signal.SIGTERM, old_sigterm) + _stop_agy_servers(servers) + cleanup() + # Print session compression summary (idempotent — won't double-print + # if _agy_sigterm already flushed it). Remove the logging handler so + # it doesn't leak into the click process. Ref: headroom-30y.15 + session_stats.print_summary(fail_open_handler) + remove_fail_open_handler(fail_open_handler) + # Clean up the throwaway savings dir (the env vars die with the process, + # which is fine — nothing else in this process consumes them). + if agy_savings_tmp is not None: + shutil.rmtree(agy_savings_tmp, ignore_errors=True) + + +# ============================================================================= +# unwrap agy +# ============================================================================= + + +@unwrap.command("agy") +def unwrap_agy() -> None: + """Undo ``headroom wrap agy`` — revert any persistent agy configuration changes. + + Removes the Headroom block from GEMINI.md and unregisters any MCP server + entry that Headroom registered in the Antigravity CLI config. All user + content outside Headroom-managed markers is preserved. + """ + from headroom.mcp_registry.agy import AgyRegistrar + + click.echo() + click.echo(" ╔═══════════════════════════════════════════════╗") + click.echo(" ║ HEADROOM UNWRAP: AGY ║") + click.echo(" ╚═══════════════════════════════════════════════╝") + click.echo() + + # 1. Remove headroom block from GEMINI.md. + gemini_md = Path.home() / ".gemini" / "GEMINI.md" + if _remove_gemini_md_block(gemini_md, verbose=True): + click.echo(" headroom block removed from GEMINI.md") + else: + click.echo(" GEMINI.md: no headroom block found (already clean)") + + # 2. Remove the headroom MCP retrieve entry only if the ledger proves + # 'wrap agy' installed it. It is now a persistent, ledger-recorded server, + # so cooperative uninstall is ledger-gated like Serena/CBM. + # BEHAVIOR: a 'headroom mcp install' fleet entry is NOT ledger-recorded by + # that path, so unwrap now leaves it in place (respecting the deliberate + # fleet-wide install) instead of clobbering it. A stable persistent entry + # is harmless to leave — it resolves from the on-disk store, never hangs. + agy_reg = AgyRegistrar() + retrieve_status = _remove_headroom_installed_retrieve_mcp(agy_reg) + if retrieve_status == "removed": + click.echo(" Removed Headroom MCP retrieve tool from agy.") + elif retrieve_status == "failed": + click.echo(" Headroom MCP retrieve tool matched Headroom ledger but could not be removed.") + else: # not_headroom_owned (absent, or a user/fleet-managed entry left untouched) + click.echo(" Headroom MCP retrieve tool left as-is (not 'wrap agy'-installed).") + + # 3. Remove Serena MCP only if the ledger proves Headroom installed it; + # a user-managed 'serena' entry is left untouched. + serena_status = _remove_headroom_installed_serena_mcp(agy_reg) + if serena_status == "removed": + click.echo(" Removed Headroom-installed Serena MCP server from agy.") + elif serena_status == "failed": + click.echo(" Serena MCP server matched Headroom ledger but could not be removed.") + elif serena_status == "not_headroom_owned": + click.echo(" Kept user-managed Serena MCP server (not Headroom-owned).") + + # 4. Remove any legacy codebase-memory-mcp entry an older build registered. + # agy no longer wires a code-graph MCP: --code-graph now drives the + # proxy-side watcher, matching `wrap claude` / `unwrap claude`. + if agy_reg.unregister_server(_CBM_MCP_SERVER_NAME): + click.echo(" Removed legacy codebase-memory-mcp code graph server from agy.") + + # 5. Remove the tokensave code-graph MCP only if the ledger proves Headroom + # installed it as the primary compressor (user-managed entries untouched). + tokensave_status = _remove_headroom_installed_tokensave_mcp(agy_reg) + if tokensave_status == "removed": + click.echo(" Removed Headroom-installed tokensave MCP server from agy.") + elif tokensave_status == "failed": + click.echo(" tokensave MCP server matched Headroom ledger but could not be removed.") + elif tokensave_status == "not_headroom_owned": + click.echo(" tokensave MCP server left as-is (not Headroom-installed).") + + click.echo() + click.echo("✓ agy headroom configuration reverted.") + click.echo() + + # ============================================================================= # Oh My Pi (omp) # ============================================================================= diff --git a/headroom/evals/session_probes.py b/headroom/evals/session_probes.py index 5574841ab..fb1ba46e9 100644 --- a/headroom/evals/session_probes.py +++ b/headroom/evals/session_probes.py @@ -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=|<]+>>") +# 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 diff --git a/headroom/mcp_registry/__init__.py b/headroom/mcp_registry/__init__.py index f19ad1ba7..076bb3b6f 100644 --- a/headroom/mcp_registry/__init__.py +++ b/headroom/mcp_registry/__init__.py @@ -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", diff --git a/headroom/mcp_registry/agy.py b/headroom/mcp_registry/agy.py new file mode 100644 index 000000000..facf0c1da --- /dev/null +++ b/headroom/mcp_registry/agy.py @@ -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": {"": {"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 ``/mcp//.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: ``/mcp``. + + agy writes ``/mcp//.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) diff --git a/headroom/mcp_registry/install.py b/headroom/mcp_registry/install.py index 4d1fbeb72..669469462 100644 --- a/headroom/mcp_registry/install.py +++ b/headroom/mcp_registry/install.py @@ -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: diff --git a/headroom/parser.py b/headroom/parser.py index a51d435a3..1849d814b 100644 --- a/headroom/parser.py +++ b/headroom/parser.py @@ -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=|<]+>>") +# 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_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 diff --git a/headroom/providers/agy/__init__.py b/headroom/providers/agy/__init__.py new file mode 100644 index 000000000..34deb4138 --- /dev/null +++ b/headroom/providers/agy/__init__.py @@ -0,0 +1,5 @@ +"""agy-specific provider helpers.""" + +from .runtime import build_agy_env + +__all__ = ["build_agy_env"] diff --git a/headroom/providers/agy/runtime.py b/headroom/providers/agy/runtime.py new file mode 100644 index 000000000..07d86158b --- /dev/null +++ b/headroom/providers/agy/runtime.py @@ -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:``). + 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 diff --git a/headroom/providers/agy/stats.py b/headroom/providers/agy/stats.py new file mode 100644 index 000000000..84298ad7a --- /dev/null +++ b/headroom/providers/agy/stats.py @@ -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 diff --git a/headroom/providers/proxy_routes.py b/headroom/providers/proxy_routes.py index 5cf11ea47..6c097a536 100644 --- a/headroom/providers/proxy_routes.py +++ b/headroom/providers/proxy_routes.py @@ -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( diff --git a/headroom/providers/proxy_targets.py b/headroom/providers/proxy_targets.py index 0ef4dfff8..94e5ad244 100644 --- a/headroom/providers/proxy_targets.py +++ b/headroom/providers/proxy_targets.py @@ -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 diff --git a/headroom/proxy/agy_ca.py b/headroom/proxy/agy_ca.py new file mode 100644 index 000000000..d411c4d83 --- /dev/null +++ b/headroom/proxy/agy_ca.py @@ -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 ``.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 ``/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 diff --git a/headroom/proxy/agy_dispatch.py b/headroom/proxy/agy_dispatch.py new file mode 100644 index 000000000..b3e362a4b --- /dev/null +++ b/headroom/proxy/agy_dispatch.py @@ -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", ) + 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() diff --git a/headroom/proxy/agy_retrieve.py b/headroom/proxy/agy_retrieve.py new file mode 100644 index 000000000..a8d2f19b9 --- /dev/null +++ b/headroom/proxy/agy_retrieve.py @@ -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", ) + 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) diff --git a/headroom/proxy/agy_savings_inbox.py b/headroom/proxy/agy_savings_inbox.py new file mode 100644 index 000000000..4ca428b3c --- /dev/null +++ b/headroom/proxy/agy_savings_inbox.py @@ -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", +] diff --git a/headroom/proxy/agy_terminator.py b/headroom/proxy/agy_terminator.py new file mode 100644 index 000000000..dadac038a --- /dev/null +++ b/headroom/proxy/agy_terminator.py @@ -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:`` 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:`` (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() diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index 79463c579..147f60298 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -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}" diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 42ac1a8c0..e4a21752a 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -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) diff --git a/headroom/proxy/handlers/streaming.py b/headroom/proxy/handlers/streaming.py index 6d0f31606..d7929bad5 100644 --- a/headroom/proxy/handlers/streaming.py +++ b/headroom/proxy/handlers/streaming.py @@ -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) diff --git a/headroom/proxy/outcome.py b/headroom/proxy/outcome.py index ba94feee7..9f2b967b2 100644 --- a/headroom/proxy/outcome.py +++ b/headroom/proxy/outcome.py @@ -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: diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 233a0367c..2a57e5903 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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, diff --git a/headroom/transforms/agy_fr_compressor.py b/headroom/transforms/agy_fr_compressor.py new file mode 100644 index 000000000..2009b56c7 --- /dev/null +++ b/headroom/transforms/agy_fr_compressor.py @@ -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"(? 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"] diff --git a/headroom/transforms/compression_units.py b/headroom/transforms/compression_units.py index 063b55b7f..3965d060b 100644 --- a/headroom/transforms/compression_units.py +++ b/headroom/transforms/compression_units.py @@ -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_MARKER_RE = re.compile(rf"(?m)^.*(?:{CCR_MARKER_ALTERNATION}).*$") _LOSSY_UNMARKED_STRATEGIES = { CompressionStrategy.KOMPRESS.value, diff --git a/llms.txt b/llms.txt index fe2412fdb..8d4c99fa6 100644 --- a/llms.txt +++ b/llms.txt @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 3e930bc2d..ae38cf7e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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). diff --git a/tests/conftest.py b/tests/conftest.py index 54334d9ac..db909f45a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 diff --git a/tests/test_agy_ca.py b/tests/test_agy_ca.py new file mode 100644 index 000000000..40f0ba9a4 --- /dev/null +++ b/tests/test_agy_ca.py @@ -0,0 +1,1299 @@ +"""Tests for headroom.proxy.agy_ca — root CA lifecycle + combined bundle. + +All tests use pytest's tmp_path; real ~/.headroom is never touched. +""" + +from __future__ import annotations + +import datetime +import os +import sys +from pathlib import Path + +import pytest +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from cryptography.x509.oid import NameOID + +from headroom.proxy.agy_ca import ( + _BUNDLE_NAME, + _CA_CERT_NAME, + _CA_KEY_NAME, + _OS_TRUST_PATHS, + _assert_perms, + _cert_near_expiry, + _collect_corporate_ca_pems, + _detect_system_bundle, + _is_ca_cert, + _load_via_mkstemp, + _not_in_os_trust, + _parse_ca_certs_from_pem, + _system_trust_pem, + _windows_trust_pem, + _write_all_fd, + build_combined_bundle, + ensure_root_ca, + load_cert_chain_in_memory, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_cert( + is_ca: bool, + days_valid: int = 3650, + path_length: int | None = 0, +) -> bytes: + """Generate a minimal PEM certificate for testing.""" + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + subject = issuer = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "test")]) + builder = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(issuer) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(datetime.datetime.now(datetime.timezone.utc)) + .not_valid_after( + datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=days_valid) + ) + .add_extension( + x509.BasicConstraints(ca=is_ca, path_length=path_length if is_ca else None), + critical=True, + ) + ) + cert = builder.sign(key, hashes.SHA256()) + return cert.public_bytes(serialization.Encoding.PEM) + + +def _fake_system_bundle(tmp_path: Path, pem_data: bytes | None = None) -> Path: + """Write a minimal fake system bundle, returning its path.""" + if pem_data is None: + pem_data = _make_cert(is_ca=True) + p = tmp_path / "system-ca-bundle.pem" + p.write_bytes(pem_data) + return p + + +# --------------------------------------------------------------------------- +# ensure_root_ca — generation +# --------------------------------------------------------------------------- + + +def test_ca_generated_on_first_call(tmp_path: Path) -> None: + """First call creates key + cert under base_dir/ca/.""" + key, cert, key_path, cert_path = ensure_root_ca(base_dir=tmp_path) + assert key_path.exists() + assert cert_path.exists() + assert _is_ca_cert(cert) + + +def test_ca_dir_is_0700(tmp_path: Path) -> None: + ensure_root_ca(base_dir=tmp_path) + ca_dir = tmp_path / "ca" + _assert_perms(ca_dir, 0o700) + + +def test_ca_key_is_0600(tmp_path: Path) -> None: + _, _, key_path, _ = ensure_root_ca(base_dir=tmp_path) + _assert_perms(key_path, 0o600) + + +def test_ca_cert_is_0600(tmp_path: Path) -> None: + _, _, _, cert_path = ensure_root_ca(base_dir=tmp_path) + _assert_perms(cert_path, 0o600) + + +def test_ca_has_basic_constraints_ca_true(tmp_path: Path) -> None: + _, cert, _, _ = ensure_root_ca(base_dir=tmp_path) + bc = cert.extensions.get_extension_for_class(x509.BasicConstraints) + assert bc.value.ca is True + assert bc.value.path_length == 0 + + +def test_ca_has_long_validity(tmp_path: Path) -> None: + """Cert must be valid for at least 9 years (allowing some clock skew).""" + _, cert, _, _ = ensure_root_ca(base_dir=tmp_path) + now = datetime.datetime.now(datetime.timezone.utc) + delta = cert.not_valid_after_utc - now + assert delta.days >= 365 * 9 + + +# --------------------------------------------------------------------------- +# ensure_root_ca — idempotency (reuse) +# --------------------------------------------------------------------------- + + +def test_second_call_reuses_existing_ca(tmp_path: Path) -> None: + """Second call with valid existing CA returns same cert (by serial).""" + _, cert1, _, _ = ensure_root_ca(base_dir=tmp_path) + _, cert2, _, _ = ensure_root_ca(base_dir=tmp_path) + assert cert1.serial_number == cert2.serial_number + + +def test_second_call_key_object_matches(tmp_path: Path) -> None: + key1, _, _, _ = ensure_root_ca(base_dir=tmp_path) + key2, _, _, _ = ensure_root_ca(base_dir=tmp_path) + pub1 = key1.public_key().public_bytes( + serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo + ) + pub2 = key2.public_key().public_bytes( + serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo + ) + assert pub1 == pub2 + + +# --------------------------------------------------------------------------- +# ensure_root_ca — regeneration on expiry +# --------------------------------------------------------------------------- + + +def _write_expiring_ca(base_dir: Path, days_valid: int = 1) -> None: + """Overwrite the CA with a cert that expires soon (within regen threshold).""" + ca_dir = base_dir / "ca" + ca_dir.mkdir(mode=0o700, parents=True, exist_ok=True) + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + subject = issuer = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "expiring")]) + now = datetime.datetime.now(datetime.timezone.utc) + cert = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(issuer) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(days=1)) + .not_valid_after(now + datetime.timedelta(days=days_valid)) + .add_extension(x509.BasicConstraints(ca=True, path_length=0), critical=True) + .sign(key, hashes.SHA256()) + ) + key_pem = key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.TraditionalOpenSSL, + serialization.NoEncryption(), + ) + cert_pem = cert.public_bytes(serialization.Encoding.PEM) + key_path = ca_dir / _CA_KEY_NAME + cert_path = ca_dir / _CA_CERT_NAME + key_path.write_bytes(key_pem) + key_path.chmod(0o600) + cert_path.write_bytes(cert_pem) + cert_path.chmod(0o600) + return cert.serial_number # type: ignore[return-value] + + +def test_regen_on_expiry_produces_new_serial(tmp_path: Path) -> None: + old_serial = _write_expiring_ca(tmp_path, days_valid=1) + _, new_cert, _, _ = ensure_root_ca(base_dir=tmp_path) + assert new_cert.serial_number != old_serial + + +def test_regen_deletes_old_bundle(tmp_path: Path) -> None: + """Stale combined bundle is removed when CA is regenerated.""" + old_bundle = tmp_path / _BUNDLE_NAME + old_bundle.write_bytes(b"stale") + old_bundle.chmod(0o600) + _write_expiring_ca(tmp_path, days_valid=1) + ensure_root_ca(base_dir=tmp_path) + # Bundle was deleted; new content would need build_combined_bundle. + assert not old_bundle.exists() + + +def test_regen_deletes_old_leaves(tmp_path: Path) -> None: + """Leaf cert directory is cleaned up on regeneration.""" + leaves_dir = tmp_path / "leaves" + leaves_dir.mkdir(mode=0o700) + (leaves_dir / "example.com.crt").write_bytes(b"leaf") + _write_expiring_ca(tmp_path, days_valid=1) + ensure_root_ca(base_dir=tmp_path) + assert not leaves_dir.exists() + + +# --------------------------------------------------------------------------- +# _is_ca_cert +# --------------------------------------------------------------------------- + + +def test_is_ca_cert_true_for_ca() -> None: + pem = _make_cert(is_ca=True) + cert = x509.load_pem_x509_certificate(pem) + assert _is_ca_cert(cert) is True + + +def test_is_ca_cert_false_for_leaf() -> None: + pem = _make_cert(is_ca=False) + cert = x509.load_pem_x509_certificate(pem) + assert _is_ca_cert(cert) is False + + +# --------------------------------------------------------------------------- +# _parse_ca_certs_from_pem — per-object filter +# --------------------------------------------------------------------------- + + +def test_parse_filters_non_ca_leaves() -> None: + """Multi-cert PEM: only CA:TRUE objects survive.""" + ca_pem = _make_cert(is_ca=True) + leaf_pem = _make_cert(is_ca=False) + combined = ca_pem + leaf_pem + results = _parse_ca_certs_from_pem(combined) + assert len(results) == 1 + cert = x509.load_pem_x509_certificate(results[0]) + assert _is_ca_cert(cert) is True + + +def test_parse_all_ca_certs_included() -> None: + ca1 = _make_cert(is_ca=True) + ca2 = _make_cert(is_ca=True) + combined = ca1 + ca2 + results = _parse_ca_certs_from_pem(combined) + assert len(results) == 2 + + +def test_parse_empty_pem_returns_empty() -> None: + assert _parse_ca_certs_from_pem(b"") == [] + + +def test_parse_skips_invalid_pem_blocks() -> None: + ca_pem = _make_cert(is_ca=True) + garbage = b"-----BEGIN CERTIFICATE-----\nZZZZZZ\n-----END CERTIFICATE-----\n" + combined = ca_pem + garbage + results = _parse_ca_certs_from_pem(combined) + # Only the valid CA cert should come through. + assert len(results) == 1 + + +# --------------------------------------------------------------------------- +# _collect_corporate_ca_pems +# --------------------------------------------------------------------------- + + +def test_collect_corp_ca_from_env_var(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Corporate CA file with one CA + one leaf → only CA returned.""" + ca_pem = _make_cert(is_ca=True) + leaf_pem = _make_cert(is_ca=False) + corp_file = tmp_path / "corp.pem" + corp_file.write_bytes(ca_pem + leaf_pem) + + monkeypatch.setenv("SSL_CERT_FILE", str(corp_file)) + results = _collect_corporate_ca_pems(("SSL_CERT_FILE",)) + assert len(results) == 1 + cert = x509.load_pem_x509_certificate(results[0]) + assert _is_ca_cert(cert) is True + + +def test_collect_corp_ca_missing_file_warns( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Missing corporate CA file → empty result (no crash).""" + monkeypatch.setenv("SSL_CERT_FILE", str(tmp_path / "nonexistent.pem")) + results = _collect_corporate_ca_pems(("SSL_CERT_FILE",)) + assert results == [] + + +def test_collect_corp_ca_unset_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("SSL_CERT_FILE", raising=False) + monkeypatch.delenv("NODE_EXTRA_CA_CERTS", raising=False) + results = _collect_corporate_ca_pems(("SSL_CERT_FILE", "NODE_EXTRA_CA_CERTS")) + assert results == [] + + +# --------------------------------------------------------------------------- +# build_combined_bundle +# --------------------------------------------------------------------------- + + +def test_bundle_is_created(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + sys_bundle = _fake_system_bundle(tmp_path) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(sys_bundle),), + ) + bundle_path = build_combined_bundle(base_dir=tmp_path, corp_env_vars=()) + assert bundle_path.exists() + assert bundle_path.stat().st_size > 0 + + +def test_bundle_is_0600(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + sys_bundle = _fake_system_bundle(tmp_path) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(sys_bundle),), + ) + bundle_path = build_combined_bundle(base_dir=tmp_path, corp_env_vars=()) + _assert_perms(bundle_path, 0o600) + + +def test_parent_dir_is_0700(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + sys_bundle = _fake_system_bundle(tmp_path) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(sys_bundle),), + ) + build_combined_bundle(base_dir=tmp_path, corp_env_vars=()) + _assert_perms(tmp_path, 0o700) + + +def test_bundle_contains_system_ca(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + sys_ca_pem = _make_cert(is_ca=True) + sys_bundle = _fake_system_bundle(tmp_path, pem_data=sys_ca_pem) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(sys_bundle),), + ) + bundle_path = build_combined_bundle(base_dir=tmp_path, corp_env_vars=()) + bundle_data = bundle_path.read_bytes() + # The system CA PEM bytes must appear verbatim in the bundle. + assert sys_ca_pem in bundle_data + + +def test_bundle_contains_headroom_ca(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + sys_bundle = _fake_system_bundle(tmp_path) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(sys_bundle),), + ) + _, ca_cert, _, _ = ensure_root_ca(base_dir=tmp_path) + headroom_pem = ca_cert.public_bytes(serialization.Encoding.PEM) + bundle_path = build_combined_bundle(base_dir=tmp_path, corp_env_vars=()) + bundle_data = bundle_path.read_bytes() + assert headroom_pem in bundle_data + + +def test_bundle_contains_corp_ca_but_not_leaf( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Corporate CA:TRUE cert appears in bundle; leaf cert does not.""" + sys_bundle = _fake_system_bundle(tmp_path) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(sys_bundle),), + ) + corp_ca_pem = _make_cert(is_ca=True) + leaf_pem = _make_cert(is_ca=False) + corp_file = tmp_path / "corp.pem" + corp_file.write_bytes(corp_ca_pem + leaf_pem) + + # First call without corp CAs to seed the CA on disk. + build_combined_bundle( + base_dir=tmp_path, + corp_env_vars=(), + ) + # Call again using a custom corp_env_vars pointing at our fixture file. + monkeypatch.setenv("_TEST_CORP_CA", str(corp_file)) + bundle_path2 = build_combined_bundle( + base_dir=tmp_path, + corp_env_vars=("_TEST_CORP_CA",), + ) + bundle_data = bundle_path2.read_bytes() + assert corp_ca_pem in bundle_data + assert leaf_pem not in bundle_data + + +def test_windows_trust_pem_filters_non_ca(monkeypatch: pytest.MonkeyPatch) -> None: + """The Windows ssl.enum_certificates path must drop non-CA (leaf) certs. + + ssl.enum_certificates returns every cert in the store including leaf certs; + _windows_trust_pem must run them through the CA:TRUE filter so only CA + anchors end up in the trust bundle. Mock-driven so it runs on every OS + (ssl.enum_certificates does not exist off Windows). + """ + ca_pem = _make_cert(is_ca=True) + leaf_pem = _make_cert(is_ca=False) + ca_cert = x509.load_pem_x509_certificate(ca_pem) + leaf_cert = x509.load_pem_x509_certificate(leaf_pem) + ca_der = ca_cert.public_bytes(serialization.Encoding.DER) + leaf_der = leaf_cert.public_bytes(serialization.Encoding.DER) + + def fake_enum(store: str) -> list[tuple[bytes, str, bool]]: + # Return the CA + leaf only for ROOT; CA store empty (avoid double count). + if store == "ROOT": + return [(ca_der, "x509_asn", True), (leaf_der, "x509_asn", True)] + return [] + + monkeypatch.setattr("ssl.enum_certificates", fake_enum, raising=False) + + result = _windows_trust_pem() + # Parse EVERY cert block in the raw result (NOT via the CA filter) so a + # regression that dropped the internal filter would surface the leaf here. + marker = b"-----BEGIN CERTIFICATE-----" + present = { + x509.load_pem_x509_certificate(marker + block).serial_number + for block in result.split(marker)[1:] + } + assert ca_cert.serial_number in present, "CA anchor must be present" + assert leaf_cert.serial_number not in present, "leaf cert must be filtered out" + + +def test_windows_trust_pem_enum_failure_raises(monkeypatch: pytest.MonkeyPatch) -> None: + """A failure reading the Windows trust store surfaces as a clear RuntimeError. + + Without this, an ssl.enum_certificates OSError would propagate raw through + _system_trust_pem (which only guards RuntimeError) and crash `wrap agy`. + """ + + def boom(store: str) -> list[tuple[bytes, str, bool]]: + raise OSError("simulated cert-store failure") + + monkeypatch.setattr("ssl.enum_certificates", boom, raising=False) + with pytest.raises(RuntimeError, match="Windows system trust store"): + _windows_trust_pem() + + +def test_windows_trust_pem_empty_raises(monkeypatch: pytest.MonkeyPatch) -> None: + """An empty CA set after filtering must fail loud, not return b''. + + A silent empty system-trust component would leave the combined bundle + trusting only the headroom MITM root. + """ + leaf_der = x509.load_pem_x509_certificate(_make_cert(is_ca=False)).public_bytes( + serialization.Encoding.DER + ) + + def only_leaf(store: str) -> list[tuple[bytes, str, bool]]: + return [(leaf_der, "x509_asn", True)] if store == "ROOT" else [] + + monkeypatch.setattr("ssl.enum_certificates", only_leaf, raising=False) + with pytest.raises(RuntimeError, match="no CA anchors"): + _windows_trust_pem() + + +def test_bundle_not_in_os_trust_paths(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Bundle path must not reside under any known OS trust store location.""" + sys_bundle = _fake_system_bundle(tmp_path) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(sys_bundle),), + ) + bundle_path = build_combined_bundle(base_dir=tmp_path, corp_env_vars=()) + resolved = str(bundle_path.resolve()) + for trust_path in _OS_TRUST_PATHS: + assert not resolved.startswith(trust_path), ( + f"Bundle {resolved} is inside OS trust path {trust_path}" + ) + + +def test_ca_never_written_to_os_trust_store( + tmp_path: Path, +) -> None: + """CA key + cert paths must not reside under OS trust store directories.""" + _, _, key_path, cert_path = ensure_root_ca(base_dir=tmp_path) + for path in (key_path, cert_path): + resolved = str(path.resolve()) + for trust_path in _OS_TRUST_PATHS: + assert not resolved.startswith(trust_path), ( + f"{path} is inside OS trust path {trust_path}" + ) + + +# --------------------------------------------------------------------------- +# fail-fast: no system bundle +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif( + sys.platform == "win32", + reason="Windows falls back to ssl.enum_certificates when no bundle file exists", +) +def test_no_system_bundle_raises(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (), + ) + with pytest.raises(RuntimeError, match="No system CA bundle found"): + build_combined_bundle(base_dir=tmp_path, corp_env_vars=()) + + +# --------------------------------------------------------------------------- +# _cert_near_expiry +# --------------------------------------------------------------------------- + + +def test_cert_near_expiry_true_for_expiring() -> None: + pem = _make_cert(is_ca=True, days_valid=1) + cert = x509.load_pem_x509_certificate(pem) + assert _cert_near_expiry(cert) is True + + +def test_cert_near_expiry_false_for_valid() -> None: + pem = _make_cert(is_ca=True, days_valid=3650) + cert = x509.load_pem_x509_certificate(pem) + assert _cert_near_expiry(cert) is False + + +# --------------------------------------------------------------------------- +# Bundle idempotency +# --------------------------------------------------------------------------- + + +def test_build_bundle_twice_same_content(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Building the bundle twice without CA regen produces identical content.""" + sys_bundle = _fake_system_bundle(tmp_path) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(sys_bundle),), + ) + path1 = build_combined_bundle(base_dir=tmp_path, corp_env_vars=()) + data1 = path1.read_bytes() + path2 = build_combined_bundle(base_dir=tmp_path, corp_env_vars=()) + data2 = path2.read_bytes() + assert data1 == data2 + + +# --------------------------------------------------------------------------- +# Regression: clean-install with nested base_dir (parents must be created) +# --------------------------------------------------------------------------- + + +def test_clean_install_nested_base_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """ensure_root_ca then build_combined_bundle on a completely fresh nested + base_dir must succeed and leave base_dir at 0o700. + + Constructs base_dir as tmp_path / "sub" / ".headroom" so the code itself + must create all intermediate directories — none are pre-created. + """ + base_dir = tmp_path / "sub" / ".headroom" + # Sanity: must not exist before the call. + assert not base_dir.exists() + + sys_bundle = _fake_system_bundle(tmp_path) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(sys_bundle),), + ) + + # This must not raise AssertionError or PermissionError. + ensure_root_ca(base_dir=base_dir) + bundle_path = build_combined_bundle(base_dir=base_dir, corp_env_vars=()) + + # base_dir itself must be 0o700 (the root cause of the original bug). + _assert_perms(base_dir, 0o700) + assert bundle_path.exists() + + +# --------------------------------------------------------------------------- +# Helpers shared by load_cert_chain_in_memory tests +# --------------------------------------------------------------------------- + + +def _make_leaf_pem_pair() -> tuple[bytes, bytes]: + """Return (cert_pem, key_pem) for a minimal self-signed leaf.""" + from cryptography.x509.oid import ExtendedKeyUsageOID, NameOID + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + now = datetime.datetime.now(datetime.timezone.utc) + cert = ( + x509.CertificateBuilder() + .subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "leaf.test")])) + .issuer_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "leaf.test")])) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now) + .not_valid_after(now + datetime.timedelta(hours=72)) + .add_extension(x509.SubjectAlternativeName([x509.DNSName("leaf.test")]), critical=False) + .add_extension(x509.ExtendedKeyUsage([ExtendedKeyUsageOID.SERVER_AUTH]), critical=True) + .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True) + .sign(key, hashes.SHA256()) + ) + cert_pem = cert.public_bytes(serialization.Encoding.PEM) + key_pem = key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.TraditionalOpenSSL, + serialization.NoEncryption(), + ) + return cert_pem, key_pem + + +# --------------------------------------------------------------------------- +# load_cert_chain_in_memory — primary path (memfd on Linux) +# --------------------------------------------------------------------------- + + +def test_load_cert_chain_in_memory_loads_usable_ctx() -> None: + """Combined cert+key is loaded into a usable SSLContext; no exception.""" + import ssl + + cert_pem, key_pem = _make_leaf_pem_pair() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + # Must not raise. + load_cert_chain_in_memory(ctx, cert_pem, key_pem) + + +@pytest.mark.skipif(sys.platform != "linux", reason="requires /proc/self/fd") +def test_load_cert_chain_in_memory_no_fd_leak() -> None: + """After load, the memfd (or temp file) is closed — no leaked descriptors.""" + import ssl + + cert_pem, key_pem = _make_leaf_pem_pair() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + + before = set(os.listdir("/proc/self/fd")) + load_cert_chain_in_memory(ctx, cert_pem, key_pem) + after = set(os.listdir("/proc/self/fd")) + + # The only new fd allowed is the /proc/self/fd dirfd opened by listdir itself. + new_fds = after - before + # Filter out the dirfd from the listdir call above (it closes immediately). + assert len(new_fds) == 0, f"Leaked file descriptors after load: {new_fds}" + + +def test_load_cert_chain_in_memory_no_tmpfile_on_linux(monkeypatch: pytest.MonkeyPatch) -> None: + """On Linux (memfd available), mkstemp and NamedTemporaryFile are NOT called.""" + import ssl + import tempfile as _tempfile + + cert_pem, key_pem = _make_leaf_pem_pair() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + + mkstemp_called = [False] + named_tmp_called = [False] + + original_mkstemp = _tempfile.mkstemp + original_named = _tempfile.NamedTemporaryFile + + def _spy_mkstemp(*args: object, **kwargs: object) -> object: + mkstemp_called[0] = True + return original_mkstemp(*args, **kwargs) + + def _spy_named(*args: object, **kwargs: object) -> object: + named_tmp_called[0] = True + return original_named(*args, **kwargs) + + monkeypatch.setattr(_tempfile, "mkstemp", _spy_mkstemp) + monkeypatch.setattr(_tempfile, "NamedTemporaryFile", _spy_named) + + if sys.platform == "linux" and hasattr(os, "memfd_create"): + load_cert_chain_in_memory(ctx, cert_pem, key_pem) + assert not mkstemp_called[0], "mkstemp must NOT be called when memfd_create is available" + assert not named_tmp_called[0], ( + "NamedTemporaryFile must NOT be called when memfd_create is available" + ) + + +# --------------------------------------------------------------------------- +# load_cert_chain_in_memory — short-write safety +# --------------------------------------------------------------------------- + + +def test_load_cert_chain_in_memory_short_write_handled() -> None: + """Helper writes all bytes even if os.write short-writes (1 byte at a time).""" + import ssl + + cert_pem, key_pem = _make_leaf_pem_pair() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + + if not hasattr(os, "memfd_create"): + pytest.skip("memfd_create not available on this platform") + + original_write = os.write + written_chunks: list[int] = [] + + def _one_byte_write(fd: int, data: bytes | bytearray) -> int: + # Only short-write to memfd fds; pass through others. + try: + path = os.readlink(f"/proc/self/fd/{fd}") + except OSError: + path = "" + if "memfd" in path or "anon" in path.lower(): + n = original_write(fd, data[:1]) + written_chunks.append(n) + return n + return original_write(fd, data) + + import unittest.mock + + with unittest.mock.patch("os.write", side_effect=_one_byte_write): + # Must succeed despite 1-byte writes. + load_cert_chain_in_memory(ctx, cert_pem, key_pem) + + total = sum(written_chunks) + expected = len(cert_pem + key_pem) + assert total == expected, f"Expected {expected} bytes written in chunks, got {total}" + + +# --------------------------------------------------------------------------- +# load_cert_chain_in_memory — fallback path (memfd absent/unavailable) +# --------------------------------------------------------------------------- + + +def test_load_cert_chain_in_memory_fallback_when_no_memfd( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """When memfd_create is absent, fallback uses mkstemp (0600) and unlinks it.""" + import ssl + import tempfile as _tempfile + + cert_pem, key_pem = _make_leaf_pem_pair() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + + # Force fallback: remove memfd_create from os. + monkeypatch.delattr(os, "memfd_create", raising=False) + + tmp_paths_created: list[str] = [] + tmp_paths_unlinked: list[str] = [] + original_mkstemp = _tempfile.mkstemp + original_unlink = os.unlink + + def _spy_mkstemp(*args: object, **kwargs: object) -> tuple[int, str]: + fd, path = original_mkstemp(*args, **kwargs) + tmp_paths_created.append(path) + return fd, path + + def _spy_unlink(path: str, *args: object, **kwargs: object) -> None: + if any(path == p for p in tmp_paths_created): + tmp_paths_unlinked.append(path) + original_unlink(path, *args, **kwargs) + + monkeypatch.setattr(_tempfile, "mkstemp", _spy_mkstemp) + monkeypatch.setattr(os, "unlink", _spy_unlink) + + load_cert_chain_in_memory(ctx, cert_pem, key_pem) + + assert tmp_paths_created, "Fallback must call mkstemp" + for p in tmp_paths_created: + assert not os.path.exists(p), f"Temp file {p} must be unlinked after load" + assert set(tmp_paths_created) == set(tmp_paths_unlinked), ( + "Every temp file created must be unlinked" + ) + + +def test_load_cert_chain_in_memory_fallback_0600(monkeypatch: pytest.MonkeyPatch) -> None: + """Fallback temp file has 0600 permissions (asserted by helper).""" + import ssl + import stat as _stat + import tempfile as _tempfile + + cert_pem, key_pem = _make_leaf_pem_pair() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + + monkeypatch.delattr(os, "memfd_create", raising=False) + + observed_modes: list[int] = [] + original_mkstemp = _tempfile.mkstemp + + def _spy_mkstemp(*args: object, **kwargs: object) -> tuple[int, str]: + fd, path = original_mkstemp(*args, **kwargs) + mode = _stat.S_IMODE(os.stat(path).st_mode) + observed_modes.append(mode) + return fd, path + + monkeypatch.setattr(_tempfile, "mkstemp", _spy_mkstemp) + + load_cert_chain_in_memory(ctx, cert_pem, key_pem) + + assert observed_modes, "Fallback must call mkstemp" + if sys.platform != "win32": + for mode in observed_modes: + assert mode == 0o600, f"Temp file mode must be 0600, got {oct(mode)}" + + +def test_load_cert_chain_in_memory_fallback_unlinks_on_load_exception( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Fallback unlinks temp file even when load_cert_chain raises.""" + import ssl + import tempfile as _tempfile + + cert_pem, key_pem = _make_leaf_pem_pair() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + + monkeypatch.delattr(os, "memfd_create", raising=False) + + tmp_paths_created: list[str] = [] + original_mkstemp = _tempfile.mkstemp + + def _spy_mkstemp(*args: object, **kwargs: object) -> tuple[int, str]: + fd, path = original_mkstemp(*args, **kwargs) + tmp_paths_created.append(path) + return fd, path + + monkeypatch.setattr(_tempfile, "mkstemp", _spy_mkstemp) + + # Patch load_cert_chain to always raise. + monkeypatch.setattr( + ctx, "load_cert_chain", lambda *a, **kw: (_ for _ in ()).throw(ssl.SSLError("injected")) + ) + + with pytest.raises(ssl.SSLError): + load_cert_chain_in_memory(ctx, cert_pem, key_pem) + + # Temp file must still be cleaned up. + assert tmp_paths_created, "mkstemp must have been called" + for p in tmp_paths_created: + assert not os.path.exists(p), f"Temp file {p} must be unlinked even after load exception" + + +def test_load_cert_chain_in_memory_fallback_via_proc_oserror( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """When memfd exists but the /proc path is missing (FileNotFoundError), + the helper falls back to mkstemp. + + Real ``/proc``-absent failure surfaces as FileNotFoundError (ENOENT), which + is what the helper catches narrowly — a bare OSError/SSLError must NOT + trigger the disk fallback (see test_..._bad_cert_propagates_without_disk). + """ + import ssl + import tempfile as _tempfile + + cert_pem, key_pem = _make_leaf_pem_pair() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + + if not hasattr(os, "memfd_create"): + pytest.skip("memfd_create not available; fallback-via-/proc path not applicable") + + # Patch load_cert_chain to raise FileNotFoundError on first call (simulating + # /proc not mounted), then succeed on the second (fallback's mkstemp path). + calls: list[int] = [0] + original_load = ctx.__class__.load_cert_chain + + def _raise_once(self: ssl.SSLContext, *args: object, **kwargs: object) -> None: + calls[0] += 1 + if calls[0] == 1: + raise FileNotFoundError("simulated /proc not mounted") + original_load(self, *args, **kwargs) + + monkeypatch.setattr(ssl.SSLContext, "load_cert_chain", _raise_once) + + tmp_paths_created: list[str] = [] + original_mkstemp = _tempfile.mkstemp + + def _spy_mkstemp(*args: object, **kwargs: object) -> tuple[int, str]: + fd, path = original_mkstemp(*args, **kwargs) + tmp_paths_created.append(path) + return fd, path + + monkeypatch.setattr(_tempfile, "mkstemp", _spy_mkstemp) + + load_cert_chain_in_memory(ctx, cert_pem, key_pem) + + assert tmp_paths_created, "Fallback (mkstemp) must trigger when /proc path is FileNotFoundError" + for p in tmp_paths_created: + assert not os.path.exists(p), f"Fallback temp {p} must be unlinked" + + +def test_load_cert_chain_in_memory_bad_cert_propagates_without_disk( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A malformed cert/key (ssl.SSLError, an OSError subclass) must propagate + and NOT silently disk-fall-back via mkstemp.""" + import ssl + import tempfile as _tempfile + + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + + if not hasattr(os, "memfd_create"): + pytest.skip("memfd_create not available; primary path not exercised") + + mkstemp_called = [False] + original_mkstemp = _tempfile.mkstemp + + def _spy_mkstemp(*args: object, **kwargs: object) -> tuple[int, str]: + mkstemp_called[0] = True + return original_mkstemp(*args, **kwargs) + + monkeypatch.setattr(_tempfile, "mkstemp", _spy_mkstemp) + + # Garbage PEM -> load_cert_chain raises ssl.SSLError (subclass of OSError). + with pytest.raises(ssl.SSLError): + load_cert_chain_in_memory(ctx, b"-----BEGIN CERTIFICATE-----\nnope\n", b"not-a-key") + + assert not mkstemp_called[0], "bad cert must NOT trigger the disk fallback" + + +# --------------------------------------------------------------------------- +# ensure_root_ca: corrupt CA key → regenerate (not crash) +# --------------------------------------------------------------------------- + + +def test_ensure_root_ca_corrupt_key_regenerates(tmp_path: Path) -> None: + """Valid cert + corrupt key file → ensure_root_ca regenerates, not raises.""" + # First call creates a valid CA on disk. + _, cert1, key_path, cert_path = ensure_root_ca(base_dir=tmp_path) + + # Overwrite the key with garbage so the parse fails. + key_path.write_bytes( + b"-----BEGIN RSA PRIVATE KEY-----\nGARBAGE\n-----END RSA PRIVATE KEY-----\n" + ) + + # Must not raise; must produce a fresh (different serial) CA. + key2, cert2, _, _ = ensure_root_ca(base_dir=tmp_path) + assert cert2.serial_number != cert1.serial_number, ( + "corrupt key must trigger regeneration, yielding a new cert" + ) + # Returned key must be usable (public_bytes does not raise). + key2.public_key().public_bytes( + serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo + ) + + +# --------------------------------------------------------------------------- +# _assert_perms: skipped on non-POSIX +# --------------------------------------------------------------------------- + + +def test_assert_perms_skipped_on_non_posix(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """On non-POSIX platforms _assert_perms must be a no-op (never raise).""" + p = tmp_path / "file.bin" + p.write_bytes(b"x") + # Monkeypatch os.name inside the module under test. + monkeypatch.setattr("headroom.proxy.agy_ca.os.name", "nt") + # Any expected_mode value; on real POSIX the mode would differ and raise. + _assert_perms(p, 0o600) # must not raise + _assert_perms(p, 0o700) # must not raise + + +# --------------------------------------------------------------------------- +# _write_secure: uses os.replace (atomic cross-platform rename) +# --------------------------------------------------------------------------- + + +def test_write_secure_uses_os_replace(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """_write_secure must call os.replace instead of Path.rename.""" + import headroom.proxy.agy_ca as _mod + + replace_calls: list[tuple[object, object]] = [] + original_replace = os.replace + + def _spy_replace(src: object, dst: object) -> None: + replace_calls.append((src, dst)) + original_replace(src, dst) # type: ignore[arg-type] + + monkeypatch.setattr(_mod.os, "replace", _spy_replace) + + dest = tmp_path / "out.key" + _mod._write_secure(dest, b"hello") + + assert replace_calls, "os.replace must have been called by _write_secure" + assert dest.read_bytes() == b"hello" + + +# --------------------------------------------------------------------------- +# _assert_perms: raises on mismatched mode (line 92) +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX mode bits not enforceable on Windows") +def test_assert_perms_raises_on_wrong_mode(tmp_path: Path) -> None: + p = tmp_path / "wrong-perms.bin" + p.write_bytes(b"x") + p.chmod(0o644) + with pytest.raises(PermissionError, match="Permission check failed"): + _assert_perms(p, 0o600) + + +# --------------------------------------------------------------------------- +# _not_in_os_trust: raises for a path under a trust root (line 139) +# --------------------------------------------------------------------------- + + +def test_not_in_os_trust_raises_for_matching_path( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + trust_root = tmp_path / "trust-root" + monkeypatch.setattr( + "headroom.proxy.agy_ca._OS_TRUST_PATHS", + (str(trust_root),), + ) + bad_path = trust_root / "ca.crt" + with pytest.raises(RuntimeError, match="inside OS trust path"): + _not_in_os_trust(bad_path) + + +# --------------------------------------------------------------------------- +# _is_ca_cert: ExtensionNotFound handling (lines 206-207) +# --------------------------------------------------------------------------- + + +def test_is_ca_cert_false_when_basic_constraints_missing() -> None: + """A cert with no BasicConstraints extension must be treated as non-CA.""" + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + subject = issuer = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "no-bc")]) + now = datetime.datetime.now(datetime.timezone.utc) + 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=365)) + .sign(key, hashes.SHA256()) + ) + assert _is_ca_cert(cert) is False + + +# --------------------------------------------------------------------------- +# _detect_system_bundle: candidate-loop fallback (line 225->223) +# --------------------------------------------------------------------------- + + +def test_detect_system_bundle_skips_missing_candidate( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A missing first candidate is skipped; the next existing one wins.""" + missing = tmp_path / "does-not-exist.crt" + real_bundle = _fake_system_bundle(tmp_path) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(missing), str(real_bundle)), + ) + result = _detect_system_bundle() + assert result == real_bundle + + +# --------------------------------------------------------------------------- +# _windows_trust_pem: a single bad-DER entry is skipped, not fatal (lines 248-250) +# --------------------------------------------------------------------------- + + +def test_windows_trust_pem_skips_bad_der_entry(monkeypatch: pytest.MonkeyPatch) -> None: + import ssl + + ca_pem = _make_cert(is_ca=True) + ca_cert = x509.load_pem_x509_certificate(ca_pem) + ca_der = ca_cert.public_bytes(serialization.Encoding.DER) + bad_der = b"not-a-real-der-cert" + + def fake_enum(store: str) -> list[tuple[bytes, str, bool]]: + if store == "ROOT": + return [(bad_der, "x509_asn", True), (ca_der, "x509_asn", True)] + return [] + + monkeypatch.setattr("ssl.enum_certificates", fake_enum, raising=False) + + original_der_to_pem = ssl.DER_cert_to_PEM_cert + + def fake_der_to_pem(der: bytes) -> str: + if der == bad_der: + raise ValueError("simulated malformed DER") + return original_der_to_pem(der) + + monkeypatch.setattr("ssl.DER_cert_to_PEM_cert", fake_der_to_pem) + + result = _windows_trust_pem() + marker = b"-----BEGIN CERTIFICATE-----" + present = { + x509.load_pem_x509_certificate(marker + block).serial_number + for block in result.split(marker)[1:] + } + assert ca_cert.serial_number in present, "a sibling bad-DER entry must not drop the good cert" + + +# --------------------------------------------------------------------------- +# _system_trust_pem: falls back to the Windows cert store (line 280) +# --------------------------------------------------------------------------- + + +def test_system_trust_pem_falls_back_to_windows_store( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr("headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", ()) + monkeypatch.setattr(sys, "platform", "win32") + + ca_pem = _make_cert(is_ca=True) + ca_cert = x509.load_pem_x509_certificate(ca_pem) + ca_der = ca_cert.public_bytes(serialization.Encoding.DER) + + def fake_enum(store: str) -> list[tuple[bytes, str, bool]]: + return [(ca_der, "x509_asn", True)] if store == "ROOT" else [] + + monkeypatch.setattr("ssl.enum_certificates", fake_enum, raising=False) + + pem_bytes, source = _system_trust_pem() + assert source == "windows-cert-store" + marker = b"-----BEGIN CERTIFICATE-----" + present = { + x509.load_pem_x509_certificate(marker + block).serial_number + for block in pem_bytes.split(marker)[1:] + } + assert ca_cert.serial_number in present + + +# --------------------------------------------------------------------------- +# _parse_ca_certs_from_pem: PEM block missing its END marker is skipped (line 296) +# --------------------------------------------------------------------------- + + +def test_parse_skips_pem_block_missing_end_marker() -> None: + good_pem = _make_cert(is_ca=True) + truncated = b"-----BEGIN CERTIFICATE-----\nMIIBnotcompletenoendmarkerhere\n" + combined = truncated + good_pem + results = _parse_ca_certs_from_pem(combined) + assert len(results) == 1 + cert = x509.load_pem_x509_certificate(results[0]) + assert _is_ca_cert(cert) is True + + +# --------------------------------------------------------------------------- +# ensure_root_ca / build_combined_bundle: Path.home() default (lines 367 & 461) +# --------------------------------------------------------------------------- + + +def test_ensure_root_ca_defaults_to_home_dir( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(Path, "home", lambda: tmp_path) + _, cert, key_path, cert_path = ensure_root_ca() + assert key_path == tmp_path / ".headroom" / "ca" / _CA_KEY_NAME + assert cert_path.exists() + assert _is_ca_cert(cert) + + +def test_build_combined_bundle_defaults_to_home_dir( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + sys_bundle = _fake_system_bundle(tmp_path) + monkeypatch.setattr(Path, "home", lambda: tmp_path) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(sys_bundle),), + ) + bundle_path = build_combined_bundle(corp_env_vars=()) + assert bundle_path == tmp_path / ".headroom" / _BUNDLE_NAME + assert bundle_path.exists() + + +# --------------------------------------------------------------------------- +# ensure_root_ca: corrupt CERT (not key) → regenerate (lines 383-385) +# --------------------------------------------------------------------------- + + +def test_ensure_root_ca_corrupt_cert_regenerates(tmp_path: Path) -> None: + """Valid key + corrupt cert file → ensure_root_ca regenerates, not raises.""" + _, cert1, key_path, cert_path = ensure_root_ca(base_dir=tmp_path) + + cert_path.write_bytes(b"-----BEGIN CERTIFICATE-----\nGARBAGE\n-----END CERTIFICATE-----\n") + cert_path.chmod(0o600) + + _, cert2, _, _ = ensure_root_ca(base_dir=tmp_path) + assert cert2.serial_number != cert1.serial_number, ( + "corrupt cert must trigger regeneration (event=ca_load_failed), yielding a new cert" + ) + + +# --------------------------------------------------------------------------- +# build_combined_bundle: missing trailing newline on system bundle (line 474) +# --------------------------------------------------------------------------- + + +def test_bundle_appends_newline_when_system_pem_missing_trailing_newline( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + sys_pem_no_nl = _make_cert(is_ca=True).rstrip(b"\n") + assert not sys_pem_no_nl.endswith(b"\n") + sys_bundle = tmp_path / "system-ca-bundle.pem" + sys_bundle.write_bytes(sys_pem_no_nl) + monkeypatch.setattr( + "headroom.proxy.agy_ca._SYSTEM_BUNDLE_CANDIDATES", + (str(sys_bundle),), + ) + bundle_path = build_combined_bundle(base_dir=tmp_path, corp_env_vars=()) + bundle_data = bundle_path.read_bytes() + assert bundle_data.startswith(sys_pem_no_nl + b"\n"), ( + "a missing trailing newline on the system bundle must be appended exactly once" + ) + + +# --------------------------------------------------------------------------- +# _write_all_fd: os.write returning 0 bytes raises OSError (line 565) +# --------------------------------------------------------------------------- + + +def test_write_all_fd_raises_on_zero_byte_write(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(os, "write", lambda fd, data: 0) + with pytest.raises(OSError, match="wrote 0 bytes"): + _write_all_fd(0, b"some bytes to write") + + +# --------------------------------------------------------------------------- +# _load_via_mkstemp: cleanup errors in the finally block are swallowed +# (lines 586-589 & 592-593) +# --------------------------------------------------------------------------- + + +def test_load_via_mkstemp_close_oserror_is_swallowed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """If the write fails (fd never reaches -1) and the finally-block os.close + also raises, the close OSError must be swallowed and the original write + failure must be the one that propagates.""" + import ssl + import tempfile as _tempfile + + cert_pem, key_pem = _make_leaf_pem_pair() + combined = cert_pem + key_pem + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + + captured_fd: list[int] = [] + original_mkstemp = _tempfile.mkstemp + + def _spy_mkstemp(*args: object, **kwargs: object) -> tuple[int, str]: + fd, path = original_mkstemp(*args, **kwargs) + captured_fd.append(fd) + return fd, path + + monkeypatch.setattr(_tempfile, "mkstemp", _spy_mkstemp) + + original_write = os.write + + def _fail_write(fd: int, data: bytes | bytearray) -> int: + if captured_fd and fd == captured_fd[0]: + return 0 + return original_write(fd, data) + + monkeypatch.setattr(os, "write", _fail_write) + + original_close = os.close + close_attempts: list[int] = [] + + def _fail_close(fd: int) -> None: + if captured_fd and fd == captured_fd[0]: + close_attempts.append(fd) + raise OSError("simulated close failure") + original_close(fd) + + monkeypatch.setattr(os, "close", _fail_close) + + with pytest.raises(OSError, match="wrote 0 bytes"): + _load_via_mkstemp(ctx, combined) + + assert close_attempts, "os.close must have been attempted in the finally block" + + +def test_load_via_mkstemp_unlink_oserror_is_swallowed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A failing os.unlink in the cleanup finally block must not propagate.""" + import ssl + import tempfile as _tempfile + + cert_pem, key_pem = _make_leaf_pem_pair() + combined = cert_pem + key_pem + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + + captured_paths: list[str] = [] + original_mkstemp = _tempfile.mkstemp + + def _spy_mkstemp(*args: object, **kwargs: object) -> tuple[int, str]: + fd, path = original_mkstemp(*args, **kwargs) + captured_paths.append(path) + return fd, path + + monkeypatch.setattr(_tempfile, "mkstemp", _spy_mkstemp) + + def _fail_unlink(path: str, *args: object, **kwargs: object) -> None: + raise OSError("simulated unlink failure") + + monkeypatch.setattr(os, "unlink", _fail_unlink) + + # Must not raise despite the unlink failure being swallowed. + _load_via_mkstemp(ctx, combined) + + # Manual cleanup: our fake unlink prevented real removal of the temp file. + monkeypatch.undo() + for p in captured_paths: + if os.path.exists(p): + os.unlink(p) diff --git a/tests/test_agy_ccr_downgrade_warning.py b/tests/test_agy_ccr_downgrade_warning.py new file mode 100644 index 000000000..2bd59865f --- /dev/null +++ b/tests/test_agy_ccr_downgrade_warning.py @@ -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//); 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] diff --git a/tests/test_agy_ccr_retrieve_roundtrip.py b/tests/test_agy_ccr_retrieve_roundtrip.py new file mode 100644 index 000000000..9e0661f32 --- /dev/null +++ b/tests/test_agy_ccr_retrieve_roundtrip.py @@ -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=]`` 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 diff --git a/tests/test_agy_dispatch.py b/tests/test_agy_dispatch.py new file mode 100644 index 000000000..9f050666d --- /dev/null +++ b/tests/test_agy_dispatch.py @@ -0,0 +1,1444 @@ +"""Tests for headroom.proxy.agy_dispatch.AgyDispatchServer. + +All tests use ephemeral loopback ports; ~/.headroom is never touched. +The upstream Gemini/CloudCode network is mocked via monkeypatching +HeadroomProxy._stream_response so no real network calls are made. + +Test coverage: + (a) TLS client (verifying against root CA, SNI=daily-cloudcode-pa.googleapis.com) + connects to hypercorn port, POSTs /v1internal:streamGenerateContent, gets 200. + (b) ALPN negotiates h2. + (c) End-to-end: agy-side CONNECT terminator → tunnel → hypercorn → app → 200. + (d) Authorization + x-goog-api-key NOT present in headroom logs (caplog). + (e) All pre-existing T8 terminator tests still pass (those remain in + test_agy_terminator.py; this file covers the dispatch-server side only). +""" + +from __future__ import annotations + +import asyncio +import datetime +import json +import logging +import ssl +from typing import Any + +import pytest +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 + +from headroom.proxy.agy_dispatch import AgyDispatchServer, make_host_guard +from headroom.proxy.agy_terminator import DEFAULT_ALLOWLIST, AgyCONNECTTerminator + +# --------------------------------------------------------------------------- +# CA fixture +# --------------------------------------------------------------------------- + +ALLOWLIST_HOST = "daily-cloudcode-pa.googleapis.com" + +_AGY_REQUEST_BODY = json.dumps( + { + "model": "gemini-2.5-pro", + "request": { + "contents": [{"role": "user", "parts": [{"text": "ping"}]}], + }, + } +).encode() + + +def _make_test_ca() -> tuple[RSAPrivateKey, Certificate, bytes]: + """Generate 2048-bit RSA root CA (never touches disk).""" + key: RSAPrivateKey = rsa.generate_private_key(public_exponent=65537, key_size=2048) + now = datetime.datetime.now(tz=datetime.timezone.utc) + subject = issuer = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "Headroom Test CA")]) + 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=365)) + .add_extension(x509.BasicConstraints(ca=True, path_length=0), critical=True) + .sign(key, hashes.SHA256()) + ) + cert_pem = cert.public_bytes(serialization.Encoding.PEM) + return key, cert, cert_pem + + +@pytest.fixture +def tmp_ca() -> tuple[RSAPrivateKey, Certificate, bytes]: + return _make_test_ca() + + +# --------------------------------------------------------------------------- +# SSL context helpers +# --------------------------------------------------------------------------- + + +def _build_client_ssl_ctx(ca_cert_pem: bytes) -> ssl.SSLContext: + """Build a TLS client context that trusts only the test CA.""" + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ctx.check_hostname = True + ctx.verify_mode = ssl.CERT_REQUIRED + ctx.load_verify_locations(cadata=ca_cert_pem.decode("ascii")) + return ctx + + +# --------------------------------------------------------------------------- +# Mock helpers +# --------------------------------------------------------------------------- + + +def _make_sse_mock_response() -> bytes: + """Minimal SSE response payload that handle_google_cloudcode_stream can relay.""" + lines = [ + b'data: {"candidates":[{"content":{"parts":[{"text":"pong"}]}}]}\r\n', + b"\r\n", + b"data: [DONE]\r\n", + b"\r\n", + ] + return b"".join(lines) + + +# --------------------------------------------------------------------------- +# Tests: (a) + (b) direct TLS → dispatch server → 200 + h2 +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dispatch_server_tls_and_route( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(a) TLS client verifying against root CA connects to hypercorn port, + POSTs /v1internal:streamGenerateContent, gets 200. + No real upstream network: _stream_response is monkeypatched. + """ + ca_key, ca_cert, ca_cert_pem = tmp_ca + + # Patch HeadroomProxy._stream_response so no upstream call is made. + from fastapi.responses import StreamingResponse + + from headroom.proxy.server import HeadroomProxy + + async def _fake_stream(self: Any, *args: Any, **kwargs: Any) -> StreamingResponse: + async def _body() -> bytes: + yield b'data: {"candidates":[]}\n\ndata: [DONE]\n\n' + + return StreamingResponse( + _body(), + status_code=200, + media_type="text/event-stream", + ) + + monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream) + + async with AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) as srv: + host, port = srv.address + assert host == "127.0.0.1" + assert port > 0 + + # Build an HTTPS client that trusts the test CA. + ssl_ctx = _build_client_ssl_ctx(ca_cert_pem) + # Use HTTP/1.1 for the direct request (simpler to compose manually). + ssl_ctx.set_alpn_protocols(["http/1.1"]) + + conn_reader, conn_writer = await asyncio.open_connection( + "127.0.0.1", + port, + ssl=ssl_ctx, + server_hostname=ALLOWLIST_HOST, + ) + try: + body = _AGY_REQUEST_BODY + request = ( + f"POST /v1internal:streamGenerateContent HTTP/1.1\r\n" + f"Host: {ALLOWLIST_HOST}\r\n" + f"Content-Type: application/json\r\n" + f"Content-Length: {len(body)}\r\n" + f"\r\n" + ).encode() + body + + conn_writer.write(request) + await conn_writer.drain() + + # Read enough response to confirm 200. + response_line = await asyncio.wait_for(conn_reader.readline(), timeout=10.0) + assert b"200" in response_line, f"Expected 200, got {response_line!r}" + finally: + conn_writer.close() + try: + await conn_writer.wait_closed() + except Exception: # noqa: BLE001 + pass + + +@pytest.mark.asyncio +async def test_dispatch_server_alpn_h2( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(b) ALPN negotiates h2 when client offers ["h2", "http/1.1"].""" + ca_key, ca_cert, ca_cert_pem = tmp_ca + + from fastapi.responses import StreamingResponse + + from headroom.proxy.server import HeadroomProxy + + async def _fake_stream(self: Any, *args: Any, **kwargs: Any) -> StreamingResponse: + async def _body() -> bytes: + yield b"data: [DONE]\n\n" + + return StreamingResponse(_body(), status_code=200, media_type="text/event-stream") + + monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream) + + async with AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) as srv: + _, port = srv.address + + ssl_ctx = _build_client_ssl_ctx(ca_cert_pem) + ssl_ctx.set_alpn_protocols(["h2", "http/1.1"]) + + conn_reader, conn_writer = await asyncio.open_connection( + "127.0.0.1", + port, + ssl=ssl_ctx, + server_hostname=ALLOWLIST_HOST, + ) + try: + ssl_obj = conn_writer.get_extra_info("ssl_object") + alpn = ssl_obj.selected_alpn_protocol() if ssl_obj else None + assert alpn == "h2", f"Expected h2 ALPN, got {alpn!r}" + finally: + conn_writer.close() + try: + await conn_writer.wait_closed() + except Exception: # noqa: BLE001 + pass + + +# --------------------------------------------------------------------------- +# Test: (c) end-to-end via terminator → tunnel → hypercorn → 200 +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_terminator_tunnel_to_dispatch_server( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(c) agy-side CONNECT terminator → byte-splice tunnel → hypercorn → app → 200.""" + ca_key, ca_cert, ca_cert_pem = tmp_ca + + from fastapi.responses import StreamingResponse + + from headroom.proxy.server import HeadroomProxy + + async def _fake_stream(self: Any, *args: Any, **kwargs: Any) -> StreamingResponse: + async def _body() -> bytes: + yield b"data: [DONE]\n\n" + + return StreamingResponse(_body(), status_code=200, media_type="text/event-stream") + + monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream) + + async with AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) as dispatch_srv: + _, dispatch_port = dispatch_srv.address + + async with AgyCONNECTTerminator( + allowlist=DEFAULT_ALLOWLIST, + ca_key=ca_key, + ca_cert=ca_cert, + dispatch_port=dispatch_port, + ) as terminator: + proxy_host, proxy_port = terminator.address + + # Step 1: TCP CONNECT to terminator. + raw_reader, raw_writer = await asyncio.open_connection(proxy_host, proxy_port) + raw_writer.write( + f"CONNECT {ALLOWLIST_HOST}:443 HTTP/1.1\r\n" + f"Host: {ALLOWLIST_HOST}:443\r\n" + "\r\n".encode() + ) + await raw_writer.drain() + resp = await asyncio.wait_for(raw_reader.readline(), timeout=5.0) + assert b"200" in resp, f"Expected 200 tunnel ACK, got {resp!r}" + + # Step 2: TLS handshake over the tunnel (to hypercorn's SNI cert). + ssl_ctx = _build_client_ssl_ctx(ca_cert_pem) + ssl_ctx.set_alpn_protocols(["http/1.1"]) + loop = asyncio.get_event_loop() + raw_writer.transport.pause_reading() + tls_transport = await asyncio.wait_for( + loop.start_tls( + raw_writer.transport, + raw_writer.transport.get_protocol(), + ssl_ctx, + server_hostname=ALLOWLIST_HOST, + ), + timeout=10.0, + ) + + # Step 3: POST through TLS tunnel and check 200. + # Re-wrap tls_transport in a StreamReader so we can readline(). + tls_reader = asyncio.StreamReader() + tls_proto = asyncio.StreamReaderProtocol(tls_reader) + tls_transport.set_protocol(tls_proto) + tls_proto.connection_made(tls_transport) + + body = _AGY_REQUEST_BODY + tls_transport.write( + ( + f"POST /v1internal:streamGenerateContent HTTP/1.1\r\n" + f"Host: {ALLOWLIST_HOST}\r\n" + f"Content-Type: application/json\r\n" + f"Content-Length: {len(body)}\r\n" + f"\r\n" + ).encode() + + body + ) + + # Read the HTTP status line through the terminator tunnel. + status_line = await asyncio.wait_for(tls_reader.readline(), timeout=10.0) + assert b"200" in status_line, ( + f"Expected HTTP 200 through CONNECT tunnel, got: {status_line!r}" + ) + + tls_transport.close() + + +# --------------------------------------------------------------------------- +# Test: (d) Authorization + x-goog-api-key NOT in headroom logs +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_secret_headers_not_logged( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """(d) Authorization and x-goog-api-key must not appear in headroom logs.""" + ca_key, ca_cert, ca_cert_pem = tmp_ca + + from fastapi.responses import StreamingResponse + + from headroom.proxy.server import HeadroomProxy + + async def _fake_stream(self: Any, *args: Any, **kwargs: Any) -> StreamingResponse: + async def _body() -> bytes: + yield b"data: [DONE]\n\n" + + return StreamingResponse(_body(), status_code=200, media_type="text/event-stream") + + monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream) + + with caplog.at_level(logging.DEBUG, logger="headroom"): + async with AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) as srv: + _, port = srv.address + + ssl_ctx = _build_client_ssl_ctx(ca_cert_pem) + ssl_ctx.set_alpn_protocols(["http/1.1"]) + + conn_reader, conn_writer = await asyncio.open_connection( + "127.0.0.1", + port, + ssl=ssl_ctx, + server_hostname=ALLOWLIST_HOST, + ) + try: + body = _AGY_REQUEST_BODY + secret_auth = "Bearer supersecret-token-xyz" + secret_api_key = "AIzaSySecret1234" + request = ( + f"POST /v1internal:streamGenerateContent HTTP/1.1\r\n" + f"Host: {ALLOWLIST_HOST}\r\n" + f"Authorization: {secret_auth}\r\n" + f"x-goog-api-key: {secret_api_key}\r\n" + f"Content-Type: application/json\r\n" + f"Content-Length: {len(body)}\r\n" + f"\r\n" + ).encode() + body + + conn_writer.write(request) + await conn_writer.drain() + # Read enough to let the handler log. + await asyncio.wait_for(conn_reader.readline(), timeout=10.0) + finally: + conn_writer.close() + try: + await conn_writer.wait_closed() + except Exception: # noqa: BLE001 + pass + + log_text = "\n".join(r.getMessage() for r in caplog.records) + # Default path (log_outbound_headers) only logs counts, never values. + # Assert no header VALUE leaks into any log record on the default path. + assert "supersecret-token-xyz" not in log_text, "Bearer token leaked into headroom logs" + assert "AIzaSySecret1234" not in log_text, "x-goog-api-key leaked into headroom logs" + + +# --------------------------------------------------------------------------- +# Test: (d2) redaction unit — _should_redact_key / redact_for_wire_debug +# --------------------------------------------------------------------------- + + +def test_redaction_is_load_bearing() -> None: + """Redaction of authorization and x-goog-api-key is structurally enforced. + + This test is deliberately coupled to _should_redact_key and + redact_for_wire_debug so that removing or weakening either function + causes a failure here, making this a load-bearing regression guard. + """ + from headroom.proxy.helpers import ( + _CODEX_WIRE_REDACTED, + _should_redact_key, + redact_for_wire_debug, + ) + + # 1. _should_redact_key must flag both sensitive header names. + assert _should_redact_key("authorization"), "authorization must be redacted" + assert _should_redact_key("Authorization"), "Authorization (mixed case) must be redacted" + assert _should_redact_key("x-goog-api-key"), "x-goog-api-key must be redacted" + assert _should_redact_key("X-Goog-Api-Key"), "X-Goog-Api-Key (mixed case) must be redacted" + + # 2. redact_for_wire_debug must replace values with _CODEX_WIRE_REDACTED. + secret_auth = "Bearer supersecret-token-xyz" + secret_api_key = "AIzaSySecret1234" + headers = { + "authorization": secret_auth, + "x-goog-api-key": secret_api_key, + "content-type": "application/json", + } + redacted = redact_for_wire_debug(headers) + assert redacted["authorization"] == _CODEX_WIRE_REDACTED, ( + f"authorization must be {_CODEX_WIRE_REDACTED!r}, got {redacted['authorization']!r}" + ) + assert redacted["x-goog-api-key"] == _CODEX_WIRE_REDACTED, ( + f"x-goog-api-key must be {_CODEX_WIRE_REDACTED!r}, got {redacted['x-goog-api-key']!r}" + ) + # Non-secret headers must pass through unchanged. + assert redacted["content-type"] == "application/json" + + # 3. Secret VALUES must not appear in the redacted output at all. + import json as _json + + redacted_str = _json.dumps(redacted) + assert secret_auth not in redacted_str, "Bearer token survived redact_for_wire_debug" + assert secret_api_key not in redacted_str, "API key survived redact_for_wire_debug" + + +# --------------------------------------------------------------------------- +# Tests: dispatch server lifecycle +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dispatch_server_loopback_only( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """AgyDispatchServer binds 127.0.0.1 only.""" + ca_key, ca_cert, _ = tmp_ca + async with AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) as srv: + host, port = srv.address + assert host == "127.0.0.1" + assert port > 0 + + +@pytest.mark.asyncio +async def test_dispatch_server_start_stop_idempotent( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """stop() after stop() does not raise.""" + ca_key, ca_cert, _ = tmp_ca + srv = AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) + await srv.start() + await srv.stop() + await srv.stop() # idempotent + + +@pytest.mark.asyncio +async def test_dispatch_server_ensure_root_ca_fallback(tmp_path: Any) -> None: + """start() with base_dir=tmp_path and NO injected CA falls back to + ensure_root_ca (local key-gen only, no network): CA key/cert are written + under base_dir/ca, and the server's leaf cache is populated from them.""" + srv = AgyDispatchServer(base_dir=tmp_path) + await srv.start() + try: + ca_key_path = tmp_path / "ca" / "ca.key" + ca_cert_path = tmp_path / "ca" / "ca.crt" + assert ca_key_path.exists(), "ensure_root_ca fallback must generate ca.key on disk" + assert ca_cert_path.exists(), "ensure_root_ca fallback must generate ca.crt on disk" + assert srv._leaf_cache is not None, "leaf cache must be populated after start()" + finally: + await srv.stop() + + +@pytest.mark.asyncio +async def test_dispatch_server_start_reraises_lifespan_startup_failure( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """start() re-raises when the hypercorn lifespan task fails during startup. + + We stub hypercorn.asyncio.run.Lifespan (imported locally inside start()) + with a fake whose handle_lifespan() signals startup THEN raises in the + same task step (no intervening await), so the task is deterministically + already done-with-exception by the time start() checks + `self._lifespan_task.done()`. + """ + import hypercorn.asyncio.run as hypercorn_run + + ca_key, ca_cert, _ = tmp_ca + + class _FailingLifespan: + def __init__(self, app: Any, config: Any, loop: Any, lifespan_state: Any) -> None: + self.startup = asyncio.Event() + self.shutdown = asyncio.Event() + + async def handle_lifespan(self) -> None: + self.startup.set() + raise RuntimeError("injected lifespan startup failure") + + async def wait_for_startup(self) -> None: + await self.startup.wait() + + async def wait_for_shutdown(self) -> None: + await self.shutdown.wait() + + monkeypatch.setattr(hypercorn_run, "Lifespan", _FailingLifespan) + + srv = AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) + with pytest.raises(RuntimeError, match="injected lifespan startup failure"): + await srv.start() + + +@pytest.mark.asyncio +async def test_dispatch_server_windows_so_exclusiveaddruse( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Non-POSIX branch: SO_EXCLUSIVEADDRUSE sockopt is set instead of SO_REUSEADDR. + + os.name must read as non-"posix" ONLY for the `if os.name == "posix":` + check inside agy_dispatch.start() — globally patching the real `os` + module's `.name` breaks unrelated code paths reached from start() + (e.g. create_app() -> Path.home(), which selects WindowsPath/PosixPath + from the live os.name and blows up on a non-Windows filesystem). So we + replace the `os` symbol bound inside the agy_dispatch module with a thin + proxy that fakes only `.name`, delegating everything else to the real + `os` module. Real asyncio.start_server is also replaced with a fake so + the forced-name window does not span any real event-loop/transport + internals. + """ + import os as os_mod + import socket as socket_mod + + import headroom.proxy.agy_dispatch as agy_dispatch_mod + + ca_key, ca_cert, _ = tmp_ca + + class _FakeOSName: + """Proxies the real `os` module except `.name`, which reads "nt".""" + + def __getattr__(self, attr: str) -> Any: + if attr == "name": + return "nt" + return getattr(os_mod, attr) + + # Inject SO_EXCLUSIVEADDRUSE on platforms (e.g. Linux) that lack it, + # aliased to a real, valid sockopt so the actual setsockopt() call succeeds. + monkeypatch.setattr(socket_mod, "SO_EXCLUSIVEADDRUSE", socket_mod.SO_REUSEADDR, raising=False) + + setsockopt_calls: list[tuple[int, int]] = [] + original_setsockopt = socket_mod.socket.setsockopt + + def _spy_setsockopt( + self: socket_mod.socket, level: int, optname: int, value: Any, *a: Any, **kw: Any + ) -> Any: + setsockopt_calls.append((level, optname)) + return original_setsockopt(self, level, optname, value, *a, **kw) + + monkeypatch.setattr(socket_mod.socket, "setsockopt", _spy_setsockopt) + + class _FakeSocketInfo: + def getsockname(self) -> tuple[str, int]: + return ("127.0.0.1", 54321) + + captured_socks: list[socket_mod.socket] = [] + + class _FakeServer: + sockets = [_FakeSocketInfo()] + + def close(self) -> None: + pass + + async def wait_closed(self) -> None: + pass + + async def _fake_start_server(*args: Any, **kwargs: Any) -> _FakeServer: + sock = kwargs.get("sock") + if sock is not None: + captured_socks.append(sock) + return _FakeServer() + + monkeypatch.setattr(asyncio, "start_server", _fake_start_server) + monkeypatch.setattr(agy_dispatch_mod, "os", _FakeOSName()) + + srv = AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) + try: + await srv.start() + assert (socket_mod.SOL_SOCKET, socket_mod.SO_EXCLUSIVEADDRUSE) in setsockopt_calls, ( + "SO_EXCLUSIVEADDRUSE setsockopt must be called when os.name != 'posix'" + ) + finally: + await srv.stop() + for s in captured_socks: + s.close() + + +@pytest.mark.asyncio +async def test_dispatch_server_stop_swallows_lifespan_shutdown_exception( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """stop() swallows an exception raised by lifespan.wait_for_shutdown() + (e.g. LifespanTimeoutError) instead of letting it escape.""" + ca_key, ca_cert, _ = tmp_ca + srv = AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) + await srv.start() + + class _BoomShutdown: + async def wait_for_shutdown(self) -> None: + raise RuntimeError("injected shutdown failure") + + srv._lifespan = _BoomShutdown() # type: ignore[assignment] + + await srv.stop() # must not raise despite the injected RuntimeError + + assert srv._lifespan is None, "stop() must clear _lifespan even after a swallowed exception" + + +@pytest.mark.asyncio +async def test_dispatch_server_stop_swallows_task_cancel_exception( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """stop() swallows a non-CancelledError exception raised while awaiting + the cancelled lifespan task.""" + ca_key, ca_cert, _ = tmp_ca + srv = AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) + await srv.start() + + async def _stubborn() -> None: + try: + await asyncio.sleep(30) + except asyncio.CancelledError: + raise RuntimeError("injected cancel-time failure") from None + + loop = asyncio.get_event_loop() + stubborn_task = loop.create_task(_stubborn()) + await asyncio.sleep(0) # let it start awaiting sleep(30) before we swap it in + + real_lifespan_task = srv._lifespan_task + srv._lifespan_task = stubborn_task + + try: + await srv.stop() # must not raise despite the injected RuntimeError + assert srv._lifespan_task is None, "stop() must clear _lifespan_task after swallowing" + finally: + # Clean up the real (now-orphaned) lifespan task so it is not left pending. + if real_lifespan_task is not None and not real_lifespan_task.done(): + real_lifespan_task.cancel() + try: + await real_lifespan_task + except (asyncio.CancelledError, Exception): # noqa: BLE001 + pass + + +def test_dispatch_server_address_raises_before_start( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """address property raises RuntimeError before start().""" + ca_key, ca_cert, _ = tmp_ca + srv = AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) + with pytest.raises(RuntimeError, match="not started"): + _ = srv.address + + +# --------------------------------------------------------------------------- +# Tests: SNI allowlist guard (headroom-oqb.1) +# --------------------------------------------------------------------------- + +_ATTACKER_HOST = "evilcloudcode-pa.googleapis.com" +_CONTROLLED_HOST = "allowed.test" +_CONTROLLED_ALLOWLIST: frozenset[str] = frozenset({_CONTROLLED_HOST}) + + +async def _try_tls_connect( + port: int, + ca_cert_pem: bytes, + server_hostname: str | None, + *, + timeout: float = 5.0, +) -> bool: + """Return True if TLS handshake succeeds, False if it fails with an SSL error.""" + ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ssl_ctx.check_hostname = server_hostname is not None + ssl_ctx.verify_mode = ssl.CERT_REQUIRED if server_hostname is not None else ssl.CERT_NONE + ssl_ctx.load_verify_locations(cadata=ca_cert_pem.decode("ascii")) + try: + _, writer = await asyncio.wait_for( + asyncio.open_connection( + "127.0.0.1", + port, + ssl=ssl_ctx, + server_hostname=server_hostname, + ), + timeout=timeout, + ) + writer.close() + try: + await writer.wait_closed() + except Exception: # noqa: BLE001 + pass + return True + except (ssl.SSLError, OSError, ConnectionResetError, asyncio.TimeoutError): + return False + + +@pytest.mark.asyncio +async def test_sni_allowlisted_still_routes( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Allowlisted SNI completes handshake (no regression).""" + ca_key, ca_cert, ca_cert_pem = tmp_ca + + from fastapi.responses import StreamingResponse + + from headroom.proxy.server import HeadroomProxy + + async def _fake_stream(self: Any, *args: Any, **kwargs: Any) -> StreamingResponse: + async def _body() -> bytes: + yield b"data: [DONE]\n\n" + + return StreamingResponse(_body(), status_code=200, media_type="text/event-stream") + + monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream) + + async with AgyDispatchServer( + ca_key=ca_key, ca_cert=ca_cert, allowlist=_CONTROLLED_ALLOWLIST + ) as srv: + _, port = srv.address + success = await _try_tls_connect(port, ca_cert_pem, _CONTROLLED_HOST) + assert success, "Allowlisted SNI must complete handshake" + + +@pytest.mark.asyncio +async def test_sni_non_allowlisted_rejected( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """Non-allowlisted SNI -> handshake aborts; server stays alive; + get_or_mint NOT called for attacker host; event=sni_refused is logged.""" + from unittest.mock import patch + + from headroom.proxy.agy_terminator import _LeafCache + + ca_key, ca_cert, ca_cert_pem = tmp_ca + + # Capture log records via a handler installed before the server starts. + # caplog cannot reliably capture records from SSL C-level callbacks, so + # we install a custom handler directly on the module logger. + sni_log_records: list[logging.LogRecord] = [] + + class _Capture(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + sni_log_records.append(record) + + capture_handler = _Capture(logging.WARNING) + _dispatch_logger = logging.getLogger("headroom.proxy.agy_dispatch") + _dispatch_logger.addHandler(capture_handler) + + try: + async with AgyDispatchServer( + ca_key=ca_key, ca_cert=ca_cert, allowlist=_CONTROLLED_ALLOWLIST + ) as srv: + _, port = srv.address + + # Spy on _LeafCache.get_or_mint to verify attacker host is never minted. + original_get_or_mint = _LeafCache.get_or_mint + call_hostnames: list[str] = [] + + def _spy_get_or_mint(self: _LeafCache, host: str, *args: Any, **kwargs: Any) -> Any: + call_hostnames.append(host) + return original_get_or_mint(self, host, *args, **kwargs) + + with patch.object(_LeafCache, "get_or_mint", _spy_get_or_mint): + rejected = not await _try_tls_connect(port, ca_cert_pem, _ATTACKER_HOST) + + # Server must still respond to further connections. + assert srv._server is not None, "Server must stay alive after rejected SNI" + finally: + _dispatch_logger.removeHandler(capture_handler) + + assert rejected, "Non-allowlisted SNI must abort handshake" + attacker_mints = [h for h in call_hostnames if h == _ATTACKER_HOST] + assert attacker_mints == [], f"get_or_mint called for attacker host: {attacker_mints}" + + # Attacker host must be absent from the leaf cache. + assert srv._leaf_cache is not None + assert _ATTACKER_HOST not in srv._leaf_cache._cache, ( + "Attacker host must not appear in leaf cache" + ) + + warned = any("event=sni_refused" in r.getMessage() for r in sni_log_records) + assert warned, ( + f"Expected event=sni_refused WARNING; got records: " + f"{[r.getMessage() for r in sni_log_records]}" + ) + + +@pytest.mark.asyncio +async def test_sni_named_attack_evilcloudcode_rejected( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """Named attack: SNI 'evilcloudcode-pa.googleapis.com' is rejected.""" + ca_key, ca_cert, ca_cert_pem = tmp_ca + + async with AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) as srv: + _, port = srv.address + rejected = not await _try_tls_connect(port, ca_cert_pem, _ATTACKER_HOST) + + assert rejected, "evilcloudcode-pa.googleapis.com must be rejected by SNI guard" + + +@pytest.mark.asyncio +async def test_sni_placeholder_headroom_internal_rejected( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """Explicit wire SNI 'headroom.internal' (the placeholder) is rejected.""" + ca_key, ca_cert, ca_cert_pem = tmp_ca + + # Use a handler to verify event=sni_refused is logged (caplog is unreliable + # in SSL C-level callbacks). + sni_log_records: list[logging.LogRecord] = [] + + class _Capture(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + sni_log_records.append(record) + + capture_handler = _Capture(logging.WARNING) + _dispatch_logger = logging.getLogger("headroom.proxy.agy_dispatch") + _dispatch_logger.addHandler(capture_handler) + + try: + async with AgyDispatchServer( + ca_key=ca_key, ca_cert=ca_cert, allowlist=_CONTROLLED_ALLOWLIST + ) as srv: + _, port = srv.address + rejected = not await _try_tls_connect(port, ca_cert_pem, "headroom.internal") + finally: + _dispatch_logger.removeHandler(capture_handler) + + assert rejected, "headroom.internal must be rejected (not in allowlist)" + warned = any("event=sni_refused" in r.getMessage() for r in sni_log_records) + assert warned, ( + f"Expected event=sni_refused WARNING for headroom.internal; " + f"got: {[r.getMessage() for r in sni_log_records]}" + ) + + +@pytest.mark.asyncio +async def test_sni_none_and_empty_rejected( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """None SNI and empty-string SNI are rejected; no headroom.internal leaf served.""" + ca_key, ca_cert, ca_cert_pem = tmp_ca + + # Capture log records via a handler (caplog is unreliable in SSL C-level callbacks). + sni_log_records: list[logging.LogRecord] = [] + + class _Capture(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + sni_log_records.append(record) + + capture_handler = _Capture(logging.WARNING) + _dispatch_logger = logging.getLogger("headroom.proxy.agy_dispatch") + _dispatch_logger.addHandler(capture_handler) + + try: + async with AgyDispatchServer( + ca_key=ca_key, ca_cert=ca_cert, allowlist=_CONTROLLED_ALLOWLIST + ) as srv: + _, port = srv.address + + # None SNI: disable hostname verification so we can send without SNI. + ssl_ctx_no_sni = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ssl_ctx_no_sni.check_hostname = False + ssl_ctx_no_sni.verify_mode = ssl.CERT_NONE + try: + _, writer = await asyncio.wait_for( + asyncio.open_connection( + "127.0.0.1", + port, + ssl=ssl_ctx_no_sni, + server_hostname=None, # no SNI extension + ), + timeout=5.0, + ) + writer.close() + try: + await writer.wait_closed() + except Exception: # noqa: BLE001 + pass + none_sni_accepted = True + except (ssl.SSLError, OSError, ConnectionResetError, asyncio.TimeoutError): + none_sni_accepted = False + finally: + _dispatch_logger.removeHandler(capture_handler) + + assert not none_sni_accepted, "None SNI must be rejected" + warned = any("event=sni_refused" in r.getMessage() for r in sni_log_records) + assert warned, ( + f"Expected event=sni_refused WARNING for None SNI; " + f"got: {[r.getMessage() for r in sni_log_records]}" + ) + + +@pytest.mark.asyncio +async def test_sni_mixed_case_host_is_allowlisted( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """A mixed-case SNI names the same host (RFC 6066) and must terminate. + + A spelling one layer accepts and another rejects is worse than a hard + failure: traffic silently skips compression with no signal. Every layer + normalizes via ``normalize_host``. + """ + ca_key, ca_cert, ca_cert_pem = tmp_ca + + # Controlled allowlist holding only the canonical form. + allowlist = frozenset({"daily-cloudcode-pa.googleapis.com"}) + async with AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert, allowlist=allowlist) as srv: + _, port = srv.address + accepted = await _try_tls_connect(port, ca_cert_pem, "Daily-CloudCode-PA.googleapis.com") + + assert accepted, "mixed-case SNI is the same host as the allowlisted form" + + +@pytest.mark.asyncio +async def test_sni_exception_inside_callback_server_stays_alive( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """Exception raised inside _sni_callback -> handshake aborts AND server stays alive.""" + from unittest.mock import patch + + from headroom.proxy.agy_terminator import _LeafCache + + ca_key, ca_cert, ca_cert_pem = tmp_ca + + async with AgyDispatchServer( + ca_key=ca_key, ca_cert=ca_cert, allowlist=_CONTROLLED_ALLOWLIST + ) as srv: + _, port = srv.address + + def _boom(self: _LeafCache, *a: Any, **kw: Any) -> Any: + raise RuntimeError("injected failure") + + with patch.object(_LeafCache, "get_or_mint", _boom): + # The handshake must fail (SSL error), not crash the server. + rejected = not await _try_tls_connect(port, ca_cert_pem, _CONTROLLED_HOST) + + # Server must still be alive. + assert srv._server is not None, "Server must stay alive after exception in callback" + + assert rejected, "Exception in SNI callback must abort handshake" + + +def test_placeholder_host_not_in_default_allowlist() -> None: + """_PLACEHOLDER_HOST 'headroom.internal' must NOT be in DEFAULT_ALLOWLIST.""" + assert "headroom.internal" not in DEFAULT_ALLOWLIST, ( + "headroom.internal must never appear in DEFAULT_ALLOWLIST" + ) + + +# --------------------------------------------------------------------------- +# Tests: post-handshake Host guard (headroom-oqb.1) +# --------------------------------------------------------------------------- + + +async def _http11_request( + port: int, + ca_cert_pem: bytes, + sni_host: str, + host_header: str, + *, + timeout: float = 10.0, +) -> int: + """Perform HTTP/1.1 GET / and return the status code (or 0 on connection failure).""" + ssl_ctx = _build_client_ssl_ctx(ca_cert_pem) + ssl_ctx.set_alpn_protocols(["http/1.1"]) + try: + reader, writer = await asyncio.wait_for( + asyncio.open_connection("127.0.0.1", port, ssl=ssl_ctx, server_hostname=sni_host), + timeout=timeout, + ) + except (ssl.SSLError, OSError, ConnectionResetError): + return 0 + try: + request = (f"GET / HTTP/1.1\r\nHost: {host_header}\r\nConnection: close\r\n\r\n").encode() + writer.write(request) + await writer.drain() + status_line = await asyncio.wait_for(reader.readline(), timeout=timeout) + if not status_line: + return 0 + parts = status_line.split() + return int(parts[1]) if len(parts) >= 2 else 0 + except (OSError, asyncio.TimeoutError, IndexError, ValueError): + return 0 + finally: + writer.close() + try: + await writer.wait_closed() + except Exception: # noqa: BLE001 + pass + + +@pytest.mark.asyncio +async def test_host_guard_non_allowlisted_returns_421( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """Post-handshake Host guard: non-allowlisted Host header -> 421 Misdirected Request. + + We spy on _send_421 to confirm the guard is what sends the 421 (not the upstream app). + """ + from unittest.mock import patch + + ca_key, ca_cert, ca_cert_pem = tmp_ca + + send_421_called = [False] + + async def _spy_send_421(send: Any) -> None: + send_421_called[0] = True + import headroom.proxy.agy_dispatch as _m + + await _m._send_421.__wrapped__(send) if hasattr( + _m._send_421, "__wrapped__" + ) else await _real_send_421(send) + + from headroom.proxy import agy_dispatch as _agy_dispatch_mod + + _real_send_421 = _agy_dispatch_mod._send_421 + + with patch("headroom.proxy.agy_dispatch._send_421", _spy_send_421): + async with AgyDispatchServer( + ca_key=ca_key, ca_cert=ca_cert, allowlist=_CONTROLLED_ALLOWLIST + ) as srv: + _, port = srv.address + status = await _http11_request(port, ca_cert_pem, _CONTROLLED_HOST, "evil.example.com") + + assert status == 421, f"Expected 421 for non-allowlisted Host, got {status}" + assert send_421_called[0], "Guard must call _send_421 for non-allowlisted Host" + + +@pytest.mark.asyncio +async def test_host_guard_allowlisted_passes( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Post-handshake Host guard: allowlisted Host passes through to the app (guard not triggered). + + We spy on _send_421 to confirm the guard does NOT refuse the allowlisted host. + The app may return any status (404, 200, …) — that is app-layer behavior, not the guard. + """ + from unittest.mock import patch + + ca_key, ca_cert, ca_cert_pem = tmp_ca + + send_421_called = [False] + + async def _spy_send_421(send: Any) -> None: + send_421_called[0] = True + from headroom.proxy.agy_dispatch import _send_421 + + await _send_421(send) + + with patch("headroom.proxy.agy_dispatch._send_421", _spy_send_421): + async with AgyDispatchServer( + ca_key=ca_key, ca_cert=ca_cert, allowlist=_CONTROLLED_ALLOWLIST + ) as srv: + _, port = srv.address + status = await _http11_request(port, ca_cert_pem, _CONTROLLED_HOST, _CONTROLLED_HOST) + + assert not send_421_called[0], ( + f"Guard must NOT refuse the allowlisted Host '{_CONTROLLED_HOST}'; got HTTP status {status}" + ) + assert status != 0, "Expected a valid HTTP response (guard passed request to app)" + + +@pytest.mark.asyncio +async def test_host_guard_port_qualified_host_passes( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """Post-handshake Host guard: 'host:443' form is normalized and passes the guard. + + Spy on _send_421 — the guard must NOT refuse 'allowed.test:443'. + """ + from unittest.mock import patch + + ca_key, ca_cert, ca_cert_pem = tmp_ca + + send_421_called = [False] + + async def _spy_send_421(send: Any) -> None: + send_421_called[0] = True + from headroom.proxy.agy_dispatch import _send_421 + + await _send_421(send) + + with patch("headroom.proxy.agy_dispatch._send_421", _spy_send_421): + async with AgyDispatchServer( + ca_key=ca_key, ca_cert=ca_cert, allowlist=_CONTROLLED_ALLOWLIST + ) as srv: + _, port = srv.address + status = await _http11_request( + port, ca_cert_pem, _CONTROLLED_HOST, f"{_CONTROLLED_HOST}:443" + ) + + assert not send_421_called[0], ( + f"Guard must NOT refuse 'host:port' form; got HTTP status {status}" + ) + assert status != 0, "Expected a valid HTTP response (guard passed request to app)" + + +@pytest.mark.asyncio +async def test_host_guard_mixed_case_host_passes( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], +) -> None: + """Post-handshake Host guard: mixed-case Host is normalized (lowercased) and passes. + + Spy on _send_421 — the guard must NOT refuse the uppercased form. + """ + from unittest.mock import patch + + ca_key, ca_cert, ca_cert_pem = tmp_ca + + send_421_called = [False] + + async def _spy_send_421(send: Any) -> None: + send_421_called[0] = True + from headroom.proxy.agy_dispatch import _send_421 + + await _send_421(send) + + with patch("headroom.proxy.agy_dispatch._send_421", _spy_send_421): + async with AgyDispatchServer( + ca_key=ca_key, ca_cert=ca_cert, allowlist=_CONTROLLED_ALLOWLIST + ) as srv: + _, port = srv.address + mixed_case = _CONTROLLED_HOST.upper() + status = await _http11_request(port, ca_cert_pem, _CONTROLLED_HOST, mixed_case) + + assert not send_421_called[0], ( + f"Guard must NOT refuse mixed-case Host (normalized to lower); got HTTP status {status}" + ) + assert status != 0, "Expected a valid HTTP response (guard passed request to app)" + + +# --------------------------------------------------------------------------- +# Tests: load_cert_chain_in_memory used in dispatch (headroom-oqb.2) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dispatch_no_tmpfile_on_linux( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """On Linux (memfd_create available), load_cert_chain is never called with + a regular filesystem path for leaf key material — only /proc/self/fd/ paths.""" + import os + import ssl as _ssl + + ca_key, ca_cert, _ = tmp_ca + + if not hasattr(os, "memfd_create"): + pytest.skip("memfd_create not available; primary path not applicable") + + # Spy on ssl.SSLContext.load_cert_chain to check which paths are passed. + leaf_fs_paths: list[str] = [] + original_load = _ssl.SSLContext.load_cert_chain + + def _spy_load( + self: _ssl.SSLContext, certfile: str, keyfile: object = None, **kwargs: object + ) -> None: + # Flag any certfile that is NOT an anonymous memfd /proc path. + if not certfile.startswith("/proc/self/fd/"): + leaf_fs_paths.append(certfile) + original_load(self, certfile, keyfile, **kwargs) # type: ignore[arg-type] + + monkeypatch.setattr(_ssl.SSLContext, "load_cert_chain", _spy_load) + + async with AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert): + pass + + assert not leaf_fs_paths, ( + f"load_cert_chain must only use /proc/self/fd/ on Linux (memfd), " + f"but got regular fs paths: {leaf_fs_paths}" + ) + + +@pytest.mark.asyncio +async def test_dispatch_handshake_still_works_via_helper( + tmp_ca: tuple[RSAPrivateKey, Certificate, bytes], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """AgyDispatchServer (using load_cert_chain_in_memory) still completes + a TLS handshake for an allowlisted SNI host — regression guard.""" + from fastapi.responses import StreamingResponse + + from headroom.proxy.server import HeadroomProxy + + ca_key, ca_cert, ca_cert_pem = tmp_ca + + async def _fake_stream(self: Any, *args: Any, **kwargs: Any) -> StreamingResponse: + async def _body() -> bytes: + yield b"data: [DONE]\n\n" + + return StreamingResponse(_body(), status_code=200, media_type="text/event-stream") + + monkeypatch.setattr(HeadroomProxy, "_stream_response", _fake_stream) + + async with AgyDispatchServer(ca_key=ca_key, ca_cert=ca_cert) as srv: + _, port = srv.address + ssl_ctx = _build_client_ssl_ctx(ca_cert_pem) + ssl_ctx.set_alpn_protocols(["http/1.1"]) + conn_reader, conn_writer = await asyncio.open_connection( + "127.0.0.1", + port, + ssl=ssl_ctx, + server_hostname=ALLOWLIST_HOST, + ) + try: + request = (f"GET / HTTP/1.1\r\nHost: {ALLOWLIST_HOST}\r\n\r\n").encode() + conn_writer.write(request) + await conn_writer.drain() + response_line = await asyncio.wait_for(conn_reader.readline(), timeout=10.0) + finally: + conn_writer.close() + + # Any HTTP response (even 404/421) confirms the TLS handshake succeeded. + assert response_line.startswith(b"HTTP/"), ( + f"Expected HTTP response; TLS handshake must succeed via helper. Got: {response_line!r}" + ) + + +# --------------------------------------------------------------------------- +# make_host_guard — unit tests (synthetic ASGI scopes, no TLS) +# --------------------------------------------------------------------------- + +_GUARD_ALLOW = frozenset({"daily-cloudcode-pa.googleapis.com"}) + + +async def _run_host_guard( + allowlist: frozenset[str], scope: dict[str, Any] +) -> tuple[bool, int | None]: + """Drive make_host_guard over *scope*; return (downstream_app_called, status).""" + app_called = [False] + status: list[int | None] = [None] + + async def _app(s: Any, r: Any, sd: Any) -> None: + app_called[0] = True + + async def _send(msg: dict[str, Any]) -> None: + if msg.get("type") == "http.response.start": + status[0] = msg["status"] + + async def _receive() -> dict[str, Any]: + return {"type": "http.request", "body": b"", "more_body": False} + + await make_host_guard(_app, allowlist)(scope, _receive, _send) + return app_called[0], status[0] + + +@pytest.mark.asyncio +async def test_host_guard_allowlisted_passes_to_app() -> None: + called, status = await _run_host_guard( + _GUARD_ALLOW, + {"type": "http", "headers": [(b"host", b"daily-cloudcode-pa.googleapis.com")]}, + ) + assert called and status is None + + +@pytest.mark.asyncio +async def test_host_guard_non_allowlisted_421() -> None: + called, status = await _run_host_guard( + _GUARD_ALLOW, {"type": "http", "headers": [(b"host", b"evil.example.com")]} + ) + assert not called and status == 421 + + +@pytest.mark.asyncio +async def test_host_guard_duplicate_host_421() -> None: + """Two Host headers (smuggling vector) -> 421, app never reached.""" + called, status = await _run_host_guard( + _GUARD_ALLOW, + { + "type": "http", + "headers": [ + (b"host", b"daily-cloudcode-pa.googleapis.com"), + (b"host", b"evil.example.com"), + ], + }, + ) + assert not called and status == 421 + + +@pytest.mark.asyncio +async def test_host_guard_zero_host_421() -> None: + called, status = await _run_host_guard(_GUARD_ALLOW, {"type": "http", "headers": []}) + assert not called and status == 421 + + +@pytest.mark.asyncio +async def test_host_guard_empty_host_value_421() -> None: + """A single Host header present but with an empty value -> 421 (blank-host branch).""" + called, status = await _run_host_guard( + _GUARD_ALLOW, {"type": "http", "headers": [(b"host", b"")]} + ) + assert not called and status == 421 + + +@pytest.mark.asyncio +async def test_host_guard_non_digit_port_suffix_kept_as_is() -> None: + """Host 'example.com:abc' has a non-digit suffix after ':' so it is NOT + stripped (the `if right.isdigit()` branch is False) and the literal + string (including the bogus ':abc' suffix) is checked against the + allowlist as-is. + + Proof this covers the False branch (not just a passthrough): if the + suffix were incorrectly stripped, `normalized` would become + 'example.com', which is absent from this test's allowlist, and the + request would be refused (421) instead of passed through. + """ + allowlist = frozenset({"example.com:abc"}) + called, status = await _run_host_guard( + allowlist, {"type": "http", "headers": [(b"host", b"example.com:abc")]} + ) + assert called and status is None + + +@pytest.mark.asyncio +async def test_host_guard_uppercase_and_port_passes() -> None: + """RFC-compliant mixed-case + port-qualified Host normalizes and passes.""" + called, status = await _run_host_guard( + _GUARD_ALLOW, + {"type": "http", "headers": [(b"host", b"DAILY-CloudCode-PA.googleapis.com:443")]}, + ) + assert called and status is None + + +@pytest.mark.asyncio +async def test_host_guard_lifespan_scope_passes() -> None: + """Non-http/websocket scopes (e.g. lifespan) are not guarded.""" + called, status = await _run_host_guard(_GUARD_ALLOW, {"type": "lifespan", "headers": []}) + assert called and status is None + + +# --------------------------------------------------------------------------- +# make_host_guard — project header injection (synthetic ASGI scopes, no TLS) +# --------------------------------------------------------------------------- + + +async def _run_host_guard_capture( + allowlist: frozenset[str], + scope: dict[str, Any], + project: str | None, +) -> tuple[dict[str, Any] | None, int | None]: + """Drive make_host_guard(app, allowlist, project) over *scope*. + + Returns (captured_scope, status). captured_scope is the scope the inner + app was invoked with (or None if the app was never called, e.g. refused). + """ + captured: dict[str, Any] = {} + status: list[int | None] = [None] + + async def inner(s: Any, r: Any, sd: Any) -> None: + captured["scope"] = s + + async def _send(msg: dict[str, Any]) -> None: + if msg.get("type") == "http.response.start": + status[0] = msg["status"] + + async def _receive() -> dict[str, Any]: + return {"type": "http.request", "body": b"", "more_body": False} + + await make_host_guard(inner, allowlist, project)(scope, _receive, _send) + return captured.get("scope"), status[0] + + +def _project_header_values(scope: dict[str, Any]) -> list[bytes]: + return [v for name, v in scope.get("headers", ()) if name.lower() == b"x-headroom-project"] + + +@pytest.mark.asyncio +async def test_host_guard_injects_project_header() -> None: + """(a) project set + allowlisted Host -> exactly one x-headroom-project header.""" + scope, status = await _run_host_guard_capture( + _GUARD_ALLOW, + {"type": "http", "headers": [(b"host", b"daily-cloudcode-pa.googleapis.com")]}, + "myproj", + ) + assert status is None + assert scope is not None, "inner app must be called for allowlisted host" + assert _project_header_values(scope) == [b"myproj"] + + +@pytest.mark.asyncio +async def test_host_guard_replaces_forged_project_header() -> None: + """(b) client-forged x-headroom-project is replaced (not duplicated).""" + scope, status = await _run_host_guard_capture( + _GUARD_ALLOW, + { + "type": "http", + "headers": [ + (b"host", b"daily-cloudcode-pa.googleapis.com"), + (b"x-headroom-project", b"attacker"), + ], + }, + "myproj", + ) + assert status is None + assert scope is not None + assert _project_header_values(scope) == [b"myproj"], "forged value must be replaced, not kept" + + +@pytest.mark.asyncio +async def test_host_guard_no_project_leaves_headers_untouched() -> None: + """(c) project=None -> no x-headroom-project header; scope headers unchanged.""" + original_headers = [(b"host", b"daily-cloudcode-pa.googleapis.com")] + scope, status = await _run_host_guard_capture( + _GUARD_ALLOW, + {"type": "http", "headers": list(original_headers)}, + None, + ) + assert status is None + assert scope is not None + assert _project_header_values(scope) == [] + assert scope["headers"] == original_headers + + +@pytest.mark.asyncio +async def test_host_guard_non_allowlisted_refused_with_project() -> None: + """(d) non-allowlisted Host still refused (421) and inner app NOT called.""" + scope, status = await _run_host_guard_capture( + _GUARD_ALLOW, + {"type": "http", "headers": [(b"host", b"evil.example.com")]}, + "myproj", + ) + assert status == 421 + assert scope is None, "inner app must NOT be called for non-allowlisted host" diff --git a/tests/test_agy_fr_compressor_unit.py b/tests/test_agy_fr_compressor_unit.py new file mode 100644 index 000000000..f3096ad5e --- /dev/null +++ b/tests/test_agy_fr_compressor_unit.py @@ -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 diff --git a/tests/test_agy_fr_mode_default.py b/tests/test_agy_fr_mode_default.py new file mode 100644 index 000000000..74c367d06 --- /dev/null +++ b/tests/test_agy_fr_mode_default.py @@ -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" diff --git a/tests/test_agy_fr_perf_37g35.py b/tests/test_agy_fr_perf_37g35.py new file mode 100644 index 000000000..f19f49bb9 --- /dev/null +++ b/tests/test_agy_fr_perf_37g35.py @@ -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 diff --git a/tests/test_agy_fr_retrieve_envelope_exempt.py b/tests/test_agy_fr_retrieve_envelope_exempt.py new file mode 100644 index 000000000..f460f796a --- /dev/null +++ b/tests/test_agy_fr_retrieve_envelope_exempt.py @@ -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) diff --git a/tests/test_agy_functionresponse_compression.py b/tests/test_agy_functionresponse_compression.py new file mode 100644 index 000000000..34b6654a6 --- /dev/null +++ b/tests/test_agy_functionresponse_compression.py @@ -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=]`` 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() diff --git a/tests/test_agy_functionresponse_compression_matrix.py b/tests/test_agy_functionresponse_compression_matrix.py new file mode 100644 index 000000000..2ff56e43b --- /dev/null +++ b/tests/test_agy_functionresponse_compression_matrix.py @@ -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() diff --git a/tests/test_agy_print_mode_version_gate.py b/tests/test_agy_print_mode_version_gate.py new file mode 100644 index 000000000..0874cbfc3 --- /dev/null +++ b/tests/test_agy_print_mode_version_gate.py @@ -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 diff --git a/tests/test_agy_provider_env.py b/tests/test_agy_provider_env.py new file mode 100644 index 000000000..2debf2831 --- /dev/null +++ b/tests/test_agy_provider_env.py @@ -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"} diff --git a/tests/test_agy_registrar.py b/tests/test_agy_registrar.py new file mode 100644 index 000000000..01c30bd2c --- /dev/null +++ b/tests/test_agy_registrar.py @@ -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 diff --git a/tests/test_agy_retrieve.py b/tests/test_agy_retrieve.py new file mode 100644 index 000000000..cd454e025 --- /dev/null +++ b/tests/test_agy_retrieve.py @@ -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 diff --git a/tests/test_agy_retrieve_exemption.py b/tests/test_agy_retrieve_exemption.py new file mode 100644 index 000000000..2e84fdcfd --- /dev/null +++ b/tests/test_agy_retrieve_exemption.py @@ -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 diff --git a/tests/test_agy_retrieve_exposure_gate.py b/tests/test_agy_retrieve_exposure_gate.py new file mode 100644 index 000000000..ad2280ee6 --- /dev/null +++ b/tests/test_agy_retrieve_exposure_gate.py @@ -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 (``/mcp//.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 diff --git a/tests/test_agy_retrieve_persistent.py b/tests/test_agy_retrieve_persistent.py new file mode 100644 index 000000000..1ed2e30f9 --- /dev/null +++ b/tests/test_agy_retrieve_persistent.py @@ -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}' diff --git a/tests/test_agy_savings_inbox.py b/tests/test_agy_savings_inbox.py new file mode 100644 index 000000000..2db63a302 --- /dev/null +++ b/tests/test_agy_savings_inbox.py @@ -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 diff --git a/tests/test_agy_savings_integration.py b/tests/test_agy_savings_integration.py new file mode 100644 index 000000000..8dcc9f394 --- /dev/null +++ b/tests/test_agy_savings_integration.py @@ -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 diff --git a/tests/test_agy_sse_usage_envelope.py b/tests/test_agy_sse_usage_envelope.py new file mode 100644 index 000000000..05a13e7b3 --- /dev/null +++ b/tests/test_agy_sse_usage_envelope.py @@ -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": {}}) == {} diff --git a/tests/test_agy_stats.py b/tests/test_agy_stats.py new file mode 100644 index 000000000..9f38ea5e6 --- /dev/null +++ b/tests/test_agy_stats.py @@ -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 diff --git a/tests/test_agy_terminator.py b/tests/test_agy_terminator.py new file mode 100644 index 000000000..d9b09de84 --- /dev/null +++ b/tests/test_agy_terminator.py @@ -0,0 +1,1624 @@ +"""Tests for headroom.proxy.agy_terminator. + +All tests use ephemeral ports and tmp_path; real ~/.headroom is never touched. +Tests use real asyncio connections over loopback to verify behavior. +""" + +from __future__ import annotations + +import asyncio +import datetime +import ssl + +import pytest +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_terminator import ( + DEFAULT_ALLOWLIST, + AgyCONNECTTerminator, + _is_loopback, + _LeafCache, + _parse_connect, + mint_leaf, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +ALLOWLIST_HOST = "daily-cloudcode-pa.googleapis.com" +NON_ALLOWLIST_HOST = "example.com" + + +def _make_test_ca() -> tuple[RSAPrivateKey, Certificate, bytes]: + """Generate a fast 2048-bit RSA root CA for tests (never touches disk).""" + key: RSAPrivateKey = rsa.generate_private_key(public_exponent=65537, key_size=2048) + now = datetime.datetime.now(tz=datetime.timezone.utc) + subject = issuer = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "Headroom Test CA")]) + 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=365)) + .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, + ) + .sign(key, hashes.SHA256()) + ) + ca_cert_pem = cert.public_bytes(serialization.Encoding.PEM) + return key, cert, ca_cert_pem + + +@pytest.fixture(scope="module") +def tmp_ca() -> tuple[RSAPrivateKey, Certificate, bytes]: + """Return (ca_key, ca_cert, ca_cert_pem) — module-scoped; generated once.""" + return _make_test_ca() + + +# --------------------------------------------------------------------------- +# Unit: _parse_connect +# --------------------------------------------------------------------------- + + +def test_parse_connect_basic() -> None: + host, port = _parse_connect("CONNECT example.com:443 HTTP/1.1") + assert host == "example.com" + assert port == 443 + + +def test_parse_connect_lowercase() -> None: + host, port = _parse_connect("connect api.example.com:8443 HTTP/1.1") + assert host == "api.example.com" + assert port == 8443 + + +@pytest.mark.parametrize( + "target", + [ + "CloudCode-PA.googleapis.com:443", # mixed case + "cloudcode-pa.googleapis.com.:443", # trailing root dot + ], +) +def test_parse_connect_normalizes_equivalent_host_forms(target: str) -> None: + """Equivalent spellings must reach the allowlist in one canonical form. + + The allowlist check is exact match, so an un-normalized target would fall + through to the blind tunnel: the request still works but silently skips TLS + termination and compression, with no signal that it was bypassed. + """ + host, port = _parse_connect(f"CONNECT {target} HTTP/1.1") + assert host == "cloudcode-pa.googleapis.com" + assert port == 443 + assert host in DEFAULT_ALLOWLIST + + +def test_parse_connect_invalid_raises() -> None: + with pytest.raises(ValueError): + _parse_connect("GET / HTTP/1.1") + + +def test_parse_connect_missing_port_raises() -> None: + with pytest.raises(ValueError): + _parse_connect("CONNECT example.com HTTP/1.1") + + +# --------------------------------------------------------------------------- +# Unit: _is_loopback +# --------------------------------------------------------------------------- + + +def test_is_loopback_127() -> None: + assert _is_loopback("127.0.0.1") is True + + +def test_is_loopback_localhost() -> None: + assert _is_loopback("localhost") is True + + +def test_is_loopback_ipv6() -> None: + assert _is_loopback("::1") is True + + +def test_is_loopback_public() -> None: + assert _is_loopback("8.8.8.8") is False + + +def test_is_loopback_hostname() -> None: + assert _is_loopback("example.com") is False + + +def test_is_loopback_ipv4_shorthand_dotted() -> None: + """127.1 is a valid inet_aton shorthand for 127.0.0.1.""" + assert _is_loopback("127.1") is True + + +def test_is_loopback_ipv4_decimal() -> None: + """2130706433 is the decimal encoding of 127.0.0.1.""" + assert _is_loopback("2130706433") is True + + +def test_is_loopback_zero_shorthand() -> None: + """0 is inet_aton shorthand for 0.0.0.0 (unspecified, treated as loopback).""" + assert _is_loopback("0") is True + + +def test_is_loopback_unspecified() -> None: + """0.0.0.0 is is_unspecified, not is_loopback, but Linux connect() reaches localhost.""" + assert _is_loopback("0.0.0.0") is True + + +def test_is_loopback_localhost_trailing_dot() -> None: + assert _is_loopback("localhost.") is True + + +def test_is_loopback_ipv4_mapped_ipv6() -> None: + """Must not depend on interpreter version (CPython gh-103365, fixed in 3.13).""" + assert _is_loopback("::ffff:127.0.0.1") is True + + +def test_is_loopback_still_no_dns_for_shorthand_lookalike() -> None: + """example.com must still return False — the function stays DNS-free.""" + assert _is_loopback("example.com") is False + + +# --------------------------------------------------------------------------- +# Unit: mint_leaf +# --------------------------------------------------------------------------- + + +def test_mint_leaf_san(tmp_ca: tuple) -> None: + """Minted leaf must have SAN=dNSName for the host. (f)""" + ca_key, ca_cert, _ = tmp_ca + cert_pem, _ = mint_leaf("api.example.com", ca_key, ca_cert) + cert = x509.load_pem_x509_certificate(cert_pem) + san = cert.extensions.get_extension_for_class(x509.SubjectAlternativeName) + dns_names = san.value.get_values_for_type(x509.DNSName) + # Exact SAN match (not substring/membership) — the leaf carries exactly one dNSName. + assert dns_names == ["api.example.com"] + + +def test_mint_leaf_eku_server_auth(tmp_ca: tuple) -> None: + """Minted leaf must have EKU=serverAuth only. (f)""" + ca_key, ca_cert, _ = tmp_ca + cert_pem, _ = mint_leaf("api.example.com", ca_key, ca_cert) + cert = x509.load_pem_x509_certificate(cert_pem) + eku = cert.extensions.get_extension_for_class(x509.ExtendedKeyUsage) + assert list(eku.value) == [ExtendedKeyUsageOID.SERVER_AUTH] + + +def test_mint_leaf_validity_lte_72h(tmp_ca: tuple) -> None: + """Minted leaf validity must be <= 72 hours. (f)""" + ca_key, ca_cert, _ = tmp_ca + cert_pem, _ = mint_leaf("api.example.com", ca_key, ca_cert) + cert = x509.load_pem_x509_certificate(cert_pem) + delta = cert.not_valid_after_utc - cert.not_valid_before_utc + assert delta <= datetime.timedelta(hours=72) + + +def test_mint_leaf_not_ca(tmp_ca: tuple) -> None: + """Minted leaf must not have CA:TRUE.""" + ca_key, ca_cert, _ = tmp_ca + cert_pem, _ = mint_leaf("api.example.com", ca_key, ca_cert) + cert = x509.load_pem_x509_certificate(cert_pem) + bc = cert.extensions.get_extension_for_class(x509.BasicConstraints) + assert bc.value.ca is False + + +def test_mint_leaf_signed_by_root(tmp_ca: tuple) -> None: + """Leaf issuer must match the root CA subject.""" + ca_key, ca_cert, _ = tmp_ca + cert_pem, _ = mint_leaf("api.example.com", ca_key, ca_cert) + cert = x509.load_pem_x509_certificate(cert_pem) + assert cert.issuer == ca_cert.subject + + +# --------------------------------------------------------------------------- +# Unit: _LeafCache +# --------------------------------------------------------------------------- + + +def test_leaf_cache_reuse(tmp_ca: tuple) -> None: + """Same host returns same cert PEM (serial equality). (b)""" + ca_key, ca_cert, _ = tmp_ca + cache = _LeafCache(max_size=10) + cert1, _ = cache.get_or_mint("api.example.com", ca_key, ca_cert) + cert2, _ = cache.get_or_mint("api.example.com", ca_key, ca_cert) + obj1 = x509.load_pem_x509_certificate(cert1) + obj2 = x509.load_pem_x509_certificate(cert2) + assert obj1.serial_number == obj2.serial_number + + +def test_leaf_cache_different_hosts(tmp_ca: tuple) -> None: + """Different hosts get different leaf certs.""" + ca_key, ca_cert, _ = tmp_ca + cache = _LeafCache(max_size=10) + cert1, _ = cache.get_or_mint("host-a.example.com", ca_key, ca_cert) + cert2, _ = cache.get_or_mint("host-b.example.com", ca_key, ca_cert) + obj1 = x509.load_pem_x509_certificate(cert1) + obj2 = x509.load_pem_x509_certificate(cert2) + assert obj1.serial_number != obj2.serial_number + + +def test_leaf_cache_bound_evicts(tmp_ca: tuple) -> None: + """Cache with max_size=1 evicts oldest on second host.""" + ca_key, ca_cert, _ = tmp_ca + cache = _LeafCache(max_size=1) + cache.get_or_mint("host-a.example.com", ca_key, ca_cert) + cache.get_or_mint("host-b.example.com", ca_key, ca_cert) + assert len(cache._cache) == 1 + # After max_size=1 eviction the sole cached key is exactly host-b. + assert list(cache._cache) == ["host-b.example.com"] + + +# --------------------------------------------------------------------------- +# Integration: listener bind address +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_listener_bound_to_loopback_only(tmp_ca: tuple) -> None: + """Listener must be bound to 127.0.0.1, not 0.0.0.0. (d)""" + ca_key, ca_cert, _ = tmp_ca + terminator = AgyCONNECTTerminator( + allowlist=DEFAULT_ALLOWLIST, + ca_key=ca_key, + ca_cert=ca_cert, + dispatch_port=1, + ) + await terminator.start() + try: + bound_host, bound_port = terminator.address + assert bound_host == "127.0.0.1", f"Expected 127.0.0.1 but got {bound_host}" + assert bound_port > 0 + + # Connecting via 127.0.0.1 succeeds. + reader, writer = await asyncio.open_connection("127.0.0.1", bound_port) + writer.close() + await writer.wait_closed() + + # 0.0.0.0 is NOT a valid bind address assertion; + # verify sockets don't list 0.0.0.0. + for sock in terminator._server.sockets: + sock_host = sock.getsockname()[0] + assert sock_host != "0.0.0.0", "Server must not bind to 0.0.0.0" + finally: + await terminator.stop() + + +# --------------------------------------------------------------------------- +# Integration: non-allowlist → blind tunnel (c) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_blind_tunnel_byte_faithful(tmp_ca: tuple) -> None: + """Non-allowlisted CONNECT: bytes round-trip unmodified via plain TCP echo server. (c)""" + ca_key, ca_cert, _ = tmp_ca + + # Spin up a plain TCP echo server. + echo_host = "127.0.0.1" + + async def echo_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + data = await asyncio.wait_for(reader.read(1024), timeout=5.0) + if data: + writer.write(data) + await writer.drain() + finally: + writer.close() + + echo_server = await asyncio.start_server(echo_handler, echo_host, 0) + echo_port = echo_server.sockets[0].getsockname()[1] + + terminator = AgyCONNECTTerminator( + allowlist=frozenset({ALLOWLIST_HOST}), # echo host NOT in allowlist + ca_key=ca_key, + ca_cert=ca_cert, + dispatch_port=1, + ) + await terminator.start() + + try: + proxy_host, proxy_port = terminator.address + + raw_reader, raw_writer = await asyncio.open_connection(proxy_host, proxy_port) + connect_req = ( + f"CONNECT {echo_host}:{echo_port} HTTP/1.1\r\nHost: {echo_host}:{echo_port}\r\n\r\n" + ) + raw_writer.write(connect_req.encode()) + await raw_writer.drain() + response = await raw_reader.readline() + assert b"200" in response, f"Expected 200 for blind tunnel, got {response!r}" + # Drain the blank line separating HTTP status from body. + await raw_reader.readline() + + # Send payload and expect it echoed back verbatim — no TLS wrapping. + payload = b"hello blind tunnel \x00\x01\x02" + raw_writer.write(payload) + await raw_writer.drain() + + received = await asyncio.wait_for(raw_reader.read(len(payload)), timeout=5.0) + assert received == payload, f"Echo mismatch: {received!r} != {payload!r}" + finally: + await terminator.stop() + echo_server.close() + await echo_server.wait_closed() + + +# --------------------------------------------------------------------------- +# Integration: self-loop guard (e) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_self_loop_guard_via_https_proxy_env( + tmp_ca: tuple, monkeypatch: pytest.MonkeyPatch +) -> None: + """HTTPS_PROXY pointing at loopback must be refused. (e)""" + ca_key, ca_cert, _ = tmp_ca + monkeypatch.setenv("HTTPS_PROXY", "http://127.0.0.1:3128") + + terminator = AgyCONNECTTerminator( + allowlist=frozenset({ALLOWLIST_HOST}), + ca_key=ca_key, + ca_cert=ca_cert, + dispatch_port=1, + ) + await terminator.start() + + try: + proxy_host, proxy_port = terminator.address + raw_reader, raw_writer = await asyncio.open_connection(proxy_host, proxy_port) + connect_req = ( + f"CONNECT {NON_ALLOWLIST_HOST}:443 HTTP/1.1\r\nHost: {NON_ALLOWLIST_HOST}:443\r\n\r\n" + ) + raw_writer.write(connect_req.encode()) + await raw_writer.drain() + response = await raw_reader.readline() + assert b"403" in response, f"Expected 403 when HTTPS_PROXY is loopback, got {response!r}" + finally: + await terminator.stop() + + +@pytest.mark.asyncio +async def test_non_http_upstream_proxy_scheme_rejected( + tmp_ca: tuple, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """A non-http(s) HTTPS_PROXY (e.g. socks5://) must be refused before any + + HTTP CONNECT text is written into it, and the refusal must never leak + the credential embedded in the proxy URL's userinfo. + """ + import headroom.proxy.agy_terminator as _mod + + ca_key, ca_cert, _ = tmp_ca + monkeypatch.setenv("HTTPS_PROXY", "socks5://user:s3cr3t@proxy:1080") + + called = False + + async def _spy(*args: object, **kwargs: object) -> tuple[object, object]: + nonlocal called + called = True + raise AssertionError("_connect_via_upstream_proxy must not be reached") + + monkeypatch.setattr(_mod, "_connect_via_upstream_proxy", _spy) + + terminator = AgyCONNECTTerminator( + allowlist=frozenset({ALLOWLIST_HOST}), + ca_key=ca_key, + ca_cert=ca_cert, + dispatch_port=1, + ) + await terminator.start() + + try: + with caplog.at_level("WARNING"): + proxy_host, proxy_port = terminator.address + raw_reader, raw_writer = await asyncio.open_connection(proxy_host, proxy_port) + connect_req = ( + f"CONNECT {NON_ALLOWLIST_HOST}:443 HTTP/1.1\r\n" + f"Host: {NON_ALLOWLIST_HOST}:443\r\n\r\n" + ) + raw_writer.write(connect_req.encode()) + await raw_writer.drain() + response = await raw_reader.readline() + assert b"403" in response, f"Expected 403 for socks5:// upstream, got {response!r}" + assert called is False + assert "s3cr3t" not in caplog.text + finally: + await terminator.stop() + + +# --------------------------------------------------------------------------- +# Integration: AgyCONNECTTerminator context manager +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_terminator_context_manager(tmp_ca: tuple) -> None: + """async with AgyCONNECTTerminator works correctly.""" + ca_key, ca_cert, _ = tmp_ca + async with AgyCONNECTTerminator(dispatch_port=1, ca_key=ca_key, ca_cert=ca_cert) as t: + host, port = t.address + assert host == "127.0.0.1" + assert port > 0 + assert t._server is None + + +# --------------------------------------------------------------------------- +# Integration: bad CONNECT request → 400 +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_bad_connect_returns_400(tmp_ca: tuple) -> None: + """Malformed (non-CONNECT) request returns 400.""" + ca_key, ca_cert, _ = tmp_ca + async with AgyCONNECTTerminator(dispatch_port=1, ca_key=ca_key, ca_cert=ca_cert) as t: + proxy_host, proxy_port = t.address + reader, writer = await asyncio.open_connection(proxy_host, proxy_port) + writer.write(b"GET / HTTP/1.1\r\n\r\n") + await writer.drain() + response = await reader.readline() + assert b"400" in response + writer.close() + + +# --------------------------------------------------------------------------- +# Regression: header-drain timeout aborts (no splice) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_connect_header_timeout_aborts() -> None: + """Client stalls mid-headers after CONNECT line → connection aborted, no splice. + + Verifies defect fix: asyncio.TimeoutError in header drain must close + client_writer and return, never proceeding to _handle_mitm/_handle_blind_tunnel. + """ + import unittest.mock as mock + + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _handle_connect + + mitm_called = False + blind_called = False + + async def _fake_mitm(*args: object, **kwargs: object) -> None: + nonlocal mitm_called + mitm_called = True + + async def _fake_blind(*args: object, **kwargs: object) -> None: + nonlocal blind_called + blind_called = True + + # Feed CONNECT line, then nothing — header-drain readline will block. + client_reader = asyncio.StreamReader() + client_reader.feed_data(b"CONNECT notallowlisted.example.com:443 HTTP/1.1\r\n") + + close_called = False + + class _TrackingWriter: + def get_extra_info(self, key: str, default: object = None) -> object: # noqa: ANN401 + if key == "peername": + return ("127.0.0.1", 1234) + return default + + def write(self, data: bytes) -> None: + pass + + async def drain(self) -> None: + pass + + def close(self) -> None: + nonlocal close_called + close_called = True + + async def wait_closed(self) -> None: + pass + + client_writer = _TrackingWriter() # type: ignore[assignment] + + # Tiny timeout so the header-drain readline genuinely times out fast. + with ( + mock.patch.object(_mod, "_handle_mitm", _fake_mitm), + mock.patch.object(_mod, "_handle_blind_tunnel", _fake_blind), + mock.patch.object(_mod, "_CONNECT_TIMEOUT", 0.01), + ): + await _handle_connect( + client_reader, + client_writer, # type: ignore[arg-type] + allowlist=frozenset(), + dispatch_port=1, + ) + + assert close_called, "client_writer.close() must be called on header timeout" + assert not mitm_called, "_handle_mitm must NOT be called on header timeout" + assert not blind_called, "_handle_blind_tunnel must NOT be called on header timeout" + + +# --------------------------------------------------------------------------- +# Regression: upstream proxy header-drain timeout closes upstream writer +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_upstream_proxy_timeout_closes_writer() -> None: + """Header-drain readline in _connect_via_upstream_proxy times out → upstream writer closed. + + A fake upstream proxy sends the 200 response line then stalls (never sends + the blank-line header terminator). With a tiny _CONNECT_TIMEOUT the + header-drain readline times out and the upstream writer must be closed. + """ + import unittest.mock as mock + + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _connect_via_upstream_proxy + + closed_writers: list[object] = [] + # Gate that the proxy releases when it has sent the 200 response. + proxy_sent_200 = asyncio.Event() + # Gate the proxy waits on so teardown can unblock it cleanly. + proxy_release = asyncio.Event() + + async def _stall_proxy_handler( + reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + """Accept CONNECT, reply 200, then stall without the blank-line terminator.""" + try: + await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), timeout=2.0) + except (asyncio.IncompleteReadError, asyncio.TimeoutError): + pass + writer.write(b"HTTP/1.1 200 Connection Established\r\n") + await writer.drain() + proxy_sent_200.set() + # Block until teardown releases us (or until cancelled). + try: + await proxy_release.wait() + except asyncio.CancelledError: + pass + finally: + writer.close() + + proxy_server = await asyncio.start_server(_stall_proxy_handler, host="127.0.0.1", port=0) + proxy_addr = proxy_server.sockets[0].getsockname() + proxy_host, proxy_port = proxy_addr[0], proxy_addr[1] + + orig_open_conn = asyncio.open_connection + + async def _spy_open_conn( + host: str, port: int, **kwargs: object + ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + r, w = await orig_open_conn(host, port, **kwargs) + orig_close = w.close + + def _tracked_close() -> None: + closed_writers.append(w) + orig_close() + + w.close = _tracked_close # type: ignore[method-assign] + return r, w + + try: + with ( + mock.patch.object(_mod, "_CONNECT_TIMEOUT", 0.2), + mock.patch("headroom.proxy.agy_terminator.asyncio.open_connection", _spy_open_conn), + ): + try: + r, w = await _connect_via_upstream_proxy( + proxy_host, proxy_port, "target.example.com", 443, None, None + ) + w.close() + pytest.fail("Expected asyncio.TimeoutError from stalled header drain") + except (asyncio.TimeoutError, OSError): + pass # expected path + finally: + proxy_release.set() # unblock any stalled handler + proxy_server.close() + try: + await asyncio.wait_for(proxy_server.wait_closed(), timeout=2.0) + except asyncio.TimeoutError: + pass + + assert closed_writers, "upstream writer must be closed when header-drain readline times out" + + +# --------------------------------------------------------------------------- +# Regression: blind tunnel drain error closes target writer +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_blind_tunnel_drain_error_closes_target() -> None: + """client_writer.drain() raises before _blind_splice → target_writer is closed. + + Verifies defect fix: if the 200-response drain raises (client disconnected), + target_writer must be closed to avoid fd leak. + """ + import unittest.mock as mock + + from headroom.proxy.agy_terminator import _handle_blind_tunnel + + target_release = asyncio.Event() + + async def _idle_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + await target_release.wait() + except asyncio.CancelledError: + pass + finally: + writer.close() + + target_server = await asyncio.start_server(_idle_handler, host="127.0.0.1", port=0) + target_addr = target_server.sockets[0].getsockname() + target_host, target_port = target_addr[0], target_addr[1] + + target_writer_closed = False + orig_open_conn = asyncio.open_connection + + async def _spy_target_conn( + host: str, port: int, **kwargs: object + ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + r, w = await orig_open_conn(host, port, **kwargs) + orig_close = w.close + + def _tracked_close() -> None: + nonlocal target_writer_closed + target_writer_closed = True + orig_close() + + w.close = _tracked_close # type: ignore[method-assign] + return r, w + + class _DrainFailWriter: + """client_writer stub whose drain() always raises ConnectionResetError.""" + + def get_extra_info(self, key: str, default: object = None) -> object: # noqa: ANN401 + return ("127.0.0.1", 9999) if key == "peername" else default + + def write(self, data: bytes) -> None: + pass + + async def drain(self) -> None: + raise ConnectionResetError("client gone") + + def close(self) -> None: + pass + + async def wait_closed(self) -> None: + pass + + client_reader = asyncio.StreamReader() + + try: + with mock.patch("headroom.proxy.agy_terminator.asyncio.open_connection", _spy_target_conn): + try: + await _handle_blind_tunnel( + client_reader, + _DrainFailWriter(), # type: ignore[arg-type] + target_host, + target_port, + None, + ) + except Exception: # noqa: BLE001 + pass # any propagated exception is acceptable + finally: + target_release.set() + target_server.close() + try: + await asyncio.wait_for(target_server.wait_closed(), timeout=2.0) + except asyncio.TimeoutError: + pass + + assert target_writer_closed, ( + "target_writer.close() must be called when client_writer.drain() raises before splice" + ) + + +# --------------------------------------------------------------------------- +# Coverage: _LeafCache re-mints an expired leaf in place +# --------------------------------------------------------------------------- + + +def test_leaf_cache_expired_entry_remints(tmp_ca: tuple) -> None: + """An expired cache entry (not_valid_after in the past) is re-minted.""" + ca_key, ca_cert, _ = tmp_ca + cache = _LeafCache(max_size=10) + cert1, _ = cache.get_or_mint("expiring.example.com", ca_key, ca_cert) + obj1 = x509.load_pem_x509_certificate(cert1) + + # Force the cached entry to look expired. + cert_pem, key_pem, _ = cache._cache["expiring.example.com"] + past = datetime.datetime.now(tz=datetime.timezone.utc) - datetime.timedelta(hours=1) + cache._cache["expiring.example.com"] = (cert_pem, key_pem, past) + + cert2, _ = cache.get_or_mint("expiring.example.com", ca_key, ca_cert) + obj2 = x509.load_pem_x509_certificate(cert2) + assert obj1.serial_number != obj2.serial_number, "Expired leaf must be re-minted" + + +# --------------------------------------------------------------------------- +# Coverage: _splice_half swallows writer.write_eof() exceptions +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_splice_half_write_eof_exception_swallowed() -> None: + """_splice_half must swallow exceptions raised by writer.write_eof().""" + from headroom.proxy.agy_terminator import _splice_half + + reader = asyncio.StreamReader() + reader.feed_data(b"payload") + reader.feed_eof() + + written = bytearray() + + class _EofRaisingWriter: + def write(self, data: bytes) -> None: + written.extend(data) + + async def drain(self) -> None: + pass + + def write_eof(self) -> None: + raise RuntimeError("eof boom") + + # Must not raise, despite write_eof() raising internally. + await _splice_half(reader, _EofRaisingWriter()) # type: ignore[arg-type] + assert bytes(written) == b"payload" + + +# --------------------------------------------------------------------------- +# Coverage: _blind_splice except-branch cancels both pump tasks +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_blind_splice_wait_exception_cancels_both_tasks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """If awaiting the pump tasks raises, both tasks are cancelled and both + writers are still closed via the finally block.""" + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _blind_splice + + client_reader = asyncio.StreamReader() + target_reader = asyncio.StreamReader() + closed = {"client": False, "target": False} + + class _W: + def __init__(self, name: str) -> None: + self._name = name + + def write(self, data: bytes) -> None: + pass + + async def drain(self) -> None: + pass + + def write_eof(self) -> None: + pass + + def close(self) -> None: + closed[self._name] = True + + async def wait_closed(self) -> None: + pass + + client_writer = _W("client") + target_writer = _W("target") + + async def _raise_wait(*args: object, **kwargs: object) -> None: + raise RuntimeError("pump wait boom") + + monkeypatch.setattr(_mod.asyncio, "wait", _raise_wait) + + await _blind_splice( + client_reader, + client_writer, # type: ignore[arg-type] + target_reader, + target_writer, # type: ignore[arg-type] + ) + + assert closed["client"] is True + assert closed["target"] is True + + +# --------------------------------------------------------------------------- +# Coverage: _handle_connect first-line CONNECT read timeout +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_handle_connect_first_line_timeout_closes_writer() -> None: + """First readline() (the CONNECT line itself) times out -> client_writer + is closed and neither MITM nor blind-tunnel dispatch runs.""" + import unittest.mock as mock + + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _handle_connect + + client_reader = asyncio.StreamReader() # No data fed -> readline() blocks forever. + + close_called = False + + class _TrackingWriter: + def get_extra_info(self, key: str, default: object = None) -> object: # noqa: ANN401 + return ("127.0.0.1", 1234) if key == "peername" else default + + def write(self, data: bytes) -> None: + pass + + async def drain(self) -> None: + pass + + def close(self) -> None: + nonlocal close_called + close_called = True + + async def wait_closed(self) -> None: + pass + + client_writer = _TrackingWriter() + + with mock.patch.object(_mod, "_CONNECT_TIMEOUT", 0.01): + await _handle_connect( + client_reader, + client_writer, # type: ignore[arg-type] + allowlist=frozenset(), + dispatch_port=1, + ) + + assert close_called, "client_writer.close() must be called on first-line CONNECT timeout" + + +# --------------------------------------------------------------------------- +# Coverage: Proxy-Authorization header is parsed off the CONNECT request +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_proxy_authorization_header_parsed(tmp_ca: tuple) -> None: + """CONNECT with a Proxy-Authorization header is accepted and tunnels bytes.""" + ca_key, ca_cert, _ = tmp_ca + echo_host = "127.0.0.1" + + async def echo_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + data = await asyncio.wait_for(reader.read(1024), timeout=5.0) + if data: + writer.write(data) + await writer.drain() + finally: + writer.close() + + echo_server = await asyncio.start_server(echo_handler, echo_host, 0) + echo_port = echo_server.sockets[0].getsockname()[1] + + terminator = AgyCONNECTTerminator( + allowlist=frozenset({ALLOWLIST_HOST}), # echo host NOT allowlisted -> blind tunnel + ca_key=ca_key, + ca_cert=ca_cert, + dispatch_port=1, + ) + await terminator.start() + try: + proxy_host, proxy_port = terminator.address + raw_reader, raw_writer = await asyncio.open_connection(proxy_host, proxy_port) + connect_req = ( + f"CONNECT {echo_host}:{echo_port} HTTP/1.1\r\n" + f"Host: {echo_host}:{echo_port}\r\n" + "Proxy-Authorization: Basic dXNlcjpwYXNz\r\n" + "\r\n" + ) + raw_writer.write(connect_req.encode()) + await raw_writer.drain() + response = await raw_reader.readline() + assert b"200" in response, f"Expected 200, got {response!r}" + await raw_reader.readline() # Drain the blank line separating status from body. + + payload = b"auth header parsed ok" + raw_writer.write(payload) + await raw_writer.drain() + received = await asyncio.wait_for(raw_reader.read(len(payload)), timeout=5.0) + assert received == payload + finally: + await terminator.stop() + echo_server.close() + await echo_server.wait_closed() + + +# --------------------------------------------------------------------------- +# Coverage: dispatch_port SUCCESS — ACK + blind-splice to loopback dispatch server +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dispatch_port_success_splices_to_dispatch_server(tmp_ca: tuple) -> None: + """Allowlisted CONNECT with a reachable dispatch_port: ACK written and raw + bytes are byte-spliced to the loopback dispatch server (no TLS).""" + ca_key, ca_cert, _ = tmp_ca + echo_host = "127.0.0.1" + + async def echo_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + data = await asyncio.wait_for(reader.read(1024), timeout=5.0) + if data: + writer.write(data) + await writer.drain() + finally: + writer.close() + + dispatch_server = await asyncio.start_server(echo_handler, echo_host, 0) + dispatch_port = dispatch_server.sockets[0].getsockname()[1] + + terminator = AgyCONNECTTerminator( + allowlist=frozenset({ALLOWLIST_HOST}), + ca_key=ca_key, + ca_cert=ca_cert, + dispatch_port=dispatch_port, + ) + await terminator.start() + try: + proxy_host, proxy_port = terminator.address + raw_reader, raw_writer = await asyncio.open_connection(proxy_host, proxy_port) + connect_req = f"CONNECT {ALLOWLIST_HOST}:443 HTTP/1.1\r\nHost: {ALLOWLIST_HOST}:443\r\n\r\n" + raw_writer.write(connect_req.encode()) + await raw_writer.drain() + response = await raw_reader.readline() + assert b"200" in response, f"Expected ACK 200, got {response!r}" + await raw_reader.readline() # Drain the blank line separating status from body. + + payload = b"raw bytes over dispatch splice" + raw_writer.write(payload) + await raw_writer.drain() + received = await asyncio.wait_for(raw_reader.read(len(payload)), timeout=5.0) + assert received == payload, f"Splice mismatch: {received!r} != {payload!r}" + # Let the target-side echo connection close, unblocking the server-side + # _blind_splice call so it runs to completion (its own return statement) + # before teardown — otherwise the background handler task may be torn + # down mid-flight. + await asyncio.sleep(0.05) + finally: + await terminator.stop() + dispatch_server.close() + await dispatch_server.wait_closed() + + +@pytest.mark.asyncio +async def test_dispatch_port_connect_failed_close_exception_swallowed() -> None: + """dispatch_connect_failed handling: if client_writer.close() itself also + raises, the inner except swallows it (headroom-vro.2: lines 461-462).""" + from headroom.proxy.agy_terminator import _handle_mitm + + # Bind then immediately close an ephemeral port so connecting to it + # deterministically raises ConnectionRefusedError (an OSError subclass). + probe = await asyncio.start_server(lambda r, w: None, "127.0.0.1", 0) + dead_port = probe.sockets[0].getsockname()[1] + probe.close() + await probe.wait_closed() + + close_called = False + + class _RaisingCloseWriter: + def write(self, data: bytes) -> None: + pass + + async def drain(self) -> None: + pass + + def close(self) -> None: + nonlocal close_called + close_called = True + raise RuntimeError("close boom") + + async def wait_closed(self) -> None: + pass + + client_reader = asyncio.StreamReader() + client_writer = _RaisingCloseWriter() + + # Must not raise, despite client_writer.close() raising inside the handler. + await _handle_mitm( + client_reader, + client_writer, # type: ignore[arg-type] + dead_port, + ) + assert close_called + + +# --------------------------------------------------------------------------- +# Coverage: dispatch_port connect failure closes the client after ACK +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_dispatch_port_unreachable_closes_client_after_ack(tmp_ca: tuple) -> None: + """Allowlisted CONNECT with an unreachable dispatch_port: ACK is still sent, + then the connect attempt fails and client_writer is closed.""" + ca_key, ca_cert, _ = tmp_ca + + # Bind then immediately close an ephemeral port so connecting to it + # deterministically raises ConnectionRefusedError (an OSError subclass). + probe = await asyncio.start_server(lambda r, w: None, "127.0.0.1", 0) + dead_port = probe.sockets[0].getsockname()[1] + probe.close() + await probe.wait_closed() + + terminator = AgyCONNECTTerminator( + allowlist=frozenset({ALLOWLIST_HOST}), + ca_key=ca_key, + ca_cert=ca_cert, + dispatch_port=dead_port, + ) + await terminator.start() + try: + proxy_host, proxy_port = terminator.address + raw_reader, raw_writer = await asyncio.open_connection(proxy_host, proxy_port) + connect_req = f"CONNECT {ALLOWLIST_HOST}:443 HTTP/1.1\r\nHost: {ALLOWLIST_HOST}:443\r\n\r\n" + raw_writer.write(connect_req.encode()) + await raw_writer.drain() + response = await raw_reader.readline() + assert b"200" in response, ( + f"Expected ACK 200 before dispatch-connect attempt, got {response!r}" + ) + await raw_reader.readline() # Drain the blank line separating status from body. + + # dispatch connect failed -> client_writer.close() -> EOF, no more data. + data = await asyncio.wait_for(raw_reader.read(10), timeout=5.0) + assert data == b"" + finally: + await terminator.stop() + + +# --------------------------------------------------------------------------- +# Coverage: blind tunnel upstream connect failure -> 502 Bad Gateway +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_blind_tunnel_connect_failure_returns_502( + tmp_ca: tuple, monkeypatch: pytest.MonkeyPatch +) -> None: + """Non-allowlisted CONNECT whose upstream open_connection raises OSError + must result in a 502 Bad Gateway response.""" + import headroom.proxy.agy_terminator as _mod + + ca_key, ca_cert, _ = tmp_ca + + terminator = AgyCONNECTTerminator( + allowlist=frozenset({ALLOWLIST_HOST}), # target NOT allowlisted -> blind tunnel + ca_key=ca_key, + ca_cert=ca_cert, + dispatch_port=1, + ) + await terminator.start() + try: + proxy_host, proxy_port = terminator.address + # Establish the client<->proxy connection BEFORE patching open_connection, + # since that patch also covers this very call target. + raw_reader, raw_writer = await asyncio.open_connection(proxy_host, proxy_port) + + async def _raise_oserror( + host: str, port: int, **kwargs: object + ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + raise OSError("connection refused") + + monkeypatch.setattr(_mod.asyncio, "open_connection", _raise_oserror) + + connect_req = ( + f"CONNECT {NON_ALLOWLIST_HOST}:443 HTTP/1.1\r\nHost: {NON_ALLOWLIST_HOST}:443\r\n\r\n" + ) + raw_writer.write(connect_req.encode()) + await raw_writer.drain() + response = await raw_reader.readline() + assert b"502 Bad Gateway" in response, f"Expected 502, got {response!r}" + finally: + await terminator.stop() + + +# --------------------------------------------------------------------------- +# Coverage: AgyCONNECTTerminator lifecycle edges +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_terminator_start_without_ca_key_uses_ensure_root_ca(tmp_path: object) -> None: + """Omitting ca_key/ca_cert triggers the ensure_root_ca(base_dir=...) start path + (local CA key generation under tmp_path; never touches real ~/.headroom).""" + terminator = AgyCONNECTTerminator(dispatch_port=1, base_dir=tmp_path) # type: ignore[arg-type] + await terminator.start() + try: + host, port = terminator.address + assert host == "127.0.0.1" + assert port > 0 + assert terminator._ca_key is not None + assert terminator._ca_cert is not None + finally: + await terminator.stop() + + +def test_address_before_start_raises_runtime_error() -> None: + """Reading .address before .start() raises RuntimeError.""" + terminator = AgyCONNECTTerminator(dispatch_port=1) + with pytest.raises(RuntimeError): + _ = terminator.address + + +@pytest.mark.asyncio +async def test_stop_before_start_is_noop() -> None: + """Calling .stop() before .start() (no server) is a no-op and does not raise.""" + terminator = AgyCONNECTTerminator(dispatch_port=1) + await terminator.stop() + assert terminator._server is None + + +# --------------------------------------------------------------------------- +# Regression: blind-tunnel target guard +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_blind_tunnel_refuses_self_connect() -> None: + """CONNECT to the terminator's own port must be refused, not tunnelled. + + Without the guard each nesting level costs two fds: a client that keeps + re-CONNECTing through the terminator to itself exhausts them. + """ + async with AgyCONNECTTerminator(dispatch_port=1) as term: + _, port = term.address + reader, writer = await asyncio.open_connection("127.0.0.1", port) + writer.write(f"CONNECT 127.0.0.1:{port} HTTP/1.1\r\n\r\n".encode()) + await writer.drain() + response = await asyncio.wait_for(reader.read(64), timeout=5) + writer.close() + + assert b"403" in response, f"expected 403 for self-connect, got {response!r}" + + +@pytest.mark.asyncio +async def test_blind_tunnel_refuses_link_local_metadata_host() -> None: + """169.254.169.254 (cloud instance metadata) must not be reachable.""" + async with AgyCONNECTTerminator(dispatch_port=1) as term: + _, port = term.address + reader, writer = await asyncio.open_connection("127.0.0.1", port) + writer.write(b"CONNECT 169.254.169.254:80 HTTP/1.1\r\n\r\n") + await writer.drain() + response = await asyncio.wait_for(reader.read(64), timeout=5) + writer.close() + + assert b"403" in response, f"expected 403 for link-local target, got {response!r}" + + +@pytest.mark.parametrize( + ("proxy_url", "expected_port"), + [ + ("http://proxy.corp", 80), + ("https://proxy.corp", 443), + ("http://proxy.corp:3128", 3128), + ], +) +@pytest.mark.asyncio +async def test_upstream_proxy_port_defaults_follow_scheme( + monkeypatch: pytest.MonkeyPatch, proxy_url: str, expected_port: int +) -> None: + """A port-less HTTPS_PROXY must be dialled per scheme, not always on :443.""" + import headroom.proxy.agy_terminator as _mod + + dialled: list[int] = [] + + async def _spy( + proxy_host: str, + proxy_port: int, + target_host: str, + target_port: int, + proxy_auth: str | None, + ssl_context: ssl.SSLContext | None, + ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + dialled.append(proxy_port) + raise OSError("stop here — the dialled port is what matters") + + monkeypatch.setenv("HTTPS_PROXY", proxy_url) + monkeypatch.setattr(_mod, "_connect_via_upstream_proxy", _spy) + + async with AgyCONNECTTerminator(dispatch_port=1) as term: + _, port = term.address + reader, writer = await asyncio.open_connection("127.0.0.1", port) + writer.write(b"CONNECT example.com:443 HTTP/1.1\r\n\r\n") + await writer.drain() + await asyncio.wait_for(reader.read(64), timeout=5) + writer.close() + + assert dialled == [expected_port] + + +# --------------------------------------------------------------------------- +# Coverage: upstream CONNECT wire bytes + https-proxy TLS parameters +# --------------------------------------------------------------------------- + + +class _SpliceCompatWriter: + """client_writer stub satisfying the subset of StreamWriter used by + _handle_blind_tunnel and _blind_splice: write/drain/close/wait_closed for + the 200-response, write_eof for the splice's finally-block once the + target side reaches EOF. + """ + + def write(self, data: bytes) -> None: + pass + + async def drain(self) -> None: + pass + + def write_eof(self) -> None: + pass + + def close(self) -> None: + pass + + async def wait_closed(self) -> None: + pass + + +async def _start_fake_upstream_proxy() -> tuple[asyncio.AbstractServer, str, int, list[bytes]]: + """Bind a fake upstream proxy on an ephemeral loopback port. + + Records the raw bytes of the CONNECT request it receives, replies + 200 Connection Established, then closes -- the resulting target-side EOF + is what lets _blind_splice's FIRST_COMPLETED race return promptly instead + of hanging on an idle tunnel. + """ + received: list[bytes] = [] + + async def _handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + try: + data = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), timeout=5.0) + except (asyncio.IncompleteReadError, asyncio.TimeoutError): + data = b"" + received.append(data) + writer.write(b"HTTP/1.1 200 Connection Established\r\n\r\n") + await writer.drain() + writer.close() + + server = await asyncio.start_server(_handler, host="127.0.0.1", port=0) + host, port = server.sockets[0].getsockname()[:2] + return server, host, port, received + + +@pytest.mark.asyncio +async def test_blind_tunnel_upstream_proxy_url_credential_emitted( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """(a) HTTPS_PROXY URL userinfo is derived into Proxy-Authorization.""" + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _handle_blind_tunnel + + server, proxy_host, proxy_port, received = await _start_fake_upstream_proxy() + # HARNESS TRAP: a fake proxy on 127.0.0.1 is refused by the self-loop + # guard (_is_loopback) before _connect_via_upstream_proxy is ever + # reached. This test targets auth-derivation, not the loopback guard + # (which has its own dedicated tests), so bypass it explicitly. + monkeypatch.setattr(_mod, "_is_loopback", lambda host: False) + monkeypatch.setenv("HTTPS_PROXY", f"http://user:secret@{proxy_host}:{proxy_port}") + client_reader = asyncio.StreamReader() + client_reader.feed_eof() + + try: + with caplog.at_level("DEBUG"): + await asyncio.wait_for( + _handle_blind_tunnel( + client_reader, _SpliceCompatWriter(), "target.example.com", 443, None + ), + timeout=5.0, + ) + finally: + server.close() + await server.wait_closed() + + assert received, "fake proxy must have received a CONNECT request" + assert b"Proxy-Authorization: Basic dXNlcjpzZWNyZXQ=\r\n" in received[0] + assert "secret" not in caplog.text + + +@pytest.mark.asyncio +async def test_blind_tunnel_upstream_proxy_percent_decodes_credential( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """(b) Percent-encoded userinfo in HTTPS_PROXY is decoded before Basic auth.""" + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _handle_blind_tunnel + + server, proxy_host, proxy_port, received = await _start_fake_upstream_proxy() + monkeypatch.setattr(_mod, "_is_loopback", lambda host: False) + # user%40corp:p%40ss decodes to "user@corp:p@ss". + monkeypatch.setenv("HTTPS_PROXY", f"http://user%40corp:p%40ss@{proxy_host}:{proxy_port}") + client_reader = asyncio.StreamReader() + client_reader.feed_eof() + + try: + with caplog.at_level("DEBUG"): + await asyncio.wait_for( + _handle_blind_tunnel( + client_reader, _SpliceCompatWriter(), "target.example.com", 443, None + ), + timeout=5.0, + ) + finally: + server.close() + await server.wait_closed() + + assert received, "fake proxy must have received a CONNECT request" + assert b"Proxy-Authorization: Basic dXNlckBjb3JwOnBAc3M=\r\n" in received[0] + assert "p@ss" not in caplog.text + + +@pytest.mark.asyncio +async def test_blind_tunnel_upstream_proxy_url_credential_beats_inbound_header( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """(c) HTTPS_PROXY URL userinfo takes precedence over an inbound header.""" + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _handle_blind_tunnel + + server, proxy_host, proxy_port, received = await _start_fake_upstream_proxy() + monkeypatch.setattr(_mod, "_is_loopback", lambda host: False) + monkeypatch.setenv("HTTPS_PROXY", f"http://user:secret@{proxy_host}:{proxy_port}") + client_reader = asyncio.StreamReader() + client_reader.feed_eof() + inbound = "Basic aW5ib3VuZDpzZWNyZXQ=" # "inbound:secret" -- must be shadowed by the URL cred + + try: + with caplog.at_level("DEBUG"): + await asyncio.wait_for( + _handle_blind_tunnel( + client_reader, _SpliceCompatWriter(), "target.example.com", 443, inbound + ), + timeout=5.0, + ) + finally: + server.close() + await server.wait_closed() + + assert received, "fake proxy must have received a CONNECT request" + assert b"Proxy-Authorization: Basic dXNlcjpzZWNyZXQ=\r\n" in received[0] + assert inbound.encode() not in received[0] + assert "secret" not in caplog.text + + +@pytest.mark.asyncio +async def test_blind_tunnel_upstream_proxy_inbound_header_used_when_no_userinfo( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """(d) No userinfo in HTTPS_PROXY -> the inbound header is forwarded as-is.""" + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _handle_blind_tunnel + + server, proxy_host, proxy_port, received = await _start_fake_upstream_proxy() + monkeypatch.setattr(_mod, "_is_loopback", lambda host: False) + monkeypatch.setenv("HTTPS_PROXY", f"http://{proxy_host}:{proxy_port}") + client_reader = asyncio.StreamReader() + client_reader.feed_eof() + inbound = "Basic aW5ib3VuZDpzZWNyZXQ=" # "inbound:secret" + + try: + with caplog.at_level("DEBUG"): + await asyncio.wait_for( + _handle_blind_tunnel( + client_reader, _SpliceCompatWriter(), "target.example.com", 443, inbound + ), + timeout=5.0, + ) + finally: + server.close() + await server.wait_closed() + + assert received, "fake proxy must have received a CONNECT request" + assert f"Proxy-Authorization: {inbound}\r\n".encode() in received[0] + assert "secret" not in caplog.text + + +@pytest.mark.asyncio +async def test_blind_tunnel_upstream_proxy_no_credential_no_header( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(e) No URL userinfo and no inbound header -> no Proxy-Authorization line at all.""" + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _handle_blind_tunnel + + server, proxy_host, proxy_port, received = await _start_fake_upstream_proxy() + monkeypatch.setattr(_mod, "_is_loopback", lambda host: False) + monkeypatch.setenv("HTTPS_PROXY", f"http://{proxy_host}:{proxy_port}") + client_reader = asyncio.StreamReader() + client_reader.feed_eof() + + try: + await asyncio.wait_for( + _handle_blind_tunnel( + client_reader, _SpliceCompatWriter(), "target.example.com", 443, None + ), + timeout=5.0, + ) + finally: + server.close() + await server.wait_closed() + + assert received, "fake proxy must have received a CONNECT request" + assert b"Proxy-Authorization" not in received[0] + + +@pytest.mark.asyncio +async def test_blind_tunnel_upstream_proxy_empty_password_no_crash( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + """(e2) Username with an EMPTY password -> Basic b64("user:"), no crash from `password or ''`.""" + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _handle_blind_tunnel + + server, proxy_host, proxy_port, received = await _start_fake_upstream_proxy() + monkeypatch.setattr(_mod, "_is_loopback", lambda host: False) + monkeypatch.setenv("HTTPS_PROXY", f"http://user:@{proxy_host}:{proxy_port}") + client_reader = asyncio.StreamReader() + client_reader.feed_eof() + + try: + with caplog.at_level("DEBUG"): + await asyncio.wait_for( + _handle_blind_tunnel( + client_reader, _SpliceCompatWriter(), "target.example.com", 443, None + ), + timeout=5.0, + ) + finally: + server.close() + await server.wait_closed() + + assert received, "fake proxy must have received a CONNECT request" + assert b"Proxy-Authorization: Basic dXNlcjo=\r\n" in received[0] + assert "dXNlcjo=" not in caplog.text + + +@pytest.mark.asyncio +async def test_blind_tunnel_https_upstream_proxy_tls_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(f) https:// upstream proxy dials over TLS with verified defaults, SNI to the + proxy host, and ALPN pinned to http/1.1. + + ALPN cannot be read back off a real ssl.SSLContext (set_alpn_protocols() + forwards to the C layer and stores nothing readable; selected_alpn_protocol() + only exists on a post-handshake SSLObject, which a mocked open_connection + never produces). So this test proves the two halves separately: + verify_mode/check_hostname/server_hostname off a REAL default context + (first act), and the ALPN pin via set_alpn_protocols() call-args on a + MOCKED context (second act). + """ + import unittest.mock as mock + + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _handle_blind_tunnel + + monkeypatch.setattr(_mod, "_is_loopback", lambda host: False) + monkeypatch.setenv("HTTPS_PROXY", "https://proxy.corp.example:9443") + + # --- act 1: real default context -> verify_mode / check_hostname / SNI --- + captured: dict[str, object] = {} + + async def _fake_open_connection_real_ctx( + host: str, port: int, *, ssl: object = None, server_hostname: object = None, **_: object + ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + captured["ssl"] = ssl + captured["server_hostname"] = server_hostname + raise OSError("stop before real I/O -- the TLS context passed in is what's under test") + + with mock.patch.object(_mod.asyncio, "open_connection", _fake_open_connection_real_ctx): + client_reader = asyncio.StreamReader() + client_reader.feed_eof() + await asyncio.wait_for( + _handle_blind_tunnel( + client_reader, _SpliceCompatWriter(), "target.example.com", 443, None + ), + timeout=5.0, + ) + + ctx = captured["ssl"] + assert isinstance(ctx, ssl.SSLContext) + assert ctx.verify_mode == ssl.CERT_REQUIRED + assert ctx.check_hostname is True + assert captured["server_hostname"] == "proxy.corp.example" + + # --- act 2: mocked context -> ALPN pin asserted via call-args --- + mock_ctx = mock.MagicMock(spec=ssl.SSLContext) + + async def _fake_open_connection_mock_ctx( + host: str, port: int, **_: object + ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + raise OSError("stop before real I/O -- ALPN pinning is what's under test") + + with ( + mock.patch.object(_mod.ssl, "create_default_context", lambda: mock_ctx), + mock.patch.object(_mod.asyncio, "open_connection", _fake_open_connection_mock_ctx), + ): + client_reader = asyncio.StreamReader() + client_reader.feed_eof() + await asyncio.wait_for( + _handle_blind_tunnel( + client_reader, _SpliceCompatWriter(), "target.example.com", 443, None + ), + timeout=5.0, + ) + + mock_ctx.set_alpn_protocols.assert_called_once_with(["http/1.1"]) + + +@pytest.mark.asyncio +async def test_blind_tunnel_http_upstream_proxy_dials_without_tls( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """(g) http:// upstream proxy dials plaintext -- ssl=None, no server_hostname/SNI.""" + import unittest.mock as mock + + import headroom.proxy.agy_terminator as _mod + from headroom.proxy.agy_terminator import _handle_blind_tunnel + + monkeypatch.setattr(_mod, "_is_loopback", lambda host: False) + monkeypatch.setenv("HTTPS_PROXY", "http://proxy.corp.example:3128") + + captured: dict[str, object] = {} + + async def _fake_open_connection( + host: str, port: int, *, ssl: object = None, server_hostname: object = None, **_: object + ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: + captured["ssl"] = ssl + captured["server_hostname"] = server_hostname + raise OSError("stop before real I/O -- the ssl kwarg is what's under test") + + client_reader = asyncio.StreamReader() + client_reader.feed_eof() + + with mock.patch.object(_mod.asyncio, "open_connection", _fake_open_connection): + await asyncio.wait_for( + _handle_blind_tunnel( + client_reader, _SpliceCompatWriter(), "target.example.com", 443, None + ), + timeout=5.0, + ) + + assert captured["ssl"] is None + assert captured["server_hostname"] is None + + +def test_upstream_proxy_auth_rejects_socks5_scheme() -> None: + """(h) socks5:// upstream proxy URL never derives a Proxy-Authorization header. + + _upstream_proxy_auth is pure string/URL logic (no socket I/O), per its own + docstring, so it is unit-tested directly rather than through + _handle_blind_tunnel (which 403s a non-http(s) scheme before this helper's + return value would even be used). + """ + import urllib.parse + + from headroom.proxy.agy_terminator import _upstream_proxy_auth + + parsed = urllib.parse.urlparse("socks5://user:secret@proxy.corp.example:1080") + assert _upstream_proxy_auth(parsed, None) is None + assert _upstream_proxy_auth(parsed, "Basic aW5ib3VuZDpzZWNyZXQ=") is None diff --git a/tests/test_cli/test_wrap_agy_port_alias.py b/tests/test_cli/test_wrap_agy_port_alias.py new file mode 100644 index 000000000..179b965b9 --- /dev/null +++ b/tests/test_cli/test_wrap_agy_port_alias.py @@ -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 == [] diff --git a/tests/test_provider_proxy_routes.py b/tests/test_provider_proxy_routes.py index 5f561ae06..316852d09 100644 --- a/tests/test_provider_proxy_routes.py +++ b/tests/test_provider_proxy_routes.py @@ -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, ...]]] = [] diff --git a/tests/test_proxy_agy_compression.py b/tests/test_proxy_agy_compression.py new file mode 100644 index 000000000..c7743111d --- /dev/null +++ b/tests/test_proxy_agy_compression.py @@ -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 +# --------------------------------------------------------------------------- diff --git a/tests/test_proxy_google_cloudcode_route_aliases.py b/tests/test_proxy_google_cloudcode_route_aliases.py index b94e918fd..1b7431e11 100644 --- a/tests/test_proxy_google_cloudcode_route_aliases.py +++ b/tests/test_proxy_google_cloudcode_route_aliases.py @@ -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" diff --git a/tests/test_wrap_agy.py b/tests/test_wrap_agy.py new file mode 100644 index 000000000..e03d6e0d5 --- /dev/null +++ b/tests/test_wrap_agy.py @@ -0,0 +1,1541 @@ +"""Tests for headroom wrap agy / unwrap agy.""" + +from __future__ import annotations + +import logging +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +from click.testing import CliRunner + +from headroom.cli.wrap import _PROXY_URL_REDACTED_PLACEHOLDER, redact_proxy_url + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +_WRAP_MODULE = "headroom.cli.wrap" + + +# --------------------------------------------------------------------------- +# headroom wrap agy — CLI integration tests +# --------------------------------------------------------------------------- + + +def _get_main(): + from headroom.cli.main import main + + return main + + +@pytest.fixture(autouse=True) +def _never_start_a_real_proxy(monkeypatch: pytest.MonkeyPatch) -> None: + """Stub ``_ensure_proxy`` for every test in this module. + + ``_ensure_proxy`` -> ``_start_proxy`` calls ``subprocess.Popen`` directly, + which none of this file's per-test ``subprocess.run`` stubs touch. Left + unstubbed, any test that drives the full ``agy``/``unwrap`` CLI spawns a + real ``headroom.cli proxy`` subprocess that binds a real port — on a dev + machine with a live proxy already on that port, this evicts it. No test + in this file exercises ``_ensure_proxy`` itself (that's covered + elsewhere), so stubbing it here is safe for all of them. + """ + import headroom.cli.wrap as wrap_mod + + def _fake_ensure_proxy(port, no_proxy=False, **_kwargs): + return None, port + + monkeypatch.setattr(wrap_mod, "_ensure_proxy", _fake_ensure_proxy) + + +class TestWrapAgyBinaryMissing: + """Binary-missing path must exit 1 with install hint.""" + + def test_exits_1_when_agy_not_found(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("shutil.which", lambda _: None) + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy"]) + assert result.exit_code == 1 + + def test_prints_install_hint_when_agy_not_found(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("shutil.which", lambda _: None) + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy"]) + assert "agy" in result.output.lower() or "install" in result.output.lower() + + +class TestWrapAgyRustBackendFails: + """Rust backend must hard-fail with a clear message.""" + + def _run_with_rust_backend(self, monkeypatch: pytest.MonkeyPatch, via_env: bool): + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/agy" if name == "agy" else None) + if via_env: + monkeypatch.setenv("HEADROOM_BACKEND", "rust") + runner = CliRunner() + args = ["wrap", "agy"] if not via_env else ["wrap", "agy"] + if not via_env: + args += ["--backend", "rust"] + return runner.invoke(_get_main(), args) + + def test_rust_backend_flag_exits_1(self, monkeypatch: pytest.MonkeyPatch) -> None: + result = self._run_with_rust_backend(monkeypatch, via_env=False) + assert result.exit_code == 1 + + def test_rust_backend_flag_prints_clear_message(self, monkeypatch: pytest.MonkeyPatch) -> None: + result = self._run_with_rust_backend(monkeypatch, via_env=False) + output = result.output.lower() + assert "rust" in output or "python" in output or "not supported" in output + + def test_rust_backend_env_exits_1(self, monkeypatch: pytest.MonkeyPatch) -> None: + result = self._run_with_rust_backend(monkeypatch, via_env=True) + assert result.exit_code == 1 + + +class TestWrapAgyDisclosureBanner: + """TLS interception disclosure banner must name the intercepted host.""" + + _INTERCEPTED_HOST = "daily-cloudcode-pa.googleapis.com" + + def _invoke_agy(self, monkeypatch: pytest.MonkeyPatch, extra_args: list[str] | None = None): + """Invoke wrap agy with servers and subprocess fully stubbed out.""" + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/agy" if name == "agy" else None) + + # Stub the lifecycle helper so no real servers start + import headroom.cli.wrap as wrap_mod + + fake_servers = MagicMock() + fake_servers.terminator.address = ("127.0.0.1", 54321) + fake_servers.dispatch.address = ("127.0.0.1", 54322) + # No retrieve listener here: this test only checks the disclosure + # banner, and a real port would trigger MCP registration against the + # real ~/.gemini. retrieve_port=None makes agy() skip registration. + fake_servers.retrieve_port = None + + def fake_start_agy_servers( + ca_key, ca_cert, base_dir=None, *, start_retrieve=False, project=None + ): + return fake_servers + + monkeypatch.setattr(wrap_mod, "_start_agy_servers", fake_start_agy_servers) + monkeypatch.setattr(wrap_mod, "_stop_agy_servers", lambda s: None) + + # Stub ensure_root_ca + build_combined_bundle + import datetime + + from cryptography import x509 + from cryptography.hazmat.primitives import hashes + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.x509.oid import NameOID + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + now = datetime.datetime.now(tz=datetime.timezone.utc) + subject = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "Test CA")]) + cert = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(subject) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now) + .not_valid_after(now + datetime.timedelta(days=1)) + .add_extension(x509.BasicConstraints(ca=True, path_length=0), critical=True) + .sign(key, hashes.SHA256()) + ) + + monkeypatch.setattr( + "headroom.proxy.agy_ca.ensure_root_ca", + lambda base_dir=None: (key, cert, Path("/tmp/ca.key"), Path("/tmp/ca.crt")), + ) + monkeypatch.setattr( + "headroom.proxy.agy_ca.build_combined_bundle", + lambda base_dir=None, corp_env_vars=None: Path("/tmp/bundle.pem"), + ) + + # Stub subprocess.run so agy never actually launches + monkeypatch.setattr("subprocess.run", lambda *a, **kw: MagicMock(returncode=0)) + + runner = CliRunner() + args = ["wrap", "agy"] + (extra_args or []) + return runner.invoke(_get_main(), args, catch_exceptions=False) + + def test_disclosure_banner_names_intercepted_host( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + result = self._invoke_agy(monkeypatch) + assert self._INTERCEPTED_HOST in result.output + + def test_disclosure_banner_names_every_allowlisted_host( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Consent surface must not understate interception: the banner must + name EVERY host the terminator's allowlist will TLS-terminate, not + just the primary one.""" + from headroom.proxy.agy_terminator import DEFAULT_ALLOWLIST + + result = self._invoke_agy(monkeypatch) + for host in DEFAULT_ALLOWLIST: + assert host in result.output, f"disclosure omits intercepted host {host}" + + def test_disclosure_banner_mentions_no_intercept_option( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + result = self._invoke_agy(monkeypatch) + assert "--no-intercept" in result.output + + def test_disclosure_banner_mentions_unwrap(self, monkeypatch: pytest.MonkeyPatch) -> None: + result = self._invoke_agy(monkeypatch) + assert "unwrap" in result.output.lower() + + +class TestWrapAgyMcpFlagParity: + """agy exposes the same MCP opt-out surface as its sibling subcommands.""" + + def test_no_mcp_is_offered_like_the_siblings(self) -> None: + result = CliRunner().invoke(_get_main(), ["wrap", "agy", "--help"]) + + assert result.exit_code == 0 + assert "--no-mcp" in result.output + assert "--no-serena" in result.output + + def test_no_mcp_promises_the_same_thing_as_the_siblings(self) -> None: + """Same flag, same promise — drift between siblings is the bug being fixed.""" + agy_help = " ".join( + CliRunner().invoke(_get_main(), ["wrap", "agy", "--help"]).output.split() + ) + assert "--no-mcp Skip headroom MCP server registration" in agy_help + + +class TestWrapAgyNoIntercept: + """--no-intercept flag must change behavior (no MITM server startup).""" + + def test_no_intercept_does_not_start_servers(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/agy" if name == "agy" else None) + + import headroom.cli.wrap as wrap_mod + + server_started = [] + + def fake_start(ca_key, ca_cert, base_dir=None, *, start_retrieve=False, project=None): + server_started.append(True) + raise AssertionError("Servers must NOT start in --no-intercept mode") + + monkeypatch.setattr(wrap_mod, "_start_agy_servers", fake_start) + monkeypatch.setattr(wrap_mod, "_stop_agy_servers", lambda s: None) + monkeypatch.setattr("subprocess.run", lambda *a, **kw: MagicMock(returncode=0)) + + runner = CliRunner() + runner.invoke(_get_main(), ["wrap", "agy", "--no-intercept"]) + # Must not have started servers (no AssertionError bubbled = no start call) + assert not server_started + + +class TestWrapAgySignalTeardown: + """SIGTERM during the agy run must tear the MITM servers down (and the + pre-existing handlers must be restored afterwards).""" + + def _stub_ca(self, monkeypatch: pytest.MonkeyPatch) -> None: + import datetime + + from cryptography import x509 + from cryptography.hazmat.primitives import hashes + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.x509.oid import NameOID + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + now = datetime.datetime.now(tz=datetime.timezone.utc) + subject = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "Test CA")]) + cert = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(subject) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now) + .not_valid_after(now + datetime.timedelta(days=1)) + .add_extension(x509.BasicConstraints(ca=True, path_length=0), critical=True) + .sign(key, hashes.SHA256()) + ) + monkeypatch.setattr( + "headroom.proxy.agy_ca.ensure_root_ca", + lambda base_dir=None: (key, cert, Path("/tmp/ca.key"), Path("/tmp/ca.crt")), + ) + monkeypatch.setattr( + "headroom.proxy.agy_ca.build_combined_bundle", + lambda base_dir=None, corp_env_vars=None: Path("/tmp/bundle.pem"), + ) + + def test_sigterm_during_run_tears_down_and_restores_handlers( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + import signal + + import headroom.cli.wrap as wrap_mod + + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/agy" if name == "agy" else None) + self._stub_ca(monkeypatch) + + fake_servers = MagicMock() + fake_servers.terminator.address = ("127.0.0.1", 54321) + fake_servers.dispatch.address = ("127.0.0.1", 54322) + # No retrieve listener: keep this signal-teardown test focused and avoid + # touching the real ~/.gemini via MCP registration. + fake_servers.retrieve_port = None + monkeypatch.setattr( + wrap_mod, + "_start_agy_servers", + lambda ca_key, ca_cert, base_dir=None, *, start_retrieve=False, project=None: ( + fake_servers + ), + ) + + stop_calls: list[object] = [] + monkeypatch.setattr(wrap_mod, "_stop_agy_servers", lambda s: stop_calls.append(s)) + + captured: dict[str, object] = {} + + def fake_run(*_a, **_kw): + # Simulate agy receiving SIGTERM mid-run: invoke the handler that + # production installed. It must stop the servers and raise SystemExit(143). + captured["sigterm"] = signal.getsignal(signal.SIGTERM) + captured["sigint"] = signal.getsignal(signal.SIGINT) + handler = captured["sigterm"] + assert callable(handler) + handler(signal.SIGTERM, None) # raises SystemExit(143) + raise AssertionError("SIGTERM handler did not raise") # pragma: no cover + + monkeypatch.setattr("subprocess.run", fake_run) + + original_sigterm = signal.getsignal(signal.SIGTERM) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy"]) + + # The SIGTERM handler raised SystemExit(143) -> that is the exit code. + assert result.exit_code == 143 + # SIGINT was delegated to agy via the ignore-child handler. + assert captured["sigint"] is wrap_mod._ignore_child_sigint + # The installed SIGTERM handler was a real handler (not default/ignore). + assert captured["sigterm"] not in (signal.SIG_DFL, signal.SIG_IGN) + # Servers were stopped (handler + finally both call _stop_agy_servers). + assert len(stop_calls) >= 1 + # Prior SIGTERM handler restored — no leak into the host process. + assert signal.getsignal(signal.SIGTERM) is original_sigterm + + +# --------------------------------------------------------------------------- +# headroom unwrap agy +# --------------------------------------------------------------------------- + + +class TestUnwrapAgy: + """unwrap agy reverts GEMINI.md block and MCP registration.""" + + def test_unwrap_agy_exits_0(self) -> None: + runner = CliRunner() + result = runner.invoke(_get_main(), ["unwrap", "agy"]) + assert result.exit_code == 0 + + def test_unwrap_agy_prints_status_message(self) -> None: + runner = CliRunner() + result = runner.invoke(_get_main(), ["unwrap", "agy"]) + # Should have some output acknowledging the command ran + assert result.output.strip() != "" + + +# --------------------------------------------------------------------------- +# T9: GEMINI.md block removal (legacy blocks from pre-2677 installs) +# --------------------------------------------------------------------------- + + +class TestGeminiMdBlock: + """_remove_gemini_md_block deletes only the Headroom block. + + `wrap agy` no longer writes a GEMINI.md block (the rtk context-tool + instructions it carried were removed upstream), but `unwrap agy` must still + clean a block an older install left behind. + """ + + def _get_helpers(self): + from headroom.cli.wrap import ( + _AGY_GEMINI_BLOCK_END, + _AGY_GEMINI_BLOCK_START, + _remove_gemini_md_block, + ) + + return (_remove_gemini_md_block, _AGY_GEMINI_BLOCK_START, _AGY_GEMINI_BLOCK_END) + + def _write_legacy(self, gemini_md: Path, user_text: str = "") -> None: + """Write a GEMINI.md exactly as an older `wrap agy` left it.""" + _, start, end = self._get_helpers() + block = f"{start}\n## Headroom\nContext.\n{end}\n" + gemini_md.parent.mkdir(parents=True, exist_ok=True) + gemini_md.write_text(f"{user_text}\n\n{block}" if user_text else block) + + def test_remove_deletes_only_headroom_block(self, tmp_path: Path) -> None: + remove, start, end = self._get_helpers() + gemini_md = tmp_path / "GEMINI.md" + self._write_legacy(gemini_md, "# User content\nKeep this.") + assert remove(gemini_md, verbose=False) is True + text = gemini_md.read_text() + assert "# User content" in text + assert "Keep this." in text + assert start not in text + assert end not in text + + def test_remove_is_idempotent(self, tmp_path: Path) -> None: + remove, _, _ = self._get_helpers() + gemini_md = tmp_path / "GEMINI.md" + self._write_legacy(gemini_md) + assert remove(gemini_md, verbose=False) is True + assert remove(gemini_md, verbose=False) is False + + def test_remove_returns_false_when_file_absent(self, tmp_path: Path) -> None: + remove, _, _ = self._get_helpers() + gemini_md = tmp_path / "GEMINI.md" + assert remove(gemini_md, verbose=False) is False + + def test_remove_returns_false_when_no_block(self, tmp_path: Path) -> None: + remove, _, _ = self._get_helpers() + gemini_md = tmp_path / "GEMINI.md" + gemini_md.write_text("# User content only\n") + assert remove(gemini_md, verbose=False) is False + + +# --------------------------------------------------------------------------- +# T9: unwrap agy reverts GEMINI.md block (integration via CLI runner) +# --------------------------------------------------------------------------- + + +class TestUnwrapAgyReverts: + """unwrap agy removes headroom block; preserves user content; is idempotent.""" + + def test_unwrap_removes_gemini_md_block( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + from headroom.cli.wrap import ( + _AGY_GEMINI_BLOCK_END, + _AGY_GEMINI_BLOCK_START, + ) + + gemini_md = tmp_path / ".gemini" / "GEMINI.md" + gemini_md.parent.mkdir(parents=True, exist_ok=True) + gemini_md.write_text( + f"# User content\n\n{_AGY_GEMINI_BLOCK_START}\n## Headroom\n{_AGY_GEMINI_BLOCK_END}\n" + ) + monkeypatch.setattr(Path, "home", lambda: tmp_path) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["unwrap", "agy"]) + assert result.exit_code == 0 + text = gemini_md.read_text() + assert _AGY_GEMINI_BLOCK_START not in text + assert "# User content" in text + + def test_unwrap_is_idempotent_when_already_clean( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + gemini_md = tmp_path / ".gemini" / "GEMINI.md" + gemini_md.parent.mkdir(parents=True, exist_ok=True) + gemini_md.write_text("# User content only\n") + monkeypatch.setattr(Path, "home", lambda: tmp_path) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["unwrap", "agy"]) + assert result.exit_code == 0 + + +# --------------------------------------------------------------------------- +# T9: MCP retrieve tool wiring (N/A-v1 for per-run ephemeral port) +# --------------------------------------------------------------------------- + + +class TestAgyMcpRetrieveNa: + """Verify wrap agy does NOT register a retrieve MCP entry outside the + interactive MITM path. + + Interactive MITM registers a persistent, ledger-recorded headroom MCP + retrieve entry (stable spec, on-disk store resolution — see + TestAgyRetrieveMcpWiring). But --no-intercept (passthrough) starts no + servers, so on a fresh config it registers nothing. + """ + + def test_agy_mcp_config_not_written_during_wrap_no_intercept( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """--no-intercept path: no MCP registration should happen.""" + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/agy" if name == "agy" else None) + + # Redirect HOME so we never touch the real ~/.gemini. + monkeypatch.setattr(Path, "home", lambda: tmp_path) + monkeypatch.setattr("subprocess.run", lambda *a, **kw: MagicMock(returncode=0)) + + runner = CliRunner() + runner.invoke(_get_main(), ["wrap", "agy", "--no-intercept"]) + + # agy 1.1.x read-path (migrated from .gemini/antigravity-cli/). + mcp_config = tmp_path / ".gemini" / "config" / "mcp_config.json" + # No per-run registration: file must not exist OR must not contain an + # ephemeral headroom entry (port range check omitted; just assert no + # ephemeral entry was written for "headroom"). + if mcp_config.exists(): + import json + + cfg = json.loads(mcp_config.read_text()) + assert "headroom" not in cfg.get("mcpServers", {}), ( + "wrap agy must not register an ephemeral headroom MCP entry" + ) + + +# --------------------------------------------------------------------------- +# T9 Fix 1: Serena MCP WIRED for agy (full MITM path, all servers stubbed) +# --------------------------------------------------------------------------- + + +def _stub_agy_mitm_run( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + *, + with_uvx: bool = True, +): + """Stub the full agy MITM run so wrap agy reaches the MCP wiring. + + Redirects HOME to tmp_path (isolating ~/.gemini and ~/.headroom ledger), + stubs server lifecycle + CA + subprocess so nothing real launches. When + ``with_uvx`` is True, shutil.which("uvx") resolves so _setup_serena_mcp + proceeds. Pre-creates ~/.gemini/antigravity-cli so AgyRegistrar.detect() + returns True. + """ + import datetime + + from cryptography import x509 + from cryptography.hazmat.primitives import hashes + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.x509.oid import NameOID + + import headroom.cli.wrap as wrap_mod + + monkeypatch.setattr(Path, "home", lambda: tmp_path) + # Pre-create the agy config dir so AgyRegistrar.detect() is True. + (tmp_path / ".gemini" / "antigravity-cli").mkdir(parents=True, exist_ok=True) + + def fake_which(name: str): + if name == "agy": + return "/usr/bin/agy" + if name == "uvx" and with_uvx: + return "/usr/bin/uvx" + return None + + monkeypatch.setattr("shutil.which", fake_which) + + fake_servers = MagicMock() + fake_servers.terminator.address = ("127.0.0.1", 54321) + fake_servers.dispatch.address = ("127.0.0.1", 54322) + # Interactive-mode retrieve listener port (a real int so the headroom MCP + # spec gets a well-formed loopback URL). _agy_start_calls records the + # start_retrieve flag each call so tests can assert print-mode skips it. + fake_servers.retrieve_port = 54323 + fake_servers.retrieve = MagicMock() + + def _fake_start_agy_servers( + ca_key, ca_cert, base_dir=None, *, start_retrieve=False, project=None + ): + _agy_start_calls.append(start_retrieve) + # In print mode the real server starts no retrieve listener: model that + # so the agy() guard (servers.retrieve_port is not None) holds. + if not start_retrieve: + fake_servers.retrieve = None + fake_servers.retrieve_port = None + else: + fake_servers.retrieve = MagicMock() + fake_servers.retrieve_port = 54323 + return fake_servers + + _agy_start_calls: list[bool] = [] + fake_servers._agy_start_calls = _agy_start_calls + monkeypatch.setattr(wrap_mod, "_start_agy_servers", _fake_start_agy_servers) + monkeypatch.setattr(wrap_mod, "_stop_agy_servers", lambda s: None) + # Default the MCP handshake smoke check to PASS so interactive registrations + # survive; individual tests override this when they exercise the failure path. + monkeypatch.setattr(wrap_mod, "_smoke_verify_mcp_handshake", lambda *a, **kw: True) + # Default the stubbed agy to a known-good version so print-mode MCP wiring is + # exercised (headroom-37g.37 gates print-mode MCP on agy >= 1.0.16). The + # suppress/purge path for older/unknown agy is covered separately in + # tests/test_agy_print_mode_version_gate.py. + monkeypatch.setattr(wrap_mod, "_detect_agy_version", lambda _agy_bin: (1, 0, 16)) + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + now = datetime.datetime.now(tz=datetime.timezone.utc) + subject = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "Test CA")]) + cert = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(subject) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now) + .not_valid_after(now + datetime.timedelta(days=1)) + .add_extension(x509.BasicConstraints(ca=True, path_length=0), critical=True) + .sign(key, hashes.SHA256()) + ) + monkeypatch.setattr( + "headroom.proxy.agy_ca.ensure_root_ca", + lambda base_dir=None: (key, cert, Path("/tmp/ca.key"), Path("/tmp/ca.crt")), + ) + monkeypatch.setattr( + "headroom.proxy.agy_ca.build_combined_bundle", + lambda base_dir=None, corp_env_vars=None: Path("/tmp/bundle.pem"), + ) + monkeypatch.setattr("subprocess.run", lambda *a, **kw: MagicMock(returncode=0)) + + +class TestAgySerenaWired: + """wrap agy registers Serena via AgyRegistrar; --no-serena removes/skips it.""" + + def test_wrap_agy_registers_serena( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + from headroom.mcp_registry.agy import AgyRegistrar + + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy"], catch_exceptions=False) + assert result.exit_code == 0 + + reg = AgyRegistrar(home_dir=tmp_path) + spec = reg.get_server("serena") + assert spec is not None, "wrap agy must register a 'serena' MCP entry" + assert spec.command == "uvx" + assert "ide-assistant" in spec.args + + def test_wrap_agy_no_serena_does_not_register( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + from headroom.mcp_registry.agy import AgyRegistrar + + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy", "--no-serena"], catch_exceptions=False) + assert result.exit_code == 0 + + reg = AgyRegistrar(home_dir=tmp_path) + assert reg.get_server("serena") is None, "--no-serena must not leave a Serena MCP entry" + + def test_wrap_agy_no_serena_removes_prior_headroom_entry( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """--no-serena actively removes a Headroom-installed Serena entry.""" + from headroom.mcp_registry.agy import AgyRegistrar + from headroom.mcp_registry.install import build_serena_spec + from headroom.mcp_registry.ledger import record_install + + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + + # Seed a Headroom-installed Serena entry + ledger record. + reg = AgyRegistrar(home_dir=tmp_path) + serena_spec = build_serena_spec("ide-assistant") + reg.register_server(serena_spec) + record_install("agy", serena_spec) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy", "--no-serena"], catch_exceptions=False) + assert result.exit_code == 0 + assert AgyRegistrar(home_dir=tmp_path).get_server("serena") is None + + +# --------------------------------------------------------------------------- +# WU1: current main retired tokensave; AGY uses Serena code memory. +# Print mode wires MCP identically to interactive once the version gate allows it. +# --------------------------------------------------------------------------- + + +class TestAgySerenaCodeMemory: + """wrap agy retires stale tokensave entries and registers Serena by default.""" + + def _spy_helpers(self, monkeypatch: pytest.MonkeyPatch): + """Replace code-memory helpers with call-recording spies.""" + import headroom.cli.wrap as wrap_mod + + calls: dict[str, list] = { + "disable_tokensave": [], + "setup_serena": [], + "disable_serena": [], + } + + def _disable_tokensave(registrar, *, verbose=False): + calls["disable_tokensave"].append(verbose) + + def _setup_serena(registrar, *, context, verbose=False, force=False): + calls["setup_serena"].append(context) + + def _disable_serena(registrar, *, verbose=False, reason="--no-serena"): + calls["disable_serena"].append(reason) + + monkeypatch.setattr(wrap_mod, "_disable_tokensave_mcp", _disable_tokensave) + monkeypatch.setattr(wrap_mod, "_setup_serena_mcp", _setup_serena) + monkeypatch.setattr(wrap_mod, "_disable_serena_mcp", _disable_serena) + return calls + + def test_interactive_retires_tokensave_and_registers_serena( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + calls = self._spy_helpers(monkeypatch) + + result = CliRunner().invoke(_get_main(), ["wrap", "agy"], catch_exceptions=False) + assert result.exit_code == 0 + assert calls["disable_tokensave"], "stale Headroom-installed tokensave must be retired" + assert calls["setup_serena"] == ["ide-assistant"] + assert not calls["disable_serena"] + + def test_interactive_no_tokensave_flag_is_compat_and_uses_serena( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + calls = self._spy_helpers(monkeypatch) + + result = CliRunner().invoke( + _get_main(), ["wrap", "agy", "--no-tokensave"], catch_exceptions=False + ) + assert result.exit_code == 0 + assert calls["disable_tokensave"], "--no-tokensave still cleans up stale tokensave" + assert calls["setup_serena"] == ["ide-assistant"] + + def test_print_mode_retires_tokensave_and_registers_serena( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + calls = self._spy_helpers(monkeypatch) + + result = CliRunner().invoke( + _get_main(), ["wrap", "agy", "--", "--print", "hi"], catch_exceptions=False + ) + assert result.exit_code == 0 + assert calls["disable_tokensave"] + assert calls["setup_serena"] == ["ide-assistant"] + assert not calls["disable_serena"] + + +# --------------------------------------------------------------------------- +# T9 Fix 2: unwrap_agy Serena removal is ledger-gated (falsification guard) +# --------------------------------------------------------------------------- + + +class TestUnwrapAgySerena: + """unwrap_agy removes only Headroom-installed Serena; preserves user entries.""" + + def test_unwrap_removes_headroom_installed_serena( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + from headroom.mcp_registry.agy import AgyRegistrar + from headroom.mcp_registry.install import build_serena_spec + from headroom.mcp_registry.ledger import record_install + + monkeypatch.setattr(Path, "home", lambda: tmp_path) + + reg = AgyRegistrar(home_dir=tmp_path) + serena_spec = build_serena_spec("ide-assistant") + reg.register_server(serena_spec) + record_install("agy", serena_spec) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["unwrap", "agy"]) + assert result.exit_code == 0 + assert AgyRegistrar(home_dir=tmp_path).get_server("serena") is None + + def test_unwrap_preserves_user_managed_serena( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A user-managed serena entry (absent from ledger) must survive unwrap.""" + from headroom.mcp_registry.agy import AgyRegistrar + from headroom.mcp_registry.base import ServerSpec + + monkeypatch.setattr(Path, "home", lambda: tmp_path) + + reg = AgyRegistrar(home_dir=tmp_path) + # User-managed entry: different command, NOT recorded in ledger. + user_spec = ServerSpec( + name="serena", + command="/opt/my-serena/bin/serena", + args=("custom",), + env={}, + ) + reg.register_server(user_spec) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["unwrap", "agy"]) + assert result.exit_code == 0 + survived = AgyRegistrar(home_dir=tmp_path).get_server("serena") + assert survived is not None, "user-managed serena must not be removed" + assert survived.command == "/opt/my-serena/bin/serena" + + +class TestUnwrapAgyTokensave: + """unwrap_agy removes only Headroom-installed tokensave; preserves user entries.""" + + def test_unwrap_removes_headroom_installed_tokensave( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + from headroom.mcp_registry.agy import AgyRegistrar + from headroom.mcp_registry.base import ServerSpec + from headroom.mcp_registry.ledger import record_install + + monkeypatch.setattr(Path, "home", lambda: tmp_path) + + reg = AgyRegistrar(home_dir=tmp_path) + spec = ServerSpec(name="tokensave", command="tokensave", args=("serve",)) + reg.register_server(spec) + record_install("agy", spec) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["unwrap", "agy"]) + assert result.exit_code == 0 + assert AgyRegistrar(home_dir=tmp_path).get_server("tokensave") is None + + def test_unwrap_preserves_user_managed_tokensave( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A user-managed tokensave entry (absent from ledger) must survive unwrap.""" + from headroom.mcp_registry.agy import AgyRegistrar + from headroom.mcp_registry.base import ServerSpec + + monkeypatch.setattr(Path, "home", lambda: tmp_path) + + reg = AgyRegistrar(home_dir=tmp_path) + user_spec = ServerSpec( + name="tokensave", + command="/opt/my-tokensave/bin/tokensave", + args=("serve",), + env={}, + ) + reg.register_server(user_spec) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["unwrap", "agy"]) + assert result.exit_code == 0 + survived = AgyRegistrar(home_dir=tmp_path).get_server("tokensave") + assert survived is not None, "user-managed tokensave must not be removed" + assert survived.command == "/opt/my-tokensave/bin/tokensave" + + +# --------------------------------------------------------------------------- +# agy print-mode MCP hang fix +# --------------------------------------------------------------------------- + + +class TestAgyPrintModeDetection: + """_agy_print_mode flags single-shot non-interactive invocations.""" + + def _fn(self): + from headroom.cli.wrap import _agy_print_mode + + return _agy_print_mode + + def test_detects_print(self) -> None: + assert self._fn()(("--print", "hello")) is True + + def test_detects_short_p(self) -> None: + assert self._fn()(("-p", "hello")) is True + + def test_detects_prompt_alias(self) -> None: + assert self._fn()(("--prompt", "hello")) is True + + def test_detects_print_equals_joined(self) -> None: + # agy accepts `--print=hi` (live-verified) — must be treated as print mode, + # else the interactive branch persists an MCP and the hang returns. + assert self._fn()(("--print=hi",)) is True + + def test_detects_prompt_equals_joined(self) -> None: + assert self._fn()(("--prompt=hi",)) is True + + def test_detects_short_p_equals_joined(self) -> None: + # agy accepts `-p=hi` (live-verified). + assert self._fn()(("-p=hi",)) is True + + def test_attached_short_p_value_is_false(self) -> None: + # agy REJECTS `-pVALUE` (exit 2, "flags provided but not defined") — it + # never reaches MCP init, so it must NOT be treated as print mode. + assert self._fn()(("-pHI",)) is False + + def test_interactive_is_false(self) -> None: + assert self._fn()(()) is False + assert self._fn()(("--model", "x")) is False + assert self._fn()(("--model=x",)) is False + + +class TestAgyPrintModeSuppressesMcp: + """Print-mode wrap agy skips a context tool only when its binary is absent. + + (Print mode otherwise wires MCP identically to interactive — see the + tokensave/retrieve/code-graph parity tests; agy no longer hangs on MCP.) + """ + + +class TestAgyRetrieveMcpWiring: + """Headroom retrieve MCP: persistent, local-store-backed, ledger-recorded. + + The retrieve entry is a stable ``headroom mcp serve`` server (no ephemeral + port; ``env={}`` — it resolves markers from the on-disk CCR store). Started + in BOTH print and interactive mode, it is registered PERSISTENTLY and + recorded in the install ledger (like Serena/CBM), NOT reverted on teardown, + so agy can cache and expose ``headroom_retrieve`` across sessions. + """ + + def test_interactive_registers_persistent_retrieve_entry( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Interactive: headroom entry registered DURING the run and PERSISTS + after teardown (ledger-recorded, resolves from the on-disk store).""" + from headroom.mcp_registry.agy import AgyRegistrar + + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + + # Capture whether the headroom entry was live AT THE MOMENT agy ran + # (i.e. while subprocess.run executes), proving it existed mid-session. + seen: dict[str, object] = {} + + def _capture_run(cmd, *a, **kw): + spec = AgyRegistrar(home_dir=tmp_path).get_server("headroom") + seen["spec"] = spec + return MagicMock(returncode=0) + + monkeypatch.setattr("subprocess.run", _capture_run) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy"], catch_exceptions=False) + assert result.exit_code == 0 + + live_spec = seen["spec"] + assert live_spec is not None, "interactive run must register a headroom retrieve entry" + # The entry invokes `headroom mcp serve`, resolved via + # resolve_headroom_command() — either the resolved `headroom` binary or + # ` -m headroom.cli` when the binary is not on PATH. Assert + # against the actual resolution rather than a hard-coded "headroom" so + # the test is robust across dev (editable) and CI installs. + from headroom.install.runtime import resolve_headroom_command + + expected = resolve_headroom_command() + assert live_spec.command == expected[0] + assert live_spec.args == (*expected[1:], "mcp", "serve") + # Stable, port-independent spec: no ephemeral HEADROOM_PROXY_URL — the + # child resolves markers from the shared on-disk CCR store. + assert dict(live_spec.env) == {} + + # The persistent entry SURVIVES teardown (like Serena/CBM) so agy caches + # and exposes it next session. + assert AgyRegistrar(home_dir=tmp_path).get_server("headroom") is not None, ( + "the persistent retrieve entry must survive teardown" + ) + + def test_print_mode_registers_persistent_retrieve_entry( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Print mode wires the retrieve MCP like interactive: a stable headroom + entry is live mid-run and PERSISTS after teardown.""" + from headroom.mcp_registry.agy import AgyRegistrar + + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + + # Capture mid-session: the headroom entry must exist DURING the run. + seen: dict[str, object] = {} + + def _capture_run(cmd, *a, **kw): + seen["spec"] = AgyRegistrar(home_dir=tmp_path).get_server("headroom") + return MagicMock(returncode=0) + + monkeypatch.setattr("subprocess.run", _capture_run) + + runner = CliRunner() + result = runner.invoke( + _get_main(), ["wrap", "agy", "--", "--print", "hi"], catch_exceptions=False + ) + assert result.exit_code == 0 + live_spec = seen["spec"] + assert live_spec is not None, "print mode must register a headroom retrieve entry mid-run" + # Stable, port-independent spec (on-disk store resolution). + assert dict(live_spec.env) == {} + # Persistent: survives teardown. + assert AgyRegistrar(home_dir=tmp_path).get_server("headroom") is not None, ( + "the persistent retrieve entry must survive teardown" + ) + + def test_print_mode_starts_retrieve_listener( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Print mode: _start_agy_servers is called with start_retrieve=True (parity).""" + import headroom.cli.wrap as wrap_mod + + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + captured: list[bool] = [] + real_stub = wrap_mod._start_agy_servers + + def _spy(ca_key, ca_cert, base_dir=None, *, start_retrieve=False, project=None): + captured.append(start_retrieve) + return real_stub( + ca_key, ca_cert, base_dir, start_retrieve=start_retrieve, project=project + ) + + monkeypatch.setattr(wrap_mod, "_start_agy_servers", _spy) + + runner = CliRunner() + result = runner.invoke( + _get_main(), ["wrap", "agy", "--", "-p", "hi"], catch_exceptions=False + ) + assert result.exit_code == 0 + assert captured == [True], "print mode must start the retrieve listener (parity)" + + def test_interactive_starts_retrieve_listener( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Interactive: _start_agy_servers is called with start_retrieve=True.""" + import headroom.cli.wrap as wrap_mod + + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + captured: list[bool] = [] + real_stub = wrap_mod._start_agy_servers + + def _spy(ca_key, ca_cert, base_dir=None, *, start_retrieve=False, project=None): + captured.append(start_retrieve) + return real_stub( + ca_key, ca_cert, base_dir, start_retrieve=start_retrieve, project=project + ) + + monkeypatch.setattr(wrap_mod, "_start_agy_servers", _spy) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy"], catch_exceptions=False) + assert result.exit_code == 0 + assert captured == [True], "interactive mode must start the retrieve listener" + + def test_failed_smoke_handshake_removes_retrieve_entry( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A retrieve entry that fails the MCP handshake must not persist.""" + import headroom.cli.wrap as wrap_mod + from headroom.mcp_registry.agy import AgyRegistrar + + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + # Handshake FAILS -> verify-then-remove path for the headroom entry. + monkeypatch.setattr(wrap_mod, "_smoke_verify_mcp_handshake", lambda *a, **kw: False) + + seen: dict[str, object] = {} + + def _capture_run(cmd, *a, **kw): + seen["spec"] = AgyRegistrar(home_dir=tmp_path).get_server("headroom") + return MagicMock(returncode=0) + + monkeypatch.setattr("subprocess.run", _capture_run) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy"], catch_exceptions=False) + assert result.exit_code == 0 + assert seen["spec"] is None, ( + "a headroom entry that fails the handshake must be removed before agy runs" + ) + assert AgyRegistrar(home_dir=tmp_path).get_server("headroom") is None + + +class TestUnwrapAgyUserEntries: + """unwrap agy leaves MCP entries Headroom never installed untouched.""" + + def test_unwrap_preserves_unrelated_user_entry( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + from headroom.mcp_registry.agy import AgyRegistrar + from headroom.mcp_registry.base import ServerSpec + + monkeypatch.setattr(Path, "home", lambda: tmp_path) + reg = AgyRegistrar(home_dir=tmp_path) + reg.register_server(ServerSpec(name="my-tool", command="/opt/my-tool", args=(), env={})) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["unwrap", "agy"]) + assert result.exit_code == 0 + survived = AgyRegistrar(home_dir=tmp_path).get_server("my-tool") + assert survived is not None, "unrelated user MCP entries must survive unwrap" + + +class TestSmokeVerifyMcpHandshake: + """_smoke_verify_mcp_handshake: pass on a real responder, fail on a broken one.""" + + def test_returns_true_for_responding_server(self, tmp_path: Path) -> None: + from headroom.cli.wrap import _smoke_verify_mcp_handshake + + # A tiny stdio server that echoes a JSON-RPC initialize response. + server = tmp_path / "fake_mcp.py" + server.write_text( + "import sys, json\n" + "line = sys.stdin.readline()\n" + "req = json.loads(line)\n" + "print(json.dumps({'jsonrpc': '2.0', 'id': req['id'], 'result': {}}))\n" + "sys.stdout.flush()\n" + ) + import sys as _sys + + ok = _smoke_verify_mcp_handshake(_sys.executable, [str(server)], {}, timeout=10.0) + assert ok is True + + def test_returns_false_for_nonexistent_command(self) -> None: + from headroom.cli.wrap import _smoke_verify_mcp_handshake + + assert _smoke_verify_mcp_handshake("/nonexistent/mcp-bin", [], {}, timeout=5.0) is False + + def test_returns_false_when_no_response_in_time(self, tmp_path: Path) -> None: + from headroom.cli.wrap import _smoke_verify_mcp_handshake + + # A server that reads but never replies — must time out -> False. + server = tmp_path / "silent_mcp.py" + server.write_text("import sys, time\nsys.stdin.readline()\ntime.sleep(30)\n") + import sys as _sys + + ok = _smoke_verify_mcp_handshake(_sys.executable, [str(server)], {}, timeout=2.0) + assert ok is False + + +# --------------------------------------------------------------------------- +# headroom-30y.15: fail-open observability + session compression summary +# --------------------------------------------------------------------------- + + +class TestAgySessionCompressionSummary: + """Integration: wrap agy prints a session compression summary on normal exit.""" + + def test_summary_line_appears_on_normal_exit_mixed( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Summary appears in combined output when mix_stderr=True (default).""" + from unittest.mock import patch + + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + + _empty_stats = { + "entry_count": 0, + "total_original_tokens": 0, + "total_compressed_tokens": 0, + } + + with patch( + "headroom.providers.agy.stats._get_compression_stats", + return_value=_empty_stats, + ): + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy"], catch_exceptions=False) + + assert result.exit_code == 0 + assert "Headroom agy session" in result.output + + def test_fail_open_handler_removed_after_session( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The FailOpenWarnHandler must NOT remain on the logger after agy exits.""" + import logging + from unittest.mock import patch + + from headroom.providers.agy.stats import _GEMINI_LOGGER + + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + + _empty_stats = { + "entry_count": 0, + "total_original_tokens": 0, + "total_compressed_tokens": 0, + } + + logger = logging.getLogger(_GEMINI_LOGGER) + handlers_before = list(logger.handlers) + + with patch( + "headroom.providers.agy.stats._get_compression_stats", + return_value=_empty_stats, + ): + runner = CliRunner() + runner.invoke(_get_main(), ["wrap", "agy"], catch_exceptions=False) + + # No new handlers leaked + assert logger.handlers == handlers_before + + +# --------------------------------------------------------------------------- +# WU-s04.3: print-mode purges stale "headroom" retrieve MCP entry +# --------------------------------------------------------------------------- + + +class TestPrintModePurgesStaleHeadroomEntry: + """Print-mode wrap agy must not remove user-managed MCP entries. + + (Print mode now wires the headroom retrieve MCP like interactive; it no + longer scrubs a stale 'headroom' entry, since MCP no longer hangs agy in + --print mode. User-managed entries are still left untouched.) + """ + + def test_print_mode_purge_does_not_remove_user_managed_entries( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Print mode only removes the 'headroom' retrieve entry; user entries survive.""" + from headroom.mcp_registry.agy import AgyRegistrar + from headroom.mcp_registry.base import ServerSpec + + _stub_agy_mitm_run(tmp_path, monkeypatch, with_uvx=True) + + reg = AgyRegistrar(home_dir=tmp_path) + user_spec = ServerSpec(name="my-tool", command="/opt/my-tool", args=(), env={}) + reg.register_server(user_spec) + + runner = CliRunner() + result = runner.invoke( + _get_main(), ["wrap", "agy", "--", "--print", "hi"], catch_exceptions=False + ) + assert result.exit_code == 0 + survived = AgyRegistrar(home_dir=tmp_path).get_server("my-tool") + assert survived is not None, "user-managed entries must not be removed by print-mode purge" + + +# --------------------------------------------------------------------------- +# WU s04.4: graceful failure modes +# --------------------------------------------------------------------------- + + +class TestAgyGracefulFailures: + """agy launch must fail loud and clean on every expected error path. + + WU s04.4: watchdog, preflight, port-in-use, terminal restore. + """ + + # ------------------------------------------------------------------ + # Shared CA stubs (avoid real cert generation in every test). + # ------------------------------------------------------------------ + + @staticmethod + def _patch_ca(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "headroom.proxy.agy_ca.ensure_root_ca", + lambda base_dir=None: (object(), object(), None, None), + ) + monkeypatch.setattr( + "headroom.proxy.agy_ca.build_combined_bundle", + lambda base_dir=None, corp_env_vars=None: "/tmp/fake-bundle.pem", + ) + + # ------------------------------------------------------------------ + # 1. agy-not-installed: clear, actionable error — no raw traceback. + # ------------------------------------------------------------------ + + def test_agy_not_installed_clear_error(self, monkeypatch: pytest.MonkeyPatch) -> None: + """agy binary unresolvable → click.ClickException, nonzero exit, no traceback.""" + monkeypatch.setattr("shutil.which", lambda _: None) + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy"]) + # Must exit non-zero. + assert result.exit_code != 0 + output = result.output + # Discriminating: this exact text is produced ONLY by the binary + # preflight (ClickException). If the preflight were removed, the run + # would fail elsewhere without this message and the test would fail — + # so it is NOT satisfied by the wrap banner or downstream errors. + assert "'agy' not found in PATH" in output + assert "github.com/google/agy" in output + # click.ClickException formats with an "Error: " prefix. + assert "error" in output.lower() + # Must NOT contain a raw Python traceback. + assert "Traceback" not in output + assert "FileNotFoundError" not in output + + # ------------------------------------------------------------------ + # 2. Watchdog: MITM thread death → abort before subprocess, clear message. + # ------------------------------------------------------------------ + + def test_agy_mitm_thread_death_aborts_launch(self, monkeypatch: pytest.MonkeyPatch) -> None: + """MITM server startup failure → subprocess NOT invoked, clear error message.""" + import headroom.cli.wrap as wrap_mod + + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/agy" if name == "agy" else None) + self._patch_ca(monkeypatch) + + # Simulate _start_agy_servers failing (e.g. the daemon thread dies). + def _fail_startup(*a, **kw): + raise RuntimeError( + "agy MITM server startup failed: connection refused on dispatch bind" + ) + + monkeypatch.setattr(wrap_mod, "_start_agy_servers", _fail_startup) + + subprocess_called: list[list[str]] = [] + + def _capture_run(cmd, *a, **kw): + subprocess_called.append(list(cmd)) + return MagicMock(returncode=0) + + monkeypatch.setattr("subprocess.run", _capture_run) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy"]) + + # Subprocess (agy) must NOT have been invoked. + assert subprocess_called == [], ( + "agy subprocess must NOT be launched when the MITM servers fail to start; " + f"got calls: {subprocess_called}" + ) + # Must exit non-zero. + assert result.exit_code != 0 + # Must produce a clear, human-readable message — not a raw exception chain. + output = result.output.lower() + assert "error" in output or "failed" in output + assert "Traceback" not in result.output + + # ------------------------------------------------------------------ + # 3. Port-in-use: OSError(EADDRINUSE) → explicit "port" mention in error. + # ------------------------------------------------------------------ + + def test_agy_port_in_use_message(self, monkeypatch: pytest.MonkeyPatch) -> None: + """MITM bind failure (EADDRINUSE) → message explicitly names port-in-use problem.""" + import errno + + import headroom.cli.wrap as wrap_mod + + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/agy" if name == "agy" else None) + self._patch_ca(monkeypatch) + + bind_error = OSError(errno.EADDRINUSE, "Address already in use") + + def _fail_with_bind_error(*a, **kw): + # Simulate what _start_agy_servers raises when the async bind fails. + raise RuntimeError(f"agy MITM server startup failed: {bind_error}") from bind_error + + monkeypatch.setattr(wrap_mod, "_start_agy_servers", _fail_with_bind_error) + monkeypatch.setattr("subprocess.run", lambda *a, **kw: MagicMock(returncode=0)) + + runner = CliRunner() + result = runner.invoke(_get_main(), ["wrap", "agy"]) + + assert result.exit_code != 0 + output = result.output.lower() + # The error message must name the problem as a port conflict, not just + # re-echo the raw OSError. The word "port" must appear in isolation + # (i.e., not just as part of "transport"). + import re + + assert re.search(r"\bport\b", output), ( + f"Expected 'port' (as a word) in output; got: {output!r}" + ) + # And must still mention that it's in use / unavailable. + assert "in use" in output or "unavailable" in output or "address already in use" in output + # Must not be a raw traceback. + assert "Traceback" not in result.output + + # ------------------------------------------------------------------ + # 4. Terminal/env restore: _stop_agy_servers called in finally on error. + # ------------------------------------------------------------------ + + def test_agy_server_stop_called_on_error_path(self, monkeypatch: pytest.MonkeyPatch) -> None: + """_stop_agy_servers is called in the finally block even when startup raises.""" + import headroom.cli.wrap as wrap_mod + + monkeypatch.setattr("shutil.which", lambda name: "/usr/bin/agy" if name == "agy" else None) + self._patch_ca(monkeypatch) + + monkeypatch.setattr( + wrap_mod, + "_start_agy_servers", + lambda *a, **kw: (_ for _ in ()).throw(RuntimeError("boom")), + ) + + stop_calls: list[object] = [] + original_stop = wrap_mod._stop_agy_servers + + def _spy_stop(servers: object) -> None: + stop_calls.append(servers) + original_stop(servers) # type: ignore[arg-type] + + monkeypatch.setattr(wrap_mod, "_stop_agy_servers", _spy_stop) + monkeypatch.setattr("subprocess.run", lambda *a, **kw: MagicMock(returncode=0)) + + runner = CliRunner() + runner.invoke(_get_main(), ["wrap", "agy"]) + + # The finally block must have called _stop_agy_servers. + assert len(stop_calls) >= 1, ( + "_stop_agy_servers must run in finally even when startup raises" + ) + + +# --------------------------------------------------------------------------- +# Regression: unwrap agy removes ALL Headroom-added agy config entries +# --------------------------------------------------------------------------- + + +class TestUnwrapAgyRemovesAllHeadroomConfig: + """unwrap agy removes every entry Headroom wrote; user entries survive.""" + + def test_unwrap_agy_removes_all_headroom_config( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """All Headroom-added entries gone after unwrap; user entry preserved. + + Arrange a temp HOME with: + - GEMINI.md containing a headroom-marked block (plus user content) + - AgyRegistrar config with a ledger-recorded persistent "headroom" retrieve entry + - AgyRegistrar config with a ledger-recorded serena entry + - AgyRegistrar config with a user-managed "my-tool" entry (no ledger) + + Act: run `unwrap agy` via CliRunner. + + Assert: + - GEMINI.md headroom block is removed; user content survives + - "headroom" retrieve entry is gone + - serena entry is gone (was ledger-recorded) + - "my-tool" entry is preserved (never in ledger) + - ~/.headroom/ca directory is NOT removed (shared CA is headroom state, + not reverted by unwrap — by design) + """ + from headroom.cli.wrap import _AGY_GEMINI_BLOCK_END, _AGY_GEMINI_BLOCK_START + from headroom.mcp_registry.agy import AgyRegistrar + from headroom.mcp_registry.base import ServerSpec + from headroom.mcp_registry.install import ( + build_headroom_spec, + build_serena_spec, + ) + from headroom.mcp_registry.ledger import record_install + + monkeypatch.setattr(Path, "home", lambda: tmp_path) + + # --- Arrange GEMINI.md with headroom block + user content --- + gemini_md = tmp_path / ".gemini" / "GEMINI.md" + gemini_md.parent.mkdir(parents=True, exist_ok=True) + gemini_md.write_text( + f"# User Instructions\nKeep this.\n\n" + f"{_AGY_GEMINI_BLOCK_START}\n## Headroom\nContext.\n{_AGY_GEMINI_BLOCK_END}\n" + ) + + # --- Arrange AgyRegistrar entries --- + reg = AgyRegistrar(home_dir=tmp_path) + + # Persistent, ledger-recorded "headroom" retrieve entry (as wrap agy now + # installs it: stable spec, env={}, recorded in the ledger). Ledger-gated + # unwrap removes it because it is Headroom-owned. (A NON-ledgered headroom + # entry — e.g. a `headroom mcp install` fleet entry — is left in place; + # that path is covered by test_agy_retrieve_persistent.) + headroom_spec = build_headroom_spec() + reg.register_server(headroom_spec) + record_install("agy", headroom_spec) + + # Headroom-installed serena entry (recorded in ledger). + serena_spec = build_serena_spec("ide-assistant") + reg.register_server(serena_spec) + record_install("agy", serena_spec) + + # User-managed entry: NOT in ledger — must survive. + user_spec = ServerSpec(name="my-tool", command="/opt/my-tool", args=(), env={}) + reg.register_server(user_spec) + + # Arrange a fake ~/.headroom/ca dir to prove unwrap does NOT touch it. + ca_dir = tmp_path / ".headroom" / "ca" + ca_dir.mkdir(parents=True, exist_ok=True) + (ca_dir / "ca.crt").write_text("fake cert") + + # --- Act --- + runner = CliRunner() + result = runner.invoke(_get_main(), ["unwrap", "agy"]) + assert result.exit_code == 0, f"unwrap agy failed:\n{result.output}" + + # --- Assert: GEMINI.md --- + gemini_text = gemini_md.read_text() + assert _AGY_GEMINI_BLOCK_START not in gemini_text, ( + "headroom block START marker must be removed from GEMINI.md" + ) + assert _AGY_GEMINI_BLOCK_END not in gemini_text, ( + "headroom block END marker must be removed from GEMINI.md" + ) + assert "# User Instructions" in gemini_text, "user content must survive GEMINI.md cleanup" + assert "Keep this." in gemini_text, "user content body must survive GEMINI.md cleanup" + + # --- Assert: AgyRegistrar entries removed --- + reg2 = AgyRegistrar(home_dir=tmp_path) + assert reg2.get_server("headroom") is None, ( + "the ledger-recorded persistent 'headroom' retrieve entry must be removed by unwrap" + ) + assert reg2.get_server("serena") is None, ( + "the Headroom-installed serena MCP entry must be removed by unwrap" + ) + # --- Assert: user-managed entry preserved --- + survived = reg2.get_server("my-tool") + assert survived is not None, "user-managed 'my-tool' entry must survive unwrap" + assert survived.command == "/opt/my-tool" + + # --- Assert: CA directory intentionally NOT removed (by design) --- + assert ca_dir.exists(), "unwrap must NOT remove ~/.headroom/ca (shared headroom CA state)" + assert (ca_dir / "ca.crt").exists(), "CA certificate must remain intact after unwrap" + + +# --------------------------------------------------------------------------- +# headroom-n0i.7 — corporate proxy credentials must never leak +# --------------------------------------------------------------------------- + + +class TestWrapAgyCorpProxyRedaction(TestWrapAgyDisclosureBanner): + """HTTPS_PROXY / https_proxy userinfo must never reach the launch banner + or logs, while the host:port is still surfaced for operator visibility.""" + + _CORP_PROXY_USER = "user" + _CORP_PROXY_PASS = "s3cr3t-pw" + _CORP_PROXY_HOSTPORT = "proxy.example:3128" + _CORP_PROXY_USERINFO = f"{_CORP_PROXY_USER}:{_CORP_PROXY_PASS}@" + _CORP_PROXY_URL = f"http://{_CORP_PROXY_USERINFO}{_CORP_PROXY_HOSTPORT}" + + def _invoke_with_corp_proxy( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, *, lowercase: bool + ): + caplog.set_level(logging.DEBUG) + monkeypatch.delenv("HTTPS_PROXY", raising=False) + monkeypatch.delenv("https_proxy", raising=False) + monkeypatch.setenv("https_proxy" if lowercase else "HTTPS_PROXY", self._CORP_PROXY_URL) + return self._invoke_agy(monkeypatch) + + def test_uppercase_https_proxy_credentials_absent_from_banner_and_logs( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + result = self._invoke_with_corp_proxy(monkeypatch, caplog, lowercase=False) + assert self._CORP_PROXY_USERINFO not in result.output + assert self._CORP_PROXY_PASS not in result.output + assert self._CORP_PROXY_USERINFO not in caplog.text + assert self._CORP_PROXY_PASS not in caplog.text + assert self._CORP_PROXY_HOSTPORT in result.output + + def test_lowercase_https_proxy_credentials_absent_from_banner_and_logs( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + result = self._invoke_with_corp_proxy(monkeypatch, caplog, lowercase=True) + assert self._CORP_PROXY_USERINFO not in result.output + assert self._CORP_PROXY_PASS not in result.output + assert self._CORP_PROXY_USERINFO not in caplog.text + assert self._CORP_PROXY_PASS not in caplog.text + assert self._CORP_PROXY_HOSTPORT in result.output + + +class TestRedactProxyUrl: + """Parametrized table over every ``redact_proxy_url`` edge case.""" + + @pytest.mark.parametrize( + ("url", "expected"), + [ + ( + "http://user:secret@proxy.example:3128", + "http://proxy.example:3128", + ), + ("http://proxy.corp", "http://proxy.corp:80"), + ("https://proxy.corp", "https://proxy.corp:443"), + ("http://proxy.corp:abc", _PROXY_URL_REDACTED_PLACEHOLDER), + ("user:pass@proxy.corp:3128", _PROXY_URL_REDACTED_PLACEHOLDER), + ("http://u:p@[::1]:8080", "http://[::1]:8080"), + ("http://u:p@a@h:80", "http://h:80"), + ("http://ho\x1bst:80", "http://host:80"), + ], + ) + def test_redact_proxy_url_table(self, url: str, expected: str) -> None: + assert redact_proxy_url(url) == expected + + def test_schemeless_url_never_leaks_password(self) -> None: + assert "pass" not in redact_proxy_url("user:pass@proxy.corp:3128") + + def test_control_bytes_never_reach_result(self) -> None: + assert "\x1b" not in redact_proxy_url("http://ho\x1bst:80") diff --git a/tests/test_wrap_agy_proxy_wiring.py b/tests/test_wrap_agy_proxy_wiring.py new file mode 100644 index 000000000..2816c99e1 --- /dev/null +++ b/tests/test_wrap_agy_proxy_wiring.py @@ -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)" diff --git a/uv.lock b/uv.lock index 39ad607f6..d5243c56f 100644 --- a/uv.lock +++ b/uv.lock @@ -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"