mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
ceae879e79
## Description Codex `/v1/responses` WebSocket disconnects can trigger a known `websockets` callback failure before `connection_made()` initializes `recv_messages`. When that happens, the proxy process can stay alive while `/livez` keeps advertising a clean healthy state. This change contains that known callback failure in the proxy runtime, records loop callback health, and makes `/livez` report the degraded state instead of always returning a clean process-alive payload. Closes #1720 ## 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 - Add a proxy-owned asyncio loop exception handler that recognizes the known `websockets` `connection_lost` `ClientConnection.recv_messages` `AttributeError`, records it in bounded runtime health state, and leaves unrelated loop exceptions delegated to the previous or default handler. - Extend `/livez` so the route remains cheap and unauthenticated while reflecting recorded event-loop callback health instead of always reporting a clean process-alive payload. - Preserve existing Codex WebSocket relay, fallback, session deregistration, and termination-cause behavior for normal handler-owned failures. - Add focused regression coverage for the known callback failure, the negative-space delegation path, and the health route response after loop callback degradation. ## Testing - [x] Unit tests pass (`uv run pytest tests/test_proxy_healthchecks.py tests/test_openai_codex_ws_lifecycle.py tests/test_proxy_loop_exception_health.py -q`) - [x] Linting passes (`uv run ruff check headroom/proxy/server.py tests/test_proxy_healthchecks.py tests/test_openai_codex_ws_lifecycle.py tests/test_proxy_loop_exception_health.py` and `uv run ruff format --check headroom/proxy/server.py tests/test_proxy_healthchecks.py tests/test_openai_codex_ws_lifecycle.py tests/test_proxy_loop_exception_health.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_healthchecks.py tests/test_openai_codex_ws_lifecycle.py tests/test_proxy_loop_exception_health.py -q ============================= test session starts ============================= platform win32 -- Python 3.12.13, pytest-9.0.3, pluggy-1.6.0 rootdir: D:\Repos\headroom-pr-1720-responses-ws-livez-wedge configfile: pyproject.toml plugins: anyio-4.12.1, langsmith-0.9.3, asyncio-1.3.0, cov-7.0.0 asyncio: mode=Mode.AUTO, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function collected 33 items tests\test_proxy_healthchecks.py ............ [ 36%] tests\test_openai_codex_ws_lifecycle.py ................... [ 93%] tests\test_proxy_loop_exception_health.py .. [100%] ============================== warnings summary =============================== .venv\Lib\site-packages\fastapi\testclient.py:1 D:\Repos\headroom-pr-1720-responses-ws-livez-wedge\.venv\Lib\site-packages\fastapi\testclient.py:1: StarletteDeprecationWarning: Using `httpx` with `starlette.testclient` is deprecated; install `httpx2` instead. from starlette.testclient import TestClient as TestClient # noqa -- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html ======================== 33 passed, 1 warning in 9.78s ======================== uv run ruff check headroom/proxy/server.py tests/test_proxy_healthchecks.py tests/test_openai_codex_ws_lifecycle.py tests/test_proxy_loop_exception_health.py All checks passed! uv run ruff format --check headroom/proxy/server.py tests/test_proxy_healthchecks.py tests/test_openai_codex_ws_lifecycle.py tests/test_proxy_loop_exception_health.py 4 files already formatted ``` ## Real Behavior Proof - Environment: Python proxy runtime with FastAPI TestClient, no external OpenAI credentials required. - Exact command / steps: invoke the installed loop exception handler with an asyncio context matching `Connection.connection_lost` plus `AttributeError("'ClientConnection' object has no attribute 'recv_messages'")`, then request `/livez`. - Observed result: the known `websockets` callback failure is recorded without delegating to the noisy default handler, `/livez` reports degraded loop callback health (HTTP 503, `"status": "unhealthy"`, `"alive": false`), and unrelated callback exceptions still reach the delegated handler. - Not tested: the nondeterministic upstream CPython or `websockets` timing edge against a live network connection; the focused regression pins the callback shape reported in #1720. ## 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 - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] 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 `CHANGELOG.md` is release-managed from conventional commits, so this PR does not edit it manually. The scope stays inside the proxy runtime and Codex WebSocket dispatch path; it does not change compression, CCR, provider-neutral pipeline behavior, or generic transform modules.
96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
"""Unit: event-loop callback handling for Codex WS disconnect regressions."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
from fastapi.testclient import TestClient
|
|
|
|
from headroom.proxy.server import ProxyConfig, create_app
|
|
|
|
pytest.importorskip("fastapi")
|
|
pytest.importorskip("httpx")
|
|
|
|
|
|
def _known_loop_callback_context() -> dict[str, object]:
|
|
return {
|
|
"message": "Exception in callback Connection.connection_lost(ConnectionResetError())",
|
|
"exception": AttributeError("'ClientConnection' object has no attribute 'recv_messages'"),
|
|
}
|
|
|
|
|
|
def _make_client(app):
|
|
return TestClient(app, base_url="http://127.0.0.1", client=("127.0.0.1", 12345))
|
|
|
|
|
|
def test_livez_reports_known_websockets_callback_degradation():
|
|
config = ProxyConfig(
|
|
optimize=False,
|
|
cache_enabled=False,
|
|
rate_limit_enabled=False,
|
|
cost_tracking_enabled=False,
|
|
)
|
|
app = create_app(config)
|
|
|
|
with _make_client(app) as client:
|
|
before = client.get("/livez")
|
|
assert before.status_code == 200
|
|
assert before.json()["status"] == "healthy"
|
|
assert before.json()["alive"] is True
|
|
|
|
assert app.state.loop_exception_handler is not None
|
|
mock_loop = MagicMock(spec=asyncio.AbstractEventLoop)
|
|
app.state.loop_exception_handler(mock_loop, _known_loop_callback_context())
|
|
|
|
after = client.get("/livez")
|
|
assert after.status_code == 503
|
|
payload = after.json()
|
|
assert payload["status"] == "unhealthy"
|
|
assert payload["alive"] is False
|
|
loop_health = payload["loop_health"]
|
|
assert loop_health["status"] == "unhealthy"
|
|
assert loop_health["known_failures"] == 1
|
|
assert (
|
|
loop_health["last_known_failure"]["exception"]
|
|
== "'ClientConnection' object has no attribute 'recv_messages'"
|
|
)
|
|
|
|
|
|
def test_unrelated_loop_callback_is_delegated_to_previous_handler():
|
|
delegate_calls: list[dict[str, object]] = []
|
|
config = ProxyConfig(
|
|
optimize=False,
|
|
cache_enabled=False,
|
|
rate_limit_enabled=False,
|
|
cost_tracking_enabled=False,
|
|
)
|
|
app = create_app(config)
|
|
|
|
with _make_client(app) as client:
|
|
client.get("/livez")
|
|
assert app.state.loop_exception_handler is not None
|
|
|
|
def _previous(_loop: object, context: dict[str, object]) -> None:
|
|
delegate_calls.append(dict(context))
|
|
|
|
app.state.previous_loop_exception_handler = _previous
|
|
|
|
mock_loop = MagicMock(spec=asyncio.AbstractEventLoop)
|
|
app.state.loop_exception_handler(
|
|
mock_loop,
|
|
{
|
|
"message": "random callback failed",
|
|
"exception": RuntimeError("not known failure"),
|
|
},
|
|
)
|
|
|
|
assert len(delegate_calls) == 1
|
|
assert delegate_calls[0]["message"] == "random callback failed"
|
|
assert app.state.loop_callback_health["status"] == "healthy"
|
|
assert app.state.loop_callback_health["known_failures"] == 0
|
|
|
|
health = client.get("/livez").json()
|
|
assert health["status"] == "healthy"
|
|
assert health["alive"] is True
|