fix(deriver): bound the shutdown drain and release claims before exit

Shutdown waited on asyncio.gather(*active_tasks) with no timeout, and released
work-unit claims only via the finally in initialize(). Both assume the process
gets as long as it needs.

It does not. A platform can terminate the process on a budget shorter than a
single work unit takes -- a deriver LLM call alone can run for tens of seconds
-- and when that budget expires the process is killed outright. The unbounded
wait is then still in progress, initialize() never returns, its finally never
runs, and the claims are left held. The work is not lost, but nothing may touch
it until cleanup_stale_work_units notices, which takes
STALE_SESSION_TIMEOUT_MINUTES plus a jittered poll interval.

So shutdown is a race to hand the claims back, not an opportunity to finish the
work. Two changes:

Drain is bounded by SHUTDOWN_DRAIN_TIMEOUT_SECONDS (default 8). Work that fits
in the budget still completes; work that does not is cancelled rather than
waited on, since the work is reclaimable and the remaining budget is not.

cleanup() is called at the end of shutdown() rather than relying on the finally
in initialize(). Claims are given up while the process is still alive, so
another worker repicks the work on its next poll instead of waiting out the
stale sweep. cleanup() clears worker_ownership, so the existing call in the
finally becomes a no-op and the ordering is idempotent.

The drain deliberately runs before the release rather than after. Releasing
first would leave a window in which a task that completes during shutdown has
already had its claim taken by another worker, and the same queue items get
derived twice. Cancelling first and releasing after keeps that window closed,
and both steps fit inside a short budget.

Cancelling also keeps the error reporting honest without special-casing it.
asyncio.CancelledError derives from BaseException rather than Exception, so an
abandoned task passes straight through the `except Exception` handlers that
report to Sentry. Shutdown stops generating spurious errors on its own, while a
task that genuinely fails during the drain, or a cleanup() that cannot reach the
database, still reports as it should.

Tests cover a task that outlives the budget being abandoned with cleanup still
awaited, and a task that fits in the budget still finishing.

src/config.py

@@ -908,6 +908,11 @@ class DeriverSettings(HonchoSettings):
     # to 0.0 to disable.
     POLLING_JITTER_RATIO: Annotated[float, Field(default=0.5, ge=0.0, le=1.0)] = 0.5
     STALE_SESSION_TIMEOUT_MINUTES: Annotated[int, Field(default=5, gt=0, le=1440)] = 5
+    # How long shutdown waits for in-flight work before abandoning it. Must
+    # leave room in the shutdown budget for the claim release that follows.
+    SHUTDOWN_DRAIN_TIMEOUT_SECONDS: Annotated[
+        float, Field(default=8.0, ge=0.0, le=600.0)
+    ] = 8.0
     # Minimum (jittered) spacing between stale-work-unit cleanup runs
     STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS: Annotated[
         float, Field(default=60.0, ge=0.0, le=3600.0)

src/deriver/queue_manager.py

@@ -224,7 +224,12 @@ class QueueManager:
             await self.cleanup()
 
     async def shutdown(self, sig: signal.Signals) -> None:
-        """Handle graceful shutdown"""
+        """Hand work-unit claims back before the process dies.
+
+        The shutdown budget can be shorter than one work unit takes, so finishing
+        the work is not the goal. A claim still held at exit strands its work
+        until the stale sweep reclaims it.
+        """
         logger.info(f"Received exit signal {sig.name}...")
         self.shutdown_event.set()
 
@@ -236,9 +241,24 @@ class QueueManager:
 
         if self.active_tasks:
             logger.info(
-                f"Waiting for {len(self.active_tasks)} active tasks to complete..."
+                "Waiting up to %ss for %d active tasks to complete...",
+                settings.DERIVER.SHUTDOWN_DRAIN_TIMEOUT_SECONDS,
+                len(self.active_tasks),
             )
-            await asyncio.gather(*self.active_tasks, return_exceptions=True)
+            drain = asyncio.gather(*self.active_tasks, return_exceptions=True)
+            try:
+                await asyncio.wait_for(
+                    drain, timeout=settings.DERIVER.SHUTDOWN_DRAIN_TIMEOUT_SECONDS
+                )
+            except (TimeoutError, asyncio.TimeoutError):
+                # wait_for already cancelled the gather. The work is
+                # reclaimable; the remaining budget is not.
+                logger.warning(
+                    "Drain timed out; abandoning in-flight work units for reclaim"
+                )
+
+        # Not left to the finally in initialize(): a hard kill skips it.
+        await self.cleanup()
 
     async def cleanup(self) -> None:
         """Clean up owned work units"""

tests/deriver/test_shutdown_drain.py

@@ -0,0 +1,73 @@
+"""Shutdown must bound the drain and give claims back before the process dies.
+
+A claim still held at exit strands its work until the stale sweep reclaims it.
+"""
+
+import asyncio
+import signal
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from src.deriver.queue_manager import QueueManager
+
+
+@pytest.mark.asyncio
+async def test_drain_is_bounded_and_claims_are_released() -> None:
+    """A task outliving the budget is abandoned; cleanup still runs."""
+    qm = QueueManager()
+
+    started = asyncio.Event()
+
+    async def never_finishes() -> None:
+        started.set()
+        await asyncio.sleep(3600)
+
+    task = asyncio.create_task(never_finishes())
+    qm.add_task(task)
+    await started.wait()
+
+    with (
+        patch.object(qm, "cleanup", new=AsyncMock()) as cleanup,
+        patch.object(qm.dream_scheduler, "shutdown", new=AsyncMock()),
+        patch.object(qm.reconciler_scheduler, "shutdown", new=AsyncMock()),
+        patch("src.deriver.queue_manager.settings") as settings,
+    ):
+        settings.DERIVER.SHUTDOWN_DRAIN_TIMEOUT_SECONDS = 0.05
+
+        loop = asyncio.get_running_loop()
+        began = loop.time()
+        await qm.shutdown(signal.SIGTERM)
+        elapsed = loop.time() - began
+
+    assert elapsed < 2, f"shutdown blocked on the stuck task for {elapsed:.1f}s"
+    assert qm.shutdown_event.is_set()
+    cleanup.assert_awaited_once()  # claims released before exit
+
+    assert task.cancelled() or task.done()
+
+
+@pytest.mark.asyncio
+async def test_fast_work_still_completes_before_claims_are_released() -> None:
+    """Bounding the drain must not cut short work that fits."""
+    qm = QueueManager()
+    finished = False
+
+    async def quick() -> None:
+        nonlocal finished
+        await asyncio.sleep(0.01)
+        finished = True
+
+    qm.add_task(asyncio.create_task(quick()))
+
+    with (
+        patch.object(qm, "cleanup", new=AsyncMock()) as cleanup,
+        patch.object(qm.dream_scheduler, "shutdown", new=AsyncMock()),
+        patch.object(qm.reconciler_scheduler, "shutdown", new=AsyncMock()),
+        patch("src.deriver.queue_manager.settings") as settings,
+    ):
+        settings.DERIVER.SHUTDOWN_DRAIN_TIMEOUT_SECONDS = 5.0
+        await qm.shutdown(signal.SIGTERM)
+
+    assert finished
+    cleanup.assert_awaited_once()