fix(tui_gateway): drain in-flight turns before finalizing sessions on compute-host shutdown - #77053
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.
There was a problem hiding this comment.
🟡 Not ready to approve
The drain loop’s fixed 50ms sleep can overshoot the computed deadline and encroach on the reserved flush budget, undermining the intended “drain + flush within wait” behavior.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Pull request overview
Fixes TUI gateway compute-host shutdown durability by ensuring in-flight turns are drained before sessions are finalized, so final transcript persistence and memory commit are based on the completed turn output during shutdown (dashboard quit / SIGTERM / stdin-close paths).
Changes:
- Reorders
ComputeHost.shutdown()to drain pending turn futures before callingflush_all_sessions(). - Introduces a small reserved shutdown-time slice (
_FLUSH_RESERVE_SECS, capped to ≤ 50% ofwait) to ensure finalization still runs even when turns outlast the drain window. - Adds regression tests covering (a) drain-before-finalize ordering and (b) finalize still happening within the shutdown budget.
File summaries
| File | Description |
|---|---|
| tui_gateway/compute_host.py | Reorders shutdown drain/finalize and adds a reserved time slice to reach flush_all_sessions() reliably. |
| tests/tui_gateway/test_compute_host_phase1.py | Adds regression coverage for shutdown ordering and budget preservation. |
Review details
- Files reviewed: 2/2 changed files
- Comments generated: 1
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
| 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) |
There was a problem hiding this comment.
Good catch — fixed in 80b663c.
You're right that the flat tick could overshoot deadline by up to one 0.05s interval and spend part of the very reserve it was meant to protect, and that the small-wait case is where it actually bites: the reserve is clamped to budget / 2, so at wait=0.1 the reserve is 0.05s and one overshooting tick eats all of it.
The drain now clamps each tick to the time remaining:
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(min(0.05, remaining))Added test_shutdown_drain_sleep_never_overshoots_the_reserve to pin it. It asserts on the summed requested sleep rather than wall-clock so it can't flake: each sleep is bounded by a strictly decreasing remainder, so the total can never exceed the drain budget however the scheduler interleaves. At wait=0.34 (deliberately not a multiple of the tick) the drain budget is 0.17s — the test fails on the previous commit with assert 0.2 <= (0.17 + 1e-06) / sum([0.05, 0.05, 0.05, 0.05]), and passes now at exactly 0.17.
All three regression tests fail on clean origin/main and pass with the change; tests/tui_gateway/ + tests/test_tui_gateway_server.py are 836 passed.
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.
alt-glitch
added
type/bug
teknium1
left a comment
teknium1
left a comment
There was a problem hiding this comment.
Thanks for identifying the flush-before-drain ordering defect: current main still calls flush_all_sessions() before waiting in tui_gateway/compute_host.py:167-177, and _finalize_session() is a one-shot lifecycle boundary in tui_gateway/server.py:652-664.
Problems
- Blocking —
tui_gateway/compute_host.py:211: the new deadline path still callsflush_all_sessions()when a future is pending. The followingThreadPoolExecutor.shutdown(wait=False, cancel_futures=True)at:212does not wait for a running future, so a long turn is finalized while live. That recreates the one-shot-finalization race for deadline-exceeding turns.tests/tui_gateway/test_compute_host_phase1.py:177-204explicitly holds a turn open and accepts this outcome.
Suggested changes
- Make deadline expiry either bring the affected turn to a terminal state before finalization or retain it as an unfinalized, resumable in-flight session. Add an integration regression proving finalization side effects occur only after the associated turn has settled.
Automated hermes-sweeper review.
| # 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) |
There was a problem hiding this comment.
If the drain deadline expires with a future still running, this invokes the one-shot _finalize_session while that turn remains live; shutdown(wait=False) below does not join it. That retains the same persistence/lifecycle race for long turns. Please establish a terminal turn state before finalizing this session, or retain it unfinalized for recovery.
There was a problem hiding this comment.
You're right, and thanks — the reserve only relocated the race past the deadline rather than removing it. Fixed in 036195851, taking your second option (retain unfinalized for recovery).
The blocker was that _turn_futures was a bare set[Future], so on deadline expiry the drain could tell that something was still running but not whose — leaving flush_all_sessions no choice but to finalize everything. It's now a dict[Future, str] mapping each in-flight turn to its sid, captured at both submit sites (_handle_turn_start, _handle_spike_turn_start) from the sid already in scope. That's the same key space as server._sessions, which _ensure_server_session looks up by sid directly, so no new identity plumbing was needed.
On expiry, shutdown now computes the sids whose futures are still running and passes them to flush_all_sessions(skip_sids=...). Those sessions are left unfinalized — one-shot latch unspent, active_session_lease unreleased — so they stay recoverable. Sessions with no live turn finalize in the same pass exactly as before.
Unchanged: wait semantics, the min(_FLUSH_RESERVE_SECS, budget / 2.0) reserve, and the min(0.05, remaining) sleep bound. No added shutdown latency. The done-callback now pops under _turn_futures_lock rather than relying on set.discard as a bare callback, so entries still can't accumulate.
Tests in tests/tui_gateway/test_compute_host_phase1.py:
test_shutdown_retains_a_live_turns_session_when_the_drain_deadline_expires— replacestest_shutdown_still_finalizes_when_the_drain_deadline_expires, which asserted the behaviour you flagged. Two sessions, one with a turn outliving the window: the live one is retained, the idle one still flushes, and the call still returns insidewait.test_shutdown_retains_live_sessions_within_the_stdin_closed_budget— new, pinned to thewait=2.0stdin_closedpath, the tightest budget any caller uses and where the reserve has least room. Also asserts the drain still consumes its full window.test_shutdown_drains_in_flight_turn_before_finalizing_sessions— extended to guard the other direction: a turn that did drain must still finalize, so the skip can't over-reach into completed work. It also now asserts_turn_futuresempties, covering the container change.
Both new assertions were verified to fail against the previous flush behaviour and pass with the fix. flush_all_sessions has exactly one call site repo-wide, and all three shutdown callers (orphan, sigterm, stdin_closed) funnel through it, so this covers every path of the concern.
teknium1
added
sweeper:risk-message-delivery
…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.
|
Merged via #77330. Your commits were cherry-picked onto current Thank you for the excellent fix — the drain-before-finalize reordering, the |
This is a sibling follow-up to commit
1927b60771927b6077(fix(tui): extend deferred context-engine finalize to the compute-host compress routes) extended the deferred context-engine finalize acrosscompute_host.py'ssession.compress/slash.compressroutes.ComputeHost.shutdown()itself.What does this PR do?
Quit the TUI/dashboard while the agent is mid-answer and the tail of that answer — including tool results — is missing from the transcript on resume, and the memory summary is written from the truncated conversation.
ComputeHost.shutdown()finalizes every session before it drains in-flight turns:server._finalize_sessionis a one-shot latch — it returns early onsession["_finalized"]and sets it — so the flush gets exactly one chance to snapshot the session, andshutdown()spends that chance up to 10 seconds before the work it was supposed to capture exists. For any turn still running in the drain window that means:commit_memory_session(history)commits the mid-turn history, so long-term memory is written from a truncated transcript;on_session_endfirescompleted=False, interrupted=True(tui_gateway/server.py:709-710) against a session that is still running;_release_active_session_slotreleases the lease out from under a live turn;The invariant is stated in the repo against itself.
_finalize_session's own docstring: it "attempts to persist any unflushed messages before closing the session … prevents data loss when the TUI is force-quit … while the agent is mid-turn." The drain loop exists for exactly that reason; finalizing first defeats it. A turn that completes inside the window persists through its own normal path, so running the finalize afterwards lets the backstop snapshot a complete session instead of a partial one.All three teardown paths reach it:
_parent_guard_loop(dashboard/parent dies)shutdown(reason="orphan"), thenos._exit(0)immediately after — the drain is the last chance to save workshutdown(reason="sigterm")shutdown(reason="stdin_closed", wait=2.0)No new kill exposure, no added latency. Moving the flush behind the drain naively would land the durability write exactly on the supervisor's kill boundary:
host_supervisor._SHUTDOWN_TIMEOUT_SECSis10.0, the same value asshutdown()'s defaultwait, and_terminate_pidSIGKILLs the host that long after SIGTERM. So a slice of the existing budget (_FLUSH_RESERVE_SECS, capped at half of it so a short explicitwaitstill gets a real drain) is withheld from the drain, keepingdrain + flush <= wait.waitis not extended; the total shutdown window and the SIGTERM→SIGKILL margin are unchanged.Deliberately not double-flushing (once early, once late) as a hedge — the
_finalizedlatch means an early flush poisons the late one, which is precisely the bug.Sibling-site sweep
git grep -n "flush_all_sessions\|_finalize_session"over non-test sources, plus everydef shutdown/close/stopintui_gateway/:compute_host.py— flush before the drain inshutdown()compute_host.py—handle_frame,kind == "shutdown"compute_host.py—close()server.py—_shutdown_sessions()plugins/memory/openviking/__init__.py—_finalize_session_asyncRelated Issue
N/A — no filed issue; found by reading the compute-host teardown paths.
Type of Change
Changes Made
tui_gateway/compute_host.py—ComputeHost.shutdown()now runs the in-flight-turn drain loop beforeflush_all_sessions(), and withholds_FLUSH_RESERVE_SECS(new module constant,1.0, clamped to at most half the budget) from the drain deadline so the flush is still reached when the drain times out. Each drain tick is clamped to the time remaining, so it cannot overshoot the deadline and spend the reserve it exists to protect. Added a docstring recording the ordering invariant and why the reserve exists.tests/tui_gateway/test_compute_host_phase1.py— three regression tests.How to Test
Fails-before / passes-after, verified in both directions by reverting only the production hunk to
origin/mainand re-running:origin/maintest_shutdown_drains_in_flight_turn_before_finalizing_sessionsassert ['finalize:compute_host_sigterm', 'turn_end'] == ['turn_end', 'finalize:compute_host_sigterm'],At index 0 difftest_shutdown_still_finalizes_when_the_drain_deadline_expiresassert 1.030385416932404 < 1.0(the drain ate the whole budget)test_shutdown_drain_sleep_never_overshoots_the_reserveassert 0.2 <= (0.17 + 1e-06)/sum([0.05, 0.05, 0.05, 0.05])The second test is the anti-regression guard on the reserve: it holds a turn open past the window and asserts the finalize still runs and that
shutdown()returns inside the caller's budget. The third pins the per-tick sleep bound; it asserts on the summed requested sleep rather than wall-clock so it is deterministic — each tick is clamped to a strictly decreasing remainder, so the total cannot exceed the drain budget however the scheduler interleaves.Adjacent suites, all green with the change:
Checklist
Code
fix(scope):,feat(scope):, etc.)pytest tests/ -qand all tests passDocumentation & Housekeeping
docs/, docstrings) — or N/Acli-config.yaml.exampleif I added/changed config keys — or N/ACONTRIBUTING.mdorAGENTS.mdif I changed architecture or workflows — or N/A