mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
3be2526b76
## Description Buffered Anthropic `/v1/messages` requests still use Headroom's generic 300-second read timeout, which can produce proxy-generated `502 ReadTimeout` errors on long turns. This adds a dedicated buffered Anthropic timeout, keeps it applied across CCR and memory continuations plus batch paths, and makes the direct server entrypoint enforce the same positive-integer contract as the Click CLI. Closes #1261. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made - Added `anthropic_buffered_request_timeout_seconds` for buffered Anthropic reads. - Routed `/v1/messages`, CCR continuation, memory continuation, batch create, batch passthrough, and batch results through that timeout. - Enforced the same positive-integer validation for `HEADROOM_ANTHROPIC_BUFFERED_REQUEST_TIMEOUT_SECONDS` and `--anthropic-buffered-request-timeout-seconds` in both startup paths. - Added focused regressions and updated `CHANGELOG.md`. ## Testing - [x] `uv run pytest tests/test_proxy/test_anthropic_buffered_timeout.py tests/test_cli_proxy_improvements.py::TestNewEnvVarWiring` - [x] `uv run ruff check .` - [x] `uv run ruff format . --check` ### Test Output ```text $ uv run pytest tests/test_proxy/test_anthropic_buffered_timeout.py tests/test_cli_proxy_improvements.py::TestNewEnvVarWiring 17 passed in 3.42s $ uv run ruff check . All checks passed! $ uv run ruff format . --check 966 files already formatted ``` ## Real Behavior Proof - Environment: local FastAPI `TestClient` with stubbed retry and HTTP client seams - Exact command / steps: run `uv run pytest tests/test_proxy/test_anthropic_buffered_timeout.py tests/test_cli_proxy_improvements.py::TestNewEnvVarWiring`; the tests build `ProxyConfig(request_timeout_seconds=7, connect_timeout_seconds=3, anthropic_buffered_request_timeout_seconds=19)`, drive `/v1/messages`, `/v1/messages/batches`, `/v1/messages/batches/{batch_id}/results`, a CCR continuation, and a memory continuation through `TestClient`, then verify `HEADROOM_ANTHROPIC_BUFFERED_REQUEST_TIMEOUT_SECONDS=0` falls back to `600`, `--anthropic-buffered-request-timeout-seconds 0` is rejected, and default proxy timeouts stay `read=300` and `write=300` - Observed result: buffered Anthropic paths use `httpx.Timeout(connect=3, read=19, write=7, pool=3)`, continuation requests stay on that same budget, invalid zero-valued startup config is rejected or ignored back to the default, and unrelated proxy timeout defaults stay unchanged - Not tested: live upstream Anthropic latency beyond the focused stubbed-timeout regression ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have added tests that prove the fix - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable
412 lines
13 KiB
Python
412 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
fastapi = pytest.importorskip("fastapi")
|
|
httpx = pytest.importorskip("httpx")
|
|
|
|
from fastapi.testclient import TestClient # noqa: E402
|
|
|
|
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
|
|
|
|
|
|
class _FakePrefixTracker:
|
|
def get_frozen_message_count(self) -> int:
|
|
return 0
|
|
|
|
def get_last_original_messages(self) -> list[dict]:
|
|
return []
|
|
|
|
def get_last_forwarded_messages(self) -> list[dict]:
|
|
return []
|
|
|
|
def update_from_response(self, **kwargs): # noqa: ANN003
|
|
return None
|
|
|
|
|
|
class _BufferedPassthroughClient:
|
|
def __init__(self, response: httpx.Response) -> None:
|
|
self.response = response
|
|
self.calls: list[dict[str, object]] = []
|
|
|
|
async def request(self, method, url, headers=None, content=None, timeout=None): # noqa: ANN001
|
|
self.calls.append(
|
|
{
|
|
"method": method,
|
|
"url": url,
|
|
"headers": headers,
|
|
"content": content,
|
|
"timeout": timeout,
|
|
}
|
|
)
|
|
_assert_buffered_timeout(timeout)
|
|
return self.response
|
|
|
|
async def get(self, url, headers=None, timeout=None): # noqa: ANN001
|
|
return await self.request("GET", url, headers=headers, timeout=timeout)
|
|
|
|
async def post(self, url, headers=None, content=None, timeout=None): # noqa: ANN001
|
|
return await self.request("POST", url, headers=headers, content=content, timeout=timeout)
|
|
|
|
async def aclose(self) -> None:
|
|
return None
|
|
|
|
|
|
def _assert_buffered_timeout(timeout: httpx.Timeout | None) -> None:
|
|
assert isinstance(timeout, httpx.Timeout)
|
|
assert timeout.connect == 3.0
|
|
assert timeout.read == 19.0
|
|
assert timeout.write == 7.0
|
|
assert timeout.pool == 3.0
|
|
|
|
|
|
def _make_config() -> ProxyConfig:
|
|
return ProxyConfig(
|
|
optimize=False,
|
|
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,
|
|
connect_timeout_seconds=3,
|
|
request_timeout_seconds=7,
|
|
anthropic_buffered_request_timeout_seconds=19,
|
|
)
|
|
|
|
|
|
def _install_prefix_tracker(proxy) -> None:
|
|
tracker = _FakePrefixTracker()
|
|
proxy.session_tracker_store.compute_session_id = lambda request, model, messages: "s1"
|
|
proxy.session_tracker_store.get_or_create = lambda session_id, provider: tracker
|
|
|
|
|
|
def _anthropic_message_response() -> dict[str, object]:
|
|
return {
|
|
"id": "msg_test_1",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "ok"}],
|
|
"usage": {
|
|
"input_tokens": 12,
|
|
"output_tokens": 3,
|
|
"cache_read_input_tokens": 0,
|
|
"cache_creation_input_tokens": 0,
|
|
},
|
|
}
|
|
|
|
|
|
def _anthropic_batch_response() -> dict[str, object]:
|
|
return {"id": "batch_test_1", "object": "batch", "status": "in_progress"}
|
|
|
|
|
|
def _anthropic_list_response() -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
json={"object": "list", "data": [], "first_id": None, "last_id": None},
|
|
headers={"content-type": "application/json"},
|
|
)
|
|
|
|
|
|
def test_anthropic_messages_buffered_timeout_override_reaches_retry_request():
|
|
config = _make_config()
|
|
app = create_app(config)
|
|
with TestClient(app) as client:
|
|
proxy = client.app.state.proxy
|
|
_install_prefix_tracker(proxy)
|
|
captured: dict[str, object] = {}
|
|
|
|
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
|
|
timeout = kwargs.get("timeout")
|
|
captured["timeout"] = timeout
|
|
_assert_buffered_timeout(timeout)
|
|
return httpx.Response(200, json=_anthropic_message_response())
|
|
|
|
proxy._retry_request = _fake_retry # type: ignore[assignment]
|
|
|
|
response = client.post(
|
|
"/v1/messages",
|
|
headers={
|
|
"x-api-key": "test-key",
|
|
"anthropic-version": "2023-06-01",
|
|
"content-type": "application/json",
|
|
},
|
|
json={
|
|
"model": "claude-sonnet-4-6",
|
|
"max_tokens": 64,
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
},
|
|
)
|
|
|
|
assert "timeout" in captured, response.text
|
|
assert isinstance(captured["timeout"], httpx.Timeout)
|
|
|
|
|
|
def test_anthropic_batch_create_buffered_timeout_override_reaches_retry_request():
|
|
config = _make_config()
|
|
app = create_app(config)
|
|
with TestClient(app) as client:
|
|
proxy = client.app.state.proxy
|
|
captured: dict[str, object] = {}
|
|
|
|
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
|
|
timeout = kwargs.get("timeout")
|
|
captured["timeout"] = timeout
|
|
_assert_buffered_timeout(timeout)
|
|
return httpx.Response(200, json=_anthropic_batch_response())
|
|
|
|
proxy._retry_request = _fake_retry # type: ignore[assignment]
|
|
|
|
response = client.post(
|
|
"/v1/messages/batches",
|
|
headers={
|
|
"x-api-key": "test-key",
|
|
"anthropic-version": "2023-06-01",
|
|
"content-type": "application/json",
|
|
},
|
|
json={
|
|
"requests": [
|
|
{
|
|
"custom_id": "req-1",
|
|
"params": {
|
|
"model": "claude-sonnet-4-6",
|
|
"max_tokens": 64,
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
},
|
|
}
|
|
]
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200, response.text
|
|
assert isinstance(captured["timeout"], httpx.Timeout)
|
|
|
|
|
|
def test_anthropic_batch_passthrough_buffered_timeout_override_reaches_http_client():
|
|
config = _make_config()
|
|
app = create_app(config)
|
|
with TestClient(app) as client:
|
|
proxy = client.app.state.proxy
|
|
http_client = _BufferedPassthroughClient(_anthropic_list_response())
|
|
proxy.http_client = http_client
|
|
|
|
response = client.get(
|
|
"/v1/messages/batches",
|
|
headers={
|
|
"x-api-key": "test-key",
|
|
"anthropic-version": "2023-06-01",
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200, response.text
|
|
assert len(http_client.calls) == 1
|
|
assert isinstance(http_client.calls[0]["timeout"], httpx.Timeout)
|
|
|
|
|
|
def test_anthropic_batch_results_buffered_timeout_override_reaches_http_client_get():
|
|
config = _make_config()
|
|
app = create_app(config)
|
|
with TestClient(app) as client:
|
|
proxy = client.app.state.proxy
|
|
http_client = _BufferedPassthroughClient(
|
|
httpx.Response(
|
|
200,
|
|
content=b'{"custom_id":"req-1","result":{"type":"succeeded"}}\n',
|
|
headers={"content-type": "application/jsonl"},
|
|
)
|
|
)
|
|
proxy.http_client = http_client
|
|
|
|
response = client.get(
|
|
"/v1/messages/batches/batch_test_1/results",
|
|
headers={
|
|
"x-api-key": "test-key",
|
|
"anthropic-version": "2023-06-01",
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200, response.text
|
|
assert len(http_client.calls) == 1
|
|
assert http_client.calls[0]["method"] == "GET"
|
|
assert isinstance(http_client.calls[0]["timeout"], httpx.Timeout)
|
|
|
|
|
|
def test_anthropic_ccr_continuation_uses_buffered_timeout() -> None:
|
|
config = _make_config()
|
|
config.ccr_inject_tool = True
|
|
config.ccr_handle_responses = True
|
|
app = create_app(config)
|
|
|
|
class _CCRHandler:
|
|
def has_ccr_tool_calls(self, response, provider): # noqa: ANN001
|
|
return True
|
|
|
|
async def handle_response( # noqa: ANN001
|
|
self,
|
|
response,
|
|
optimized_messages,
|
|
tools,
|
|
api_call_fn,
|
|
provider,
|
|
):
|
|
return await api_call_fn(
|
|
optimized_messages
|
|
+ [{"role": "assistant", "content": response.get("content", [])}],
|
|
tools,
|
|
)
|
|
|
|
with TestClient(app) as client:
|
|
proxy = client.app.state.proxy
|
|
_install_prefix_tracker(proxy)
|
|
proxy.ccr_response_handler = _CCRHandler()
|
|
http_client = _BufferedPassthroughClient(
|
|
httpx.Response(200, json=_anthropic_message_response())
|
|
)
|
|
proxy.http_client = http_client
|
|
|
|
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
|
|
return httpx.Response(200, json=_anthropic_message_response())
|
|
|
|
proxy._retry_request = _fake_retry # type: ignore[assignment]
|
|
|
|
client.post(
|
|
"/v1/messages",
|
|
headers={
|
|
"x-api-key": "test-key",
|
|
"anthropic-version": "2023-06-01",
|
|
"content-type": "application/json",
|
|
},
|
|
json={
|
|
"model": "claude-sonnet-4-6",
|
|
"max_tokens": 64,
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
},
|
|
)
|
|
|
|
assert len(http_client.calls) == 1
|
|
assert http_client.calls[0]["method"] == "POST"
|
|
assert isinstance(http_client.calls[0]["timeout"], httpx.Timeout)
|
|
|
|
|
|
def test_anthropic_memory_continuation_uses_buffered_timeout() -> None:
|
|
config = _make_config()
|
|
config.memory_enabled = True
|
|
app = create_app(config)
|
|
|
|
class _MemoryHandler:
|
|
def __init__(self) -> None:
|
|
self.config = type(
|
|
"MemoryConfig",
|
|
(),
|
|
{
|
|
"inject_context": False,
|
|
"inject_tools": False,
|
|
"project_root_override": "",
|
|
},
|
|
)()
|
|
self.initialized = False
|
|
self.backend = None
|
|
|
|
def get_beta_headers(self) -> dict[str, str]:
|
|
return {}
|
|
|
|
def has_memory_tool_calls(self, response, provider): # noqa: ANN001
|
|
return True
|
|
|
|
async def handle_memory_tool_calls( # noqa: ANN001
|
|
self,
|
|
response,
|
|
user_id,
|
|
provider,
|
|
**kwargs,
|
|
):
|
|
return [{"type": "tool_result", "tool_use_id": "mem_1", "content": "memory"}]
|
|
|
|
with TestClient(app) as client:
|
|
proxy = client.app.state.proxy
|
|
_install_prefix_tracker(proxy)
|
|
proxy.memory_handler = _MemoryHandler()
|
|
captured_timeouts: list[httpx.Timeout | None] = []
|
|
|
|
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
|
|
timeout = kwargs.get("timeout")
|
|
captured_timeouts.append(timeout)
|
|
if len(body["messages"]) > 1:
|
|
_assert_buffered_timeout(timeout)
|
|
return httpx.Response(200, json=_anthropic_message_response())
|
|
|
|
proxy._retry_request = _fake_retry # type: ignore[assignment]
|
|
|
|
client.post(
|
|
"/v1/messages",
|
|
headers={
|
|
"x-api-key": "test-key",
|
|
"anthropic-version": "2023-06-01",
|
|
"content-type": "application/json",
|
|
"x-headroom-user-id": "user-1",
|
|
},
|
|
json={
|
|
"model": "claude-sonnet-4-6",
|
|
"max_tokens": 64,
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
},
|
|
)
|
|
|
|
assert len(captured_timeouts) == 2
|
|
assert all(isinstance(timeout, httpx.Timeout) for timeout in captured_timeouts)
|
|
|
|
|
|
def test_retry_request_without_override_uses_client_default_timeout():
|
|
config = _make_config()
|
|
app = create_app(config)
|
|
with TestClient(app) as client:
|
|
proxy = client.app.state.proxy
|
|
captured: list[dict[str, object]] = []
|
|
|
|
async def _fake_retry(method, url, headers, body, stream=False, **kwargs): # noqa: ANN001
|
|
captured.append(kwargs)
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"id": "chatcmpl-1",
|
|
"object": "chat.completion",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": "ok"},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
|
|
},
|
|
)
|
|
|
|
proxy._retry_request = _fake_retry # type: ignore[assignment]
|
|
|
|
response = client.post(
|
|
"/v1/chat/completions",
|
|
headers={"authorization": "Bearer test-key", "content-type": "application/json"},
|
|
json={
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200, response.text
|
|
assert len(captured) == 1
|
|
assert "timeout" not in captured[0], (
|
|
"_retry_request without an override must not pass timeout=None to httpx"
|
|
)
|
|
|
|
|
|
def test_generic_proxy_timeout_defaults_stay_unchanged():
|
|
app = create_app(ProxyConfig())
|
|
with TestClient(app) as client:
|
|
timeout = client.app.state.proxy.http_client.timeout
|
|
|
|
assert timeout.connect == 10.0
|
|
assert timeout.read == 300.0
|
|
assert timeout.write == 300.0
|
|
assert timeout.pool == 10.0
|