fix(tui_gateway): spare live-turn sessions from the shutdown finalize sweep

[briandevans]

What does this PR do?

This fix is specified verbatim by our own merged code. tui_gateway/compute_host.py:219-222 on main — merged in #77330 — ends the HostRuntime.shutdown docstring with:

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. […] A follow-up could gate _shutdown_sessions on not session.get("_finalized") and not session.get("running") to close the gap.

This PR is that follow-up, with the exact predicate the NOTE names, plus the NOTE rewritten to describe the closed gap instead of the open one.

The gap

tui_gateway/server.py:1179-1187 on main:

def _shutdown_sessions() -> None:
    try:
        _release_gateway_wake_owner()
    except Exception:
        pass
    with _sessions_lock:
        sids = list(_sessions)
    for sid in sids:
        _close_session_by_id(sid, end_reason="tui_shutdown")

No gate on turn state. _finalize_session (server.py:729-740) is a one-shot latch — it sets session["_finalized"] = True and every later call returns immediately. Spending that latch while a turn is running runs each step against a live session:

step	server.py	effect mid-turn
agent._persist_session(snapshot)	:765	persists a truncated transcript
commit_memory_session(history)	:795	commits long-term memory from that truncated history
db.end_session(sid, end_reason)	:827	ends a live durable row
invoke_hook("on_session_end", …)	:783	fires completed=False, interrupted=True on a running turn
_release_active_session_slot(session)	:741	drops the lease out from under live work
interrupt_for_session(...)	:851	kills in-flight async delegations

Nothing joins the turn thread, so the tail it goes on to produce is thereafter unpersistable — the session is permanently un-finalizable.

Why it fires for everyone

atexit.register(_shutdown_sessions) (server.py:1368) runs on stdin EOF and on any sys.exit(0), and tui_gateway/entry.py:155-160 calls it directly from _log_signal on SIGTERM / SIGHUP / SIGBREAK, inside the grace-window hard-exit timer. No config key gates it. So quitting the TUI, closing the desktop app, or stopping the dashboard backend while the assistant is mid-reply loses that reply and any tool tail, on the default configuration — and MEMORY has already been committed from the truncated transcript.

The in-file contrast

Every other automatic reclaim path in server.py already refuses to reclaim mid-turn. The one that runs on every exit did not:

path	line	guard
_ws_session_is_orphaned	:1021	if session.get("running"): return False
_session_is_evictable	:1212	if session.get("running") or _session_pending_kind(sid)
_session_is_lru_evictable	:1300	"never evict a session mid-turn" (:1301)
_shutdown_sessions	:1187	none

The fix

_close_session_by_id already accepts predicate= and re-validates it under _sessions_lock immediately before the ownership claim (server.py:983-1009), exactly as the idle and LRU reapers use it. So the change adds no new locking and no new teardown path:

for sid in sids:
    _close_session_by_id(
        sid,
        end_reason="tui_shutdown",
        predicate=_shutdown_session_is_reclaimable,
    )

Leaving a live-turn session unfinalized is this repo's established answer, not a new judgement call: eb4f514b2d3 ("retain live-turn sessions unfinalized when the drain deadline expires") makes the same trade one layer up. Unfinalized stays recoverable via the turn-end flush; latched mid-turn is permanent.

The _finalized half of the predicate covers the sibling case: compute_host.flush_all_sessions finalizes sessions in place without popping them from server._sessions, so a session the drain already flushed would otherwise be torn down a second time at exit — re-announcing a reclaim and calling agent.close() on a dying process, with a finalize that is already a no-op.

Related Issue

