From 346cecf7228424944ac8bce543b09913efc078ab Mon Sep 17 00:00:00 2001 From: briandevans <252620095+briandevans@users.noreply.github.com> Date: Sun, 2 Aug 2026 12:33:52 -0700 Subject: [PATCH 1/4] fix(tui_gateway): drain in-flight turns before finalizing sessions on compute-host shutdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ComputeHost.shutdown() called flush_all_sessions() before its own in-flight turn drain loop. server._finalize_session latches on session["_finalized"] and every later call returns immediately, so that one flush was spent while turns were still producing output: the unflushed tail was never persisted, commit_memory_session wrote long-term memory from a truncated transcript, the session's DB row was marked ended while it was live, on_session_end fired with completed=False/interrupted=True against a running session, and the active-session lease was released out from under a turn. The drain loop exists precisely so that mid-turn work survives a teardown; finalizing first defeated it. Reachable from all three teardown paths: the parent/orphan guard (which os._exit(0)s immediately after), the SIGTERM/SIGINT handler, and stdin close. Drain first, then flush. A slice of the caller's budget (_FLUSH_RESERVE_SECS, 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 turns outlast the window: HostSupervisor SIGKILLs the host _SHUTDOWN_TIMEOUT_SECS after SIGTERM — 10.0s, the same value as shutdown()'s default wait — so a drain allowed to consume the whole budget would leave the durability write racing that kill. `wait` itself is unchanged, so total shutdown latency and the SIGTERM->SIGKILL margin are unchanged. --- tests/tui_gateway/test_compute_host_phase1.py | 69 +++++++++++++++++++ tui_gateway/compute_host.py | 31 ++++++++- 2 files changed, 98 insertions(+), 2 deletions(-) diff --git a/tests/tui_gateway/test_compute_host_phase1.py b/tests/tui_gateway/test_compute_host_phase1.py index cc1b6347137c8..c7cb621ceca16 100644 --- a/tests/tui_gateway/test_compute_host_phase1.py +++ b/tests/tui_gateway/test_compute_host_phase1.py @@ -8,6 +8,7 @@ import pytest +from tui_gateway import server from tui_gateway.compute_host import ComputeHost, _default_workers from tui_gateway.host_supervisor import ( MUTATOR_ROUTE_TABLE, @@ -132,3 +133,71 @@ class _Agent: } +def _record_finalize(monkeypatch, events: list[str]) -> None: + """Give ``flush_all_sessions`` one session and record when it finalizes.""" + monkeypatch.setattr(server, "_sessions", {"s1": {"session_key": "s1"}}, raising=False) + monkeypatch.setattr( + server, + "_finalize_session", + lambda _session, end_reason="tui_close": events.append(f"finalize:{end_reason}"), + raising=False, + ) + + +def _register_turn(host: ComputeHost, fn) -> None: + """Submit a turn exactly the way ``_handle_turn_start`` does.""" + future = host._executor.submit(fn) + with host._turn_futures_lock: + host._turn_futures.add(future) + future.add_done_callback(host._turn_futures.discard) + + +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) + 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. + assert events == ["turn_end", "finalize:compute_host_sigterm"] + + +def test_shutdown_still_finalizes_when_the_drain_deadline_expires(monkeypatch): + events: list[str] = [] + _record_finalize(monkeypatch, events) + + 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) + assert running.wait(timeout=5.0) + + try: + started = time.monotonic() + host.shutdown(reason="sigterm", wait=1.0) + elapsed = time.monotonic() - started + finally: + release.set() + + # A turn that outlives the window must not cost the flush entirely: the + # supervisor's SIGKILL lands on the same deadline this budget comes from, + # so the drain has to stop short and leave the finalize room to run. + assert events == ["finalize:compute_host_sigterm"] + assert elapsed < 1.0 diff --git a/tui_gateway/compute_host.py b/tui_gateway/compute_host.py index 1f255533bd7ba..e9fbe93ec036a 100644 --- a/tui_gateway/compute_host.py +++ b/tui_gateway/compute_host.py @@ -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, @@ -167,15 +175,34 @@ 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. + """ self._closed.set() - self.flush_all_sessions(reason=reason) - deadline = time.monotonic() + max(0.0, wait) + budget = max(0.0, wait) + deadline = time.monotonic() + budget - min(_FLUSH_RESERVE_SECS, budget / 2.0) while time.monotonic() < deadline: 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) + self.flush_all_sessions(reason=reason) self._executor.shutdown(wait=False, cancel_futures=True) def flush_all_sessions(self, *, reason: str = "shutdown") -> None: From 261c7a8eb661c43e26b3f60987464966b60a206c Mon Sep 17 00:00:00 2001 From: briandevans <252620095+briandevans@users.noreply.github.com> Date: Sun, 2 Aug 2026 12:42:53 -0700 Subject: [PATCH 2/4] fix(tui_gateway): bound the shutdown drain sleep by the time left to it The drain loop slept a flat 0.05s per tick, so it could overshoot its deadline by up to one tick and spend part of the reserve withheld for flush_all_sessions(). For a small `wait` the reserve is itself half the budget, so a single overshoot can consume all of it: at wait=0.34 the drain budget is 0.17s but the loop requested 4 x 0.05 = 0.20s of sleep. Clamp each tick to the remaining time. The new test asserts on the summed *requested* sleep rather than wall-clock, which is deterministic: every sleep is bounded by the strictly-decreasing remainder, so the total can never exceed the drain budget regardless of how the scheduler interleaves. --- tests/tui_gateway/test_compute_host_phase1.py | 53 +++++++++++++++++-- tui_gateway/compute_host.py | 10 +++- 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/tests/tui_gateway/test_compute_host_phase1.py b/tests/tui_gateway/test_compute_host_phase1.py index c7cb621ceca16..07e1efd564924 100644 --- a/tests/tui_gateway/test_compute_host_phase1.py +++ b/tests/tui_gateway/test_compute_host_phase1.py @@ -8,7 +8,7 @@ import pytest -from tui_gateway import server +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, @@ -175,6 +175,7 @@ def _turn() -> None: def test_shutdown_still_finalizes_when_the_drain_deadline_expires(monkeypatch): + wait = 1.0 events: list[str] = [] _record_finalize(monkeypatch, events) @@ -191,7 +192,7 @@ def _stuck_turn() -> None: try: started = time.monotonic() - host.shutdown(reason="sigterm", wait=1.0) + host.shutdown(reason="sigterm", wait=wait) elapsed = time.monotonic() - started finally: release.set() @@ -200,4 +201,50 @@ def _stuck_turn() -> None: # supervisor's SIGKILL lands on the same deadline this budget comes from, # so the drain has to stop short and leave the finalize room to run. assert events == ["finalize:compute_host_sigterm"] - assert elapsed < 1.0 + 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) + + 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) + assert running.wait(timeout=5.0) + + try: + host.shutdown(reason="sigterm", wait=wait) + finally: + release.set() + + assert events == ["finalize:compute_host_sigterm"] + assert slept, "the drain loop should have ticked at least once" + assert sum(slept) <= drain_budget + 1e-6 diff --git a/tui_gateway/compute_host.py b/tui_gateway/compute_host.py index e9fbe93ec036a..ed951fa97a1d4 100644 --- a/tui_gateway/compute_host.py +++ b/tui_gateway/compute_host.py @@ -196,12 +196,18 @@ def shutdown(self, *, reason: str = "shutdown", wait: float = 10.0) -> None: self._closed.set() budget = max(0.0, wait) deadline = time.monotonic() + budget - min(_FLUSH_RESERVE_SECS, budget / 2.0) - while time.monotonic() < deadline: + 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)) self.flush_all_sessions(reason=reason) self._executor.shutdown(wait=False, cancel_futures=True) From 04f9f81ad0f3dd21713ee64556351775ba05d722 Mon Sep 17 00:00:00 2001 From: briandevans <252620095+briandevans@users.noreply.github.com> Date: Sun, 2 Aug 2026 16:34:08 -0700 Subject: [PATCH 3/4] fix(tui_gateway): retain live-turn sessions unfinalized when the drain deadline expires MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The drain reserves a slice of the shutdown budget so flush_all_sessions still runs when in-flight turns outlast the window. But that flush was unconditional: a session whose turn was still running got its one-shot _finalize_session spent mid-turn, and the executor.shutdown(wait=False, cancel_futures=True) immediately after does not join the turn. The session was then permanently un-finalizable and its active-session lease had been released out from under live work — the same persistence and lifecycle race the drain exists to close, just relocated past the deadline instead of removed. Give _turn_futures a session association (Future -> sid, the same key space as server._sessions) at both submit sites, and on deadline expiry exclude the sids whose futures are still running from the flush. Those sessions are retained unfinalized and therefore recoverable; sessions with no live turn finalize exactly as before. The done-callback now pops under the lock, since a bare dict.pop is not the drop-in set.discard was. wait semantics, the reserve math and the bounded per-tick sleep are unchanged, so this adds no shutdown latency. All three shutdown callers (orphan, sigterm, and the tight stdin_closed wait=2.0 path) funnel through this one function and are covered. --- tests/tui_gateway/test_compute_host_phase1.py | 101 ++++++++++++++---- tui_gateway/compute_host.py | 60 +++++++++-- 2 files changed, 128 insertions(+), 33 deletions(-) diff --git a/tests/tui_gateway/test_compute_host_phase1.py b/tests/tui_gateway/test_compute_host_phase1.py index 07e1efd564924..5b9ccb68c5ad2 100644 --- a/tests/tui_gateway/test_compute_host_phase1.py +++ b/tests/tui_gateway/test_compute_host_phase1.py @@ -133,23 +133,28 @@ class _Agent: } -def _record_finalize(monkeypatch, events: list[str]) -> None: - """Give ``flush_all_sessions`` one session and record when it finalizes.""" - monkeypatch.setattr(server, "_sessions", {"s1": {"session_key": "s1"}}, raising=False) +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:{end_reason}"), + lambda _session, end_reason="tui_close": events.append( + f"finalize:{_session['session_key']}:{end_reason}" + ), raising=False, ) -def _register_turn(host: ComputeHost, fn) -> None: +def _register_turn(host: ComputeHost, fn, sid: str = "s1") -> None: """Submit a turn exactly the way ``_handle_turn_start`` does.""" - future = host._executor.submit(fn) - with host._turn_futures_lock: - host._turn_futures.add(future) - future.add_done_callback(host._turn_futures.discard) + host._track_turn_future(host._executor.submit(fn), sid) def test_shutdown_drains_in_flight_turn_before_finalizing_sessions(monkeypatch): @@ -164,20 +169,29 @@ def _turn() -> None: time.sleep(0.3) events.append("turn_end") - _register_turn(host, _turn) + _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. - assert events == ["turn_end", "finalize:compute_host_sigterm"] + # 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_still_finalizes_when_the_drain_deadline_expires(monkeypatch): +def test_shutdown_retains_a_live_turns_session_when_the_drain_deadline_expires(monkeypatch): wait = 1.0 events: list[str] = [] - _record_finalize(monkeypatch, events) + _record_finalize(monkeypatch, events, "live", "idle") host = ComputeHost(stdout=io.StringIO(), heartbeat_secs=0) release = threading.Event() @@ -187,7 +201,7 @@ def _stuck_turn() -> None: running.set() release.wait(timeout=30.0) - _register_turn(host, _stuck_turn) + _register_turn(host, _stuck_turn, sid="live") assert running.wait(timeout=5.0) try: @@ -197,10 +211,53 @@ def _stuck_turn() -> None: finally: release.set() - # A turn that outlives the window must not cost the flush entirely: the - # supervisor's SIGKILL lands on the same deadline this budget comes from, - # so the drain has to stop short and leave the finalize room to run. - assert events == ["finalize:compute_host_sigterm"] + # ``_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 @@ -218,7 +275,7 @@ def test_shutdown_drain_sleep_never_overshoots_the_reserve(monkeypatch): drain_budget = wait - min(compute_host._FLUSH_RESERVE_SECS, wait / 2.0) events: list[str] = [] - _record_finalize(monkeypatch, events) + _record_finalize(monkeypatch, events, "idle") slept: list[float] = [] real_sleep = time.sleep @@ -237,7 +294,7 @@ def _stuck_turn() -> None: running.set() release.wait(timeout=30.0) - _register_turn(host, _stuck_turn) + _register_turn(host, _stuck_turn, sid="live") assert running.wait(timeout=5.0) try: @@ -245,6 +302,6 @@ def _stuck_turn() -> None: finally: release.set() - assert events == ["finalize:compute_host_sigterm"] + 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 diff --git a/tui_gateway/compute_host.py b/tui_gateway/compute_host.py index ed951fa97a1d4..706be07b0860b 100644 --- a/tui_gateway/compute_host.py +++ b/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 @@ -152,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 = ( @@ -192,6 +195,15 @@ def shutdown(self, *, reason: str = "shutdown", wait: float = 10.0) -> None: 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. """ self._closed.set() budget = max(0.0, wait) @@ -208,15 +220,30 @@ def shutdown(self, *, reason: str = "shutdown", wait: float = 10.0) -> None: # 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)) - self.flush_all_sessions(reason=reason) + 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: @@ -262,15 +289,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 "") @@ -284,9 +324,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 "") From a620c81601c1fc7c376fc1e8248e758b11721fd3 Mon Sep 17 00:00:00 2001 From: kshitij <82637225+kshitijk4poor@users.noreply.github.com> Date: Mon, 3 Aug 2026 09:54:59 +0530 Subject: [PATCH 4/4] docs: note atexit re-finalization interaction in shutdown() docstring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The PR's skip-live-sessions optimization is partially defeated by server._shutdown_sessions() registered via atexit (server.py:1172), which runs on SystemExit after shutdown() returns for the SIGTERM and stdin_closed paths. The orphan path (os._exit(0)) bypasses atexit. This is a pre-existing issue — the old finalize-first order had the same atexit interaction. The comment documents the gap and suggests a follow-up: gate _shutdown_sessions on not session.get('running'). --- tui_gateway/compute_host.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tui_gateway/compute_host.py b/tui_gateway/compute_host.py index 706be07b0860b..d90024557aad5 100644 --- a/tui_gateway/compute_host.py +++ b/tui_gateway/compute_host.py @@ -204,6 +204,20 @@ def shutdown(self, *, reason: str = "shutdown", wait: float = 10.0) -> None: 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() budget = max(0.0, wait)