From aa99db8f06003214a30775e8521382dc2c7a42a1 Mon Sep 17 00:00:00 2001 From: JerrettDavis Date: Fri, 10 Apr 2026 12:09:10 -0500 Subject: [PATCH 1/3] feat: add proxy healthcheck endpoints Add /livez and /readyz, keep /health backward-compatible, and wire readiness into Docker artifacts and docs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- CHANGELOG.md | 5 + Dockerfile | 3 + docker-compose.yml | 6 + docs/proxy.md | 62 ++++++++- headroom/cli/proxy.py | 4 +- headroom/memory/backends/direct_mem0.py | 4 + headroom/proxy/memory_handler.py | 15 ++ headroom/proxy/server.py | 175 ++++++++++++++++++++---- tests/test_proxy_healthchecks.py | 129 +++++++++++++++++ 9 files changed, 371 insertions(+), 32 deletions(-) create mode 100644 tests/test_proxy_healthchecks.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 50d9aa3bc..5845da1ba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- **Proxy liveness and readiness health checks** + - Adds `GET /livez` for process liveness and `GET /readyz` for traffic readiness + - Keeps `GET /health` backward compatible while expanding it with readiness details and subsystem checks + - Eagerly initializes configured memory backends during proxy startup so readiness reflects real serving capability + - Wires `/readyz` into the Docker image `HEALTHCHECK` and the example `docker-compose.yml` - **Durable proxy savings history** - Persists proxy compression savings history locally at `~/.headroom/proxy_savings.json` - Supports `HEADROOM_SAVINGS_PATH` to override the storage location diff --git a/Dockerfile b/Dockerfile index 6a5981a6a..d374c4eb7 100644 --- a/Dockerfile +++ b/Dockerfile @@ -60,6 +60,9 @@ ENV HEADROOM_HOST=0.0.0.0 \ EXPOSE 8787 +HEALTHCHECK --interval=30s --timeout=5s --start-period=20s --retries=3 \ + CMD ["curl", "--fail", "--silent", "http://127.0.0.1:8787/readyz"] + ENTRYPOINT ["headroom", "proxy"] CMD ["--host", "0.0.0.0", "--port", "8787"] diff --git a/docker-compose.yml b/docker-compose.yml index 66331836d..633ac0ad9 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -9,6 +9,12 @@ services: # - OPENAI_TARGET_API_URL=https://api.x.ai ports: - "8787:8787" + healthcheck: + test: ["CMD", "curl", "--fail", "--silent", "http://127.0.0.1:8787/readyz"] + interval: 30s + timeout: 5s + retries: 3 + start_period: 20s depends_on: - qdrant - neo4j diff --git a/docs/proxy.md b/docs/proxy.md index d33b7ea3a..e479b295e 100644 --- a/docs/proxy.md +++ b/docs/proxy.md @@ -125,7 +125,52 @@ headroom proxy --llmlingua --llmlingua-rate 0.5 ## API Endpoints -### Health Check +### Liveness + +```bash +curl http://localhost:8787/livez +``` + +Response: +```json +{ + "service": "headroom-proxy", + "status": "healthy", + "alive": true, + "version": "0.5.21", + "timestamp": "2026-04-10T16:36:25Z", + "uptime_seconds": 12.483 +} +``` + +### Readiness + +```bash +curl http://localhost:8787/readyz +``` + +Response: +```json +{ + "service": "headroom-proxy", + "status": "healthy", + "ready": true, + "version": "0.5.21", + "timestamp": "2026-04-10T16:36:25Z", + "uptime_seconds": 12.483, + "checks": { + "startup": {"enabled": true, "ready": true, "status": "healthy"}, + "http_client": {"enabled": true, "ready": true, "status": "healthy"}, + "cache": {"enabled": true, "ready": true, "status": "healthy"}, + "rate_limiter": {"enabled": true, "ready": true, "status": "healthy"}, + "memory": {"enabled": false, "ready": true, "status": "disabled"} + } +} +``` + +`/readyz` returns HTTP 503 when Headroom has not completed startup or a required enabled subsystem is unavailable. This is the endpoint used by the container health checks. + +### Aggregate Health ```bash curl http://localhost:8787/health @@ -135,11 +180,16 @@ Response: ```json { "status": "healthy", - "optimize": true, - "stats": { - "total_requests": 42, - "tokens_saved": 15000, - "savings_percent": 45.2 + "ready": true, + "version": "0.5.21", + "config": { + "optimize": true, + "cache": true, + "rate_limit": true + }, + "checks": { + "startup": {"enabled": true, "ready": true, "status": "healthy"}, + "http_client": {"enabled": true, "ready": true, "status": "healthy"} } } ``` diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index 43ab571c8..80cc83ace 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -371,7 +371,9 @@ Usage: Codex / OpenAI: OPENAI_BASE_URL=http://{config.host}:{config.port}/v1 your-app {memory_section} Endpoints: - GET /health Health check + GET /livez Process liveness + GET /readyz Traffic readiness + GET /health Aggregate health GET /stats Detailed statistics GET /stats-history Durable compression history + display session GET /metrics Prometheus metrics diff --git a/headroom/memory/backends/direct_mem0.py b/headroom/memory/backends/direct_mem0.py index 7d07ec0de..39b8124c1 100644 --- a/headroom/memory/backends/direct_mem0.py +++ b/headroom/memory/backends/direct_mem0.py @@ -226,6 +226,10 @@ class DirectMem0Adapter: self._initialized = True + async def ensure_initialized(self) -> None: + """Public initialization hook for callers that need readiness guarantees.""" + await self._ensure_initialized() + def _embed(self, text: str) -> list[float]: """Generate embedding for text using OpenAI.""" response = self._openai_client.embeddings.create( diff --git a/headroom/proxy/memory_handler.py b/headroom/proxy/memory_handler.py index 2c3b10600..77a3bc5ea 100644 --- a/headroom/proxy/memory_handler.py +++ b/headroom/proxy/memory_handler.py @@ -156,6 +156,7 @@ class MemoryHandler: enable_graph=True, ) self._backend = DirectMem0Adapter(mem0_config) + await self._backend.ensure_initialized() logger.info( f"Memory: Initialized Qdrant+Neo4j backend " f"({self.config.qdrant_host}:{self.config.qdrant_port})" @@ -1422,6 +1423,20 @@ To SAVE: create /memories/.txt "content" """Whether the backend has been initialized.""" return self._initialized + async def ensure_initialized(self) -> None: + """Initialize the configured backend so readiness checks can be accurate.""" + await self._ensure_initialized() + + def health_status(self) -> dict[str, Any]: + """Return a lightweight health snapshot for readiness endpoints.""" + return { + "enabled": self.config.enabled, + "backend": self.config.backend, + "initialized": self._initialized, + "native_tool": self.config.use_native_tool, + "bridge_enabled": self.config.bridge_enabled, + } + async def close(self) -> None: """Close the memory backend.""" if self._backend and hasattr(self._backend, "close"): diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 1767a480d..23cf456c6 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -31,6 +31,7 @@ import os import random import sys import time +from datetime import datetime, timezone from pathlib import Path from typing import TYPE_CHECKING, Any, Literal @@ -45,7 +46,7 @@ try: import uvicorn from fastapi import FastAPI, HTTPException, Request, Response, WebSocket from fastapi.middleware.cors import CORSMiddleware - from fastapi.responses import HTMLResponse, PlainTextResponse + from fastapi.responses import HTMLResponse, JSONResponse, PlainTextResponse FASTAPI_AVAILABLE = True except ImportError: @@ -601,6 +602,16 @@ class HeadroomProxy( if eager_status.get("magika") == "enabled": logger.info("Magika: ENABLED (ML content detection)") + if self.memory_handler: + await self.memory_handler.ensure_initialized() + memory_status = self.memory_handler.health_status() + 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: @@ -630,6 +641,10 @@ class HeadroomProxy( """Cleanup async resources.""" if self.http_client: await self.http_client.aclose() + self.http_client = None + + if self.memory_handler: + await self.memory_handler.close() # Print final stats self._print_summary() @@ -916,24 +931,34 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: LangfuseTracingConfig.from_env(default_service_name="headroom-proxy") ) + app.state.started_at = time.time() + app.state.ready = False + app.state.startup_error = None + try: - # Startup - await proxy.startup() - 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() + try: + # Startup + await proxy.startup() + 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() - # Only start beacon if we acquire the lock (first worker wins) - _beacon_is_owner[0] = _try_acquire_beacon_lock() - if _beacon_is_owner[0]: - await _beacon.start() - else: - logger.debug("Beacon: skipping (another worker owns the lock)") + # Only start beacon if we acquire the lock (first worker wins) + _beacon_is_owner[0] = _try_acquire_beacon_lock() + if _beacon_is_owner[0]: + await _beacon.start() + else: + logger.debug("Beacon: skipping (another worker owns the lock)") - yield + app.state.ready = True + yield + except Exception as exc: + app.state.startup_error = str(exc) + raise finally: + app.state.ready = False # Shutdown if _beacon_is_owner[0]: await _beacon.stop() @@ -953,6 +978,92 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: lifespan=lifespan, ) app.state.proxy = proxy + app.state.started_at = None + app.state.ready = False + app.state.startup_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 started_at is None: + return 0.0 + return round(max(0.0, time.time() - 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, + } + ) + 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_status["enabled"], + ready=memory_status["initialized"], + backend=memory_status["backend"], + initialized=memory_status["initialized"], + native_tool=memory_status["native_tool"], + bridge_enabled=memory_status["bridge_enabled"], + ), + } + + def _health_payload(*, include_config: bool) -> dict[str, Any]: + checks = _health_checks() + ready = all(check["ready"] for check in checks.values()) + 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, + } + if include_config: + payload["config"] = { + "optimize": config.optimize, + "cache": config.cache_enabled, + "rate_limit": config.rate_limit_enabled, + } + return payload # CORS app.add_middleware( @@ -964,17 +1075,29 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: ) # Health & Metrics + @app.get("/livez") + async def livez(): + return JSONResponse( + status_code=200, + content={ + "service": "headroom-proxy", + "status": "healthy", + "alive": True, + "version": __version__, + "timestamp": _iso_utc_now(), + "uptime_seconds": _uptime_seconds(), + }, + ) + + @app.get("/readyz") + async def readyz(): + payload = _health_payload(include_config=False) + return JSONResponse(status_code=200 if payload["ready"] else 503, content=payload) + @app.get("/health") async def health(): - return { - "status": "healthy", - "version": __version__, - "config": { - "optimize": config.optimize, - "cache": config.cache_enabled, - "rate_limit": config.rate_limit_enabled, - }, - } + payload = _health_payload(include_config=True) + return JSONResponse(status_code=200, content=payload) @app.get("/dashboard", response_class=HTMLResponse) async def dashboard(): @@ -2166,7 +2289,9 @@ def run_server( ║ Cursor: Set base URL in settings ║ ╠══════════════════════════════════════════════════════════════════════╣ ║ ENDPOINTS: ║ -║ /health Health check ║ +║ /livez Process liveness ║ +║ /readyz Traffic readiness ║ +║ /health Aggregate health ║ ║ /stats Detailed statistics ║ ║ /metrics Prometheus metrics ║ ║ /cache/clear Clear response cache ║ diff --git a/tests/test_proxy_healthchecks.py b/tests/test_proxy_healthchecks.py new file mode 100644 index 000000000..f9c75b1de --- /dev/null +++ b/tests/test_proxy_healthchecks.py @@ -0,0 +1,129 @@ +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from fastapi.testclient import TestClient + +from headroom.proxy.server import ProxyConfig, create_app + + +@pytest.fixture +def client(): + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + with TestClient(app) as test_client: + yield test_client + + +def test_livez_reports_process_health(client): + response = client.get("/livez") + + assert response.status_code == 200 + data = response.json() + assert data["service"] == "headroom-proxy" + assert data["status"] == "healthy" + assert data["alive"] is True + assert data["uptime_seconds"] >= 0 + + +def test_readyz_reports_core_subsystem_checks(client): + response = client.get("/readyz") + + assert response.status_code == 200 + data = response.json() + assert data["ready"] is True + assert data["status"] == "healthy" + assert "config" not in data + assert data["checks"]["startup"]["status"] == "healthy" + assert data["checks"]["http_client"]["status"] == "healthy" + assert data["checks"]["cache"]["status"] == "disabled" + assert data["checks"]["rate_limiter"]["status"] == "disabled" + assert data["checks"]["memory"]["status"] == "disabled" + + +def test_health_preserves_backwards_compatible_config_payload(client): + response = client.get("/health") + + assert response.status_code == 200 + data = response.json() + assert data["status"] == "healthy" + assert data["ready"] is True + assert data["config"] == { + "optimize": False, + "cache": False, + "rate_limit": False, + } + + +def test_health_remains_200_when_proxy_is_not_ready(client): + client.app.state.ready = False + + response = client.get("/health") + + assert response.status_code == 200 + assert response.json()["ready"] is False + + +def test_readyz_reports_memory_backend_when_enabled(tmp_path): + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + memory_enabled=True, + memory_backend="local", + memory_db_path=str(tmp_path / "headroom_memory.db"), + memory_inject_tools=True, + memory_inject_context=True, + ) + app = create_app(config) + + with TestClient(app) as client: + response = client.get("/readyz") + + assert response.status_code == 200 + data = response.json() + assert data["checks"]["memory"]["status"] == "healthy" + assert data["checks"]["memory"]["backend"] == "local" + assert data["checks"]["memory"]["initialized"] is True + + +def test_readyz_initializes_qdrant_memory_backend(monkeypatch): + from headroom.memory.backends import direct_mem0 + + init_calls: list[str] = [] + + class FakeDirectMem0Adapter: + def __init__(self, config): + self.config = config + + async def ensure_initialized(self): + init_calls.append("initialized") + + monkeypatch.setattr(direct_mem0, "DirectMem0Adapter", FakeDirectMem0Adapter) + + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + memory_enabled=True, + memory_backend="qdrant-neo4j", + ) + app = create_app(config) + + with TestClient(app) as client: + response = client.get("/readyz") + + assert response.status_code == 200 + data = response.json() + assert init_calls == ["initialized"] + assert data["checks"]["memory"]["status"] == "healthy" + assert data["checks"]["memory"]["backend"] == "qdrant-neo4j" + assert data["checks"]["memory"]["initialized"] is True From bd1ac2285535bbfb15a633fca3f84618208af736 Mon Sep 17 00:00:00 2001 From: JerrettDavis Date: Fri, 10 Apr 2026 12:17:35 -0500 Subject: [PATCH 2/3] fix: tighten healthcheck typing Address the new mypy regression in the healthcheck helpers without changing endpoint behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- headroom/proxy/server.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 23cf456c6..3f3813633 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -987,9 +987,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: def _uptime_seconds() -> float: started_at = getattr(app.state, "started_at", None) - if started_at is None: + if not isinstance(started_at, (int, float)): return 0.0 - return round(max(0.0, time.time() - started_at), 3) + return round(max(0.0, time.time() - float(started_at)), 3) def _component_health( *, @@ -1017,6 +1017,8 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: "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, @@ -1036,12 +1038,12 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: ready=(proxy.rate_limiter is not None), ), "memory": _component_health( - enabled=memory_status["enabled"], - ready=memory_status["initialized"], + enabled=memory_enabled, + ready=memory_initialized, backend=memory_status["backend"], - initialized=memory_status["initialized"], - native_tool=memory_status["native_tool"], - bridge_enabled=memory_status["bridge_enabled"], + initialized=memory_initialized, + native_tool=bool(memory_status.get("native_tool", False)), + bridge_enabled=bool(memory_status.get("bridge_enabled", False)), ), } From 2b549c5c5972773ed3e8fb54d372c2709f183eb6 Mon Sep 17 00:00:00 2001 From: JerrettDavis Date: Fri, 10 Apr 2026 12:25:53 -0500 Subject: [PATCH 3/3] fix: tolerate stubbed memory handlers on shutdown Guard proxy cleanup so tests and integrations that swap in lightweight memory-handler stubs do not fail during FastAPI lifespan shutdown. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- headroom/proxy/server.py | 2 +- tests/test_proxy_healthchecks.py | 26 ++++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 3f3813633..d8837fbaa 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -643,7 +643,7 @@ class HeadroomProxy( await self.http_client.aclose() self.http_client = None - if self.memory_handler: + if self.memory_handler and hasattr(self.memory_handler, "close"): await self.memory_handler.close() # Print final stats diff --git a/tests/test_proxy_healthchecks.py b/tests/test_proxy_healthchecks.py index f9c75b1de..d517f24f2 100644 --- a/tests/test_proxy_healthchecks.py +++ b/tests/test_proxy_healthchecks.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + import pytest pytest.importorskip("fastapi") @@ -127,3 +129,27 @@ def test_readyz_initializes_qdrant_memory_backend(monkeypatch): assert data["checks"]["memory"]["status"] == "healthy" assert data["checks"]["memory"]["backend"] == "qdrant-neo4j" assert data["checks"]["memory"]["initialized"] is True + + +def test_shutdown_tolerates_stubbed_memory_handler(): + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + + with TestClient(app) as client: + client.app.state.proxy.memory_handler = SimpleNamespace( + health_status=lambda: { + "enabled": False, + "backend": None, + "initialized": False, + "native_tool": False, + "bridge_enabled": False, + } + ) + response = client.get("/health") + + assert response.status_code == 200