from __future__ import annotations import json import time from typing import Any import pytest pytest.importorskip("fastapi") pytest.importorskip("headroom._core") from fastapi.testclient import TestClient from headroom.config import TransformResult from headroom.proxy.server import ProxyConfig, create_app def _proxy_config(**overrides: Any) -> ProxyConfig: defaults: dict[str, Any] = { "optimize": True, "cache_enabled": False, "rate_limit_enabled": False, "cost_tracking_enabled": False, "log_requests": False, "ccr_inject_tool": False, "ccr_handle_responses": False, "ccr_context_tracking": False, "image_optimize": False, "disable_kompress": True, "compression_max_workers": 1, } defaults.update(overrides) return ProxyConfig(**defaults) def test_proxy_health_surfaces_compression_runtime_metrics(monkeypatch) -> None: monkeypatch.setenv("HEADROOM_SKIP_UPSTREAM_CHECK", "1") app = create_app(_proxy_config(optimize=False)) with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client: live = client.get("/livez") health = client.get("/health") assert live.status_code == 200 assert live.json()["alive"] is True assert health.status_code == 200 runtime = health.json()["runtime"] assert runtime["compression_executor"]["max_workers"] == 1 assert runtime["compression_executor"]["queued"] == 0 assert runtime["compression_executor"]["queue_timeouts_total"] == 0 def test_v1_compress_success_reports_actual_metrics(monkeypatch) -> None: monkeypatch.setenv("HEADROOM_SKIP_UPSTREAM_CHECK", "1") app = create_app(_proxy_config()) proxy = app.state.proxy request_messages = [{"role": "user", "content": "summarize this repeated payload"}] compressed_messages = [{"role": "user", "content": "summary payload"}] def fake_apply(**kwargs): assert kwargs["messages"] == request_messages assert kwargs["model"] == "gpt-4o" return TransformResult( messages=compressed_messages, tokens_before=100, tokens_after=40, transforms_applied=["test:compress"], markers_inserted=["marker-1"], ) monkeypatch.setattr(proxy.openai_pipeline, "apply", fake_apply) with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client: response = client.post( "/v1/compress", json={"model": "gpt-4o", "messages": request_messages}, ) body = response.json() assert response.status_code == 200 assert body["messages"] == compressed_messages assert body["tokens_before"] == 100 assert body["tokens_after"] == 40 assert body["tokens_saved"] == 60 assert body["compression_ratio"] == 0.4 assert body["transforms_applied"] == ["test:compress"] assert body["transforms_summary"] == {"test:compress": 1} assert body["ccr_hashes"] == ["marker-1"] def test_v1_compress_timeout_fails_open_quickly(monkeypatch) -> None: monkeypatch.setenv("HEADROOM_SKIP_UPSTREAM_CHECK", "1") app = create_app(_proxy_config()) proxy = app.state.proxy request_messages = [{"role": "user", "content": "do not mutate me"}] async def timeout_executor(fn, *, timeout): # noqa: ANN001 raise TimeoutError("compression deadline exceeded") monkeypatch.setattr(proxy, "_run_compression_in_executor", timeout_executor) with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client: started = time.perf_counter() response = client.post( "/v1/compress", json={"model": "gpt-4o", "messages": request_messages}, ) elapsed = time.perf_counter() - started body = response.json() assert response.status_code == 200 assert elapsed < 0.5 assert body["messages"] == request_messages assert body["tokens_saved"] == 0 assert body["compression_ratio"] == 1.0 assert body["transforms_applied"] == [] assert body["compression_skipped"] is True assert body["skip_reason"] == "compression_timeout" def test_v1_compress_real_json_tool_payload_reduces_tokens(monkeypatch) -> None: monkeypatch.setenv("HEADROOM_SKIP_UPSTREAM_CHECK", "1") app = create_app( _proxy_config( ccr_inject_marker=False, min_tokens_to_crush=20, max_items_after_crush=10, ) ) items = [ { "id": i, "status": "ok", "score": i % 5, "message": "same repeated value " * 20, } for i in range(80) ] request = { "model": "gpt-4o", "messages": [ {"role": "user", "content": "summarize rows"}, { "role": "assistant", "content": None, "tool_calls": [ { "id": "call-1", "type": "function", "function": {"name": "list_rows", "arguments": "{}"}, } ], }, {"role": "tool", "tool_call_id": "call-1", "content": json.dumps(items)}, ], } with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as client: response = client.post("/v1/compress", json=request) body = response.json() assert response.status_code == 200, response.text assert body["tokens_before"] > body["tokens_after"], body assert body["tokens_saved"] > 0 assert body["compression_ratio"] < 1.0 assert body["transforms_applied"], body