Address review: bound cleanup, always reset context, add real-DB cancel test

Review findings on the cancellation-shield fix:

- Bound _run_to_completion with _CLEANUP_TIMEOUT_SECONDS. Shielded cleanup is
  uninterruptible by design, so a wedged connection (dead socket, no libpq
  timeout) would otherwise pin the task forever; past the bound we abandon
  cleanup rather than hang. Switched the wait to asyncio.wait (shields fut,
  supports a timeout) and guarded fut.exception()/fut.cancelled() so loop
  shutdown can't raise from the tail.
- tracked_db: reset request_context in its own finally so a cancelled cleanup
  path can't leak the task-scoped contextvar into a reused long-lived task
  (deriver).
- Tests: reset-context-on-cancellation, abandon-wedged-cleanup, and a real-DB
  wire cancellation test asserting no `idle in transaction` backend lingers
  after a cancelled in-flight write transaction (pg_stat_activity).

Full suite: 21 passed live (incl. 4 pre-existing + 4 new real-DB tests).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

src/dependencies.py

@@ -13,25 +13,45 @@ from src.db import ReadSessionLocal, SessionLocal, request_context
 
 logger = logging.getLogger(__name__)
 
+# Upper bound on how long shielded cleanup may run before we stop waiting on it.
+# Healthy ROLLBACK/close take milliseconds; this only bites a wedged connection
+# (e.g. a black-holed socket with no libpq timeout) so it can't pin the task
+# forever. Cleanup is abandoned (not cancelled) past this — the pool/OS reclaims
+# the connection — since a shielded task cannot be interrupted anyway.
+_CLEANUP_TIMEOUT_SECONDS = 10.0
+
 
 async def _run_to_completion(coro: Coroutine[Any, Any, None]) -> None:
     """Run ``coro`` to completion even if the current task is cancelled (DEV-1861).
 
-    Runs the coro as a detached task and keeps awaiting it through
-    ``asyncio.shield`` (which never cancels the inner task), re-raising the
-    cancellation only afterward. Survives repeated re-delivery. NB:
-    ``anyio.CancelScope(shield=True)`` does NOT work here — it ignores a native
-    ``asyncio.Task.cancel()``, which is how Starlette and the deriver cancel.
+    Runs the coro as a detached task and waits on it via ``asyncio.wait`` (which
+    never cancels the inner task), re-raising the cancellation only afterward.
+    Survives repeated re-delivery. NB: ``anyio.CancelScope(shield=True)`` does
+    NOT work here — it ignores a native ``asyncio.Task.cancel()``, which is how
+    Starlette and the deriver cancel.
     """
     fut = asyncio.ensure_future(coro)
     cancelled = False
     while not fut.done():
         try:
-            await asyncio.shield(fut)
+            # wait() shields fut from our cancellation but still bounds the wait.
+            await asyncio.wait({fut}, timeout=_CLEANUP_TIMEOUT_SECONDS)
         except asyncio.CancelledError:
             cancelled = True  # keep waiting so ROLLBACK/close reach Postgres
-    if (exc := fut.exception()) is not None:  # cleanup swallows its own errors
-        logger.exception("session cleanup task failed", exc_info=exc)
+            continue
+        if not fut.done():
+            # Wedged cleanup (dead socket). Stop pinning this task; leave fut to
+            # error out / be reclaimed rather than block indefinitely.
+            logger.error(
+                "DB session cleanup exceeded %.0fs; abandoning to avoid pinning "
+                "the task",
+                _CLEANUP_TIMEOUT_SECONDS,
+            )
+            break
+    if fut.cancelled():
+        cancelled = True  # cleanup itself was cancelled (close() still ran)
+    elif fut.done() and (exc := fut.exception()) is not None:
+        logger.error("session cleanup task failed", exc_info=exc)
     if cancelled:
         raise asyncio.CancelledError()
 
