From be14a4bee3d5c134fb903b4aabc4e8a7e89c8ab6 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] 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 6f89ed237d..df98acc856 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 84ee8849f7..63824c6eac 100644 --- a/run_agent.py +++ b/run_agent.py @@ -613,6 +613,9 @@ class AIAgent: 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) @@ -4340,15 +4343,33 @@ class AIAgent: # 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 0000000000..47a5aa413e --- /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 d90024557a..c4d5a6ae7c 100644 --- a/tui_gateway/compute_host.py +++ b/tui_gateway/compute_host.py @@ -9,6 +9,7 @@ from __future__ import annotations import argparse import concurrent.futures +import contextlib import json import os import signal @@ -540,6 +541,7 @@ class ComputeHost: 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 @@ class ComputeHost: 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 @@ class ComputeHost: 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 42fb9bae32..214e84c425 100644 --- a/tui_gateway/methods_session.py +++ b/tui_gateway/methods_session.py @@ -690,6 +690,22 @@ def _(rid, params: 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: @@ -2786,6 +2802,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 @@ -2795,11 +2815,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 ) @@ -2836,6 +2859,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) @@ -2847,6 +2877,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 ed8f1dcc0c..462b89ff24 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 _start_agent_build(sid: str, session: dict) -> 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 _start_agent_build(sid: str, session: dict) -> 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 _start_agent_build(sid: str, session: dict) -> 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)