diff --git a/contributors/emails/theunathi@gmail.com b/contributors/emails/theunathi@gmail.com new file mode 100644 index 00000000000..da16c209d44 --- /dev/null +++ b/contributors/emails/theunathi@gmail.com @@ -0,0 +1,2 @@ +UnathiCodex +# PR contribution (tui_gateway: keep busy submits resume-safe) diff --git a/tests/test_tui_gateway_queue_on_busy.py b/tests/test_tui_gateway_queue_on_busy.py index e1c5050295d..804a3b5505c 100644 --- a/tests/test_tui_gateway_queue_on_busy.py +++ b/tests/test_tui_gateway_queue_on_busy.py @@ -10,6 +10,7 @@ next turn, drained in ``run``'s tail. """ import threading +import time import types from tui_gateway import server @@ -52,11 +53,14 @@ def test_busy_interrupt_mode_interrupts_and_queues(monkeypatch): monkeypatch.setattr(server, "_load_busy_input_mode", lambda: "interrupt") calls = {"interrupt": 0} agent = types.SimpleNamespace(interrupt=lambda *a, **k: calls.__setitem__("interrupt", calls["interrupt"] + 1)) - session = _session(agent=agent) + session = _session(agent=agent, running=True) resp = server._handle_busy_submit("r1", "sid", session, "redirect", "ws-1") assert resp["result"]["status"] == "queued" + deadline = time.monotonic() + 1 + while calls["interrupt"] != 1 and time.monotonic() < deadline: + time.sleep(0.01) assert calls["interrupt"] == 1 assert session["queued_prompt"]["text"] == "redirect" @@ -65,7 +69,7 @@ def test_busy_queue_mode_queues_without_interrupting(monkeypatch): monkeypatch.setattr(server, "_load_busy_input_mode", lambda: "queue") calls = {"interrupt": 0} agent = types.SimpleNamespace(interrupt=lambda *a, **k: calls.__setitem__("interrupt", calls["interrupt"] + 1)) - session = _session(agent=agent) + session = _session(agent=agent, running=True) resp = server._handle_busy_submit("r1", "sid", session, "later", "ws-1") @@ -77,7 +81,7 @@ def test_busy_queue_mode_queues_without_interrupting(monkeypatch): def test_busy_steer_mode_injects_when_accepted(monkeypatch): monkeypatch.setattr(server, "_load_busy_input_mode", lambda: "steer") agent = types.SimpleNamespace(steer=lambda text: True, interrupt=lambda *a, **k: None) - session = _session(agent=agent) + session = _session(agent=agent, running=True) resp = server._handle_busy_submit("r1", "sid", session, "nudge", "ws-1") @@ -88,7 +92,7 @@ def test_busy_steer_mode_injects_when_accepted(monkeypatch): def test_busy_steer_mode_falls_back_to_queue_when_rejected(monkeypatch): monkeypatch.setattr(server, "_load_busy_input_mode", lambda: "steer") agent = types.SimpleNamespace(steer=lambda text: False, interrupt=lambda *a, **k: None) - session = _session(agent=agent) + session = _session(agent=agent, running=True) resp = server._handle_busy_submit("r1", "sid", session, "nudge", "ws-1") @@ -96,6 +100,40 @@ def test_busy_steer_mode_falls_back_to_queue_when_rejected(monkeypatch): assert session["queued_prompt"]["text"] == "nudge" +def test_busy_interrupt_does_not_hold_history_lock_or_delay_queue(monkeypatch): + monkeypatch.setattr(server, "_load_busy_input_mode", lambda: "interrupt") + interrupt_started = threading.Event() + release_interrupt = threading.Event() + + def blocking_interrupt(): + interrupt_started.set() + release_interrupt.wait(timeout=2) + + session = _session( + agent=types.SimpleNamespace(interrupt=blocking_interrupt), + running=True, + ) + + started = time.monotonic() + resp = server._handle_busy_submit("r1", "sid", session, "keep this", "ws-1") + + assert resp["result"]["status"] == "queued" + assert time.monotonic() - started < 0.25 + assert session["queued_prompt"]["text"] == "keep this" + assert interrupt_started.wait(timeout=1) + assert session["history_lock"].acquire(timeout=0.25) + session["history_lock"].release() + release_interrupt.set() + + +def test_busy_helper_retries_when_turn_finished(monkeypatch): + monkeypatch.setattr(server, "_load_busy_input_mode", lambda: "interrupt") + session = _session(running=False) + + assert server._handle_busy_submit("r1", "sid", session, "run now", "ws-1") is None + assert session.get("queued_prompt") is None + + # ── _drain_queued_prompt ─────────────────────────────────────────────────── def test_drain_fires_queued_prompt_and_claims_running(monkeypatch): diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index a23804d5daa..30b6bcbfb0d 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -2278,6 +2278,50 @@ def test_session_close_commits_memory_and_fires_finalize_hook(monkeypatch): server._sessions.pop("sid", None) +def test_session_close_releases_resume_lock_before_slow_teardown(monkeypatch): + """One slow session finalizer must not stall unrelated session.resume RPCs.""" + teardown_started = threading.Event() + release_teardown = threading.Event() + response = {} + + def _slow_teardown(_session, *, end_reason="tui_close"): + assert end_reason == "tui_close" + teardown_started.set() + assert release_teardown.wait(timeout=2.0) + + monkeypatch.setattr(server, "_teardown_session", _slow_teardown) + server._sessions["slow-close"] = _session() + + def _close(): + response.update( + server.handle_request( + { + "id": "close", + "method": "session.close", + "params": {"session_id": "slow-close"}, + } + ) + ) + + thread = threading.Thread(target=_close) + thread.start() + acquired = False + try: + assert teardown_started.wait(timeout=1.0) + assert "slow-close" not in server._sessions + acquired = server._session_resume_lock.acquire(timeout=0.2) + assert acquired, "slow teardown kept the global resume lock held" + finally: + if acquired: + server._session_resume_lock.release() + release_teardown.set() + thread.join(timeout=2.0) + server._sessions.pop("slow-close", None) + + assert not thread.is_alive() + assert response["result"] == {"closed": True} + + def test_ws_orphan_reap_closes_worker_when_session_stays_detached(monkeypatch): """A detached WS session past its grace window has its slash_worker closed. @@ -2308,6 +2352,51 @@ def test_ws_orphan_reap_closes_worker_when_session_stays_detached(monkeypatch): server._sessions.pop("orphan-sid", None) +def test_ws_orphan_reap_releases_resume_lock_before_slow_teardown(monkeypatch): + """Grace reaping claims under the lock but finalizes after releasing it.""" + scheduled = {} + teardown_started = threading.Event() + release_teardown = threading.Event() + + class _Timer: + def __init__(self, _delay, callback): + scheduled["callback"] = callback + + def start(self): + return None + + def _slow_teardown(_session, *, end_reason="tui_close"): + assert end_reason == "ws_orphan_reap" + teardown_started.set() + assert release_teardown.wait(timeout=2.0) + + monkeypatch.setattr(server, "_WS_ORPHAN_REAP_GRACE_S", 0.01) + monkeypatch.setattr(server.threading, "Timer", _Timer) + monkeypatch.setattr(server, "_teardown_session", _slow_teardown) + server._sessions["slow-orphan"] = _session( + transport=server._detached_ws_transport, + running=False, + ) + + server._schedule_ws_orphan_reap("slow-orphan") + thread = threading.Thread(target=scheduled["callback"]) + thread.start() + acquired = False + try: + assert teardown_started.wait(timeout=1.0) + assert "slow-orphan" not in server._sessions + acquired = server._session_resume_lock.acquire(timeout=0.2) + assert acquired, "orphan teardown kept the global resume lock held" + finally: + if acquired: + server._session_resume_lock.release() + release_teardown.set() + thread.join(timeout=2.0) + server._sessions.pop("slow-orphan", None) + + assert not thread.is_alive() + + def test_finalize_session_closes_slash_worker(monkeypatch): """_finalize_session closes the slash_worker subprocess itself. @@ -9773,17 +9862,20 @@ def test_restart_slash_worker_stores_on_live_session(monkeypatch): server._sessions.pop("live-restart", None) -def test_session_close_rpc_delegates_to_close_session_by_id(monkeypatch): +def test_session_close_rpc_claims_then_tears_down(monkeypatch): seen = [] + claimed = {"session_key": "k"} + monkeypatch.setattr(server, "_pop_session_by_id", lambda sid: seen.append(sid) or claimed) monkeypatch.setattr( - server, "_close_session_by_id", - lambda sid, *, end_reason: bool(seen.append((sid, end_reason))) or True, + server, + "_teardown_popped_session", + lambda session, *, end_reason: seen.append((session, end_reason)) or True, ) resp = server.handle_request( {"id": "1", "method": "session.close", "params": {"session_id": "s9"}} ) assert resp["result"] == {"closed": True} - assert seen == [("s9", "tui_close")] + assert seen == ["s9", (claimed, "tui_close")] def test_close_sessions_for_transport_closes_flagged_repoints_rest(monkeypatch): diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 280469f2ff5..a79c06947da 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -749,24 +749,49 @@ def _attach_worker(sid: str, session: dict, worker) -> None: worker.close() -def _close_session_by_id(sid: str, *, end_reason: str = "tui_close") -> bool: - """Single idempotent teardown for one session: pop it under the sessions - lock, then finalize, unregister notify, close agent + slash worker via the - shared ``_teardown_session`` path. Returns True iff it closed a live - session. The ``_finalized`` / worker ``_closed`` guards make concurrent or - repeat calls (e.g. session.close racing the WS-orphan reaper) harmless.""" +def _pop_session_by_id(sid: str) -> dict | None: + """Atomically detach one live session from the registry. + + Detaching is the ownership claim for teardown: once the record is no + longer in ``_sessions``, a concurrent close/reaper becomes a no-op. Keep + this operation separate from ``_teardown_session`` because finalization can + flush SQLite state, invoke plugins, commit memory, interrupt delegations, + and close agents/workers. None of that slow external work belongs under + the global ``_session_resume_lock``. + """ with _sessions_lock: session = _sessions.pop(sid, None) if session is None: - return False + return None # The session is already out of _sessions here, so downstream teardown # (e.g. _finalize_session's per-session async-delegation interrupt) can't # recover its live id by scanning the dict — stamp it on the record. session["_sid"] = sid + return session + + +def _teardown_popped_session( + session: dict | None, *, end_reason: str = "tui_close" +) -> bool: + """Finish a close after the caller has atomically detached the session.""" + if session is None: + return False _teardown_session(session, end_reason=end_reason) return True +def _close_session_by_id(sid: str, *, end_reason: str = "tui_close") -> bool: + """Single idempotent teardown funnel for callers needing no resume race. + + Resume-sensitive callers first pop under ``_session_resume_lock`` and then + call ``_teardown_popped_session`` after releasing it. Other reapers can use + this convenience wrapper directly. The pop remains the single atomic + ownership claim, so concurrent/repeat close attempts stay harmless. + """ + return _teardown_popped_session( + _pop_session_by_id(sid), end_reason=end_reason + ) + def _ws_session_is_orphaned(session: dict | None) -> bool: """True if a WS session has no live transport and no in-flight turn. @@ -795,9 +820,10 @@ def _schedule_ws_orphan_reap(sid: str) -> None: def _reap() -> None: # Serialize the orphan re-check against session.resume (which re-binds a # live transport under _session_resume_lock and would make this session - # non-orphaned). The actual pop + teardown then goes through the shared - # _close_session_by_id funnel so the dict mutation happens under - # _sessions_lock — consistent with every other _sessions mutator + # non-orphaned). Claim teardown by popping under both lifecycle locks, + # then release the global resume lock before the slow finalization work. + # The dict mutation still happens under _sessions_lock — consistent + # with every other _sessions mutator # (#39591: _reap previously popped under _session_resume_lock, giving no # mutual exclusion against _init_session / _close_session_by_id, which # guard with _sessions_lock). _sessions_lock is an RLock and the global @@ -805,7 +831,8 @@ def _schedule_ws_orphan_reap(sid: str) -> None: with _session_resume_lock: if not _ws_session_is_orphaned(_sessions.get(sid)): return - _close_session_by_id(sid, end_reason="ws_orphan_reap") + session = _pop_session_by_id(sid) + _teardown_popped_session(session, end_reason="ws_orphan_reap") timer = threading.Timer(_WS_ORPHAN_REAP_GRACE_S, _reap) timer.daemon = True @@ -5441,7 +5468,44 @@ def _enqueue_prompt(session: dict, text: Any, transport: Any) -> None: session["queued_prompt"] = {"text": text, "transport": transport} -def _handle_busy_submit(rid, sid: str, session: dict, text: Any, transport: Any) -> dict: +def _interrupt_busy_session(sid: str, session: dict, agent: Any) -> None: + """Interrupt a busy turn without blocking the RPC reader or session lock. + + Some providers cannot apply ``interrupt()`` until a synchronous tool or + network call returns. Running that call inline used to leave + ``prompt.submit`` holding ``history_lock`` for the whole wait, which in turn + blocked ``session.resume`` and delayed the queued prompt itself. Keep at + most one interrupt worker per session so repeated steering cannot leak an + unbounded number of blocked threads. + """ + use_agent = agent is not None and hasattr(agent, "interrupt") + use_compute_host = not use_agent and _session_uses_compute_host(session) + if not use_agent and not use_compute_host: + return + + with session["history_lock"]: + if session.get("_busy_interrupt_pending"): + return + session["_busy_interrupt_pending"] = True + + def interrupt() -> None: + try: + if use_agent: + agent.interrupt() + else: + _get_compute_host_supervisor().interrupt(sid) + except Exception: + pass + finally: + with session["history_lock"]: + session["_busy_interrupt_pending"] = False + + threading.Thread(target=interrupt, daemon=True, name=f"busy-interrupt-{sid}").start() + + +def _handle_busy_submit( + rid, sid: str, session: dict, text: Any, transport: Any +) -> dict | None: """Apply the ``display.busy_input_mode`` policy to a prompt that lands while a turn is in flight, instead of rejecting it with ``session busy``. @@ -5458,25 +5522,30 @@ def _handle_busy_submit(rid, sid: str, session: dict, text: Any, transport: Any) """ mode = _load_busy_input_mode() agent = session.get("agent") + with session["history_lock"]: + if not session.get("running"): + # The turn ended between prompt.submit's first busy check and this + # helper. Let the caller retry and claim the now-idle session. + return None if mode == "steer" and agent is not None and hasattr(agent, "steer"): try: if agent.steer(text): - session["last_active"] = time.time() + with session["history_lock"]: + session["last_active"] = time.time() return _ok(rid, {"status": "steered"}) except Exception: pass # fall through to queue - if mode != "queue" and agent is not None and hasattr(agent, "interrupt"): - try: - agent.interrupt() - except Exception: - pass - elif mode != "queue" and _session_uses_compute_host(session): - try: - _get_compute_host_supervisor().interrupt(sid) - except Exception: - pass - _enqueue_prompt(session, text, transport) - session["last_active"] = time.time() + # Queue before asking the live turn to stop. In particular, never call a + # provider or compute-host method while holding history_lock: an interrupt + # can wait behind the very operation it is trying to cancel. + with session["history_lock"]: + if not session.get("running"): + return None + _enqueue_prompt(session, text, transport) + session["last_active"] = time.time() + + if mode != "queue": + _interrupt_busy_session(sid, session, agent) return _ok(rid, {"status": "queued"}) @@ -8725,13 +8794,13 @@ def _(rid, params: dict) -> dict: @method("session.close") def _(rid, params: dict) -> dict: sid = params.get("session_id", "") - # Serialize against the WS-orphan reaper (which also pops under - # _session_resume_lock) so a disconnect-reap and an explicit close can't - # both tear the same session down. _close_session_by_id is the single - # idempotent teardown path (pop + _teardown_session) and returns False - # when the session is already gone. + # Serialize only the ownership claim against session.resume / the orphan + # reaper. Finalization may run arbitrary plugin/agent cleanup and must not + # keep every unrelated session.resume waiting behind it. with _session_resume_lock: - return _ok(rid, {"closed": _close_session_by_id(sid, end_reason="tui_close")}) + session = _pop_session_by_id(sid) + closed = _teardown_popped_session(session, end_reason="tui_close") + return _ok(rid, {"closed": closed}) @method("session.branch") @@ -9163,13 +9232,25 @@ def _(rid, params: dict) -> dict: # or fallback moved the session transport to stdio. if (t := current_transport()) is not None: session["transport"] = t + while True: + busy_transport = None + with session["history_lock"]: + if session.get("running"): + # Don't reject a mid-turn prompt — queue it (and, by default, + # interrupt the live turn) so it runs as the next turn. The + # provider interrupt itself must happen after this lock is + # released: a non-interruptible tool may keep it waiting. + busy_transport = t or session.get("transport") + else: + break + busy_response = _handle_busy_submit(rid, sid, session, text, busy_transport) + if busy_response is not None: + return busy_response + # The old turn finished between the two lock acquisitions. Retry the + # claim so this prompt starts normally instead of being stranded in a + # queue whose drain already ran. + with session["history_lock"]: - if session.get("running"): - # Don't reject a mid-turn prompt — queue it (and, by default, - # interrupt the live turn) so it runs as the next turn. See - # _handle_busy_submit for why the old "session busy" rejection - # dropped messages when teardown outlived the client's retry window. - return _handle_busy_submit(rid, sid, session, text, t or session.get("transport")) # A watch session's run lives in the PARENT turn, so its own running # flag is False — without this, typing mid-run builds a second agent # racing the in-flight child on the same stored session (interleaved