fix(tui_gateway): drain in-flight turns before finalizing sessions on compute-host shutdown - #77330
Merged
kshitijk4poor merged 4 commits intoAug 3, 2026
Conversation
… compute-host shutdown 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.
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.
…n deadline expires 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.
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').
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Drains in-flight turns before finalizing sessions on
ComputeHost.shutdown(), preventing data loss when the TUI/dashboard is quit mid-answer. Previouslyshutdown()finalized first (spending the one-shot_finalize_sessionlatch), then drained — so the unflushed message tail produced during the drain window was never persisted.Changes
tui_gateway/compute_host.py—ComputeHost.shutdown()now runs the in-flight-turn drain loop beforeflush_all_sessions(), withholds_FLUSH_RESERVE_SECS(1.0, capped at half the budget) from the drain deadline so the flush still runs when turns outlast the window, and skips finalizing sessions whose turn is still live at drain timeout._turn_futureschanged fromset[Future]todict[Future, str](needed to know which session a live turn belongs to). Extracted_track_turn_future/_untrack_turn_futurehelpers.tests/tui_gateway/test_compute_host_phase1.py— 4 regression tests covering drain-before-finalize ordering, deadline-expiry retain, stdin_closed budget, and sleep clamping.Follow-up commit (ours)
Added docstring note documenting the
atexitinteraction:server._shutdown_sessions()(registered viaatexitatserver.py:1172) runs onSystemExitaftershutdown()returns and may re-finalize sessions skipped here whose turn is still running. The orphan path (os._exit(0)) bypasses atexit. Pre-existing issue — the old finalize-first order had the same interaction.Validation
Closes #77053
Credit
Cherry-picked from @briandevans's PR #77053 with authorship preserved. 3 contributor commits + 1 follow-up docstring commit.