fix(proxy): cancel retry backoff on shutdown (#1834)
Some checks failed
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime name:]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code name:code]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim name:code-slim]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Waiting to run
CI / changes (push) Failing after 50s
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-nonroot name:nonroot]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim name:slim]) (push) Waiting to run
Docker / docker-build (map[name:amd64 platform:linux/amd64 runs_on:ubuntu-24.04], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime name:]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code name:code]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim name:code-slim]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-nonroot name:nonroot]) (push) Waiting to run
CI / commitlint (push) Failing after 2s
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim name:slim]) (push) Waiting to run
Docker / docker-build (map[name:arm64 platform:linux/arm64 runs_on:ubuntu-24.04-arm], map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Waiting to run
Docker / docker-manifest (map[bake_target:runtime name:]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-code name:code]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-code-nonroot name:code-nonroot]) (push) Blocked by required conditions
CI / lint (push) Failing after 4s
CI / build-wheel (push) Failing after 2s
Docker / docker-manifest (map[bake_target:runtime-code-slim name:code-slim]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-code-slim-nonroot name:code-slim-nonroot]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-nonroot name:nonroot]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-slim name:slim]) (push) Blocked by required conditions
Docker / docker-manifest (map[bake_target:runtime-slim-nonroot name:slim-nonroot]) (push) Blocked by required conditions
CI / prefetch-model (push) Failing after 4s
Docker / promote-latest (push) Blocked by required conditions
CI / test-agno (push) Failing after 4s
CI / build (push) Failing after 3s
CI / test-extras (push) Failing after 7s
CI / workflow-validation (push) Failing after 4s
CI / docker-native-e2e (push) Failing after 4s
Deploy Documentation / validate (push) Failing after 6s
Deploy Documentation / deploy (push) Failing after 6s
CI / test-dashboard-ui (push) Failing after 3s
CI / test (1) (push) Failing after 3s
CI / test (2) (push) Failing after 5s
CI / test (3) (push) Failing after 5s
CI / test (4) (push) Failing after 5s
Init E2E / docker-init-e2e (push) Failing after 6s
Install Native E2E / install-native (ubuntu-latest) (push) Failing after 6s
Merge Conflicts / merge-conflicts (push) Failing after 5s
Release Please / release-please (push) Failing after 4s
Security / CodeQL (javascript-typescript) (push) Failing after 31s
Security / CodeQL (python) (push) Failing after 28s
Security / Dependency audit (pip-audit) (push) Failing after 1m3s
Wrap E2E / docker-wrap-e2e (push) Failing after 21s
Wrap Native E2E / wrap-native (ubuntu-latest) (push) Failing after 6s
Security / Secret scan (gitleaks) (push) Failing after 36s
Init Native E2E / init-native (ubuntu-latest, copilot) (push) Failing after 12m32s
Init Native E2E / init-native (ubuntu-latest, codex) (push) Failing after 12m35s
Init Native E2E / init-native (ubuntu-latest, claude) (push) Failing after 12m36s
CI / windows-native-wrapper (push) Has been cancelled
CI / macos-native-wrapper (push) Has been cancelled
Init Native E2E / init-native (macos-latest, claude) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, codex) (push) Has been cancelled
Init Native E2E / init-native (macos-latest, copilot) (push) Has been cancelled
Install Native E2E / install-native (macos-latest) (push) Has been cancelled
Wrap Native E2E / wrap-native (macos-latest) (push) Has been cancelled

## 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.
This commit is contained in:
Rod Boev 2026-07-06 09:24:47 -04:00 committed by GitHub
parent afd9cbdfaf
commit da2d8dc9db
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 127 additions and 6 deletions

View file

@ -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 "<upstream-url>",
)
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 "<upstream-url>",
)
return self._shutdown_retry_response(method, url)
if last_error is None:
raise RuntimeError(

View file

@ -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"}},

View file

@ -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(

View file

@ -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": []}))