No filed issue — this is the follow-up named in the merged compute_host.shutdown NOTE (#77330).

Related: eb4f514b2d3 (drain-deadline live-turn skip), #77330 (the NOTE this closes).

Type of Change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 🔒 Security fix
• 📝 Documentation update
• ✅ Tests (adding or improving test coverage)
• ♻️ Refactor (no behavior change)
• 🎯 New skill (bundled or hub)

Changes Made

• tui_gateway/server.py — add _shutdown_session_is_reclaimable(session) and pass it as _shutdown_sessions' predicate=, so the atexit sweep skips sessions that are mid-turn or already finalized.
• tui_gateway/compute_host.py — rewrite the now-stale HostRuntime.shutdown NOTE. It described the atexit re-finalize as an open accepted gap and proposed this predicate; it now records the contract that exists.
• tests/test_tui_gateway_server.py — two regression tests (below), plus the predicate=None keyword on the existing test_shutdown_sessions_closes_every_session_via_helper stub so it keeps asserting that every eligible session is routed through the helper.

Sibling sweep — all 6 reclaim sites, and the one deliberately excluded

Enumerated _close_session_by_id( / _teardown_session( / _teardown_popped_session( / _finalize_session( across tui_gateway/:

1. _reap_idle_sessions (:1233) — already gates via _session_is_evictable ✓
2. _enforce_session_cap (:1333) — already gates via _session_is_lru_evictable ✓
3. _schedule_ws_orphan_reap — already gates via _ws_session_is_orphaned ✓
4. _shutdown_sessions (:1187) — fixed here
5. the close_on_disconnect branch of _close_sessions_for_transport (:1163) — same missing gate, DELIBERATELY EXCLUDED. Three open PRs already have hunks inside _close_sessions_for_transport, so touching it here would guarantee a conflict with work in flight: fix(tui-gateway): close WS disconnect/reconnect session race #77129 (@JoaoMarcos44) rewrites that for-loop body outright (@@ -1077,23 +1076,48 @@, including the close_on_disconnect branch itself), and fix(gateway): re-bind session transport to a surviving window on pop-out close #86039 (@ayushnangia) and fix(tui_gateway): close slash_worker on WS detach to prevent memory leak #57687 (@yingliang-zhang) each add hunks inside the same function. Flagging it rather than silently omitting it: it is a real sibling and it wants the same gate.
6. the session.close RPC — user-initiated, not an automatic reclaim. Out of scope by design.

Related open PRs, disclosed up front

• fix(tui-gateway): finalize sessions on clean pipe exit #60302 (@ruanbarroso, "finalize sessions on clean pipe exit", open, cold since 07-24) — its test asserts the behaviour this PR changes. Its production diff touches tui_gateway/entry.py::_log_exit only (adds a server._shutdown_sessions() call), so there is no textual conflict with this diff. But test_shutdown_sessions_persists_tool_tail_and_marks_tui_shutdown builds a running=True session and asserts _shutdown_sessions ends it as tui_shutdown. Reconciliation: that PR's goal is to stop an interrupted turn being left for a later WS-orphan reap, and eb4f514b2d3 has since established unfinalized-but-recoverable as the desired state for a live turn — the two goals meet by keeping the session registered and recoverable, not by latching it. If fix(tui-gateway): finalize sessions on clean pipe exit #60302 lands first, that assertion is the thing to update.
• fix(tui-gateway): retry and log unflushed-message persist failures before teardown discards them #79109 (@sakshamzip2-sys) — edits _finalize_session's body (persist retry). Different function, but adjacent enough that a textual conflict is possible.
• fix(tui): finalize closed sessions in background #25027 (@jamesdag, May) — rewrites _shutdown_sessions' loop body to background cleanup for session.close. Different concern (blocking /new), and its base predates the _close_session_by_id funnel entirely, so it conflicts with main independently of this PR.
• fix(tui): reschedule WS orphan reap while a turn is still running #85598 (@686f6c61) and fix(tui-gateway): re-arm WS orphan reap while a detached session is mid-turn #44102 (@AIalliAI) — apply the same never-reclaim-mid-turn principle to the sibling _schedule_ws_orphan_reap path. Corroborating, not colliding.
• fix(desktop): keep and recover interrupted chats across update/restart #81319 — closest by user-visible harm ("keep and recover interrupted chats across update/restart"), but a complementary mechanism (turn-marker preservation + boot recovery); it does not prevent the latch from being spent.

How to Test

1. Reproduce: start a TUI/desktop session, send a prompt, and quit (or SIGTERM the gateway) while the assistant is still streaming. On main the session's durable row is already ended, the reply tail is missing on reopen, and memory has been committed from the truncated transcript.
2. Regression tests — both fail on main, pass with this change:

uv run --with pytest --with pytest-asyncio python3 -m pytest \
  tests/test_tui_gateway_server.py -k shutdown_sessions -v

test_shutdown_sessions_spares_a_session_whose_turn_is_running — registers one running=True and one running=False session, stubs _teardown_session so the real _close_session_by_id predicate path runs. Asserts the idle session is still reclaimed as tui_shutdown and popped, the live one is still in _sessions, and its _finalized latch is unspent.
Red on main: assert [('live', 'tui_shutdown'), ('idle', 'tui_shutdown')] == [('idle', 'tui_shutdown')].

test_shutdown_sessions_skips_a_session_the_drain_already_finalized — registers a session with _finalized already set (the state compute_host.flush_all_sessions leaves behind) and asserts the sweep does not tear it down again.
Red on main: assert ['flushed'] == [].

Both directions were verified by restoring tui_gateway/server.py from origin/main and re-running each test individually.

3. Adjacent suites, all green:

uv run --with pytest --with pytest-xdist --with pytest-asyncio python3 -m pytest \
  tests/test_tui_gateway_server.py tests/tui_gateway/ tests/test_tui_gateway_ws.py \
  tests/test_tui_gateway_queue_on_busy.py tests/test_tui_gateway_server_crash_history.py -q
# 1060 passed, 1 skipped

Checklist

Code

• I've read the Contributing Guide
• My commit messages follow Conventional Commits (fix(scope):, feat(scope):, etc.)
• I searched for existing PRs to make sure this isn't a duplicate
• My PR contains only changes related to this fix/feature (no unrelated commits)
• I've run pytest tests/ -q and all tests pass — ran the tui_gateway suites listed above (1060 passed, 1 skipped), not the full tree
• I've added tests for my changes (required for bug fixes, strongly encouraged for features)
• I've tested on my platform: macOS 15 (Darwin 25.4)

Documentation & Housekeeping

• I've updated relevant documentation (README, docs/, docstrings) — the compute_host.shutdown NOTE and the new _shutdown_session_is_reclaimable docstring
• N/A — no config keys added or changed
• N/A — no architecture or workflow change
• I've considered cross-platform impact (Windows, macOS) — pure-Python dict predicate, no platform branches; the SIGHUP/SIGBREAK entry path is unchanged
• N/A — no tool behaviour changed

[Copilot]

Pull request overview

This PR fixes a shutdown-time data-loss bug in the TUI gateway by ensuring the atexit session sweep does not finalize/tear down sessions that are mid-turn (or already finalized). This aligns the atexit path with the existing “never reclaim mid-turn” policy already enforced by other reapers in tui_gateway/server.py, preventing the one-shot _finalize_session latch from being spent while a turn is still producing output.

Changes:

• Add _shutdown_session_is_reclaimable() and pass it as predicate= to _close_session_by_id() inside _shutdown_sessions(), so the atexit sweep skips running and _finalized sessions.
• Update HostRuntime.shutdown()’s docstring NOTE to reflect the closed gap and the new predicate-based guarantee.
• Add regression tests proving _shutdown_sessions() skips (a) running sessions and (b) sessions already finalized by the drain flush.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
tui_gateway/server.py	Gate the atexit shutdown sweep via predicate= to avoid mid-turn (or already-finalized) teardown.
tui_gateway/compute_host.py	Update shutdown NOTE to document the new atexit behavior contract and closed gap.
tests/test_tui_gateway_server.py	Add regression coverage for “skip running” and “skip already-finalized” during shutdown sweep.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[686f6c61]

#85598 is the WS-orphan sibling only (_schedule_ws_orphan_reap / mid-turn re-arm). It does not touch _shutdown_sessions or the atexit predicate. Agree this is corroborating, not colliding — happy to rebase #85598 if anything in the shared predicate= path moves when this lands.

[briandevans]

Thanks for checking, and confirming from this side: your read is right. This PR touches tui_gateway/server.py, tui_gateway/compute_host.py and tests/test_tui_gateway_server.py, and there is no reference to _schedule_ws_orphan_reap or the mid-turn re-arm anywhere in the diff.

On the shared predicate= path specifically — it does not move, so I don't think #85598 needs a rebase on this PR's account:

_close_session_by_id's predicate: Callable[[dict], bool] | None = None parameter already exists on main (tui_gateway/server.py:987), and this diff does not modify that signature. All it does is add a new caller of the existing keyword: _shutdown_sessions now passes predicate=_shutdown_session_is_reclaimable instead of calling with end_reason alone. Anything already relying on the parameter keeps working unchanged.

Your _ws_session_is_orphaned work is actually part of the argument here rather than just adjacent to it. The case for gating shutdown is that every other automatic reclaim path in the module already gates on running — _ws_session_is_orphaned, _session_is_evictable, _session_is_lru_evictable — and _shutdown_sessions was the one that ran on every exit without doing so. So the sibling paths being consistent is what makes the shutdown path's omission legible as a gap rather than a deliberate choice.

If the signature does end up moving I'll flag it here before pushing. Likewise, if you'd rather land #85598 first, this side rebases cleanly either way.

[Enough1122]

AI code review — automated review for reference, author can ignore or act on any point.

fix(tui_gateway): spare live-turn sessions from the shutdown finalize sweep

The predicate is exactly the gap the old comment flagged ("A follow-up could gate _shutdown_sessions on not _finalized and not running"), and the docstrings correctly lay out the one-shot-latch reasoning and the trade-off for live-turn sessions. Tests are clear and deterministic. Minor observations:

1. tui_gateway/server.py — _shutdown_session_is_reclaimable leaves a running session in _sessions at atexit time. The docstring says the turn-end flush will persist it, but on the SIGTERM/stdin_closed paths the process exits right after the sweep — if the turn thread is still mid-flight, the flush may never run and the session's tail is lost. This is a deliberate trade (finalizing mid-turn persists a truncated transcript + commits memory from it, which is worse), and compute_host.shutdown(wait=...) gives the drain some time — worth one sentence noting the residual data-loss window on hard exits so it isn't rediscovered as a bug.

2. tui_gateway/server.py:1184-1190 — sessions the drain already finalized (_finalized=True) are skipped and left registered in _sessions. On a normal exit that's harmless, but it means _shutdown_sessions no longer empties the registry; any future code that asserts _sessions is empty after shutdown would break. The existing tests already server._sessions.clear() in finally, which is good hygiene.

3. Minor: _shutdown_session_is_reclaimable(session) gates on plain session.get(...), but _close_session_by_id presumably handles missing keys already. The predicate silently treats a session missing _finalized/running keys as reclaimable — consistent with the pre-change behavior, so no regression; just noting the implicit default.

4. Test note: test_shutdown_sessions_spares_a_session_whose_turn_is_running relies on dict insertion order for torn == [("idle", "tui_shutdown")] — deterministic in CPython, fine as written.