From 17cdb185bc79d8cfec104e781a7e555af3ef11e1 Mon Sep 17 00:00:00 2001
From: Patrick A <141967+neogenix@users.noreply.github.com>
Date: Wed, 5 Aug 2026 23:33:34 -0400
Subject: [PATCH] fix(proxy): graceful shutdown and reliable Ctrl+C exit (#621)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Problems

### 1. Noisy CancelledError traceback on Ctrl+C

Every Ctrl+C produced one or more "Exception in ASGI application" ERROR
log entries with a CancelledError traceback:

```
ERROR:    Exception in ASGI application
Traceback (most recent call last):
  ...
  File "uvicorn/protocols/http/h11_impl.py", line 410, in run_asgi
    result = await app(...)
  ...
asyncio.exceptions.CancelledError
```

### 2. Inconsistent / hung shutdown in multi-worker mode (`--workers 8`)

Workers blocked in a C-extension call (hnswlib, tree-sitter, ONNX
inference) could prevent Ctrl+C from completing because
`timeout_graceful_shutdown` defaulted to `None` (wait forever).

---

## Root causes

**Root cause A (CancelledError noise)**
uvicorn 0.40.0's `h11_impl.run_asgi()` (line 413) catches
`BaseException` — not just `Exception` — so `asyncio.CancelledError`
raised on every in-flight request at shutdown is unconditionally logged
as `ERROR: Exception in ASGI application`. This is expected behaviour
during shutdown, not a bug.

**Root cause B (hung multi-worker shutdown)**
`uvicorn.run()` was called without `timeout_graceful_shutdown`, which
defaults to `None`. This means the supervisor waits indefinitely for
in-flight requests to drain. A single request blocked in a C-extension
(e.g. hnswlib nearest-neighbour search, tree-sitter parse, ONNX
inference) prevents the whole process group from exiting.

**Root cause C (hung single-worker shutdown — lifespan unbounded
awaits)**
The lifespan `finally` block contained unbounded `await` calls to
`_beacon.stop()`, `proxy.usage_reporter.stop()`,
`proxy.traffic_learner.stop()`, and `proxy.shutdown()`. uvicorn's
`lifespan.shutdown()` calls `await self.shutdown_event.wait()` with no
timeout — that event is only set once the lifespan `finally` block
returns. Any of these awaits hanging (e.g. a reporter making a network
call) therefore requires a second Ctrl+C to force-exit.

---

## Changes

### `headroom/proxy/server.py`

1. **`_SuppressCancelledErrorFilter`** (new class, ~10 lines): a
`logging.Filter` that returns `False` for ERROR records on
`uvicorn.error` whose `exc_info[0]` is a subclass of
`asyncio.CancelledError`. Installed on
`logging.getLogger("uvicorn.error")` at the start of `run_server()`.

2. **`timeout_graceful_shutdown=10`** added to `uvicorn.run()`: forces
cancellation of any tasks still running 10 seconds after the shutdown
signal, ensuring workers blocked in C-extensions are reaped promptly.

3. **Bounded awaits in lifespan `finally` block**: a local `_timed(coro,
label, timeout)` helper wraps each shutdown step with
`asyncio.wait_for()`. Timeouts: beacon.stop 3s, usage_reporter.stop 3s,
traffic_learner.stop 3s, proxy.shutdown 5s. Each step logs a warning on
timeout/error and continues — the teardown path is now deterministic and
completes within ~15s on a single Ctrl+C.

4. **Shutdown log message** in the lifespan `finally` block:
`event=proxy_shutdown reason=signal pid=<n>` is logged as the first
action on teardown.

### `tests/test_graceful_shutdown.py` (new)

9 tests:
- 6 unit tests for `_SuppressCancelledErrorFilter` (suppresses
CancelledError at ERROR level, passes through WARNING-level
CancelledError, passes through other exceptions, handles
`exc_info=None`, handles `(None,None,None)` tuple, suppresses
subclasses)
- 1 integration test: `run_server()` installs the filter on
`uvicorn.error`
- 1 integration test: `run_server()` passes
`timeout_graceful_shutdown=10` to `uvicorn.run()`
- 1 integration test: lifespan emits `event=proxy_shutdown` on teardown

---

## Files changed

- `headroom/proxy/server.py` — filter class, bounded lifespan awaits,
graceful shutdown timeout
- `tests/test_graceful_shutdown.py` (new) — 9 tests
- `uv.lock` — dependency lockfile updated (routine sync, no dependency
changes)
- `CHANGELOG.md` — changelog entry

---

## How to verify

1. Start the proxy: `headroom proxy --port 8787 --workers 8 --memory
--code-aware ...`
2. Press Ctrl+C
3. Before: ERROR traceback for each in-flight request; second Ctrl+C
sometimes required
4. After: clean `event=proxy_shutdown reason=signal pid=...` log, then
process exits within ~15s regardless of stuck C-extensions or slow
reporters

---------

Co-authored-by: JD Davis <jd@jds-macbook-air.tail2a279.ts.net>
---
 headroom/proxy/server.py        |  49 +++++-
 tests/test_graceful_shutdown.py | 285 ++++++++++++++++++++++++++++++++
 2 files changed, 328 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_graceful_shutdown.py

diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py
index c03c76ee8..1fb97ffe8 100644
--- a/headroom/proxy/server.py
+++ b/headroom/proxy/server.py
@@ -470,6 +470,18 @@ logging.basicConfig(
 )
 logger = logging.getLogger("headroom.proxy")
 
+
+class _SuppressCancelledErrorFilter(logging.Filter):
+    """Hide expected uvicorn CancelledError tracebacks during shutdown."""
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        if record.levelno == logging.ERROR and record.exc_info:
+            exc_type = record.exc_info[0]
+            if exc_type is not None and issubclass(exc_type, asyncio.CancelledError):
+                return False
+        return True
+
+
 LoopExceptionHandler = Callable[[asyncio.AbstractEventLoop, dict[str, Any]], object]
 
 
@@ -2565,21 +2577,39 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
                 loop.set_exception_handler(previous)
 
             app.state.ready = False
-            # Shutdown
+            logger.info("event=proxy_shutdown reason=signal pid=%d", os.getpid())
+
+            async def _timed(coro: Any, *, label: str, timeout: float) -> None:
+                try:
+                    await asyncio.wait_for(coro, timeout=timeout)
+                except Exception as exc:
+                    logger.warning(
+                        "event=shutdown_step_timeout_or_error label=%s timeout=%.1fs exc=%r",
+                        label,
+                        timeout,
+                        exc,
+                    )
+
             if _cc_reconciler is not None:
-                await _cc_reconciler.stop()
+                await _timed(_cc_reconciler.stop(), label="cc_reconciler.stop", timeout=3.0)
             if _beacon_is_owner[0]:
                 _release_beacon_lock()
             if proxy.usage_reporter:
-                await proxy.usage_reporter.stop()
+                await _timed(proxy.usage_reporter.stop(), label="usage_reporter.stop", timeout=3.0)
             if proxy.traffic_learner:
-                await proxy.traffic_learner.stop()
+                await _timed(
+                    proxy.traffic_learner.stop(), label="traffic_learner.stop", timeout=3.0
+                )
             if proxy._background_compression_enabled:
-                await proxy._background_compressor.stop()
+                await _timed(
+                    proxy._background_compressor.stop(),
+                    label="background_compressor.stop",
+                    timeout=3.0,
+                )
             proxy._background_compression_executor.shutdown(wait=False)
             if proxy.code_graph_watcher:
                 proxy.code_graph_watcher.stop()
-            await proxy.shutdown()
+            await _timed(proxy.shutdown(), label="proxy.shutdown", timeout=5.0)
             shutdown_headroom_tracing()
             shutdown_otel_metrics()
 
@@ -5095,6 +5125,12 @@ def run_server(
 ╚══════════════════════════════════════════════════════════════════════╝
 """)
 
+    uvicorn_error_logger = logging.getLogger("uvicorn.error")
+    if not any(
+        isinstance(item, _SuppressCancelledErrorFilter) for item in uvicorn_error_logger.filters
+    ):
+        uvicorn_error_logger.addFilter(_SuppressCancelledErrorFilter())
+
     app_target: Any
     uvicorn_kwargs: dict[str, Any] = {}
     if sys.platform == "win32":
@@ -5150,6 +5186,7 @@ def run_server(
         # default. Disabling proxy_headers here guarantees the guard sees the
         # real peer address regardless of env.
         proxy_headers=False,
+        timeout_graceful_shutdown=10,
         **uvicorn_kwargs,
     )
 
diff --git a/tests/test_graceful_shutdown.py b/tests/test_graceful_shutdown.py
new file mode 100644
index 000000000..7bff1e5a1
--- /dev/null
+++ b/tests/test_graceful_shutdown.py
@@ -0,0 +1,285 @@
+"""Tests for graceful shutdown and Ctrl+C signal handling.
+
+Covers:
+- _SuppressCancelledErrorFilter suppresses "Exception in ASGI application"
+  log records whose exc_info is CancelledError
+- _SuppressCancelledErrorFilter passes through unrelated error records
+- timeout_graceful_shutdown is present in the uvicorn.run() call path
+- The lifespan shutdown branch logs the proxy_shutdown event
+- Lifespan shutdown completes even when individual steps block/raise
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+
+import pytest
+
+from headroom.proxy.server import (
+    ProxyConfig,
+    _SuppressCancelledErrorFilter,
+    create_app,
+)
+
+# ---------------------------------------------------------------------------
+# Unit tests for the logging filter
+# ---------------------------------------------------------------------------
+
+
+class TestSuppressCancelledErrorFilter:
+    """_SuppressCancelledErrorFilter silences CancelledError noise from uvicorn."""
+
+    def _make_record(
+        self,
+        level: int = logging.ERROR,
+        exc_type: type | None = None,
+    ) -> logging.LogRecord:
+        record = logging.LogRecord(
+            name="uvicorn.error",
+            level=level,
+            pathname="",
+            lineno=0,
+            msg="Exception in ASGI application",
+            args=(),
+            exc_info=(exc_type, exc_type() if exc_type else None, None) if exc_type else None,
+        )
+        return record
+
+    def test_suppresses_cancelled_error_at_error_level(self) -> None:
+        f = _SuppressCancelledErrorFilter()
+        record = self._make_record(logging.ERROR, asyncio.CancelledError)
+        assert f.filter(record) is False
+
+    def test_passes_through_cancelled_error_at_warning_level(self) -> None:
+        # Only suppress ERROR, not lower-severity records
+        f = _SuppressCancelledErrorFilter()
+        record = self._make_record(logging.WARNING, asyncio.CancelledError)
+        assert f.filter(record) is True
+
+    def test_passes_through_other_exception_at_error_level(self) -> None:
+        f = _SuppressCancelledErrorFilter()
+        record = self._make_record(logging.ERROR, ValueError)
+        assert f.filter(record) is True
+
+    def test_passes_through_record_without_exc_info(self) -> None:
+        f = _SuppressCancelledErrorFilter()
+        record = self._make_record(logging.ERROR, None)
+        # exc_info is set to None tuple when exc_type is None
+        record.exc_info = None
+        assert f.filter(record) is True
+
+    def test_passes_through_record_with_none_exc_type(self) -> None:
+        f = _SuppressCancelledErrorFilter()
+        record = self._make_record(logging.ERROR, None)
+        record.exc_info = (None, None, None)
+        assert f.filter(record) is True
+
+    def test_suppresses_subclass_of_cancelled_error(self) -> None:
+        """BaseException subclasses of CancelledError are also suppressed."""
+
+        class MyCancelled(asyncio.CancelledError):
+            pass
+
+        f = _SuppressCancelledErrorFilter()
+        record = self._make_record(logging.ERROR, MyCancelled)
+        assert f.filter(record) is False
+
+
+# ---------------------------------------------------------------------------
+# Integration: filter is installed on uvicorn.error in run_server()
+# ---------------------------------------------------------------------------
+
+
+def test_run_server_installs_cancelled_error_filter(monkeypatch: pytest.MonkeyPatch) -> None:
+    """run_server() attaches _SuppressCancelledErrorFilter to uvicorn.error logger."""
+    installed_filters: list = []
+
+    original_add_filter = logging.Logger.addFilter
+
+    def capturing_add_filter(self: logging.Logger, f: logging.Filter) -> None:
+        if self.name == "uvicorn.error" and isinstance(f, _SuppressCancelledErrorFilter):
+            installed_filters.append(f)
+        original_add_filter(self, f)
+
+    monkeypatch.setattr(logging.Logger, "addFilter", capturing_add_filter)
+
+    # Intercept uvicorn.run so we don't actually start a server
+    monkeypatch.setattr("uvicorn.run", lambda *a, **kw: None)
+
+    from headroom.proxy.server import run_server
+
+    run_server(ProxyConfig(), print_banner=False)
+
+    assert len(installed_filters) == 1, "Expected exactly one _SuppressCancelledErrorFilter"
+
+
+# ---------------------------------------------------------------------------
+# Integration: timeout_graceful_shutdown is forwarded to uvicorn.run()
+# ---------------------------------------------------------------------------
+
+
+def test_run_server_passes_timeout_graceful_shutdown(monkeypatch: pytest.MonkeyPatch) -> None:
+    """run_server() passes timeout_graceful_shutdown=10 to uvicorn.run()."""
+    captured: dict = {}
+
+    def fake_uvicorn_run(*args: object, **kwargs: object) -> None:
+        captured.update(kwargs)
+
+    monkeypatch.setattr("uvicorn.run", fake_uvicorn_run)
+
+    from headroom.proxy.server import run_server
+
+    run_server(ProxyConfig(), print_banner=False)
+
+    assert "timeout_graceful_shutdown" in captured, (
+        "uvicorn.run() must receive timeout_graceful_shutdown kwarg"
+    )
+    assert captured["timeout_graceful_shutdown"] == 10
+
+
+# ---------------------------------------------------------------------------
+# Integration: lifespan logs proxy_shutdown event on teardown
+# ---------------------------------------------------------------------------
+
+
+def test_lifespan_logs_shutdown_event(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The lifespan finally-block logs event=proxy_shutdown when the app tears down.
+
+    caplog cannot capture records from loggers that emit before propagation is
+    configured, so this test installs a custom handler directly on
+    ``headroom.proxy`` and checks that handler's records.
+    """
+    # Collect log records manually because caplog propagation is unreliable
+    # when the root logger has pre-existing basicConfig handlers.
+    captured: list[logging.LogRecord] = []
+
+    class _Capture(logging.Handler):
+        def emit(self, record: logging.LogRecord) -> None:
+            captured.append(record)
+
+    proxy_logger = logging.getLogger("headroom.proxy")
+    capture_handler = _Capture()
+    proxy_logger.addHandler(capture_handler)
+
+    try:
+        # Prevent sys.exit(78) from _check_rust_core when Rust extension absent
+        monkeypatch.setattr(
+            "headroom.proxy.server._check_rust_core", lambda: ("disabled", "test-mock")
+        )
+
+        config = ProxyConfig(
+            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,
+            image_optimize=False,
+        )
+        app = create_app(config)
+
+        from fastapi.testclient import TestClient
+
+        with TestClient(app, raise_server_exceptions=False):
+            pass  # lifespan shutdown runs when the context manager exits
+
+    finally:
+        proxy_logger.removeHandler(capture_handler)
+
+    shutdown_records = [r for r in captured if "event=proxy_shutdown" in r.getMessage()]
+    assert shutdown_records, "Expected at least one log record containing 'event=proxy_shutdown'"
+
+
+# ---------------------------------------------------------------------------
+# Lifespan shutdown: bounded await (_timed helper)
+# ---------------------------------------------------------------------------
+
+
+def test_lifespan_shutdown_completes_when_proxy_shutdown_hangs(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Lifespan shutdown must complete even if proxy.shutdown() never returns.
+
+    Before the fix, an unbounded ``await _beacon.stop()`` would block the
+    lifespan finally-block forever, requiring a second Ctrl+C.  The fix wraps
+    every shutdown await with asyncio.wait_for so a slow step is skipped
+    after its timeout and teardown continues.
+    """
+    import asyncio
+
+    async def hanging_stop() -> None:
+        await asyncio.sleep(9999)  # simulate a blocked network call
+
+    # Prevent sys.exit(78) from the Rust-core check
+    monkeypatch.setattr("headroom.proxy.server._check_rust_core", lambda: ("disabled", "test-mock"))
+
+    config = ProxyConfig(
+        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,
+        image_optimize=False,
+    )
+    app = create_app(config)
+
+    import headroom.proxy.server as server_mod
+
+    monkeypatch.setattr(server_mod.HeadroomProxy, "shutdown", lambda self: hanging_stop())
+
+    # If the fix is absent this would hang; with the fix it returns quickly.
+    import time
+
+    from fastapi.testclient import TestClient
+
+    start = time.monotonic()
+    with TestClient(app, raise_server_exceptions=False):
+        pass
+    elapsed = time.monotonic() - start
+    # Teardown should complete well within 15 s even with the timeout; hanging
+    # without the fix would block until the test runner times out (~60 s).
+    assert elapsed < 15.0, f"Lifespan shutdown took too long: {elapsed:.1f}s"
+
+
+def test_lifespan_shutdown_completes_when_proxy_shutdown_raises(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Lifespan shutdown must complete even if proxy.shutdown() raises.
+
+    The _timed wrapper catches both TimeoutError and arbitrary exceptions,
+    logs a warning, and continues so all subsequent teardown steps still run.
+    """
+
+    async def raising_shutdown() -> None:
+        raise RuntimeError("simulated shutdown failure")
+
+    monkeypatch.setattr("headroom.proxy.server._check_rust_core", lambda: ("disabled", "test-mock"))
+
+    config = ProxyConfig(
+        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,
+        image_optimize=False,
+    )
+    app = create_app(config)
+
+    import headroom.proxy.server as server_mod
+
+    monkeypatch.setattr(server_mod.HeadroomProxy, "shutdown", lambda self: raising_shutdown())
+
+    from fastapi.testclient import TestClient
+
+    # Should not raise — the _timed helper swallows the exception with a warning
+    with TestClient(app, raise_server_exceptions=False):
+        pass