From 914a60a2b07caad8488c1e19a5465726b95f83d3 Mon Sep 17 00:00:00 2001 From: Focused Instability <70747559+MrAshRhodes@users.noreply.github.com> Date: Thu, 11 Jun 2026 04:04:45 +0200 Subject: [PATCH] feat(proxy): per-project savings breakdown on the dashboard (claude, codex, aider, copilot, cursor) (#803) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds a **per-project savings breakdown** to the proxy dashboard, covering **all wrap-supported agents: Claude Code, Codex, aider, Copilot, and Cursor**. Spec and rationale in #802 (feature-request issue — happy to adjust scope per maintainer feedback). How it works — two attribution channels, by client capability: **Header channel** (clients that can send custom headers): - `headroom wrap claude` appends `X-Headroom-Project: ` to `ANTHROPIC_CUSTOM_HEADERS` (a user-supplied x-headroom-project header always wins; other user headers preserved). - `headroom wrap codex` extends the injected `[model_providers.headroom]` block with `env_http_headers = { "X-Headroom-Project" = "HEADROOM_PROJECT" }` and sets `HEADROOM_PROJECT` per launch — Codex only sends the header when the env var is set, so the static config stays inert outside wrap. **Base-URL prefix channel** (clients that cannot send custom headers): - The proxy accepts `/p//...`; middleware strips the prefix before routing (before Starlette caches the URL) and binds the project. The explicit header wins over the prefix. - `headroom wrap aider` points `OPENAI_API_BASE` / `ANTHROPIC_BASE_URL` at the prefixed URL. - `headroom wrap copilot` points `COPILOT_PROVIDER_BASE_URL` at the prefixed URL (BYOK anthropic/openai provider types and the GitHub-subscription path). - `headroom wrap cursor` prints the prefixed Override Base URL in its setup instructions, with a note explaining the attribution. **Shared plumbing:** - New `headroom/proxy/project_context.py`: header classification, `/p/` prefix split/strip, URL-prefix builder, and a contextvar bound per request (HTTP middleware + WS accept for the Codex responses bridge); the outcome funnel resolves it (explicit `RequestOutcome.project` wins), stamps `RequestLog.tags["project"]`, and forwards it through `PrometheusMetrics.record_request` into the `SavingsTracker`. - `SavingsTracker`: persisted state gains a `projects` map (requests, tokens saved, savings USD, input tokens/cost, last activity). Schema v2 → v3 with transparent forward migration (v2 files load cleanly, `projects` starts empty). Names sanitized (printable-only, 128-char cap); map capped at 50 projects, evicting the smallest bucket. - `/stats` exposes `savings.per_project` (and `projects`/`projects_limit` inside `persistent_savings`); `/stats-history` exposes `projects`. - Dashboard gains a "Per-Project Savings" table mirroring the per-model table (Alpine `x-text` only — names are user-supplied, no HTML injection). **Behavior changes:** none for unattributed traffic — no header and no prefix means no bucket, aggregate totals exactly as before (regression-tested: legacy `/stats` and `/stats-history` shapes are pinned by tests). ## Real behavior proof **Setup tested on:** macOS (Darwin 25.5.0), Python 3.11.9, `pip install -e ".[dev]"`, proxy on spare ports with isolated `HEADROOM_SAVINGS_PATH`; real Claude Code CLI (subscription auth) and real Codex CLI (ChatGPT auth) as clients. **Header channel — exact steps:** ``` HEADROOM_SAVINGS_PATH=/tmp/headroom-proof-savings.json .venv/bin/python -m headroom.cli proxy --port 9123 # background cd /tmp/proof-alpha && headroom wrap claude --port 9123 --no-rtk --no-mcp --no-serena -- -p "Reply with exactly: OK" cd /tmp/proof-beta && headroom wrap claude --port 9123 --no-rtk --no-mcp --no-serena -- -p "Reply with exactly: OK" cd /tmp/proof-beta && headroom wrap codex --port 9123 --no-rtk --no-mcp --no-serena -- exec --skip-git-repo-check "Reply with exactly: OK" ``` **Observed** (copied live `/stats` output; all runs replied `OK` through the proxy): ```json { "proof-beta": { "requests": 3, "tokens_saved": 140, "compression_savings_usd": 0.0007, "total_input_tokens": 38581, "total_input_cost_usd": 0.360433, "last_activity_at": "2026-06-10T09:19:17Z", "savings_percent": 0.36 }, "proof-alpha": { "requests": 1, "tokens_saved": 0, "compression_savings_usd": 0.0, "total_input_tokens": 9050, "total_input_cost_usd": 0.32245, "last_activity_at": "2026-06-10T09:18:09Z", "savings_percent": 0.0 } } ``` `proof-beta` aggregates one Claude Code turn + one Codex `exec` run (Codex attribution flows through `env_http_headers`). **Prefix channel — exact steps** (the same mechanism the aider/copilot/cursor wraps emit, driven by a real client): ``` .venv/bin/python -m headroom.cli proxy --port 9124 # background, isolated savings path ANTHROPIC_BASE_URL="http://127.0.0.1:9124/p/aider-style-project" claude -p "Reply with exactly: OK" ``` **Observed:** ```json { "aider-style-project": { "requests": 1, "tokens_saved": 0, "compression_savings_usd": 0.0, "total_input_tokens": 8571, "total_input_cost_usd": 1.018575, "last_activity_at": "2026-06-10T10:10:13Z", "savings_percent": 0.0 } } ``` `/stats-history` returned `schema_version: 3` with the same `projects` keys; `GET /dashboard` HTML contains the new "Per-Project Savings" table; persisted state survived a tracker reload; a hand-written v2 savings file loaded cleanly with an empty `projects` map. **What I did NOT test live:** the actual aider/Copilot/Cursor binaries end-to-end (their wraps emit exactly the prefixed URLs exercised above — unit tests pin the emitted env/URLs); Codex subscription-mode routing via the built-in `openai` provider (no provider headers there — such traffic simply stays unattributed); multi-worker uvicorn; Windows. ## Tests - `tests/test_proxy_project_savings.py` (17 tests): sanitization, header classification, `/p/` prefix split + URL-builder round-trip, tracker aggregation/persistence/migration/cardinality-cap/state-sanitization, funnel→`/stats` end-to-end, middleware binding for header + prefix + precedence, plus regression tests pinning legacy `/stats`/`/stats-history` shape and unattributed-traffic totals. - Wrap/provider tests extended: `test_cli/test_wrap_helpers.py`, `test_cli/test_wrap_codex.py`, `test_provider_aider.py`, `test_provider_cursor.py`, `test_provider_copilot_wrap.py` (prefixed env/URLs, user-override wins, no duplicate header, TOML block contents, block strip, setup-line note). - Full `pytest` suite run locally; `ruff check` + `ruff format` clean on all touched files. - `CHANGELOG.md` updated. ## Dependencies None added or bumped. Closes #802 (pending maintainer 👍 per CONTRIBUTING — raised there first with the full spec). --------- Co-authored-by: Ash Rhodes --- CHANGELOG.md | 7 + e2e/wrap/run.py | 17 +- headroom/cli/wrap.py | 73 ++++- headroom/dashboard/templates/dashboard.html | 38 +++ headroom/providers/aider/runtime.py | 16 +- headroom/providers/copilot/wrap.py | 14 +- headroom/providers/cursor/runtime.py | 28 +- headroom/proxy/handlers/openai.py | 4 + headroom/proxy/outcome.py | 9 + headroom/proxy/project_context.py | 106 ++++++ headroom/proxy/prometheus_metrics.py | 2 + headroom/proxy/savings_tracker.py | 141 +++++++- headroom/proxy/server.py | 13 + tests/test_cli/test_wrap_aider.py | 10 +- tests/test_cli/test_wrap_codex.py | 61 ++++ tests/test_cli/test_wrap_copilot.py | 24 +- tests/test_cli/test_wrap_helpers.py | 111 +++++++ tests/test_provider_aider.py | 18 ++ tests/test_provider_copilot_wrap.py | 29 ++ tests/test_provider_cursor.py | 18 ++ tests/test_proxy_project_savings.py | 341 ++++++++++++++++++++ tests/test_proxy_savings_history.py | 4 +- 22 files changed, 1051 insertions(+), 33 deletions(-) create mode 100644 headroom/proxy/project_context.py create mode 100644 tests/test_proxy_project_savings.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c284fb3ba..46cf18156 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,13 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## Unreleased + +### Features + +* **proxy:** per-project savings breakdown on the dashboard for all wrapped agents — Claude Code, Codex, aider, Copilot, and Cursor ([#802](https://github.com/chopratejas/headroom/issues/802)). `headroom wrap claude`/`codex` tag requests with an `X-Headroom-Project` header (launch-directory name); `wrap aider`/`copilot`/`cursor` — whose clients cannot send custom headers — use a `/p/` base-URL prefix the proxy strips. Savings are aggregated per project (persisted, schema v3 with transparent v2 migration), exposed as `savings.per_project` in `/stats` and `projects` in `/stats-history`, and shown in a Per-Project Savings dashboard table. + + ## [0.24.0](https://github.com/chopratejas/headroom/compare/v0.23.0...v0.24.0) (2026-06-08) diff --git a/e2e/wrap/run.py b/e2e/wrap/run.py index 4597f68d9..38e57a870 100644 --- a/e2e/wrap/run.py +++ b/e2e/wrap/run.py @@ -14,6 +14,7 @@ import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import Any +from urllib.parse import quote import httpx @@ -654,19 +655,23 @@ def verify_aider_wrap(base_env: dict[str, str], project_dir: Path, log_dir: Path entries = read_jsonl(log_dir / "aider.jsonl") assert_true(len(entries) > 0, "Aider shim should have been invoked") env_vars = entries[-1]["env"] + # Aider cannot send custom headers, so its wrap embeds the launch + # directory as a /p/ base-URL prefix for per-project savings; + # the proxy strips it before routing, so the probes still succeed. + project_prefix = f"/p/{quote(project_dir.name, safe='')}" assert_true( - env_vars.get("OPENAI_API_BASE") == f"http://127.0.0.1:{port}/v1", + env_vars.get("OPENAI_API_BASE") == f"http://127.0.0.1:{port}{project_prefix}/v1", "Aider wrap should set OPENAI_API_BASE", ) assert_true( - env_vars.get("ANTHROPIC_BASE_URL") == f"http://127.0.0.1:{port}", + env_vars.get("ANTHROPIC_BASE_URL") == f"http://127.0.0.1:{port}{project_prefix}", "Aider wrap should set ANTHROPIC_BASE_URL", ) assert_true( entries[-1]["probes"] == [ - {"url": f"http://127.0.0.1:{port}/v1/models", "status": 200}, - {"url": f"http://127.0.0.1:{port}/health", "status": 200}, + {"url": f"http://127.0.0.1:{port}{project_prefix}/v1/models", "status": 200}, + {"url": f"http://127.0.0.1:{port}{project_prefix}/health", "status": 200}, ], "Aider shim should prove both configured base URLs point at a live proxy", ) @@ -686,8 +691,10 @@ def verify_cursor_wrap(base_env: dict[str, str], project_dir: Path) -> None: ) try: output = wait_for_output(proc, "Press Ctrl+C to stop the proxy.", timeout=30) + # Cursor setup lines embed the /p/ per-project prefix. + cursor_prefix = f"/p/{quote(project_dir.name, safe='')}" assert_true( - f"http://127.0.0.1:{port}/v1" in output, + f"http://127.0.0.1:{port}{cursor_prefix}/v1" in output, "Cursor wrap should print the OpenAI base URL override", ) assert_true( diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 353383d95..da00b4064 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -87,6 +87,7 @@ from headroom.providers.openclaw import ( from headroom.providers.openclaw import ( normalize_gateway_provider_ids as _normalize_openclaw_gateway_provider_ids_impl, ) +from headroom.proxy.project_context import with_project_prefix as _with_project_prefix from .main import main @@ -944,6 +945,47 @@ def _prepare_wrap_rtk(verbose: bool = False, *, label: str | None = None) -> Pat return _ensure_rtk_binary(verbose=verbose) +# Canonical casing for the proxy's per-project savings header (matched +# case-insensitively by headroom.proxy.project_context.PROJECT_HEADER). +_PROJECT_HEADER_NAME = "X-Headroom-Project" + + +def _project_name_from_cwd() -> str | None: + """Project label for X-Headroom-Project: basename of the launch directory. + + The proxy sanitizes and caps the value server-side + (headroom.proxy.savings_tracker.sanitize_project_name), so the raw + directory name is safe to send as-is. + """ + name = Path.cwd().name.strip() + return name or None + + +def _apply_project_header_env(env: dict[str, str]) -> None: + """Inject X-Headroom-Project into ``ANTHROPIC_CUSTOM_HEADERS``. + + Claude Code reads ``ANTHROPIC_CUSTOM_HEADERS`` as newline-separated + ``Name: value`` lines and attaches them to every API request; the + Headroom proxy uses the X-Headroom-Project header for per-project + savings attribution. An existing user-supplied x-headroom-project + header (any casing) always wins — we never duplicate or overwrite it, + and any other user headers are preserved by appending. + """ + project = _project_name_from_cwd() + if not project: + return + header_line = f"{_PROJECT_HEADER_NAME}: {project}" + existing = env.get("ANTHROPIC_CUSTOM_HEADERS") + if existing: + for line in existing.splitlines(): + name = line.split(":", 1)[0].strip() + if name.lower() == _PROJECT_HEADER_NAME.lower(): + return # user override wins + env["ANTHROPIC_CUSTOM_HEADERS"] = f"{existing}\n{header_line}" + else: + env["ANTHROPIC_CUSTOM_HEADERS"] = header_line + + def _inject_codex_provider_config(port: int) -> None: """Inject a Headroom model provider into Codex's config.toml. @@ -985,6 +1027,11 @@ def _inject_codex_provider_config(port: int) -> None: 'name = "OpenAI via Headroom proxy"\n' f'base_url = "http://127.0.0.1:{port}/v1"\n' f"supports_websockets = true\n" + # Per-project savings: Codex sends the header only when the mapped + # env var (HEADROOM_PROJECT, set by `headroom wrap codex`) exists at + # Codex runtime. Inline table keeps the key inside this section so + # _strip_codex_headroom_blocks removes it with the rest of the block. + f'env_http_headers = {{ "{_PROJECT_HEADER_NAME}" = "HEADROOM_PROJECT" }}\n' f"{_CODEX_END_MARKER}\n" ) @@ -2494,6 +2541,10 @@ def claude( else: env["ANTHROPIC_BASE_URL"] = proxy_url + # Per-project savings attribution: tag every request with the launch + # directory's name via X-Headroom-Project (user override wins). + _apply_project_header_env(env) + # Issue #746: keep Claude Code's on-demand tool loading on through the # proxy so tool schemas are not eagerly materialized into local context. _tool_search_value = _configure_tool_search_env(env, tool_search) @@ -2743,7 +2794,11 @@ def copilot( _copilot_default_wire_api_for_model(selected_model) if subscription else "completions" ) env["COPILOT_PROVIDER_TYPE"] = "openai" - env["COPILOT_PROVIDER_BASE_URL"] = f"http://127.0.0.1:{port}/v1" + # Per-project savings: the Copilot CLI cannot send custom headers, so + # the project rides as a /p/ base-URL prefix the proxy strips. + env["COPILOT_PROVIDER_BASE_URL"] = _with_project_prefix( + f"http://127.0.0.1:{port}/v1", _project_name_from_cwd() + ) env["COPILOT_PROVIDER_WIRE_API"] = effective_wire_api env["COPILOT_PROVIDER_BEARER_TOKEN"] = client_bearer env["GITHUB_COPILOT_USE_TOKEN_EXCHANGE"] = "false" @@ -2760,7 +2815,7 @@ def copilot( copilot_proxy_token = client_bearer env_vars_display = [ "COPILOT_PROVIDER_TYPE=openai", - f"COPILOT_PROVIDER_BASE_URL=http://127.0.0.1:{port}/v1", + f"COPILOT_PROVIDER_BASE_URL={env['COPILOT_PROVIDER_BASE_URL']}", f"COPILOT_PROVIDER_WIRE_API={effective_wire_api}", ( "COPILOT_AUTH_MODE=github-subscription-experimental" @@ -2787,6 +2842,7 @@ def copilot( provider_type=effective_provider_type, wire_api=wire_api, environ=env, + project=_project_name_from_cwd(), ) if not env.get("COPILOT_PROVIDER_API_KEY"): @@ -3008,6 +3064,13 @@ def codex( env, env_vars_display = _build_codex_launch_env(port, os.environ) + # Per-project savings attribution: the injected provider config maps the + # X-Headroom-Project header to HEADROOM_PROJECT via env_http_headers, so + # Codex sends it only when this var is set. A user-set value wins. + _codex_project = _project_name_from_cwd() + if _codex_project and "HEADROOM_PROJECT" not in env: + env["HEADROOM_PROJECT"] = _codex_project + # Inject Headroom provider into Codex config so WebSocket traffic also # routes through the proxy. Codex ignores OPENAI_BASE_URL for its WS # transport unless a custom provider declares supports_websockets = true. @@ -3119,7 +3182,9 @@ def aider( click.echo("Install aider: pip install aider-chat") raise SystemExit(1) - env, env_vars_display = _build_aider_launch_env(port, os.environ) + env, env_vars_display = _build_aider_launch_env( + port, os.environ, project=_project_name_from_cwd() + ) _launch_tool( binary=aider_bin, @@ -3202,7 +3267,7 @@ def cursor( return def _print_cursor_setup() -> None: - for line in _render_cursor_setup_lines(port): + for line in _render_cursor_setup_lines(port, project=_project_name_from_cwd()): click.echo(line) if not no_rtk: click.echo() diff --git a/headroom/dashboard/templates/dashboard.html b/headroom/dashboard/templates/dashboard.html index 58a72988a..a2efe9369 100644 --- a/headroom/dashboard/templates/dashboard.html +++ b/headroom/dashboard/templates/dashboard.html @@ -886,6 +886,44 @@ + + +
diff --git a/headroom/providers/aider/runtime.py b/headroom/providers/aider/runtime.py index e382b282a..8ced698b4 100644 --- a/headroom/providers/aider/runtime.py +++ b/headroom/providers/aider/runtime.py @@ -7,15 +7,23 @@ from collections.abc import Mapping from headroom.providers.claude import proxy_base_url as claude_proxy_base_url from headroom.providers.codex import proxy_base_url as codex_proxy_base_url +from headroom.proxy.project_context import with_project_prefix def build_launch_env( - port: int, environ: Mapping[str, str] | None = None + port: int, + environ: Mapping[str, str] | None = None, + project: str | None = None, ) -> tuple[dict[str, str], list[str]]: - """Build environment variables for Aider through the local proxy.""" + """Build environment variables for Aider through the local proxy. + + ``project`` (the wrap launch directory) is encoded as a ``/p/`` + base-URL prefix because aider cannot send custom headers; the proxy + strips it and attributes savings per project. + """ env = dict(environ or os.environ) - openai_base_url = codex_proxy_base_url(port) - anthropic_base_url = claude_proxy_base_url(port) + openai_base_url = with_project_prefix(codex_proxy_base_url(port), project) + anthropic_base_url = with_project_prefix(claude_proxy_base_url(port), project) env["OPENAI_API_BASE"] = openai_base_url env["ANTHROPIC_BASE_URL"] = anthropic_base_url return env, [ diff --git a/headroom/providers/copilot/wrap.py b/headroom/providers/copilot/wrap.py index cd5cfe313..672e50954 100644 --- a/headroom/providers/copilot/wrap.py +++ b/headroom/providers/copilot/wrap.py @@ -11,6 +11,8 @@ from typing import Any import click +from headroom.proxy.project_context import with_project_prefix + def resolve_provider_type( backend: str | None, provider_type: str, environ: Mapping[str, str] | None = None @@ -113,8 +115,14 @@ def build_launch_env( provider_type: str, wire_api: str | None, environ: Mapping[str, str] | None = None, + project: str | None = None, ) -> tuple[dict[str, str], list[str]]: - """Build the Copilot BYOK environment for the selected provider type.""" + """Build the Copilot BYOK environment for the selected provider type. + + ``project`` (the wrap launch directory) is encoded as a ``/p/`` + base-URL prefix because the Copilot CLI cannot send custom headers; the + proxy strips it and attributes savings per project. + """ # Distinguish "caller passed nothing" (use os.environ) from "caller # explicitly passed an empty dict" (start fresh — the test/CLI is in # charge of which keys to seed). The previous `environ or os.environ` @@ -129,7 +137,7 @@ def build_launch_env( env["COPILOT_PROVIDER_API_KEY"] = key if provider_type == "anthropic": - base_url = f"http://127.0.0.1:{port}" + base_url = with_project_prefix(f"http://127.0.0.1:{port}", project) env["COPILOT_PROVIDER_BASE_URL"] = base_url return env, [ "COPILOT_PROVIDER_TYPE=anthropic", @@ -137,7 +145,7 @@ def build_launch_env( ] effective_wire_api = wire_api or "completions" - base_url = f"http://127.0.0.1:{port}/v1" + base_url = with_project_prefix(f"http://127.0.0.1:{port}/v1", project) env["COPILOT_PROVIDER_BASE_URL"] = base_url env["COPILOT_PROVIDER_WIRE_API"] = effective_wire_api return env, [ diff --git a/headroom/providers/cursor/runtime.py b/headroom/providers/cursor/runtime.py index 7a2f33135..b67f36456 100644 --- a/headroom/providers/cursor/runtime.py +++ b/headroom/providers/cursor/runtime.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from headroom.providers.claude import proxy_base_url as claude_proxy_base_url from headroom.providers.codex import proxy_base_url as codex_proxy_base_url +from headroom.proxy.project_context import with_project_prefix @dataclass(frozen=True) @@ -16,18 +17,23 @@ class CursorProxyTargets: anthropic_base_url: str -def build_proxy_targets(port: int) -> CursorProxyTargets: - """Build the local proxy URLs shown to Cursor users.""" +def build_proxy_targets(port: int, project: str | None = None) -> CursorProxyTargets: + """Build the local proxy URLs shown to Cursor users. + + ``project`` (the wrap launch directory) is encoded as a ``/p/`` + base-URL prefix because Cursor cannot send custom headers; the proxy + strips it and attributes savings per project. + """ return CursorProxyTargets( - openai_base_url=codex_proxy_base_url(port), - anthropic_base_url=claude_proxy_base_url(port), + openai_base_url=with_project_prefix(codex_proxy_base_url(port), project), + anthropic_base_url=with_project_prefix(claude_proxy_base_url(port), project), ) -def render_setup_lines(port: int) -> list[str]: +def render_setup_lines(port: int, project: str | None = None) -> list[str]: """Render the Cursor setup instructions for the local proxy.""" - targets = build_proxy_targets(port) - return [ + targets = build_proxy_targets(port, project) + lines = [ " Headroom proxy is running. Configure Cursor:", "", " For OpenAI models:", @@ -42,3 +48,11 @@ def render_setup_lines(port: int) -> list[str]: " Settings > Models > OpenAI API Key > Override OpenAI Base URL", f" Set to: {targets.openai_base_url}", ] + if project: + lines += [ + "", + f" Dashboard savings will be attributed to project '{project}'", + " (the directory this command was run from). Re-run from another", + " project directory to get that project's URL.", + ] + return lines diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 80802a2b8..be9b0a85a 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -47,6 +47,7 @@ from headroom.proxy.auth_mode import classify_auth_mode, classify_client from headroom.proxy.compression_decision import CompressionDecision from headroom.proxy.cost import _summarize_transforms from headroom.proxy.outcome import RequestOutcome +from headroom.proxy.project_context import classify_project, set_current_project logger = logging.getLogger("headroom.proxy") @@ -3381,6 +3382,9 @@ class OpenAIHandlerMixin: # Identify the WS harness before downstream auth/header rewrites. # Captured in closure so per-turn RequestOutcome can stamp it. client = classify_client(ws_headers) + # WS sessions bypass the HTTP middleware, so bind the project here; + # per-turn outcome emission inside this task inherits the context. + set_current_project(classify_project(ws_headers)) _ws_url_obj = getattr(websocket, "url", None) _ws_url = str(_ws_url_obj) if _ws_url_obj is not None else "" _ws_path = getattr(_ws_url_obj, "path", "") if _ws_url_obj is not None else "" diff --git a/headroom/proxy/outcome.py b/headroom/proxy/outcome.py index 37f5d2643..ebd17516d 100644 --- a/headroom/proxy/outcome.py +++ b/headroom/proxy/outcome.py @@ -129,6 +129,7 @@ class RequestOutcome: request_messages: list[dict[str, Any]] | None = None tags: dict[str, str] = field(default_factory=dict) client: str | None = None + project: str | None = None # ── Derived (computed once, no caching needed — properties are cheap) ─ @@ -303,6 +304,11 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None: """ from headroom.proxy.cost import _summarize_transforms from headroom.proxy.models import RequestLog + from headroom.proxy.project_context import get_current_project + + # Project attribution: explicit outcome field wins, else the value the + # HTTP middleware / WS accept captured from ``X-Headroom-Project``. + project = outcome.project or get_current_project() # 1. Prometheus / SavingsTracker. await handler.metrics.record_request( @@ -323,6 +329,7 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None: 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, ) # 2. Cost tracker (optional). @@ -349,6 +356,8 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None: log_tags = dict(outcome.tags) if outcome.client: log_tags["client"] = outcome.client + if project: + log_tags["project"] = project request_logger.log( RequestLog( request_id=outcome.request_id, diff --git a/headroom/proxy/project_context.py b/headroom/proxy/project_context.py new file mode 100644 index 000000000..fe48a7163 --- /dev/null +++ b/headroom/proxy/project_context.py @@ -0,0 +1,106 @@ +"""Per-request project attribution for the proxy. + +``headroom wrap`` launches agents with an ``X-Headroom-Project`` header +(via ``ANTHROPIC_CUSTOM_HEADERS`` for Claude Code and ``env_http_headers`` +for Codex) naming the project directory the agent is working in. The proxy +captures that header once per request — in the HTTP middleware for regular +requests and at the WebSocket accept for Codex responses-WS sessions — +into a :mod:`contextvars` variable, so the outcome funnel can attribute +savings to a project without threading a parameter through every handler. + +The value is sanitized (printable characters only, length-capped) before it +is stored; an absent or unusable header simply leaves attribution off for +that request, matching pre-feature behavior. +""" + +from __future__ import annotations + +from collections.abc import Mapping, MutableMapping +from contextvars import ContextVar +from typing import Any +from urllib.parse import quote, unquote, urlsplit, urlunsplit + +from headroom.proxy.savings_tracker import sanitize_project_name + +PROJECT_HEADER = "x-headroom-project" +PROJECT_PATH_PREFIX = "/p/" + +_current_project: ContextVar[str | None] = ContextVar("headroom_current_project", default=None) + + +def classify_project(headers: Mapping[str, Any] | Any) -> str | None: + """Extract a sanitized project name from request headers, if present.""" + get = getattr(headers, "get", None) + if get is None: + return None + value = get(PROJECT_HEADER) or get("X-Headroom-Project") + return sanitize_project_name(value) + + +def set_current_project(project: str | None) -> None: + """Bind the active request's project for downstream outcome recording.""" + _current_project.set(sanitize_project_name(project)) + + +def get_current_project() -> str | None: + """Project bound to the current request context, or ``None``.""" + return _current_project.get() + + +def split_project_path(path: str) -> tuple[str | None, str]: + """Split ``/p//rest`` into ``(name, /rest)``. + + Clients that cannot send custom headers (aider, Copilot BYOK, Cursor) + are pointed at a project-prefixed base URL instead; the first path + segment after ``/p/`` is the URL-encoded project name. Returns + ``(None, path)`` unchanged when the prefix is absent or unusable. + """ + if not path.startswith(PROJECT_PATH_PREFIX): + return None, path + remainder = path[len(PROJECT_PATH_PREFIX) :] + segment, sep, rest = remainder.partition("/") + project = sanitize_project_name(unquote(segment)) if segment else None + if project is None: + return None, path + return project, ("/" + rest) if sep else "/" + + +def strip_project_path_prefix(scope: MutableMapping[str, Any]) -> str | None: + """Strip a ``/p/`` prefix from an ASGI scope, returning the name. + + Mutates ``scope["path"]`` (and ``raw_path``) so routing sees the + canonical path. Must run before anything caches the request URL. + """ + project, stripped = split_project_path(scope.get("path", "")) + if project is not None: + scope["path"] = stripped + if "raw_path" in scope: + scope["raw_path"] = quote(stripped).encode("ascii") + return project + + +def with_project_prefix(base_url: str, project: str | None) -> str: + """Insert ``/p/`` ahead of the path of a local proxy base URL. + + Producer-side counterpart of :func:`split_project_path`, used by + ``headroom wrap`` for clients that cannot send custom headers. + Returns ``base_url`` unchanged when the project name is unusable. + """ + name = sanitize_project_name(project) + if name is None: + return base_url + parts = urlsplit(base_url) + prefixed = f"{PROJECT_PATH_PREFIX}{quote(name, safe='')}{parts.path}" + return urlunsplit(parts._replace(path=prefixed.rstrip("/"))) + + +__all__ = [ + "PROJECT_HEADER", + "PROJECT_PATH_PREFIX", + "classify_project", + "get_current_project", + "set_current_project", + "split_project_path", + "strip_project_path_prefix", + "with_project_prefix", +] diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index edfce23c9..a1bf43019 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -562,6 +562,7 @@ class PrometheusMetrics: cache_write_1h_tokens: int = 0, uncached_input_tokens: int = 0, attempted_input_tokens: int = 0, + project: str | None = None, ): """Record metrics for a request.""" async with self._lock: @@ -649,6 +650,7 @@ class PrometheusMetrics: input_tokens=input_tokens, tokens_saved=tokens_saved, provider=provider, + project=project, cache_read_tokens=cache_read_tokens, cache_write_tokens=cache_write_tokens, uncached_input_tokens=uncached_input_tokens, diff --git a/headroom/proxy/savings_tracker.py b/headroom/proxy/savings_tracker.py index 30e6deb8f..6a7ebfba2 100644 --- a/headroom/proxy/savings_tracker.py +++ b/headroom/proxy/savings_tracker.py @@ -26,8 +26,10 @@ logger = logging.getLogger(__name__) HEADROOM_SAVINGS_PATH_ENV_VAR = _paths.HEADROOM_SAVINGS_PATH_ENV DEFAULT_SAVINGS_DIR = ".headroom" DEFAULT_SAVINGS_FILE = "proxy_savings.json" -SCHEMA_VERSION = 2 +SCHEMA_VERSION = 3 DEFAULT_MAX_HISTORY_POINTS = 5000 +DEFAULT_MAX_PROJECTS = 50 +PROJECT_NAME_MAX_LENGTH = 128 DEFAULT_MAX_HISTORY_AGE_DAYS = 365 DEFAULT_MAX_RESPONSE_HISTORY_POINTS = 500 DEFAULT_DISPLAY_SESSION_INACTIVITY_MINUTES = 60 @@ -287,6 +289,64 @@ def _empty_display_session() -> dict[str, Any]: } +def sanitize_project_name(value: Any) -> str | None: + """Normalize a client-supplied project name; ``None`` when unusable. + + Strips control characters, trims whitespace, and caps length so a + misbehaving client cannot bloat the persisted state or the dashboard. + """ + if not isinstance(value, str): + return None + cleaned = "".join(ch for ch in value if ch.isprintable()).strip() + if not cleaned: + return None + return cleaned[:PROJECT_NAME_MAX_LENGTH] + + +def _empty_project_entry() -> dict[str, Any]: + return { + "requests": 0, + "tokens_saved": 0, + "compression_savings_usd": 0.0, + "total_input_tokens": 0, + "total_input_cost_usd": 0.0, + "last_activity_at": None, + } + + +def _normalize_projects(raw: Any) -> dict[str, dict[str, Any]]: + if not isinstance(raw, dict): + return {} + projects: dict[str, dict[str, Any]] = {} + for name, entry in raw.items(): + cleaned_name = sanitize_project_name(name) + if cleaned_name is None or not isinstance(entry, dict): + continue + normalized = _empty_project_entry() + normalized["requests"] = _coerce_int(entry.get("requests")) + normalized["tokens_saved"] = _coerce_int(entry.get("tokens_saved")) + normalized["compression_savings_usd"] = round( + _coerce_float(entry.get("compression_savings_usd")), 6 + ) + normalized["total_input_tokens"] = _coerce_int(entry.get("total_input_tokens")) + normalized["total_input_cost_usd"] = round( + _coerce_float(entry.get("total_input_cost_usd")), 6 + ) + last_activity = _parse_timestamp(entry.get("last_activity_at")) + normalized["last_activity_at"] = _to_utc_iso(last_activity) if last_activity else None + projects[cleaned_name] = normalized + if len(projects) > DEFAULT_MAX_PROJECTS: + # Oversized persisted maps (hand-edited or future versions) would + # otherwise shrink only one entry per recorded request. + kept = sorted( + projects.items(), + key=lambda item: (item[1]["tokens_saved"], item[1]["last_activity_at"] or ""), + reverse=True, + )[:DEFAULT_MAX_PROJECTS] + projects = dict(kept) + return projects + + def _normalize_display_session(entry: Any) -> dict[str, Any]: if not isinstance(entry, dict): return _empty_display_session() @@ -427,6 +487,7 @@ class SavingsTracker: input_tokens: int, tokens_saved: int, provider: str | None = None, + project: str | None = None, cache_read_tokens: int = 0, cache_write_tokens: int = 0, uncached_input_tokens: int = 0, @@ -526,6 +587,16 @@ class SavingsTracker: if session.get("started_at") is None: session["started_at"] = session["last_activity_at"] + self._record_project_locked( + project, + timestamp_dt=timestamp_dt, + requests_delta=1, + tokens_saved_delta=delta_tokens_saved, + savings_usd_delta=delta_savings_usd, + input_tokens_delta=delta_input_tokens, + input_cost_usd_delta=delta_input_cost_usd, + ) + if delta_tokens_saved > 0: self._state["history"].append( { @@ -542,6 +613,67 @@ class SavingsTracker: self._save_locked() return True + def _record_project_locked( + self, + project: str | None, + *, + timestamp_dt: datetime, + requests_delta: int = 0, + tokens_saved_delta: int = 0, + savings_usd_delta: float = 0.0, + input_tokens_delta: int = 0, + input_cost_usd_delta: float = 0.0, + ) -> None: + """Accumulate per-project savings. Caller must hold ``self._lock``. + + Unattributed traffic (``project`` missing or unusable) is skipped so + existing aggregate behavior is unchanged. The map is capped at + ``DEFAULT_MAX_PROJECTS`` entries, evicting the smallest/oldest bucket. + """ + name = sanitize_project_name(project) + if name is None: + return + projects: dict[str, dict[str, Any]] = self._state.setdefault("projects", {}) + entry = projects.setdefault(name, _empty_project_entry()) + entry["requests"] += max(requests_delta, 0) + entry["tokens_saved"] += max(tokens_saved_delta, 0) + entry["compression_savings_usd"] = round( + entry["compression_savings_usd"] + max(savings_usd_delta, 0.0), 6 + ) + entry["total_input_tokens"] += max(input_tokens_delta, 0) + entry["total_input_cost_usd"] = round( + entry["total_input_cost_usd"] + max(input_cost_usd_delta, 0.0), 6 + ) + entry["last_activity_at"] = _to_utc_iso(timestamp_dt) + if len(projects) > DEFAULT_MAX_PROJECTS: + evict = min( + (key for key in projects if key != name), + key=lambda key: ( + projects[key]["tokens_saved"], + projects[key]["last_activity_at"] or "", + ), + ) + del projects[evict] + + def _projects_snapshot_locked(self) -> dict[str, dict[str, Any]]: + """Per-project stats with a derived ``savings_percent``, sorted by savings.""" + projects = self._state.get("projects", {}) + ranked = sorted( + projects.items(), + key=lambda item: item[1]["tokens_saved"], + reverse=True, + ) + result: dict[str, dict[str, Any]] = {} + for name, entry in ranked: + view = dict(entry) + total_before = entry["tokens_saved"] + entry["total_input_tokens"] + view["savings_percent"] = round( + (entry["tokens_saved"] / total_before * 100) if total_before > 0 else 0.0, + 2, + ) + result[name] = view + return result + def stats_preview(self, recent_points: int = 20) -> dict[str, Any]: """Return a compact preview for `/stats`.""" snapshot = self.snapshot() @@ -554,6 +686,8 @@ class SavingsTracker: "history_points": len(snapshot["history"]), "recent_history": snapshot["history"][-recent_points:], "retention": snapshot["retention"], + "projects": snapshot["projects"], + "projects_limit": DEFAULT_MAX_PROJECTS, } def history_response(self, history_mode: str = "compact") -> dict[str, Any]: @@ -582,6 +716,7 @@ class SavingsTracker: "available_series": ["history", *series.keys()], }, "retention": snapshot["retention"], + "projects": snapshot["projects"], "history_summary": { "mode": history_mode, "stored_points": len(raw_history), @@ -645,6 +780,7 @@ class SavingsTracker: "max_history_age_days": self._max_history_age_days, "max_response_history_points": self._max_response_history_points, }, + "projects": self._projects_snapshot_locked(), } def _default_state(self) -> dict[str, Any]: @@ -659,6 +795,7 @@ class SavingsTracker: }, "display_session": _empty_display_session(), "history": [], + "projects": {}, } def _load_state(self) -> dict[str, Any]: @@ -731,6 +868,7 @@ class SavingsTracker: }, "display_session": _normalize_display_session(raw.get("display_session")), "history": normalized_history, + "projects": _normalize_projects(raw.get("projects")), } if normalized_history: @@ -822,6 +960,7 @@ class SavingsTracker: "lifetime": self._state["lifetime"], "display_session": self._state["display_session"], "history": self._state["history"], + "projects": self._state.get("projects", {}), } json_data = json.dumps(payload, indent=2) diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 613fe6fb6..afac7fa9d 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -139,6 +139,11 @@ from headroom.proxy.modes import ( is_token_mode, normalize_proxy_mode, ) +from headroom.proxy.project_context import ( + classify_project, + set_current_project, + strip_project_path_prefix, +) from headroom.proxy.prometheus_metrics import PrometheusMetrics # noqa: F401 from headroom.proxy.rate_limiter import TokenBucketRateLimiter # noqa: F401 from headroom.proxy.request_logger import RequestLogger # noqa: F401 @@ -1709,10 +1714,17 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: async def _record_headroom_stack(request, call_next): started = time.perf_counter() inbound_id = f"inbound-{time.time_ns()}" + # Project attribution: an explicit X-Headroom-Project header wins + # (claude/codex wraps); otherwise a /p/ base-URL prefix (aider, + # Copilot BYOK, Cursor — clients that cannot send custom headers). + # The prefix strip mutates the scope, so it must happen before + # request.url is first accessed (Starlette caches the URL). + prefix_project = strip_project_path_prefix(request.scope) path = request.url.path method = request.method query = request.url.query headers = dict(request.headers.items()) + set_current_project(classify_project(headers) or prefix_project) client = getattr(request, "client", None) client_addr = "" if client is not None: @@ -2024,6 +2036,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: "summary": summary, "savings": { "total_tokens": total_tokens_all_layers, + "per_project": persistent_savings.get("projects", {}), "by_layer": { "cli_filtering": { "tool": cli_filtering_tool, diff --git a/tests/test_cli/test_wrap_aider.py b/tests/test_cli/test_wrap_aider.py index 83a6486bb..9ac91fdfe 100644 --- a/tests/test_cli/test_wrap_aider.py +++ b/tests/test_cli/test_wrap_aider.py @@ -4,6 +4,7 @@ from __future__ import annotations from pathlib import Path from unittest.mock import patch +from urllib.parse import quote import pytest from click.testing import CliRunner @@ -11,6 +12,11 @@ from click.testing import CliRunner from headroom.cli.main import main +def _expected_project_prefix() -> str: + """The /p/ prefix the wrap now embeds (launch-directory basename).""" + return f"/p/{quote(Path.cwd().name, safe='')}" + + @pytest.fixture def runner() -> CliRunner: return CliRunner() @@ -32,8 +38,8 @@ def test_wrap_aider_sets_provider_envs( assert result.exit_code == 0, result.output env = captured["env"] assert isinstance(env, dict) - assert env["OPENAI_API_BASE"] == "http://127.0.0.1:8787/v1" - assert env["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:8787" + assert env["OPENAI_API_BASE"] == f"http://127.0.0.1:8787{_expected_project_prefix()}/v1" + assert env["ANTHROPIC_BASE_URL"] == f"http://127.0.0.1:8787{_expected_project_prefix()}" assert captured["tool_label"] == "AIDER" assert captured["agent_type"] == "aider" assert captured["args"] == ("--model", "gpt-4o") diff --git a/tests/test_cli/test_wrap_codex.py b/tests/test_cli/test_wrap_codex.py index f6aedefd2..a0f3f813e 100644 --- a/tests/test_cli/test_wrap_codex.py +++ b/tests/test_cli/test_wrap_codex.py @@ -741,3 +741,64 @@ def test_unwrap_codex_preserves_unrelated_sections( assert result.exit_code == 0, result.output restored = config_file.read_text() assert restored == original + + +# --------------------------------------------------------------------------- +# Per-project savings: env_http_headers in the injected provider block +# --------------------------------------------------------------------------- + + +class TestCodexProjectHeaderConfig: + """The injected provider maps X-Headroom-Project to HEADROOM_PROJECT. + + Codex's ``env_http_headers`` sends a header only when the mapped env var + is set at Codex runtime, so `headroom wrap codex` exports + ``HEADROOM_PROJECT`` and the proxy attributes savings per project. + """ + + def test_inject_writes_env_http_headers_mapping( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + _set_test_home(monkeypatch, tmp_path) + + wrap_mod._inject_codex_provider_config(8787) + + content = (tmp_path / ".codex" / "config.toml").read_text() + assert 'env_http_headers = { "X-Headroom-Project" = "HEADROOM_PROJECT" }' in content + + def test_env_http_headers_inside_provider_section( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """The mapping must live inside [model_providers.headroom], before + the closing marker, so it applies to the Headroom provider.""" + _set_test_home(monkeypatch, tmp_path) + + wrap_mod._inject_codex_provider_config(8787) + + content = (tmp_path / ".codex" / "config.toml").read_text() + section_start = content.index("[model_providers.headroom]") + mapping_pos = content.index("env_http_headers") + end_marker_pos = content.index(wrap_mod._CODEX_END_MARKER, section_start) + assert section_start < mapping_pos < end_marker_pos + + def test_strip_removes_block_with_env_http_headers( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + """_strip_codex_headroom_blocks removes the whole injected block, + including the new env_http_headers line, leaving user content.""" + _set_test_home(monkeypatch, tmp_path) + config_dir = tmp_path / ".codex" + config_dir.mkdir() + config_file = config_dir / "config.toml" + original = '[profiles.default]\nmodel = "gpt-4o"\n' + config_file.write_text(original) + + wrap_mod._inject_codex_provider_config(8787) + wrapped = config_file.read_text() + assert "env_http_headers" in wrapped + + cleaned = wrap_mod._strip_codex_headroom_blocks(wrapped) + assert "env_http_headers" not in cleaned + assert "X-Headroom-Project" not in cleaned + assert "[model_providers.headroom]" not in cleaned + assert 'model = "gpt-4o"' in cleaned diff --git a/tests/test_cli/test_wrap_copilot.py b/tests/test_cli/test_wrap_copilot.py index 0e44977b7..cc2f0c173 100644 --- a/tests/test_cli/test_wrap_copilot.py +++ b/tests/test_cli/test_wrap_copilot.py @@ -8,6 +8,7 @@ import sys import types from pathlib import Path from unittest.mock import patch +from urllib.parse import quote import click import pytest @@ -16,6 +17,11 @@ from click.testing import CliRunner from headroom.copilot_auth import DEFAULT_API_URL +def _expected_project_prefix() -> str: + """The /p/ prefix the wrap now embeds (launch-directory basename).""" + return f"/p/{quote(Path.cwd().name, safe='')}" + + @pytest.fixture def runner() -> CliRunner: return CliRunner() @@ -98,7 +104,7 @@ def test_wrap_copilot_auto_anthropic_injects_instructions( env = captured["env"] assert isinstance(env, dict) assert env["COPILOT_PROVIDER_TYPE"] == "anthropic" - assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787" + assert env["COPILOT_PROVIDER_BASE_URL"] == f"http://127.0.0.1:8787{_expected_project_prefix()}" assert "COPILOT_PROVIDER_WIRE_API" not in env assert captured["agent_type"] == "copilot" assert captured["tool_label"] == "COPILOT" @@ -145,7 +151,9 @@ def test_wrap_copilot_openai_backend_sets_completions_env( env = captured["env"] assert isinstance(env, dict) assert env["COPILOT_PROVIDER_TYPE"] == "openai" - assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1" + assert env["COPILOT_PROVIDER_BASE_URL"] == ( + f"http://127.0.0.1:8787{_expected_project_prefix()}/v1" + ) assert env["COPILOT_PROVIDER_WIRE_API"] == "completions" assert captured["backend"] == "anyllm" assert captured["anyllm_provider"] == "groq" @@ -181,7 +189,9 @@ def test_wrap_copilot_auto_detects_running_proxy_backend( env = captured["env"] assert isinstance(env, dict) assert env["COPILOT_PROVIDER_TYPE"] == "openai" - assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1" + assert env["COPILOT_PROVIDER_BASE_URL"] == ( + f"http://127.0.0.1:8787{_expected_project_prefix()}/v1" + ) assert env["COPILOT_PROVIDER_WIRE_API"] == "completions" @@ -210,7 +220,9 @@ def test_wrap_copilot_prefers_existing_oauth_session( env = captured["env"] assert isinstance(env, dict) assert env["COPILOT_PROVIDER_TYPE"] == "openai" - assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1" + assert env["COPILOT_PROVIDER_BASE_URL"] == ( + f"http://127.0.0.1:8787{_expected_project_prefix()}/v1" + ) assert env["COPILOT_PROVIDER_WIRE_API"] == "completions" assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "gho-existing" assert env["GITHUB_COPILOT_API_URL"] == DEFAULT_API_URL @@ -249,7 +261,9 @@ def test_wrap_copilot_subscription_uses_github_auth_without_provider_key( env = captured["env"] assert isinstance(env, dict) assert env["COPILOT_PROVIDER_TYPE"] == "openai" - assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1" + assert env["COPILOT_PROVIDER_BASE_URL"] == ( + f"http://127.0.0.1:8787{_expected_project_prefix()}/v1" + ) assert env["COPILOT_PROVIDER_WIRE_API"] == "completions" assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "gho-existing" assert "COPILOT_PROVIDER_API_KEY" not in env diff --git a/tests/test_cli/test_wrap_helpers.py b/tests/test_cli/test_wrap_helpers.py index cc3b9b12f..fbb56460d 100644 --- a/tests/test_cli/test_wrap_helpers.py +++ b/tests/test_cli/test_wrap_helpers.py @@ -405,3 +405,114 @@ def test_run_proxy_only_watcher_calls_cleanup_on_finally( inv = runner.invoke(_cmd) assert inv.exit_code == 1 assert cleanup_calls["n"] >= 1, "cleanup must run via the finally block" + + +# --------------------------------------------------------------------------- +# _project_name_from_cwd / _apply_project_header_env — per-project savings +# header injection for `headroom wrap claude` (issue: per-project savings). +# --------------------------------------------------------------------------- + + +class TestApplyProjectHeaderEnv: + """X-Headroom-Project injection into ANTHROPIC_CUSTOM_HEADERS.""" + + def test_sets_header_from_cwd_basename( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + project_dir = tmp_path / "my-project" + project_dir.mkdir() + monkeypatch.chdir(project_dir) + + env: dict[str, str] = {} + wrap_mod._apply_project_header_env(env) + + assert env["ANTHROPIC_CUSTOM_HEADERS"] == "X-Headroom-Project: my-project" + + def test_appends_to_existing_custom_headers( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + project_dir = tmp_path / "proj" + project_dir.mkdir() + monkeypatch.chdir(project_dir) + + env = {"ANTHROPIC_CUSTOM_HEADERS": "X-Custom-Trace: abc123"} + wrap_mod._apply_project_header_env(env) + + # User header preserved verbatim, ours appended on a new line. + assert env["ANTHROPIC_CUSTOM_HEADERS"] == ( + "X-Custom-Trace: abc123\nX-Headroom-Project: proj" + ) + + @pytest.mark.parametrize( + "user_value", + [ + "X-Headroom-Project: their-name", + "x-headroom-project: their-name", + "X-HEADROOM-PROJECT: their-name", + "X-Other: 1\nx-Headroom-Project: their-name", + ], + ) + def test_existing_project_header_wins_case_insensitive( + self, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + user_value: str, + ) -> None: + project_dir = tmp_path / "proj" + project_dir.mkdir() + monkeypatch.chdir(project_dir) + + env = {"ANTHROPIC_CUSTOM_HEADERS": user_value} + wrap_mod._apply_project_header_env(env) + + # Untouched: no duplicate header, user override wins. + assert env["ANTHROPIC_CUSTOM_HEADERS"] == user_value + + @pytest.mark.parametrize( + "user_value", + [ + "X-Headroom-Project-Id: other", + "X-Trace: mentions x-headroom-project in the value", + ], + ) + def test_similar_header_names_do_not_suppress_injection( + self, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + user_value: str, + ) -> None: + project_dir = tmp_path / "proj" + project_dir.mkdir() + monkeypatch.chdir(project_dir) + + env = {"ANTHROPIC_CUSTOM_HEADERS": user_value} + wrap_mod._apply_project_header_env(env) + + # Only an exact header-name match counts as a user override. + assert env["ANTHROPIC_CUSTOM_HEADERS"] == (f"{user_value}\nX-Headroom-Project: proj") + + def test_empty_cwd_name_sets_nothing(self, monkeypatch: pytest.MonkeyPatch) -> None: + """A degenerate cwd (e.g. filesystem root → empty basename) is a no-op.""" + monkeypatch.setattr(wrap_mod.Path, "cwd", classmethod(lambda cls: Path("/"))) + + env: dict[str, str] = {} + wrap_mod._apply_project_header_env(env) + + assert "ANTHROPIC_CUSTOM_HEADERS" not in env + + def test_whitespace_only_cwd_name_sets_nothing(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(wrap_mod.Path, "cwd", classmethod(lambda cls: Path("/tmp/ "))) + + env: dict[str, str] = {} + wrap_mod._apply_project_header_env(env) + + assert "ANTHROPIC_CUSTOM_HEADERS" not in env + + def test_project_name_from_cwd_returns_basename( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + project_dir = tmp_path / "vibe-headroom" + project_dir.mkdir() + monkeypatch.chdir(project_dir) + + assert wrap_mod._project_name_from_cwd() == "vibe-headroom" diff --git a/tests/test_provider_aider.py b/tests/test_provider_aider.py index 542edab62..08a0d1d22 100644 --- a/tests/test_provider_aider.py +++ b/tests/test_provider_aider.py @@ -31,3 +31,21 @@ def test_aider_build_install_env_returns_only_persistent_proxy_variables() -> No "OPENAI_API_BASE": "http://127.0.0.1:8787/v1", "ANTHROPIC_BASE_URL": "http://127.0.0.1:8787", } + + +def test_aider_build_launch_env_applies_project_path_prefix() -> None: + env, lines = build_launch_env(port=9999, environ={}, project="my repo") + + assert env["OPENAI_API_BASE"] == "http://127.0.0.1:9999/p/my%20repo/v1" + assert env["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:9999/p/my%20repo" + assert lines == [ + "OPENAI_API_BASE=http://127.0.0.1:9999/p/my%20repo/v1", + "ANTHROPIC_BASE_URL=http://127.0.0.1:9999/p/my%20repo", + ] + + +def test_aider_build_launch_env_ignores_unusable_project() -> None: + env, _lines = build_launch_env(port=9999, environ={}, project=" ") + + assert env["OPENAI_API_BASE"] == "http://127.0.0.1:9999/v1" + assert env["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:9999" diff --git a/tests/test_provider_copilot_wrap.py b/tests/test_provider_copilot_wrap.py index 720386070..3defacfa4 100644 --- a/tests/test_provider_copilot_wrap.py +++ b/tests/test_provider_copilot_wrap.py @@ -164,3 +164,32 @@ def test_model_configured_detects_env_and_cli_variants() -> None: assert model_configured(("--model", "gpt-4o"), {}) is True assert model_configured(("--model=gpt-4o",), {}) is True assert model_configured(("--other", "value"), {}) is False + + +def test_build_launch_env_applies_project_path_prefix() -> None: + anthropic_env, _ = build_launch_env( + port=8787, + provider_type="anthropic", + wire_api=None, + environ={"ANTHROPIC_API_KEY": "sk-ant-test"}, + project="api server", + ) + openai_env, openai_lines = build_launch_env( + port=8787, + provider_type="openai", + wire_api=None, + environ={"OPENAI_API_KEY": "sk-proj-test"}, + project="api server", + ) + + assert anthropic_env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/p/api%20server" + assert openai_env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/p/api%20server/v1" + assert "COPILOT_PROVIDER_BASE_URL=http://127.0.0.1:8787/p/api%20server/v1" in openai_lines + + plain_env, _ = build_launch_env( + port=8787, + provider_type="openai", + wire_api=None, + environ={"OPENAI_API_KEY": "sk-proj-test"}, + ) + assert plain_env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1" diff --git a/tests/test_provider_cursor.py b/tests/test_provider_cursor.py index b16671524..3652e21dd 100644 --- a/tests/test_provider_cursor.py +++ b/tests/test_provider_cursor.py @@ -28,3 +28,21 @@ def test_cursor_build_install_env_returns_both_proxy_urls() -> None: "OPENAI_BASE_URL": "http://127.0.0.1:7654/v1", "ANTHROPIC_BASE_URL": "http://127.0.0.1:7654", } + + +def test_cursor_proxy_targets_apply_project_path_prefix() -> None: + targets = build_proxy_targets(9999, project="frontend") + + assert targets.openai_base_url == "http://127.0.0.1:9999/p/frontend/v1" + assert targets.anthropic_base_url == "http://127.0.0.1:9999/p/frontend" + + +def test_cursor_setup_lines_mention_project_attribution() -> None: + lines = render_setup_lines(8787, project="frontend") + joined = "\n".join(lines) + + assert "http://127.0.0.1:8787/p/frontend/v1" in joined + assert "attributed to project 'frontend'" in joined + + plain = "\n".join(render_setup_lines(8787)) + assert "attributed" not in plain diff --git a/tests/test_proxy_project_savings.py b/tests/test_proxy_project_savings.py new file mode 100644 index 000000000..086819f59 --- /dev/null +++ b/tests/test_proxy_project_savings.py @@ -0,0 +1,341 @@ +"""Tests for per-project savings attribution (X-Headroom-Project).""" + +import asyncio +import json + +import pytest + +pytest.importorskip("fastapi") +from fastapi.testclient import TestClient # noqa: E402 + +from headroom.proxy.outcome import RequestOutcome, emit_request_outcome # noqa: E402 +from headroom.proxy.project_context import ( # noqa: E402 + classify_project, + get_current_project, + set_current_project, + split_project_path, + with_project_prefix, +) +from headroom.proxy.savings_tracker import ( # noqa: E402 + DEFAULT_MAX_PROJECTS, + SavingsTracker, + sanitize_project_name, +) +from headroom.proxy.server import ProxyConfig, create_app # noqa: E402 + +# --------------------------------------------------------------------------- +# sanitize_project_name / classify_project +# --------------------------------------------------------------------------- + + +def test_sanitize_project_name_normalizes_and_caps(): + assert sanitize_project_name(" api-server ") == "api-server" + assert sanitize_project_name("a" * 300) == "a" * 128 + assert sanitize_project_name("x\x00\x1by") == "xy" + assert sanitize_project_name("") is None + assert sanitize_project_name(" ") is None + assert sanitize_project_name(None) is None + assert sanitize_project_name(42) is None + + +def test_classify_project_reads_header_case_insensitively(): + assert classify_project({"x-headroom-project": "frontend"}) == "frontend" + assert classify_project({"X-Headroom-Project": " frontend "}) == "frontend" + assert classify_project({"user-agent": "claude-code/1.0"}) is None + assert classify_project(object()) is None + + +def test_split_project_path_extracts_and_strips(): + assert split_project_path("/p/frontend/v1/messages") == ("frontend", "/v1/messages") + assert split_project_path("/p/my%20repo/v1/chat/completions") == ( + "my repo", + "/v1/chat/completions", + ) + assert split_project_path("/p/frontend") == ("frontend", "/") + # No prefix / unusable name: path passes through untouched. + assert split_project_path("/v1/messages") == (None, "/v1/messages") + assert split_project_path("/p//v1/messages") == (None, "/p//v1/messages") + assert split_project_path("/p/%20%20/v1") == (None, "/p/%20%20/v1") + + +def test_with_project_prefix_round_trips_through_split(): + url = with_project_prefix("http://127.0.0.1:8787/v1", "my repo") + assert url == "http://127.0.0.1:8787/p/my%20repo/v1" + path = url.removeprefix("http://127.0.0.1:8787") + assert split_project_path(path) == ("my repo", "/v1") + + # Bare host (anthropic-style base) and unusable names. + assert with_project_prefix("http://127.0.0.1:8787", "api") == "http://127.0.0.1:8787/p/api" + assert with_project_prefix("http://127.0.0.1:8787/v1", " ") == "http://127.0.0.1:8787/v1" + assert with_project_prefix("http://127.0.0.1:8787/v1", None) == "http://127.0.0.1:8787/v1" + + +def test_project_contextvar_roundtrip(): + set_current_project(" demo ") + assert get_current_project() == "demo" + set_current_project(None) + assert get_current_project() is None + + +# --------------------------------------------------------------------------- +# SavingsTracker per-project aggregation +# --------------------------------------------------------------------------- + + +def test_tracker_accumulates_per_project_and_persists(tmp_path): + path = tmp_path / "savings.json" + tracker = SavingsTracker(path=str(path)) + + tracker.record_request(model="gpt-4o", input_tokens=1000, tokens_saved=400, project="api") + tracker.record_request(model="gpt-4o", input_tokens=500, tokens_saved=100, project="api") + tracker.record_request(model="gpt-4o", input_tokens=200, tokens_saved=50, project="web") + tracker.record_request(model="gpt-4o", input_tokens=99, tokens_saved=9) # unattributed + + projects = tracker.stats_preview()["projects"] + assert list(projects) == ["api", "web"] # sorted by tokens saved desc + assert projects["api"]["requests"] == 2 + assert projects["api"]["tokens_saved"] == 500 + assert projects["api"]["total_input_tokens"] == 1500 + assert projects["api"]["savings_percent"] == pytest.approx(25.0) + assert projects["web"]["requests"] == 1 + assert projects["api"]["last_activity_at"] is not None + + # Unattributed traffic still lands in the lifetime totals. + assert tracker.stats_preview()["lifetime"]["requests"] == 4 + + # Survives a restart via the persisted JSON state. + reloaded = SavingsTracker(path=str(path)) + assert reloaded.stats_preview()["projects"]["api"]["tokens_saved"] == 500 + + +def test_tracker_migrates_v2_state_without_projects(tmp_path): + path = tmp_path / "savings.json" + path.write_text( + json.dumps( + { + "schema_version": 2, + "lifetime": { + "requests": 3, + "tokens_saved": 77, + "compression_savings_usd": 0.1, + "total_input_tokens": 500, + "total_input_cost_usd": 0.2, + }, + "display_session": None, + "history": [], + } + ) + ) + tracker = SavingsTracker(path=str(path)) + preview = tracker.stats_preview() + assert preview["projects"] == {} + assert preview["lifetime"]["tokens_saved"] == 77 + + +def test_tracker_caps_project_cardinality(tmp_path): + tracker = SavingsTracker(path=str(tmp_path / "savings.json")) + for i in range(DEFAULT_MAX_PROJECTS + 5): + tracker.record_request( + model="gpt-4o", + input_tokens=10, + tokens_saved=i + 1, + project=f"proj-{i:03d}", + ) + projects = tracker.stats_preview()["projects"] + assert len(projects) == DEFAULT_MAX_PROJECTS + # The smallest buckets were evicted; the biggest savers survive. + assert "proj-000" not in projects + assert f"proj-{DEFAULT_MAX_PROJECTS + 4:03d}" in projects + + +def test_tracker_sanitizes_persisted_project_state(tmp_path): + path = tmp_path / "savings.json" + path.write_text( + json.dumps( + { + "schema_version": 3, + "lifetime": {}, + "display_session": None, + "history": [], + "projects": { + "ok": {"requests": "2", "tokens_saved": 10}, + "": {"requests": 1}, + "bad-entry": "not-a-dict", + }, + } + ) + ) + projects = SavingsTracker(path=str(path)).stats_preview()["projects"] + assert set(projects) == {"ok"} + assert projects["ok"]["requests"] == 2 + assert projects["ok"]["tokens_saved"] == 10 + assert projects["ok"]["compression_savings_usd"] == 0.0 + + +def test_tracker_caps_persisted_projects_on_load(tmp_path): + path = tmp_path / "savings.json" + oversized = { + f"proj-{i:03d}": {"requests": 1, "tokens_saved": i} + for i in range(DEFAULT_MAX_PROJECTS + 10) + } + path.write_text( + json.dumps( + { + "schema_version": 3, + "lifetime": {}, + "display_session": None, + "history": [], + "projects": oversized, + } + ) + ) + projects = SavingsTracker(path=str(path)).stats_preview()["projects"] + assert len(projects) == DEFAULT_MAX_PROJECTS + # Lowest tokens_saved entries are dropped, highest kept. + assert "proj-000" not in projects + assert f"proj-{DEFAULT_MAX_PROJECTS + 9:03d}" in projects + + +# --------------------------------------------------------------------------- +# End-to-end: outcome funnel -> tracker -> /stats payload +# --------------------------------------------------------------------------- + + +def _emit_outcome(proxy, *, project_field=None): + outcome = RequestOutcome( + request_id="req-1", + provider="openai", + model="gpt-4o", + original_tokens=1000, + optimized_tokens=600, + output_tokens=20, + tokens_saved=400, + attempted_input_tokens=1000, + project=project_field, + ) + asyncio.run(emit_request_outcome(proxy, outcome)) + + +def test_funnel_attributes_savings_from_context_and_stats_exposes_them(tmp_path, monkeypatch): + monkeypatch.setenv("HEADROOM_SAVINGS_PATH", str(tmp_path / "savings.json")) + config = ProxyConfig(cache_enabled=False, rate_limit_enabled=False, log_requests=False) + + with TestClient(create_app(config)) as client: + proxy = client.app.state.proxy + + set_current_project("ctx-project") + try: + _emit_outcome(proxy) + finally: + set_current_project(None) + + # Explicit outcome.project wins over the bound context. + _emit_outcome(proxy, project_field="field-project") + + stats = client.get("/stats").json() + per_project = stats["savings"]["per_project"] + assert per_project["ctx-project"]["tokens_saved"] == 400 + assert per_project["field-project"]["tokens_saved"] == 400 + assert stats["persistent_savings"]["projects"] == per_project + assert stats["persistent_savings"]["projects_limit"] == DEFAULT_MAX_PROJECTS + + history = client.get("/stats-history").json() + assert history["schema_version"] == 3 + assert history["projects"]["ctx-project"]["requests"] == 1 + + +# --------------------------------------------------------------------------- +# Regression: pre-feature behavior must be unchanged +# --------------------------------------------------------------------------- + + +def test_record_request_without_project_matches_legacy_totals(tmp_path): + """No-header traffic produces exactly the pre-v3 aggregates.""" + path = tmp_path / "savings.json" + tracker = SavingsTracker(path=str(path)) + tracker.record_request(model="gpt-4o", input_tokens=100, tokens_saved=40) + tracker.record_request(model="gpt-4o", input_tokens=200, tokens_saved=60) + + preview = tracker.stats_preview() + assert preview["projects"] == {} + assert preview["lifetime"]["requests"] == 2 + assert preview["lifetime"]["tokens_saved"] == 100 + assert preview["display_session"]["tokens_saved"] == 100 + + persisted = json.loads(path.read_text()) + # Every legacy top-level key survives alongside the new projects map. + assert set(persisted) >= {"schema_version", "lifetime", "display_session", "history"} + assert persisted["projects"] == {} + + +def test_stats_payload_keeps_legacy_shape(tmp_path, monkeypatch): + """Dashboard consumers of the old /stats keys must not break.""" + monkeypatch.setenv("HEADROOM_SAVINGS_PATH", str(tmp_path / "savings.json")) + config = ProxyConfig(cache_enabled=False, rate_limit_enabled=False, log_requests=False) + + with TestClient(create_app(config)) as client: + proxy = client.app.state.proxy + _emit_outcome(proxy) # unattributed: no header, no context, no field + + stats = client.get("/stats").json() + assert stats["savings"]["per_project"] == {} + for legacy_key in ("requests", "savings", "persistent_savings", "cost"): + assert legacy_key in stats, f"legacy /stats key {legacy_key!r} disappeared" + assert stats["persistent_savings"]["lifetime"]["requests"] == 1 + + history = client.get("/stats-history").json() + for legacy_key in ("schema_version", "lifetime", "display_session", "retention"): + assert legacy_key in history, f"legacy /stats-history key {legacy_key!r} disappeared" + + +def test_metrics_record_request_works_without_project_kwarg(tmp_path, monkeypatch): + """Existing callers that never pass ``project=`` keep working.""" + monkeypatch.setenv("HEADROOM_SAVINGS_PATH", str(tmp_path / "savings.json")) + config = ProxyConfig(cache_enabled=False, rate_limit_enabled=False, log_requests=False) + + with TestClient(create_app(config)) as client: + proxy = client.app.state.proxy + asyncio.run( + proxy.metrics.record_request( + provider="openai", + model="gpt-4o", + input_tokens=120, + output_tokens=24, + tokens_saved=30, + latency_ms=15.0, + ) + ) + preview = proxy.metrics.savings_tracker.stats_preview() + assert preview["lifetime"]["tokens_saved"] == 30 + assert preview["projects"] == {} + + +def test_middleware_binds_project_header_to_context(tmp_path, monkeypatch): + monkeypatch.setenv("HEADROOM_SAVINGS_PATH", str(tmp_path / "savings.json")) + config = ProxyConfig(cache_enabled=False, rate_limit_enabled=False, log_requests=False) + + captured: list[str | None] = [] + + import headroom.proxy.server as server_module + + def _capture(project: str | None) -> None: + captured.append(project) + set_current_project(project) + + monkeypatch.setattr(server_module, "set_current_project", _capture) + + with TestClient(create_app(config)) as client: + assert client.get("/health", headers={"X-Headroom-Project": " my repo "}).status_code == 200 + assert client.get("/health").status_code == 200 + # /p/ base-URL prefix (aider/copilot/cursor wraps): stripped + # before routing, so the request still reaches /health. + assert client.get("/p/my%20repo/health").status_code == 200 + # An explicit header wins over the path prefix. + assert ( + client.get( + "/p/prefix-project/health", headers={"X-Headroom-Project": "header-project"} + ).status_code + == 200 + ) + + assert captured == ["my repo", None, "my repo", "header-project"] diff --git a/tests/test_proxy_savings_history.py b/tests/test_proxy_savings_history.py index ba9391f81..421505d8c 100644 --- a/tests/test_proxy_savings_history.py +++ b/tests/test_proxy_savings_history.py @@ -111,7 +111,7 @@ def test_savings_tracker_sanitizes_legacy_state_and_applies_retention(tmp_path): ) snapshot = tracker.snapshot() - assert snapshot["schema_version"] == 2 + assert snapshot["schema_version"] == 3 assert snapshot["lifetime"] == { "requests": 0, "tokens_saved": 30, @@ -686,7 +686,7 @@ def test_stats_history_persists_across_restarts_and_stats_stays_compatible(tmp_p history = client.get("/stats-history") assert history.status_code == 200 history_data = history.json() - assert history_data["schema_version"] == 2 + assert history_data["schema_version"] == 3 assert history_data["storage_path"] == str(savings_path) assert history_data["lifetime"]["tokens_saved"] == 40 assert history_data["lifetime"]["total_input_tokens"] == 120