fix(tui_gateway): drain in-flight turns before finalizing sessions on compute-host shutdown

tests/tui_gateway/test_compute_host_phase1.py

@@ -8,6 +8,7 @@
 
 import pytest
 
+from tui_gateway import compute_host, server
 from tui_gateway.compute_host import ComputeHost, _default_workers
 from tui_gateway.host_supervisor import (
     MUTATOR_ROUTE_TABLE,
@@ -132,3 +133,175 @@ class _Agent:
     }
 
 
+def _record_finalize(monkeypatch, events: list[str], *sids: str) -> None:
+    """Give ``flush_all_sessions`` sessions and record which ones finalize."""
+    keys = sids or ("s1",)
+    monkeypatch.setattr(
+        server,
+        "_sessions",
+        {sid: {"session_key": sid} for sid in keys},
+        raising=False,
+    )
+    monkeypatch.setattr(
+        server,
+        "_finalize_session",
+        lambda _session, end_reason="tui_close": events.append(
+            f"finalize:{_session['session_key']}:{end_reason}"
+        ),
+        raising=False,
+    )
+
+
+def _register_turn(host: ComputeHost, fn, sid: str = "s1") -> None:
+    """Submit a turn exactly the way ``_handle_turn_start`` does."""
+    host._track_turn_future(host._executor.submit(fn), sid)
+
+
+def test_shutdown_drains_in_flight_turn_before_finalizing_sessions(monkeypatch):
+    events: list[str] = []
+    _record_finalize(monkeypatch, events)
+
+    host = ComputeHost(stdout=io.StringIO(), heartbeat_secs=0)
+    running = threading.Event()
+
+    def _turn() -> None:
+        running.set()
+        time.sleep(0.3)
+        events.append("turn_end")
+
+    _register_turn(host, _turn, sid="s1")
+    assert running.wait(timeout=5.0)
+
+    host.shutdown(reason="sigterm", wait=3.0)
+
+    # ``_finalize_session`` latches on ``session["_finalized"]``, so its single
+    # run has to observe the finished turn or the tail is unpersistable. A turn
+    # that *did* drain must still finalize — the live-turn skip must not
+    # over-reach into sessions whose work is done.
+    assert events == ["turn_end", "finalize:s1:compute_host_sigterm"]
+
+    # The done-callback still has to remove the entry now that the container is
+    # a dict: ``set.discard`` was a valid bare callback, ``dict.pop`` is not.
+    deadline = time.monotonic() + 2.0
+    while host._turn_futures and time.monotonic() < deadline:
+        time.sleep(0.01)
+    assert host._turn_futures == {}, "in-flight turns must not accumulate"
+
+
+def test_shutdown_retains_a_live_turns_session_when_the_drain_deadline_expires(monkeypatch):
+    wait = 1.0
+    events: list[str] = []
+    _record_finalize(monkeypatch, events, "live", "idle")
+
+    host = ComputeHost(stdout=io.StringIO(), heartbeat_secs=0)
+    release = threading.Event()
+    running = threading.Event()
+
+    def _stuck_turn() -> None:
+        running.set()
+        release.wait(timeout=30.0)
+
+    _register_turn(host, _stuck_turn, sid="live")
+    assert running.wait(timeout=5.0)
+
+    try:
+        started = time.monotonic()
+        host.shutdown(reason="sigterm", wait=wait)
+        elapsed = time.monotonic() - started
+    finally:
+        release.set()
+
+    # ``_finalize_session`` is one-shot, and the ``shutdown(wait=False)`` that
+    # follows does not join the turn. Spending "live"'s single latch mid-turn
+    # would leave it permanently un-finalizable and release its active-session
+    # lease out from under running work — the same lifecycle race the drain
+    # exists to close, just moved past the deadline. It is retained unfinalized
+    # for recovery instead. A turn outliving the window must not cost the flush
+    # for anyone else, so "idle" still finalizes in the same pass.
+    assert events == ["finalize:idle:compute_host_sigterm"]
+    assert elapsed < wait
+
+
+def test_shutdown_retains_live_sessions_within_the_stdin_closed_budget(monkeypatch):
+    """The tightest real budget any caller uses is ``wait=2.0``.
+
+    ``run_host`` finalizes through ``host.shutdown(reason="stdin_closed",
+    wait=2.0)``, which is where the reserve — ``wait`` minus
+    ``min(_FLUSH_RESERVE_SECS, wait / 2)`` — has the least room to work with.
+    The retain-live-sessions rule must hold there without costing the flush for
+    idle sessions and without pushing the call past the budget the supervisor's
+    kill escalation is timed against.
+    """
+    wait = 2.0
+    drain_budget = wait - min(compute_host._FLUSH_RESERVE_SECS, wait / 2.0)
+
+    events: list[str] = []
+    _record_finalize(monkeypatch, events, "live", "idle")
+
+    host = ComputeHost(stdout=io.StringIO(), heartbeat_secs=0)
+    release = threading.Event()
+    running = threading.Event()
+
+    def _stuck_turn() -> None:
+        running.set()
+        release.wait(timeout=30.0)
+
+    _register_turn(host, _stuck_turn, sid="live")
+    assert running.wait(timeout=5.0)
+
+    try:
+        started = time.monotonic()
+        host.shutdown(reason="stdin_closed", wait=wait)
+        elapsed = time.monotonic() - started
+    finally:
+        release.set()
+
+    assert events == ["finalize:idle:compute_host_stdin_closed"]
+    assert elapsed >= drain_budget - 1e-6, "the drain must use its full window"
+    assert elapsed < wait
+
+
+def test_shutdown_drain_sleep_never_overshoots_the_reserve(monkeypatch):
+    """The drain's per-tick sleep must be bounded by the time left to it.
+
+    A flat tick overshoots the drain deadline by up to one tick, eating the
+    reserve held back for ``flush_all_sessions``; for a small ``wait`` that is
+    the whole reserve. Asserting on the *requested* sleep totals rather than on
+    wall-clock keeps this deterministic: each sleep is clamped to the remaining
+    time, so the sum can never exceed the drain budget however the scheduler
+    interleaves.
+    """
+    wait = 0.34
+    drain_budget = wait - min(compute_host._FLUSH_RESERVE_SECS, wait / 2.0)
+
+    events: list[str] = []
+    _record_finalize(monkeypatch, events, "idle")
+
+    slept: list[float] = []
+    real_sleep = time.sleep
+
+    def _recording_sleep(seconds: float) -> None:
+        slept.append(seconds)
+        real_sleep(seconds)
+
+    monkeypatch.setattr(compute_host.time, "sleep", _recording_sleep)
+
+    host = ComputeHost(stdout=io.StringIO(), heartbeat_secs=0)
+    release = threading.Event()
+    running = threading.Event()
+
+    def _stuck_turn() -> None:
+        running.set()
+        release.wait(timeout=30.0)
+
+    _register_turn(host, _stuck_turn, sid="live")
+    assert running.wait(timeout=5.0)
+
+    try:
+        host.shutdown(reason="sigterm", wait=wait)
+    finally:
+        release.set()
+
+    assert events == ["finalize:idle:compute_host_sigterm"]
+    assert slept, "the drain loop should have ticked at least once"
+    assert sum(slept) <= drain_budget + 1e-6

