Files
wehub-resource-sync 0ef5fcb1c5
Security / Dependency audit (pip-audit) (push) Has been cancelled
Security / CodeQL (javascript-typescript) (push) Has been cancelled
Security / CodeQL (python) (push) Has been cancelled
Security / Secret scan (gitleaks) (push) Has been cancelled
rust / test (ubuntu) (push) Has been cancelled
rust / simulator e2e (macos-latest) (push) Has been cancelled
rust / simulator e2e (ubuntu-latest) (push) Has been cancelled
rust / simulator e2e (windows-latest) (push) Has been cancelled
rust / wheels (aarch64-apple-darwin) (push) Has been cancelled
rust / wheels (x86_64-unknown-linux-gnu) (push) Has been cancelled
rust / wheels (x86_64-apple-darwin) (push) Has been cancelled
rust / audit (push) Has been cancelled
rust / parity (nightly, allowed to fail during Phase 0) (push) Has been cancelled
CI / commitlint (push) Has been skipped
Dev Containers / validate (.devcontainer/devcontainer.json, default) (push) Failing after 0s
Dev Containers / validate (.devcontainer/memory-stack/devcontainer.json, memory-stack) (push) Failing after 0s
Dev Containers / validate-worktree (push) Failing after 0s
CI / changes (push) Failing after 4s
Deploy Documentation / validate (push) Has been skipped
Deploy Documentation / deploy (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, claude) (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, codex) (push) Failing after 1s
Install Native E2E / install-native (ubuntu-latest) (push) Failing after 1s
OpenCode Plugin / typecheck + build + test (push) Failing after 1s
Init Native E2E / init-native (ubuntu-latest, copilot) (push) Failing after 1s
Release Please / release-please (push) Failing after 1s
Wrap E2E / docker-wrap-e2e (push) Failing after 1s
Wrap Native E2E / wrap-native (ubuntu-latest) (push) Failing after 1s
Init E2E / docker-init-e2e (push) Failing after 4s
Merge Conflicts / merge-conflicts (push) Failing after 4s
CI / lint (push) Has been cancelled
CI / build-wheel (push) Has been cancelled
CI / build-wheel-windows (push) Has been cancelled
CI / prefetch-model (push) Has been cancelled
CI / test-dashboard-ui (push) Has been cancelled
CI / test (1) (push) Has been cancelled
CI / test (2) (push) Has been cancelled
CI / test (3) (push) Has been cancelled
CI / test (4) (push) Has been cancelled
CI / test-extras (push) Has been cancelled
CI / test-agno (push) Has been cancelled
CI / build (push) Has been cancelled
CI / workflow-validation (push) Has been cancelled
CI / docker-native-e2e (push) Has been cancelled
CI / windows-native-wrapper (push) Has been cancelled
CI / macos-native-wrapper (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-manifest (map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim name:code-slim]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-nonroot name:nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim name:slim]) (push) Has been cancelled
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime name:]) (push) Has been cancelled
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code name:code]) (push) Has been cancelled
Docker / promote-latest (push) Has been cancelled
Init Native E2E / init-native (macos-latest, claude) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, codex) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, copilot) (push) Has been cancelled
Install Native E2E / install-native (macos-latest) (push) Has been cancelled
Wrap Native E2E / wrap-native (macos-latest) (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:03:20 +08:00

5075 lines
225 KiB
Python

"""Headroom Proxy Server - Production Ready.
A full-featured LLM proxy with optimization, caching, rate limiting,
and observability.
Features:
- Context optimization (SmartCrusher, CacheAligner — live-zone-only after Phase B)
- Semantic caching (save costs on repeated queries)
- Rate limiting (token bucket)
- Retry with exponential backoff
- Cost tracking and budgets
- Request tagging and metadata
- Provider fallback
- Prometheus metrics
- Full request/response logging
Usage:
python -m headroom.proxy.server --port 8787
# With Claude Code:
ANTHROPIC_BASE_URL=http://localhost:8787 claude
"""
from __future__ import annotations
import argparse
import asyncio
import concurrent.futures
import contextlib
import hmac
import json
import logging
import os
import sys
import threading
import time
from collections.abc import Callable
from dataclasses import fields, is_dataclass, replace
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast
if TYPE_CHECKING:
from ..backends.base import Backend
from ..cache.compression_cache import CompressionCache
from ..memory.tracker import MemoryTracker
from .outcome import RequestOutcome
import httpx
try:
import uvicorn
from fastapi import Depends, FastAPI, HTTPException, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse, JSONResponse, PlainTextResponse
FASTAPI_AVAILABLE = True
except ImportError:
FASTAPI_AVAILABLE = False
# Add parent to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from headroom._version import __version__
from headroom.agent_savings import proxy_pipeline_kwargs
from headroom.cache.compression_feedback import get_compression_feedback
from headroom.cache.compression_store import format_retrieval_miss_detail, get_compression_store
from headroom.ccr import (
CCR_TOOL_NAME,
# Batch processing
CCRResponseHandler,
CCRToolInjector,
ContextTracker,
ContextTrackerConfig,
ResponseHandlerConfig,
parse_tool_call,
)
from headroom.config import (
DEFAULT_EXCLUDE_TOOLS,
CacheAlignerConfig,
ReadLifecycleConfig,
)
from headroom.dashboard import get_dashboard_html
from headroom.observability import (
LangfuseTracingConfig,
OTelMetricsConfig,
configure_langfuse_tracing,
configure_otel_metrics,
get_langfuse_tracing_status,
get_otel_metrics_status,
shutdown_headroom_tracing,
shutdown_otel_metrics,
)
from headroom.offline import apply_offline_env, is_offline
from headroom.pipeline import PipelineExtensionManager, PipelineStage
from headroom.providers.proxy_routes import register_provider_routes
from headroom.providers.registry import (
DEFAULT_ANTHROPIC_API_URL,
DEFAULT_CLOUDCODE_API_URL,
DEFAULT_GEMINI_API_URL,
DEFAULT_OPENAI_API_URL,
DEFAULT_VERTEX_API_URL,
build_proxy_provider_runtime,
create_proxy_backend,
format_backend_status,
resolve_api_targets,
)
from headroom.proxy import runtime_env
from headroom.proxy.audit import is_auditable_path, record_admin_action
from headroom.proxy.auth_mode import should_stamp_codex_client
from headroom.proxy.background_compression import BackgroundCompressor
# =============================================================================
# Extracted modules (re-exported for backward compatibility)
# =============================================================================
from headroom.proxy.cost import (
_CACHE_ECONOMICS, # noqa: F401
CostTracker, # noqa: F401
_summarize_transforms, # noqa: F401
build_prefix_cache_stats, # noqa: F401
build_session_summary, # noqa: F401
merge_cost_stats, # noqa: F401
)
from headroom.proxy.helpers import (
COMPRESSION_TIMEOUT_SECONDS, # noqa: F401
EAGER_PRELOAD_TIMEOUT_SECONDS,
MAX_COMPRESSION_CACHE_SESSIONS, # noqa: F401
MAX_MESSAGE_ARRAY_LENGTH, # noqa: F401
MAX_REQUEST_BODY_SIZE, # noqa: F401
MAX_SSE_BUFFER_SIZE, # noqa: F401
RETRYABLE_OVERLOAD_STATUSES,
_get_context_tool_stats,
_get_image_compressor, # noqa: F401
_get_rtk_stats, # noqa: F401
_read_request_json, # noqa: F401
_setup_file_logging, # noqa: F401
initialize_context_tool_session_baseline,
is_anthropic_auth, # noqa: F401
jitter_delay_ms,
retry_after_ms,
)
from headroom.proxy.loop_callback_failure_policy import is_known_websocket_callback_failure
from headroom.proxy.loopback_guard import is_loopback_host
from headroom.proxy.memory_handler import MemoryConfig, MemoryHandler
# Data models (extracted to headroom/proxy/models.py for maintainability)
from headroom.proxy.models import CacheEntry, ProxyConfig, RateLimitState, RequestLog # noqa: F401
from headroom.proxy.modes import (
PROXY_MODE_CACHE,
PROXY_MODE_TOKEN,
is_token_mode,
normalize_proxy_mode,
)
from headroom.proxy.probe_recorder import probe_recorder_from_env
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
from headroom.proxy.savings_tracker import LITELLM_AVAILABLE
from headroom.proxy.semantic_cache import SemanticCache # noqa: F401
from headroom.proxy.ssl_context import build_httpx_verify
from headroom.proxy.tool_schema_savings_policy import tool_schema_saved_from_tags
from headroom.proxy.warmup import WarmupRegistry
from headroom.proxy.ws_session_registry import WebSocketSessionRegistry
from headroom.subscription.base import get_quota_registry, reset_quota_registry
from headroom.subscription.codex_rate_limits import get_codex_rate_limit_state
from headroom.subscription.copilot_quota import get_copilot_quota_tracker
from headroom.subscription.tracker import (
configure_subscription_tracker,
get_subscription_tracker,
)
from headroom.telemetry import get_telemetry_collector
from headroom.telemetry.beacon import is_telemetry_enabled
from headroom.telemetry.toin import get_toin
from headroom.transforms import (
CacheAligner,
CodeAwareCompressor,
CodeCompressorConfig,
CompressionStrategy,
ContentRouter,
ContentRouterConfig,
TransformPipeline,
is_tree_sitter_available,
)
AnyLLMBackend: Any = None
LiteLLMBackend: Any = None
fcntl: Any = None
try:
import fcntl as _fcntl
fcntl = _fcntl
HAS_FCNTL = True
except ImportError:
HAS_FCNTL = False
_build_prefix_cache_stats = build_prefix_cache_stats
_build_session_summary = build_session_summary
_merge_cost_stats = merge_cost_stats
_AGENT_LABELS: dict[str, str] = {
"claude": "Claude",
"claude-code": "Claude",
"claude_cli": "Claude",
"claude-code-cli": "Claude",
"codex": "Codex",
"codex-cli": "Codex",
"cursor": "Cursor",
"copilot": "GitHub Copilot",
"github-copilot": "GitHub Copilot",
"aider": "Aider",
"zed": "Zed",
"opencode": "OpenCode",
"openclaw": "OpenClaw",
"gemini": "Gemini",
"google": "Gemini",
"vertex:google": "Gemini",
"anthropic": "Claude",
"openai": "OpenAI",
"unknown": "Unidentified",
}
_AGENT_SOURCE_PRIORITY: dict[str, int] = {
"unknown": 0,
"provider": 1,
"model": 2,
"stack": 3,
"client": 4,
}
def _normalize_agent_key(raw: Any) -> str | None:
if raw is None:
return None
value = str(raw).strip().lower()
if not value:
return None
value = value.replace(" ", "-").replace("_", "-")
if value.startswith("wrap-"):
value = value.removeprefix("wrap-")
if value in {"claude-cli", "claude-code", "claude-code-cli"}:
return "claude-code"
if value in {"codex-cli", "codex"}:
return "codex"
if value in {"github-copilot", "copilot"}:
return "copilot"
if value in {"google", "vertex-google", "vertex:google"}:
return "gemini"
return value
def _agent_label(agent_key: str) -> str:
if agent_key in _AGENT_LABELS:
return _AGENT_LABELS[agent_key]
return agent_key.replace("-", " ").replace("_", " ").title()
def _classify_agent_from_log(entry: dict[str, Any]) -> tuple[str, str, str]:
raw_tags = entry.get("tags")
tags = raw_tags if isinstance(raw_tags, dict) else {}
for source, candidate in (
("client", tags.get("client")),
("stack", tags.get("stack") or tags.get("headroom-stack")),
):
key = _normalize_agent_key(candidate)
if key:
return key, _agent_label(key), source
model = str(entry.get("model") or "").lower()
if "codex" in model:
return "codex", _agent_label("codex"), "model"
if "claude" in model:
return "claude-code", _agent_label("claude-code"), "model"
if "gemini" in model:
return "gemini", _agent_label("gemini"), "model"
key = _normalize_agent_key(entry.get("provider"))
if key:
return key, _agent_label(key), "provider"
return "unknown", _agent_label("unknown"), "unknown"
def _build_agent_usage_summary(
logs: list[dict[str, Any]],
*,
requests_by_provider: dict[str, int],
requests_by_model: dict[str, int],
global_before_tokens: int,
global_after_tokens: int,
global_tokens_saved: int,
global_output_tokens: int,
) -> dict[str, Any]:
agents: dict[str, dict[str, Any]] = {}
def _agent_row(agent_key: str, label: str, source: str) -> dict[str, Any]:
row = agents.setdefault(
agent_key,
{
"agent": agent_key,
"label": label,
"source": source,
"requests": 0,
"before_tokens": 0,
"after_tokens": 0,
"output_tokens": 0,
"tokens_saved": 0,
"models": {},
"providers": {},
"has_exact_tokens": False,
},
)
if _AGENT_SOURCE_PRIORITY.get(source, 0) > _AGENT_SOURCE_PRIORITY.get(
str(row.get("source") or "unknown"), 0
):
row["source"] = source
return row
for entry in logs:
agent_key, label, source = _classify_agent_from_log(entry)
row = _agent_row(agent_key, label, source)
before = max(0, int(entry.get("input_tokens_original") or 0))
after = max(0, int(entry.get("input_tokens_optimized") or 0))
saved = max(0, int(entry.get("tokens_saved") or 0))
output = max(0, int(entry.get("output_tokens") or 0))
provider = str(entry.get("provider") or "unknown")
model = str(entry.get("model") or "unknown")
row["requests"] += 1
row["before_tokens"] += before
row["after_tokens"] += after
row["output_tokens"] += output
row["tokens_saved"] += saved
row["providers"][provider] = int(row["providers"].get(provider, 0)) + 1
row["models"][model] = int(row["models"].get(model, 0)) + 1
if before > 0 or after > 0 or saved > 0:
row["has_exact_tokens"] = True
if not agents:
inferred_model_counts: dict[str, int] = {}
for model, count in requests_by_model.items():
model_lower = str(model).lower()
if "codex" in model_lower:
key = "codex"
elif "claude" in model_lower:
key = "claude-code"
elif "gemini" in model_lower:
key = "gemini"
else:
continue
inferred_model_counts[str(model)] = int(count)
provider_request_count = sum(max(0, int(count)) for count in requests_by_provider.values())
inferred_request_count = sum(max(0, count) for count in inferred_model_counts.values())
use_model_fallback = (
inferred_request_count > 0 and inferred_request_count == provider_request_count
)
if not use_model_fallback:
for provider, count in requests_by_provider.items():
key = _normalize_agent_key(provider) or "unknown"
row = _agent_row(key, _agent_label(key), "provider")
row["requests"] += int(count)
row["providers"][provider] = int(row["providers"].get(provider, 0)) + int(count)
for model, count in requests_by_model.items():
model_lower = str(model).lower()
if "codex" in model_lower:
key = "codex"
elif "claude" in model_lower:
key = "claude-code"
elif "gemini" in model_lower:
key = "gemini"
else:
continue
if not use_model_fallback:
continue
row = _agent_row(key, _agent_label(key), "model")
row["requests"] += int(count)
row["models"][str(model)] = int(row["models"].get(str(model), 0)) + int(count)
rows: list[dict[str, Any]] = []
for row in agents.values():
before = int(row["before_tokens"])
saved = int(row["tokens_saved"])
after = int(row["after_tokens"])
if before == 0 and (after > 0 or saved > 0):
before = after + saved
savings_percent = round((saved / before) * 100.0, 2) if before else 0.0
row["before_tokens"] = before
row["savings_percent"] = savings_percent
row["after_percent"] = round((after / before) * 100.0, 2) if before else 0.0
row["share_of_saved_percent"] = (
round((saved / global_tokens_saved) * 100.0, 2) if global_tokens_saved else 0.0
)
row["share_of_requests_percent"] = 0.0
rows.append(row)
total_requests = sum(int(row["requests"]) for row in rows)
for row in rows:
row["share_of_requests_percent"] = (
round((int(row["requests"]) / total_requests) * 100.0, 2) if total_requests else 0.0
)
rows.sort(
key=lambda row: (
int(row.get("tokens_saved", 0)),
int(row.get("before_tokens", 0)),
int(row.get("requests", 0)),
),
reverse=True,
)
return {
"agents": rows,
"totals": {
"requests": total_requests,
"before_tokens": global_before_tokens,
"after_tokens": global_after_tokens,
"output_tokens": global_output_tokens,
"tokens_saved": global_tokens_saved,
"savings_percent": (
round((global_tokens_saved / global_before_tokens) * 100.0, 2)
if global_before_tokens
else 0.0
),
},
"coverage": {
"logged_requests": len(logs),
"exact_token_rows": sum(1 for row in rows if row.get("has_exact_tokens")),
"mode": "request_logs" if logs else "aggregate_fallback",
},
}
# Suppress "[transformers] PyTorch was not found" warning emitted when
# transformers is imported for availability checks (e.g. kompress ONNX probe).
# PyTorch is optional in headroom; the warning is not actionable for operators.
os.environ.setdefault("TRANSFORMERS_VERBOSITY", "error")
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("headroom.proxy")
LoopExceptionHandler = Callable[[asyncio.AbstractEventLoop, dict[str, Any]], object]
class LoopFailureDetails(TypedDict):
message: Any | None
exception: str | None
class LoopHealthState(TypedDict):
status: str
known_failures: int
last_known_failure: LoopFailureDetails | None
_MULTI_WORKER_CONFIG_ENV = "HEADROOM_PROXY_CONFIG_JSON"
# Env var that opts out of the Rust core deployment smoke test (Hotfix-A0).
# Default behavior: hard-fail at startup if `headroom._core` is unimportable
# (Finding #2 in HEADROOM_PROXY_LOG_FINDINGS_2026_05_03.md — production
# deployment was silently running without the Rust extension and degrading
# every compressed request to a Python-only path or a no-op).
#
# Set to the literal string "false" to start the proxy in degraded
# Python-only mode. Any other value (including unset) keeps the
# fail-loud behavior.
_RUST_CORE_REQUIRED_ENV = "HEADROOM_REQUIRE_RUST_CORE"
# sysexits.h(3) — EX_CONFIG. Process supervisors (systemd, k8s, docker)
# treat this as a deliberate configuration failure rather than a crash, so
# they won't restart-loop on a broken deployment.
_EXIT_CONFIG = 78
def _check_rust_core() -> tuple[str, str | None]:
"""Verify the Rust extension `headroom._core` is loadable at startup.
Returns a `(status, error)` tuple:
- ``("loaded", None)`` — `headroom._core.hello()` returned the
expected sentinel.
- ``("disabled", reason)`` — opt-out env var was set; proxy starts
in Python-only degraded mode. `reason` carries the underlying
import error (or ``None`` if the import actually succeeded).
- ``("missing", reason)`` — never returned: this branch calls
``sys.exit(78)`` so the proxy refuses to start. The branch exists
only as a typed sentinel for callers that want to reason about
all three states (e.g. health endpoints).
Behavior is gated by the ``HEADROOM_REQUIRE_RUST_CORE`` env var:
any value other than ``"false"`` (case-insensitive) keeps the
fail-loud default.
"""
require = os.environ.get(_RUST_CORE_REQUIRED_ENV, "true").strip().lower() != "false"
try:
from headroom._core import hello as _rust_hello
marker = _rust_hello()
except Exception as exc: # ImportError, but also any init-time PyO3 failure
reason = f"{type(exc).__name__}: {exc}"
if not require:
logger.warning(
"event=rust_core_disabled reason=%r opt_out_env=%s=false mode=python_only_degraded",
reason,
_RUST_CORE_REQUIRED_ENV,
)
return ("disabled", reason)
# Fail loud. Print to stderr in addition to logging so operators
# see it even if the logging handler is mis-configured.
msg = (
f"FATAL: Rust extension `headroom._core` not loadable.\n"
f" error: {reason}\n"
f" fix: `make build-wheel && pip install --force-reinstall "
f"target/wheels/headroom_*.whl`\n"
f" opt-out: set {_RUST_CORE_REQUIRED_ENV}=false to start in "
f"degraded Python-only mode\n"
)
logger.error("event=rust_core_missing reason=%r action=exit_78", reason)
print(msg, file=sys.stderr, flush=True)
sys.exit(_EXIT_CONFIG)
# Import succeeded; sanity-check the marker so we catch a stale or
# mis-linked .so where the symbol name resolves but returns garbage.
if marker != "headroom-core":
reason = f"unexpected marker {marker!r}"
if not require:
logger.warning(
"event=rust_core_disabled reason=%r opt_out_env=%s=false",
reason,
_RUST_CORE_REQUIRED_ENV,
)
return ("disabled", reason)
msg = (
f"FATAL: Rust extension `headroom._core` is loaded but the "
f"marker function returned {marker!r}; expected 'headroom-core'.\n"
f" fix: rebuild: `make build-wheel && pip install "
f"--force-reinstall target/wheels/headroom_*.whl`\n"
)
logger.error("event=rust_core_marker_mismatch marker=%r action=exit_78", marker)
print(msg, file=sys.stderr, flush=True)
sys.exit(_EXIT_CONFIG)
logger.info("event=rust_core_loaded marker=%r", marker)
return ("loaded", None)
# Compression pipeline timeout in seconds
from headroom.proxy.handlers import ( # noqa: E402
AnthropicHandlerMixin,
BatchHandlerMixin,
BedrockHandlerMixin,
GeminiHandlerMixin,
OpenAIHandlerMixin,
StreamingMixin,
)
def _apply_stateless_persistence(config: ProxyConfig) -> None:
"""When the proxy runs stateless, force global persisters to in-memory so no
files are written to the workspace.
Covers TOIN (the always-on serving writer): it keeps learning patterns
in-memory but never reads or writes ``toin.json``. An empty ``storage_path``
makes the backend ``None``, which no-ops load/save/auto-save. The savings
subsystem is handled separately via ``PrometheusMetrics(stateless=...)``.
Note: setting ``HEADROOM_TOIN_BACKEND=none`` is NOT sufficient on its own —
``ToolIntelligenceNetwork`` falls back to ``config.storage_path`` when no
backend is passed, so we must clear the path explicitly here.
Concurrency: ``stateless`` is a per-process flag (set once at ``headroom
proxy`` launch), never a per-request/per-session value — every session a
process serves shares it, and two proxies with different settings run as
separate OS processes with independent TOIN singletons. In the rare case
where two HeadroomProxy instances with different ``stateless`` settings live
in ONE process (e.g. tests), this fails closed: the reset forces the
process-global TOIN in-memory, so a stateless proxy never persists (the safe
direction). A co-resident stateful proxy would then also stop persisting
TOIN — acceptable, since not-writing can never leak data.
"""
if not getattr(config, "stateless", False):
return
from headroom.telemetry.toin import TOINConfig, get_toin, reset_toin
# Reset first so this wins regardless of whether the singleton was already
# created with a filesystem backend earlier in the process.
reset_toin()
get_toin(TOINConfig(storage_path=""))
def _provider_httpx_client_options(
config: ProxyConfig,
verify: Any,
) -> tuple[bool, dict[str, Any]]:
client_kwargs: dict[str, Any] = {
"timeout": httpx.Timeout(
connect=config.connect_timeout_seconds,
read=config.request_timeout_seconds,
write=config.request_timeout_seconds,
pool=config.connect_timeout_seconds,
),
"limits": httpx.Limits(
max_connections=config.max_connections,
max_keepalive_connections=config.max_keepalive_connections,
keepalive_expiry=config.keepalive_expiry,
),
"verify": verify,
}
if config.http_proxy:
client_kwargs["proxy"] = config.http_proxy
return config.http2 and not config.http_proxy, client_kwargs
class HeadroomProxy(
StreamingMixin,
AnthropicHandlerMixin,
OpenAIHandlerMixin,
GeminiHandlerMixin,
BatchHandlerMixin,
BedrockHandlerMixin,
):
"""Production-ready Headroom optimization proxy."""
ANTHROPIC_API_URL = DEFAULT_ANTHROPIC_API_URL
OPENAI_API_URL = DEFAULT_OPENAI_API_URL
GEMINI_API_URL = DEFAULT_GEMINI_API_URL
CLOUDCODE_API_URL = DEFAULT_CLOUDCODE_API_URL
VERTEX_API_URL = DEFAULT_VERTEX_API_URL
def __init__(self, config: ProxyConfig):
self.config = config
self.config.mode = normalize_proxy_mode(self.config.mode)
# Record process-wide stateless mode so module-level persisters
# (output-savings recorder, etc.) can skip workspace writes.
from headroom import paths as _hr_paths
_hr_paths.set_process_stateless(config.stateless)
# Stateless: keep TOIN learning in-memory; never touch toin.json.
_apply_stateless_persistence(self.config)
pipeline_extensions = list(config.pipeline_extensions or [])
probe_recorder = probe_recorder_from_env()
if probe_recorder is not None:
pipeline_extensions.append(probe_recorder)
self.pipeline_extensions = PipelineExtensionManager(
hooks=config.hooks,
extensions=pipeline_extensions,
discover=config.discover_pipeline_extensions,
)
self.provider_runtime = build_proxy_provider_runtime(config)
api_targets = self.provider_runtime.api_targets
# Preserve the long-standing proxy compatibility surface while keeping
# provider_runtime as the source of truth for resolved upstream targets.
HeadroomProxy.ANTHROPIC_API_URL = api_targets.anthropic
HeadroomProxy.OPENAI_API_URL = api_targets.openai
HeadroomProxy.GEMINI_API_URL = api_targets.gemini
HeadroomProxy.CLOUDCODE_API_URL = api_targets.cloudcode
HeadroomProxy.VERTEX_API_URL = api_targets.vertex
self.anthropic_provider = self.provider_runtime.pipeline_provider("anthropic")
self.openai_provider = self.provider_runtime.pipeline_provider("openai")
# `metrics` is hoisted ahead of transform construction so the
# transforms can receive `self.metrics` as their compression
# observer at __init__ time. The forcing function for catching
# silent strategy regressions: per-strategy counters increment
# only when wired up here, so the wiring is mandatory, not
# something we patch in later. (See `RUST_DEV.md` audit notes.)
self.cost_tracker = (
CostTracker(
budget_limit_usd=config.budget_limit_usd,
budget_period=config.budget_period,
)
if config.cost_tracking_enabled
else None
)
self.metrics = PrometheusMetrics(cost_tracker=self.cost_tracker, stateless=config.stateless)
# Initialize transforms based on routing mode.
#
# Phase B PR-B1 retired the IntelligentContextManager / RollingWindow
# message-dropping branch. Live-zone-only compression (PR-B2..B7) does
# not drop messages — it operates on content blocks within messages —
# so the proxy no longer needs a "context manager" transform stage.
# Reported via metrics as `_context_manager_status = "passthrough"`.
self._context_manager_status = "passthrough"
# ContentRouter is the single proxy routing surface. Provider handlers
# normalize their request shapes into messages or CompressionUnits, and
# the router chooses SmartCrusher, log/search/diff/code, or Kompress.
profile_kwargs = proxy_pipeline_kwargs(config)
router_config = ContentRouterConfig(
enable_code_aware=config.code_aware_enabled,
prefer_code_aware_for_code=_get_env_bool("HEADROOM_PREFER_CODE_AWARE_FOR_CODE", True),
tool_profiles=config.tool_profiles,
read_lifecycle=ReadLifecycleConfig(enabled=config.read_lifecycle),
smart_crusher_max_items_after_crush=cast(
int | None,
profile_kwargs.get("max_items_after_crush"),
),
smart_crusher_with_compaction=cast(
bool,
profile_kwargs.get("smart_crusher_with_compaction", True),
),
ccr_inject_marker=config.ccr_inject_marker,
force_kompress_all=config.force_kompress_all,
lossless=config.lossless,
)
# No-CCR lossless mode: compress tool outputs with format-native
# lossless compaction and marker-free SmartCrusher, and suppress every
# retrieval marker + the retrieve-tool injection so no MCP round-trip is
# needed. Mirrors the force_kompress_all wiring precedent.
if config.lossless:
router_config.lossless = True
router_config.smart_crusher_lossless_only = True
router_config.ccr_inject_marker = False
if hasattr(config, "ccr_inject_tool"):
config.ccr_inject_tool = False
if config.disable_kompress:
router_config.enable_kompress = False
# Opt-in restore of the legacy behaviour: send fall-through content
# to PASSTHROUGH instead of the default KOMPRESS fallback strategy.
if config.disable_kompress_fallback:
router_config.fallback_strategy = CompressionStrategy.PASSTHROUGH
# `HEADROOM_LOSSLESS_ONLY=1` routes SmartCrusher through strict
# marker-free mode: lossless tabular compaction still applies, but
# any path that would emit a `<<ccr:…>>` marker (row-drop or
# opaque-blob offload) leaves the content uncompacted instead — so
# the session needs no CCR retrieval round-trips to stay recoverable.
if "HEADROOM_LOSSLESS_ONLY" in os.environ:
router_config.smart_crusher_lossless_only = _get_env_bool(
"HEADROOM_LOSSLESS_ONLY", False
)
# A non-None exclude_tools replaces DEFAULT_EXCLUDE_TOOLS in
# ContentRouter, so merge rather than assign.
if config.exclude_tools:
router_config.exclude_tools = set(DEFAULT_EXCLUDE_TOOLS) | config.exclude_tools
# protect_tool_results: force-merge named tools into the exclude set
# so their results are never lossy-compressed, regardless of mode.
if config.protect_tool_results:
base = (
router_config.exclude_tools
if router_config.exclude_tools is not None
else set(DEFAULT_EXCLUDE_TOOLS)
)
router_config.exclude_tools = base | config.protect_tool_results
# Token mode: allow compression of older excluded-tool results,
# and emit search results grouped by file (path once per file
# instead of repeated on every match line).
if is_token_mode(config.mode):
router_config.protect_recent_reads_fraction = 0.3
router_config.search_group_by_file = True
if config.protect_tool_results:
router_config.protect_recent_reads_fraction = 0.0
# `--compress-user-messages` flips the router's default skip rule.
# Off by default for prefix-cache safety; enabled for workloads where
# user-message content dominates input (OpenAI/Azure chat with pasted
# code/RAG context — see issue #454).
if profile_kwargs.get("compress_user_messages"):
router_config.skip_user_messages = False
# Kompress (lossy ML text compression) is resolved per provider. The
# global `disable_kompress` above is the baseline for both; a per-
# provider override (disable_kompress_{anthropic,openai}) wins when set.
# Only `enable_kompress` differs between providers — routing, tool
# exclusion, and read-protection are identical — so when both resolve
# the same we reuse ONE ContentRouter instance and the Kompress model
# still loads once (startup warmup dedupes transforms by id()).
base_kompress_disabled = not router_config.enable_kompress
anthropic_kompress_disabled = (
base_kompress_disabled
if config.disable_kompress_anthropic is None
else config.disable_kompress_anthropic
)
openai_kompress_disabled = (
base_kompress_disabled
if config.disable_kompress_openai is None
else config.disable_kompress_openai
)
def _router_config_for(kompress_disabled: bool) -> ContentRouterConfig:
if kompress_disabled == base_kompress_disabled:
return router_config
return replace(router_config, enable_kompress=not kompress_disabled)
cache_aligner = CacheAligner(CacheAlignerConfig(enabled=False))
anthropic_router = ContentRouter(
_router_config_for(anthropic_kompress_disabled), observer=self.metrics
)
openai_router = (
anthropic_router
if openai_kompress_disabled == anthropic_kompress_disabled
else ContentRouter(_router_config_for(openai_kompress_disabled), observer=self.metrics)
)
self._code_aware_status = "lazy" if config.code_aware_enabled else "disabled"
_intercept_prefix: list = []
if os.environ.get("HEADROOM_INTERCEPT_ENABLED"):
from headroom.proxy.interceptors import ToolResultInterceptorTransform
_intercept_prefix = [ToolResultInterceptorTransform()]
self.anthropic_pipeline = TransformPipeline(
transforms=[*_intercept_prefix, cache_aligner, anthropic_router],
provider=self.anthropic_provider,
)
self.openai_pipeline = TransformPipeline(
transforms=[*_intercept_prefix, cache_aligner, openai_router],
provider=self.openai_provider,
)
# Initialize components
self.cache = (
SemanticCache(
max_entries=config.cache_max_entries,
ttl_seconds=config.cache_ttl_seconds,
)
if config.cache_enabled
else None
)
self.rate_limiter = (
TokenBucketRateLimiter(
requests_per_minute=config.rate_limit_requests_per_minute,
tokens_per_minute=config.rate_limit_tokens_per_minute,
)
if config.rate_limit_enabled
else None
)
# `cost_tracker` and `metrics` were hoisted to before transforms so
# ContentRouter / SmartCrusher could take `self.metrics` as their
# compression observer at __init__ time.
# Prefix cache tracking: freeze already-cached messages to avoid
# invalidating the provider's prefix cache with our transforms
from headroom.cache.prefix_tracker import PrefixFreezeConfig, SessionTrackerStore
self.session_tracker_store = SessionTrackerStore(
default_config=PrefixFreezeConfig(
enabled=config.prefix_freeze_enabled,
session_ttl_seconds=config.prefix_freeze_session_ttl,
)
)
# Compression cache store for token mode (session-scoped). The dict
# itself is mutated under `_compression_caches_lock`; the per-session
# `CompressionCache` instances have their own internal lock guarding
# `_cache`/`_stable_hashes`/`_first_seen` against concurrent
# async-dispatched requests for the same session.
self._compression_caches: dict[str, CompressionCache] = {}
self._compression_caches_lock = threading.RLock()
self.logger = (
RequestLogger(
log_file=config.log_file,
log_full_messages=config.log_full_messages,
)
if config.log_requests
else None
)
# Enterprise security plugin (loaded dynamically if available + licensed)
self.security = None
# HTTP client
self.http_client: httpx.AsyncClient | None = None
# HTTP/1.1-only client for ChatGPT passthrough (Cloudflare challenges
# our HTTP/2 fingerprint on its sensitive account endpoints).
self.http_client_h1: httpx.AsyncClient | None = None
self._shutdown_event: asyncio.Event | None = None
# Shared cold-start warmup registry (populated by startup()).
# Holds typed slots with loaded / loading / null / error status for
# each preloaded heavy asset. Exposed as ``proxy.warmup`` and
# serialized by the /debug/warmup route (Unit 5).
self.warmup: WarmupRegistry = WarmupRegistry()
# Unit 3: live registry of Codex WS sessions. Populated by
# ``handle_openai_responses_ws`` on accept; drained in its
# outermost ``finally``. Consumed by ``/debug/ws-sessions``.
self.ws_sessions: WebSocketSessionRegistry = WebSocketSessionRegistry()
# Unit 4: bounded pre-upstream concurrency for the Anthropic HTTP
# path. Caps how many ``handle_anthropic_messages`` calls may be
# running deep-copy / first-stage compression / memory-context
# lookup / upstream connect concurrently. ``/livez``, ``/readyz``,
# ``/health``, ``/metrics``, ``/stats``, and the Codex WS path are
# intentionally NOT gated by this semaphore.
#
# A value of ``0`` or negative disables the semaphore (unbounded
# mode); this is useful for the Unit 6 counter-factual where we
# deliberately reproduce the original starvation. The default is
# ``max(2, min(8, os.cpu_count() or 4))``.
_pre_upstream_cfg = config.anthropic_pre_upstream_concurrency
if _pre_upstream_cfg is None:
_pre_upstream_resolved = max(2, min(8, os.cpu_count() or 4))
else:
_pre_upstream_resolved = _pre_upstream_cfg
self.anthropic_pre_upstream_concurrency: int = _pre_upstream_resolved
self.anthropic_pre_upstream_acquire_timeout_seconds = float(
config.anthropic_pre_upstream_acquire_timeout_seconds
)
self.anthropic_pre_upstream_memory_context_timeout_seconds = float(
config.anthropic_pre_upstream_memory_context_timeout_seconds
)
if _pre_upstream_resolved > 0:
self.anthropic_pre_upstream_sem: asyncio.Semaphore | None = asyncio.Semaphore(
_pre_upstream_resolved
)
else:
self.anthropic_pre_upstream_sem = None
# Dedicated compression executor — see C3 in the audit followup.
# Replaces ``asyncio.to_thread(...)`` for ``pipeline.apply()`` calls
# so that:
# 1. Compression work is bounded — CPU-bound Rust runs here, and
# bursts cannot starve other ``asyncio.to_thread`` callers
# sharing the loop's default executor (file IO, etc.).
# 2. Tasks that exceed ``COMPRESSION_TIMEOUT_SECONDS`` and complete
# *after* the asyncio future was cancelled are counted in the
# ``compression_leaked_threads`` gauge — Python cannot preempt
# the worker, so this is the only signal that some pool slots
# are sitting on stuck work.
_compression_max_cfg = config.compression_max_workers
if _compression_max_cfg is None:
_compression_max = max(1, os.cpu_count() or 1)
else:
_compression_max = max(1, _compression_max_cfg)
self.compression_max_workers: int = _compression_max
self._compression_executor = concurrent.futures.ThreadPoolExecutor(
max_workers=_compression_max,
thread_name_prefix="headroom-compress",
)
# Phase 3 (#1171): off-path background compression. When enabled, a
# cold-start-large request (frozen=0 + large live zone) forwards
# uncompressed immediately and enqueues the compression here instead of
# blocking the request thread under the 30s budget (which leaks a
# non-preemptible worker -> executor saturation -> cascade). Default
# off (opt-in), fail-open. Per-process, matching _compression_caches.
self._background_compression_enabled: bool = os.environ.get(
"HEADROOM_BACKGROUND_COMPRESSION", ""
).strip().lower() in ("1", "true", "yes", "on")
try:
self._background_compression_min_tokens: int = int(
os.environ.get("HEADROOM_BACKGROUND_COMPRESSION_MIN_TOKENS", "50000")
)
except ValueError:
self._background_compression_min_tokens = 50000
# Dedicated single thread: no-timeout background jobs never contend with
# the request-path executor (Phase 3, #1171). Lazy -- no thread spawns
# until the first off-path job is submitted.
self._background_compression_executor = concurrent.futures.ThreadPoolExecutor(
max_workers=1, thread_name_prefix="headroom-bg-compress"
)
self._background_compressor = BackgroundCompressor(self._run_compression_background)
# Gauge: currently-running compression tasks. Mutated under
# ``_compression_metrics_lock`` from worker threads + the asyncio
# event loop.
self._compression_queued: int = 0
self._compression_queued_max: int = 0
self._compression_queue_timeouts: int = 0
self._compression_queue_wait_seconds_total: float = 0.0
self._compression_queue_wait_seconds_max: float = 0.0
self._compression_in_flight: int = 0
# High-water mark for in-flight count.
self._compression_in_flight_max: int = 0
self._compression_run_seconds_total: float = 0.0
self._compression_run_seconds_max: float = 0.0
# Counter: threads that finished AFTER their asyncio future hit the
# timeout. Stuck-thread leak indicator.
self._compression_leaked_threads: int = 0
self._compression_metrics_lock = threading.Lock()
# Backend for Anthropic API (direct, LiteLLM, or any-llm)
# Supports: "anthropic" (direct), "bedrock", "vertex", "litellm-<provider>", or "anyllm"
self.anthropic_backend: Backend | None = create_proxy_backend(
backend=config.backend,
anyllm_provider=config.anyllm_provider,
bedrock_region=config.bedrock_region,
bedrock_profile=config.bedrock_profile,
logger=logger,
openai_api_url=config.openai_api_url,
anyllm_backend_cls=AnyLLMBackend,
litellm_backend_cls=LiteLLMBackend,
)
# Request counter for IDs
self._request_counter = 0
self._request_counter_lock = asyncio.Lock()
# CCR tool injectors (one per provider)
self.anthropic_tool_injector = CCRToolInjector(
provider="anthropic",
inject_tool=config.ccr_inject_tool,
inject_system_instructions=config.ccr_inject_system_instructions,
)
self.openai_tool_injector = CCRToolInjector(
provider="openai",
inject_tool=config.ccr_inject_tool,
inject_system_instructions=config.ccr_inject_system_instructions,
)
# CCR Response Handler (handles CCR tool calls automatically)
self.ccr_response_handler = (
CCRResponseHandler(
ResponseHandlerConfig(
enabled=True,
max_retrieval_rounds=config.ccr_max_retrieval_rounds,
)
)
if config.ccr_handle_responses
else None
)
# CCR Context Tracker (tracks compressed content across turns)
self.ccr_context_tracker = (
ContextTracker(
ContextTrackerConfig(
enabled=True,
proactive_expansion=config.ccr_proactive_expansion,
max_proactive_expansions=config.ccr_max_proactive_expansions,
)
)
if config.ccr_context_tracking
else None
)
# Turn counter for context tracking
self._turn_counter = 0
# Memory Handler (persistent user memory)
self.memory_handler: MemoryHandler | None = None
if config.memory_enabled and config.stateless:
# Persistent memory writes a SQLite DB + markdown files to disk,
# which stateless mode forbids. Memory is cross-session learning and
# is contradictory with an ephemeral/read-only deployment, so we
# disable it rather than persist. Run without --stateless to use it.
logger.warning(
"Memory is disabled in stateless mode (it persists to disk). "
"Run without --stateless to enable persistent memory."
)
elif config.memory_enabled:
# Resolve memory DB path: empty → project-scoped default
_mem_db_path = config.memory_db_path
if not _mem_db_path:
_mem_dir = Path.cwd() / ".headroom"
_mem_dir.mkdir(parents=True, exist_ok=True)
_mem_db_path = str(_mem_dir / "memory.db")
logger.info(f"Memory: Project-scoped DB at {_mem_db_path}")
# PR-B6: translate the string-typed ``ProxyConfig.memory_mode``
# into the typed ``MemoryMode`` enum. Unknown values raise
# loudly per the no-silent-fallback policy.
from headroom.proxy.memory_handler import MemoryMode
try:
_memory_mode = MemoryMode(config.memory_mode)
except ValueError as exc:
raise ValueError(
f"Invalid memory_mode={config.memory_mode!r}; "
f"expected one of {[m.value for m in MemoryMode]}"
) from exc
from headroom.memory.storage_router import MemoryStorageMode
try:
_storage_mode = MemoryStorageMode(config.memory_storage_mode)
except ValueError as exc:
raise ValueError(
f"Invalid memory_storage_mode={config.memory_storage_mode!r}; "
f"expected one of {[m.value for m in MemoryStorageMode]}"
) from exc
memory_config = MemoryConfig(
enabled=True,
backend=config.memory_backend,
db_path=_mem_db_path,
inject_tools=config.memory_inject_tools,
use_native_tool=config.memory_use_native_tool,
inject_context=config.memory_inject_context,
top_k=config.memory_top_k,
min_similarity=config.memory_min_similarity,
mode=_memory_mode,
storage_mode=_storage_mode,
project_root_override=config.memory_project_root_override,
qdrant_url=config.memory_qdrant_url,
qdrant_host=config.memory_qdrant_host,
qdrant_port=config.memory_qdrant_port,
qdrant_api_key=config.memory_qdrant_api_key,
neo4j_uri=config.memory_neo4j_uri,
neo4j_user=config.memory_neo4j_user,
neo4j_password=config.memory_neo4j_password,
bridge_enabled=config.memory_bridge_enabled,
bridge_md_paths=config.memory_bridge_md_paths,
bridge_md_format=config.memory_bridge_md_format,
bridge_auto_import=config.memory_bridge_auto_import,
bridge_export_path=config.memory_bridge_export_path,
)
self.memory_handler = MemoryHandler(
memory_config,
agent_type=config.traffic_learning_agent_type,
)
# Migration UX (GH #462). When the user is on the new
# project-scoped default but a legacy single-file DB exists
# with prior memories, surface that clearly so it doesn't
# look like an upgrade ate their data.
if _storage_mode is MemoryStorageMode.PROJECT:
_legacy_path = Path(_mem_db_path)
if _legacy_path.exists() and _legacy_path.stat().st_size > 0:
logger.info(
"event=memory_storage_legacy_detected path=%s mode=project "
"hint=pass_--memory-storage=global_to_reach_pre-fix_memories",
_legacy_path,
)
# The Memory Bridge binds to the single legacy backend at
# init time; it doesn't (yet) follow per-project routing.
# Warn so users running bridge + project mode aren't
# surprised that only the legacy DB syncs with markdown.
if config.memory_bridge_enabled and _storage_mode is MemoryStorageMode.PROJECT:
logger.warning(
"event=memory_bridge_global_only mode=project "
"hint=bridge_syncs_only_the_legacy_DB_today_per-project_bridge_follow-up_planned"
)
# Usage Reporter (license validation + phone-home for managed/enterprise).
# Suppressed entirely in offline mode — the air-gap switch must stop all
# egress, including license phone-home, even when a key is configured.
self.usage_reporter: UsageReporter | None = None
if config.license_key and not (config.offline or is_offline()):
from headroom.telemetry.reporter import UsageReporter
self.usage_reporter = UsageReporter(
license_key=config.license_key,
cloud_url=config.license_cloud_url,
report_interval=config.license_report_interval,
)
# Traffic Learner (live pattern extraction from proxy traffic)
# Only activates with --learn flag; requires --memory for backend
self.traffic_learner: TrafficLearner | None = None
self.traffic_learning_agent_type: str = config.traffic_learning_agent_type
if config.traffic_learning_enabled:
from headroom.memory.traffic_learner import TrafficLearner
self.traffic_learner = TrafficLearner(
user_id=os.environ.get("HEADROOM_USER_ID", os.environ.get("USER", "default")),
agent_type=config.traffic_learning_agent_type,
min_evidence=config.traffic_learning_min_evidence,
)
# Code graph file watcher (live reindex on file changes)
self.code_graph_watcher: CodeGraphWatcher | None = None # type: ignore[annotation-unchecked]
if config.code_graph_watcher:
from headroom.graph.watcher import CodeGraphWatcher
self.code_graph_watcher = CodeGraphWatcher(project_dir=Path.cwd())
if self.code_graph_watcher.start():
logger.info("Code graph: file watcher started")
else:
self.code_graph_watcher = None
self.pipeline_extensions.emit(
PipelineStage.SETUP,
operation="proxy.setup",
metadata={
"mode": self.config.mode,
"optimize": self.config.optimize,
"backend": self.config.backend,
"memory_enabled": self.config.memory_enabled,
},
)
async def _run_compression_in_executor(
self,
fn, # noqa: ANN001 — caller-supplied no-arg sync callable
*,
timeout: float,
):
"""Run a synchronous compression callable on the bounded executor
with cancel-aware metrics.
Replaces ``asyncio.wait_for(asyncio.to_thread(fn), timeout=...)``.
Why a dedicated executor: the proxy's compression path is CPU-bound
Rust work that releases the GIL via ``py.allow_threads``. Sharing
the loop's default executor (used by ``asyncio.to_thread``) means
a burst of slow compressions can starve unrelated ``to_thread``
callers (file IO, etc.). The compression executor is sized
independently via ``config.compression_max_workers``.
Why "cancel-aware metrics": when ``asyncio.wait_for`` times out, it
cancels the *asyncio future*. The underlying
``concurrent.futures.Future`` from ``run_in_executor`` cannot
actually cancel a thread that has started — Python has no way to
preempt running CPython bytecode or in-flight Rust calls. The
worker keeps running to completion, ignored. We detect this by
marking the call timed out on the asyncio side and incrementing
``_compression_leaked_threads`` from the worker's ``finally``
block after it eventually finishes. Jobs that time out before a
worker starts are removed from the queued gauge instead. Operators
can see leaked-thread rate and queue pressure climbing in
``/stats`` before the pool fills up.
Args:
fn: A no-arg sync callable that runs the compression. Must not
raise asyncio Cancellation; if it does, the wrapper still
decrements the in-flight gauge but the leaked-thread
counter may double-count.
timeout: Wall-clock timeout for the asyncio side. The
executor worker keeps running past this (Python limitation
— see above), but at least the awaiter unblocks.
Returns:
Whatever ``fn()`` returns.
Raises:
``asyncio.TimeoutError`` if the callable doesn't return within
``timeout``. Any exception raised by ``fn`` propagates
unchanged.
"""
loop = asyncio.get_running_loop()
queued_at = time.monotonic()
state = {"queued": True, "timed_out": False}
with self._compression_metrics_lock:
self._compression_queued += 1
if self._compression_queued > self._compression_queued_max:
self._compression_queued_max = self._compression_queued
def _wrapped(): # noqa: ANN202
started_at = time.monotonic()
queue_wait = started_at - queued_at
with self._compression_metrics_lock:
if state["queued"]:
self._compression_queued -= 1
state["queued"] = False
self._compression_queue_wait_seconds_total += queue_wait
if queue_wait > self._compression_queue_wait_seconds_max:
self._compression_queue_wait_seconds_max = queue_wait
self._compression_in_flight += 1
if self._compression_in_flight > self._compression_in_flight_max:
self._compression_in_flight_max = self._compression_in_flight
try:
return fn()
finally:
elapsed = time.monotonic() - started_at
with self._compression_metrics_lock:
self._compression_in_flight -= 1
self._compression_run_seconds_total += elapsed
if elapsed > self._compression_run_seconds_max:
self._compression_run_seconds_max = elapsed
if state["timed_out"]:
self._compression_leaked_threads += 1
future = loop.run_in_executor(self._compression_executor, _wrapped)
try:
return await asyncio.wait_for(future, timeout=timeout)
except asyncio.TimeoutError:
with self._compression_metrics_lock:
state["timed_out"] = True
if state["queued"]:
self._compression_queued -= 1
state["queued"] = False
self._compression_queue_timeouts += 1
raise
async def _run_compression_background(self, fn): # noqa: ANN001, ANN201
"""Run a compression callable on the shared executor with NO request-
coupled deadline (Phase 3 off-path, #1171).
Unlike ``_run_compression_in_executor`` there is no ``asyncio.wait_for``
and no leaked-thread accounting: no caller is waiting, so a slow run
backs up the background queue rather than starving the request executor.
Runs on the dedicated single-thread background executor.
"""
loop = asyncio.get_running_loop()
return await loop.run_in_executor(self._background_compression_executor, fn)
def _get_compression_cache(self, session_id: str) -> CompressionCache:
"""Get or create a CompressionCache for a session.
Thread-safe under `_compression_caches_lock`: a concurrent pair of
`_get_compression_cache(session_id)` calls (e.g. two async requests
for the same conversation) must return the **same** instance,
otherwise the per-session cache state splits and the two halves
diverge across requests.
"""
with self._compression_caches_lock:
if session_id not in self._compression_caches:
from headroom.cache.compression_cache import CompressionCache
# Evict oldest caches if at capacity
if len(self._compression_caches) >= MAX_COMPRESSION_CACHE_SESSIONS:
# Remove oldest quarter to amortize cleanup cost
oldest_keys = list(self._compression_caches.keys())[
: MAX_COMPRESSION_CACHE_SESSIONS // 4
]
for key in oldest_keys:
del self._compression_caches[key]
logger.info(
"Evicted %d compression caches (exceeded %d max sessions)",
len(oldest_keys),
MAX_COMPRESSION_CACHE_SESSIONS,
)
self._compression_caches[session_id] = CompressionCache()
return self._compression_caches[session_id]
def _setup_code_aware(self, config: ProxyConfig, transforms: list) -> str:
"""Set up code-aware compression if enabled.
Args:
config: Proxy configuration
transforms: Transform list to append to
Returns:
Status string for logging: 'enabled', 'disabled', 'available', 'unavailable'
"""
if config.code_aware_enabled:
if is_tree_sitter_available():
code_config = CodeCompressorConfig(
preserve_imports=True,
preserve_signatures=True,
preserve_type_annotations=True,
)
# CodeAware runs after the content/structure transforms.
# Phase B PR-B1 retired the trailing context_manager so we
# append rather than insert(-1).
transforms.append(CodeAwareCompressor(code_config))
return "enabled"
else:
logger.warning(
"Code-aware compression requested but tree-sitter not installed. "
"Install with: pip install headroom-ai[code]"
)
return "unavailable"
else:
if is_tree_sitter_available():
return "available" # Available but not enabled
return "disabled"
def _eager_preload_transforms(self) -> tuple[dict[str, str], list[dict[str, str]]]:
"""Eagerly load every compressor/parser/detector once (dedup by ``id()``).
Pure load: returns the merged ``eager_status`` plus the per-transform
status dicts for the caller to merge into ``self.warmup`` on the main
thread (``WarmupRegistry`` is not written off-thread). This runs via
``asyncio.to_thread`` so a slow or hung native model load cannot keep
startup from binding the port (#790).
"""
eager_status: dict[str, str] = {}
transform_statuses: list[dict[str, str]] = []
seen_transform_ids: set[int] = set()
for pipeline in (self.anthropic_pipeline, self.openai_pipeline):
for transform in pipeline.transforms:
if id(transform) in seen_transform_ids:
continue
seen_transform_ids.add(id(transform))
if not hasattr(transform, "eager_load_compressors"):
continue
try:
transform_status = transform.eager_load_compressors()
except Exception as exc:
logger.warning(
"Eager preload failed for %s: %s",
type(transform).__name__,
exc,
)
continue
if not isinstance(transform_status, dict):
continue
# Merge: later writers win only if the key wasn't set. Preload a
# transform ONCE — if another pipeline also has
# ``eager_load_compressors`` it contributes only new keys.
for key, value in transform_status.items():
eager_status.setdefault(key, value)
transform_statuses.append(transform_status)
return eager_status, transform_statuses
async def startup(self):
"""Initialize async resources."""
self._get_shutdown_event().clear()
self.pipeline_extensions.emit(
PipelineStage.PRE_START,
operation="proxy.startup",
metadata={"port": self.config.port, "host": self.config.host},
)
# Resolve TLS verification: a custom CA bundle (corporate PKI) if one
# is configured, else a strict-relaxed default context when
# HEADROOM_TLS_STRICT=0, else httpx's default strict verification.
_verify = build_httpx_verify()
_http2, _client_kwargs = _provider_httpx_client_options(self.config, _verify)
self.http_client = httpx.AsyncClient(http2=_http2, **_client_kwargs)
# Reuse the primary client when HTTP/2 is already off; otherwise keep a
# dedicated HTTP/1.1 client for ChatGPT passthrough.
self.http_client_h1 = (
self.http_client if not _http2 else httpx.AsyncClient(http2=False, **_client_kwargs)
)
logger.info("Headroom Proxy started (version %s)", __version__)
logger.info(f"Optimization: {'ENABLED' if self.config.optimize else 'DISABLED'}")
self.config.mode = normalize_proxy_mode(self.config.mode)
logger.info(f"Mode: {self.config.mode}")
if self.config.mode == PROXY_MODE_TOKEN:
logger.info(" Prefix freeze: re-freeze after compression")
logger.info(" Read protection window: 30%% of excluded-tool messages")
logger.info(" CCR TTL: extended for session lifetime")
logger.info(" Compression cache: active")
if self.config.mode == PROXY_MODE_CACHE:
logger.info(" Prefix freeze: strict (all prior turns immutable)")
logger.info(" Mutations: latest turn only")
logger.info(f"Caching: {'ENABLED' if self.config.cache_enabled else 'DISABLED'}")
logger.info(f"Rate Limiting: {'ENABLED' if self.config.rate_limit_enabled else 'DISABLED'}")
logger.info(
f"Connection Pool: max_connections={self.config.max_connections}, "
f"max_keepalive={self.config.max_keepalive_connections}, "
f"http2={'ENABLED' if _http2 else 'DISABLED'}"
)
# Unit 4 pre-upstream concurrency announcement. Report the resolved
# value (auto-detected vs. explicit) so operators can correlate
# ``pre_upstream_wait_ms`` log lines with the configured cap.
if self.anthropic_pre_upstream_sem is None:
logger.info("Anthropic pre-upstream concurrency: unbounded (explicitly disabled)")
else:
_explicit = self.config.anthropic_pre_upstream_concurrency
_origin = "auto-detected" if _explicit is None else "explicit"
logger.info(
"Anthropic pre-upstream concurrency: %d (%s)",
self.anthropic_pre_upstream_concurrency,
_origin,
)
logger.info(
"Anthropic pre-upstream timeouts: acquire=%.1fs compression=%.1fs memory_context=%.1fs",
self.anthropic_pre_upstream_acquire_timeout_seconds,
float(COMPRESSION_TIMEOUT_SECONDS),
self.anthropic_pre_upstream_memory_context_timeout_seconds,
)
logger.info("Smart Routing: ENABLED (ContentRouter is always active)")
# Eagerly load ALL compressors, parsers, and detectors at startup
# This eliminates cold-start latency spikes on first requests.
# Iterate BOTH pipelines (Anthropic + OpenAI) and dedupe transforms
# by id() so shared-transform instances never load twice. The
# resulting status dict is merged into ``self.warmup`` so /debug/warmup
# (Unit 5) and /readyz have a single source of truth.
self._kompress_status = "not installed"
eager_status: dict[str, str] = {}
if self.config.optimize:
logger.info("Pre-loading compressors and parsers...")
# Run the preload OFF the event loop with a bound. The loop body
# already swallows per-transform Exceptions, so the only thing that
# can still block ASGI lifespan startup (and therefore the socket
# bind) is a hang or an uncatchable native stall during a model load
# on Windows — the "never opens its port" failure in #790. Capping it
# means startup always returns and uvicorn binds; on timeout the
# transforms simply fall back to lazy loading on first use.
transform_statuses: list[dict[str, str]] = []
try:
eager_status, transform_statuses = await asyncio.wait_for(
asyncio.to_thread(self._eager_preload_transforms),
timeout=EAGER_PRELOAD_TIMEOUT_SECONDS,
)
except Exception as exc:
logger.warning(
"Eager preload exceeded %.0fs or failed (%s); continuing so "
"the proxy still binds — transforms load lazily on first use.",
EAGER_PRELOAD_TIMEOUT_SECONDS,
exc,
)
eager_status, transform_statuses = {}, []
# Merge warmup status on the main thread (WarmupRegistry is not
# written off-thread).
for transform_status in transform_statuses:
self.warmup.merge_transform_status(transform_status)
# Update internal status from eager loading results
if eager_status.get("kompress") == "enabled":
self._kompress_status = "enabled"
if eager_status.get("code_aware") == "enabled":
self._code_aware_status = "enabled"
# Log component status
if self._kompress_status == "enabled":
logger.info("Kompress: ENABLED (ModernBERT token compressor)")
elif self.config.optimize:
logger.info("Kompress: not installed (pip install headroom-ai[ml] for ML compression)")
if self._code_aware_status == "enabled":
logger.info("Code-Aware: ENABLED (AST-based compression)")
if "tree_sitter" in eager_status:
logger.info(f"Tree-Sitter: {eager_status['tree_sitter']}")
elif self._code_aware_status == "lazy":
logger.info("Code-Aware: LAZY (will load when code content detected)")
elif self._code_aware_status == "available":
logger.info("Code-Aware: available but disabled (use --code-aware)")
elif self._code_aware_status == "unavailable":
logger.info("Code-Aware: not installed (pip install headroom-ai[code])")
elif self._code_aware_status == "disabled":
logger.info("Code-Aware: DISABLED")
if eager_status.get("magika") == "enabled":
logger.info("Magika: ENABLED (ML content detection)")
if self.memory_handler:
if (
self.config.memory_backend == "qdrant-neo4j"
and not self.config.memory_neo4j_password
):
logger.warning(
"NEO4J password is not set — using default credentials is insecure in production"
)
self.warmup.memory_backend.mark_loading()
try:
await self.memory_handler.ensure_initialized()
except Exception as exc: # pragma: no cover - defensive
self.warmup.memory_backend.mark_error(str(exc))
logger.warning("Memory: backend initialization failed (startup continues): %s", exc)
memory_status = self.memory_handler.health_status()
if memory_status.get("initialized"):
self.warmup.memory_backend.mark_loaded(
handle=self.memory_handler,
backend=memory_status.get("backend"),
)
# Force one embed call so the ONNX graph is compiled now,
# not lazily during the first request. Best-effort — any
# failure is swallowed inside warmup_embedder.
self.warmup.memory_embedder.mark_loading()
warmed = await self.memory_handler.warmup_embedder()
if warmed:
self.warmup.memory_embedder.mark_loaded()
else:
# Not an error — e.g. qdrant-neo4j has no embedder slot
# we can reach, or the backend simply exposes no handle.
self.warmup.memory_embedder.mark_null()
else:
if self.warmup.memory_backend.status != "error":
self.warmup.memory_backend.mark_null()
self.warmup.memory_embedder.mark_null()
logger.info(
"Memory: ENABLED "
f"(backend={memory_status['backend']}, initialized={memory_status['initialized']})"
)
else:
logger.info("Memory: DISABLED")
# CCR status
ccr_features = []
if self.config.ccr_inject_tool:
ccr_features.append("tool_injection")
if self.config.ccr_handle_responses:
ccr_features.append("response_handling")
if self.config.ccr_context_tracking:
ccr_features.append("context_tracking")
if self.config.ccr_proactive_expansion:
ccr_features.append("proactive_expansion")
if ccr_features:
logger.info(f"CCR (Compress-Cache-Retrieve): ENABLED ({', '.join(ccr_features)})")
else:
logger.info("CCR: DISABLED")
logger.info(f"Savings history: {self.metrics.savings_tracker.storage_path}")
# Reset and rebuild the quota tracker registry for this server instance.
# reset_quota_registry() ensures a clean slate when the proxy is restarted
# (e.g. in tests that spin up multiple app instances in the same process).
reset_quota_registry()
registry = get_quota_registry()
tracker = configure_subscription_tracker(
poll_interval_s=self.config.subscription_poll_interval_s,
active_window_s=self.config.subscription_active_window_s,
enabled=self.config.subscription_tracking_enabled,
)
registry.register(tracker)
registry.register(get_codex_rate_limit_state())
registry.register(get_copilot_quota_tracker())
await registry.start_all()
if self.config.subscription_tracking_enabled:
logger.info(
"Subscription tracking: ENABLED "
f"(poll_interval={self.config.subscription_poll_interval_s}s, "
f"active_window={self.config.subscription_active_window_s}s)"
)
else:
logger.info("Subscription tracking: DISABLED")
copilot_tracker = get_copilot_quota_tracker()
if copilot_tracker.is_available():
logger.info("GitHub Copilot quota tracking: ENABLED")
else:
logger.info(
"GitHub Copilot quota tracking: DISABLED "
"(set GITHUB_TOKEN or GITHUB_COPILOT_GITHUB_TOKEN to enable)"
)
# Log local telemetry status so operators can see it in the log stream.
# Nothing is sent externally — telemetry is collected locally only (the
# anonymous telemetry beacon was removed); operational metrics export
# only to your own OTEL collector via HEADROOM_OTEL_METRICS_*.
if is_telemetry_enabled():
logger.info(
"Local telemetry: ENABLED (aggregate stats, local only — nothing sent "
"externally). Opt out: HEADROOM_TELEMETRY=off or --no-telemetry"
)
else:
logger.info(
"Local telemetry: DISABLED (off by default — opt in: "
"HEADROOM_TELEMETRY=on or --telemetry)"
)
self.pipeline_extensions.emit(
PipelineStage.POST_START,
operation="proxy.startup",
metadata={
"port": self.config.port,
"host": self.config.host,
"warmup": self.warmup.to_dict(),
},
)
async def shutdown(self):
"""Cleanup async resources."""
self._get_shutdown_event().set()
if self.http_client_h1 and self.http_client_h1 is not self.http_client:
await self.http_client_h1.aclose()
self.http_client_h1 = None
if self.http_client:
await self.http_client.aclose()
self.http_client = None
if self.memory_handler and hasattr(self.memory_handler, "close"):
await self.memory_handler.close()
with contextlib.suppress(Exception):
from headroom.models.ml_models import MLModelRegistry
released_models = []
released_models.extend(MLModelRegistry.unload_prefix("technique_router:"))
released_models.extend(MLModelRegistry.unload_prefix("siglip:"))
if released_models:
logger.info("Released image optimizer models: %s", ", ".join(released_models))
# Stop all quota trackers via the registry
await get_quota_registry().stop_all()
# Persist any savings the tracker's write throttle is still holding, so
# a graceful shutdown doesn't drop the last few requests' totals.
with contextlib.suppress(Exception):
self.metrics.savings_tracker.flush()
# Print final stats
self._print_summary()
def _print_summary(self):
"""Print session summary."""
m = self.metrics
logger.info("=" * 70)
logger.info("HEADROOM PROXY SESSION SUMMARY")
logger.info("=" * 70)
logger.info(f"Total requests: {m.requests_total}")
logger.info(f"Cached responses: {m.requests_cached}")
logger.info(f"Rate limited: {m.requests_rate_limited}")
logger.info(f"Failed: {m.requests_failed}")
logger.info(f"Input tokens: {m.tokens_input_total:,}")
logger.info(f"Output tokens: {m.tokens_output_total:,}")
logger.info(f"Tokens saved: {m.tokens_saved_total:,}")
# Active-compression ratio: savings as a fraction of what we
# *attempted* to compress (extracted units + tool schema),
# NOT the whole request. The full-request denominator is
# dominated by frozen prefix bytes (instructions, user msgs,
# prior turns) that we never touch — including them collapses
# the headline number even on sessions where every attempted
# compression succeeded.
attempted = getattr(m, "attempted_input_tokens_total", 0)
if attempted > 0:
# `attempted` is pre-compression; savings rate is plain
# saved / attempted.
savings_pct = (m.tokens_saved_total / attempted) * 100
logger.info(f"Active compression: {savings_pct:.1f}%")
logger.info(f" (attempted tokens: {attempted:,})")
if m.tokens_input_total > 0:
whole_request_pct = (
m.tokens_saved_total / (m.tokens_input_total + m.tokens_saved_total)
) * 100
logger.info(f"Of total wire traffic: {whole_request_pct:.2f}%")
if m.latency_count > 0:
avg_latency = m.latency_sum_ms / m.latency_count
logger.info(f"Avg latency: {avg_latency:.0f}ms")
logger.info("=" * 70)
async def _record_request_outcome(self, outcome: RequestOutcome) -> None:
"""Single funnel for per-request bookkeeping.
Thin wrapper around :func:`headroom.proxy.outcome.emit_request_outcome`
so call sites can write ``await self._record_request_outcome(outcome)``
(idiomatic) instead of ``await emit_request_outcome(self, outcome)``.
The real implementation lives in ``outcome.py`` as a free function so
test dummies and provider mixins can call it without inheriting from
``HeadroomProxy``.
See ``docs/superpowers/specs/P0-proxy-pipeline-audit.md`` for the
divergence catalog this funnel collapses.
"""
from headroom.proxy.outcome import emit_request_outcome
await emit_request_outcome(self, outcome)
async def _next_request_id(self) -> str:
"""Generate unique request ID."""
async with self._request_counter_lock:
self._request_counter += 1
return f"hr_{int(time.time())}_{self._request_counter:06d}"
def _extract_tags(self, headers: dict) -> dict[str, str]:
"""Backwards-compat wrapper around :func:`extract_tags`.
Handlers call ``extract_tags(headers)`` directly. Kept here for
any external caller still using ``proxy._extract_tags(headers)``.
"""
from headroom.proxy.helpers import extract_tags
return extract_tags(headers)
def _get_shutdown_event(self) -> asyncio.Event:
event = getattr(self, "_shutdown_event", None)
if event is None:
event = asyncio.Event()
self._shutdown_event = event
return event
async def _wait_for_retry_delay_or_shutdown(self, delay_seconds: float) -> bool:
try:
await asyncio.wait_for(self._get_shutdown_event().wait(), timeout=delay_seconds)
return True
except asyncio.TimeoutError:
return False
def _shutdown_retry_response(self, method: str, url: str) -> httpx.Response:
return httpx.Response(
503,
request=httpx.Request(method, url),
headers={"content-type": "application/json", "retry-after": "0"},
json={
"error": {
"type": "shutdown",
"message": "Proxy is shutting down; retry backoff cancelled.",
}
},
)
async def _retry_request(
self,
method: str,
url: str,
headers: dict,
body: dict,
stream: bool = False,
*,
original_body_bytes: bytes | None = None,
body_mutated: bool = True,
mutation_reasons: list[str] | None = None,
request_id: str | None = None,
forwarder_name: str = "server",
path_for_log: str | None = None,
timeout: httpx.Timeout | float | None = None,
) -> httpx.Response:
"""Make request with retry and exponential backoff.
Byte-faithful forwarding (PR-A3, fixes P0-2):
* If ``original_body_bytes`` is provided AND ``body_mutated`` is
``False``, the original bytes are forwarded verbatim. SHA-256
of upstream-received bytes equals client-sent bytes.
* Otherwise the body dict is canonically re-serialized via
``serialize_body_canonical`` (compact separators, ensure_ascii=False).
* ``HEADROOM_PROXY_PYTHON_FORWARDER_MODE=legacy_json_kwarg`` is an
explicit operator opt-in for emergency rollback to the old
``httpx ... json=body`` behavior.
The default ``body_mutated=True`` preserves backward compatibility
for callers that still pass only ``body`` (e.g. CCR continuations
construct their body from scratch, so canonical serialization is
correct and original bytes do not exist).
"""
from headroom.proxy.body_forwarding import prepare_outbound_body_bytes
from headroom.proxy.helpers import log_outbound_request
last_error = None
reasons = list(mutation_reasons or [])
outbound_bytes, source = prepare_outbound_body_bytes(
body=body,
original_body_bytes=original_body_bytes,
body_mutated=body_mutated,
)
outbound_headers = {**headers, "content-type": "application/json"}
log_outbound_request(
forwarder=forwarder_name,
method=method,
path=path_for_log or url,
body_bytes_count=len(outbound_bytes),
body_mutated=body_mutated,
mutation_reasons=reasons,
request_id=request_id,
source=source,
)
post_kwargs: dict = {"content": outbound_bytes, "headers": outbound_headers}
if timeout is not None:
post_kwargs["timeout"] = timeout
for attempt in range(self.config.retry_max_attempts):
try:
if stream:
# For streaming, we return early - retry happens at higher level
return await self.http_client.post( # type: ignore[union-attr]
url, **post_kwargs
)
else:
response = await self.http_client.post( # type: ignore[union-attr]
url, **post_kwargs
)
# Transient overloads (429 rate-limit, 529 overloaded):
# retry honoring Retry-After, but return verbatim once
# exhausted — a clean overload signal beats a synthesized 5xx
# (extends #1221 to 529, Anthropic's overloaded_error).
if response.status_code in RETRYABLE_OVERLOAD_STATUSES:
if (
not self.config.retry_enabled
or attempt >= self.config.retry_max_attempts - 1
):
return response
delay_ms = retry_after_ms(
response, self.config.retry_max_delay_ms
) or jitter_delay_ms(
self.config.retry_base_delay_ms,
self.config.retry_max_delay_ms,
attempt,
)
logger.warning(
f"Upstream {response.status_code} (attempt {attempt + 1}), "
f"retrying in {delay_ms:.0f}ms"
)
if await self._wait_for_retry_delay_or_shutdown(delay_ms / 1000):
logger.info(
"Shutdown interrupted retry backoff for %s %s",
method,
path_for_log or "<upstream-url>",
)
return self._shutdown_retry_response(method, url)
continue
# Don't retry other client errors (4xx)
if 400 <= response.status_code < 500:
return response
# Retry other server errors (5xx)
if response.status_code >= 500:
raise httpx.HTTPStatusError(
f"Server error: {response.status_code}",
request=response.request,
response=response,
)
return response
# httpx.TransportError covers ConnectError, the timeout family, and —
# crucially — the protocol errors (Local/RemoteProtocolError, e.g. an
# HTTP/2 `StreamReset`) that a poisoned shared h2 connection raises on
# every in-flight request. Retrying drops the bad connection and
# re-sends on a fresh one instead of collapsing to a 502. (#1639)
except (httpx.TransportError, httpx.HTTPStatusError) as e:
last_error = e
if not self.config.retry_enabled or attempt >= self.config.retry_max_attempts - 1:
# On exhaustion, preserve the upstream 5xx status (e.g. 503
# Service Unavailable, 500, 502, 504) so the client can apply
# its own retry/backoff. Collapsing every exhausted 5xx into a
# generic 502 hides the retryable signal and makes clients give
# up. The 429/529 overload statuses are already returned
# verbatim by the RETRYABLE_OVERLOAD_STATUSES branch above and
# never reach here. ConnectError/TimeoutException carry no
# response, so those still raise.
if isinstance(e, httpx.HTTPStatusError) and e.response is not None:
return e.response
raise
# Exponential backoff with jitter
delay_with_jitter = jitter_delay_ms(
self.config.retry_base_delay_ms,
self.config.retry_max_delay_ms,
attempt,
)
logger.warning(
f"Request failed (attempt {attempt + 1}), retrying in {delay_with_jitter:.0f}ms: {e}"
)
if await self._wait_for_retry_delay_or_shutdown(delay_with_jitter / 1000):
logger.info(
"Shutdown interrupted retry backoff for %s %s",
method,
path_for_log or "<upstream-url>",
)
return self._shutdown_retry_response(method, url)
if last_error is None:
raise RuntimeError(
"retry loop exhausted with no error recorded; retry_max_attempts must be >= 1"
)
raise last_error
async def _log_toin_stats_periodically(interval_seconds: int = 300) -> None:
"""Background task that logs TOIN stats periodically.
Args:
interval_seconds: How often to log stats (default: 5 minutes).
"""
while True:
await asyncio.sleep(interval_seconds)
try:
toin = get_toin()
stats = toin.get_stats()
total_compressions = stats.get("total_compressions", 0)
if total_compressions > 0:
patterns = stats.get("patterns_tracked", 0)
retrievals = stats.get("total_retrievals", 0)
retrieval_rate = stats.get("global_retrieval_rate", 0.0)
logger.info(
"TOIN: %d patterns, %d compressions, %d retrievals, %.1f%% retrieval rate",
patterns,
total_compressions,
retrievals,
retrieval_rate * 100,
)
except Exception as e:
logger.debug("Failed to log TOIN stats: %s", e)
def _register_memory_components(proxy: HeadroomProxy, tracker: MemoryTracker) -> None:
"""Register all memory-tracked components with the tracker.
This function is idempotent - it checks if components are already registered.
Args:
proxy: The HeadroomProxy instance.
tracker: The MemoryTracker instance.
"""
# Register compression store (global singleton)
if "compression_store" not in tracker.registered_components:
store = get_compression_store()
tracker.register("compression_store", store.get_memory_stats)
# Register semantic cache (instance on proxy)
if proxy.cache and "semantic_cache" not in tracker.registered_components:
tracker.register("semantic_cache", proxy.cache.get_memory_stats)
# Register request logger (instance on proxy)
if proxy.logger and "request_logger" not in tracker.registered_components:
tracker.register("request_logger", proxy.logger.get_memory_stats)
# Register batch context store (global singleton)
if "batch_context_store" not in tracker.registered_components:
try:
from ..ccr.batch_store import get_batch_context_store
batch_store = get_batch_context_store()
if hasattr(batch_store, "get_memory_stats"):
tracker.register("batch_context_store", batch_store.get_memory_stats)
except ImportError:
pass
# Note: graph_store and vector_index are created per-user within the
# LocalMemoryBackend, not as global singletons. They would need to be
# registered when the memory system is initialized with specific backends.
def _request_is_loopback(request: Request) -> bool:
"""Return True iff the caller is on loopback by *both* peer IP and Host header.
Mirrors the two-gate check in :func:`loopback_guard.require_loopback`
(loopback client IP + loopback ``Host`` header, the DNS-rebinding defence)
but returns a bool instead of raising. Endpoints use it to vary their
payload — serving sensitive sub-blocks (upstream URLs, per-request logs)
only to loopback callers — rather than 404ing network callers that still
have a legitimate use for the non-sensitive aggregate fields.
"""
from headroom.proxy.loopback_guard import is_loopback_host, is_loopback_host_header
client = getattr(request, "client", None)
client_host = getattr(client, "host", None) if client is not None else None
try:
host_header = request.headers.get("host")
except AttributeError:
host_header = None
# The Host-header gate is the DNS-rebinding defence and always applies.
if not is_loopback_host_header(host_header):
return False
# Genuine loopback peer (native run, or curl inside the container).
if is_loopback_host(client_host):
return True
# Containerized dashboards: when Headroom runs in a bridge-network
# container, a browser on the host reaches it via the container's
# gateway, so ``request.client.host`` is the gateway IP, not 127.0.0.1
# — and the per-request logs / upstream URLs get stripped even though
# the operator is local. Treat a peer inside an operator-configured
# trusted-gateway CIDR as loopback-equivalent. Opt-in and empty by
# default (HEADROOM_PROXY_TRUSTED_GATEWAY_CIDRS), so this is a no-op
# unless the operator explicitly allow-lists their container gateway.
from headroom.proxy.forwarded_headers import (
load_trusted_gateway_cidrs,
peer_is_trusted_gateway,
)
return peer_is_trusted_gateway(client_host, load_trusted_gateway_cidrs())
_is_known_websocket_callback_failure = is_known_websocket_callback_failure
_tool_schema_saved_from_tags = tool_schema_saved_from_tags
def create_app(config: ProxyConfig | None = None) -> FastAPI:
"""Create FastAPI application."""
if not FASTAPI_AVAILABLE:
raise ImportError("FastAPI required. Install: pip install fastapi uvicorn httpx")
from contextlib import asynccontextmanager
# Always-on file logging to ~/.headroom/logs/ for `headroom perf` analysis.
# Installed here (not at module import) so importing headroom.proxy.server
# in tests or library contexts does not silently attach a RotatingFileHandler
# to the user's live proxy.log.
_setup_file_logging()
config = config or ProxyConfig()
# Air-gap master switch. Propagate config.offline to the env so the
# env-based egress predicates (telemetry, update check, license) all honor
# it, force HF/transformers offline before any model code loads, and
# announce that every outbound path is disabled.
if config.offline:
os.environ.setdefault("HEADROOM_OFFLINE", "1")
if is_offline():
apply_offline_env()
logger.warning(
"event=proxy_offline_mode air-gap active — all outbound egress disabled "
"(telemetry, update check, license reporter, HuggingFace downloads)"
)
proxy = HeadroomProxy(config)
# cc-switch reconciler (opt-in: HEADROOM_CC_SWITCH_RECONCILE=1).
# Keeps Headroom in the request path while cc-switch overwrites
# ~/.claude/settings.json on every provider switch. See
# headroom/proxy/cc_switch_reconciler.py for the full rationale.
from headroom.proxy.cc_switch_reconciler import (
CCSwitchReconciler,
reconciler_enabled,
)
_cc_reconciler: CCSwitchReconciler | None = None
if reconciler_enabled():
_cc_proxy_port = config.port if hasattr(config, "port") else 8787
def _set_anthropic_upstream(url: str) -> None:
from headroom.providers.registry import _normalize_api_url
HeadroomProxy.ANTHROPIC_API_URL = _normalize_api_url(
url, default=DEFAULT_ANTHROPIC_API_URL
)
_cc_reconciler = CCSwitchReconciler(
proxy_url=f"http://127.0.0.1:{_cc_proxy_port}",
default_upstream=DEFAULT_ANTHROPIC_API_URL,
set_upstream=_set_anthropic_upstream,
)
# Single-worker-owner lock. With uvicorn workers > 1, each worker runs the
# lifespan independently. A file lock elects ONE owner worker so that
# single-instance background tasks (currently the cc-switch reconciler) run
# once across all workers instead of N times.
from headroom import paths as _hr_paths
_beacon_lock_path = _hr_paths.beacon_lock_path(config.port)
_beacon_lock_fd: list = [None] # mutable holder for the lock file descriptor
_beacon_is_owner: list = [False]
def _try_acquire_beacon_lock() -> bool:
"""Try to acquire the beacon file lock (non-blocking).
Returns True if this process is the beacon owner.
"""
if not HAS_FCNTL:
return True
fd = None
try:
_beacon_lock_path.parent.mkdir(parents=True, exist_ok=True)
fd = open(_beacon_lock_path, "w") # noqa: SIM115
fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
fd.write(str(os.getpid()))
fd.flush()
_beacon_lock_fd[0] = fd
return True
except OSError:
if fd is not None:
fd.close()
return False
def _release_beacon_lock() -> None:
"""Release the beacon file lock."""
fd = _beacon_lock_fd[0]
if fd:
try:
if HAS_FCNTL:
fcntl.flock(fd, fcntl.LOCK_UN)
fd.close()
except Exception:
pass
_beacon_lock_fd[0] = None
try:
_beacon_lock_path.unlink(missing_ok=True)
except Exception:
pass
@asynccontextmanager
async def lifespan(app: FastAPI): # type: ignore[no-untyped-def]
# Hotfix-A0: Rust core deployment smoke test. Refuse to accept
# traffic if the Rust extension is missing unless the operator
# explicitly opted out with HEADROOM_REQUIRE_RUST_CORE=false. See
# Finding #2 in HEADROOM_PROXY_LOG_FINDINGS_2026_05_03.md.
# `_check_rust_core` either returns ("loaded"|"disabled", _) or
# calls `sys.exit(78)` — execution past this line implies the
# rust_core_status is recorded.
_rust_core_status, _rust_core_error = _check_rust_core()
app.state.rust_core_status = _rust_core_status
app.state.rust_core_error = _rust_core_error
configure_otel_metrics(OTelMetricsConfig.from_env(default_service_name="headroom-proxy"))
configure_langfuse_tracing(
LangfuseTracingConfig.from_env(default_service_name="headroom-proxy")
)
app.state.started_at = time.time()
app.state.ready = False
app.state.startup_error = None
await initialize_context_tool_session_baseline()
try:
try:
previous_handler = _install_loop_exception_handler()
# Startup
await proxy.startup()
if config.periodic_toin_stats_enabled:
asyncio.create_task(_log_toin_stats_periodically())
if proxy.usage_reporter:
await proxy.usage_reporter.start(proxy)
if proxy.traffic_learner:
await proxy.traffic_learner.start()
if proxy._background_compression_enabled:
await proxy._background_compressor.start()
# Elect the single owner worker (first worker wins the lock).
_beacon_is_owner[0] = _try_acquire_beacon_lock()
# Only the owner worker runs the reconciler. With uvicorn
# workers > 1 each worker runs this lifespan; without this guard
# every worker would watch + rewrite settings.json concurrently
# and each process would hold its own HeadroomProxy.ANTHROPIC_API_URL,
# so workers could disagree on the upstream.
if _cc_reconciler is not None and _beacon_is_owner[0]:
await _cc_reconciler.start()
app.state.ready = True
yield
except Exception as exc:
app.state.startup_error = str(exc)
raise
finally:
loop: asyncio.AbstractEventLoop | None
previous: LoopExceptionHandler | None
try:
loop = asyncio.get_running_loop()
previous = previous_handler
except RuntimeError:
loop = None
previous = app.state.previous_loop_exception_handler
if loop is not None:
loop.set_exception_handler(previous)
app.state.ready = False
# Shutdown
if _cc_reconciler is not None:
await _cc_reconciler.stop()
if _beacon_is_owner[0]:
_release_beacon_lock()
if proxy.usage_reporter:
await proxy.usage_reporter.stop()
if proxy.traffic_learner:
await proxy.traffic_learner.stop()
if proxy._background_compression_enabled:
await proxy._background_compressor.stop()
proxy._background_compression_executor.shutdown(wait=False)
if proxy.code_graph_watcher:
proxy.code_graph_watcher.stop()
await proxy.shutdown()
shutdown_headroom_tracing()
shutdown_otel_metrics()
app = FastAPI(
title="Headroom Proxy",
description="Production-ready LLM optimization proxy",
version=__version__,
lifespan=lifespan,
)
loop_health_state: LoopHealthState = {
"status": "healthy",
"known_failures": 0,
"last_known_failure": None,
}
app.state.proxy = proxy
app.state.started_at = None
app.state.ready = False
app.state.startup_error = None
app.state.loop_callback_health = loop_health_state
app.state.loop_exception_handler = None
app.state.previous_loop_exception_handler = None
# Set by the lifespan startup smoke test (`_check_rust_core`). Default
# "missing" means lifespan hasn't run yet — anything reading /health
# before startup completes (rare; lifespan runs before the first
# request) sees an honest "missing" rather than a stale "loaded".
app.state.rust_core_status = "missing"
app.state.rust_core_error = None
def _iso_utc_now() -> str:
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
def _uptime_seconds() -> float:
started_at = getattr(app.state, "started_at", None)
if not isinstance(started_at, int | float):
return 0.0
return round(max(0.0, time.time() - float(started_at)), 3)
def _component_health(
*,
enabled: bool,
ready: bool,
**details: Any,
) -> dict[str, Any]:
status = "disabled" if not enabled else ("healthy" if ready else "unhealthy")
return {
"enabled": enabled,
"ready": (ready if enabled else True),
"status": status,
**details,
}
def _health_checks() -> dict[str, dict[str, Any]]:
memory_status = (
proxy.memory_handler.health_status()
if proxy.memory_handler
else {
"enabled": False,
"backend": None,
"initialized": False,
"native_tool": False,
"bridge_enabled": False,
}
)
memory_enabled = bool(memory_status.get("enabled", False))
memory_initialized = bool(memory_status.get("initialized", False))
return {
"startup": _component_health(
enabled=True,
ready=bool(getattr(app.state, "ready", False)),
error=getattr(app.state, "startup_error", None),
),
"http_client": _component_health(
enabled=True,
ready=proxy.http_client is not None,
),
"cache": _component_health(
enabled=config.cache_enabled,
ready=(proxy.cache is not None),
),
"rate_limiter": _component_health(
enabled=config.rate_limit_enabled,
ready=(proxy.rate_limiter is not None),
),
"memory": _component_health(
enabled=memory_enabled,
ready=memory_initialized,
backend=memory_status["backend"],
initialized=memory_initialized,
native_tool=bool(memory_status.get("native_tool", False)),
bridge_enabled=bool(memory_status.get("bridge_enabled", False)),
),
"upstream": _component_health(
enabled=os.environ.get("HEADROOM_SKIP_UPSTREAM_CHECK", "").strip() != "1",
ready=bool(_upstream_check_cache["ok"]),
url=_upstream_check_cache["url"],
error=_upstream_check_cache["error"],
),
"kompress": _component_health(
enabled=not config.disable_kompress,
ready=proxy.warmup.kompress.status == "loaded",
backend=proxy.warmup.kompress.info.get("backend", None),
),
}
def _runtime_payload() -> dict[str, Any]:
ws_registry = getattr(proxy, "ws_sessions", None)
ws_active_sessions = ws_registry.active_count() if ws_registry is not None else 0
ws_active_relay_tasks = (
ws_registry.active_relay_task_count() if ws_registry is not None else 0
)
# Snapshot compression executor metrics under their lock (gauges
# mutated by worker threads; not safe to read without).
with proxy._compression_metrics_lock:
_comp_queued = proxy._compression_queued
_comp_queued_max = proxy._compression_queued_max
_comp_queue_timeouts = proxy._compression_queue_timeouts
_comp_queue_wait_total = proxy._compression_queue_wait_seconds_total
_comp_queue_wait_max = proxy._compression_queue_wait_seconds_max
_comp_in_flight = proxy._compression_in_flight
_comp_in_flight_max = proxy._compression_in_flight_max
_comp_run_total = proxy._compression_run_seconds_total
_comp_run_max = proxy._compression_run_seconds_max
_comp_leaked = proxy._compression_leaked_threads
return {
"anthropic_pre_upstream": {
"enabled": proxy.anthropic_pre_upstream_sem is not None,
"resolved_concurrency": proxy.anthropic_pre_upstream_concurrency,
"source": (
"auto" if config.anthropic_pre_upstream_concurrency is None else "explicit"
),
"acquire_timeout_seconds": proxy.anthropic_pre_upstream_acquire_timeout_seconds,
"compression_timeout_seconds": float(COMPRESSION_TIMEOUT_SECONDS),
"memory_context_timeout_seconds": (
proxy.anthropic_pre_upstream_memory_context_timeout_seconds
),
"codex_ws_gated": False,
},
"compression_executor": {
"max_workers": proxy.compression_max_workers,
"queued": _comp_queued,
"queued_max": _comp_queued_max,
"queue_timeouts_total": _comp_queue_timeouts,
"queue_wait_seconds_total": _comp_queue_wait_total,
"queue_wait_seconds_max": _comp_queue_wait_max,
"running": _comp_in_flight,
"in_flight": _comp_in_flight,
"in_flight_max": _comp_in_flight_max,
"run_seconds_total": _comp_run_total,
"run_seconds_max": _comp_run_max,
"leaked_threads_total": _comp_leaked,
"source": ("auto" if config.compression_max_workers is None else "explicit"),
},
"websocket_sessions": {
"active_sessions": ws_active_sessions,
"active_relay_tasks": ws_active_relay_tasks,
},
}
def _loop_callback_payload() -> LoopHealthState:
return {
"status": loop_health_state["status"],
"known_failures": loop_health_state["known_failures"],
"last_known_failure": loop_health_state["last_known_failure"],
}
def _install_loop_exception_handler() -> LoopExceptionHandler | None:
try:
loop = asyncio.get_running_loop()
except RuntimeError:
return None
previous_handler = loop.get_exception_handler()
def _loop_exception_handler(
_loop: asyncio.AbstractEventLoop, context: dict[str, Any]
) -> None:
if _is_known_websocket_callback_failure(context):
loop_health_state["status"] = "unhealthy"
loop_health_state["known_failures"] += 1
loop_health_state["last_known_failure"] = {
"message": context.get("message"),
"exception": str(context.get("exception"))
if context.get("exception")
else None,
}
return
delegate_handler = app.state.previous_loop_exception_handler
if delegate_handler is not None:
delegate_handler(_loop, context)
return
_loop.default_exception_handler(context)
loop.set_exception_handler(_loop_exception_handler)
app.state.loop_exception_handler = _loop_exception_handler
app.state.previous_loop_exception_handler = previous_handler
return previous_handler
def _health_payload(*, include_config: bool) -> dict[str, Any]:
checks = _health_checks()
# Kompress is an optional soft component: model downloads lazily on
# first use, so "not ready" (cold cache) must not degrade overall health.
ready = all(check["ready"] for name, check in checks.items() if name != "kompress")
payload: dict[str, Any] = {
"service": "headroom-proxy",
"status": "healthy" if ready else "unhealthy",
"ready": ready,
"version": __version__,
"timestamp": _iso_utc_now(),
"uptime_seconds": _uptime_seconds(),
"checks": checks,
"runtime": _runtime_payload(),
# Hotfix-A0: surface rust core load state so operators can alert
# on `rust_core != "loaded"` (Finding #2).
"rust_core": getattr(app.state, "rust_core_status", "missing"),
}
rust_core_error = getattr(app.state, "rust_core_error", None)
if rust_core_error:
payload["rust_core_error"] = rust_core_error
deployment_profile = os.environ.get("HEADROOM_DEPLOYMENT_PROFILE")
if deployment_profile:
payload["deployment"] = {
"profile": deployment_profile,
"preset": os.environ.get("HEADROOM_DEPLOYMENT_PRESET"),
"runtime": os.environ.get("HEADROOM_DEPLOYMENT_RUNTIME"),
"supervisor": os.environ.get("HEADROOM_DEPLOYMENT_SUPERVISOR"),
"scope": os.environ.get("HEADROOM_DEPLOYMENT_SCOPE"),
}
if include_config:
profile_kwargs = proxy_pipeline_kwargs(config)
effective_target_ratio = cast(
float | None,
profile_kwargs.get("target_ratio", config.target_ratio),
)
payload["config"] = {
"backend": config.backend,
"optimize": config.optimize,
"cache": config.cache_enabled,
"rate_limit": config.rate_limit_enabled,
"disable_kompress": config.disable_kompress,
"disable_kompress_fallback": config.disable_kompress_fallback,
"disable_kompress_anthropic": config.disable_kompress_anthropic,
"disable_kompress_openai": config.disable_kompress_openai,
"memory": config.memory_enabled,
"learn": config.traffic_learning_enabled,
"code_graph": config.code_graph_watcher,
"anthropic_api_url": config.anthropic_api_url,
"openai_api_url": config.openai_api_url,
"gemini_api_url": config.gemini_api_url,
"cloudcode_api_url": config.cloudcode_api_url,
"vertex_api_url": config.vertex_api_url,
"savings_profile": config.savings_profile,
"target_ratio": effective_target_ratio,
"target_savings_percent": (
round(max(0.0, min(1.0, 1.0 - float(effective_target_ratio))) * 100, 1)
if effective_target_ratio is not None
else None
),
"compress_user_messages": bool(
profile_kwargs.get("compress_user_messages", config.compress_user_messages)
),
"compress_system_messages": bool(
profile_kwargs.get(
"compress_system_messages",
config.compress_system_messages,
)
),
"protect_recent": profile_kwargs.get(
"read_protection_window",
config.protect_recent,
),
"protect_analysis_context": profile_kwargs.get(
"protect_analysis_context",
config.protect_analysis_context,
),
"min_tokens_to_crush": profile_kwargs.get(
"min_tokens_to_compress",
config.min_tokens_to_crush,
),
"max_items_after_crush": profile_kwargs.get(
"max_items_after_crush",
config.max_items_after_crush,
),
"smart_crusher_with_compaction": profile_kwargs.get(
"smart_crusher_with_compaction",
config.smart_crusher_with_compaction,
),
"force_kompress": bool(profile_kwargs.get("force_kompress", False)),
"accuracy_guard": config.accuracy_guard,
# Live (per-request) env knobs the proxy reads after startup.
# Surfaced so `headroom wrap` can see what a reused proxy is
# actually using and hot-sync it via /admin/runtime-env.
"runtime_env": runtime_env.effective_runtime_env(),
"pid": os.getpid(),
}
return payload
# ---------------------------------------------------------------------------
# Upstream connectivity check — cached to avoid hammering the upstream on
# every /readyz poll. Set HEADROOM_SKIP_UPSTREAM_CHECK=1 to opt out (e.g.
# in air-gapped or test environments where the upstream isn't reachable at
# startup time).
# ---------------------------------------------------------------------------
_UPSTREAM_CHECK_TTL = 30.0 # seconds
_upstream_check_cache: dict[str, Any] = {
"expires_at": 0.0,
"ok": True,
"error": None,
"url": None,
}
_upstream_check_lock = asyncio.Lock()
def _upstream_target_url() -> str:
"""Return the primary upstream base URL to probe."""
# Use the resolved API target from the provider runtime so we respect
# any overrides set by ProxyConfig.anthropic_api_url / env vars.
return proxy.provider_runtime.api_targets.anthropic
async def _check_upstream() -> None:
"""Probe the upstream API endpoint and update the cached result.
Uses a HEAD request with a 5-second timeout — just enough to verify
TLS + TCP reachability without triggering an inference call.
"""
if os.environ.get("HEADROOM_SKIP_UPSTREAM_CHECK", "").strip() == "1":
# Opt-out: treat upstream as always reachable.
_upstream_check_cache["ok"] = True
_upstream_check_cache["error"] = None
_upstream_check_cache["expires_at"] = time.monotonic() + _UPSTREAM_CHECK_TTL
return
now = time.monotonic()
# Fast-path: return if the cached result is still fresh (no lock needed
# for a simple float comparison — worst case we re-check twice).
if now < _upstream_check_cache["expires_at"]:
return
async with _upstream_check_lock:
# Re-check inside the lock to handle concurrent waiters.
if time.monotonic() < _upstream_check_cache["expires_at"]:
return
url = _upstream_target_url()
_upstream_check_cache["url"] = url
client = proxy.http_client
if client is None:
_upstream_check_cache["ok"] = False
_upstream_check_cache["error"] = "proxy client not initialised"
_upstream_check_cache["expires_at"] = time.monotonic() + _UPSTREAM_CHECK_TTL
return
try:
resp = await client.head(url, timeout=5.0)
# Any HTTP response (even 4xx/5xx) means TLS+TCP worked.
_ = resp.status_code
_upstream_check_cache["ok"] = True
_upstream_check_cache["error"] = None
except Exception as exc: # noqa: BLE001
_upstream_check_cache["ok"] = False
_upstream_check_cache["error"] = str(exc)
_upstream_check_cache["expires_at"] = time.monotonic() + _UPSTREAM_CHECK_TTL
# CORS: scoped to localhost by default. The old wildcard origin combined
# with allow_credentials=True let any web page the user had open read the
# proxy's content endpoints (e.g. /v1/retrieve returns raw, uncompressed
# tool outputs) via a cross-origin fetch to 127.0.0.1 (CWE-346).
#
# The default matches any loopback origin on any port via a regex, so it
# works regardless of the --port the proxy was started on without the app
# needing to know its own bound port (the port lives in the CLI/uvicorn
# layer, not in ProxyConfig). Set HEADROOM_CORS_ORIGINS (comma-separated)
# to pin an explicit allowlist for Docker or remote-dashboard deployments;
# "*" restores the old wildcard behaviour if the operator accepts the risk.
_default_loopback_origin_regex = r"https?://(localhost|127\.0\.0\.1|\[::1\])(:\d+)?"
_cors_origins_env = os.environ.get("HEADROOM_CORS_ORIGINS", "").strip()
if _cors_origins_env:
_cors_allow_origins = [o.strip() for o in _cors_origins_env.split(",") if o.strip()]
_cors_allow_origin_regex: str | None = None
else:
_cors_allow_origins = []
_cors_allow_origin_regex = _default_loopback_origin_regex
app.add_middleware(
CORSMiddleware,
allow_origins=_cors_allow_origins,
allow_origin_regex=_cors_allow_origin_regex,
allow_credentials=False,
allow_methods=["GET", "POST"],
allow_headers=["Content-Type", "Authorization"],
)
# X-Headroom-Stack: SDK adapters (TS openai/anthropic/etc.) tag their
# requests so telemetry can segment by integration surface. Registered
# before extension middleware so any extension-level auth/guards run
# outermost and we don't count requests they reject.
@app.middleware("http")
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)
# Path-based Codex identification: stamp X-Client: codex on the
# Responses endpoint for callers that don't otherwise classify (e.g.
# Codex Desktop, whose User-Agent isn't a known codex UA). Without it
# the backend refuses oversized
# requests with a 413 on a compression timeout, which Codex treats as a
# hard connection failure. Mutating scope["headers"] before call_next
# makes every downstream classify_client(headers) read "codex".
if should_stamp_codex_client(path, headers):
request.scope["headers"].append((b"x-client", b"codex"))
client = getattr(request, "client", None)
client_addr = ""
if client is not None:
client_host = getattr(client, "host", None)
client_port = getattr(client, "port", None)
client_addr = f"{client_host}:{client_port}" if client_port else str(client_host)
try:
proxy.metrics.record_inbound_request(method=method, path=path)
except Exception:
logger.debug("record_inbound_request failed", exc_info=True)
try:
from headroom.proxy.helpers import redact_for_wire_debug
safe_headers = redact_for_wire_debug(headers)
except Exception:
safe_headers = {"redaction_error": True}
logger.info(
"event=proxy_inbound_request id=%s method=%s path=%s query=%s client=%s "
"content_length=%s headers=%s",
inbound_id,
method,
path,
query,
client_addr,
request.headers.get("content-length", ""),
json.dumps(safe_headers, ensure_ascii=False, default=str),
)
if request.url.path.startswith("/v1/"):
stack = request.headers.get("x-headroom-stack")
if stack:
try:
proxy.metrics.record_stack(stack)
except Exception:
logger.debug("record_stack failed", exc_info=True)
try:
response = await call_next(request)
except asyncio.CancelledError:
try:
proxy.metrics.record_inbound_aborted(reason="cancelled")
except Exception:
logger.debug("record_inbound_aborted failed", exc_info=True)
logger.info(
"event=proxy_inbound_request_aborted id=%s method=%s path=%s reason=cancelled "
"duration_ms=%.2f",
inbound_id,
method,
path,
(time.perf_counter() - started) * 1000.0,
)
raise
except Exception as exc:
try:
proxy.metrics.record_inbound_aborted(reason=type(exc).__name__)
except Exception:
logger.debug("record_inbound_aborted failed", exc_info=True)
logger.error(
"event=proxy_inbound_request_aborted id=%s method=%s path=%s reason=%s "
"duration_ms=%.2f",
inbound_id,
method,
path,
type(exc).__name__,
(time.perf_counter() - started) * 1000.0,
exc_info=True,
)
raise
try:
proxy.metrics.record_inbound_response(status_code=response.status_code)
except Exception:
logger.debug("record_inbound_response failed", exc_info=True)
logger.info(
"event=proxy_inbound_response id=%s method=%s path=%s status=%s duration_ms=%.2f",
inbound_id,
method,
path,
response.status_code,
(time.perf_counter() - started) * 1000.0,
)
return response
# ── Security gate (registered last → runs outermost) ──────────────────
# Three concerns, kept together because they all wrap every inbound
# request: optional inbound auth on the data plane, response security
# headers, and an audit trail for state-mutating admin endpoints.
_proxy_token = config.proxy_token or os.environ.get("HEADROOM_PROXY_TOKEN") or None
# Pre-encode once for constant-time comparison (compare_digest on str raises
# TypeError for non-ASCII input, which would turn a 401 into a 500).
_proxy_token_bytes = _proxy_token.encode("utf-8") if _proxy_token else b""
# Health/readiness probes must stay reachable without a token so
# orchestrators can check a container that binds non-loopback.
_AUTH_EXEMPT_PATHS = frozenset({"/health", "/healthz", "/livez", "/readyz"})
# Loud warning when a non-loopback bind has no token configured: that is the
# exact shape (e.g. the Docker 0.0.0.0 image) that exposes unauthenticated
# /v1/* routes to the surrounding network.
if not _proxy_token and not is_loopback_host(getattr(config, "host", None)):
logger.warning(
"event=proxy_open_bind host=%s — proxy is bound to a non-loopback "
"interface with no HEADROOM_PROXY_TOKEN set; the /v1/* data-plane "
"routes are reachable WITHOUT authentication. Set HEADROOM_PROXY_TOKEN "
"to require a bearer token from non-loopback callers.",
getattr(config, "host", None),
)
def _apply_security_headers(response) -> None:
# setdefault: never clobber a header an upstream/handler already set.
response.headers.setdefault("X-Content-Type-Options", "nosniff")
response.headers.setdefault("X-Frame-Options", "DENY")
response.headers.setdefault("Referrer-Policy", "no-referrer")
response.headers.setdefault(
"Strict-Transport-Security", "max-age=31536000; includeSubDomains"
)
def _extract_proxy_token(headers) -> str | None:
auth = str(headers.get("authorization") or "")
if auth.lower().startswith("bearer "):
return auth[7:].strip() or None
raw = headers.get("x-headroom-proxy-token")
return str(raw) if raw else None
@app.middleware("http")
async def _security_gate(request, call_next):
# 1) Optional inbound auth. When a token is configured, require it on
# non-loopback requests; loopback callers and health probes are
# exempt. Loopback is the same trust boundary the admin/debug
# endpoints already use (see loopback_guard).
if _proxy_token:
path = request.url.path
client = getattr(request, "client", None)
client_host = getattr(client, "host", None) if client is not None else None
if path not in _AUTH_EXEMPT_PATHS and not is_loopback_host(client_host):
provided = _extract_proxy_token(request.headers)
if provided is None or not hmac.compare_digest(
provided.encode("utf-8", "replace"), _proxy_token_bytes
):
logger.warning(
"event=proxy_auth_rejected path=%s client=%s reason=%s",
path,
client_host,
"missing_token" if provided is None else "bad_token",
)
rejection = JSONResponse(status_code=401, content={"error": "unauthorized"})
_apply_security_headers(rejection)
return rejection
response = await call_next(request)
_apply_security_headers(response)
# 2) Audit trail for admin / state-mutating endpoints.
try:
if is_auditable_path(request.url.path):
record_admin_action(
request=request,
action="admin_request",
status_code=response.status_code,
)
except Exception:
logger.debug("admin audit emission failed", exc_info=True)
return response
# Third-party proxy extensions (Enterprise, custom plugins). Discovered via
# the `headroom.proxy_extension` entry-point group, but **opt-in only**:
# only names listed in config.proxy_extensions (CLI: --proxy-extension,
# env: HEADROOM_PROXY_EXTENSIONS) actually get installed. Discovery alone
# never runs third-party code. An extension that raises from its install()
# is a deliberate fail-closed signal and aborts startup.
from headroom.proxy.extensions import install_all as _install_extensions
_install_extensions(app, config, enabled=getattr(config, "proxy_extensions", None))
# Health & Metrics
@app.get("/livez")
async def livez():
callback_state = _loop_callback_payload()
healthy = callback_state["status"] == "healthy"
return JSONResponse(
status_code=200 if healthy else 503,
content={
"service": "headroom-proxy",
"status": "healthy" if healthy else "unhealthy",
"alive": healthy,
"version": __version__,
"timestamp": _iso_utc_now(),
"uptime_seconds": _uptime_seconds(),
"loop_health": callback_state,
},
)
@app.get("/readyz")
async def readyz():
await _check_upstream()
payload = _health_payload(include_config=False)
return JSONResponse(status_code=200 if payload["ready"] else 503, content=payload)
@app.get("/health")
async def health(request: Request):
await _check_upstream()
# /health echoes upstream API URLs + backend config (the `config`
# block). That is operational detail an external scanner should not
# see, so include it only for loopback callers; network callers get the
# same body as /readyz (status + checks, no config). /livez and /readyz
# remain the unauthenticated probes for orchestration health.
payload = _health_payload(include_config=_request_is_loopback(request))
return JSONResponse(status_code=200, content=payload)
# Loopback-only debug introspection (Unit 5). A remote IP gets 404 —
# debug endpoints are invisible to external scanners.
from headroom.proxy.debug_introspection import (
collect_tasks as _collect_tasks,
)
from headroom.proxy.loopback_guard import require_loopback as _require_loopback
@app.get("/admin/upstream", dependencies=[Depends(_require_loopback)])
async def get_upstream():
"""Current Anthropic upstream + cc-switch reconciler state (loopback-only).
Read-only. The upstream is mutated only via the in-process cc-switch
reconciler (driven by ~/.claude/settings.json) — there is deliberately
no HTTP write route, so a local process cannot redirect credential-
bearing traffic to an arbitrary URL through this surface.
"""
return {
"anthropic": HeadroomProxy.ANTHROPIC_API_URL,
"cc_switch_reconcile": _cc_reconciler is not None,
"captured_upstream": getattr(_cc_reconciler, "current_upstream", None),
}
@app.get("/debug/tasks", dependencies=[Depends(_require_loopback)])
async def debug_tasks(stack: bool = False):
"""Enumerate running asyncio tasks.
Default is cheap — ``stack_depth`` is ``null`` in every entry so
a storm snapshot does not walk 50+ coroutine frames synchronously.
Pass ``?stack=true`` to compute ``stack_depth`` for each task
(useful for single-shot human debugging).
"""
ws_registry = getattr(proxy, "ws_sessions", None)
return JSONResponse(
status_code=200,
content=_collect_tasks(ws_registry, with_stack_depth=stack),
)
@app.get("/debug/ws-sessions", dependencies=[Depends(_require_loopback)])
async def debug_ws_sessions():
ws_registry = getattr(proxy, "ws_sessions", None)
snapshot = ws_registry.snapshot() if ws_registry is not None else []
return JSONResponse(status_code=200, content=snapshot)
@app.get("/debug/warmup", dependencies=[Depends(_require_loopback)])
async def debug_warmup():
warmup_registry = getattr(proxy, "warmup", None)
payload = warmup_registry.to_dict() if warmup_registry is not None else {}
payload["runtime"] = _runtime_payload()
return JSONResponse(status_code=200, content=payload)
@app.post("/admin/runtime-env", dependencies=[Depends(_require_loopback)])
async def admin_runtime_env(request: Request):
"""Hot-reload live env knobs (the output-shaper family, the ast-grep
read threshold) without restarting the proxy.
Live knobs are read from the proxy's *process* environment, so a proxy
that ``headroom wrap`` reused — rather than started — never sees values
a user exported afterwards. Instead of a disruptive restart (cold ML
load, dropped requests, lost caches), ``wrap`` POSTs the values here and
the proxy applies them in memory, effective on the next request.
Loopback-only. The body is a flat ``{ENV_NAME: "value"}`` map; unknown
keys and non-string values are ignored. Returns what was applied plus
the resulting live config. Last writer wins (overrides are global to the
proxy, which is inherent — every wrapper shares one process).
"""
try:
body = await request.json()
except (ValueError, UnicodeDecodeError):
body = None
if not isinstance(body, dict):
return JSONResponse(
status_code=400,
content={"error": "expected a JSON object of {ENV_NAME: value}"},
)
applied = runtime_env.set_overrides(body)
if applied:
logger.info("runtime-env hot-reload applied: %s", sorted(applied))
# Record which runtime-env keys changed (the "what" of a config
# change) in addition to the generic admin-request audit emitted by
# the security middleware. Values are intentionally omitted — keys
# alone avoid logging any secret values that were set.
record_admin_action(
request=request,
action="runtime_env_update",
status_code=200,
details={"changed_keys": sorted(applied)},
)
return JSONResponse(
status_code=200,
content={"applied": applied, "runtime_env": runtime_env.effective_runtime_env()},
)
@app.get("/dashboard", response_class=HTMLResponse)
async def dashboard():
"""Serve the Headroom dashboard UI."""
return get_dashboard_html()
@app.get("/favicon.ico")
async def favicon() -> Response:
# Registered before register_provider_routes' catch-all passthrough
# route so browsers' automatic favicon requests for /dashboard are
# answered locally instead of being tunneled to the wrapped upstream
# provider (GH #1787).
return Response(status_code=204)
DASHBOARD_STATS_CACHE_TTL_SECONDS = 5.0
_stats_snapshot_lock = asyncio.Lock()
_stats_snapshot: dict[str, Any] = {"expires_at": 0.0, "value": None}
THROUGHPUT_CACHE_TTL_SECONDS = 10.0
_throughput_cache_lock = asyncio.Lock()
_throughput_cache: dict[str, Any] = {"expires_at": 0.0, "value": None}
RECENT_REQUEST_LOG_WINDOW = 100
def _build_recent_request_payload(limit: int = RECENT_REQUEST_LOG_WINDOW) -> dict[str, Any]:
recent_request_logs = proxy.logger.get_recent(limit) if proxy.logger else []
dashboard_recent_requests = [
{
"request_id": log.get("request_id"),
"timestamp": log.get("timestamp"),
"provider": log.get("provider"),
"model": log.get("model"),
"input_tokens_original": log.get("input_tokens_original"),
"input_tokens_optimized": log.get("input_tokens_optimized"),
"output_tokens": log.get("output_tokens"),
"tokens_saved": log.get("tokens_saved"),
"savings_percent": log.get("savings_percent"),
"optimization_latency_ms": log.get("optimization_latency_ms"),
"total_latency_ms": log.get("total_latency_ms"),
"transforms_applied": log.get("transforms_applied", []),
"waste_signals": log.get("waste_signals"),
"tool_schema_saved_tokens": _tool_schema_saved_from_tags(log.get("tags")),
}
for log in recent_request_logs
if log.get("input_tokens_original") is not None
and log.get("input_tokens_optimized") is not None
][-10:]
return {
"request_logs": recent_request_logs[-10:],
"recent_requests": dashboard_recent_requests,
}
async def _build_stats_payload() -> dict[str, Any]:
"""Build the full `/stats` response payload.
This is the main stats endpoint - it aggregates data from all subsystems:
- Request metrics (total, cached, failed, by model/provider)
- Token usage and savings
- Cost tracking
- Canonical persisted display_session metrics for downstream dashboards
- Compression (CCR) statistics
- Telemetry/TOIN (data flywheel) statistics
- Cache and rate limiter stats
"""
m = proxy.metrics
import time
async with _throughput_cache_lock:
now = time.time()
if _throughput_cache["expires_at"] < now or _throughput_cache["value"] is None:
def _compute_throughput():
from headroom.perf.analyzer import build_perf_summary, parse_log_files
perf_report = parse_log_files(last_n_hours=1.0)
return build_perf_summary(perf_report).get("throughput")
try:
throughput = await asyncio.to_thread(_compute_throughput)
_throughput_cache["value"] = throughput
_throughput_cache["expires_at"] = now + THROUGHPUT_CACHE_TTL_SECONDS
except Exception as e:
logger.warning("Failed to calculate throughput for stats: %s", e, exc_info=True)
if _throughput_cache["value"] is None:
_throughput_cache["value"] = None
throughput = _throughput_cache["value"]
# Calculate average latency
avg_latency_ms = round(m.latency_sum_ms / m.latency_count, 2) if m.latency_count > 0 else 0
min_latency_ms = (
round(m.latency_min_ms, 2)
if m.latency_count > 0 and m.latency_min_ms != float("inf")
else 0
)
max_latency_ms = round(m.latency_max_ms, 2) if m.latency_count > 0 else 0
# Calculate Headroom overhead (optimization time only, excludes pass-through requests)
avg_overhead_ms = (
round(m.overhead_sum_ms / m.overhead_count, 2) if m.overhead_count > 0 else 0
)
min_overhead_ms = (
round(m.overhead_min_ms, 2)
if m.overhead_count > 0 and m.overhead_min_ms != float("inf")
else 0
)
max_overhead_ms = round(m.overhead_max_ms, 2) if m.overhead_count > 0 else 0
# Calculate TTFB (time to first byte)
avg_ttfb_ms = round(m.ttfb_sum_ms / m.ttfb_count, 2) if m.ttfb_count > 0 else 0
min_ttfb_ms = (
round(m.ttfb_min_ms, 2) if m.ttfb_count > 0 and m.ttfb_min_ms != float("inf") else 0
)
max_ttfb_ms = round(m.ttfb_max_ms, 2) if m.ttfb_count > 0 else 0
def _pct(part: int | float, whole: int | float) -> float:
return round((float(part) / float(whole)) * 100.0, 2) if whole else 0.0
# Get compression store stats
store = get_compression_store()
compression_stats = store.get_stats()
# Get telemetry/TOIN stats
telemetry = get_telemetry_collector()
telemetry_stats = telemetry.get_stats()
# Get feedback loop stats
feedback = get_compression_feedback()
feedback_stats = feedback.get_stats()
# Build prefix cache stats once (used in both prefix_cache and cost)
prefix_cache_stats = _build_prefix_cache_stats(m, proxy.cost_tracker)
# Fetch CLI filtering savings from the selected context tool. These
# tokens are avoided before they reach model context.
cli_filtering_stats = await asyncio.to_thread(_get_context_tool_stats)
cli_filtering_tool = (
str(cli_filtering_stats.get("tool", "rtk")) if cli_filtering_stats else "rtk"
)
cli_filtering_label = (
str(cli_filtering_stats.get("label", "RTK")) if cli_filtering_stats else "RTK"
)
cli_tokens_avoided = (
cli_filtering_stats.get("tokens_saved", 0) if cli_filtering_stats else 0
)
cli_filtering_session = (
cli_filtering_stats.get("session", {}) if cli_filtering_stats else {}
)
cli_filtering_lifetime = (
cli_filtering_stats.get("lifetime", {}) if cli_filtering_stats else {}
)
rtk_tokens_avoided = cli_tokens_avoided if cli_filtering_tool == "rtk" else 0
lean_ctx_tokens_avoided = cli_tokens_avoided if cli_filtering_tool == "lean-ctx" else 0
cli_filtering_available = bool(
cli_filtering_stats and cli_filtering_stats.get("installed", False)
)
# Calculate total tokens before Headroom-side reduction. Proxy
# compression and the configured context tool both remove tokens before
# they reach model context, so dashboard-facing savings combines them.
proxy_compression_tokens = m.tokens_saved_total
all_layers_tokens_saved = proxy_compression_tokens + cli_tokens_avoided
total_tokens_before = m.tokens_input_total + all_layers_tokens_saved
proxy_total_before_compression = m.tokens_input_total + proxy_compression_tokens
# `attempted_input_tokens` is the compressible-only denominator
# (extracted units + tool schema). The "active compression"
# ratio is what fraction of the tokens we *tried* to compress
# actually got compressed. Excludes prefix-frozen content
# (user/system messages, prior turns) we never touched —
# otherwise the ratio is dominated by content we deliberately
# avoided changing for prefix-cache safety.
# `attempted_input_tokens_total` is already pre-compression: it
# accumulates `unit.tokens_before` for each eligible unit that
# reached the router, plus the original (pre-compaction) tool
# schema size. So the savings rate is plain `saved / attempted`
# — adding `saved` again would double-count.
attempted_input_tokens = getattr(m, "attempted_input_tokens_total", 0)
# New-content denominator: what the provider actually billed as
# non-cache-read input (uncached + cache-write tokens, summed
# across providers from response usage). Unlike
# `proxy_total_before_compression`, this does NOT recount the
# full transcript on every turn — a long session's history is
# served from prefix cache, not re-billed, so it doesn't belong
# in a denominator that claims to measure what compression had
# any power over. Tokens Headroom removed never reached the
# provider at all, so they're added back to form the baseline.
_pc_totals = prefix_cache_stats.get("totals", {})
new_input_tokens = int(_pc_totals.get("uncached_input_tokens", 0) or 0) + int(
_pc_totals.get("cache_write_tokens", 0) or 0
)
# Build human-readable summary
summary = _build_session_summary(
proxy, m, prefix_cache_stats, cli_tokens_avoided, total_tokens_before
)
# DEBUG: log the summary payload for external upsert consumers
try:
logger.debug("/stats summary data: %r", summary)
except Exception:
logger.warning("Failed to log /stats summary payload")
# Compression cache stats (token mode). Snapshot the cache list under
# the dict lock so a concurrent eviction can't mutate the dict while
# we iterate. Each per-session `get_stats()` is independently
# thread-safe via the cache's own internal lock.
compression_cache_stats: dict = {}
if proxy.config.mode == PROXY_MODE_TOKEN and proxy._compression_caches:
with proxy._compression_caches_lock:
_caches_snapshot = list(proxy._compression_caches.values())
_active_sessions = len(proxy._compression_caches)
total_entries = 0
total_hits = 0
total_misses = 0
total_tokens_saved = 0
for cache in _caches_snapshot:
s = cache.get_stats()
total_entries += s.get("entries", 0)
total_hits += s.get("hits", 0)
total_misses += s.get("misses", 0)
total_tokens_saved += s.get("total_tokens_saved", 0)
compression_cache_stats = {
"mode": PROXY_MODE_TOKEN,
"active_sessions": _active_sessions,
"total_entries": total_entries,
"total_hits": total_hits,
"total_misses": total_misses,
"hit_rate": round(total_hits / max(1, total_hits + total_misses) * 100, 1),
"total_tokens_saved": total_tokens_saved,
}
else:
compression_cache_stats = {"mode": proxy.config.mode}
# Build unified savings summary (all layers)
cache_net_usd = prefix_cache_stats.get("totals", {}).get("net_savings_usd", 0.0)
total_tokens_all_layers = all_layers_tokens_saved
persistent_savings = m.savings_tracker.stats_preview()
display_session = persistent_savings.get("display_session", {})
recent_request_logs = proxy.logger.get_recent(10_000) if proxy.logger else []
recent_request_payload = _build_recent_request_payload()
# Tool-schema deferral savings: tool-definition tokens kept out of the
# model's context by deferring heavy schemas until they're needed
# (native tool-search injection + any registered turn-hook tools
# rewrite). Attributed to Headroom only — see _tool_schema_saved_from_tags.
# Aggregated over the recent request-log window.
tool_schema_tokens = 0
tool_schema_requests = 0
for _ts_log in recent_request_logs:
_ts_saved = _tool_schema_saved_from_tags(_ts_log.get("tags"))
if _ts_saved > 0:
tool_schema_tokens += _ts_saved
tool_schema_requests += 1
agent_usage = _build_agent_usage_summary(
recent_request_logs,
requests_by_provider=dict(m.requests_by_provider),
requests_by_model=dict(m.requests_by_model),
global_before_tokens=proxy_total_before_compression,
global_after_tokens=m.tokens_input_total,
global_tokens_saved=proxy_compression_tokens,
global_output_tokens=m.tokens_output_total,
)
# Output-side reduction (counterfactual estimate from the shaper's
# ledger). Distinct from input compression above: these are OUTPUT
# tokens the model didn't emit because we steered verbosity / routed
# effort down. Always labelled estimated-vs-measured + a CI so it's
# never mistaken for an exact count. Best-effort — never break /stats.
output_reduction: dict[str, Any] = {"available": False}
try:
from headroom.proxy.output_savings import get_recorder
_oest = get_recorder().estimate()
if _oest.n_requests > 0:
output_reduction = {
"available": True,
"method": _oest.kind, # "measured" | "estimated"
"tokens_saved": round(_oest.tokens_saved),
"baseline_tokens": round(_oest.baseline_tokens),
"reduction_percent": round(_oest.pct, 1),
"ci_low_percent": round(_oest.ci_low_pct, 1),
"ci_high_percent": round(_oest.ci_high_pct, 1),
"requests": _oest.n_requests,
}
except Exception: # pragma: no cover - defensive
pass
return {
"summary": summary,
"agent_usage": agent_usage,
"savings": {
"total_tokens": total_tokens_all_layers,
"per_project": persistent_savings.get("projects", {}),
"by_layer": {
"cli_filtering": {
"tool": cli_filtering_tool,
"label": cli_filtering_label,
"available": cli_filtering_available,
"tokens": cli_tokens_avoided,
"tokens_saved": cli_tokens_avoided,
"session": cli_filtering_session,
"lifetime": cli_filtering_lifetime,
"session_savings_pct": (
cli_filtering_stats.get("session_savings_pct")
if cli_filtering_stats
else None
),
"lifetime_savings_pct": (
cli_filtering_stats.get("lifetime_avg_savings_pct")
if cli_filtering_stats
else None
),
"refresh_interval_seconds": (
cli_filtering_stats.get("refresh_interval_seconds")
if cli_filtering_stats
else None
),
"included_in": "tokens.saved",
"description": (
f"Tokens avoided by CLI output filtering ({cli_filtering_label}) "
"before reaching context. "
"Included in dashboard token savings, but not in dollar savings."
),
},
"compression": {
"tokens": proxy_compression_tokens,
"proxy_tokens": proxy_compression_tokens,
"cli_filtering_tokens": cli_tokens_avoided,
"rtk_tokens": rtk_tokens_avoided,
"lean_ctx_tokens": lean_ctx_tokens_avoided,
"all_layers_tokens": all_layers_tokens_saved,
"description": (
"Tokens removed by Headroom proxy compression. "
"Dashboard token savings also includes CLI context-tool filtering."
),
},
"prefix_cache": {
"discount_usd": round(cache_net_usd, 4),
"description": (
"Cost discount from provider prefix caching. "
"Headroom's CacheAligner improves hit rates; "
"baseline caching is provider-native."
),
},
"output_shaping": {
**output_reduction,
"description": (
"OUTPUT tokens the model didn't emit because the shaper "
"steered verbosity / routed effort down. Counterfactual — "
"shown as an estimate (vs a learned baseline) or measured "
"(A/B holdout), always with a confidence band."
),
},
"tool_search": {
"tokens": tool_schema_tokens,
"tokens_saved": tool_schema_tokens,
"requests": tool_schema_requests,
"window": len(recent_request_logs),
"description": (
"Tool-definition tokens kept out of the model's context "
"by deferring heavy tool schemas until they're searched "
"for. Counted only when Headroom performed the deferral — "
"not when the client (e.g. Claude Code / Codex) already "
"had tool search enabled. Aggregated over the recent "
"request window."
),
},
},
},
"requests": {
"total": m.requests_total,
"cached": m.requests_cached,
"rate_limited": m.requests_rate_limited,
"failed": m.requests_failed,
"by_provider": dict(m.requests_by_provider),
"by_model": dict(m.requests_by_model),
"by_stack": dict(m.requests_by_stack),
},
"tokens": {
"input": m.tokens_input_total,
"output": m.tokens_output_total,
"output_saved": output_reduction.get("tokens_saved", 0),
"output_reduction_percent": output_reduction.get("reduction_percent", 0),
"output_reduction": output_reduction,
"saved": all_layers_tokens_saved,
"proxy_compression_saved": proxy_compression_tokens,
"cli_filtering_saved": cli_tokens_avoided,
"rtk_saved": rtk_tokens_avoided,
"lean_ctx_saved": lean_ctx_tokens_avoided,
"cli_tokens_avoided": cli_tokens_avoided,
"proxy_total_before_compression": proxy_total_before_compression,
"total_before_compression": total_tokens_before,
"all_layers_saved": all_layers_tokens_saved,
# Compressible-only denominator: tokens we extracted as
# candidates + tool-schema tokens we compacted. Excludes
# frozen-prefix content (user msgs, system prompt, prior
# turns) that we deliberately don't touch. Already
# pre-compression — do NOT add `tokens_saved` again.
"proxy_attempted_tokens": attempted_input_tokens,
# Active compression: savings as a fraction of what we
# *tried* to compress. The number the dashboard headline
# should show — it answers "are we doing well *when we
# have something to compress?*" rather than diluting the
# win by frozen-prefix bytes we never touched.
"active_savings_percent": round(
(proxy_compression_tokens / attempted_input_tokens * 100)
if attempted_input_tokens > 0
else 0,
2,
),
# Whole-request ratio kept for transparency. Heavily
# diluted by frozen prefix on Codex-style requests
# where most input is non-compressible by design.
"proxy_savings_percent": round(
(proxy_compression_tokens / proxy_total_before_compression * 100)
if proxy_total_before_compression > 0
else 0,
2,
),
# New-content-relative rate: savings as a fraction of the
# input that would have newly entered context (provider-
# billed uncached + cache-write tokens, plus the tokens
# compression removed before they could be billed). The
# whole-request ratios above recount the FULL transcript
# every turn, so a 200-turn session counts its history
# 200x into the denominator and long-running sessions
# (1M-context models never compact) read as ~0% no
# matter how well compression performs on new content.
# Guarded on new_input_tokens > 0 (not the full sum): the
# cache accumulators only see requests with cache
# activity, so a deployment with no cache metrics (e.g.
# Bedrock) would otherwise divide savings by themselves
# and report ~100%. No usage data -> report 0, not a lie.
"new_input_tokens": new_input_tokens,
"new_input_savings_percent": round(
(proxy_compression_tokens / (new_input_tokens + proxy_compression_tokens) * 100)
if new_input_tokens > 0
else 0,
2,
),
"savings_percent": round(
(all_layers_tokens_saved / total_tokens_before * 100)
if total_tokens_before > 0
else 0,
2,
),
"all_layers_savings_percent": round(
(all_layers_tokens_saved / total_tokens_before * 100)
if total_tokens_before > 0
else 0,
2,
),
},
"latency": {
"average_ms": avg_latency_ms,
"min_ms": min_latency_ms,
"max_ms": max_latency_ms,
"total_requests": m.latency_count,
},
"overhead": {
"average_ms": avg_overhead_ms,
"min_ms": min_overhead_ms,
"max_ms": max_overhead_ms,
},
"ttfb": {
"average_ms": avg_ttfb_ms,
"min_ms": min_ttfb_ms,
"max_ms": max_ttfb_ms,
},
"pipeline_timing": {
name: {
"average_ms": round(
m.transform_timing_sum[name] / m.transform_timing_count[name], 2
),
"max_ms": round(m.transform_timing_max[name], 2),
"count": m.transform_timing_count[name],
}
for name in sorted(m.transform_timing_sum.keys())
}
if m.transform_timing_sum
else {},
"compressions_by_strategy": dict(m.compressions_by_strategy),
"tokens_saved_by_strategy": dict(m.tokens_saved_by_strategy),
"codex_ws": {
"units_total": m.codex_ws_units_total,
"units_modified_total": m.codex_ws_units_modified_total,
"units_by_strategy": dict(m.codex_ws_units_by_strategy),
"units_by_category": dict(m.codex_ws_units_by_category),
"units_by_content_type": dict(m.codex_ws_units_by_content_type),
"units_by_text_shape": dict(m.codex_ws_units_by_text_shape),
"units_to_kompress_total": m.codex_ws_units_to_kompress_total,
"units_kompress_attempted_total": m.codex_ws_units_kompress_attempted_total,
"units_to_kompress_percent": _pct(
m.codex_ws_units_to_kompress_total,
m.codex_ws_units_total,
),
"units_kompress_attempted_percent": _pct(
m.codex_ws_units_kompress_attempted_total,
m.codex_ws_units_total,
),
"unit_elapsed_ms": {
"average": round(
m.codex_ws_unit_elapsed_ms_sum / m.codex_ws_units_total,
2,
)
if m.codex_ws_units_total
else 0.0,
"max": round(m.codex_ws_unit_elapsed_ms_max, 2),
},
"unit_bytes_sum": m.codex_ws_unit_bytes_sum,
"unit_tokens_before_sum": m.codex_ws_unit_tokens_before_sum,
"unit_tokens_after_sum": m.codex_ws_unit_tokens_after_sum,
"unit_tokens_saved_sum": m.codex_ws_unit_tokens_saved_sum,
"frames_attempted_total": m.codex_ws_frames_attempted_total,
"frames_compressed_total": m.codex_ws_frames_compressed_total,
"frames_failed_total": m.codex_ws_frames_failed_total,
"frames_to_kompress_total": m.codex_ws_frames_to_kompress_total,
"frames_kompress_attempted_total": (m.codex_ws_frames_kompress_attempted_total),
"frames_to_kompress_percent": _pct(
m.codex_ws_frames_to_kompress_total,
m.codex_ws_frames_attempted_total,
),
"frames_kompress_attempted_percent": _pct(
m.codex_ws_frames_kompress_attempted_total,
m.codex_ws_frames_attempted_total,
),
"frame_elapsed_ms": {
"average": round(
m.codex_ws_frame_elapsed_ms_sum / m.codex_ws_frames_attempted_total,
2,
)
if m.codex_ws_frames_attempted_total
else 0.0,
"max": round(m.codex_ws_frame_elapsed_ms_max, 2),
},
"frame_bytes_before_sum": m.codex_ws_frame_bytes_before_sum,
"frame_bytes_after_sum": m.codex_ws_frame_bytes_after_sum,
"frame_attempted_tokens_sum": m.codex_ws_frame_attempted_tokens_sum,
"frame_tokens_saved_sum": m.codex_ws_frame_tokens_saved_sum,
},
"waste_signals": dict(m.waste_signals_total) if m.waste_signals_total else {},
# ContentRouter protection categories aggregated across the
# session. Lets operators see, e.g., that 80% of messages
# were `user_msg` (protected) and only 5% reached the
# compressor — explains why compression rate is low and
# whether `--compress-user-messages` would help (#454).
"router": {
"route_counts": dict(m.router_route_counts) if m.router_route_counts else {},
},
"savings_history": m.savings_history[-100:], # Last 100 data points
"display_session": display_session,
# Whether LiteLLM is importable. Pricing (the "$ Saved" tile) is
# derived entirely from LiteLLM's cost tables, and LiteLLM is gated
# off on Python >=3.14 in pyproject — so when this is False the
# dashboard tells the user to reinstall on 3.13 instead of just
# showing $0.00 forever.
"litellm_available": LITELLM_AVAILABLE,
"persistent_savings": persistent_savings,
"prefix_cache": prefix_cache_stats,
"cost": _merge_cost_stats(
proxy.cost_tracker.stats() if proxy.cost_tracker else None,
prefix_cache_stats,
cli_tokens_avoided=cli_tokens_avoided,
),
"compression": {
"ccr_entries": compression_stats.get("entry_count", 0),
"ccr_max_entries": compression_stats.get("max_entries", 0),
"original_tokens_cached": compression_stats.get("total_original_tokens", 0),
"compressed_tokens_cached": compression_stats.get("total_compressed_tokens", 0),
"ccr_retrievals": compression_stats.get("total_retrievals", 0),
},
"compression_cache": compression_cache_stats,
# Always False: the anonymous telemetry beacon was removed, so no
# telemetry is ever shipped externally (local collection only).
"anon_telemetry_shipping": False,
"telemetry": {
"enabled": telemetry_stats.get("enabled", False),
"total_compressions": telemetry_stats.get("total_compressions", 0),
"total_retrievals": telemetry_stats.get("total_retrievals", 0),
"global_retrieval_rate": round(telemetry_stats.get("global_retrieval_rate", 0), 4),
"tool_signatures_tracked": telemetry_stats.get("tool_signatures_tracked", 0),
"avg_compression_ratio": round(telemetry_stats.get("avg_compression_ratio", 0), 4),
"avg_token_reduction": round(telemetry_stats.get("avg_token_reduction", 0), 4),
},
"otel": get_otel_metrics_status(),
"langfuse": get_langfuse_tracing_status(),
"feedback_loop": {
"tools_tracked": feedback_stats.get("tools_tracked", 0),
"total_compressions": feedback_stats.get("total_compressions", 0),
"total_retrievals": feedback_stats.get("total_retrievals", 0),
"global_retrieval_rate": round(feedback_stats.get("global_retrieval_rate", 0), 4),
"tools_with_high_retrieval": sum(
1
for p in feedback_stats.get("tool_patterns", {}).values()
if p.get("retrieval_rate", 0) > 0.3
),
},
"toin": get_toin().get_stats(),
"context_tool": {
"configured": cli_filtering_tool,
"label": cli_filtering_label,
"available": cli_filtering_available,
"stats": cli_filtering_stats,
},
"cli_filtering": cli_filtering_stats,
"proxy_inbound": proxy.metrics.inbound_snapshot(),
"cache": await proxy.cache.stats() if proxy.cache else None,
"rate_limiter": await proxy.rate_limiter.stats() if proxy.rate_limiter else None,
**recent_request_payload,
"log_full_messages": proxy.config.log_full_messages if proxy else False,
**get_quota_registry().get_all_stats(),
"throughput": throughput,
}
def _dashboard_config_payload() -> dict[str, Any]:
profile_kwargs = proxy_pipeline_kwargs(config)
target_ratio = profile_kwargs.get("target_ratio", config.target_ratio)
target_savings_percent = None
if isinstance(target_ratio, int | float):
target_savings_percent = round(max(0.0, min(1.0, 1.0 - float(target_ratio))) * 100, 1)
return {
"savings_profile": config.savings_profile,
"target_ratio": target_ratio,
"target_savings_percent": target_savings_percent,
"compress_user_messages": bool(
profile_kwargs.get("compress_user_messages", config.compress_user_messages)
),
"compress_system_messages": bool(
profile_kwargs.get("compress_system_messages", config.compress_system_messages)
),
"protect_recent": profile_kwargs.get("read_protection_window", config.protect_recent),
"protect_analysis_context": config.protect_analysis_context,
"min_tokens_to_crush": profile_kwargs.get(
"min_tokens_to_compress", config.min_tokens_to_crush
),
"max_items_after_crush": profile_kwargs.get(
"max_items_after_crush", config.max_items_after_crush
),
"smart_crusher_with_compaction": profile_kwargs.get(
"smart_crusher_with_compaction",
config.smart_crusher_with_compaction,
),
"force_kompress": bool(profile_kwargs.get("force_kompress", False)),
"accuracy_guard": config.accuracy_guard,
}
async def _get_cached_stats_payload() -> dict[str, Any]:
"""Return a short-TTL cached `/stats` snapshot for dashboard polling."""
now = time.monotonic()
cached_payload = cast(dict[str, Any] | None, _stats_snapshot.get("value"))
if cached_payload is not None and now < float(_stats_snapshot["expires_at"]):
return cached_payload
async with _stats_snapshot_lock:
now = time.monotonic()
cached_payload = cast(dict[str, Any] | None, _stats_snapshot.get("value"))
if cached_payload is not None and now < float(_stats_snapshot["expires_at"]):
return cached_payload
payload = await _build_stats_payload()
_stats_snapshot["value"] = payload
_stats_snapshot["expires_at"] = time.monotonic() + DASHBOARD_STATS_CACHE_TTL_SECONDS
return payload
@app.get("/stats")
async def stats(request: Request, cached: bool = False):
"""Get comprehensive proxy statistics.
This is the main stats endpoint - it aggregates data from all subsystems:
- Request metrics (total, cached, failed, by model/provider)
- Token usage and savings
- Cost tracking
- Canonical persisted display_session metrics for downstream dashboards
- Compression (CCR) statistics
- Telemetry/TOIN (data flywheel) statistics
- Cache and rate limiter stats
Use ``?cached=1`` for the dashboard fast path. That returns a short-TTL
snapshot to avoid rebuilding the full payload on every UI poll.
``recent_requests`` / ``request_logs`` (per-request ids, providers,
models, errors) and ``config`` (backend + savings profile) are embedded
only for loopback callers — the local dashboard. Network callers still
get the aggregate counters but never the per-request metadata.
"""
include_sensitive = _request_is_loopback(request)
if cached:
payload = dict(await _get_cached_stats_payload())
if include_sensitive:
# Refresh the per-request tail on top of the cached snapshot.
payload.update(_build_recent_request_payload())
payload["config"] = _dashboard_config_payload()
else:
payload = await _build_stats_payload()
if include_sensitive:
payload["config"] = _dashboard_config_payload()
if not include_sensitive:
# _build_stats_payload bakes these in; strip for network callers.
payload.pop("recent_requests", None)
payload.pop("request_logs", None)
return payload
@app.post("/stats/reset", dependencies=[Depends(_require_loopback)])
async def stats_reset():
"""Reset in-memory proxy stats for local test/debug isolation."""
await proxy.metrics.reset_runtime()
if proxy.cost_tracker:
proxy.cost_tracker.reset_runtime()
await initialize_context_tool_session_baseline()
async with _stats_snapshot_lock:
_stats_snapshot["value"] = None
_stats_snapshot["expires_at"] = 0.0
return JSONResponse(status_code=200, content={"status": "reset"})
@app.get("/stats-history")
async def stats_history(
format: Literal["json", "csv"] = "json",
series: Literal["history", "hourly", "daily", "weekly", "monthly"] = "history",
history_mode: Literal["compact", "full", "none"] = "compact",
):
"""Get durable proxy compression history plus display-session state.
The JSON payload also carries a ``cli_filtering`` key with live RTK
stats. This is a curated subset (``tool``, ``label``, ``available``,
``lifetime``, ``session``) tailored to the Historical tab, not the
full ``_get_context_tool_stats()`` payload that ``/stats`` exposes.
It is ``None`` only when the stats read hard-fails; when the tool is
merely absent, ``cli_filtering`` stays populated with
``available: False`` and zeroed counters so the tab can distinguish
"not installed" from "installed, no data yet."
"""
if format == "csv":
filename = f"headroom-stats-history-{series}.csv"
return Response(
content=proxy.metrics.savings_tracker.export_csv(series=series),
media_type="text/csv; charset=utf-8",
headers={"Content-Disposition": f'attachment; filename="{filename}"'},
)
history = proxy.metrics.savings_tracker.history_response(history_mode=history_mode)
# Augment with live RTK/cli-filtering lifetime stats so the Historical
# tab can display them. These live in the context-tool's own stats file
# and survive proxy restarts — exactly what the Historical tab needs.
# Best-effort: if the RTK stats file can't be read (missing, parse error,
# IO), fall back to None so the Historical tab stays available instead of
# 500ing. The tab hides the card when cli_filtering is None.
try:
cli_stats = await asyncio.to_thread(_get_context_tool_stats)
except Exception:
logger.debug("stats-history: RTK stats unavailable", exc_info=True)
cli_stats = None
if cli_stats:
history["cli_filtering"] = {
"tool": str(cli_stats.get("tool", "rtk")),
"label": str(cli_stats.get("label", "RTK")),
"available": bool(cli_stats.get("installed", False)),
"lifetime": cli_stats.get("lifetime", {}),
"session": cli_stats.get("session", {}),
}
else:
history["cli_filtering"] = None
return history
@app.get("/transformations/feed", dependencies=[Depends(_require_loopback)])
async def transformations_feed(limit: int = 20):
"""Get recent message transformations for the live feed.
Loopback-only: when ``log_full_messages`` is enabled this returns the
full request/response message bodies (prompt content and completions)
via ``request_messages`` / ``compressed_messages`` / ``response_content``.
With the default ``--host 0.0.0.0`` Docker bind, leaving it open would
expose chat history to anyone able to reach the proxy port. The
dashboard runs in the user's browser on loopback, so this gate does not
break legitimate use.
Returns empty list if log_full_messages is disabled (messages are not stored).
"""
if limit > 100:
limit = 100
transformations = []
log_full_messages = proxy.config.log_full_messages if proxy else False
if proxy and proxy.logger:
logs = proxy.logger.get_recent_with_messages(limit)
for log in logs:
transformations.append(
{
"request_id": log.get("request_id"),
"timestamp": log.get("timestamp"),
"provider": log.get("provider"),
"model": log.get("model"),
"input_tokens_original": log.get("input_tokens_original"),
"input_tokens_optimized": log.get("input_tokens_optimized"),
"tokens_saved": log.get("tokens_saved"),
"savings_percent": log.get("savings_percent"),
"transforms_applied": log.get("transforms_applied", []),
"request_messages": log.get("request_messages"),
"compressed_messages": log.get("compressed_messages"),
"response_content": log.get("response_content"),
"turn_id": log.get("turn_id"),
}
)
return {"transformations": transformations, "log_full_messages": log_full_messages}
@app.get("/subscription-window")
async def subscription_window():
"""Current Anthropic subscription window utilisation and Headroom contribution.
Issue #281: the Anthropic OAuth usage API is polled every 5 minutes
(aggressive polling risks 429s / OAuth-token flagging), so the cached
``utilization_pct`` lags reality by up to one poll interval. When the
user's 5-hour window rolls over between two polls the dashboard would
otherwise render the OLD window's percentage. We:
1. Optionally trigger a 60s-floored singleton poll on dashboard load
(bounded across users, well within Anthropic tolerance).
2. Render via :meth:`SubscriptionTracker.render_state`, which
synthesizes post-reset windows from local transcript-derived
token counts when ``now >= window.resets_at``.
"""
tracker = get_subscription_tracker()
if tracker is None:
return JSONResponse(
status_code=503,
content={"error": "Subscription tracking is not enabled"},
)
await tracker.maybe_poll_on_demand()
return JSONResponse(content=tracker.render_state())
@app.get("/quota")
async def quota():
"""Unified quota/rate-limit stats for all registered providers (Anthropic, Codex, Copilot)."""
return JSONResponse(content=get_quota_registry().get_all_stats())
@app.get("/metrics")
async def metrics():
"""Prometheus metrics endpoint."""
return PlainTextResponse(
await proxy.metrics.export(),
media_type="text/plain; version=0.0.4",
)
# Debug endpoints
@app.get("/debug/memory", dependencies=[Depends(_require_loopback)])
async def debug_memory():
"""Get detailed memory usage statistics.
Returns memory usage for all tracked components including:
- Process-level memory (RSS, VMS, percent)
- Per-component memory usage and budgets
- Cache hit/miss statistics
- Total tracked vs target budget
This endpoint is useful for debugging memory issues and
monitoring memory budgets.
"""
from ..memory.tracker import MemoryTracker
tracker = MemoryTracker.get()
# Register components if not already registered
_register_memory_components(proxy, tracker)
report = tracker.get_report()
return report.to_dict()
@app.post("/cache/clear", dependencies=[Depends(_require_loopback)])
async def clear_cache():
"""Clear the response cache.
Loopback-only: this mutates server state. With the default
``--host 0.0.0.0`` Docker bind, an unauthenticated POST from any
network-reachable client would otherwise let them forcibly evict the
proxy's cached completions — a denial-of-service / cost-amplification
lever (every cleared entry forces a fresh upstream call).
"""
if proxy.cache:
await proxy.cache.clear()
return {"status": "cleared"}
return {"status": "cache disabled"}
# CCR (Compress-Cache-Retrieve) endpoints
@app.post("/v1/retrieve", dependencies=[Depends(_require_loopback)])
async def ccr_retrieve(request: Request):
"""Retrieve original content from CCR compression cache.
This is the "Retrieve" part of CCR (Compress-Cache-Retrieve).
When SmartCrusher compresses tool outputs, the original data is cached.
LLMs can call this endpoint to get more data if needed.
Request body:
hash (str): Hash key from compression marker (required)
Response:
{"hash": "...", "original_content": "...", ...}
"""
data = await request.json()
hash_key = data.get("hash")
if not hash_key:
raise HTTPException(status_code=400, detail="hash required")
store = get_compression_store()
entry_status = store.get_entry_status(hash_key, clean_expired=True)
if entry_status["status"] != "available":
raise HTTPException(
status_code=404,
detail=format_retrieval_miss_detail(entry_status),
)
# Retrieval is by hash: always return the full original content.
entry = store.retrieve(hash_key)
if entry:
return {
"hash": hash_key,
"original_content": entry.original_content,
"original_tokens": entry.original_tokens,
"original_item_count": entry.original_item_count,
"compressed_item_count": entry.compressed_item_count,
"tool_name": entry.tool_name,
"retrieval_count": entry.retrieval_count,
}
raise HTTPException(
status_code=404,
detail=format_retrieval_miss_detail(
store.get_entry_status(hash_key, clean_expired=True)
),
)
@app.get("/v1/retrieve/stats", dependencies=[Depends(_require_loopback)])
async def ccr_stats():
"""Get CCR compression store statistics."""
store = get_compression_store()
stats = store.get_stats()
events = store.get_retrieval_events(limit=20)
return {
"store": stats,
"recent_retrievals": [
{
"hash": e.hash,
"query": e.query,
"items_retrieved": e.items_retrieved,
"total_items": e.total_items,
"tool_name": e.tool_name,
"retrieval_type": e.retrieval_type,
}
for e in events
],
}
@app.get("/v1/feedback")
async def ccr_feedback():
"""Get CCR feedback loop statistics and learned patterns.
This endpoint exposes the feedback loop's learned patterns for monitoring
and debugging. It shows:
- Per-tool retrieval rates (high = compress less aggressively)
- Common search queries per tool
- Queried fields (suggest what to preserve)
Use this to understand how well compression is working and whether
the feedback loop is adjusting appropriately.
"""
feedback = get_compression_feedback()
stats = feedback.get_stats()
return {
"feedback": stats,
"hints_example": {
tool_name: {
"hints": {
"max_items": hints.max_items
if (hints := feedback.get_compression_hints(tool_name))
else 15,
"suggested_items": hints.suggested_items if hints else None,
"skip_compression": hints.skip_compression if hints else False,
"preserve_fields": hints.preserve_fields if hints else [],
"reason": hints.reason if hints else "",
}
}
for tool_name in list(stats.get("tool_patterns", {}).keys())[:5]
},
}
@app.get("/v1/feedback/{tool_name}")
async def ccr_feedback_for_tool(tool_name: str):
"""Get compression hints for a specific tool.
Returns feedback-based hints that would be used for compressing
this tool's output.
"""
feedback = get_compression_feedback()
hints = feedback.get_compression_hints(tool_name)
patterns = feedback.get_all_patterns().get(tool_name)
return {
"tool_name": tool_name,
"hints": {
"max_items": hints.max_items,
"min_items": hints.min_items,
"suggested_items": hints.suggested_items,
"aggressiveness": hints.aggressiveness,
"skip_compression": hints.skip_compression,
"preserve_fields": hints.preserve_fields,
"reason": hints.reason,
},
"pattern": {
"total_compressions": patterns.total_compressions if patterns else 0,
"total_retrievals": patterns.total_retrievals if patterns else 0,
"retrieval_rate": patterns.retrieval_rate if patterns else 0.0,
"full_retrieval_rate": patterns.full_retrieval_rate if patterns else 0.0,
"search_rate": patterns.search_rate if patterns else 0.0,
"common_queries": list(patterns.common_queries.keys())[:10] if patterns else [],
"queried_fields": list(patterns.queried_fields.keys())[:10] if patterns else [],
}
if patterns
else None,
}
# Telemetry endpoints (Data Flywheel)
@app.get("/v1/telemetry")
async def telemetry_stats():
"""Get telemetry statistics for the data flywheel.
This endpoint exposes privacy-preserving telemetry data that powers
the data flywheel - learning optimal compression strategies across
tool types based on usage patterns.
What's collected (anonymized):
- Tool output structure patterns (field types, not values)
- Compression decisions and ratios
- Retrieval patterns (rate, type, not content)
- Strategy effectiveness
What's NOT collected:
- Actual data values
- User identifiers
- Queries or search terms
- File paths or tool names (hashed by default)
"""
telemetry = get_telemetry_collector()
return telemetry.get_stats()
@app.get("/v1/telemetry/export")
async def telemetry_export():
"""Export full telemetry data for aggregation.
This endpoint exports all telemetry data in a format suitable for
cross-user aggregation. The data is privacy-preserving - no actual
values are included, only structural patterns and statistics.
Use this for:
- Building a central learning service
- Sharing learned patterns across instances
- Analysis and debugging
"""
telemetry = get_telemetry_collector()
return telemetry.export_stats()
@app.post("/v1/telemetry/import")
async def telemetry_import(request: Request):
"""Import telemetry data from another source.
This allows merging telemetry from multiple sources for cross-user
learning. The imported data is merged with existing statistics.
Request body: Telemetry export data from /v1/telemetry/export
"""
telemetry = get_telemetry_collector()
data = await request.json()
telemetry.import_stats(data)
return {"status": "imported", "current_stats": telemetry.get_stats()}
@app.get("/v1/telemetry/tools")
async def telemetry_tools():
"""Get telemetry statistics for all tracked tool signatures.
Returns statistics per tool signature (anonymized), including:
- Compression ratios and strategy usage
- Retrieval rates (high = compression too aggressive)
- Learned recommendations
"""
telemetry = get_telemetry_collector()
all_stats = telemetry.get_all_tool_stats()
return {
"tool_count": len(all_stats),
"tools": {sig_hash: stats.to_dict() for sig_hash, stats in all_stats.items()},
}
@app.get("/v1/telemetry/tools/{signature_hash}")
async def telemetry_tool_detail(signature_hash: str):
"""Get detailed telemetry for a specific tool signature.
Includes learned recommendations if enough data has been collected.
"""
telemetry = get_telemetry_collector()
stats = telemetry.get_tool_stats(signature_hash)
recommendations = telemetry.get_recommendations(signature_hash)
if stats is None:
raise HTTPException(
status_code=404, detail=f"No telemetry found for signature: {signature_hash}"
)
return {
"signature_hash": signature_hash,
"stats": stats.to_dict(),
"recommendations": recommendations,
}
# TOIN (Tool Output Intelligence Network) endpoints
@app.get("/v1/toin/stats")
async def toin_stats():
"""Get overall TOIN statistics.
Returns aggregated statistics from the Tool Output Intelligence Network,
which learns optimal compression strategies across all tool types.
Response includes:
- enabled: Whether TOIN is enabled
- patterns_tracked: Number of unique tool patterns being tracked
- total_compressions: Total compression events recorded
- total_retrievals: Total retrieval events recorded
- global_retrieval_rate: Overall retrieval rate (high = compression too aggressive)
- patterns_with_recommendations: Patterns with enough data for recommendations
"""
toin = get_toin()
return toin.get_stats()
@app.get("/v1/toin/patterns")
async def toin_patterns(limit: int = 20):
"""List TOIN patterns with most samples.
Returns patterns sorted by sample_size descending. Use this to see
which tool types have the most data and their learned behaviors.
Query params:
limit: Maximum number of patterns to return (default 20)
Response includes for each pattern:
- hash: Truncated tool signature hash (12 chars)
- compressions: Total compression events
- retrievals: Total retrieval events
- retrieval_rate: Percentage of compressions that triggered retrieval
- confidence: Confidence level in recommendations (0.0-1.0)
- skip_recommended: Whether TOIN recommends skipping compression
- optimal_max_items: Learned optimal max_items setting
"""
toin = get_toin()
exported = toin.export_patterns()
patterns_data = exported.get("patterns", {})
# Convert to list and sort by sample_size
patterns_list = []
for sig_hash, pattern_dict in patterns_data.items():
sample_size = pattern_dict.get("sample_size", 0)
total_compressions = pattern_dict.get("total_compressions", 0)
total_retrievals = pattern_dict.get("total_retrievals", 0)
retrieval_rate = (
total_retrievals / total_compressions if total_compressions > 0 else 0.0
)
patterns_list.append(
{
"hash": sig_hash[:12],
"compressions": total_compressions,
"retrievals": total_retrievals,
"retrieval_rate": f"{retrieval_rate:.1%}",
"confidence": round(pattern_dict.get("confidence", 0.0), 3),
"skip_recommended": pattern_dict.get("skip_compression_recommended", False),
"optimal_max_items": pattern_dict.get("optimal_max_items", 20),
"sample_size": sample_size,
}
)
# Sort by sample_size descending
patterns_list.sort(key=lambda p: p["sample_size"], reverse=True)
# Remove sample_size from output (used only for sorting)
for p in patterns_list:
del p["sample_size"]
return patterns_list[:limit]
@app.get("/v1/toin/pattern/{hash_prefix}")
async def toin_pattern_detail(hash_prefix: str):
"""Get detailed TOIN pattern info by hash prefix.
Searches for a pattern where the tool signature hash starts with
the provided prefix. Returns full pattern details if found.
Path params:
hash_prefix: Beginning of the tool signature hash (min 4 chars recommended)
Response: Full pattern.to_dict() with all learned statistics and recommendations.
"""
toin = get_toin()
exported = toin.export_patterns()
patterns_data = exported.get("patterns", {})
# Search for pattern with matching hash prefix
for sig_hash, pattern_dict in patterns_data.items():
if sig_hash.startswith(hash_prefix):
return pattern_dict
raise HTTPException(
status_code=404, detail=f"No TOIN pattern found with hash starting with: {hash_prefix}"
)
@app.get("/v1/retrieve/{hash_key}", dependencies=[Depends(_require_loopback)])
async def ccr_retrieve_get(hash_key: str):
"""GET version of CCR retrieve for easier testing."""
store = get_compression_store()
entry_status = store.get_entry_status(hash_key, clean_expired=True)
if entry_status["status"] != "available":
raise HTTPException(
status_code=404,
detail=format_retrieval_miss_detail(entry_status),
)
# Retrieval is by hash: always return the full original content.
entry = store.retrieve(hash_key)
if entry:
return {
"hash": hash_key,
"original_content": entry.original_content,
"original_tokens": entry.original_tokens,
"original_item_count": entry.original_item_count,
"compressed_item_count": entry.compressed_item_count,
"tool_name": entry.tool_name,
"retrieval_count": entry.retrieval_count,
}
raise HTTPException(
status_code=404,
detail=format_retrieval_miss_detail(
store.get_entry_status(hash_key, clean_expired=True)
),
)
# CCR Tool Call Handler - for agent frameworks to call when LLM uses headroom_retrieve
@app.post("/v1/retrieve/tool_call", dependencies=[Depends(_require_loopback)])
async def ccr_handle_tool_call(request: Request):
"""Handle a CCR tool call from an LLM response.
This endpoint accepts tool call formats from various providers and returns
a properly formatted tool result. Agent frameworks can use this to handle
CCR tool calls without implementing the retrieval logic themselves.
Request body (Anthropic format):
{
"tool_call": {
"id": "toolu_123",
"name": "headroom_retrieve",
"input": {"hash": "abc123"}
},
"provider": "anthropic"
}
Request body (OpenAI format):
{
"tool_call": {
"id": "call_123",
"function": {
"name": "headroom_retrieve",
"arguments": "{\"hash\": \"abc123\"}"
}
},
"provider": "openai"
}
Response:
{
"tool_result": {...}, # Formatted for the provider
"success": true,
"data": {...} # Raw retrieval data
}
"""
data = await request.json()
tool_call = data.get("tool_call", {})
provider = data.get("provider", "anthropic")
# Parse the tool call
hash_key = parse_tool_call(tool_call, provider)
if hash_key is None:
raise HTTPException(
status_code=400, detail=f"Invalid tool call or not a {CCR_TOOL_NAME} call"
)
# Perform retrieval
store = get_compression_store()
entry_status = store.get_entry_status(hash_key, clean_expired=True)
if entry_status["status"] != "available":
retrieval_data = {
"error": format_retrieval_miss_detail(entry_status),
"hash": hash_key,
"status": entry_status["status"],
"ttl_seconds": entry_status.get("ttl_seconds", entry_status["default_ttl_seconds"]),
}
else:
# Retrieval is by hash: always return the full original content.
entry = store.retrieve(hash_key)
if entry:
retrieval_data = {
"hash": hash_key,
"original_content": entry.original_content,
"original_item_count": entry.original_item_count,
"compressed_item_count": entry.compressed_item_count,
}
else:
miss_status = store.get_entry_status(hash_key, clean_expired=True)
retrieval_data = {
"error": format_retrieval_miss_detail(miss_status),
"hash": hash_key,
"status": miss_status["status"],
"ttl_seconds": miss_status.get(
"ttl_seconds", miss_status["default_ttl_seconds"]
),
}
# Format tool result for provider
tool_call_id = tool_call.get("id", "")
result_content = json.dumps(retrieval_data, indent=2)
if provider == "anthropic":
tool_result = {
"type": "tool_result",
"tool_use_id": tool_call_id,
"content": result_content,
}
elif provider == "openai":
tool_result = {
"role": "tool",
"tool_call_id": tool_call_id,
"content": result_content,
}
else:
tool_result = {
"tool_call_id": tool_call_id,
"content": result_content,
}
return {
"tool_result": tool_result,
"success": "error" not in retrieval_data,
"data": retrieval_data,
}
# Compression-only endpoint (for TypeScript SDK and other HTTP clients)
@app.post("/v1/compress", dependencies=[Depends(_require_loopback)])
async def compress_messages(request: Request):
return await proxy.handle_compress(request)
register_provider_routes(app, proxy)
return app
def _json_ready(value: Any) -> Any:
if is_dataclass(value) and not isinstance(value, type):
return {field.name: _json_ready(getattr(value, field.name)) for field in fields(value)}
if isinstance(value, dict):
return {str(key): _json_ready(item) for key, item in value.items()}
if isinstance(value, list | tuple | set):
return [_json_ready(item) for item in value]
return value
def _proxy_config_payload(config: ProxyConfig) -> dict[str, Any]:
payload: dict[str, Any] = {}
for field in fields(config):
value = _json_ready(getattr(config, field.name))
try:
json.dumps(value)
except TypeError:
continue
payload[field.name] = value
return payload
def _proxy_config_from_env() -> ProxyConfig:
raw_config = os.environ.get(_MULTI_WORKER_CONFIG_ENV)
if raw_config:
try:
return ProxyConfig(**json.loads(raw_config))
except (TypeError, ValueError, json.JSONDecodeError):
logger.warning(
"Invalid %s; falling back to HEADROOM_* env vars", _MULTI_WORKER_CONFIG_ENV
)
return ProxyConfig(
host=_get_env_str("HEADROOM_HOST", "127.0.0.1"),
port=_get_env_int("HEADROOM_PORT", 8787),
openai_api_url=os.environ.get("OPENAI_TARGET_API_URL"),
anthropic_api_url=os.environ.get("ANTHROPIC_TARGET_API_URL"),
anthropic_buffered_request_timeout_seconds=_get_env_int(
"HEADROOM_ANTHROPIC_BUFFERED_REQUEST_TIMEOUT_SECONDS",
600,
min_value=1,
),
vertex_api_url=os.environ.get("VERTEX_TARGET_API_URL"),
backend=_get_env_str("HEADROOM_BACKEND", "anthropic"),
bedrock_region=_get_env_str("HEADROOM_BEDROCK_REGION", "us-west-2"),
bedrock_profile=os.environ.get("AWS_PROFILE"),
bedrock_api_url=os.environ.get("BEDROCK_TARGET_API_URL"),
anyllm_provider=_get_env_str("HEADROOM_ANYLLM_PROVIDER", "openai"),
disable_kompress=_get_env_bool("HEADROOM_DISABLE_KOMPRESS", False),
disable_kompress_fallback=_get_env_bool("HEADROOM_DISABLE_KOMPRESS_FALLBACK", False),
disable_kompress_anthropic=_get_env_optional_bool("HEADROOM_DISABLE_KOMPRESS_ANTHROPIC"),
disable_kompress_openai=_get_env_optional_bool("HEADROOM_DISABLE_KOMPRESS_OPENAI"),
force_kompress_all=_get_env_bool("HEADROOM_FORCE_KOMPRESS_ALL", False),
lossless=_get_env_bool("HEADROOM_LOSSLESS", False),
max_connections=_get_env_int("HEADROOM_MAX_CONNECTIONS", 500),
max_keepalive_connections=_get_env_int("HEADROOM_MAX_KEEPALIVE", 100),
keepalive_expiry=_get_env_float("HEADROOM_KEEPALIVE_EXPIRY", 90.0),
http2=_get_env_bool("HEADROOM_HTTP2", True),
http_proxy=os.environ.get("HEADROOM_HTTP_PROXY") or None,
periodic_toin_stats_enabled=_get_env_bool("HEADROOM_PERIODIC_TOIN_STATS", True),
proxy_token=os.environ.get("HEADROOM_PROXY_TOKEN") or None,
offline=_get_env_bool("HEADROOM_OFFLINE", False),
# Default mode is CACHE (Headroom's coding posture): delta-only compression
# at ~0 prefix-cache busts. HEADROOM_MODE overrides.
mode=normalize_proxy_mode(_get_env_str("HEADROOM_MODE", PROXY_MODE_CACHE)),
# Default savings profile is "coding" so proxy_pipeline_kwargs applies its
# posture (compress_user, protect_recent, min_tokens). HEADROOM_SAVINGS_PROFILE
# overrides.
savings_profile=os.environ.get("HEADROOM_SAVINGS_PROFILE") or "coding",
read_maturation=_get_env_bool("HEADROOM_READ_MATURATION", False),
read_maturation_quiesce_turns=_get_env_int("HEADROOM_READ_MATURATION_QUIESCE_TURNS", 5),
read_maturation_max_hold_turns=_get_env_int("HEADROOM_READ_MATURATION_MAX_HOLD_TURNS", 25),
read_maturation_min_size_bytes=_get_env_int(
"HEADROOM_READ_MATURATION_MIN_SIZE_BYTES", 2048
),
)
def create_app_from_env() -> FastAPI:
# Seed the coding-profile defaults into the process env BEFORE reading config,
# so the uvicorn factory launch gets Headroom's out-of-box posture (cache mode,
# tool-search, dedupe, read protection, …). setdefault → explicit env wins.
from headroom.agent_savings import seed_proxy_env_defaults
seed_proxy_env_defaults()
return create_app(_proxy_config_from_env())
def _get_code_aware_banner_status(config: ProxyConfig) -> str:
"""Get code-aware compression status line for banner."""
if config.code_aware_enabled:
if is_tree_sitter_available():
return "ENABLED (AST-based)"
else:
return "NOT INSTALLED (pip install headroom-ai[code])"
else:
if is_tree_sitter_available():
return "DISABLED (--code-aware or HEADROOM_CODE_AWARE_ENABLED=1 to enable)"
return "DISABLED (install headroom-ai[code] to enable)"
def run_server(
config: ProxyConfig | None = None,
workers: int = 1,
limit_concurrency: int = 1000,
print_banner: bool = True,
):
"""Run the proxy server.
Args:
config: Proxy configuration
workers: Number of worker processes (use N for multi-core scaling)
limit_concurrency: Max concurrent connections before 503 response
print_banner: When False, skip the legacy ASCII banner. The
Click CLI (`headroom proxy`) prints its own startup banner
before calling this — printing a second banner here is the
"dual banner" UX issue. Direct `python -m headroom.proxy.server`
still gets the banner since it has no other startup output.
"""
if not FASTAPI_AVAILABLE:
print("ERROR: FastAPI required. Install: pip install fastapi uvicorn httpx")
sys.exit(1)
# Seed the request-time coding-profile toggles (tool-search, dedupe, read
# protection, lossless→lossy, effort-router, block-char floor) into the
# process env before serving, so downstream per-request readers pick them up.
# Done here (not in the CLI command) so unit tests that mock run_server never
# mutate os.environ. setdefault → explicit env still wins. MODE / profile are
# already resolved into `config` above via their inline defaults.
from headroom.agent_savings import seed_proxy_env_defaults
seed_proxy_env_defaults()
config = config or ProxyConfig()
code_aware_status = _get_code_aware_banner_status(config)
# Format connection pool info
pool_info = f"max={config.max_connections}, keepalive={config.max_keepalive_connections}"
http2_status = "ENABLED" if (config.http2 and not config.http_proxy) else "DISABLED"
backend_status = format_backend_status(
backend=config.backend,
anyllm_provider=config.anyllm_provider,
bedrock_region=config.bedrock_region,
)
# Resolve upstream API targets for display in the banner (#583).
api_targets = resolve_api_targets(config.provider_api_overrides)
if print_banner:
print(f"""
╔══════════════════════════════════════════════════════════════════════╗
║ HEADROOM PROXY SERVER ║
╠══════════════════════════════════════════════════════════════════════╣
║ Version: 1.0.0 ║
║ Listening: http://{config.host}:{config.port:<5}
║ Workers: {workers:<3} Concurrency Limit: {limit_concurrency:<5}
║ Backend: {backend_status:<59}
╠══════════════════════════════════════════════════════════════════════╣
║ UPSTREAM TARGETS: ║
║ Anthropic: {api_targets.anthropic:<57}
║ OpenAI: {api_targets.openai:<57}
║ Gemini: {api_targets.gemini:<57}
║ Cloud Code: {api_targets.cloudcode:<57}
║ Vertex AI: {api_targets.vertex:<57}
╠══════════════════════════════════════════════════════════════════════╣
║ FEATURES: ║
║ Optimization: {"ENABLED " if config.optimize else "DISABLED"}
║ Caching: {"ENABLED " if config.cache_enabled else "DISABLED"} (TTL: {config.cache_ttl_seconds}s) ║
║ Rate Limiting: {"ENABLED " if config.rate_limit_enabled else "DISABLED"} ({config.rate_limit_requests_per_minute} req/min, {config.rate_limit_tokens_per_minute:,} tok/min) ║
║ Retry: {"ENABLED " if config.retry_enabled else "DISABLED"} (max {config.retry_max_attempts} attempts) ║
║ Cost Tracking: {"ENABLED " if config.cost_tracking_enabled else "DISABLED"} (budget: {"$" + str(config.budget_limit_usd) + "/" + config.budget_period if config.budget_limit_usd else "unlimited"}) ║
║ Code-Aware: {code_aware_status:<52}
║ HTTP/2: {http2_status:<52}
║ Conn Pool: {pool_info:<52}
╠══════════════════════════════════════════════════════════════════════╣
║ USAGE: ║
║ Claude Code: ANTHROPIC_BASE_URL=http://{config.host}:{config.port} claude ║
║ Cursor: Set base URL in settings ║
╠══════════════════════════════════════════════════════════════════════╣
║ ENDPOINTS: ║
║ /livez Process liveness ║
║ /readyz Traffic readiness ║
║ /health Aggregate health ║
║ /stats Detailed statistics ║
║ /metrics Prometheus metrics ║
║ /cache/clear Clear response cache ║
║ /v1/retrieve CCR: Retrieve compressed content ║
║ /v1/retrieve/stats CCR: Compression store stats ║
║ /v1/retrieve/tool_call CCR: Handle LLM tool calls ║
║ /v1/feedback CCR: Feedback loop stats & patterns ║
║ /v1/feedback/{{tool}} CCR: Compression hints for a tool ║
║ /v1/telemetry Data flywheel: Telemetry stats ║
║ /v1/telemetry/export Data flywheel: Export for aggregation ║
║ /v1/telemetry/tools Data flywheel: Per-tool stats ║
║ /v1/toin/stats TOIN: Overall intelligence stats ║
║ /v1/toin/patterns TOIN: List learned patterns ║
║ /v1/toin/pattern/{{hash}} TOIN: Pattern details by hash ║
╚══════════════════════════════════════════════════════════════════════╝
""")
app_target: Any
uvicorn_kwargs: dict[str, Any] = {}
if sys.platform == "win32":
# ProactorEventLoop can close the listening socket on transient
# AcceptEx failures (for example WinError 64 from keep-alive RSTs).
# The selector loop keeps accept errors scoped to the connection.
uvicorn_kwargs["loop"] = "asyncio:SelectorEventLoop"
if workers > 1:
# CompressionCache and PrefixTracker are always per-worker instance vars.
# Python CompressionStore defaults to InMemoryBackend (per-process), so
# CCR markers written on worker A are invisible to worker B unless a
# cross-worker backend is configured via HEADROOM_CCR_BACKEND.
# See RUST_DEV.md -> "Multi-worker deployment -- CCR fragmentation".
if os.environ.get("HEADROOM_CCR_BACKEND", "").strip():
logger.warning(
"Headroom is running with workers=%d. Compression cache, "
"prefix tracker, TOIN state, and CostTracker are all per-process; "
"multi-worker deployments produce avoidable cache busts and an "
"unstable dashboard 'Proxy $ Saved' hero tile (each /stats poll "
"hits a different worker's partial total) when sessions land on "
"different workers. Run --workers 1 or place a sticky-session load "
"balancer in front of multiple --workers 1 processes. "
"See RUST_DEV.md -> 'Multi-worker deployment -- CCR fragmentation'.",
workers,
)
else:
logger.warning(
"Headroom is running with workers=%d. The in-memory CCR store, "
"compression cache (incl. off-path background compression), prefix "
"tracker, TOIN state, and CostTracker are all "
"per-process; multi-worker deployments produce silent CCR retrieval "
"failures, avoidable cache busts, and an unstable dashboard 'Proxy $ Saved' "
"hero tile (each /stats poll hits a different worker's partial total) when "
"sessions land on different workers. Set HEADROOM_CCR_BACKEND=sqlite for a "
"persistent cross-worker CCR store, run --workers 1, or place a "
"sticky-session load balancer in front of multiple --workers 1 processes. "
"See RUST_DEV.md -> 'Multi-worker deployment -- CCR fragmentation'.",
workers,
)
os.environ[_MULTI_WORKER_CONFIG_ENV] = json.dumps(_proxy_config_payload(config))
app_target = "headroom.proxy.server:create_app_from_env"
uvicorn_kwargs["factory"] = True
else:
app_target = create_app(config)
uvicorn.run(
app_target,
host=config.host,
port=config.port,
log_level="warning",
workers=workers if workers > 1 else None, # None = single process (default)
limit_concurrency=limit_concurrency,
# Defense-in-depth: the loopback guard for /debug/* endpoints trusts
# request.client.host. uvicorn's ProxyHeadersMiddleware rewrites that
# from X-Forwarded-For when FORWARDED_ALLOW_IPS is broader than the
# default. Disabling proxy_headers here guarantees the guard sees the
# real peer address regardless of env.
proxy_headers=False,
**uvicorn_kwargs,
)
def _get_env_bool(name: str, default: bool) -> bool:
"""Get boolean from environment variable."""
val = os.environ.get(name)
if val is None:
return default
return val.lower() in ("true", "1", "yes", "on")
def _get_env_optional_bool(name: str) -> bool | None:
"""Tristate boolean env var: unset/empty -> None, truthy -> True, falsy -> False."""
val = os.environ.get(name)
if val is None or val == "":
return None
return val.lower() in ("true", "1", "yes", "on")
def _get_env_int(name: str, default: int, *, min_value: int | None = None) -> int:
"""Get integer from environment variable."""
val = os.environ.get(name)
if val is None:
return default
try:
parsed = int(val)
except ValueError:
return default
if min_value is not None and parsed < min_value:
return default
return parsed
def _positive_int_arg(value: str) -> int:
parsed = int(value)
if parsed < 1:
raise argparse.ArgumentTypeError("must be >= 1")
return parsed
def _get_env_float(name: str, default: float) -> float:
"""Get float from environment variable."""
val = os.environ.get(name)
if val is None:
return default
try:
return float(val)
except ValueError:
return default
def _get_env_str(name: str, default: str) -> str:
"""Get string from environment variable."""
return os.environ.get(name, default)
def _parse_exclude_tools(cli_excludes: str | None) -> set[str]:
"""Parse extra never-compress tool names from CLI args and env var.
Both --exclude-tools and HEADROOM_EXCLUDE_TOOLS are comma-separated
(e.g. "WebSearch,WebFetch"). Each name is added in both original and
lowercase form for case-insensitive matching, mirroring
DEFAULT_EXCLUDE_TOOLS. Unset/empty -> empty set (DEFAULT_EXCLUDE_TOOLS
used unchanged). Entries may contain glob patterns (e.g. "mcp__*"); see
config.is_tool_excluded for the matching semantics.
"""
raw = ",".join(s for s in (cli_excludes, os.environ.get("HEADROOM_EXCLUDE_TOOLS")) if s)
names: set[str] = set()
for entry in raw.split(","):
name = entry.strip()
if name:
names.add(name)
names.add(name.lower())
return names
def _parse_csv_tools(raw: str | None) -> set[str]:
"""Parse a bare CSV tool-name string without merging HEADROOM_EXCLUDE_TOOLS."""
names: set[str] = set()
if not raw:
return names
for entry in raw.split(","):
name = entry.strip()
if name:
names.add(name)
names.add(name.lower())
return names
def _parse_tool_profiles(cli_profiles: list[str]) -> dict[str, Any]:
"""Parse tool profiles from CLI args and HEADROOM_TOOL_PROFILES env var.
Format: ToolName:level (e.g., Grep:conservative, Bash:moderate)
Env var format: comma-separated (e.g., "Grep:conservative,Bash:moderate")
Returns:
Dict mapping tool names to CompressionProfile instances.
"""
from headroom.config import PROFILE_PRESETS, CompressionProfile
profiles: dict[str, CompressionProfile] = {}
raw_entries: list[str] = list(cli_profiles)
# Also check env var
env_val = os.environ.get("HEADROOM_TOOL_PROFILES", "")
if env_val:
raw_entries.extend(e.strip() for e in env_val.split(",") if e.strip())
for entry in raw_entries:
if ":" not in entry:
logger.warning("Invalid tool profile format (expected ToolName:level): %s", entry)
continue
tool_name, level = entry.split(":", 1)
tool_name = tool_name.strip()
level = level.strip().lower()
if level in PROFILE_PRESETS:
profiles[tool_name] = PROFILE_PRESETS[level]
else:
logger.warning(
"Unknown profile level '%s' for tool '%s'. Use: conservative, moderate, aggressive",
level,
tool_name,
)
return profiles
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Headroom Proxy Server")
# Server
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=8787)
parser.add_argument(
"--openai-api-url", help=f"Custom OpenAI API URL (default: {DEFAULT_OPENAI_API_URL})"
)
parser.add_argument(
"--anthropic-api-url",
help=f"Custom Anthropic API URL (default: {DEFAULT_ANTHROPIC_API_URL})",
)
parser.add_argument(
"--anthropic-buffered-request-timeout-seconds",
type=_positive_int_arg,
default=600,
help=(
"Anthropic buffered read timeout in seconds for non-streaming "
"message and batch paths (default: 600)"
),
)
parser.add_argument(
"--vertex-api-url",
help=f"Custom Vertex AI regional API URL (default: {DEFAULT_VERTEX_API_URL})",
)
# Backend (anthropic direct, bedrock, openrouter, anyllm, or litellm-<provider>)
parser.add_argument(
"--backend",
default="anthropic",
help=(
"Backend: 'anthropic' (direct), 'bedrock' (AWS), 'openrouter', "
"'anyllm' (any-llm), or 'litellm-<provider>' (e.g., litellm-hosted_vllm, litellm-vertex)"
),
)
parser.add_argument(
"--bedrock-region",
default="us-west-2",
help="AWS region for Bedrock backend (default: us-west-2)",
)
parser.add_argument(
"--bedrock-profile",
help="AWS profile for Bedrock backend (default: use default credentials)",
)
parser.add_argument(
"--bedrock-api-url",
help=(
"Custom Bedrock InvokeModel upstream for the /model/{id}/invoke "
"passthrough routes — point at a re-signing gateway, not raw AWS "
"(env: BEDROCK_TARGET_API_URL)"
),
)
parser.add_argument(
"--openrouter-api-key",
help="OpenRouter API key (or set OPENROUTER_API_KEY env var)",
)
parser.add_argument(
"--anyllm-provider",
default="openai",
help="any-llm provider: openai, anthropic, mistral, groq, ollama, bedrock, etc. (default: openai)",
)
# Connection pool (scalability)
parser.add_argument(
"--max-connections",
type=int,
default=500,
help="Max connections to upstream APIs (default: 500)",
)
parser.add_argument(
"--max-keepalive", type=int, default=100, help="Max keepalive connections (default: 100)"
)
parser.add_argument(
"--keepalive-expiry",
type=float,
default=90.0,
help="Seconds an idle upstream keep-alive connection is kept open (default: 90)",
)
parser.add_argument(
"--no-http2",
action="store_true",
help="Disable HTTP/2 (enabled by default for better throughput)",
)
parser.add_argument(
"--http-proxy",
help=("HTTP proxy URL for upstream provider requests only (env: HEADROOM_HTTP_PROXY)"),
)
parser.add_argument(
"--workers",
type=int,
default=1,
help="Number of worker processes (default: 1, use N for multi-core)",
)
parser.add_argument(
"--limit-concurrency",
type=int,
default=1000,
help="Max concurrent connections before 503 (default: 1000)",
)
# Optimization
parser.add_argument("--no-optimize", action="store_true", help="Disable optimization")
parser.add_argument("--min-tokens", type=int, default=500, help="Min tokens to crush")
parser.add_argument("--max-items", type=int, default=50, help="Max items after crush")
parser.add_argument(
"--tool-profile",
action="append",
default=[],
help="Per-tool compression profile: ToolName:level (e.g., Grep:conservative, Bash:moderate, WebFetch:aggressive). "
"Can be specified multiple times. Also settable via HEADROOM_TOOL_PROFILES env var.",
)
parser.add_argument(
"--compress-user-messages",
action="store_true",
help=(
"Opt in to compressing `user` role messages. Default is off because "
"user content is typically the subject of the request and is part of "
"the prefix-cache zone. Enable this for OpenAI/Azure chat workloads "
"where the bulk of input lives in user messages (pasted content, "
"RAG context, etc.) and you want the router to consider it eligible. "
"Also settable via HEADROOM_COMPRESS_USER_MESSAGES=1."
),
)
parser.add_argument(
"--disable-kompress",
action="store_true",
help=(
"Disable Kompress ML compression while keeping structural compression enabled. "
"Also settable via HEADROOM_DISABLE_KOMPRESS=1."
),
)
parser.add_argument(
"--disable-kompress-fallback",
action="store_true",
help=(
"With --disable-kompress, route fall-through content to PASSTHROUGH instead of "
"the default KOMPRESS fallback (restores legacy --disable-kompress behaviour). "
"Also settable via HEADROOM_DISABLE_KOMPRESS_FALLBACK=1."
),
)
parser.add_argument(
"--disable-kompress-anthropic",
dest="disable_kompress_anthropic",
action="store_const",
const=True,
default=None,
help=(
"Disable Kompress for the Anthropic pipeline only, overriding --disable-kompress. "
"Also settable via HEADROOM_DISABLE_KOMPRESS_ANTHROPIC=1."
),
)
parser.add_argument(
"--enable-kompress-anthropic",
dest="disable_kompress_anthropic",
action="store_const",
const=False,
help="Force-enable Kompress for the Anthropic pipeline, overriding --disable-kompress.",
)
parser.add_argument(
"--disable-kompress-openai",
dest="disable_kompress_openai",
action="store_const",
const=True,
default=None,
help=(
"Disable Kompress for the OpenAI/Codex pipeline only, overriding --disable-kompress. "
"Also settable via HEADROOM_DISABLE_KOMPRESS_OPENAI=1."
),
)
parser.add_argument(
"--enable-kompress-openai",
dest="disable_kompress_openai",
action="store_const",
const=False,
help="Force-enable Kompress for the OpenAI/Codex pipeline, overriding --disable-kompress.",
)
parser.add_argument(
"--force-kompress-all",
action="store_true",
help=(
"Route ALL compressible content through Kompress (kompress-v2-base), "
"bypassing per-type compressor selection. Tool ground truth "
"(Read/Glob/... and reversibility-gated output) is still never touched. "
"Also settable via HEADROOM_FORCE_KOMPRESS_ALL=1."
),
)
parser.add_argument(
"--lossless",
action="store_true",
help=(
"No-CCR lossless mode: compress LOG/SEARCH/DIFF tool outputs with "
"format-native lossless compaction (and marker-free SmartCrusher) "
"without emitting any CCR retrieval marker, so no MCP retrieve tool "
"is needed. Also settable via HEADROOM_LOSSLESS=1."
),
)
parser.add_argument(
"--exclude-tools",
default=None,
help="Comma-separated tool names whose output is never compressed, "
"merged with the built-in defaults (e.g., WebSearch,WebFetch). "
"Entries may use glob patterns, e.g. 'mcp__*' to exclude every MCP tool. "
"Also settable via HEADROOM_EXCLUDE_TOOLS env var.",
)
parser.add_argument(
"--protect-tool-results",
default=None,
help="Comma-separated tool names whose results are never lossy-compressed, "
"merged with the built-in defaults (e.g. Bash,WebFetch). "
"Also settable via HEADROOM_PROTECT_TOOL_RESULTS env var.",
)
# Caching
parser.add_argument("--no-cache", action="store_true", help="Disable caching")
parser.add_argument("--cache-ttl", type=int, default=3600, help="Cache TTL seconds")
# Rate limiting
parser.add_argument("--no-rate-limit", action="store_true", help="Disable rate limiting")
parser.add_argument("--rpm", type=int, default=60, help="Requests per minute")
parser.add_argument("--tpm", type=int, default=100000, help="Tokens per minute")
# Cost
parser.add_argument("--budget", type=float, help="Budget limit in USD")
parser.add_argument("--budget-period", choices=["hourly", "daily", "monthly"], default="daily")
# Logging
parser.add_argument("--log-file", help="Log file path")
parser.add_argument("--log-messages", action="store_true", help="Log full messages")
# Code-aware compression
parser.add_argument(
"--code-aware",
action="store_true",
help="Enable AST-based code compression (requires: pip install headroom-ai[code])",
)
parser.add_argument(
"--no-code-aware",
action="store_true",
help="Disable code-aware compression",
)
args = parser.parse_args()
# Environment variable defaults (HEADROOM_* prefix)
# CLI args override env vars, env vars override ProxyConfig defaults
env_code_aware = _get_env_bool("HEADROOM_CODE_AWARE_ENABLED", True)
env_optimize = _get_env_bool("HEADROOM_OPTIMIZE", True)
env_cache = _get_env_bool("HEADROOM_CACHE_ENABLED", True)
env_rate_limit = _get_env_bool("HEADROOM_RATE_LIMIT_ENABLED", True)
# Determine settings: CLI flags override env vars
# --no-X explicitly disables, --X explicitly enables, neither uses env var
code_aware_enabled = (
env_code_aware
if not (args.code_aware or args.no_code_aware)
else (args.code_aware or not args.no_code_aware)
)
optimize = env_optimize if not args.no_optimize else False
cache_enabled = env_cache if not args.no_cache else False
rate_limit_enabled = env_rate_limit if not args.no_rate_limit else False
disable_kompress = args.disable_kompress or _get_env_bool("HEADROOM_DISABLE_KOMPRESS", False)
disable_kompress_fallback = args.disable_kompress_fallback or _get_env_bool(
"HEADROOM_DISABLE_KOMPRESS_FALLBACK", False
)
disable_kompress_anthropic = (
args.disable_kompress_anthropic
if args.disable_kompress_anthropic is not None
else _get_env_optional_bool("HEADROOM_DISABLE_KOMPRESS_ANTHROPIC")
)
disable_kompress_openai = (
args.disable_kompress_openai
if args.disable_kompress_openai is not None
else _get_env_optional_bool("HEADROOM_DISABLE_KOMPRESS_OPENAI")
)
force_kompress_all = args.force_kompress_all or _get_env_bool(
"HEADROOM_FORCE_KOMPRESS_ALL", False
)
lossless = getattr(args, "lossless", False) or _get_env_bool("HEADROOM_LOSSLESS", False)
# Set OpenRouter API key from CLI if provided
if hasattr(args, "openrouter_api_key") and args.openrouter_api_key:
os.environ["OPENROUTER_API_KEY"] = args.openrouter_api_key
# Parse per-tool compression profiles from CLI and env var
tool_profiles = _parse_tool_profiles(args.tool_profile)
# Parse extra never-compress tools from CLI and env var
exclude_tools = _parse_exclude_tools(args.exclude_tools)
protect_tool_results = _parse_csv_tools(
args.protect_tool_results or os.environ.get("HEADROOM_PROTECT_TOOL_RESULTS")
)
config = ProxyConfig(
host=_get_env_str("HEADROOM_HOST", args.host),
port=_get_env_int("HEADROOM_PORT", args.port),
openai_api_url=_get_env_str("OPENAI_TARGET_API_URL", args.openai_api_url),
anthropic_api_url=_get_env_str("ANTHROPIC_TARGET_API_URL", args.anthropic_api_url),
anthropic_buffered_request_timeout_seconds=_get_env_int(
"HEADROOM_ANTHROPIC_BUFFERED_REQUEST_TIMEOUT_SECONDS",
args.anthropic_buffered_request_timeout_seconds,
min_value=1,
),
vertex_api_url=_get_env_str("VERTEX_TARGET_API_URL", args.vertex_api_url),
# Backend settings
backend=_get_env_str("HEADROOM_BACKEND", args.backend), # type: ignore[arg-type]
bedrock_region=_get_env_str("HEADROOM_BEDROCK_REGION", args.bedrock_region),
bedrock_profile=args.bedrock_profile or os.environ.get("AWS_PROFILE"),
bedrock_api_url=_get_env_str("BEDROCK_TARGET_API_URL", args.bedrock_api_url),
anyllm_provider=_get_env_str("HEADROOM_ANYLLM_PROVIDER", args.anyllm_provider),
optimize=optimize,
min_tokens_to_crush=_get_env_int("HEADROOM_MIN_TOKENS", args.min_tokens),
max_items_after_crush=_get_env_int("HEADROOM_MAX_ITEMS", args.max_items),
smart_crusher_with_compaction=(
_get_env_bool("HEADROOM_SMART_CRUSHER_COMPACTION", False)
if "HEADROOM_SMART_CRUSHER_COMPACTION" in os.environ
else None
),
cache_enabled=cache_enabled,
cache_ttl_seconds=_get_env_int("HEADROOM_CACHE_TTL", args.cache_ttl),
rate_limit_enabled=rate_limit_enabled,
rate_limit_requests_per_minute=_get_env_int("HEADROOM_RPM", args.rpm),
rate_limit_tokens_per_minute=_get_env_int("HEADROOM_TPM", args.tpm),
budget_limit_usd=args.budget,
budget_period=args.budget_period,
log_file=_get_env_str("HEADROOM_LOG_FILE", args.log_file)
if args.log_file
else os.environ.get("HEADROOM_LOG_FILE"),
log_full_messages=args.log_messages or _get_env_bool("HEADROOM_LOG_MESSAGES", False),
code_aware_enabled=code_aware_enabled,
disable_kompress=disable_kompress,
disable_kompress_fallback=disable_kompress_fallback,
disable_kompress_anthropic=disable_kompress_anthropic,
disable_kompress_openai=disable_kompress_openai,
force_kompress_all=force_kompress_all,
lossless=lossless,
# Connection pool settings
max_connections=_get_env_int("HEADROOM_MAX_CONNECTIONS", args.max_connections),
max_keepalive_connections=_get_env_int("HEADROOM_MAX_KEEPALIVE", args.max_keepalive),
keepalive_expiry=_get_env_float("HEADROOM_KEEPALIVE_EXPIRY", args.keepalive_expiry),
http2=not args.no_http2 and _get_env_bool("HEADROOM_HTTP2", True),
http_proxy=_get_env_str("HEADROOM_HTTP_PROXY", args.http_proxy or "") or None,
read_maturation=_get_env_bool("HEADROOM_READ_MATURATION", False),
read_maturation_quiesce_turns=_get_env_int("HEADROOM_READ_MATURATION_QUIESCE_TURNS", 5),
read_maturation_max_hold_turns=_get_env_int("HEADROOM_READ_MATURATION_MAX_HOLD_TURNS", 25),
read_maturation_min_size_bytes=_get_env_int(
"HEADROOM_READ_MATURATION_MIN_SIZE_BYTES", 2048
),
tool_profiles=tool_profiles if tool_profiles else None,
exclude_tools=exclude_tools if exclude_tools else None,
protect_tool_results=frozenset(protect_tool_results)
if protect_tool_results
else frozenset(),
mode=normalize_proxy_mode(_get_env_str("HEADROOM_MODE", PROXY_MODE_CACHE)),
compress_user_messages=args.compress_user_messages
or _get_env_bool("HEADROOM_COMPRESS_USER_MESSAGES", False),
savings_profile=os.environ.get("HEADROOM_SAVINGS_PROFILE") or "coding",
# Default 0.4 keep-ratio so the Kompress text (prose/code) path compresses
# meaningfully out of the box; HEADROOM_TARGET_RATIO overrides.
target_ratio=(
float(os.environ["HEADROOM_TARGET_RATIO"])
if os.environ.get("HEADROOM_TARGET_RATIO")
else 0.4
),
compress_system_messages=(
_get_env_bool("HEADROOM_COMPRESS_SYSTEM_MESSAGES", False)
if "HEADROOM_COMPRESS_SYSTEM_MESSAGES" in os.environ
else None
),
protect_recent=(
int(os.environ["HEADROOM_PROTECT_RECENT"])
if os.environ.get("HEADROOM_PROTECT_RECENT")
else None
),
protect_analysis_context=(
_get_env_bool("HEADROOM_PROTECT_ANALYSIS_CONTEXT", False)
if "HEADROOM_PROTECT_ANALYSIS_CONTEXT" in os.environ
else None
),
accuracy_guard=os.environ.get("HEADROOM_ACCURACY_GUARD") or None,
)
# Get worker and concurrency settings
workers = _get_env_int("HEADROOM_WORKERS", args.workers)
limit_concurrency = _get_env_int("HEADROOM_LIMIT_CONCURRENCY", args.limit_concurrency)
run_server(config, workers=workers, limit_concurrency=limit_concurrency)