From 6d3f39f213f4eb2d1c6c814b34e1bf6fe2a5c959 Mon Sep 17 00:00:00 2001 From: Michael Sam Date: Fri, 12 Jun 2026 22:12:22 +0300 Subject: [PATCH] feat: add dashboard agent usage stats (#814) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Add a clear dashboard view for per-agent token usage so end users can see Cursor, Claude, Codex, and other detected clients with before/after token counts, tokens saved, and savings percentages. The stats API now exposes a stable `agent_usage` object that the dashboard renders near the top of the session view. Fixes # ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made ### New Files **Tests:** - `tests/test_dashboard_agent_usage.py` — Covers agent classification, exact per-request aggregation, and aggregate fallback behavior. ### Modified Files - `headroom/proxy/server.py` — Adds per-agent usage aggregation to `/stats` with before tokens, after tokens, output tokens, saved tokens, savings percentage, source, providers, and models. - `headroom/dashboard/templates/dashboard.html` — Adds a prominent Agent Usage panel with totals, coverage status, per-agent token-flow bars, request counts, before/after tokens, saved tokens, and share of savings. ## Testing - [x] Unit tests pass: `.venv312/bin/pytest tests/test_dashboard_agent_usage.py` - [x] Linting passes: `.venv312/bin/ruff check headroom/proxy/server.py tests/test_dashboard_agent_usage.py` - [x] Diff whitespace check passes: `git diff --check origin/main...HEAD` - [x] Dashboard smoke render: local proxy on `127.0.0.1:8790`, captured Chrome headless screenshot of `/dashboard` - [x] New tests added for new functionality ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing relevant unit tests pass locally with my changes - [ ] I have made corresponding changes to the documentation - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes The agent usage panel uses exact request-log data when available. If detailed request logs are empty, it falls back to aggregate provider/model request counts and labels the coverage as aggregate fallback so users are not misled. --- .github/workflows/ci.yml | 11 + codecov.yml | 15 +- headroom/cli/proxy.py | 12 +- headroom/cli/wrap.py | 81 ++++- headroom/compress.py | 9 +- headroom/dashboard/templates/dashboard.html | 154 +++++++++ headroom/perf/analyzer.py | 7 +- headroom/proxy/auth_mode.py | 4 +- headroom/proxy/handlers/anthropic.py | 12 +- headroom/proxy/handlers/openai.py | 36 +- headroom/proxy/server.py | 347 +++++++++++++++++++- tests/test_agent_savings.py | 105 ++++++ tests/test_auth_mode.py | 7 + tests/test_cli/test_wrap_codex.py | 61 ++++ tests/test_cli_perf_format.py | 48 +++ tests/test_cli_proxy_env.py | 18 + tests/test_dashboard_agent_usage.py | 276 ++++++++++++++++ tests/test_openai_codex_routing.py | 65 ++++ tests/test_proxy_stats_recent_requests.py | 162 +++++++++ tests/test_release_workflows.py | 45 ++- 20 files changed, 1424 insertions(+), 51 deletions(-) create mode 100644 tests/test_dashboard_agent_usage.py create mode 100644 tests/test_proxy_stats_recent_requests.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1ce2b8f67..23e80cec1 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -195,8 +195,19 @@ jobs: run: | pytest tests scripts/tests \ --splits 4 --group ${{ matrix.shard }} \ + --cov=headroom --cov-branch \ + --cov-report=xml:coverage-${{ matrix.shard }}.xml \ + --cov-report= \ --tb=short -q + - name: Upload coverage shard ${{ matrix.shard }} to Codecov + uses: codecov/codecov-action@v5 + with: + files: coverage-${{ matrix.shard }}.xml + flags: python + name: python-shard-${{ matrix.shard }} + fail_ci_if_error: true + test-extras: needs: [changes, build-wheel] if: needs.changes.outputs.code == 'true' diff --git a/codecov.yml b/codecov.yml index 6776b3f84..937781b6a 100644 --- a/codecov.yml +++ b/codecov.yml @@ -10,10 +10,11 @@ coverage: default: target: auto -ignore: - - "tests/**" - - "scripts/tests/**" - - ".github/**" - - ".claude-plugin/**" - - "plugins/headroom-agent-hooks/.claude-plugin/**" - - "plugins/headroom-agent-hooks/.github/**" +ignore: + - "tests/**" + - "scripts/tests/**" + - ".github/**" + - ".claude-plugin/**" + - "headroom/dashboard/templates/**" + - "plugins/headroom-agent-hooks/.claude-plugin/**" + - "plugins/headroom-agent-hooks/.github/**" diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index 0cf5e1ed4..1c126997e 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -783,12 +783,6 @@ def proxy( optimize=not no_optimize, cache_enabled=not no_cache, rate_limit_enabled=not no_rate_limit, - # CCR opt-outs for compression-only deployments (streaming / non-MCP - # clients that can't resolve the injected retrieve tool). Defaults keep - # CCR fully on; each flag flips one dataclass default to False. - ccr_inject_tool=not no_ccr_inject_tool, - ccr_inject_marker=not no_ccr_marker, - ccr_proactive_expansion=not no_ccr_proactive_expansion, compress_user_messages=_get_env_bool("HEADROOM_COMPRESS_USER_MESSAGES", False), min_tokens_to_crush=_get_env_int_optional("HEADROOM_MIN_TOKENS") or 500, max_items_after_crush=_get_env_int_optional("HEADROOM_MAX_ITEMS") or 50, @@ -799,6 +793,12 @@ def proxy( protect_recent=_get_env_int_optional("HEADROOM_PROTECT_RECENT"), protect_analysis_context=_get_env_bool_optional("HEADROOM_PROTECT_ANALYSIS_CONTEXT"), accuracy_guard=os.environ.get("HEADROOM_ACCURACY_GUARD") or None, + # CCR opt-outs for compression-only deployments (streaming / non-MCP + # clients that can't resolve the injected retrieve tool). Defaults keep + # CCR fully on; each flag flips one dataclass default to False. + ccr_inject_tool=not no_ccr_inject_tool, + ccr_inject_marker=not no_ccr_marker, + ccr_proactive_expansion=not no_ccr_proactive_expansion, # Flatten repeat-flag tuple AND any comma-separated values inside it. # `--proxy-extension a,b --proxy-extension c` and `HEADROOM_PROXY_EXTENSIONS=a,b,c` # both yield ["a", "b", "c"]. None when nothing was supplied. diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index 050a6f642..9e21589b0 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -38,6 +38,9 @@ if sys.platform == "win32" and hasattr(sys.stdout, "buffer"): import click from headroom._version import __version__ as _HEADROOM_VERSION +from headroom.agent_savings import ( + apply_agent_savings_env_defaults, +) from headroom.copilot_auth import ( has_oauth_auth, resolve_client_bearer_token, @@ -95,6 +98,7 @@ _CONTEXT_TOOL_ENV = "HEADROOM_CONTEXT_TOOL" _CONTEXT_TOOL_RTK = "rtk" _CONTEXT_TOOL_LEAN_CTX = "lean-ctx" _VALID_CONTEXT_TOOLS = {_CONTEXT_TOOL_RTK, _CONTEXT_TOOL_LEAN_CTX} +_AGENT_SAVINGS_TARGET_AGENTS = {"claude", "codex", "cursor"} _WRAP_PROXY_TIMEOUT_ENV = "HEADROOM_WRAP_PROXY_TIMEOUT" _WRAP_PROXY_TIMEOUT_DEFAULT_SECONDS = 45 _WRAP_PROXY_TIMEOUT_ML_DEFAULT_SECONDS = 90 @@ -360,6 +364,8 @@ def _start_proxy( # Ensure proxy subprocess uses UTF-8 (Windows defaults to cp1252) proxy_env = os.environ.copy() proxy_env["PYTHONIOENCODING"] = "utf-8" + if agent_type in {"claude", "codex", "cursor"}: + apply_agent_savings_env_defaults(proxy_env) # Tell the proxy which agent is being wrapped (for traffic learning output) if agent_type != "unknown": @@ -367,8 +373,6 @@ def _start_proxy( proxy_env.setdefault("HEADROOM_STACK", f"wrap_{agent_type}") savings_profile = _wrap_agent_savings_profile(agent_type) if savings_profile is not None: - from headroom.agent_savings import apply_agent_savings_env_defaults - apply_agent_savings_env_defaults(proxy_env, savings_profile) if openai_api_url: proxy_env["OPENAI_TARGET_API_URL"] = openai_api_url @@ -1553,6 +1557,77 @@ def _proxy_health_config(payload: dict[str, Any] | None) -> dict[str, Any] | Non return config if isinstance(config, dict) else None +def _env_bool_value(value: str) -> bool: + return value.strip().lower() in {"1", "true", "yes", "on"} + + +def _agent_savings_config_mismatches( + running_config: dict[str, Any], + agent_type: str, +) -> list[str]: + """Return restart reasons when a running proxy lacks target agent savings.""" + + if agent_type not in _AGENT_SAVINGS_TARGET_AGENTS: + return [] + + desired_env = os.environ.copy() + apply_agent_savings_env_defaults(desired_env) + checks: tuple[tuple[str, str, str, str], ...] = ( + ("HEADROOM_SAVINGS_PROFILE", "savings_profile", "savings-profile", "str"), + ("HEADROOM_TARGET_RATIO", "target_ratio", "target-ratio", "float"), + ( + "HEADROOM_COMPRESS_USER_MESSAGES", + "compress_user_messages", + "compress-user-messages", + "bool", + ), + ( + "HEADROOM_COMPRESS_SYSTEM_MESSAGES", + "compress_system_messages", + "compress-system-messages", + "bool", + ), + ("HEADROOM_PROTECT_RECENT", "protect_recent", "protect-recent", "int"), + ( + "HEADROOM_PROTECT_ANALYSIS_CONTEXT", + "protect_analysis_context", + "protect-analysis-context", + "bool", + ), + ("HEADROOM_MIN_TOKENS", "min_tokens_to_crush", "min-tokens", "int"), + ("HEADROOM_MAX_ITEMS", "max_items_after_crush", "max-items", "int"), + ( + "HEADROOM_SMART_CRUSHER_COMPACTION", + "smart_crusher_with_compaction", + "smart-crusher-compaction", + "bool", + ), + ("HEADROOM_ACCURACY_GUARD", "accuracy_guard", "accuracy-guard", "str"), + ) + + mismatches: list[str] = [] + for env_key, config_key, label, value_type in checks: + expected = desired_env.get(env_key) + if expected is None: + continue + actual = running_config.get(config_key) + try: + if value_type == "float": + matches = actual is not None and abs(float(actual) - float(expected)) < 1e-9 + elif value_type == "int": + matches = actual is not None and int(actual) == int(expected) + elif value_type == "bool": + matches = actual is not None and bool(actual) is _env_bool_value(expected) + else: + matches = str(actual or "").strip().lower() == expected.strip().lower() + except (TypeError, ValueError): + matches = False + if not matches: + mismatches.append(label) + + return mismatches + + def _proxy_active_session_count(payload: dict[str, Any] | None) -> int: """Return active session count from /health runtime metadata.""" if payload is None: @@ -4310,6 +4385,6 @@ def unwrap_codex(port: int, no_stop_proxy: bool) -> None: click.echo() click.echo("✓ Codex is no longer routed through the Headroom proxy.") - if not no_stop_proxy: + if not no_stop_proxy and status != "noop": _echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(port), port) click.echo() diff --git a/headroom/compress.py b/headroom/compress.py index be258958b..47ab8898e 100644 --- a/headroom/compress.py +++ b/headroom/compress.py @@ -58,9 +58,10 @@ from __future__ import annotations import logging import threading -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any +from .agent_savings import apply_agent_savings_profile from .observability import get_otel_metrics from .pipeline import PipelineExtensionManager, PipelineStage, summarize_routing_markers from .utils import extract_user_query as _extract_user_query @@ -133,6 +134,9 @@ class CompressConfig: Set to 'disabled' to skip ML compression entirely (only SmartCrusher + CacheAligner will run).""" + savings_profile: str | None = None + """Named high-savings profile, e.g. 'agent-90' for Codex/Claude/Cursor.""" + @dataclass class CompressResult: @@ -204,6 +208,9 @@ def compress( for key, value in kwargs.items(): if key in config_fields: setattr(cfg, key, value) + if cfg.savings_profile: + cfg = replace(cfg) + apply_agent_savings_profile(cfg, cfg.savings_profile) pipeline = _get_pipeline() pipeline_extensions = PipelineExtensionManager(hooks=hooks, discover=False) diff --git a/headroom/dashboard/templates/dashboard.html b/headroom/dashboard/templates/dashboard.html index f3f74578d..389fa7c38 100644 --- a/headroom/dashboard/templates/dashboard.html +++ b/headroom/dashboard/templates/dashboard.html @@ -133,6 +133,18 @@ Anon Telemetry +
Status @@ -229,6 +241,106 @@
+ +
+
+
+
Agent Usage
+
Before and after token usage by detected client
+
+
+ + +
+
+ +
+
+
+
Before
+
+
+
+
After
+
+
+
+
Saved
+
+
+
+
Savings
+
+
+
+ + + + +
+
+