fix: add anthropic pre-upstream timeouts

This commit is contained in:
Adryan Eka Vandra 2026-04-20 22:38:26 +07:00
parent 5391761fe6
commit 4e9348dec4
No known key found for this signature in database
GPG key ID: A46A577A26A97682
11 changed files with 312 additions and 24 deletions

View file

@ -73,6 +73,30 @@ from .main import main
"Env: HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY."
),
)
@click.option(
"--anthropic-pre-upstream-acquire-timeout-seconds",
type=float,
default=None,
envvar="HEADROOM_ANTHROPIC_PRE_UPSTREAM_ACQUIRE_TIMEOUT_SECONDS",
help=(
"Fail-fast timeout for waiting on the Anthropic pre-upstream semaphore "
"before returning 503 + Retry-After. "
"Default: 15.0 seconds. "
"Env: HEADROOM_ANTHROPIC_PRE_UPSTREAM_ACQUIRE_TIMEOUT_SECONDS."
),
)
@click.option(
"--anthropic-pre-upstream-memory-context-timeout-seconds",
type=float,
default=None,
envvar="HEADROOM_ANTHROPIC_PRE_UPSTREAM_MEMORY_CONTEXT_TIMEOUT_SECONDS",
help=(
"Fail-open timeout for Anthropic memory-context lookup while the request "
"still holds a pre-upstream slot. "
"Default: 2.0 seconds. "
"Env: HEADROOM_ANTHROPIC_PRE_UPSTREAM_MEMORY_CONTEXT_TIMEOUT_SECONDS."
),
)
@click.option("--log-file", default=None, help="Path to JSONL log file")
@click.option(
"--budget",
@ -217,6 +241,8 @@ def proxy(
retry_max_attempts: int | None,
connect_timeout_seconds: int | None,
anthropic_pre_upstream_concurrency: int | None,
anthropic_pre_upstream_acquire_timeout_seconds: float | None,
anthropic_pre_upstream_memory_context_timeout_seconds: float | None,
log_file: str | None,
budget: float | None,
code_graph: bool,
@ -352,6 +378,16 @@ def proxy(
# Precedence: CLI > env > auto-compute (click's ``envvar``
# handles the env-var fallback).
anthropic_pre_upstream_concurrency=anthropic_pre_upstream_concurrency,
anthropic_pre_upstream_acquire_timeout_seconds=(
anthropic_pre_upstream_acquire_timeout_seconds
if anthropic_pre_upstream_acquire_timeout_seconds is not None
else 15.0
),
anthropic_pre_upstream_memory_context_timeout_seconds=(
anthropic_pre_upstream_memory_context_timeout_seconds
if anthropic_pre_upstream_memory_context_timeout_seconds is not None
else 2.0
),
)
memory_status = "DISABLED"

View file

@ -343,23 +343,6 @@ class AnthropicHandlerMixin:
_pre_upstream_sem_acquired = False
pre_upstream_sem.release()
if pre_upstream_sem is not None:
_wait_started_at = time.perf_counter()
await pre_upstream_sem.acquire()
_pre_upstream_sem_acquired = True
_wait_ms = (time.perf_counter() - _wait_started_at) * 1000.0
stage_timer.record("pre_upstream_wait", _wait_ms)
if _wait_ms > 100.0:
logger.info(
"[%s] pre_upstream_wait_ms=%.2f session_id=%s "
"(anthropic pre-upstream semaphore contention)",
request_id,
_wait_ms,
trace_session_id,
)
else:
stage_timer.record("pre_upstream_wait", 0.0)
async def _finalize_pre_upstream() -> None:
"""Release the pre-upstream semaphore and emit stage-timing metrics.
@ -400,6 +383,54 @@ class AnthropicHandlerMixin:
metrics=getattr(self, "metrics", None),
)
if pre_upstream_sem is not None:
_wait_started_at = time.perf_counter()
_acquire_timeout_seconds = self.config.anthropic_pre_upstream_acquire_timeout_seconds
try:
await asyncio.wait_for(
pre_upstream_sem.acquire(),
timeout=_acquire_timeout_seconds,
)
except asyncio.TimeoutError:
_wait_ms = (time.perf_counter() - _wait_started_at) * 1000.0
stage_timer.record("pre_upstream_wait", _wait_ms)
logger.warning(
"[%s] Anthropic pre-upstream queue saturated after %.2f ms "
"(timeout=%.1fs, session_id=%s)",
request_id,
_wait_ms,
_acquire_timeout_seconds,
trace_session_id,
)
await _finalize_pre_upstream()
return JSONResponse(
status_code=503,
headers={"Retry-After": str(max(1, int(_acquire_timeout_seconds) + 1))},
content={
"type": "error",
"error": {
"type": "service_unavailable",
"message": (
"Anthropic pre-upstream queue is saturated. "
"Please retry shortly."
),
},
},
)
_pre_upstream_sem_acquired = True
_wait_ms = (time.perf_counter() - _wait_started_at) * 1000.0
stage_timer.record("pre_upstream_wait", _wait_ms)
if _wait_ms > 100.0:
logger.info(
"[%s] pre_upstream_wait_ms=%.2f session_id=%s "
"(anthropic pre-upstream semaphore contention)",
request_id,
_wait_ms,
trace_session_id,
)
else:
stage_timer.record("pre_upstream_wait", 0.0)
try:
# Check request body size
content_length = request.headers.get("content-length")
@ -980,9 +1011,22 @@ class AnthropicHandlerMixin:
if self.memory_handler.config.inject_context:
try:
async with stage_timer.measure("memory_context"):
memory_context = await self.memory_handler.search_and_format_context(
memory_user_id, optimized_messages
memory_context = await asyncio.wait_for(
self.memory_handler.search_and_format_context(
memory_user_id, optimized_messages
),
timeout=(
self.config.anthropic_pre_upstream_memory_context_timeout_seconds
),
)
except asyncio.TimeoutError:
memory_context = None
logger.info(
f"[{request_id}] Memory: Context lookup exceeded "
f"{self.config.anthropic_pre_upstream_memory_context_timeout_seconds:.1f}s; "
"continuing without it"
)
try:
if memory_context:
if is_cache_mode(self.config.mode):
logger.info(

View file

@ -25,6 +25,7 @@ Usage:
from __future__ import annotations
import asyncio
import inspect
import json
import logging
from dataclasses import dataclass, field
@ -130,6 +131,22 @@ class MemoryHandler:
self._init_lock = asyncio.Lock()
return self._init_lock
async def _close_backend_instance(self, backend: Any, *, reason: str) -> None:
"""Best-effort close for a partially initialized backend."""
close = getattr(backend, "close", None)
if not callable(close):
return
try:
result = close()
if inspect.isawaitable(result):
await result
except Exception as exc:
logger.warning(
"Memory: failed to close backend during %s cleanup: %s",
reason,
exc,
)
def _init_native_memory_dir(self) -> None:
"""Initialize native memory directory."""
if self.config.native_memory_dir:
@ -190,7 +207,11 @@ class MemoryHandler:
# was cancelled by wait_for. Callers that do
# ``if self.memory_handler._backend:`` must not see a
# truthy-but-broken backend.
existing_backend = self._backend
if existing_backend is not None:
await self._close_backend_instance(existing_backend, reason="timeout")
self._backend = None
self._initialized = False
logger.error(
"Memory: backend initialization timed out after "
f"{STARTUP_INIT_TIMEOUT_SECONDS}s "
@ -205,6 +226,9 @@ class MemoryHandler:
# blocks don't either, so it propagates unconditionally.
# Reset state so any later retry starts clean, then re-raise:
# cancellation is a signal, not an error to swallow.
existing_backend = self._backend
if existing_backend is not None:
await self._close_backend_instance(existing_backend, reason="cancellation")
self._backend = None
self._initialized = False
logger.info(f"Memory: backend initialization cancelled (backend={self.config.backend})")
@ -1714,8 +1738,8 @@ To SAVE: create /memories/<topic>.txt "content"
async def close(self) -> None:
"""Close the memory backend."""
if self._backend and hasattr(self._backend, "close"):
await self._backend.close()
if self._backend is not None:
await self._close_backend_instance(self._backend, reason="handler close")
self._backend = None
self._initialized = False
logger.info("Memory: Handler closed")

View file

@ -238,3 +238,11 @@ class ProxyConfig:
# Env: ``HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY``.
# Precedence: CLI > env > auto-compute.
anthropic_pre_upstream_concurrency: int | None = None
# Upper bound for waiting on the Anthropic pre-upstream semaphore
# before failing fast with a 503 + Retry-After. Keeps the queue bounded
# when all pre-upstream slots are occupied by slow/hung work.
anthropic_pre_upstream_acquire_timeout_seconds: float = 15.0
# Fail-open timeout for Anthropic memory-context lookup while the request
# is still holding a pre-upstream slot. Compression already has its own
# COMPRESSION_TIMEOUT_SECONDS guard; this bounds the memory leg too.
anthropic_pre_upstream_memory_context_timeout_seconds: float = 2.0

View file

@ -413,6 +413,12 @@ class HeadroomProxy(
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
@ -685,6 +691,13 @@ class HeadroomProxy(
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,
)
# Smart routing status
if self.config.smart_routing:
@ -1264,6 +1277,32 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
),
}
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
)
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,
},
"websocket_sessions": {
"active_sessions": ws_active_sessions,
"active_relay_tasks": ws_active_relay_tasks,
},
}
def _health_payload(*, include_config: bool) -> dict[str, Any]:
checks = _health_checks()
ready = all(check["ready"] for check in checks.values())
@ -1275,6 +1314,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
"timestamp": _iso_utc_now(),
"uptime_seconds": _uptime_seconds(),
"checks": checks,
"runtime": _runtime_payload(),
}
deployment_profile = os.environ.get("HEADROOM_DEPLOYMENT_PROFILE")
if deployment_profile:
@ -1371,6 +1411,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
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.get("/dashboard", response_class=HTMLResponse)

View file

@ -9,6 +9,8 @@ Covers:
- N+1 contention (only the (N+1)th waiter records ``pre_upstream_wait`` > 0)
- strict serialization under concurrency=1
- unbounded mode (``anthropic_pre_upstream_concurrency=0`` -> no semaphore)
- acquire timeout fails fast with ``503`` + ``Retry-After``
- memory-context timeout fails open without leaking the semaphore
- exception-safety (semaphore released when the critical section raises)
- ``/livez`` unaffected under Anthropic backpressure
- compression is not bypassed (the Unit 4 gate is additive, not a shortcut)
@ -156,9 +158,11 @@ class _DummyAnthropicHandler(AnthropicHandlerMixin):
self.session_tracker_store = SimpleNamespace(
compute_session_id=lambda *a, **k: "sess-1",
get_or_create=lambda *a, **k: SimpleNamespace(
_cached_token_count=0,
get_frozen_message_count=lambda: 0,
get_last_original_messages=lambda: [],
get_last_forwarded_messages=lambda: [],
update_from_response=lambda *a, **k: None,
record_request=lambda *a, **k: None,
),
)
@ -449,6 +453,84 @@ def test_exception_inside_critical_section_releases_semaphore():
anyio.run(_run)
def test_acquire_timeout_returns_503_with_retry_after(stage_log_capture):
async def _run() -> None:
sem = asyncio.Semaphore(1)
await sem.acquire()
handler = _DummyAnthropicHandler(anthropic_pre_upstream_sem=sem)
handler.config.anthropic_pre_upstream_acquire_timeout_seconds = 0.01
req = _build_request(
{
"model": "claude-3-5-sonnet-latest",
"messages": [{"role": "user", "content": "hello"}],
},
{"authorization": "Bearer sk-ant-api-test"},
)
try:
response = await handler.handle_anthropic_messages(req)
assert response.status_code == 503
assert response.headers["retry-after"] == "1"
body = json.loads(response.body)
assert body["error"]["type"] == "service_unavailable"
assert sem._value == 0
finally:
sem.release()
assert sem._value == 1
with _tokenizer_patch():
anyio.run(_run)
payloads = _parse_all_stage_logs(stage_log_capture)
assert len(payloads) == 1
assert payloads[0]["stages"]["pre_upstream_wait"] >= 10.0
def test_memory_context_timeout_fails_open_and_releases_semaphore():
class _MemoryHandler:
def __init__(self) -> None:
self.config = SimpleNamespace(inject_context=True, inject_tools=False)
self.initialized = False
self.backend = None
async def search_and_format_context(self, _user_id, _messages):
await asyncio.sleep(5.0)
return "should-timeout"
def inject_tools(self, tools, _provider):
return tools, False
def get_beta_headers(self) -> dict[str, str]:
return {}
def has_memory_tool_calls(self, _response, _provider) -> bool:
return False
async def handle_memory_tool_calls(self, _response, _user_id, _provider):
return []
async def _run() -> None:
sem = asyncio.Semaphore(1)
handler = _DummyAnthropicHandler(anthropic_pre_upstream_sem=sem)
handler.memory_handler = _MemoryHandler()
handler.config.anthropic_pre_upstream_memory_context_timeout_seconds = 0.01
req = _build_request(
{
"model": "claude-3-5-sonnet-latest",
"messages": [{"role": "user", "content": "hello"}],
},
{
"authorization": "Bearer sk-ant-api-test",
"x-headroom-user-id": "user-1",
},
)
response = await handler.handle_anthropic_messages(req)
assert response.status_code == 200
assert sem._value == 1
with _tokenizer_patch():
anyio.run(_run)
# --------------------------------------------------------------------------- #
# /livez stays fast under Anthropic pre-upstream contention. #
# --------------------------------------------------------------------------- #
@ -602,6 +684,34 @@ def test_cli_flag_overrides_env_var():
assert config.anthropic_pre_upstream_concurrency == 7
def test_cli_env_sets_pre_upstream_timeouts():
env = {
"HEADROOM_ANTHROPIC_PRE_UPSTREAM_ACQUIRE_TIMEOUT_SECONDS": "9.5",
"HEADROOM_ANTHROPIC_PRE_UPSTREAM_MEMORY_CONTEXT_TIMEOUT_SECONDS": "3.25",
}
config = _run_cli_capture([], env=env)
assert config.anthropic_pre_upstream_acquire_timeout_seconds == pytest.approx(9.5)
assert config.anthropic_pre_upstream_memory_context_timeout_seconds == pytest.approx(3.25)
def test_cli_flags_override_pre_upstream_timeout_env_vars():
env = {
"HEADROOM_ANTHROPIC_PRE_UPSTREAM_ACQUIRE_TIMEOUT_SECONDS": "9.5",
"HEADROOM_ANTHROPIC_PRE_UPSTREAM_MEMORY_CONTEXT_TIMEOUT_SECONDS": "3.25",
}
config = _run_cli_capture(
[
"--anthropic-pre-upstream-acquire-timeout-seconds",
"4.5",
"--anthropic-pre-upstream-memory-context-timeout-seconds",
"1.5",
],
env=env,
)
assert config.anthropic_pre_upstream_acquire_timeout_seconds == pytest.approx(4.5)
assert config.anthropic_pre_upstream_memory_context_timeout_seconds == pytest.approx(1.5)
# --------------------------------------------------------------------------- #
# Sanity: HeadroomProxy auto-computes default when config value is None. #
# --------------------------------------------------------------------------- #

View file

@ -150,6 +150,8 @@ async def test_ensure_initialized_timeout_nulls_partially_initialized_backend(
``if self.memory_handler._backend:`` see a truthy-but-broken backend.
"""
close_hits = {"n": 0}
class SlowBackend:
def __init__(self, config):
self.config = config
@ -159,7 +161,7 @@ async def test_ensure_initialized_timeout_nulls_partially_initialized_backend(
await asyncio.sleep(5.0)
async def close(self) -> None:
pass
close_hits["n"] += 1
import headroom.memory.backends.local as local_mod
@ -175,6 +177,7 @@ async def test_ensure_initialized_timeout_nulls_partially_initialized_backend(
# Both must be consistent after timeout.
assert handler._initialized is False
assert handler._backend is None
assert close_hits["n"] == 1
@pytest.mark.asyncio
@ -183,6 +186,8 @@ async def test_ensure_initialized_cancellation_propagates_and_resets_state(tmp_p
propagate (CancelledError is BaseException not a swallowable error)
and leave the handler in a clean state."""
close_hits = {"n": 0}
class HangingBackend:
def __init__(self, config):
self.config = config
@ -191,7 +196,7 @@ async def test_ensure_initialized_cancellation_propagates_and_resets_state(tmp_p
await asyncio.Event().wait()
async def close(self) -> None:
pass
close_hits["n"] += 1
import headroom.memory.backends.local as local_mod
@ -210,6 +215,7 @@ async def test_ensure_initialized_cancellation_propagates_and_resets_state(tmp_p
assert handler._initialized is False
assert handler._backend is None
assert close_hits["n"] == 1
# -------------------------------------------------------------------

View file

@ -293,8 +293,11 @@ def test_debug_warmup_reports_registry_slots(client):
assert "magika" in data
assert "memory_backend" in data
assert "memory_embedder" in data
assert "runtime" in data
# Each slot has at least a status field.
assert "status" in data["memory_backend"]
assert data["runtime"]["anthropic_pre_upstream"]["resolved_concurrency"] >= 0
assert data["runtime"]["websocket_sessions"]["active_relay_tasks"] == 0
def test_debug_ws_sessions_reports_live_session(app_and_client):

View file

@ -1,3 +1,4 @@
import os
from types import SimpleNamespace
import pytest
@ -47,6 +48,17 @@ def test_readyz_reports_core_subsystem_checks(client):
assert data["checks"]["cache"]["status"] == "disabled"
assert data["checks"]["rate_limiter"]["status"] == "disabled"
assert data["checks"]["memory"]["status"] == "disabled"
runtime = data["runtime"]
assert runtime["anthropic_pre_upstream"]["resolved_concurrency"] == max(
2, min(8, os.cpu_count() or 4)
)
assert runtime["anthropic_pre_upstream"]["source"] == "auto"
assert runtime["anthropic_pre_upstream"]["acquire_timeout_seconds"] == 15.0
assert runtime["anthropic_pre_upstream"]["compression_timeout_seconds"] == 30.0
assert runtime["anthropic_pre_upstream"]["memory_context_timeout_seconds"] == 2.0
assert runtime["anthropic_pre_upstream"]["codex_ws_gated"] is False
assert runtime["websocket_sessions"]["active_sessions"] == 0
assert runtime["websocket_sessions"]["active_relay_tasks"] == 0
def test_health_preserves_backwards_compatible_config_payload(client):

View file

@ -247,6 +247,8 @@ headroom proxy --mode cache
| `--retry-max-attempts` | runtime default `3` | Maximum upstream retry attempts |
| `--connect-timeout-seconds` | runtime default `10` | Upstream connection timeout |
| `--anthropic-pre-upstream-concurrency` | auto `max(2, min(8, cpu_count))` | Cap simultaneous pre-upstream work on `/v1/messages` (body read, deep copy, first compression stage, memory-context lookup, upstream connect). `0` or negative disables (unbounded); any positive integer is honoured verbatim. Prevents cold-start replay storms from starving `/livez`, `/readyz`, and new Codex WS opens. |
| `--anthropic-pre-upstream-acquire-timeout-seconds` | `15.0` | Fail fast when the Anthropic pre-upstream queue is saturated. Requests that wait longer return `503` with `Retry-After` instead of parking indefinitely. |
| `--anthropic-pre-upstream-memory-context-timeout-seconds` | `2.0` | Fail-open timeout for Anthropic memory-context lookup while the request still holds a pre-upstream slot. |
| `--log-file` | unset | JSONL log output path |
| `--budget` | unset | Daily USD budget limit |
| `--no-code-aware` | off | Disable AST-aware code compression |
@ -274,7 +276,8 @@ headroom proxy --mode cache
Notes:
- `--learn` implies memory unless `--no-learn` is also set.
- Proxy startup can also read environment variables such as `HEADROOM_HOST`, `HEADROOM_PORT`, `HEADROOM_BUDGET`, `HEADROOM_MODE`, `HEADROOM_ANYLLM_PROVIDER`, `HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY`, `ANTHROPIC_TARGET_API_URL`, `OPENAI_TARGET_API_URL`, and `GEMINI_TARGET_API_URL`. CLI flags take precedence over environment variables.
- Proxy startup can also read environment variables such as `HEADROOM_HOST`, `HEADROOM_PORT`, `HEADROOM_BUDGET`, `HEADROOM_MODE`, `HEADROOM_ANYLLM_PROVIDER`, `HEADROOM_ANTHROPIC_PRE_UPSTREAM_CONCURRENCY`, `HEADROOM_ANTHROPIC_PRE_UPSTREAM_ACQUIRE_TIMEOUT_SECONDS`, `HEADROOM_ANTHROPIC_PRE_UPSTREAM_MEMORY_CONTEXT_TIMEOUT_SECONDS`, `ANTHROPIC_TARGET_API_URL`, `OPENAI_TARGET_API_URL`, and `GEMINI_TARGET_API_URL`. CLI flags take precedence over environment variables.
- The default Anthropic pre-upstream cap is intentionally conservative for CPU/ONNX-heavy work. Larger containers may want to raise it after checking the resolved runtime values on `/readyz` or `/debug/warmup`.
See also: [Proxy Server](proxy.md), [Configuration](configuration.md)

View file

@ -60,6 +60,7 @@ Without stage timings, active-task introspection, or session bookkeeping, the ne
- **OpenTelemetry / metrics exporter wiring**: the counters added here expose `prometheus_metrics` entries; OTLP export belongs in a follow-up.
- **External watchdog in the LaunchAgent**: a plist-level `WatchPaths`/`ThrottleInterval` revision belongs in the `.dotfiles` repo (see origin §"Important External Files"), not here. This plan only adds the in-process signal the watchdog would consume (§Unit 5).
- **Quantifying multi-agent reconnect budget**: tuning the `Unit 4` semaphore default via load testing is work for after merge.
- **Dedicated Codex WS-side cap**: `/v1/responses` remains intentionally ungated by the Anthropic HTTP semaphore. The current plan surfaces `active_relay_tasks` on `/readyz` so operators can see WS pressure early; add a separate WS-side cap only if those counters show the threat model has shifted.
## Context & Research