tui_gateway/compute_host.py

@@ -19,7 +19,7 @@
 import uuid
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Callable
+from typing import Any, Callable, Collection
 
 from agent.interrupt_compat import request_hard_interrupt
 
@@ -124,6 +124,14 @@ def _build_sha() -> str:
         return "unknown"
 
 
+# Slice of ``ComputeHost.shutdown``'s budget held back for the post-drain
+# finalize.  ``HostSupervisor._terminate_pid`` SIGKILLs the host
+# ``_SHUTDOWN_TIMEOUT_SECS`` (10s — the same value as ``shutdown``'s default
+# ``wait``) after SIGTERM, so a drain allowed to consume the whole budget would
+# leave the flush racing that kill and persist nothing at all.
+_FLUSH_RESERVE_SECS = 1.0
+
+
 class ComputeHost:
     def __init__(
         self,
@@ -144,7 +152,10 @@ def __init__(
         self._boot_id = uuid.uuid4().hex
         self._progress_counter = 0
         self._progress_lock = threading.Lock()
-        self._turn_futures: set[concurrent.futures.Future] = set()
+        # Future -> the ``sid`` whose turn it is running.  ``shutdown`` needs to
+        # know *whose* turn is still live, not merely that something is, so that
+        # it can leave those sessions unfinalized; a bare set cannot answer that.
+        self._turn_futures: dict[concurrent.futures.Future, str] = {}
         self._turn_futures_lock = threading.Lock()
         self._transport = _HostTransport(self.emit)
         self._heartbeat_secs = (
@@ -167,23 +178,86 @@ def close(self) -> None:
         self._executor.shutdown(wait=False, cancel_futures=True)
 
     def shutdown(self, *, reason: str = "shutdown", wait: float = 10.0) -> None:
+        """Drain in-flight turns, then finalize every session.
+
+        Order matters. ``_finalize_session`` is a one-shot latch: it sets
+        ``session["_finalized"]`` and every later call returns immediately, so
+        the flush gets exactly one chance to snapshot the session. Running it
+        before the drain meant that chance was spent while turns were still
+        producing output — the tail was unpersistable, ``on_session_end`` fired
+        with ``interrupted=True`` against a session that was still running, and
+        the active-session lease was released out from under a live turn. The
+        drain loop exists precisely so that work survives; finalizing first
+        defeated it.
+
+        ``_FLUSH_RESERVE_SECS`` of the budget — but never more than half of it,
+        so a short explicit ``wait`` still gets a real drain — is withheld from
+        the drain, so the flush still runs when in-flight turns outlast the
+        window. ``wait`` itself is unchanged, so this adds no shutdown latency
+        and no new exposure to the supervisor's kill escalation.
+
+        Sessions whose turn is *still running* when the drain deadline expires
+        are excluded from that flush. Finalizing one would spend its single
+        latch mid-turn — ``shutdown(wait=False, cancel_futures=True)`` below
+        does not join the turn — leaving the session permanently
+        un-finalizable and its active-session lease released out from under
+        live work: exactly the race the drain exists to close, just moved later.
+        Leaving them unfinalized keeps them recoverable instead. Sessions with
+        no live turn finalize here as they always have.
+
+        NOTE: ``server._shutdown_sessions`` is registered via ``atexit``
+        (``server.py``) and runs on ``SystemExit`` after ``shutdown()``
+        returns. It calls ``_finalize_session`` on any session still in
+        ``server._sessions`` — including ones skipped here whose turn is
+        still running, since ``_executor.shutdown(wait=False)`` only cancels
+        pending futures, not running ones. The orphan path (``os._exit(0)``)
+        bypasses atexit, so the skip is fully effective there. For the
+        SIGTERM and stdin_closed paths the atexit handler may re-finalize
+        skipped sessions; this is a pre-existing issue (the old finalize-
+        first order had the same atexit interaction) and does not make the
+        drain-before-finalize reordering worse. A follow-up could gate
+        ``_shutdown_sessions`` on ``not session.get("_finalized") and not
+        session.get("running")`` to close the gap.
+        """
         self._closed.set()
-        self.flush_all_sessions(reason=reason)
-        deadline = time.monotonic() + max(0.0, wait)
-        while time.monotonic() < deadline:
+        budget = max(0.0, wait)
+        deadline = time.monotonic() + budget - min(_FLUSH_RESERVE_SECS, budget / 2.0)
+        while True:
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                break
             with self._turn_futures_lock:
                 pending = [f for f in self._turn_futures if not f.done()]
             if not pending:
                 break
-            time.sleep(0.05)
+            # Bounded by ``remaining``: a flat 0.05s sleep would overshoot the
+            # deadline and eat into the reserve it is there to protect, which
+            # for a small ``wait`` can be the whole of it.
+            time.sleep(min(0.05, remaining))
+        with self._turn_futures_lock:
+            live_sids = {sid for future, sid in self._turn_futures.items() if sid and not future.done()}
+        self.flush_all_sessions(reason=reason, skip_sids=live_sids)
         self._executor.shutdown(wait=False, cancel_futures=True)
 
-    def flush_all_sessions(self, *, reason: str = "shutdown") -> None:
+    def flush_all_sessions(
+        self,
+        *,
+        reason: str = "shutdown",
+        skip_sids: Collection[str] | None = None,
+    ) -> None:
+        """Finalize every server session except the ones named in ``skip_sids``.
+
+        ``skip_sids`` carries the sessions whose turn is still live, which must
+        not spend their one-shot ``_finalize_session`` while running.
+        """
         try:
             from tui_gateway import server
         except Exception:
             return
-        for session in list(getattr(server, "_sessions", {}).values()):
+        skip = set(skip_sids or ())
+        for sid, session in list(getattr(server, "_sessions", {}).items()):
+            if sid in skip:
+                continue
             try:
                 server._finalize_session(session, end_reason=f"compute_host_{reason}")
             except Exception:
@@ -229,15 +303,28 @@ def _handle_seed(self, frame: dict[str, Any]) -> None:
         self._sessions[sid] = HostSession(sid=sid, agent=SpikeAgent(sid, list(history)))
         self.emit({"type": "session.seeded", "sid": sid, "request_id": frame.get("request_id")})
 
+    def _track_turn_future(self, future: concurrent.futures.Future, sid: str) -> None:
+        """Register an in-flight turn against the session running it.
+
+        The callback has to remove the entry under the lock — a bare
+        ``dict.pop`` bound method is not the drop-in ``set.discard`` was — or
+        the mapping grows for the life of the host.
+        """
+        with self._turn_futures_lock:
+            self._turn_futures[future] = sid
+        future.add_done_callback(self._untrack_turn_future)
+
+    def _untrack_turn_future(self, future: concurrent.futures.Future) -> None:
+        with self._turn_futures_lock:
+            self._turn_futures.pop(future, None)
+
     def _handle_turn_start(self, frame: dict[str, Any]) -> None:
         sid = str(frame.get("sid") or "")
         if sid in self._sessions:
             self._handle_spike_turn_start(frame)
             return
         future = self._executor.submit(self._run_real_turn, dict(frame))
-        with self._turn_futures_lock:
-            self._turn_futures.add(future)
-        future.add_done_callback(self._turn_futures.discard)
+        self._track_turn_future(future, sid)
 
     def _handle_spike_turn_start(self, frame: dict[str, Any]) -> None:
         sid = str(frame.get("sid") or "")
@@ -251,9 +338,7 @@ def _handle_spike_turn_start(self, frame: dict[str, Any]) -> None:
                 return
             session.running = True
         future = self._executor.submit(self._run_spike_turn, session, dict(frame))
-        with self._turn_futures_lock:
-            self._turn_futures.add(future)
-        future.add_done_callback(self._turn_futures.discard)
+        self._track_turn_future(future, sid)
 
     def _handle_interrupt(self, frame: dict[str, Any]) -> None:
         sid = str(frame.get("sid") or "")