diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 3e1164001..28c320896 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -875,6 +875,7 @@ class HeadroomProxy( # HTTP/1.1-only client for ChatGPT passthrough (Cloudflare challenges # our HTTP/2 fingerprint on its sensitive account endpoints). self.http_client_h1: httpx.AsyncClient | None = None + self._shutdown_event: asyncio.Event | None = None # Shared cold-start warmup registry (populated by startup()). # Holds typed slots with loaded / loading / null / error status for @@ -1386,6 +1387,7 @@ class HeadroomProxy( async def startup(self): """Initialize async resources.""" + self._get_shutdown_event().clear() self.pipeline_extensions.emit( PipelineStage.PRE_START, operation="proxy.startup", @@ -1627,6 +1629,7 @@ class HeadroomProxy( async def shutdown(self): """Cleanup async resources.""" + self._get_shutdown_event().set() if self.http_client_h1 and self.http_client_h1 is not self.http_client: await self.http_client_h1.aclose() self.http_client_h1 = None @@ -1727,6 +1730,33 @@ class HeadroomProxy( return extract_tags(headers) + def _get_shutdown_event(self) -> asyncio.Event: + event = getattr(self, "_shutdown_event", None) + if event is None: + event = asyncio.Event() + self._shutdown_event = event + return event + + async def _wait_for_retry_delay_or_shutdown(self, delay_seconds: float) -> bool: + try: + await asyncio.wait_for(self._get_shutdown_event().wait(), timeout=delay_seconds) + return True + except asyncio.TimeoutError: + return False + + def _shutdown_retry_response(self, method: str, url: str) -> httpx.Response: + return httpx.Response( + 503, + request=httpx.Request(method, url), + headers={"content-type": "application/json", "retry-after": "0"}, + json={ + "error": { + "type": "shutdown", + "message": "Proxy is shutting down; retry backoff cancelled.", + } + }, + ) + async def _retry_request( self, method: str, @@ -1822,7 +1852,13 @@ class HeadroomProxy( f"Upstream {response.status_code} (attempt {attempt + 1}), " f"retrying in {delay_ms:.0f}ms" ) - await asyncio.sleep(delay_ms / 1000) + if await self._wait_for_retry_delay_or_shutdown(delay_ms / 1000): + logger.info( + "Shutdown interrupted retry backoff for %s %s", + method, + path_for_log or "", + ) + return self._shutdown_retry_response(method, url) continue # Don't retry other client errors (4xx) @@ -1855,7 +1891,13 @@ class HeadroomProxy( logger.warning( f"Request failed (attempt {attempt + 1}), retrying in {delay_with_jitter:.0f}ms: {e}" ) - await asyncio.sleep(delay_with_jitter / 1000) + if await self._wait_for_retry_delay_or_shutdown(delay_with_jitter / 1000): + logger.info( + "Shutdown interrupted retry backoff for %s %s", + method, + path_for_log or "", + ) + return self._shutdown_retry_response(method, url) if last_error is None: raise RuntimeError( diff --git a/tests/test_proxy_handler_helpers.py b/tests/test_proxy_handler_helpers.py index fe2b48f6b..9b410e8b9 100644 --- a/tests/test_proxy_handler_helpers.py +++ b/tests/test_proxy_handler_helpers.py @@ -555,6 +555,51 @@ def test_retry_request_retries_connect_timeout() -> None: assert proxy.http_client.attempts == 2 +def test_retry_request_returns_503_when_shutdown_interrupts_retry_sleep() -> None: + class _Always429Client: + def __init__(self) -> None: + self.attempts = 0 + + async def post(self, url, **kwargs): # type: ignore[no-untyped-def] + self.attempts += 1 + return httpx.Response( + 429, + request=httpx.Request("POST", url), + json={"error": {"message": "slow down"}}, + headers={"retry-after": "30"}, + ) + + proxy = object.__new__(HeadroomProxy) + proxy.http_client = _Always429Client() + proxy.config = SimpleNamespace( + retry_enabled=True, + retry_max_attempts=3, + retry_base_delay_ms=30000, + retry_max_delay_ms=30000, + ) + proxy._shutdown_event = asyncio.Event() + proxy._shutdown_event.set() + + response = asyncio.run( + proxy._retry_request( + "POST", + "https://api.anthropic.test/v1/messages", + {}, + {"model": "claude-3-5-sonnet"}, + ) + ) + + assert response.status_code == 503 + assert response.json() == { + "error": { + "type": "shutdown", + "message": "Proxy is shutting down; retry backoff cancelled.", + } + } + assert response.headers["retry-after"] == "0" + assert proxy.http_client.attempts == 1 + + def test_anthropic_tool_sort_and_context_append_helpers() -> None: tools = [ {"type": "function", "function": {"name": "beta"}}, diff --git a/tests/test_proxy_pipeline_lifecycle.py b/tests/test_proxy_pipeline_lifecycle.py index 2a7a700e5..913cbb6dc 100644 --- a/tests/test_proxy_pipeline_lifecycle.py +++ b/tests/test_proxy_pipeline_lifecycle.py @@ -128,6 +128,34 @@ def test_proxy_shutdown_flushes_savings_tracker() -> None: proxy.metrics.savings_tracker.flush.assert_called_once() +def test_proxy_shutdown_signals_retry_waiters() -> None: + config = ProxyConfig( + optimize=False, + image_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, + ) + app = create_app(config) + proxy = app.state.proxy + proxy.http_client = None + proxy.memory_handler = None + proxy._shutdown_event = asyncio.Event() + + quota_registry = SimpleNamespace(stop_all=AsyncMock()) + with ( + patch("headroom.proxy.server.get_quota_registry", return_value=quota_registry), + patch("headroom.models.ml_models.MLModelRegistry.unload_prefix"), + ): + asyncio.run(proxy.shutdown()) + + assert proxy._shutdown_event.is_set() + + def test_openai_chat_pipeline_events_cover_proxy_lifecycle(monkeypatch) -> None: recorder = _RecordingExtension() config = ProxyConfig( diff --git a/tests/test_proxy_retry_429.py b/tests/test_proxy_retry_429.py index 67f145c6e..173bc6e55 100644 --- a/tests/test_proxy_retry_429.py +++ b/tests/test_proxy_retry_429.py @@ -105,10 +105,13 @@ def test_retry_request_returns_429_verbatim_on_exhaustion() -> None: def test_retry_request_honors_retry_after(monkeypatch) -> None: slept: list[float] = [] - async def _fake_sleep(seconds: float) -> None: + async def _fake_wait(self, seconds: float) -> bool: # type: ignore[no-untyped-def] slept.append(seconds) + return False - monkeypatch.setattr("headroom.proxy.server.asyncio.sleep", _fake_sleep) + monkeypatch.setattr( + "headroom.proxy.server.HeadroomProxy._wait_for_retry_delay_or_shutdown", _fake_wait + ) transport = _RateLimitTransport(fail_status=429, fail_times=1, retry_after="2") proxy = _proxy_with(transport) asyncio.run(proxy._retry_request("POST", "https://up/v1/messages", {}, {"messages": []})) @@ -186,10 +189,13 @@ def test_retry_request_returns_529_verbatim_on_exhaustion() -> None: def test_retry_request_honors_retry_after_on_529(monkeypatch) -> None: slept: list[float] = [] - async def _fake_sleep(seconds: float) -> None: + async def _fake_wait(self, seconds: float) -> bool: # type: ignore[no-untyped-def] slept.append(seconds) + return False - monkeypatch.setattr("headroom.proxy.server.asyncio.sleep", _fake_sleep) + monkeypatch.setattr( + "headroom.proxy.server.HeadroomProxy._wait_for_retry_delay_or_shutdown", _fake_wait + ) transport = _RateLimitTransport(fail_status=529, fail_times=1, retry_after="2") proxy = _proxy_with(transport) asyncio.run(proxy._retry_request("POST", "https://up/v1/messages", {}, {"messages": []}))