From eb7ef3b06c7d8682603917a8bd7892552de0d681 Mon Sep 17 00:00:00 2001 From: darshthakkar Date: Mon, 20 Jul 2026 17:09:25 -0400 Subject: [PATCH] fix(tui-gateway): release session leases at idle boundary Schedule precise per-session idle lease timers after turn completion while keeping warm runtimes and history intact. Cancel stale timers on new turns and teardown, preserve active, queued, pending, and building work, and retain the periodic sweep as a recovery backstop. Co-Authored-By: Claude --- tests/tui_gateway/test_protocol.py | 188 +++++++++++++++++++++++++++++ tui_gateway/server.py | 152 +++++++++++++++++++++-- 2 files changed, 328 insertions(+), 12 deletions(-) diff --git a/tests/tui_gateway/test_protocol.py b/tests/tui_gateway/test_protocol.py index 39c4bf40ac5d..d91e7f070183 100644 --- a/tests/tui_gateway/test_protocol.py +++ b/tests/tui_gateway/test_protocol.py @@ -1183,6 +1183,194 @@ def test_lease_idle_release_seconds_defaults_and_coerces(server, monkeypatch): assert server._lease_idle_release_seconds() == 45.0 +class _FakeLeaseTimer: + created = [] + + def __init__(self, interval, function): + self.interval = interval + self.function = function + self.cancelled = False + self.started = False + self.daemon = False + self.__class__.created.append(self) + + def start(self): + self.started = True + + def cancel(self): + self.cancelled = True + + def fire(self): + self.function() + + +class _CountingLease: + def __init__(self): + self.release_count = 0 + + def release(self): + self.release_count += 1 + + +def test_schedule_idle_lease_release_uses_exact_delay_and_keeps_session_live( + server, monkeypatch +): + _FakeLeaseTimer.created = [] + monkeypatch.setattr(server, "_LeaseTimer", _FakeLeaseTimer) + monkeypatch.setattr(server, "_load_cfg", lambda: {"tui_lease_idle_seconds": 7}) + lease = _CountingLease() + session = _lease_test_session(active_session_lease=lease) + server._sessions["ui-idle"] = session + + server._schedule_idle_lease_release("ui-idle", session) + + timer = _FakeLeaseTimer.created[-1] + assert timer.interval == 7 + assert timer.started is True + assert timer.daemon is True + timer.fire() + + assert lease.release_count == 1 + assert session.get("active_session_lease") is None + assert session.get("_lease_idle_timer") is None + assert server._sessions["ui-idle"] is session + + +@pytest.mark.parametrize("protected", ["running", "queued", "pending", "building"]) +def test_idle_lease_timer_rechecks_protected_work_before_release( + server, monkeypatch, protected +): + _FakeLeaseTimer.created = [] + monkeypatch.setattr(server, "_LeaseTimer", _FakeLeaseTimer) + monkeypatch.setattr(server, "_load_cfg", lambda: {"tui_lease_idle_seconds": 7}) + lease = _CountingLease() + session = _lease_test_session(active_session_lease=lease) + server._sessions["ui-protected"] = session + server._schedule_idle_lease_release("ui-protected", session) + timer = _FakeLeaseTimer.created[-1] + + if protected == "running": + session["running"] = True + elif protected == "queued": + session["queued_prompt"] = {"text": "next"} + elif protected == "pending": + server._pending["pending-idle-release"] = ("ui-protected", None) + else: + session["agent_ready"] = threading.Event() + session["lazy"] = False + + try: + timer.fire() + assert lease.release_count == 0 + assert session.get("active_session_lease") is lease + assert session.get("_lease_idle_timer") is None + finally: + server._pending.pop("pending-idle-release", None) + + +def test_turn_start_invalidates_older_idle_lease_timer(server, monkeypatch): + _FakeLeaseTimer.created = [] + monkeypatch.setattr(server, "_LeaseTimer", _FakeLeaseTimer) + monkeypatch.setattr(server, "_load_cfg", lambda: {"tui_lease_idle_seconds": 7}) + lease = _CountingLease() + session = _lease_test_session(active_session_lease=lease) + + server._schedule_idle_lease_release("ui-reuse", session) + timer = _FakeLeaseTimer.created[-1] + assert server._ensure_turn_lease("ui-reuse", session) is None + + assert timer.cancelled is True + assert session.get("_lease_idle_timer") is None + timer.fire() + assert lease.release_count == 0 + assert session.get("active_session_lease") is lease + + +def test_turn_start_waits_for_atomic_idle_registry_release(server, monkeypatch): + _FakeLeaseTimer.created = [] + monkeypatch.setattr(server, "_LeaseTimer", _FakeLeaseTimer) + monkeypatch.setattr(server, "_load_cfg", lambda: {"tui_lease_idle_seconds": 7}) + release_entered = threading.Event() + allow_release = threading.Event() + release_done = threading.Event() + + class BlockingLease: + def release(self): + release_entered.set() + assert allow_release.wait(timeout=2) + release_done.set() + + replacement = _CountingLease() + + def claim(*_args, **_kwargs): + if not release_done.is_set(): + return None, "old registry lease still occupies the only slot" + return replacement, None + + monkeypatch.setattr(server, "_claim_active_session_slot", claim) + session = _lease_test_session(active_session_lease=BlockingLease()) + server._schedule_idle_lease_release("ui-race", session) + timer = _FakeLeaseTimer.created[-1] + + release_thread = threading.Thread(target=timer.fire) + release_thread.start() + assert release_entered.wait(timeout=2) + + result = {} + turn_thread = threading.Thread( + target=lambda: result.setdefault("limit", server._ensure_turn_lease("ui-race", session)) + ) + turn_thread.start() + turn_thread.join(timeout=0.05) + assert turn_thread.is_alive(), "turn admission raced ahead of registry release" + + allow_release.set() + release_thread.join(timeout=2) + turn_thread.join(timeout=2) + assert result["limit"] is None + assert session["active_session_lease"] is replacement + + +def test_active_slot_release_cancels_idle_lease_timer(server, monkeypatch): + _FakeLeaseTimer.created = [] + monkeypatch.setattr(server, "_LeaseTimer", _FakeLeaseTimer) + monkeypatch.setattr(server, "_load_cfg", lambda: {"tui_lease_idle_seconds": 7}) + lease = _CountingLease() + session = _lease_test_session(active_session_lease=lease) + + server._schedule_idle_lease_release("ui-close", session) + timer = _FakeLeaseTimer.created[-1] + server._release_active_session_slot(session) + + assert timer.cancelled is True + assert lease.release_count == 1 + timer.fire() + assert lease.release_count == 1 + + +def test_queued_dispatch_failure_still_schedules_idle_lease_release(server, monkeypatch): + session = _lease_test_session( + active_session_lease=_CountingLease(), + queued_prompt={"text": "next", "transport": None}, + ) + scheduled = [] + monkeypatch.setattr( + server, + "_run_prompt_submit", + lambda *_args: (_ for _ in ()).throw(RuntimeError("dispatch failed")), + ) + monkeypatch.setattr( + server, + "_schedule_idle_lease_release", + lambda sid, current: scheduled.append((sid, current)), + ) + + assert server._drain_queued_prompt("rid", "ui-queued-failure", session) is True + assert session["running"] is False + assert session["queued_prompt"] is None + assert scheduled == [("ui-queued-failure", session)] + + def test_session_resume_live_payload_uses_current_history_with_ancestors(server, monkeypatch): """Live resume should not reuse a stale ancestor-inclusive snapshot.""" diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 77083b74873f..3cf584da8084 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -37,6 +37,29 @@ ) logger = logging.getLogger(__name__) +_LeaseThread = threading.Thread + + +class _LeaseTimer: + """Small cancellable timer insulated from runtime Thread monkeypatches.""" + + def __init__(self, interval: float, function) -> None: + self.interval = interval + self.function = function + self.daemon = False + self._cancelled = threading.Event() + self._thread = None + + def _run(self) -> None: + if not self._cancelled.wait(self.interval): + self.function() + + def start(self) -> None: + self._thread = _LeaseThread(target=self._run, daemon=self.daemon) + self._thread.start() + + def cancel(self) -> None: + self._cancelled.set() _hermes_home = get_hermes_home() load_hermes_dotenv( @@ -443,6 +466,12 @@ def _claim_active_session_slot( def _release_active_session_slot(session: dict | None) -> None: if not session: return + timer = session.pop("_lease_idle_timer", None) + if timer is not None: + try: + timer.cancel() + except Exception: + logger.debug("Failed to cancel idle lease timer", exc_info=True) lease = session.pop("active_session_lease", None) if lease is None: return @@ -458,16 +487,26 @@ def _ensure_turn_lease(sid: str, session: dict) -> str | None: TUI/desktop sessions acquire their registry lease lazily on the first turn (not at session.create/resume, which fire for every composer paint - and sidebar switch) and keep it across turns; the idle reaper hands the + and sidebar switch) and keep it across turns; the idle timer hands the slot back after the configured idle window without conversational - activity (see ``_release_idle_session_leases``). This is the re-acquire - half: mirrors the platform gateway's claim-per-turn in + activity (see ``_schedule_idle_lease_release``; the periodic reaper is a + recovery backstop). This is the re-acquire half: mirrors the platform + gateway's claim-per-turn in ``gateway/run.py`` handle_message. Returns the limit message when the cap is reached (the caller surfaces it as the turn's error), None on success or fail-open. """ with session["history_lock"]: + # A new turn owns the current lease. Invalidate any older idle timer + # under the same lock its callback uses so a late callback cannot pull + # the slot out from under this turn. + timer = session.pop("_lease_idle_timer", None) + if timer is not None: + try: + timer.cancel() + except Exception: + logger.debug("Failed to cancel idle lease timer", exc_info=True) if session.get("active_session_lease") is not None: return None lease, limit_message = _claim_active_session_slot( @@ -943,11 +982,94 @@ def _lease_idle_release_seconds() -> float: return _LEASE_IDLE_RELEASE_DEFAULT_S +def _idle_lease_release_blocked(sid: str, session: dict) -> bool: + """Return whether work still owns this session's active-turn lease. + + Callers hold ``history_lock`` so turn start, queue drain, and timer expiry + share one atomic decision boundary. + """ + if session.get("_finalized") or session.get("running") or session.get("queued_prompt"): + return True + if _session_pending_kind(sid): + return True + ready = session.get("agent_ready") + return bool(ready is not None and not ready.is_set() and not session.get("lazy")) + + +def _release_idle_active_session_lease_locked(session: dict, *, reason: str) -> bool: + """Release registry ownership before exposing an empty session lease slot. + + The caller holds history_lock. Keeping that lock across registry release + prevents a new turn from observing ``active_session_lease=None`` while its + old registry entry still occupies the cap. + """ + lease = session.get("active_session_lease") + if lease is None: + return False + try: + lease.release() + except Exception: + logger.debug("Failed to release %s active session lease", reason, exc_info=True) + finally: + if session.get("active_session_lease") is lease: + session.pop("active_session_lease", None) + return True + + +def _schedule_idle_lease_release(sid: str, session: dict) -> None: + """Release this turn lease at the configured idle boundary. + + The periodic reaper remains a recovery sweep. This timer makes the normal + path precise and re-checks every protected state at expiry, so queued or + chained work keeps the lease without closing the warm runtime. + """ + delay = _lease_idle_release_seconds() + if delay <= 0: + return + lock = session.get("history_lock") + if lock is None: + return + + timer = None + + def release_if_still_idle() -> None: + with lock: + if session.get("_lease_idle_timer") is not timer: + return + session.pop("_lease_idle_timer", None) + if _idle_lease_release_blocked(sid, session): + return + _release_idle_active_session_lease_locked(session, reason="scheduled idle") + + with lock: + if session.get("active_session_lease") is None or _idle_lease_release_blocked( + sid, session + ): + return + previous = session.pop("_lease_idle_timer", None) + if previous is not None: + try: + previous.cancel() + except Exception: + logger.debug("Failed to cancel replaced idle lease timer", exc_info=True) + timer = _LeaseTimer(delay, release_if_still_idle) + timer.daemon = True + session["_lease_idle_timer"] = timer + try: + timer.start() + except Exception: + with lock: + if session.get("_lease_idle_timer") is timer: + session.pop("_lease_idle_timer", None) + logger.debug("Failed to start idle lease timer", exc_info=True) + + def _release_idle_session_leases(now: float) -> None: """Release the active-session LEASE (not the session) for idle sessions. - Runs on the reaper cadence. Applies to live-transport sessions too — that - is the point: connected-but-idle tabs are exactly the ones the TTL reaper + Recovery sweep on the reaper cadence; the normal path uses an exact timer + scheduled at the turn-chain idle boundary. Applies to live-transport + sessions too — connected-but-idle tabs are exactly the ones the TTL reaper can never free. Skips anything mid-turn, awaiting an input/approval prompt, holding a queued next-turn prompt, or still building its agent (imminent turn). The `running` check happens under history_lock, the same @@ -978,13 +1100,13 @@ def _release_idle_session_leases(now: float) -> None: reference = max(last_active, created_at) if (now - reference) <= idle_release_s: continue - lease = session.pop("active_session_lease", None) - if lease is None: - continue - try: - lease.release() - except Exception: - logger.debug("Failed to release idle active session lease", exc_info=True) + timer = session.pop("_lease_idle_timer", None) + if timer is not None: + try: + timer.cancel() + except Exception: + logger.debug("Failed to cancel swept idle lease timer", exc_info=True) + _release_idle_active_session_lease_locked(session, reason="swept idle") def _reap_idle_sessions() -> None: @@ -5289,6 +5411,7 @@ def _drain_queued_prompt(rid, sid: str, session: dict) -> bool: ) with session["history_lock"]: session["running"] = False + _schedule_idle_lease_release(sid, session) return True @@ -9545,6 +9668,11 @@ def _stream(delta): file=sys.stderr, ) + # This invocation reached the end of its follow-up chain. If another + # queued/user/goal/notification turn started meanwhile, the helper sees + # running=True and leaves the lease to that turn. + _schedule_idle_lease_release(sid, session) + run_thread = threading.Thread(target=run, daemon=True) session["_run_thread"] = run_thread run_thread.start()