From da2d8dc9dbf3edfcd1c3f6429db32374a6bebc64 Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Mon, 6 Jul 2026 09:24:47 -0400 Subject: [PATCH] fix(proxy): cancel retry backoff on shutdown (#1834) ## Description During proxy shutdown, an in-flight retrying request can currently stay asleep inside `_retry_request()` and keep the client socket hanging until the retry timer expires or an external supervisor kills the process. This wires retry backoff to a proxy-scoped shutdown event so shutdown interrupts those waits immediately and returns a clear `503` response instead of leaving the request stalled. Closes #1821. ## 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 - Added a proxy-scoped shutdown event in `headroom/proxy/server.py`. - Cleared that event at startup and set it at shutdown before teardown proceeds. - Replaced both retry-backoff sleeps with a helper that wakes on either timeout or shutdown. - Returned a shutdown `503` with `retry-after: 0` when shutdown interrupts retry backoff. - Stopped the shutdown interruption logs from falling back to the raw upstream URL when no safe path string is available. - Added focused regressions for retry-backoff interruption and shutdown event signaling. - Updated the existing Retry-After tests to observe the new shutdown-aware wait helper instead of the old raw sleep hook. ## Testing - [x] Unit tests pass (`uv run pytest tests/test_proxy_handler_helpers.py tests/test_proxy_pipeline_lifecycle.py -q`) - [x] Unit tests pass (`uv run pytest tests/test_proxy_retry_429.py -q`) - [x] Linting passes (`uv run ruff check headroom/proxy/server.py tests/test_proxy_handler_helpers.py tests/test_proxy_retry_429.py tests/test_proxy_pipeline_lifecycle.py`) - [ ] Type checking passes (`uv run mypy headroom`) - [x] New tests added for new functionality when applicable - [ ] Manual testing performed ### Test Output ```text uv run pytest tests/test_proxy_handler_helpers.py tests/test_proxy_pipeline_lifecycle.py -q 32 passed, 1 warning in 13.05s uv run pytest tests/test_proxy_retry_429.py -q 10 passed, 1 warning in 1.12s uv run ruff check headroom/proxy/server.py tests/test_proxy_handler_helpers.py tests/test_proxy_retry_429.py tests/test_proxy_pipeline_lifecycle.py All checks passed! ``` ## Real Behavior Proof - Environment: Windows, project `uv` environment, focused proxy retry and shutdown regressions. - Exact command / steps: copy the updated shutdown regression files into a detached `origin/main` worktree and run `tests/test_proxy_handler_helpers.py` plus `tests/test_proxy_pipeline_lifecycle.py`, then rerun those files on this branch and separately rerun `tests/test_proxy_retry_429.py` after updating the existing Retry-After tests to patch the shutdown-aware wait helper. - Observed result: base fails because retry backoff still returns the original `429` and `shutdown()` leaves the retry event unset; head passes the focused file, preserves the existing Retry-After assertions, and returns a shutdown `503` with `retry-after: 0` while signaling retry waiters during shutdown. - Not tested: live systemd-managed shutdown on Linux or a full VS Code / Claude Code session. ## 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 performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes This is intentionally scoped to retry backoff during shutdown. It does not try to cancel unrelated in-flight request work or change the broader retry policy outside shutdown. --- headroom/proxy/server.py | 46 ++++++++++++++++++++++++-- tests/test_proxy_handler_helpers.py | 45 +++++++++++++++++++++++++ tests/test_proxy_pipeline_lifecycle.py | 28 ++++++++++++++++ tests/test_proxy_retry_429.py | 14 +++++--- 4 files changed, 127 insertions(+), 6 deletions(-) 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": []}))