Mirror of headroomlabs-ai/headroom (AI context compression proxy)
Find a file
gglucass 0ce68dedd7
fix(proxy): restore Codex usage headers on WS and streaming SSE transports (#577) (#794)
## Description

Codex's subscription/rate-limit window (the `x-codex-*` headers) was
being
**stripped on every transport Codex actually uses**, so session/weekly
usage
never reached the Codex CLI's own `/status` display, Headroom
`/stats`/dashboard,
or any consumer that sniffs the client-facing handshake. This PR
restores it on
**both** the WebSocket and streaming-SSE paths — the two halves of #577
— in one
place.

Fixes #577

**Supersedes #582 and #590.** This PR incorporates #582's SSE fix
(carried verbatim
with a `Co-authored-by` trailer) and additionally forwards the window
onto the client
`101` on the WS path, which #582/#590's capture-only WS code cannot do.
Both can be
closed as superseded once this merges — GitHub closing keywords only
auto-close
issues (hence `Fixes #577` above), not PRs, so #582/#590 need a manual
close.

### WebSocket (`gpt-5.4+`)

OpenAI delivers `x-codex-*` **only** on the upstream WS handshake
response, never
in data frames. `handle_openai_responses_ws` accepted the client WS
*before* it
connected upstream and never read `upstream.response.headers`, so the
window was
dropped. This reorders the handler to **connect upstream first**,
extract the
`x-codex-*` subset, then **accept the client WS with those headers
attached** to
the `101`, and refresh the Python state for `/stats` parity.

### Streaming SSE (incorporated from #582, @m16khb)

Codex CLI almost always streams. `streaming.py` neither captured
`x-codex-*` into
`CodexRateLimitState` nor forwarded it — the forwarded-header filter
matched only
the substring `"ratelimit"`, which `x-codex-*` does not contain. This
calls
`update_from_headers()` **before** the `>=400` early-return (so a
streaming 429/5xx
still refreshes the window, matching the non-streaming handlers) and
widens the
forward filter to pass `x-codex-*`.

> Credit: the SSE fix is @m16khb's work from #582, carried here verbatim
with a
> `Co-authored-by` trailer so the maintainer gets a single PR covering
both
> transports. This supersedes #582/#590's **WS** capture (which only
writes
> `/stats`); the connect-before-accept reorder additionally forwards the
window to
> the client `101`, which capture-only cannot do. #590's optional
snapshot
> persistence is intentionally left out (separable; hot-path sync write;
doesn't
> help the `101`-sniff consumers).

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `openai.py`: add `_extract_codex_handshake_headers()` (strictly
`x-codex-*`, via
`raw_items()` to avoid `MultipleValuesError`; never
`set-cookie`/`authorization`).
- `openai.py`: reorder `handle_openai_responses_ws` — connect-only retry
loop runs
before `accept()`; `accept(headers=...)` carries the forwarded window;
first
client frame read afterward. HTTP fallback preserved; it now also
refreshes
  `/stats` from the HTTP response headers.
