mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(proxy): per-project savings breakdown on the dashboard (claude, codex, aider, copilot, cursor) (#803)
## 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: <basename(cwd)>` 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/<url-encoded-name>/...`; 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 <ashley.rhodes@king.com>
This commit is contained in:
Focused Instability
GitHub
parent
6ea6e31f09
commit
914a60a2b0
22 changed files with 1051 additions and 33 deletions
|
|
@ -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/<name>` 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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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/<name> 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/<name> 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(
|
||||
|
|
|
|||
|
|
@ -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/<name> 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()
|
||||
|
|
|
|||
|
|
@ -886,6 +886,44 @@
|
|||
</div>
|
||||
</template>
|
||||
|
||||
<!-- Per-Project Savings Breakdown -->
|
||||
<template x-if="Object.keys(stats.savings?.per_project || {}).length > 0">
|
||||
<div class="bg-surface rounded-lg border border-border overflow-hidden mb-6">
|
||||
<div class="px-4 py-3 border-b border-border flex justify-between items-center">
|
||||
<span class="text-sm font-medium text-gray-300">Per-Project Savings</span>
|
||||
<span class="text-xs text-gray-500">Lifetime totals — all time</span>
|
||||
</div>
|
||||
<div class="overflow-x-auto">
|
||||
<table class="w-full text-sm">
|
||||
<thead>
|
||||
<tr class="text-left text-xs text-gray-500 uppercase tracking-wide">
|
||||
<th class="px-4 py-3 font-medium">Project</th>
|
||||
<th class="px-4 py-3 font-medium text-right">Requests</th>
|
||||
<th class="px-4 py-3 font-medium text-right">Tokens Saved</th>
|
||||
<th class="px-4 py-3 font-medium text-right">Saved $</th>
|
||||
<th class="px-4 py-3 font-medium text-right">Savings %</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody class="divide-y divide-border">
|
||||
<template x-for="[project, info] in Object.entries(stats.savings?.per_project || {})" :key="project">
|
||||
<tr class="hover:bg-border/30 transition-colors">
|
||||
<td class="px-4 py-3">
|
||||
<span class="px-2 py-0.5 bg-border rounded text-xs" x-text="project"></span>
|
||||
</td>
|
||||
<td class="px-4 py-3 text-right font-mono" x-text="info.requests"></td>
|
||||
<td class="px-4 py-3 text-right font-mono text-accent" x-text="formatNumber(info.tokens_saved)"></td>
|
||||
<td class="px-4 py-3 text-right font-mono text-emerald-400" x-text="'$' + formatCurrency(info.compression_savings_usd || 0)"></td>
|
||||
<td class="px-4 py-3 text-right">
|
||||
<span class="text-accent font-mono" x-text="(info.savings_percent || 0).toFixed(1) + '%'"></span>
|
||||
</td>
|
||||
</tr>
|
||||
</template>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
</div>
|
||||
</template>
|
||||
|
||||
<!-- Recent Requests Table (with expandable rows) -->
|
||||
<div class="bg-surface rounded-lg border border-border overflow-hidden">
|
||||
<div class="px-4 py-3 border-b border-border flex justify-between items-center">
|
||||
|
|
|
|||
|
|
@ -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/<name>``
|
||||
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, [
|
||||
|
|
|
|||
|
|
@ -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/<name>``
|
||||
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, [
|
||||
|
|
|
|||
|
|
@ -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/<name>``
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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 ""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
106
headroom/proxy/project_context.py
Normal file
106
headroom/proxy/project_context.py
Normal file
|
|
@ -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/<name>/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/<name>`` 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/<name>`` 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",
|
||||
]
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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/<name> 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,
|
||||
|
|
|
|||
|
|
@ -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/<name> 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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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/<name> 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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
341
tests/test_proxy_project_savings.py
Normal file
341
tests/test_proxy_project_savings.py
Normal file
|
|
@ -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/<name> 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"]
|
||||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue