mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
## Description `POST /v1/compress` could hang on large payloads and freeze the whole proxy. `handle_compress()` called `self.openai_pipeline.apply()` **synchronously** inside the async handler, so a large body's CPU/Rust-bound compression blocked the single event loop for seconds — concurrent requests, even `GET /health` and `/livez`, stalled until it finished, and a pathologically large body could hang indefinitely. The fix runs the compression through the **existing bounded compression executor** (already used by the sibling OpenAI handlers in the same class), so the loop stays free and an over-long compression fails fast with a timeout instead of hanging. Closes #718 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `proxy/handlers/openai.py` (`handle_compress`): wrap `self.openai_pipeline.apply(...)` in `await self._run_compression_in_executor(lambda: ..., timeout= COMPRESSION_TIMEOUT_SECONDS)` — mirroring the existing request handlers. The bounded executor keeps the CPU/Rust work off the event loop, and the timeout makes a too-large body fail fast. - Added an explicit `except TimeoutError` arm that returns `503` with `type: "compression_timeout"` and a clear message ("payload too large"); other errors still return the existing `503 compression_error`. The bypass-header short-circuit is unchanged. - Tests: new `TestCompressEndpointDoesNotBlockLoop` — while a blocking compression is in flight, a concurrent `GET /livez` returns 200 and the compression is verifiably still running (it would already be done if `apply` had hijacked the loop). The existing happy-path compress tests now exercise the executor path. ## Testing - [x] Unit tests pass (`pytest tests/test_proxy_compress_endpoint.py` — 10 passed) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy headroom`; handler module is in the existing `proxy.handlers.*` mypy override) - [x] New tests added for new functionality - [x] Manual testing performed (live Windows large-payload smoke — see proof) ### Test Output ```text $ pytest tests/test_proxy_compress_endpoint.py -q 10 passed in 26.45s $ ruff check headroom/proxy/handlers/openai.py tests/test_proxy_compress_endpoint.py All checks passed! ``` Negative control: with the fix reverted (apply() inline) the new test fails at `assert not compress.done()` — the inline call hijacks the loop so the request finishes before `/livez` is served. With the fix it passes. ## Real Behavior Proof - Environment: Windows 11, Python 3.13.11, `headroom` 0.28.0, live `headroom proxy --port 8798 --no-telemetry`. - Exact command / steps: POST a ~2.6 MB body (≈519k tokens) to `/v1/compress` while a background thread probes `/livez` continuously. - Observed result: during a 2.36 s compression of a ~519k-token payload, `/livez` was served 155 times (mean ~5 ms) — the loop stayed responsive instead of freezing. Full output: ```text payload bytes: 2587297 compress: {'secs': 2.36, 'status': 200, 'before': 519007, 'after': 413} livez probes during compress: 155 max=178.4ms mean=4.8ms ``` During a 2.36 s compression of a half-million-token payload, `/livez` was served **155 times** with a mean latency of ~5 ms — the event loop stayed responsive instead of freezing for the whole compression. (A single 178 ms blip corresponds to a brief GIL-held pure-Python section; the bulk of the work is GIL-releasing Rust compression, which is why offloading helps.) A cold first request before warmup showed the old behavior — a single `/livez` blocked ~2.2 s for the compression duration. - Not tested: behavior on a non-Windows host (the loop-blocking is platform-independent; the regression test runs on CI/Linux). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Additional Notes - The executor and `COMPRESSION_TIMEOUT_SECONDS` already existed and are used by the other handlers; this PR only routes the compress endpoint through the same path. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
324 lines
12 KiB
Python
324 lines
12 KiB
Python
"""Tests for the /v1/compress endpoint in the proxy server.
|
|
|
|
These tests verify that the compression-only endpoint works correctly
|
|
for the TypeScript SDK and other HTTP clients.
|
|
"""
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
# Skip if fastapi not available
|
|
pytest.importorskip("fastapi")
|
|
|
|
from fastapi.testclient import TestClient
|
|
|
|
from headroom.proxy.server import ProxyConfig, create_app
|
|
|
|
|
|
@pytest.fixture
|
|
def client():
|
|
"""Create test client with optimization enabled."""
|
|
config = ProxyConfig(
|
|
optimize=True,
|
|
cache_enabled=False,
|
|
rate_limit_enabled=False,
|
|
cost_tracking_enabled=False,
|
|
)
|
|
app = create_app(config)
|
|
# /v1/compress is loopback-gated (#1227).
|
|
with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as c:
|
|
yield c
|
|
|
|
|
|
@pytest.fixture
|
|
def client_no_optimize():
|
|
"""Create test client with optimization disabled."""
|
|
config = ProxyConfig(
|
|
optimize=False,
|
|
cache_enabled=False,
|
|
rate_limit_enabled=False,
|
|
cost_tracking_enabled=False,
|
|
)
|
|
app = create_app(config)
|
|
# /v1/compress is loopback-gated (#1227).
|
|
with TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345)) as c:
|
|
yield c
|
|
|
|
|
|
class TestCompressEndpointValidation:
|
|
"""Test request validation for /v1/compress."""
|
|
|
|
def test_missing_messages_returns_400(self, client):
|
|
"""Request without messages field should return 400."""
|
|
response = client.post("/v1/compress", json={"model": "gpt-4"})
|
|
assert response.status_code == 400
|
|
data = response.json()
|
|
assert "error" in data
|
|
assert data["error"]["type"] == "invalid_request"
|
|
assert "messages" in data["error"]["message"]
|
|
|
|
def test_missing_model_returns_400(self, client):
|
|
"""Request without model field should return 400."""
|
|
response = client.post(
|
|
"/v1/compress",
|
|
json={"messages": [{"role": "user", "content": "hello"}]},
|
|
)
|
|
assert response.status_code == 400
|
|
data = response.json()
|
|
assert "error" in data
|
|
assert data["error"]["type"] == "invalid_request"
|
|
assert "model" in data["error"]["message"]
|
|
|
|
def test_invalid_json_returns_400(self, client):
|
|
"""Request with invalid JSON should return 400."""
|
|
response = client.post(
|
|
"/v1/compress",
|
|
content=b"not valid json",
|
|
headers={"content-type": "application/json"},
|
|
)
|
|
assert response.status_code == 400
|
|
data = response.json()
|
|
assert data["error"]["type"] == "invalid_request"
|
|
|
|
|
|
class TestCompressEndpointBasic:
|
|
"""Test basic compress endpoint behavior."""
|
|
|
|
def test_empty_messages_returns_empty(self, client):
|
|
"""Empty messages list should return as-is with zero metrics."""
|
|
response = client.post(
|
|
"/v1/compress",
|
|
json={"messages": [], "model": "gpt-4"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["messages"] == []
|
|
assert data["tokens_before"] == 0
|
|
assert data["tokens_after"] == 0
|
|
assert data["tokens_saved"] == 0
|
|
assert data["compression_ratio"] == 1.0
|
|
assert data["transforms_applied"] == []
|
|
assert data["ccr_hashes"] == []
|
|
|
|
def test_basic_compression_response_shape(self, client):
|
|
"""Verify the response contains all expected fields."""
|
|
response = client.post(
|
|
"/v1/compress",
|
|
json={
|
|
"messages": [{"role": "user", "content": "Hello, world!"}],
|
|
"model": "gpt-4",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
# Check all expected fields are present
|
|
assert "messages" in data
|
|
assert "tokens_before" in data
|
|
assert "tokens_after" in data
|
|
assert "tokens_saved" in data
|
|
assert "compression_ratio" in data
|
|
assert "transforms_applied" in data
|
|
assert "ccr_hashes" in data
|
|
|
|
# Messages should be a list
|
|
assert isinstance(data["messages"], list)
|
|
assert len(data["messages"]) >= 1
|
|
|
|
# Numeric fields should be non-negative
|
|
assert data["tokens_before"] >= 0
|
|
assert data["tokens_after"] >= 0
|
|
assert data["tokens_saved"] >= 0
|
|
assert data["compression_ratio"] > 0
|
|
|
|
def test_bypass_header_returns_uncompressed(self, client):
|
|
"""X-Headroom-Bypass header should skip compression."""
|
|
messages = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
{"role": "user", "content": "How are you?"},
|
|
]
|
|
response = client.post(
|
|
"/v1/compress",
|
|
json={"messages": messages, "model": "gpt-4"},
|
|
headers={"x-headroom-bypass": "true"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["messages"] == messages
|
|
assert data["tokens_before"] == 0
|
|
assert data["tokens_after"] == 0
|
|
assert data["tokens_saved"] == 0
|
|
assert data["compression_ratio"] == 1.0
|
|
assert data["transforms_applied"] == []
|
|
assert data["ccr_hashes"] == []
|
|
|
|
def test_bypass_header_case_insensitive(self, client):
|
|
"""Bypass header should be case-insensitive."""
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
response = client.post(
|
|
"/v1/compress",
|
|
json={"messages": messages, "model": "gpt-4"},
|
|
headers={"x-headroom-bypass": "TRUE"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["messages"] == messages
|
|
|
|
|
|
class TestCompressEndpointCompression:
|
|
"""Test that actual compression happens for large content."""
|
|
|
|
def test_large_tool_output_gets_compressed(self, client):
|
|
"""Large tool output content should result in tokens_saved > 0."""
|
|
# Create a large repetitive tool output that should be compressible
|
|
large_data = json.dumps(
|
|
[
|
|
{
|
|
"id": i,
|
|
"name": f"Item {i}",
|
|
"description": f"This is a detailed description for item number {i}. "
|
|
f"It contains various attributes and metadata that are typical "
|
|
f"of API responses. The item has a status of active and was "
|
|
f"created on 2024-01-{(i % 28) + 1:02d}. Additional fields "
|
|
f"include category=electronics, price={i * 10.99:.2f}, "
|
|
f"rating={4.0 + (i % 10) / 10:.1f}, stock={i * 5}.",
|
|
"tags": ["electronics", "sale", "featured", "new-arrival"],
|
|
"metadata": {
|
|
"created_by": "system",
|
|
"updated_at": "2024-01-15T00:00:00Z",
|
|
"version": i,
|
|
"source": "api",
|
|
},
|
|
}
|
|
for i in range(200)
|
|
]
|
|
)
|
|
|
|
messages = [
|
|
{"role": "user", "content": "What items are available?"},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_123",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "list_items",
|
|
"arguments": "{}",
|
|
},
|
|
}
|
|
],
|
|
},
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": "call_123",
|
|
"content": large_data,
|
|
},
|
|
{"role": "user", "content": "Summarize the first 5 items."},
|
|
]
|
|
|
|
response = client.post(
|
|
"/v1/compress",
|
|
json={"messages": messages, "model": "gpt-4"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
|
|
# With a large tool output, the pipeline should process successfully
|
|
assert data["tokens_before"] > 0
|
|
assert data["tokens_after"] > 0
|
|
assert data["tokens_after"] <= data["tokens_before"]
|
|
assert data["tokens_saved"] == data["tokens_before"] - data["tokens_after"]
|
|
assert 0 < data["compression_ratio"] <= 1.0
|
|
assert isinstance(data["transforms_applied"], list)
|
|
|
|
def test_small_content_may_not_compress(self, client):
|
|
"""Small messages may not get compressed but should still work."""
|
|
response = client.post(
|
|
"/v1/compress",
|
|
json={
|
|
"messages": [{"role": "user", "content": "Hi"}],
|
|
"model": "gpt-4",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
# Should still return valid response regardless of compression
|
|
assert data["tokens_before"] >= 0
|
|
assert data["tokens_after"] >= 0
|
|
assert isinstance(data["transforms_applied"], list)
|
|
|
|
|
|
class TestCompressEndpointDoesNotBlockLoop:
|
|
"""/v1/compress must offload to the compression executor so a slow/large
|
|
payload cannot freeze the single event loop (#718)."""
|
|
|
|
async def test_compress_does_not_block_liveness(self, monkeypatch):
|
|
import asyncio
|
|
import threading
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
|
|
config = ProxyConfig(
|
|
optimize=True,
|
|
cache_enabled=False,
|
|
rate_limit_enabled=False,
|
|
cost_tracking_enabled=False,
|
|
)
|
|
app = create_app(config)
|
|
proxy = app.state.proxy
|
|
|
|
entered = threading.Event()
|
|
release = threading.Event()
|
|
|
|
def blocking_apply(**kwargs):
|
|
# Stand in for a large CPU-bound compression: blocks its worker until
|
|
# released. If this ran inline on the loop (the bug), the loop would
|
|
# be frozen and /livez below could not be served.
|
|
entered.set()
|
|
release.wait(timeout=10)
|
|
return SimpleNamespace(
|
|
messages=kwargs["messages"],
|
|
tokens_before=10,
|
|
tokens_after=5,
|
|
transforms_applied=[],
|
|
transforms_summary={},
|
|
markers_inserted=[],
|
|
)
|
|
|
|
monkeypatch.setattr(proxy.openai_pipeline, "apply", blocking_apply)
|
|
|
|
# /v1/compress is loopback-gated (#1227) — present as 127.0.0.1.
|
|
transport = httpx.ASGITransport(app=app, client=("127.0.0.1", 12345))
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://127.0.0.1") as client:
|
|
compress = asyncio.create_task(
|
|
client.post(
|
|
"/v1/compress",
|
|
json={
|
|
"messages": [{"role": "user", "content": "hello world"}],
|
|
"model": "gpt-4",
|
|
},
|
|
)
|
|
)
|
|
# Wait until the compression is actually in flight (running in the
|
|
# executor thread), then prove the loop is still responsive.
|
|
for _ in range(200):
|
|
if entered.is_set():
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
assert entered.is_set(), "compression never started"
|
|
|
|
livez = await asyncio.wait_for(client.get("/livez"), timeout=5)
|
|
assert livez.status_code == 200
|
|
assert livez.json()["alive"] is True
|
|
# The compression is still blocked — /livez was served concurrently.
|
|
assert not compress.done()
|
|
|
|
release.set()
|
|
resp = await asyncio.wait_for(compress, timeout=5)
|
|
assert resp.status_code == 200
|
|
assert resp.json()["tokens_saved"] == 5
|