- `streaming.py`: capture `x-codex-*` on all statuses + widen the
forwarded-header
  filter (from #582).

### Diff-size note

The bulk of the `openai.py` line count is **whitespace-only
relocation**: the relay
block dedents one level out of the old per-attempt `async with`. Logical
change is
~290 lines. **Review with `?w=1`.** In API-key mode the handshake
carries no
`x-codex-*`, so the accept-header list is empty and the path behaves
exactly as
before — the fix only activates for ChatGPT-subscription auth.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

- WS: `test_ws_connect_happens_before_accept`,
`test_ws_forwards_codex_headers_to_client_accept`
(only `x-codex-*` forwarded; `set-cookie`/`authorization` excluded;
`/stats` refreshed),
`test_ws_connect_failure_falls_back_to_http`,
`test_ws_first_frame_timeout_after_connect_closes_upstream`.
- Fallback: `test_fallback_refreshes_codex_rate_limit_state`.
- SSE:
`test_codex_rate_limit_headers_captured_and_forwarded_in_streaming`,
  `test_codex_rate_limit_captured_on_streaming_429` (from #582).
- Wire-level e2e: `tests/e2e_ws_codex_usage_headers.py` boots the real
proxy + fake
upstream + real `websockets` client and reads the client `101` — closes
the gap
the unit tests stub (that uvicorn/starlette actually write
`accept(headers=...)`).

## Test Output

```
$ uv run pytest tests/test_proxy_streaming_ratelimit_headers.py \
                tests/test_ws_http_fallback.py \
                tests/test_openai_codex_ws_lifecycle.py \
                tests/test_openai_codex_ws_timings.py \
                tests/test_codex_rate_limits.py -q
63 passed in 0.83s

$ .venv/bin/python tests/e2e_ws_codex_usage_headers.py
[codex-hdr-e2e] client 101 headers:
    x-codex-primary-used-percent: 42
    x-codex-primary-window-minutes: 300
    x-codex-secondary-used-percent: 7
    x-codex-secondary-window-minutes: 10080
[codex-hdr-e2e] /stats reflects codex window (primary-used=42)
=== CODEX-HDR E2E ALL GREEN ===

$ uv run ruff check . && uv run ruff format --check <touched files>
All checks passed!
```

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [ ] I have updated the CHANGELOG.md if applicable

## Additional Notes

- **Why connect-before-accept (not capture-only).** Once `accept()`
sends the `101`,
headers can no longer be added; the `x-codex-*` window only exists after
we connect
upstream. Capturing into Python state (as #582/#590's WS code does)
fixes `/stats`
but not the Codex CLI's native display or any `101`-sniffing consumer —
those need
  the headers *on the client handshake*, which requires the reorder.
- **Security.** Forwarding is filtered strictly to `x-codex-*`;
`set-cookie`,
`authorization`, and all other upstream headers are never forwarded to
the client
  (asserted by both the unit test and the e2e).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

## Contract Schemas

Per maintainer request: a JSON Schema (draft 2020-12) artifact
enshrining the OpenAI
interaction expectations this changeset relies on, so drift is
detectable later.

Committed following the repo's parity convention:
- schema:
`tests/parity/fixtures/codex_openai_contracts/codex-openai-interaction.schema.json`
- test: `tests/test_codex_openai_contract_parity.py` binds the schema to
the **live code**
in both directions, so drift fails CI rather than living only in this
description -
every declared `x-codex-*` header must be consumed by
`parse_codex_rate_limits`, and
`_extract_codex_handshake_headers` must forward exactly the declared
subset and never
`set-cookie`/`authorization`. No new dependency (does not pull in
`jsonschema`).

It covers, as `$defs`:

- `WSUpstreamHandshakeResponse` / `StreamingUpstreamResponseHeaders` -
the upstream
`x-codex-*` header family (full superset, with per-header wire pattern +
the parsed
semantic type) the WS and SSE captures read. Source of truth:
`parse_codex_rate_limits`.
- `ClientForwardedHandshakeHeaders` - the WS-101 **allow/deny**
contract: only
`x-codex-*` may be forwarded; `set-cookie`/`authorization` are
explicitly forbidden
  (`propertyNames` + `not`).
- `ClientForwardedStreamingHeaders` - the wider SSE forward set
(`*ratelimit*` OR `x-codex*`).
- `WSClientRequestFrame` / `WSRelayEvent` / `HTTPFallbackRequestBody` -
the WS frame
  envelopes and the unwrapped HTTP-fallback POST body.
- `CodexRateLimitStatsOutput` - the headroom `/stats` shape the parity
tests assert.

Validated with `jsonschema` (Draft202012 `check_schema` passes; positive
instances from
the e2e validate; negative instances - a leaked `set-cookie`, a fallback
body still
carrying a top-level `type` - are correctly rejected).

<details>
<summary><code>codex-openai-interaction.schema.json</code> (draft
2020-12)</summary>

```json
{
  "$schema": "https://json-schema.org/draft/2020-12/schema",
  "$id": "https://github.com/chopratejas/headroom/contracts/codex-openai-interaction.schema.json",
  "title": "Codex <-> OpenAI interaction contracts (PR #794)",
  "description": "Enshrines the OpenAI interaction expectations this changeset depends on, so drift is detectable. Header values are transported as strings on the wire; the `x-headroom-parsed-type` annotation on each records the semantic type the parser (headroom/subscription/codex_rate_limits.py) coerces them to. Sources: codex_rate_limits.parse_codex_rate_limits (header family + gating), openai._extract_codex_handshake_headers (WS-101 forward filter), streaming.py (SSE forward filter).",
  "$defs": {
    "OpenAICodexWindowHeaders": {
      "title": "x-codex-*-{primary,secondary} window headers",
      "description": "A rolling rate-limit/subscription window. A window is materialized iff its `*-used-percent` header is present and numeric; `*-window-minutes` and `*-reset-at` are optional. `primary` and `secondary` are independent and either may be absent.",
      "type": "object",
      "properties": {
        "x-codex-primary-used-percent": {
          "type": "string",
          "pattern": "^\\d+(?:\\.\\d+)?$",
          "x-headroom-parsed-type": "float (0-100, NaN-guarded)",
          "description": "Percent of the primary window consumed. Gates creation of the primary window."
        },
        "x-codex-primary-window-minutes": {
          "type": "string",
          "pattern": "^\\d+$",
          "x-headroom-parsed-type": "int",
          "description": "Primary window size in minutes."
        },
        "x-codex-primary-reset-at": {
          "type": "string",
          "pattern": "^\\d+$",
          "x-headroom-parsed-type": "int (Unix epoch seconds)",
          "description": "Absolute reset time of the primary window."
        },
        "x-codex-secondary-used-percent": {
          "type": "string",
          "pattern": "^\\d+(?:\\.\\d+)?$",
          "x-headroom-parsed-type": "float (0-100, NaN-guarded)",
          "description": "Percent of the secondary window consumed. Gates creation of the secondary window."
        },
        "x-codex-secondary-window-minutes": {
          "type": "string",
          "pattern": "^\\d+$",
          "x-headroom-parsed-type": "int"
        },
        "x-codex-secondary-reset-at": {
          "type": "string",
          "pattern": "^\\d+$",
          "x-headroom-parsed-type": "int (Unix epoch seconds)"
        }
      },
      "additionalProperties": true
    },
    "OpenAICodexCreditsHeaders": {
      "title": "x-codex-credits-* headers",
      "description": "OpenAI credits balance. A credits snapshot is materialized iff `x-codex-credits-has-credits` is present; `unlimited` defaults to false; `balance` is optional.",
      "type": "object",
      "properties": {
        "x-codex-credits-has-credits": {
          "type": "string",
          "pattern": "^(?:[Tt][Rr][Uu][Ee]|[Ff][Aa][Ll][Ss][Ee]|[01])$",
          "x-headroom-parsed-type": "bool (true|false|1|0, case-insensitive)",
          "description": "Gates creation of the credits snapshot."
        },
        "x-codex-credits-unlimited": {
          "type": "string",
          "pattern": "^(?:[Tt][Rr][Uu][Ee]|[Ff][Aa][Ll][Ss][Ee]|[01])$",
          "x-headroom-parsed-type": "bool (defaults false when absent/unparseable)"
        },
        "x-codex-credits-balance": {
          "type": "string",
          "x-headroom-parsed-type": "str (empty -> null)",
          "description": "Free-form server string, e.g. \"$5.00\"."
        }
      },
      "additionalProperties": true
    },
    "OpenAICodexMetaHeaders": {
      "title": "x-codex meta headers",
      "type": "object",
      "properties": {
        "x-codex-limit-name": {
          "type": "string",
          "x-headroom-parsed-type": "str (empty -> null)",
          "description": "Active limit/model label, e.g. \"gpt-5.2-codex-sonic\"."
        },
        "x-codex-promo-message": {
          "type": "string",
          "x-headroom-parsed-type": "str (empty -> null)",
          "description": "Server announcement. Also gates snapshot creation when present."
        }
      },
      "additionalProperties": true
    },
    "OpenAICodexRateLimitHeaders": {
      "title": "Full x-codex-* header family OpenAI may emit",
      "description": "Superset of every x-codex-* header headroom reads. parse_codex_rate_limits returns a snapshot iff at least one of: a primary window, a secondary window, a credits snapshot, or a non-empty promo message is present; otherwise null (treated as a non-Codex response). All members are individually optional.",
      "type": "object",
      "allOf": [
        { "$ref": "#/$defs/OpenAICodexWindowHeaders" },
        { "$ref": "#/$defs/OpenAICodexCreditsHeaders" },
        { "$ref": "#/$defs/OpenAICodexMetaHeaders" }
      ],
      "additionalProperties": true
    },
    "WSUpstreamHandshakeResponse": {
      "title": "OpenAI WS handshake (101) response headers consumed by the WS fix",
      "description": "On the Codex WebSocket transport the x-codex-* window is delivered ONLY on the upstream handshake response (never in data frames). handle_openai_responses_ws reads upstream.response.headers here. This is the contract the connect-before-accept reorder depends on: if OpenAI ever moves these headers off the handshake (e.g. into a frame), the WS half of the fix goes stale.",
      "$ref": "#/$defs/OpenAICodexRateLimitHeaders"
    },
    "StreamingUpstreamResponseHeaders": {
      "title": "OpenAI streaming/HTTP response headers consumed by the SSE fix",
      "description": "On the streaming SSE/HTTP transport the same x-codex-* headers ride the HTTP response. streaming.py captures them on ALL statuses (including >=400) via update_from_headers, and forwards a wider set to the client (see ClientForwardedStreamingHeaders).",
      "$ref": "#/$defs/OpenAICodexRateLimitHeaders"
    },
    "ClientForwardedHandshakeHeaders": {
      "title": "Headers forwarded onto the CLIENT-facing WS 101 (allow/deny contract)",
      "description": "_extract_codex_handshake_headers forwards ONLY headers whose (lowercased) name starts with `x-codex-`. Every other upstream handshake header - notably set-cookie and authorization - MUST NOT appear on the client 101. Enforced by propertyNames below and asserted by the unit tests + tests/e2e_ws_codex_usage_headers.py.",
      "type": "object",
      "propertyNames": {
        "pattern": "^[Xx]-[Cc][Oo][Dd][Ee][Xx]-"
      },
      "not": {
        "anyOf": [
          { "required": ["set-cookie"] },
          { "required": ["Set-Cookie"] },
          { "required": ["authorization"] },
          { "required": ["Authorization"] }
        ]
      },
      "additionalProperties": { "type": "string" }
    },
    "ClientForwardedStreamingHeaders": {
      "title": "Headers forwarded to the client on the streaming SSE path",
      "description": "streaming.py forwards a header iff `\"ratelimit\" in name.lower()` OR `name.lower().startswith(\"x-codex\")`. This is a SUPERSET of the WS allow-list: it additionally passes generic *ratelimit* headers (e.g. the Anthropic streaming path) which do not contain the x-codex prefix.",
      "type": "object",
      "propertyNames": {
        "pattern": "(?:[Rr][Aa][Tt][Ee][Ll][Ii][Mm][Ii][Tt])|^[Xx]-[Cc][Oo][Dd][Ee][Xx]"
      },
      "additionalProperties": { "type": "string" }
    },
    "WSClientRequestFrame": {
      "title": "Client -> proxy WS data frame (Responses API over WS)",
      "description": "Codex sends the request as a response.create envelope. The HTTP fallback unwraps `.response` for the POST body, forces stream=true, and strips any top-level `type`. A flattened variant (no envelope, fields at top level) is also tolerated by the fallback.",
      "type": "object",
      "properties": {
        "type": { "const": "response.create" },
        "response": {
          "type": "object",
          "properties": {
            "model": { "type": "string", "description": "e.g. gpt-5.4" },
            "input": {
              "description": "String prompt or Responses-API structured input array.",
              "type": ["string", "array"]
            },
            "stream": { "type": "boolean" }
          },
          "required": ["model"],
          "additionalProperties": true
        }
      },
      "required": ["type", "response"],
      "additionalProperties": true
    },
    "WSRelayEvent": {
      "title": "proxy -> client WS data frame (relayed Responses API event)",
      "description": "SSE `data:` payloads relayed verbatim as WS text frames. `[DONE]` sentinels are dropped (not relayed). Every relayed event is a JSON object carrying a `type`. response.completed additionally carries usage under `response.usage`. anyOf (not oneOf): an error event also satisfies the looser lifecycle shape, which is fine.",
      "anyOf": [
        {
          "title": "lifecycle event",
          "type": "object",
          "properties": {
            "type": {
              "type": "string",
              "examples": [
                "response.created",
                "response.output_item.added",
                "response.completed"
              ]
            },
            "response": { "type": "object", "additionalProperties": true }
          },
          "required": ["type"],
          "additionalProperties": true
        },
        {
          "title": "error event",
          "type": "object",
          "properties": {
            "type": { "const": "error" },
            "error": {
              "type": "object",
              "properties": { "message": { "type": "string" } },
              "required": ["message"],
              "additionalProperties": true
            }
          },
          "required": ["type", "error"],
          "additionalProperties": true
        }
      ]
    },
    "HTTPFallbackRequestBody": {
      "title": "proxy -> OpenAI HTTP POST body on WS->HTTP fallback",
      "description": "Derived from WSClientRequestFrame: the inner `.response` object, with `stream` forced to true and any top-level `type` removed.",
      "type": "object",
      "properties": {
        "model": { "type": "string" },
        "stream": { "const": true },
        "input": { "type": ["string", "array"] }
      },
      "required": ["model", "stream"],
      "not": { "required": ["type"] },
      "additionalProperties": true
    },
    "CodexRateLimitStatsOutput": {
      "title": "headroom /stats output for the codex tracker (CodexRateLimitSnapshot.to_dict)",
      "description": "Internal (headroom-emitted) shape produced from the headers above; the WS and SSE update_from_headers parity tests assert this is refreshed. Included so drift in our own surface is also caught.",
      "type": "object",
      "properties": {
        "limit_id": { "const": "codex" },
        "limit_name": { "type": ["string", "null"] },
        "primary": { "$ref": "#/$defs/CodexWindowDict" },
        "secondary": { "$ref": "#/$defs/CodexWindowDict" },
        "credits": {
          "oneOf": [
            { "type": "null" },
            {
              "type": "object",
              "properties": {
                "has_credits": { "type": "boolean" },
                "unlimited": { "type": "boolean" },
                "balance": { "type": ["string", "null"] }
              },
              "required": ["has_credits", "unlimited", "balance"],
              "additionalProperties": false
            }
          ]
        },
        "promo_message": { "type": ["string", "null"] },
        "captured_at": { "type": "number", "description": "Unix epoch seconds (float)." }
      },
      "required": ["limit_id", "limit_name", "primary", "secondary", "credits", "promo_message", "captured_at"],
      "additionalProperties": false
    },
    "CodexWindowDict": {
      "oneOf": [
        { "type": "null" },
        {
          "type": "object",
          "properties": {
            "used_percent": { "type": "number" },
            "window_minutes": { "type": ["integer", "null"] },
            "window_label": { "type": "string", "description": "e.g. \"5h\", \"7d\"-style label; \"unknown\" when window_minutes is null." },
            "resets_at": { "type": ["integer", "null"], "description": "Unix epoch seconds." },
            "seconds_until_reset": { "type": ["integer", "null"] }
          },
          "required": ["used_percent", "window_minutes", "window_label", "resets_at", "seconds_until_reset"],
          "additionalProperties": false
        }
      ]
    }
  }
}
```

</details>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: m16khb <m16khb@gmail.com>
2026-06-09 15:55:53 -05:00
.claude-plugin ci(release): align manifest + pyproject + package.json to 0.22.3 2026-05-25 18:41:38 -07:00
.devcontainer ci: speed up GitHub Actions — path filters, caching, timeouts, version upgrades (#620) 2026-06-05 14:32:53 -08:00
.github feat: support Python 3.14+ via pyo3 abi3 stable ABI (#516) 2026-06-09 06:33:16 -08:00
benchmarks fix: B2 — live-zone block dispatcher skeleton 2026-05-02 12:45:43 -07:00
crates fix(observability): G3 remediation — bound cardinality + wire dead metrics 2026-05-24 10:41:56 -07:00
docker feat: add differential network capture harness (#761) 2026-06-08 22:18:31 -07:00
docs fix: preserve Claude Code tool-search deferral through the proxy (#746) (#753) 2026-06-08 11:20:48 -07:00
e2e fix(cli): wrap subcommands for cline, continue, goose, openhands 2026-05-21 21:04:50 -07:00
examples docs: fix env var discrepancies across markdown files 2026-05-29 15:04:13 +05:30
headroom fix(proxy): restore Codex usage headers on WS and streaming SSE transports (#577) (#794) 2026-06-09 15:55:53 -05:00
plugins chore: release main (#607) 2026-06-08 11:21:23 -07:00
REALIGNMENT docs: add Realignment plan (40 PRs, 9 phases) 2026-05-01 23:34:46 -07:00
scripts ci(release): fix workflow-validation for new release trigger 2026-05-25 18:36:12 -07:00
sdk/typescript chore: release main (#607) 2026-06-08 11:21:23 -07:00
sql feat(telemetry): add headroom_stack and install_mode identity fields 2026-04-17 17:12:38 +02:00
tests fix(proxy): restore Codex usage headers on WS and streaming SSE transports (#577) (#794) 2026-06-09 15:55:53 -05:00
wiki feat: add differential network capture harness (#761) 2026-06-08 22:18:31 -07:00
.actrc feat: add act testing config, fix gitignore, make workflow production-ready 2026-04-15 20:28:29 -05:00
.actrc.local.example feat: add act testing config, fix gitignore, make workflow production-ready 2026-04-15 20:28:29 -05:00
.changelog.md fix: use /tmp for changelog artifact to avoid . file matching issues 2026-04-15 22:54:34 -05:00
.commitlintrc.json ci: fix smart_crusher branch CI failures + add make ci-precheck pre-push gate 2026-04-27 11:13:47 -07:00
.dockerignore chore: normalize line endings in init diffs 2026-04-21 20:17:14 -05:00
.env.act.example feat: add act testing config, fix gitignore, make workflow production-ready 2026-04-15 20:28:29 -05:00
.env.example fix(security): patch loopback guard, retry None raise, blocking subprocess, and cache stats race 2026-06-04 16:09:53 -04:00
.git-blame-ignore-revs chore: add .git-blame-ignore-revs 2026-04-24 15:35:29 +02:00
.gitattributes chore: enforce LF checkout for Python files 2026-04-23 08:43:03 -05:00
.gitguardian.yaml fix(security): allowlist GitGuardian-flagged test fixtures 2026-05-02 18:33:22 -07:00
.gitignore feat: add differential network capture harness (#761) 2026-06-08 22:18:31 -07:00
.pre-commit-config.yaml chore: normalize line endings in init diffs 2026-04-21 20:17:14 -05:00
.release-please-config.json ci(release-please): set versioned PR title pattern to fix tagging jam 2026-06-04 09:56:35 -07:00
.release-please-manifest.json chore: release main (#607) 2026-06-08 11:21:23 -07:00
Cargo.lock Fix Windows ORT builds and Docker signing retries 2026-05-10 20:59:28 -07:00
Cargo.toml feat: support Python 3.14+ via pyo3 abi3 stable ABI (#516) 2026-06-09 06:33:16 -08:00
CHANGELOG.md feat(proxy): attribute savings history rollups per provider (#791) 2026-06-09 14:55:02 -05:00
claude_analysis_ttl.py chore: add cache TTL cost analysis script 2026-05-13 10:49:17 -07:00
CODE_OF_CONDUCT.md Prepare for OSS release v0.2.0 2026-01-07 11:36:44 -08:00
codecov.yml ci: scope codecov patch coverage 2026-04-22 00:05:47 -05:00
CONTRIBUTING.md docs: update contributing.md 2026-06-05 22:02:28 -07:00
deny.toml feat(rust): scaffold workspace + parity harness (phase-0) 2026-04-24 13:39:48 -07:00
docker-bake.hcl refactor(docker): migrate to bake with multi-variant distroless images 2026-04-04 21:51:37 +05:30
docker-compose.yml chore(deps): loosen over-pinned constraints and add upper bounds (#538) 2026-06-08 22:06:24 -08:00
Dockerfile chore(deps): loosen over-pinned constraints and add upper bounds (#538) 2026-06-08 22:06:24 -08:00
ENTERPRISE.md docs: add enterprise.md 2026-06-04 15:36:47 -07:00
Headroom-2.gif Add demo GIF to README 2026-01-20 18:57:17 -08:00
headroom-savings.png docs: rewrite README for clarity and highlight Kompress-base, leaderboard, RTK 2026-04-18 09:36:49 -07:00
headroom_learn.gif docs: add headroom learn demo GIF to README 2026-03-07 17:56:17 -08:00
HeadroomDemo-Fast.gif Replace demo GIF with HeadroomDemo-Fast.gif 2026-04-10 15:05:25 -07:00
LICENSE Prepare for OSS release v0.2.0 2026-01-07 11:36:44 -08:00
llms.txt docs: fix stale API references, retired class imports, and incorrect examples 2026-06-02 19:19:19 -04:00
Makefile fix: A0 — fail-loud rust core deployment smoke test 2026-05-02 17:52:37 -07:00
mkdocs.yml feat: add differential network capture harness (#761) 2026-06-08 22:18:31 -07:00
NOTICE Prepare for OSS release v0.2.0 2026-01-07 11:36:44 -08:00
PR.md fix: restore release and compress regressions 2026-04-18 16:01:57 -05:00
pyproject.toml feat: support Python 3.14+ via pyo3 abi3 stable ABI (#516) 2026-06-09 06:33:16 -08:00
README.md Merge pull request #592 from divyanshus2404/my-first-contribution 2026-06-06 12:08:07 -07:00
rust-toolchain.toml fix(rust): clippy 1.95 unnecessary_sort_by + pin toolchain 2026-04-27 12:11:49 -07:00
RUST_DEV.md docs: fix env var discrepancies across markdown files 2026-05-29 15:04:13 +05:30
SECURITY.md Prepare for OSS release v0.2.0 2026-01-07 11:36:44 -08:00
TESTING-copilot-subscription.md fix(copilot): restore generic endpoint for non-subscription OAuth (#610) (#612) 2026-06-04 16:27:54 -07:00
uv.lock chore(deps): loosen over-pinned constraints and add upper bounds (#538) 2026-06-08 22:06:24 -08:00

  ██╗  ██╗███████╗ █████╗ ██████╗ ██████╗  ██████╗  ██████╗ ███╗   ███╗
  ██║  ██║██╔════╝██╔══██╗██╔══██╗██╔══██╗██╔═══██╗██╔═══██╗████╗ ████║
  ███████║█████╗  ███████║██║  ██║██████╔╝██║   ██║██║   ██║██╔████╔██║
  ██╔══██║██╔══╝  ██╔══██║██║  ██║██╔══██╗██║   ██║██║   ██║██║╚██╔╝██║
  ██║  ██║███████╗██║  ██║██████╔╝██║  ██║╚██████╔╝╚██████╔╝██║ ╚═╝ ██║
  ╚═╝  ╚═╝╚══════╝╚═╝  ╚═╝╚═════╝ ╚═╝  ╚═╝ ╚═════╝  ╚═════╝ ╚═╝     ╚═╝
                  The context compression layer for AI agents

6095% fewer tokens · library · proxy · MCP · 6 algorithms · local-first · reversible

CI codecov PyPI npm Model: Kompress-base License: Apache 2.0 Docs

Docs · Install · Proof · Agents · Discord · llms.txt · Enterprise

AI agents / LLMs: read /llms.txt here, or fetch the live index / full docs blob.


chopratejas%2Fheadroom | Trendshift

Headroom compresses everything your AI agent reads — tool outputs, logs, RAG chunks, files, and conversation history — before it reaches the LLM. Same answers, fraction of the tokens.

Headroom in action
Live: 10,144 → 1,260 tokens — same FATAL found.

What it does

  • Librarycompress(messages) in Python or TypeScript, inline in any app
  • Proxyheadroom proxy --port 8787, zero code changes, any language
  • Agent wrapheadroom wrap claude|codex|cursor|aider|copilot in one command
  • MCP serverheadroom_compress, headroom_retrieve, headroom_stats for any MCP client
  • Cross-agent memory — shared store across Claude, Codex, Gemini, auto-dedup
  • headroom learn — mines failed sessions, writes corrections to CLAUDE.md / AGENTS.md
  • Reversible (CCR) — originals never deleted; LLM retrieves on demand

How it works (30 seconds)

 Your agent / app
   (Claude Code, Cursor, Codex, LangChain, Agno, Strands, your own code…)
        │   prompts · tool outputs · logs · RAG results · files
        ▼
    ┌────────────────────────────────────────────────────┐
    │  Headroom   (runs locally — your data stays here)  │
    │  ────────────────────────────────────────────────  │
    │  CacheAligner  →  ContentRouter  →  CCR            │
    │                    ├─ SmartCrusher   (JSON)        │
    │                    ├─ CodeCompressor (AST)         │
    │                    └─ Kompress-base  (text, HF)    │
    │                                                    │
    │  Cross-agent memory  ·  headroom learn  ·  MCP     │
    └────────────────────────────────────────────────────┘
        │   compressed prompt  +  retrieval tool
        ▼
 LLM provider  (Anthropic · OpenAI · Bedrock · …)
  • ContentRouter — detects content type, selects the right compressor
  • SmartCrusher / CodeCompressor / Kompress-base — compress JSON, AST, or prose
  • CacheAligner — stabilizes prefixes so provider KV caches actually hit
  • CCR — stores originals locally; LLM calls headroom_retrieve if it needs them

Architecture · CCR reversible compression · Kompress-base model card

Get started (60 seconds)

# 1 — Install
pip install "headroom-ai[all]"          # Python
npm install headroom-ai                 # Node / TypeScript

# 2 — Pick your mode
headroom wrap claude                    # wrap a coding agent
headroom proxy --port 8787              # drop-in proxy, zero code changes
# or: from headroom import compress      # inline library

# 3 — See the savings
headroom perf

Granular extras: [proxy], [mcp], [ml], [code], [memory], [relevance], [image], [agno], [langchain], [evals]. Requires Python 3.10+.

Proof

Savings on real agent workloads:

Workload Before After Savings
Code search (100 results) 17,765 1,408 92%
SRE incident debugging 65,694 5,118 92%
GitHub issue triage 54,174 14,761 73%
Codebase exploration 78,502 41,254 47%

Accuracy preserved on standard benchmarks:

Benchmark Category N Baseline Headroom Delta
GSM8K Math 100 0.870 0.870 ±0.000
TruthfulQA Factual 100 0.530 0.560 +0.030
SQuAD v2 QA 100 97% 19% compression
BFCL Tools 100 97% 32% compression

Reproduce: python -m headroom.evals suite --tier 1 · Full benchmarks & methodology

Star History Chart

Agent compatibility matrix

Agent headroom wrap Notes
Claude Code --memory · --code-graph
Codex shares memory with Claude
Cursor prints config — paste once
Aider starts proxy + launches
Copilot CLI starts proxy + launches
OpenClaw installs as ContextEngine plugin

Any OpenAI-compatible client works via headroom proxy. MCP-native: headroom mcp install.

GitHub Copilot CLI subscription mode

Headroom can route GitHub Copilot CLI subscription traffic through the local proxy:

headroom wrap copilot --subscription -- --model gpt-4o

This lets Headroom intercept OpenAI-compatible Copilot CLI requests and apply the same proxy compression pipeline before forwarding to GitHub Copilot's hosted API. The wrapper resolves the account-specific Copilot API endpoint and prints it as COPILOT_PROVIDER_API_URL=... during launch.

Platform support note: macOS auth reuse via Copilot CLI Keychain storage has been smoke-tested. Windows Credential Manager, Linux Secret Service / secret-tool, and Docker/CI token-injection paths are implemented or planned as auth-discovery paths, but still need real OS validation before they should be considered fully vetted. For Docker and CI, prefer passing an explicit GITHUB_COPILOT_TOKEN or GITHUB_COPILOT_GITHUB_TOKEN rather than relying on host keychain access.

When to use · When to skip

Great fit if you…

  • run AI coding agents daily and want savings without changing your code
  • work across multiple agents and want shared memory
  • need reversible compression — originals always retrievable via CCR

Skip it if you…

  • only use a single provider's native compaction and don't need cross-agent memory
  • work in a sandboxed environment where local processes can't run
Integrations — drop Headroom into any stack
Your setup Hook in with
Any Python app compress(messages, model=…)
Any TypeScript app await compress(messages, { model })
Anthropic / OpenAI SDK withHeadroom(new Anthropic()) · withHeadroom(new OpenAI())
Vercel AI SDK wrapLanguageModel({ model, middleware: headroomMiddleware() })
LiteLLM litellm.callbacks = [HeadroomCallback()]
LangChain HeadroomChatModel(your_llm)
Agno HeadroomAgnoModel(your_model)
Strands Strands guide
ASGI apps app.add_middleware(CompressionMiddleware)
Multi-agent SharedContext().put / .get
MCP clients headroom mcp install
What's inside
  • SmartCrusher — universal JSON: arrays of dicts, nested objects, mixed types.
  • CodeCompressor — AST-aware for Python, JS, Go, Rust, Java, C++.
  • Kompress-base — our HuggingFace model, trained on agentic traces.
  • Image compression — 4090% reduction via trained ML router.
  • CacheAligner — stabilizes prefixes so Anthropic/OpenAI KV caches actually hit.
  • IntelligentContext — score-based context fitting with learned importance.
  • CCR — reversible compression; LLM retrieves originals on demand.
  • Cross-agent memory — shared store, agent provenance, auto-dedup.
  • SharedContext — compressed context passing across multi-agent workflows.
  • headroom learn — plugin-based failure mining for Claude, Codex, Gemini.
Pipeline internals

Headroom exposes one stable request lifecycle across compress(), the SDK, and the proxy:

SetupPre-StartPost-StartInput ReceivedInput CachedInput RoutedInput CompressedInput RememberedPre-SendPost-SendResponse Received

  • Transforms do the work: CacheAligner, ContentRouter, SmartCrusher, CodeCompressor, Kompress-base, IntelligentContext / RollingWindow.
  • Pipeline extensions observe or customize lifecycle stages via on_pipeline_event(...).
  • Compression hooks sit alongside the canonical lifecycle as an additional extension seam.
  • Proxy extensions remain the server/app integration seam for ASGI middleware, routes, and startup policy.

Provider and tool-specific behavior lives under headroom/providers/ so core orchestration stays focused on lifecycle, sequencing, and policy.

  • CLI/tool slices: headroom/providers/claude, copilot, codex, openclaw
  • Provider runtime slices: headroom/providers/claude, gemini, plus shared backend/runtime dispatch in headroom/providers/registry.py
  • Core files stay orchestration-first: wrap.py, client.py, cli/proxy.py, and proxy/server.py delegate provider-specific env shaping, API target normalization, backend selection, and transport dispatch.

Install

pip install "headroom-ai[all]"          # Python, everything
npm install headroom-ai                 # TypeScript / Node
docker pull ghcr.io/chopratejas/headroom:latest

Granular extras: [proxy], [mcp], [ml] (Kompress-base), [code], [memory], [relevance], [image], [agno], [langchain], [evals]. Requires Python 3.10+.

Using pipx? Choose a supported interpreter explicitly:

pipx install --python python3.13 "headroom-ai[all]"

Installation guide — Docker tags, persistent service, PowerShell, devcontainers.

headroom learn

headroom learn in action

headroom learn — mines failed sessions, writes corrections to CLAUDE.md / AGENTS.md / GEMINI.md.

Documentation

Start here Go deeper
Quickstart Architecture
Proxy How compression works
MCP tools CCR — reversible compression
Memory Cache optimization
Failure learning Benchmarks
Configuration Limitations

Compared to

Headroom runs locally, covers every content type, works with every major framework, and is reversible.

Scope Deploy Local Reversible
Headroom All context — tools, RAG, logs, files, history Proxy · library · middleware · MCP Yes Yes
RTK CLI command outputs CLI wrapper Yes No
lean-ctx CLI commands, MCP tools, editor rules CLI wrapper · MCP Yes No
Compresr, Token Co. Text sent to their API Hosted API call No No
OpenAI Compaction Conversation history Provider-native No No

Attribution. Headroom ships with the excellent RTK binary for shell-output rewriting — git show --short, scoped ls, summarized installers. Huge thanks to the RTK team; their tool is a first-class part of our stack, and Headroom compresses everything downstream of it. Headroom can also use lean-ctx as the selected CLI context tool; set HEADROOM_CONTEXT_TOOL=lean-ctx before running headroom wrap ....

Contributing

git clone https://github.com/chopratejas/headroom.git && cd headroom
pip install -e ".[dev]" && pytest

Devcontainers in .devcontainer/ (default + memory-stack with Qdrant & Neo4j). See CONTRIBUTING.md.

Community

License

Apache 2.0 — see LICENSE.