@@ -116,9 +136,13 @@ async def tracked_db(operation_name: str | None = None, *, read_only: bool = Fal
     try:
         yield db
     finally:
-        await _run_to_completion(_finalize_session(db))  # shielded — see get_db
-        if token:  # Only reset if we set it
-            request_context.reset(token)
+        try:
+            await _run_to_completion(_finalize_session(db))  # shielded — see get_db
+        finally:
+            # Must run even when cleanup re-raises CancelledError, or a reused
+            # long-lived task (e.g. the deriver) leaks this contextvar.
+            if token:  # Only reset if we set it
+                request_context.reset(token)
 
 
 db: AsyncSession = Depends(get_db)

tests/test_dependencies.py

@@ -477,3 +477,90 @@ async def test_tracked_db_survives_cancellation_storm_during_cleanup(
 
     assert fake_db.rollback_calls == 1
     assert fake_db.close_calls == 1  # cleanup completed despite the storm
+
+
+@pytest.mark.asyncio
+async def test_tracked_db_resets_request_context_on_cancellation(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # The task-scoped contextvar must be reset even when cleanup re-raises
+    # CancelledError, else a reused long-lived task (e.g. the deriver) carries a
+    # stale `task:...` context into later work and corrupts tracing/attribution.
+    fake_db = FakeSession()
+    monkeypatch.setattr(dependencies_module, "SessionLocal", lambda: fake_db)
+
+    clear_token = request_context.set(None)
+    try:
+        with pytest.raises(asyncio.CancelledError):
+            async with real_tracked_db("ctx_cancel_op"):
+                assert (request_context.get() or "").startswith("task:ctx_cancel_op:")
+                raise asyncio.CancelledError()
+        assert request_context.get() is None  # reset despite the cancel
+    finally:
+        request_context.reset(clear_token)
+
+
+@pytest.mark.asyncio
+async def test_run_to_completion_abandons_wedged_cleanup(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # Shielded cleanup is uninterruptible by design; a wedged connection (dead
+    # socket, no libpq timeout) must therefore NOT pin the caller forever.
+    # _run_to_completion abandons cleanup past _CLEANUP_TIMEOUT_SECONDS.
+    monkeypatch.setattr(dependencies_module, "_CLEANUP_TIMEOUT_SECONDS", 0.2)
+    started = asyncio.Event()
+
+    async def wedged() -> None:
+        started.set()
+        await asyncio.sleep(3600)  # never completes
+
+    # Returns (abandons) well under the sleep instead of hanging.
+    await asyncio.wait_for(dependencies_module._run_to_completion(wedged()), timeout=5)
+    assert started.is_set()
+
+
+@pytest.mark.asyncio
+async def test_tracked_db_no_idle_in_transaction_after_real_cancellation() -> None:
+    # Wire-level chaos test (ticket verification): cancel a REAL in-flight write
+    # transaction mid-flight and assert, from an independent connection, that the
+    # backend is not left parked 'idle in transaction'. Complements the mock
+    # tests above with a real pooled connection receiving ROLLBACK under cancel.
+    from sqlalchemy import text
+
+    from src.db import read_engine
+
+    entered = asyncio.Event()
+    pid_box: dict[str, int] = {}
+
+    async def worker() -> None:
+        async with real_tracked_db("chaos_cancel") as db:
+            pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar()
+            pid_box["pid"] = int(pid)  # pyright: ignore[reportArgumentType]
+            entered.set()
+            await asyncio.sleep(3600)  # hold the open BEGIN until cancelled
+
+    async def backend_state(pid: int) -> str | None:
+        async with read_engine.connect() as obs:
+            return (
+                await obs.execute(
+                    text("SELECT state FROM pg_stat_activity WHERE pid = :p"),
+                    {"p": pid},
+                )
+            ).scalar()
+
+    task = asyncio.ensure_future(worker())
+    await entered.wait()
+    pid = pid_box["pid"]
+    assert await backend_state(pid) == "idle in transaction"  # precondition
+
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    state = await backend_state(pid)
+    for _ in range(20):  # allow shielded cleanup a moment to land
+        if state != "idle in transaction":
+            break
+        await asyncio.sleep(0.1)
+        state = await backend_state(pid)
+    assert state != "idle in transaction", f"backend left {state!r} after cancel"

uv.lock

@@ -7,10 +7,6 @@ resolution-markers = [
     "python_full_version < '3.13'",
 ]
 
-[options]
-exclude-newer = "2026-06-10T19:45:16.447512Z"
-exclude-newer-span = "P5D"
-
 [manifest]
 members = [
     "honcho",