From 45e077e2fd898fd2de2ee6cb4c6d9ee8899b10d9 Mon Sep 17 00:00:00 2001 From: Yishova <272059605+Yishova@users.noreply.github.com> Date: Sun, 2 Aug 2026 05:23:41 -0400 Subject: [PATCH 1/3] tui_gateway: session.resume abandons the profile SessionDB it opens --- .../test_session_resume_db_ownership.py | 290 +++++++ tui_gateway/methods_session.py | 714 +++++++++--------- 2 files changed, 666 insertions(+), 338 deletions(-) create mode 100644 tests/tui_gateway/test_session_resume_db_ownership.py diff --git a/tests/tui_gateway/test_session_resume_db_ownership.py b/tests/tui_gateway/test_session_resume_db_ownership.py new file mode 100644 index 000000000000..309d1e3998a0 --- /dev/null +++ b/tests/tui_gateway/test_session_resume_db_ownership.py @@ -0,0 +1,290 @@ +"""``session.resume`` must not abandon the profile-scoped SessionDB it opens. + +In app-global remote mode a resume for another local profile opens a DEDICATED +``SessionDB(db_path=/state.db)`` handle (the ``session.resume`` handler +in tui_gateway/methods_session.py). That handle is the caller's to close until +it is handed to the long-lived agent by ``_init_session`` — and +``_init_session`` never closes a caller-supplied ``session_db`` (its +``_init_owns_db`` stays False for that case). + +Every early return before that transfer used to drop the handle on the floor, +so its SQLite fds stayed open for as long as anything kept the instance +reachable — and a ``SessionDB`` pins ITSELF once its background token writer +starts (``atexit.register(self._drain_token_queue_at_exit)``, which only +``close()`` unregisters). + +Pinned here, in both directions: + +* the pre-transfer early returns (session-not-found, "resume failed", the + live-session fast path, the deferred cold-resume return) all close it; +* a resume that COMPLETES the transfer leaves it open — closing there would + fault every later turn with "Cannot operate on a closed database"; +* an ``_init_session`` that raises AFTER registering the session must drop that + half-built registration, otherwise the live-session fast path serves a + session whose db we just closed on every later resume of the same id; +* the shared launch-profile handle (``_get_db()``) is never closed, since it + outlives the RPC. +""" + +from __future__ import annotations + +import types + +import pytest + +from tui_gateway import server + + +class _RecordingDB: + """Stand-in for ``hermes_state.SessionDB`` that counts ``close()`` calls. + + Implements only the surface ``session.resume`` touches. + """ + + def __init__(self, db_path=None, **_kwargs): + self.db_path = db_path + self.closed = 0 + self.rows: dict = {} + self.reopen_error: Exception | None = None + + def close(self): + self.closed += 1 + + def get_session(self, target): + return self.rows.get(target) + + def get_session_by_title(self, _target): + return None + + def resolve_resume_session_id(self, target): + return target + + def reopen_session(self, _target): + if self.reopen_error is not None: + raise self.reopen_error + + def get_resume_conversations(self, _target): + return ([], []) + + def get_ancestor_display_prefix(self, _target): + return [] + + def get_messages_as_conversation(self, _target, **_kwargs): + return [] + + +@pytest.fixture() +def profile_dbs(monkeypatch, tmp_path): + """Route profile-scoped opens to _RecordingDB; yield the list of opens. + + ``params['profile']`` selects the profile scope; omitting it resolves to + the launch profile (``_profile_home`` -> None) and the shared handle. + """ + opened: list[_RecordingDB] = [] + profile_home = tmp_path / "work" + profile_home.mkdir() + + def _factory(db_path=None, **kwargs): + db = _RecordingDB(db_path=db_path, **kwargs) + opened.append(db) + return db + + monkeypatch.setattr("hermes_state.SessionDB", _factory) + monkeypatch.setattr( + server, "_profile_home", lambda profile: profile_home if profile else None + ) + monkeypatch.setattr(server, "_profile_configured_cwd", lambda _home: str(tmp_path)) + # The handler builds nothing on the paths under test; keep it hermetic and + # off the real agent/secret/HERMES_HOME machinery. + monkeypatch.setattr(server, "_enable_gateway_prompts", lambda: None) + monkeypatch.setattr(server, "_find_live_session_by_key", lambda _key: None) + monkeypatch.setattr(server, "_schedule_agent_build", lambda *a, **k: None) + monkeypatch.setattr(server, "_schedule_session_cap_enforcement", lambda *a, **k: None) + monkeypatch.setattr(server, "_maybe_schedule_auto_continue", lambda *a, **k: None) + monkeypatch.setattr(server, "_default_session_cwd", lambda *a, **k: str(tmp_path)) + known = set(server._sessions) + yield opened + with server._sessions_lock: + for sid in [s for s in server._sessions if s not in known]: + server._sessions.pop(sid, None) + + +def _resume(**params): + return server.handle_request( + {"id": "1", "method": "session.resume", "params": params} + ) + + +def test_resume_closes_profile_db_when_session_not_found(profile_dbs): + """The 'session not found' early return must not leak the handle.""" + resp = _resume(session_id="missing", profile="work") + + assert resp["error"]["code"] == 4007 + assert len(profile_dbs) == 1 + assert profile_dbs[0].closed == 1 + + +def test_resume_closes_profile_db_when_reopen_fails(profile_dbs, monkeypatch): + """The 'resume failed' early return must not leak the handle.""" + + def _factory(db_path=None, **kwargs): + db = _RecordingDB(db_path=db_path, **kwargs) + db.rows["s1"] = {"id": "s1", "cwd": ""} + db.reopen_error = RuntimeError("database is locked") + profile_dbs.append(db) + return db + + monkeypatch.setattr("hermes_state.SessionDB", _factory) + + resp = _resume(session_id="s1", profile="work") + + assert resp["error"]["code"] == 5000 + assert "resume failed" in resp["error"]["message"] + assert profile_dbs[0].closed == 1 + + +def test_resume_closes_profile_db_on_live_session_fast_path(profile_dbs, monkeypatch): + """Re-resuming an already-live session returns early — and must close. + + This is the hottest leak in practice: every reconnect/tile-paint resume of + a chat that is already live takes this path, so the fd growth tracked + reconnect count rather than anything rare. + """ + + def _factory(db_path=None, **kwargs): + db = _RecordingDB(db_path=db_path, **kwargs) + db.rows["s1"] = {"id": "s1", "cwd": ""} + profile_dbs.append(db) + return db + + monkeypatch.setattr("hermes_state.SessionDB", _factory) + monkeypatch.setattr(server, "_find_live_session_by_key", lambda _key: ("live-sid", {})) + monkeypatch.setattr( + server, + "_live_session_payload", + lambda sid, session, **_kwargs: {"session_id": sid}, + ) + monkeypatch.setattr(server, "_child_run_active", lambda _key: False) + + resp = _resume(session_id="s1", profile="work") + + assert resp["result"]["resumed"] == "s1" + assert profile_dbs[0].closed == 1 + + +def test_resume_closes_profile_db_on_deferred_cold_resume(profile_dbs, monkeypatch): + """The DEFAULT resume path returns before any transfer — and must close. + + A cold resume without ``eager_build`` registers a deferred session record + (no agent, no db reference) and builds the agent later off the response + path, so the handle opened here is never handed to anyone. + """ + + def _factory(db_path=None, **kwargs): + db = _RecordingDB(db_path=db_path, **kwargs) + db.rows["s1"] = {"id": "s1", "cwd": ""} + profile_dbs.append(db) + return db + + monkeypatch.setattr("hermes_state.SessionDB", _factory) + monkeypatch.setattr(server, "_stored_session_runtime_overrides", lambda _found: {}) + + resp = _resume(session_id="s1", profile="work") + + assert resp["result"]["session_key"] == "s1" + assert resp["result"]["status"] == "idle" + assert profile_dbs[0].closed == 1 + + +def test_resume_keeps_profile_db_open_after_ownership_transfer(profile_dbs, monkeypatch): + """A COMPLETED resume transfers the handle to the agent — do not close it. + + Guards the other direction: closing here would hand the live session a dead + connection and fault every subsequent turn. + """ + captured: dict = {} + + def _factory(db_path=None, **kwargs): + db = _RecordingDB(db_path=db_path, **kwargs) + db.rows["s1"] = {"id": "s1", "cwd": ""} + profile_dbs.append(db) + return db + + def _fake_make_agent(sid, key, session_db=None, **_kwargs): + captured["agent_db"] = session_db + return types.SimpleNamespace(model="test") + + def _fake_init_session(sid, key, agent, history, session_db=None, **_kwargs): + captured["init_db"] = session_db + + monkeypatch.setattr("hermes_state.SessionDB", _factory) + monkeypatch.setattr(server, "_make_agent", _fake_make_agent) + monkeypatch.setattr(server, "_init_session", _fake_init_session) + monkeypatch.setattr(server, "_set_session_context", lambda _target: []) + monkeypatch.setattr(server, "_clear_session_context", lambda _tokens: None) + monkeypatch.setattr(server, "_stored_session_runtime_overrides", lambda _found: {}) + monkeypatch.setattr(server, "_session_info", lambda agent, *a: {"model": "test"}) + + resp = _resume(session_id="s1", profile="work", eager_build=True) + + assert resp["result"]["session_key"] == "s1" + db = profile_dbs[0] + # The agent and the live session both took THIS handle... + assert captured["agent_db"] is db + assert captured["init_db"] is db + # ...so the handler must have released ownership instead of closing it. + assert db.closed == 0 + + +def test_resume_drops_half_built_session_when_init_session_raises( + profile_dbs, monkeypatch +): + """Closing the handle is only safe if the failed registration goes with it. + + ``_init_session`` publishes ``_sessions[sid]`` BEFORE its first read through + the handle. If that read raises, the handle is still ours (and gets closed), + so the half-built session must not stay registered — otherwise the + live-session fast path serves that dead session on every later resume of the + same id, forever, with "'NoneType' object has no attribute 'execute'". + """ + captured: dict = {} + + def _factory(db_path=None, **kwargs): + db = _RecordingDB(db_path=db_path, **kwargs) + db.rows["s1"] = {"id": "s1", "cwd": ""} + profile_dbs.append(db) + return db + + def _fake_init_session(sid, key, agent, history, session_db=None, **_kwargs): + # Same ordering as the real one: register, THEN read through the db. + captured["sid"] = sid + with server._sessions_lock: + server._sessions[sid] = {"agent": agent, "session_key": key} + raise RuntimeError("database is locked") + + monkeypatch.setattr("hermes_state.SessionDB", _factory) + monkeypatch.setattr( + server, "_make_agent", lambda *a, **k: types.SimpleNamespace(model="test") + ) + monkeypatch.setattr(server, "_init_session", _fake_init_session) + monkeypatch.setattr(server, "_set_session_context", lambda _target: []) + monkeypatch.setattr(server, "_clear_session_context", lambda _tokens: None) + monkeypatch.setattr(server, "_stored_session_runtime_overrides", lambda _found: {}) + + resp = _resume(session_id="s1", profile="work", eager_build=True) + + assert resp["error"]["code"] == 5000 + assert profile_dbs[0].closed == 1 + assert captured["sid"] not in server._sessions + + +def test_resume_never_closes_shared_launch_db(profile_dbs, monkeypatch): + """No profile scope -> the shared ``_get_db()`` handle, which we never close.""" + shared = _RecordingDB(db_path="launch") + monkeypatch.setattr(server, "_get_db", lambda: shared) + + resp = _resume(session_id="missing") + + assert resp["error"]["code"] == 4007 + assert profile_dbs == [] # no dedicated handle was opened + assert shared.closed == 0 diff --git a/tui_gateway/methods_session.py b/tui_gateway/methods_session.py index 8a04660557a7..073fc399f711 100644 --- a/tui_gateway/methods_session.py +++ b/tui_gateway/methods_session.py @@ -321,185 +321,276 @@ def _(rid, params: dict) -> dict: # the caller explicitly requests it; other clients keep upstream behavior. omit_messages = is_truthy_value(params.get("omit_messages", False)) - # In a profile scope, the agent OWNS a long-lived db handle bound to that - # profile (do NOT auto-close it here). Otherwise reuse the shared launch db. + # In a profile scope this opens a DEDICATED handle we own until the agent + # takes it (see the ownership transfer at _init_session below); every path + # that returns before that transfer must close it. Otherwise reuse the + # shared launch db, which outlives the RPC and is never closed here. + owns_db = False if profile_home is not None: from hermes_state import SessionDB db = SessionDB(db_path=profile_home / "state.db") + owns_db = True else: db = _get_db() - if db is None: - return _db_unavailable_error(rid, code=5000) - - found = db.get_session(target) - if not found: - found = db.get_session_by_title(target) - if found: - target = found["id"] - elif is_truthy_value(params.get("lazy", False)) and _child_run_active(target): - # Race: a watch window opened on a freshly-spawned subagent. The - # child relays `subagent.start` (which carries child_session_id and - # triggers the window) BEFORE its first run_conversation() flushes - # the DB row via _ensure_db_session, so db.get_session(target) is - # momentarily empty. On slower hosts (notably WSL2, where SQLite + - # process scheduling widen the gap) the window's resume consistently - # lands inside this window and used to hard-fail "session not found" - # — the frontend then 404'd on the REST messages fallback and the - # window spun forever. The child is provably live (_child_run_active), - # so proceed into the lazy branch with empty history; the live mirror - # streams the whole turn anyway and the row exists by upgrade time. - found = {} - else: - return _err(rid, 4007, "session not found") + try: + if db is None: + return _db_unavailable_error(rid, code=5000) + + found = db.get_session(target) + if not found: + found = db.get_session_by_title(target) + if found: + target = found["id"] + elif is_truthy_value(params.get("lazy", False)) and _child_run_active(target): + # Race: a watch window opened on a freshly-spawned subagent. The + # child relays `subagent.start` (which carries child_session_id and + # triggers the window) BEFORE its first run_conversation() flushes + # the DB row via _ensure_db_session, so db.get_session(target) is + # momentarily empty. On slower hosts (notably WSL2, where SQLite + + # process scheduling widen the gap) the window's resume consistently + # lands inside this window and used to hard-fail "session not found" + # — the frontend then 404'd on the REST messages fallback and the + # window spun forever. The child is provably live (_child_run_active), + # so proceed into the lazy branch with empty history; the live mirror + # streams the whole turn anyway and the row exists by upgrade time. + found = {} + else: + return _err(rid, 4007, "session not found") + + # Follow the compression-continuation chain to the live tip so a resume on + # a rotated-out parent id binds to the descendant that actually holds the + # post-compression turns. Auto-compression ends the session and forks a + # continuation child; without this, resuming the original id (the desktop's + # routed id when the chat was opened before it rotated) reloads the parent + # transcript and the response generated after compression is missing — the + # "I came back and the reply isn't there" bug on large sessions. Resolving + # here also re-anchors the fast path below so a still-live rotated session + # is reused (by its new key) instead of rebuilding a duplicate agent on the + # stale parent. Skipped for lazy watch windows, which intentionally attach + # to the exact child branch they were opened on. + if found and not is_truthy_value(params.get("lazy", False)): + try: + tip = db.resolve_resume_session_id(target) + except Exception: + tip = target + if tip and tip != target: + target = tip + found = db.get_session(target) or found - # Follow the compression-continuation chain to the live tip so a resume on - # a rotated-out parent id binds to the descendant that actually holds the - # post-compression turns. Auto-compression ends the session and forks a - # continuation child; without this, resuming the original id (the desktop's - # routed id when the chat was opened before it rotated) reloads the parent - # transcript and the response generated after compression is missing — the - # "I came back and the reply isn't there" bug on large sessions. Resolving - # here also re-anchors the fast path below so a still-live rotated session - # is reused (by its new key) instead of rebuilding a duplicate agent on the - # stale parent. Skipped for lazy watch windows, which intentionally attach - # to the exact child branch they were opened on. - if found and not is_truthy_value(params.get("lazy", False)): - try: - tip = db.resolve_resume_session_id(target) - except Exception: - tip = target - if tip and tip != target: - target = tip - found = db.get_session(target) or found + profile_resume_cwd = str(found.get("cwd") or "").strip() or _profile_configured_cwd( + profile_home + ) - profile_resume_cwd = str(found.get("cwd") or "").strip() or _profile_configured_cwd( - profile_home - ) + def _reuse_live_payload(sid: str, session: dict) -> dict: + payload = _live_session_payload( + sid, + session, + cols=cols, + touch=True, + transport=current_transport() or _stdio_transport, + omit_messages=omit_messages, + ) + payload["resumed"] = target + # A lazy watch session never owns a run loop, so its payload's running + # flag is always False — overlay the child-run registry so a reconnecting + # watch window keeps its busy indicator while the child is still mid-run. + if session.get("agent") is None and _child_run_active(target): + payload["running"] = True + payload["status"] = "streaming" + return payload + + # Fast path: if the session is already live, reuse it under the lock. + with _session_resume_lock: + live = _find_live_session_by_key(target) + if live is not None: + return _ok(rid, _reuse_live_payload(*live)) + + # Lazy/watch resume: register the live session WITHOUT building an agent. + # Used by the desktop's subagent windows — the child runs inside the + # parent's turn, so its window only needs the stored history plus a + # transport for the child-mirror's live events. Skipping _make_agent here + # is what keeps the window cheap while the backend is busy running the + # delegation. A later prompt.submit upgrades it via _start_agent_build + # (resume_session_id keeps the upgrade on the stored conversation). + if is_truthy_value(params.get("lazy", False)): + sid = uuid.uuid4().hex[:8] + source = _resolve_session_source(str(params.get("source") or "").strip() or None) + lease = None # claimed lazily on the first turn (_ensure_active_session_slot) + try: + db.reopen_session(target) + # The child's OWN conversation only — include_ancestors would prepend + # the parent's transcript onto the subagent's branch. + # repair_alternation: this resume feeds LIVE REPLAY (the loaded + # history becomes the resumed session record's working conversation), + # so heal a durable ``user;user`` violation once here instead of + # re-firing the pre-request repair on every subsequent turn. + history = db.get_messages_as_conversation(target, repair_alternation=True) + except Exception as e: + if lease is not None: + lease.release() + return _err(rid, 5000, f"resume failed: {e}") + cwd = profile_resume_cwd or _default_session_cwd() + record = _deferred_session_record( + target, + cols=cols, + cwd=cwd, + history=history, + lease=lease, + source=source, + close_on_disconnect=is_truthy_value(params.get("close_on_disconnect", False)), + profile_home=profile_home, + lazy=True, + ) + if (live := _claim_or_reuse_live(sid, target, record, lease)) is not None: + return _ok(rid, _reuse_live_payload(*live)) + # A delegated child mid-run emits no session events of its own — report + # its liveness from the relay registry so the window shows a busy turn. + child_running = _child_run_active(target) + # User-visible messages use the VERBATIM display projection (child-only, + # no ancestors — matching the repaired read above), so model-invisible + # rows persisted by #65919 (verification candidates collapsed by + # repair_message_sequence) survive in the watch window just as they do + # on the eager resume + REST paths. The repaired ``history`` above still + # feeds live replay. Fall back to it if the display read fails. + try: + display_history = db.get_messages_as_conversation( + target, repair_alternation=False, include_row_ids=True + ) + except Exception: + logger.debug("child-watch display projection read failed", exc_info=True) + display_history = history + messages = [] if omit_messages else _history_to_messages(display_history) + return _ok( + rid, + { + "session_id": sid, + "resumed": target, + "message_count": len(display_history) if omit_messages else len(messages), + "messages": messages, + "messages_omitted": omit_messages, + "info": _lazy_resume_info(cwd, profile=profile), + "inflight": None, + "running": child_running, + "session_key": target, + "started_at": record["created_at"], + "status": "streaming" if child_running else "idle", + }, + ) - def _reuse_live_payload(sid: str, session: dict) -> dict: - payload = _live_session_payload( - sid, - session, - cols=cols, - touch=True, - transport=current_transport() or _stdio_transport, - omit_messages=omit_messages, - ) - payload["resumed"] = target - # A lazy watch session never owns a run loop, so its payload's running - # flag is always False — overlay the child-run registry so a reconnecting - # watch window keeps its busy indicator while the child is still mid-run. - if session.get("agent") is None and _child_run_active(target): - payload["running"] = True - payload["status"] = "streaming" - return payload - - # Fast path: if the session is already live, reuse it under the lock. - with _session_resume_lock: - live = _find_live_session_by_key(target) - if live is not None: - return _ok(rid, _reuse_live_payload(*live)) - - # Lazy/watch resume: register the live session WITHOUT building an agent. - # Used by the desktop's subagent windows — the child runs inside the - # parent's turn, so its window only needs the stored history plus a - # transport for the child-mirror's live events. Skipping _make_agent here - # is what keeps the window cheap while the backend is busy running the - # delegation. A later prompt.submit upgrades it via _start_agent_build - # (resume_session_id keeps the upgrade on the stored conversation). - if is_truthy_value(params.get("lazy", False)): - sid = uuid.uuid4().hex[:8] - source = _resolve_session_source(str(params.get("source") or "").strip() or None) - lease = None # claimed lazily on the first turn (_ensure_active_session_slot) - try: - db.reopen_session(target) - # The child's OWN conversation only — include_ancestors would prepend - # the parent's transcript onto the subagent's branch. - # repair_alternation: this resume feeds LIVE REPLAY (the loaded - # history becomes the resumed session record's working conversation), - # so heal a durable ``user;user`` violation once here instead of - # re-firing the pre-request repair on every subsequent turn. - history = db.get_messages_as_conversation(target, repair_alternation=True) - except Exception as e: - if lease is not None: - lease.release() - return _err(rid, 5000, f"resume failed: {e}") - cwd = profile_resume_cwd or _default_session_cwd() - record = _deferred_session_record( - target, - cols=cols, - cwd=cwd, - history=history, - lease=lease, - source=source, - close_on_disconnect=is_truthy_value(params.get("close_on_disconnect", False)), - profile_home=profile_home, - lazy=True, - ) - if (live := _claim_or_reuse_live(sid, target, record, lease)) is not None: - return _ok(rid, _reuse_live_payload(*live)) - # A delegated child mid-run emits no session events of its own — report - # its liveness from the relay registry so the window shows a busy turn. - child_running = _child_run_active(target) - # User-visible messages use the VERBATIM display projection (child-only, - # no ancestors — matching the repaired read above), so model-invisible - # rows persisted by #65919 (verification candidates collapsed by - # repair_message_sequence) survive in the watch window just as they do - # on the eager resume + REST paths. The repaired ``history`` above still - # feeds live replay. Fall back to it if the display read fails. - try: - display_history = db.get_messages_as_conversation( - target, repair_alternation=False, include_row_ids=True + # Cold resume default: register the live session and read its stored + # transcript, but build the agent OFF the response path. _make_agent can + # block for seconds (MCP discovery, prompt/skill build, AIAgent + # construction), and every resume caller (desktop + Ink TUI) awaits this RPC + # before it paints — so building eagerly is the bulk of the multi-second + # "switching sessions is frozen" latency. Return the full display transcript + # immediately and pre-warm the agent on a short timer (the same deferred- + # build contract session.create uses); _sess() also builds on demand if the + # first prompt beats the timer. A caller that needs the agent built + # synchronously (e.g. tests of the build race) passes ``eager_build: true`` + # to fall through to the eager path below. Distinct from the lazy/watch + # branch above: a normal resume restores the full ancestor history and the + # session's persisted runtime identity, and is a real (upgradable) session. + if not is_truthy_value(params.get("eager_build", False)): + sid = uuid.uuid4().hex[:8] + source = _resolve_session_source(str(params.get("source") or "").strip() or None) + lease = None # claimed lazily on the first turn (_ensure_active_session_slot) + # Interactive resume routes approvals/clarify through gateway prompts; + # the deferred build wires the remaining per-session callbacks. + _enable_gateway_prompts() + try: + db.reopen_session(target) + # One lineage SELECT feeds both projections (#67142-adjacent perf, + # from the desktop audit): the model-fed copy is alternation-repaired + # (raw_history → sanitize_replay_history → the resumed session's + # working conversation) and the display copy stays verbatim — + # inspection/export must show what is actually stored. + if omit_messages: + raw_history = db.get_messages_as_conversation( + target, repair_alternation=True + ) + display_history = [] + else: + raw_history, display_history = db.get_resume_conversations(target) + except Exception as e: + if lease is not None: + lease.release() + return _err(rid, 5000, f"resume failed: {e}") + # Display keeps the full transcript; the model-fed history drops a + # dangling/interrupted tool-call tail so a session killed mid-loop does + # not replay the unanswered call forever (#29086). + prefix = [] if omit_messages else db.get_ancestor_display_prefix(target) + history = sanitize_replay_history(raw_history) + # Restore the model/provider/reasoning/tier this chat last used so the + # deferred build (and the info below) match the eager path — without them + # the build drops the provider ("No LLM provider configured"). + overrides = _stored_session_runtime_overrides(found) or {} + model_override = overrides.get("model_override") or {} + cwd = profile_resume_cwd or _default_session_cwd() + record = _deferred_session_record( + target, + cols=cols, + cwd=cwd, + history=history, + lease=lease, + source=source, + close_on_disconnect=is_truthy_value(params.get("close_on_disconnect", False)), + display_history_prefix=prefix, + profile_home=profile_home, + model_override=overrides.get("model_override"), + resume_runtime_overrides=overrides or None, ) - except Exception: - logger.debug("child-watch display projection read failed", exc_info=True) - display_history = history - messages = [] if omit_messages else _history_to_messages(display_history) - return _ok( - rid, - { + if (live := _claim_or_reuse_live(sid, target, record, lease)) is not None: + return _ok(rid, _reuse_live_payload(*live)) + + _schedule_agent_build(sid) + _schedule_session_cap_enforcement() # trim detached idle sessions over the cap + auto_continue = _maybe_schedule_auto_continue(sid, record, target) + + messages = [] if omit_messages else _history_to_messages(display_history) + payload = { "session_id": sid, "resumed": target, - "message_count": len(display_history) if omit_messages else len(messages), + "message_count": len(raw_history) if omit_messages else len(messages), "messages": messages, "messages_omitted": omit_messages, - "info": _lazy_resume_info(cwd, profile=profile), + "info": _lazy_resume_info( + cwd, + model=model_override.get("model") or "", + provider=overrides.get("provider_override") or "", + profile=profile, + ), "inflight": None, - "running": child_running, + "running": False, "session_key": target, "started_at": record["created_at"], - "status": "streaming" if child_running else "idle", - }, - ) + "status": "idle", + } + if auto_continue is not None: + payload["auto_continue"] = auto_continue + return _ok(rid, payload) - # Cold resume default: register the live session and read its stored - # transcript, but build the agent OFF the response path. _make_agent can - # block for seconds (MCP discovery, prompt/skill build, AIAgent - # construction), and every resume caller (desktop + Ink TUI) awaits this RPC - # before it paints — so building eagerly is the bulk of the multi-second - # "switching sessions is frozen" latency. Return the full display transcript - # immediately and pre-warm the agent on a short timer (the same deferred- - # build contract session.create uses); _sess() also builds on demand if the - # first prompt beats the timer. A caller that needs the agent built - # synchronously (e.g. tests of the build race) passes ``eager_build: true`` - # to fall through to the eager path below. Distinct from the lazy/watch - # branch above: a normal resume restores the full ancestor history and the - # session's persisted runtime identity, and is a real (upgradable) session. - if not is_truthy_value(params.get("eager_build", False)): + # Build the agent OUTSIDE the lock — _make_agent can block for seconds + # (MCP discovery, prompt/skill build, AIAgent construction). Holding + # _session_resume_lock across it would stall session.close on the main + # dispatch thread (it's not a _LONG_HANDLER), blocking fast-path RPCs. sid = uuid.uuid4().hex[:8] source = _resolve_session_source(str(params.get("source") or "").strip() or None) lease = None # claimed lazily on the first turn (_ensure_active_session_slot) - # Interactive resume routes approvals/clarify through gateway prompts; - # the deferred build wires the remaining per-session callbacks. _enable_gateway_prompts() + home_token = ( + set_hermes_home_override(str(profile_home)) if profile_home is not None else None + ) + secret_token = ( + set_secret_scope(build_profile_secret_scope(Path(str(profile_home)))) + if profile_home is not None + else None + ) try: db.reopen_session(target) - # One lineage SELECT feeds both projections (#67142-adjacent perf, - # from the desktop audit): the model-fed copy is alternation-repaired - # (raw_history → sanitize_replay_history → the resumed session's - # working conversation) and the display copy stays verbatim — - # inspection/export must show what is actually stored. + # One lineage SELECT feeds both projections (see the interactive resume + # above): the model-fed copy is alternation-repaired for LIVE REPLAY, the + # display copy stays verbatim. if omit_messages: raw_history = db.get_messages_as_conversation( target, repair_alternation=True @@ -507,200 +598,147 @@ def _reuse_live_payload(sid: str, session: dict) -> dict: display_history = [] else: raw_history, display_history = db.get_resume_conversations(target) - except Exception as e: - if lease is not None: - lease.release() - return _err(rid, 5000, f"resume failed: {e}") - # Display keeps the full transcript; the model-fed history drops a - # dangling/interrupted tool-call tail so a session killed mid-loop does - # not replay the unanswered call forever (#29086). - prefix = [] if omit_messages else db.get_ancestor_display_prefix(target) - history = sanitize_replay_history(raw_history) - # Restore the model/provider/reasoning/tier this chat last used so the - # deferred build (and the info below) match the eager path — without them - # the build drops the provider ("No LLM provider configured"). - overrides = _stored_session_runtime_overrides(found) or {} - model_override = overrides.get("model_override") or {} - cwd = profile_resume_cwd or _default_session_cwd() - record = _deferred_session_record( - target, - cols=cols, - cwd=cwd, - history=history, - lease=lease, - source=source, - close_on_disconnect=is_truthy_value(params.get("close_on_disconnect", False)), - display_history_prefix=prefix, - profile_home=profile_home, - model_override=overrides.get("model_override"), - resume_runtime_overrides=overrides or None, - ) - if (live := _claim_or_reuse_live(sid, target, record, lease)) is not None: - return _ok(rid, _reuse_live_payload(*live)) - - _schedule_agent_build(sid) - _schedule_session_cap_enforcement() # trim detached idle sessions over the cap - auto_continue = _maybe_schedule_auto_continue(sid, record, target) - - messages = [] if omit_messages else _history_to_messages(display_history) - payload = { - "session_id": sid, - "resumed": target, - "message_count": len(raw_history) if omit_messages else len(messages), - "messages": messages, - "messages_omitted": omit_messages, - "info": _lazy_resume_info( - cwd, - model=model_override.get("model") or "", - provider=overrides.get("provider_override") or "", - profile=profile, - ), - "inflight": None, - "running": False, - "session_key": target, - "started_at": record["created_at"], - "status": "idle", - } - if auto_continue is not None: - payload["auto_continue"] = auto_continue - return _ok(rid, payload) - - # Build the agent OUTSIDE the lock — _make_agent can block for seconds - # (MCP discovery, prompt/skill build, AIAgent construction). Holding - # _session_resume_lock across it would stall session.close on the main - # dispatch thread (it's not a _LONG_HANDLER), blocking fast-path RPCs. - sid = uuid.uuid4().hex[:8] - source = _resolve_session_source(str(params.get("source") or "").strip() or None) - lease = None # claimed lazily on the first turn (_ensure_active_session_slot) - _enable_gateway_prompts() - home_token = ( - set_hermes_home_override(str(profile_home)) if profile_home is not None else None - ) - secret_token = ( - set_secret_scope(build_profile_secret_scope(Path(str(profile_home)))) - if profile_home is not None - else None - ) - try: - db.reopen_session(target) - # One lineage SELECT feeds both projections (see the interactive resume - # above): the model-fed copy is alternation-repaired for LIVE REPLAY, the - # display copy stays verbatim. - if omit_messages: - raw_history = db.get_messages_as_conversation( - target, repair_alternation=True - ) - display_history = [] - else: - raw_history, display_history = db.get_resume_conversations(target) - # The display transcript keeps every row so the user still sees their - # full history. The model-fed history is sanitized: a session whose - # last turn died mid-tool-loop persists a dangling assistant(tool_calls) - # (or interrupted assistant→tool) tail; replaying it makes the model - # re-issue the unanswered call forever — the permanent-"thinking" stuck - # session in #29086. The messaging gateway already strips this; this is - # the WebUI/TUI resume path picking up the same cleanup. - display_history_prefix = ( - [] if omit_messages else db.get_ancestor_display_prefix(target) - ) - history = sanitize_replay_history(raw_history) - messages = [] if omit_messages else _history_to_messages(display_history) - tokens = _set_session_context(target) - try: - # Pass the profile's db so the agent persists turns to the right - # state.db; home override is active here so config/skills/model - # resolve to the profile too. Runtime identity is restored from the - # stored session row so switching chats does not inherit whatever - # global model another chat last selected. - stored_runtime_overrides = _stored_session_runtime_overrides(found) - agent = _make_agent( - sid, - target, - session_id=target, - session_db=db, - platform_override=source, - **stored_runtime_overrides, - ) - finally: - _clear_session_context(tokens) - except Exception as e: - if lease is not None: - lease.release() - return _err(rid, 5000, f"resume failed: {e}") - finally: - if home_token is not None: - reset_hermes_home_override(home_token) - if secret_token is not None: - reset_secret_scope(secret_token) - - # Double-checked locking: another concurrent resume may have created the - # live session while we were building. Re-check under the lock; if it won, - # discard our just-built agent and reuse theirs (no worker/poller wired yet). - with _session_resume_lock: - live = _find_live_session_by_key(target) - if live is not None: - try: - if hasattr(agent, "close"): - agent.close() - except Exception: - pass - if lease is not None: - lease.release() - other_sid, other_session = live - payload = _live_session_payload( - other_sid, - other_session, - cols=cols, - touch=True, - transport=current_transport() or _stdio_transport, - omit_messages=omit_messages, - ) - payload["resumed"] = target - return _ok(rid, payload) - try: - init_home_token = ( - set_hermes_home_override(str(profile_home)) - if profile_home is not None - else None - ) - init_secret_token = ( - set_secret_scope(build_profile_secret_scope(Path(str(profile_home)))) - if profile_home is not None - else None + # The display transcript keeps every row so the user still sees their + # full history. The model-fed history is sanitized: a session whose + # last turn died mid-tool-loop persists a dangling assistant(tool_calls) + # (or interrupted assistant→tool) tail; replaying it makes the model + # re-issue the unanswered call forever — the permanent-"thinking" stuck + # session in #29086. The messaging gateway already strips this; this is + # the WebUI/TUI resume path picking up the same cleanup. + display_history_prefix = ( + [] if omit_messages else db.get_ancestor_display_prefix(target) ) + history = sanitize_replay_history(raw_history) + messages = [] if omit_messages else _history_to_messages(display_history) + tokens = _set_session_context(target) try: - _init_session( + # Pass the profile's db so the agent persists turns to the right + # state.db; home override is active here so config/skills/model + # resolve to the profile too. Runtime identity is restored from the + # stored session row so switching chats does not inherit whatever + # global model another chat last selected. + stored_runtime_overrides = _stored_session_runtime_overrides(found) + agent = _make_agent( sid, target, - agent, - history, - cols=cols, - cwd=profile_resume_cwd, + session_id=target, session_db=db, - source=source, + platform_override=source, + **stored_runtime_overrides, ) finally: - if init_home_token is not None: - reset_hermes_home_override(init_home_token) - if init_secret_token is not None: - reset_secret_scope(init_secret_token) - if sid in _sessions: - if stored_runtime_overrides.get("model_override") is not None: - _sessions[sid]["model_override"] = stored_runtime_overrides[ - "model_override" - ] - _sessions[sid]["display_history_prefix"] = display_history_prefix - # Remember the profile home so each turn re-binds HERMES_HOME (the - # agent persists to its own db, but mid-turn home reads — memory, - # skills — must resolve to the resumed profile too). - if profile_home is not None: - _sessions[sid]["profile_home"] = str(profile_home) - _sessions[sid]["active_session_lease"] = lease + _clear_session_context(tokens) except Exception as e: if lease is not None: lease.release() return _err(rid, 5000, f"resume failed: {e}") - session = _sessions.get(sid) or {} + finally: + if home_token is not None: + reset_hermes_home_override(home_token) + if secret_token is not None: + reset_secret_scope(secret_token) + + # Double-checked locking: another concurrent resume may have created the + # live session while we were building. Re-check under the lock; if it won, + # discard our just-built agent and reuse theirs (no worker/poller wired yet). + with _session_resume_lock: + live = _find_live_session_by_key(target) + if live is not None: + try: + if hasattr(agent, "close"): + agent.close() + except Exception: + pass + if lease is not None: + lease.release() + other_sid, other_session = live + payload = _live_session_payload( + other_sid, + other_session, + cols=cols, + touch=True, + transport=current_transport() or _stdio_transport, + omit_messages=omit_messages, + ) + payload["resumed"] = target + return _ok(rid, payload) + try: + init_home_token = ( + set_hermes_home_override(str(profile_home)) + if profile_home is not None + else None + ) + init_secret_token = ( + set_secret_scope(build_profile_secret_scope(Path(str(profile_home)))) + if profile_home is not None + else None + ) + try: + _init_session( + sid, + target, + agent, + history, + cols=cols, + cwd=profile_resume_cwd, + session_db=db, + source=source, + ) + # Ownership TRANSFER — the registered session's agent now + # holds this handle for its whole life, and _init_session + # never closes a caller-supplied session_db (its + # _init_owns_db stays False). Closing it in the finally + # below would fault every later turn on this session with + # "Cannot operate on a closed database". + owns_db = False + finally: + if init_home_token is not None: + reset_hermes_home_override(init_home_token) + if init_secret_token is not None: + reset_secret_scope(init_secret_token) + if sid in _sessions: + if stored_runtime_overrides.get("model_override") is not None: + _sessions[sid]["model_override"] = stored_runtime_overrides[ + "model_override" + ] + _sessions[sid]["display_history_prefix"] = display_history_prefix + # Remember the profile home so each turn re-binds HERMES_HOME (the + # agent persists to its own db, but mid-turn home reads — memory, + # skills — must resolve to the resumed profile too). + if profile_home is not None: + _sessions[sid]["profile_home"] = str(profile_home) + _sessions[sid]["active_session_lease"] = lease + except Exception as e: + # _init_session registers _sessions[sid] BEFORE its first read + # through this handle. If it raised in between — "database is + # locked" is the realistic trigger — the half-built session is + # still registered while the finally below closes the handle it + # holds, and the live-session fast path above would then serve + # that dead session on every later resume of this id + # ("'NoneType' object has no attribute 'execute'", permanently). + # owns_db still True means ownership never transferred, so the + # registration is ours to undo. + if owns_db: + with _sessions_lock: + _sessions.pop(sid, None) + if lease is not None: + lease.release() + return _err(rid, 5000, f"resume failed: {e}") + session = _sessions.get(sid) or {} + finally: + # Every return that does NOT reach the transfer above abandons this + # handle — session-not-found, both "resume failed" paths, the live-session + # fast path (the hot one: reconnects re-resume live chats through it), the + # deferred cold/lazy returns, and the double-checked-locking discard. + # Dropping it merely relied on refcounting to release the sqlite fds; that + # stops being true the moment anything pins the instance — SessionDB pins + # ITSELF once its background token writer starts, via + # atexit.register(_drain_token_queue_at_exit) (hermes_state.py), which only + # close() unregisters. A pinned handle keeps its db/-wal/-shm fds and its + # writer thread for the life of the process. + if owns_db and db is not None: + with contextlib.suppress(Exception): + db.close() auto_continue = ( _maybe_schedule_auto_continue(sid, session, target) if session else None ) From 24ae03f6a9847a53bcb76bfb91ae4a0a3fbb1372 Mon Sep 17 00:00:00 2001 From: Yishova <272059605+Yishova@users.noreply.github.com> Date: Sun, 2 Aug 2026 07:21:20 -0400 Subject: [PATCH 2/3] tui_gateway: close dedicated profile SessionDB handles at teardown too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the review on the session.resume ownership fix. Closing the pre-transfer early returns left two gaps, both real. 1. The transfer had no owner on the other side. Once ownership moved to the agent, teardown ran AIAgent.close() (via _teardown_session on session.close and the orphaned-session reaper), which called session_db.end_session() — that finalizes the session ROW, not the connection. A successfully resumed profile session kept its dedicated handle, its db/-wal/-shm fds and its background token-writer thread for the life of the gateway. AIAgent now carries an explicit _owns_session_db, defaulting False so the SHARED launch handle — which outlives every agent and backs every other live session — is still never closed there. Only the dedicated-open sites set it, at the point ownership actually changes hands. 2. session.resume was not the only profile-scoped open with no close on its failure paths. Covered here with the same flag, via a _transfer_db_to_agent helper that refuses the transfer unless the agent really holds that handle: - the deferred builder (_start_agent_build), including the session-reaped- mid-build case, where the built agent is discarded and never torn down, so transferring to it would leak exactly as before; - session.branch's branch_db; - the compute host's per-profile open; - AIAgent's own lazy open in _get_session_db_for_recall, which no other object ever references and so was unconditionally abandoned. Where a handle has already reached a registered session, the drop is unconditional and the transfer is best-effort on top: a refused transfer leaves the old leak, which is survivable, whereas closing under a live session is the permanent "Cannot operate on a closed database" break the original patch exists to avoid. Tests: tests/tui_gateway/test_session_db_ownership_teardown.py (new, 14). 11 of the 14 fail without this change; the 3 that pass are the "must NOT close" guards, which hold in both directions by design. --- agent/agent_init.py | 9 + run_agent.py | 23 +- .../test_session_db_ownership_teardown.py | 364 ++++++++++++++++++ tui_gateway/compute_host.py | 13 + tui_gateway/methods_session.py | 36 +- tui_gateway/server.py | 47 +++ 6 files changed, 490 insertions(+), 2 deletions(-) create mode 100644 tests/tui_gateway/test_session_db_ownership_teardown.py diff --git a/agent/agent_init.py b/agent/agent_init.py index 6f89ed237dca..df98acc856ef 100644 --- a/agent/agent_init.py +++ b/agent/agent_init.py @@ -1566,6 +1566,15 @@ def init_agent( # SQLite session store (optional -- provided by CLI or gateway) agent._session_db = session_db + # Whether close() must also close that handle. Default False: a + # caller-supplied session_db is almost always the SHARED launch handle, + # which outlives every agent and must never be closed here. Callers that + # hand over a DEDICATED handle (the gateway's per-profile state.db opens) + # set this True at the point ownership transfers, so teardown releases the + # sqlite fds and the token-writer thread instead of leaking them for the + # life of the process. Also set True on the lazy self-open in + # _get_session_db_for_recall, where nothing else holds a reference. + agent._owns_session_db = False agent._parent_session_id = parent_session_id # A close flush and the worker's turn-start flush can overlap. The durable # marker is attached to each in-memory message dict, so its test-and-append diff --git a/run_agent.py b/run_agent.py index 67cb22a71e29..4f945847de6b 100644 --- a/run_agent.py +++ b/run_agent.py @@ -613,6 +613,9 @@ def _get_session_db_for_recall(self): from hermes_state import SessionDB self._session_db = SessionDB() + # We opened it here, so nothing else holds a reference — this agent + # is its only owner and close() must release it. + self._owns_session_db = True return self._session_db except Exception: logger.debug("SessionDB unavailable for recall", exc_info=True) @@ -4324,15 +4327,33 @@ def close(self) -> None: # must leave it open). end_session() is first-reason-wins and no-ops on # an already-ended row, so this never clobbers a 'compression' / # 'cron_complete' / 'cli_close' reason set by an earlier terminal path. + session_db = getattr(self, "_session_db", None) try: if getattr(self, "_end_session_on_close", True): - session_db = getattr(self, "_session_db", None) session_id = getattr(self, "session_id", None) if session_db and session_id: session_db.end_session(session_id, "agent_close") except Exception: pass + # 9. Close the SQLite handle itself, but ONLY when this agent owns it. + # end_session() above finalizes the session ROW; it does not release the + # connection. For the shared launch handle that is correct — it outlives + # every agent — so _owns_session_db defaults False and this is a no-op. + # A DEDICATED handle (the gateway's per-profile state.db opens, and the + # lazy self-open in _get_session_db_for_recall) has no other owner: left + # unclosed it keeps its db/-wal/-shm fds and its background token-writer + # thread, and once that writer has started the instance pins ITSELF via + # atexit.register(_drain_token_queue_at_exit) — which only close() + # unregisters — so it survives for the life of the process. + # Cleared first so the documented idempotency of close() holds. + try: + if getattr(self, "_owns_session_db", False) and session_db is not None: + self._owns_session_db = False + session_db.close() + except Exception: + pass + def _hydrate_todo_store(self, history: List[Dict[str, Any]]) -> None: """ Recover todo state from conversation history. diff --git a/tests/tui_gateway/test_session_db_ownership_teardown.py b/tests/tui_gateway/test_session_db_ownership_teardown.py new file mode 100644 index 000000000000..47a5aa413e6f --- /dev/null +++ b/tests/tui_gateway/test_session_db_ownership_teardown.py @@ -0,0 +1,364 @@ +"""Dedicated profile ``SessionDB`` handles must be closed by whoever ends up owning them. + +Companion to ``test_session_resume_db_ownership.py``, which pins the paths that +return BEFORE a handle reaches an agent. This file pins the other half — the two +gaps that survived that change: + +1. **Teardown.** Once ownership transfers, the agent is the owner, and + ``AIAgent.close()`` (reached from ``_teardown_session`` on ``session.close`` + and the orphaned-session reaper) only called ``session_db.end_session()`` — + which finalizes the session ROW, not the connection. A successfully resumed + profile session therefore kept its dedicated SQLite handle, its db/-wal/-shm + fds and its background token-writer thread alive for the life of the gateway. + Ownership is explicit (``_owns_session_db``) precisely so this close can + happen without ever touching the SHARED launch handle, which outlives every + agent. + +2. **The other pre-transfer build paths.** ``session.resume`` was not the only + profile-scoped open. The deferred builder (``_start_agent_build``), the + branch handler and the compute host all open a dedicated handle, pass it to + ``_make_agent``, and had no close on their failure paths. + +The direction that must NOT regress is asserted everywhere: the shared launch +handle is never closed, and a handle that WAS transferred is not closed twice. +""" + +from __future__ import annotations + +import threading +import types + +import pytest + +from tui_gateway import server + + +class _RecordingDB: + """Stand-in for ``hermes_state.SessionDB`` that counts ``close()`` calls.""" + + def __init__(self, db_path=None, **_kwargs): + self.db_path = db_path + self.closed = 0 + + def close(self): + self.closed += 1 + + def end_session(self, *_a, **_k): + pass + + +# --------------------------------------------------------------------------- +# 1. AIAgent.close() — the teardown owner +# --------------------------------------------------------------------------- + + +def _bare_agent(**attrs): + """An AIAgent with __init__ bypassed, carrying only what close() reads.""" + from unittest.mock import patch + + with patch("run_agent.AIAgent.__init__", return_value=None): + from run_agent import AIAgent + + agent = AIAgent.__new__(AIAgent) + agent.session_id = "sid" + agent._active_children = [] + agent._active_children_lock = threading.Lock() + agent.client = None + for key, value in attrs.items(): + setattr(agent, key, value) + return agent + + +def test_close_closes_a_dedicated_handle_it_owns(): + """The gap the review found: end_session() is not close().""" + db = _RecordingDB() + agent = _bare_agent(_session_db=db, _owns_session_db=True) + + agent.close() + + assert db.closed == 1 + + +def test_close_never_closes_a_shared_handle(): + """The direction that must not regress. + + Almost every agent is handed the SHARED launch handle, which outlives it and + is used by every other live session. Closing that on teardown would break + every other chat in the gateway, so ownership defaults to False and only the + dedicated-open sites set it. + """ + db = _RecordingDB() + agent = _bare_agent(_session_db=db, _owns_session_db=False) + + agent.close() + + assert db.closed == 0 + + +def test_close_without_ownership_attribute_does_not_close(): + """Agents built before this flag existed (and test doubles) must be safe.""" + db = _RecordingDB() + agent = _bare_agent(_session_db=db) # no _owns_session_db at all + + agent.close() # must not raise + + assert db.closed == 0 + + +def test_close_is_idempotent_for_an_owned_handle(): + """close() is documented as safe to call repeatedly. + + Teardown can genuinely reach an agent twice (session.close racing the + orphaned-session reaper), so the second call must not double-close. + """ + db = _RecordingDB() + agent = _bare_agent(_session_db=db, _owns_session_db=True) + + agent.close() + agent.close() + + assert db.closed == 1 + + +def test_close_still_ends_the_session_row_before_closing(): + """Ordering matters: the row is finalized THROUGH the handle we then close.""" + calls: list[str] = [] + + class _Ordered(_RecordingDB): + def end_session(self, *_a, **_k): + calls.append("end_session") + + def close(self): + calls.append("close") + super().close() + + db = _Ordered() + agent = _bare_agent( + _session_db=db, _owns_session_db=True, _end_session_on_close=True + ) + + agent.close() + + assert calls == ["end_session", "close"] + + +def test_lazy_recall_open_is_owned_by_the_agent(monkeypatch): + """The agent's own lazy open has no other owner, so close() must release it. + + ``_get_session_db_for_recall`` opens a handle when no frontend supplied one. + Nothing else ever holds a reference to it, so before this change it was + unconditionally abandoned. + """ + opened: list[_RecordingDB] = [] + + def _factory(*_a, **_k): + db = _RecordingDB() + opened.append(db) + return db + + monkeypatch.setattr("hermes_state.SessionDB", _factory) + + agent = _bare_agent(_session_db=None, _persist_disabled=False) + got = agent._get_session_db_for_recall() + + assert got is opened[0] + assert agent._owns_session_db is True + + agent.close() + assert opened[0].closed == 1 + + +# --------------------------------------------------------------------------- +# 2. _transfer_db_to_agent — the transfer contract +# --------------------------------------------------------------------------- + + +def test_transfer_marks_the_agent_that_holds_the_handle(): + db = _RecordingDB() + agent = types.SimpleNamespace(_session_db=db, _owns_session_db=False) + + assert server._transfer_db_to_agent(agent, db) is True + assert agent._owns_session_db is True + + +def test_transfer_is_refused_when_the_agent_holds_a_different_handle(): + """A refusal is the signal that the caller still owns the handle. + + If the build handed the agent some other db, marking it would make the agent + close a handle it does not hold while the real one leaks. + """ + db, other = _RecordingDB(), _RecordingDB() + agent = types.SimpleNamespace(_session_db=other, _owns_session_db=False) + + assert server._transfer_db_to_agent(agent, db) is False + assert agent._owns_session_db is False + + +@pytest.mark.parametrize( + "agent, db", + [ + (None, _RecordingDB()), + (types.SimpleNamespace(_session_db=None), None), + ], +) +def test_transfer_is_refused_for_missing_operands(agent, db): + assert server._transfer_db_to_agent(agent, db) is False + + +# --------------------------------------------------------------------------- +# 3. The deferred builder — _start_agent_build +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def build_env(monkeypatch, tmp_path): + """Neutralize everything the deferred build touches except db ownership.""" + profile_home = tmp_path / "work" + profile_home.mkdir() + + opened: list[_RecordingDB] = [] + + def _factory(db_path=None, **kwargs): + db = _RecordingDB(db_path=db_path, **kwargs) + opened.append(db) + return db + + monkeypatch.setattr("hermes_state.SessionDB", _factory) + for name, value in [ + ("_set_session_context", lambda _key: []), + ("_clear_session_context", lambda _tokens: None), + ("_wire_callbacks", lambda _sid: None), + ("_config_model_target", lambda: None), + ("_load_memory_notifications", lambda: False), + ("_start_notification_poller", lambda _sid, _session: None), + ("_notify_session_boundary", lambda *a, **k: None), + ("_session_info", lambda *a, **k: {}), + ("_probe_config_health", lambda _cfg: None), + ("_load_cfg", lambda: {}), + ("_emit", lambda *a, **k: None), + ("_schedule_mcp_late_refresh", lambda *a, **k: None), + ("_session_source", lambda _current: None), + ("_child_run_active", lambda _key: False), + ]: + if hasattr(server, name): + monkeypatch.setattr(server, name, value) + monkeypatch.setattr(server, "set_hermes_home_override", lambda _home: None) + monkeypatch.setattr(server, "reset_hermes_home_override", lambda _tok: None) + yield types.SimpleNamespace(opened=opened, profile_home=str(profile_home)) + + +def _run_build(sid, session): + """Drive _start_agent_build to completion (it builds on a daemon thread).""" + server._start_agent_build(sid, session) + assert session["agent_ready"].wait(timeout=10), "build thread did not finish" + + +def _session(profile_home): + return { + "session_key": "key-1", + "agent_ready": threading.Event(), + "profile_home": profile_home, + } + + +@pytest.fixture() +def registered(monkeypatch): + """Register/unregister sessions in the module-global _sessions map.""" + added: list[str] = [] + + def _add(sid, session): + with server._sessions_lock: + server._sessions[sid] = session + added.append(sid) + + yield _add + with server._sessions_lock: + for sid in added: + server._sessions.pop(sid, None) + + +def test_deferred_build_closes_the_handle_when_the_build_fails( + build_env, registered, monkeypatch +): + """The failure path the review named: nothing takes the handle, so close it.""" + monkeypatch.setattr( + server, + "_make_agent", + lambda *a, **k: (_ for _ in ()).throw(RuntimeError("no provider")), + ) + sid, session = "sid-fail", _session(build_env.profile_home) + registered(sid, session) + + _run_build(sid, session) + + assert session.get("agent") is None + assert len(build_env.opened) == 1 + assert build_env.opened[0].closed == 1 + + +def test_deferred_build_transfers_the_handle_on_success( + build_env, registered, monkeypatch +): + """A built, retained agent becomes the owner — and the builder must not close.""" + captured: dict = {} + + def _fake_make_agent(sid, key, session_db=None, **_kwargs): + captured["db"] = session_db + return types.SimpleNamespace(_session_db=session_db, _owns_session_db=False) + + monkeypatch.setattr(server, "_make_agent", _fake_make_agent) + sid, session = "sid-ok", _session(build_env.profile_home) + registered(sid, session) + + _run_build(sid, session) + + db = build_env.opened[0] + assert captured["db"] is db + assert db.closed == 0 + # Ownership landed on the agent, so _teardown_session releases it later. + assert session["agent"]._owns_session_db is True + + +def test_deferred_build_closes_the_handle_when_the_session_is_reaped_midbuild( + build_env, registered, monkeypatch +): + """A discarded agent is never torn down, so transferring to it would leak. + + ``_build`` already computes ``replaced`` for the approval-notifier cleanup. + When the session was swapped out from under the build, the agent it produced + is unreachable — ``_teardown_session`` will never call close() on it — so the + handle has to be closed right here instead of handed over. + """ + + def _fake_make_agent(sid, key, session_db=None, **_kwargs): + # Simulate a concurrent reap landing while the agent was being built. + with server._sessions_lock: + server._sessions[sid] = {"session_key": "someone-else"} + return types.SimpleNamespace(_session_db=session_db, _owns_session_db=False) + + monkeypatch.setattr(server, "_make_agent", _fake_make_agent) + sid, session = "sid-reaped", _session(build_env.profile_home) + registered(sid, session) + + _run_build(sid, session) + + db = build_env.opened[0] + assert db.closed == 1 + assert session["agent"]._owns_session_db is False + + +def test_deferred_build_never_opens_or_closes_for_the_launch_profile( + build_env, registered, monkeypatch +): + """No profile scope -> no dedicated handle; the shared one is untouched.""" + monkeypatch.setattr( + server, + "_make_agent", + lambda *a, **k: types.SimpleNamespace(_session_db=None, _owns_session_db=False), + ) + sid, session = "sid-launch", _session(None) + registered(sid, session) + + _run_build(sid, session) + + assert build_env.opened == [] diff --git a/tui_gateway/compute_host.py b/tui_gateway/compute_host.py index d90024557aad..c4d5a6ae7c7b 100644 --- a/tui_gateway/compute_host.py +++ b/tui_gateway/compute_host.py @@ -9,6 +9,7 @@ import argparse import concurrent.futures +import contextlib import json import os import signal @@ -540,6 +541,7 @@ def _ensure_server_session(self, server: Any, frame: dict[str, Any]) -> dict: history = frame.get("history") if isinstance(frame.get("history"), list) else [] profile_home = str(frame.get("profile_home") or "") session_db = None + owns_db = False home_token = None secret_token = None try: @@ -550,7 +552,13 @@ def _ensure_server_session(self, server: Any, frame: dict[str, Any]) -> dict: home_token = set_hermes_home_override(profile_home) secret_token = set_secret_scope(build_profile_secret_scope(Path(profile_home))) + # DEDICATED handle — ours only until _make_agent succeeds. Every + # path after that keeps the agent registered in + # server._sessions[sid] (via _init_session, or the fallback dict + # in the except below), so the agent is the right owner; a + # _make_agent that RAISES is the one path where nothing takes it. session_db = SessionDB(db_path=Path(profile_home) / "state.db") + owns_db = True agent = server._make_agent( sid, key, @@ -561,7 +569,12 @@ def _ensure_server_session(self, server: Any, frame: dict[str, Any]) -> dict: platform_override=frame.get("source"), session_db=session_db, ) + if server._transfer_db_to_agent(agent, session_db): + owns_db = False finally: + if owns_db and session_db is not None: + with contextlib.suppress(Exception): + session_db.close() if home_token is not None: try: from hermes_constants import reset_hermes_home_override diff --git a/tui_gateway/methods_session.py b/tui_gateway/methods_session.py index 073fc399f711..de575732b79c 100644 --- a/tui_gateway/methods_session.py +++ b/tui_gateway/methods_session.py @@ -690,6 +690,22 @@ def _reuse_live_payload(sid: str, session: dict) -> dict: # _init_owns_db stays False). Closing it in the finally # below would fault every later turn on this session with # "Cannot operate on a closed database". + # + # Ownership moves ONTO the agent rather than just being + # dropped: AIAgent.close() (reached from _teardown_session + # on session.close and the orphaned-session reaper) closes + # a handle it owns, so the dedicated fds and the token + # writer are released at teardown instead of living as + # long as the gateway process. + # + # The drop is UNCONDITIONAL and the transfer is best-effort + # on top of it, deliberately. Past this line the session is + # registered and holding this handle, so the finally must + # not close it even if the transfer was refused — a refusal + # leaves the old leak, which is survivable; closing under a + # live session is the permanent "Cannot operate on a closed + # database" break this patch exists to avoid. + _transfer_db_to_agent(agent, db) owns_db = False finally: if init_home_token is not None: @@ -2783,6 +2799,10 @@ def _(rid, params: dict) -> dict: if lease is not None: lease.release() return _err(rid, 5008, f"branch failed: {e}") + # Bound before the try so the ownership finally below can never see them + # unbound, whatever raises inside. + branch_db = None + branch_owns_db = False try: # Bind the branched AGENT to the parent's profile, mirroring # session.create/resume: home override so config/skills/memory resolve @@ -2792,11 +2812,14 @@ def _(rid, params: dict) -> dict: # parent's db while the agent stayed on the launch handle would # recreate the cross-profile split one turn later. parent_home = session.get("profile_home") - branch_db = None if parent_home: from hermes_state import SessionDB + # DEDICATED handle, same ownership rule as session.resume: ours + # until the branched agent takes it below. _make_agent raising, or + # _init_session raising, both leave here without that transfer. branch_db = SessionDB(db_path=Path(parent_home) / "state.db") + branch_owns_db = True home_token = ( set_hermes_home_override(parent_home) if parent_home else None ) @@ -2833,6 +2856,13 @@ def _(rid, params: dict) -> dict: source=source, profile_home=parent_home, ) + # Ownership TRANSFER — the branched session's agent holds this + # handle for its whole life and closes it on teardown. Drop is + # unconditional for the same reason as session.resume: past + # _init_session the branched session is registered against this + # handle, so the finally must not close it. + _transfer_db_to_agent(agent, branch_db) + branch_owns_db = False finally: if secret_token is not None: reset_secret_scope(secret_token) @@ -2844,6 +2874,10 @@ def _(rid, params: dict) -> dict: if lease is not None: lease.release() return _err(rid, 5000, f"agent init failed on branch: {e}") + finally: + if branch_owns_db and branch_db is not None: + with contextlib.suppress(Exception): + branch_db.close() branched_session = _sessions.get(new_sid) return _ok( rid, diff --git a/tui_gateway/server.py b/tui_gateway/server.py index ed8f1dcc0ce0..462b89ff249c 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -1348,6 +1348,34 @@ def _db_for_profile(profile: str | None = None): return None, False +def _transfer_db_to_agent(agent, db) -> bool: + """Hand a DEDICATED profile handle to *agent*, which closes it on teardown. + + The build sites open a per-profile ``state.db`` handle, pass it to + ``_make_agent``, and own it until the built agent is the one that will be + torn down. This marks that transfer: from here ``AIAgent.close()`` (reached + via :func:`_teardown_session`) releases the handle, so the caller must stop + closing it. + + Returns True only when the transfer actually happened. It is refused when + *agent* is not holding *this* handle — the build failed before + ``_make_agent``, or the agent was given a different db — because a False + return is what tells the caller the handle is still its own to close. + Never called for the shared launch handle: that one is opened by + ``_get_db()``, outlives every agent, and stays at ``_owns_session_db`` + False. + """ + if agent is None or db is None: + return False + try: + if getattr(agent, "_session_db", None) is not db: + return False + agent._owns_session_db = True + return True + except Exception: + return False + + @contextlib.contextmanager def _profile_db(params: dict | None = None): """Yield the SessionDB for ``params['profile']`` (app-global remote mode). @@ -2139,6 +2167,8 @@ def _build() -> None: notify_registered = False home_token = None secret_token = None + session_db = None + owns_db = False profile_home = current.get("profile_home") try: tokens = _set_session_context(key) @@ -2157,7 +2187,12 @@ def _build() -> None: try: from hermes_state import SessionDB + # DEDICATED handle — ours until _transfer_db_to_agent hands + # it to the built agent in the finally below. Every path + # that leaves this build without that transfer (the except + # below, and a session reaped mid-build) must close it. session_db = SessionDB(db_path=Path(profile_home) / "state.db") + owns_db = True except Exception: session_db = None @@ -2293,6 +2328,18 @@ def _build() -> None: unregister_gateway_notify(key) except Exception: pass + # Dedicated profile handle: hand it to the agent that will actually + # be torn down, or close it here when no such agent exists. Both + # non-transfer cases are real: the except above (build raised, so + # nothing holds the handle) and `replaced` (the session was reaped + # mid-build, so this agent is discarded and _teardown_session will + # never reach it). Transferring to a discarded agent would leak the + # handle exactly as before. + if owns_db and session_db is not None: + built = None if replaced else current.get("agent") + if not _transfer_db_to_agent(built, session_db): + with contextlib.suppress(Exception): + session_db.close() ready.set() build_thread = threading.Thread(target=_build, daemon=True) From c10557a8c2d1ce4f6303fbed840045b7231af403 Mon Sep 17 00:00:00 2001 From: kshitij <82637225+kshitijk4poor@users.noreply.github.com> Date: Fri, 7 Aug 2026 19:22:45 +0530 Subject: [PATCH 3/3] test(tui_gateway): pin the raising-close swallow + no-retry contract Review finding: the close block's comment promises a raising session_db.close() is swallowed with the flag already cleared (no re-close on a second agent.close()), but nothing pinned it. One test with a raising _RecordingDB proves both halves. --- .../test_session_db_ownership_teardown.py | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/tui_gateway/test_session_db_ownership_teardown.py b/tests/tui_gateway/test_session_db_ownership_teardown.py index 47a5aa413e6f..89b30bce3724 100644 --- a/tests/tui_gateway/test_session_db_ownership_teardown.py +++ b/tests/tui_gateway/test_session_db_ownership_teardown.py @@ -120,6 +120,28 @@ def test_close_is_idempotent_for_an_owned_handle(): assert db.closed == 1 +def test_raising_close_is_swallowed_and_not_retried(): + """A raising ``session_db.close()`` must not escape ``agent.close()``, + and the flag stays cleared so a second ``agent.close()`` does not + re-attempt the close (the flag is dropped BEFORE the close call — + the documented-idempotency ordering).""" + attempts: list[int] = [] + + class _Raising(_RecordingDB): + def close(self): + attempts.append(1) + raise RuntimeError("disk gone") + + db = _Raising() + agent = _bare_agent(_session_db=db, _owns_session_db=True) + + agent.close() # must not raise + agent.close() # flag already cleared — no second attempt + + assert attempts == [1] + assert getattr(agent, "_owns_session_db") is False + + def test_close_still_ends_the_session_row_before_closing(): """Ordering matters: the row is finalized THROUGH the handle we then close.""" calls: list[str] = []