diff --git a/agent/auxiliary_client.py b/agent/auxiliary_client.py index e6b26493db..39daa4deed 100644 --- a/agent/auxiliary_client.py +++ b/agent/auxiliary_client.py @@ -7523,17 +7523,71 @@ def _force_close_async_httpx(client: Any) -> None: pass -def _close_cached_client(client: Any) -> None: - """Apply the canonical best-effort close policy to one cached client.""" +def _schedule_async_close(close_result: Any, client: Any) -> None: + """Finish an async close without leaking an unawaited coroutine.""" + async def _await_close() -> None: + try: + await close_result + except Exception: + pass + finally: + _force_close_async_httpx(client) + + runner = _await_close() + try: + import asyncio as _aio + + try: + loop = _aio.get_running_loop() + except RuntimeError: + _aio.run(runner) + else: + task = loop.create_task(runner) + + def _consume(completed_task) -> None: + try: + completed_task.exception() + except BaseException: + pass + + task.add_done_callback(_consume) + runner = None + except Exception: + if runner is not None: + try: + runner.close() + except Exception: + pass + _force_close_async_httpx(client) + + +def _close_cached_client(client: Any, *, close_async: bool = False) -> None: + """Close one cached client, awaiting async transports only when safe.""" if client is None: return - _force_close_async_httpx(client) + close_fn = getattr(client, "close", None) + if not callable(close_fn): + _force_close_async_httpx(client) + return try: - close_fn = getattr(client, "close", None) - if callable(close_fn) and not inspect.iscoroutinefunction(close_fn): - close_fn() + close_result = close_fn() except Exception: - pass + _force_close_async_httpx(client) + return + if inspect.isawaitable(close_result): + if close_async: + _schedule_async_close(close_result, client) + else: + # Do not await a client owned by another live event loop. + # Closing the coroutine avoids an unawaited-coroutine warning; + # the transport is still neutered for safe eventual GC. + try: + close_result.close() + except Exception: + pass + _force_close_async_httpx(client) + return + _force_close_async_httpx(client) def shutdown_cached_clients() -> None: @@ -7541,14 +7595,34 @@ def shutdown_cached_clients() -> None: Call this during CLI shutdown, *before* the event loop is closed, to avoid ``AsyncHttpxClientWrapper.__del__`` raising on a dead loop. + + Snapshot and clear the cache under the lock, then close transports outside + it. Async transport shutdown may block while an owner loop drains; holding + the global cache lock during that wait stalls unrelated auxiliary callers + and can turn teardown into a process-wide lock convoy. """ with _client_cache_lock: - for key, entry in list(_client_cache.items()): - client = entry[0] - if client is None: - continue - _close_cached_client(client) + clients = [ + (entry[0], entry[2]) + for entry in _client_cache.values() + if entry[0] is not None + ] _client_cache.clear() + try: + import asyncio as _aio + + running_loop = _aio.get_running_loop() + except RuntimeError: + running_loop = None + for client, owner_loop in clients: + # A live foreign loop owns its async transport. Calling its coroutine + # on this thread can bind/close sockets from the wrong loop; neuter it + # and let that owner finish teardown. Closed loops are safe to drain + # locally, and the current loop can await its own client. + close_async = owner_loop is not None and ( + owner_loop.is_closed() or owner_loop is running_loop + ) + _close_cached_client(client, close_async=close_async) def cleanup_stale_async_clients() -> None: @@ -7559,15 +7633,18 @@ def cleanup_stale_async_clients() -> None: This is defense-in-depth — the primary fix is ``neuter_async_httpx_del`` which disables ``__del__`` entirely. """ + stale_clients = [] with _client_cache_lock: stale_keys = [] for key, entry in _client_cache.items(): client, _default, cached_loop = entry if cached_loop is not None and cached_loop.is_closed(): - _force_close_async_httpx(client) stale_keys.append(key) + stale_clients.append(client) for key in stale_keys: del _client_cache[key] + for client in stale_clients: + _close_cached_client(client, close_async=True) def _is_openrouter_client(client: Any) -> bool: @@ -7660,7 +7737,12 @@ def _get_cached_client( effective = _compat_model(cached_client, model, cached_default) return cached_client, effective # Stale — evict and fall through to create a new client. - _force_close_async_httpx(cached_client) + # Only a client whose owner loop is closed may be awaited from + # this thread; a live foreign loop remains force-neutered. + owner_loop_closed = ( + cached_loop is not None and cached_loop.is_closed() + ) + _close_cached_client(cached_client, close_async=owner_loop_closed) del _client_cache[cache_key] else: effective = _compat_model(cached_client, model, cached_default) @@ -7710,7 +7792,7 @@ def _get_cached_client( client, default_model, _ = _client_cache[cache_key] # This concurrently built loser was never exposed to a caller, # so it is safe to close immediately. - _close_cached_client(built_client) + _close_cached_client(built_client, close_async=async_mode) return client, model or default_model diff --git a/agent/trace_upload.py b/agent/trace_upload.py index 404d9be70b..d2c97bdc7f 100644 --- a/agent/trace_upload.py +++ b/agent/trace_upload.py @@ -336,10 +336,16 @@ def load_session_messages( """ from hermes_state import SessionDB db = SessionDB(db_path=db_path) if db_path else SessionDB() - resolved = db.resolve_session_id(session_id) or session_id - meta = db.get_session(resolved) or {} - messages = db.get_messages_as_conversation(resolved) - return messages, meta + try: + resolved = db.resolve_session_id(session_id) or session_id + meta = db.get_session(resolved) or {} + messages = db.get_messages_as_conversation(resolved) + return messages, meta + finally: + try: + db.close() + except Exception: + logger.debug("Failed to close trace-upload SessionDB", exc_info=True) def upload_session_trace( diff --git a/cli.py b/cli.py index 8fe7720d2b..eb7c0da950 100644 --- a/cli.py +++ b/cli.py @@ -12510,10 +12510,12 @@ class HermesCLI(CLIAgentSetupMixin, CLICommandsMixin, CLIBillingMixin): from agent.insights import InsightsEngine db = SessionDB() - engine = InsightsEngine(db) - report = engine.generate(days=days, source=source) - print(engine.format_terminal(report)) - db.close() + try: + engine = InsightsEngine(db) + report = engine.generate(days=days, source=source) + print(engine.format_terminal(report)) + finally: + db.close() except Exception as e: print(f" Error generating insights: {e}") diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index 13c59ab0de..841d7609fe 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -1437,6 +1437,9 @@ class APIServerAdapter(BasePlatformAdapter): # in-flight run by run_id. self._run_approval_sessions: Dict[str, str] = {} self._session_db: Optional[Any] = None # Lazy-init SessionDB for session continuity + self._session_dbs: Dict[str, Any] = {} + self._session_db_cache_lock = threading.Lock() + self._session_db_cache_closed = False # Last-known-good resolved model per session (keyed by gateway_session_key # ONLY — never session_id, which rotates/is ephemeral for one-off API # server requests; "*" is the process-wide fallback), mirroring @@ -2182,15 +2185,29 @@ class APIServerAdapter(BasePlatformAdapter): from hermes_state import SessionDB key = str(home) - cache = getattr(self, "_session_dbs", None) - if cache is None: - cache = {} - self._session_dbs = cache - db = cache.get(key) - if db is None: - db = SessionDB(db_path=home / "state.db") - cache[key] = db - return db + with self._session_db_cache_lock: + if self._session_db_cache_closed: + return None + db = self._session_dbs.get(key) + if db is None: + db = SessionDB(db_path=home / "state.db") + self._session_dbs[key] = db + return db + + def _close_cached_session_dbs(self) -> None: + """Close SessionDB handles owned by this adapter's profile cache.""" + with self._session_db_cache_lock: + self._session_db_cache_closed = True + cached = list(self._session_dbs.values()) + self._session_dbs.clear() + shared_db = getattr(self, "_session_db", None) + for db in cached: + if db is shared_db: + continue + try: + db.close() + except Exception: + logger.debug("Failed to close API-server SessionDB", exc_info=True) def _ensure_session_db(self): """Lazily initialise and return the SessionDB for the active profile home. @@ -2232,15 +2249,17 @@ class APIServerAdapter(BasePlatformAdapter): home = get_hermes_home() key = str(home) - cache = getattr(self, "_session_dbs", None) - if cache is not None and cache.get(key) is not None: - return cache[key] + with self._session_db_cache_lock: + cached = self._session_dbs.get(key) + if cached is not None: + return cached if self._session_db_lock is None: self._session_db_lock = asyncio.Lock() async with self._session_db_lock: - cache = getattr(self, "_session_dbs", None) - if cache is not None and cache.get(key) is not None: - return cache[key] + with self._session_db_cache_lock: + cached = self._session_dbs.get(key) + if cached is not None: + return cached return await asyncio.to_thread(self._open_and_cache_session_db, home) except Exception as e: logger.debug("SessionDB unavailable for API server: %s", e) @@ -7368,6 +7387,9 @@ class APIServerAdapter(BasePlatformAdapter): logger.warning("[%s] aiohttp not installed", self.name) return False + with self._session_db_cache_lock: + self._session_db_cache_closed = False + if not self._api_key_passes_startup_guard(): # A rejected API_SERVER_KEY is a configuration error, not a # transient blip — the key will not become valid on its own. A @@ -7538,13 +7560,16 @@ class APIServerAdapter(BasePlatformAdapter): logger.debug( "Failed to close response store for %s", self.name, exc_info=True, ) - if self._site: - await self._site.stop() - self._site = None - if self._runner: - await self._runner.cleanup() - self._runner = None - self._app = None + try: + if self._site: + await self._site.stop() + self._site = None + if self._runner: + await self._runner.cleanup() + self._runner = None + finally: + self._close_cached_session_dbs() + self._app = None logger.info("[%s] API server stopped", self.name) async def send( diff --git a/gateway/shutdown_flush.py b/gateway/shutdown_flush.py index 14b04cd722..a727dd975a 100644 --- a/gateway/shutdown_flush.py +++ b/gateway/shutdown_flush.py @@ -316,6 +316,14 @@ def recover_pending_to_db( session_db = SessionDB() own_db = True + def _close_owned_db() -> None: + if not own_db: + return + try: + session_db.close() + except Exception: + pass + recovered = 0 for path in flush_files: try: @@ -389,6 +397,10 @@ def recover_pending_to_db( ) recovered += 1 path.unlink(missing_ok=True) + except BaseException: + # Shutdown cancellation/interrupt must not strand an owned DB. + _close_owned_db() + raise except Exception as exc: logger.warning( "Failed to recover pending message from %s: %s", @@ -396,11 +408,7 @@ def recover_pending_to_db( ) # Leave the file for next startup retry. - if own_db: - try: - session_db.close() - except Exception: - pass + _close_owned_db() if recovered: logger.info( diff --git a/hermes_cli/main.py b/hermes_cli/main.py index 08c5014f8c..58eac00527 100644 --- a/hermes_cli/main.py +++ b/hermes_cli/main.py @@ -1558,6 +1558,7 @@ def _resolve_session_by_name_or_id(name_or_id: str) -> Optional[str]: from an exit summary printed before the bug fix, or from notes) get resumed at the live tip instead of a stale parent with no messages. """ + db = None try: from hermes_state import SessionDB @@ -1580,10 +1581,15 @@ def _resolve_session_by_name_or_id(name_or_id: str) -> Optional[str]: except Exception: pass - db.close() return resolved_id except Exception: pass + finally: + if db is not None: + try: + db.close() + except Exception: + pass return None @@ -2681,10 +2687,12 @@ def cmd_chat(args): and not getattr(args, "no_restore_cwd", False) and not getattr(args, "worktree", False) ): + _resume_db = None try: from hermes_state import SessionDB - _saved_cwd = ((SessionDB().get_session(args.resume) or {}).get("cwd") or "").strip() + _resume_db = SessionDB() + _saved_cwd = ((_resume_db.get_session(args.resume) or {}).get("cwd") or "").strip() if _saved_cwd and not os.path.isdir(_saved_cwd): print(f"⚠ session's recorded dir is gone ({_saved_cwd}); staying in {os.getcwd()}") elif _saved_cwd and os.path.realpath(_saved_cwd) != os.path.realpath(os.getcwd()): @@ -2692,6 +2700,12 @@ def cmd_chat(args): print(f"↪ restored workspace dir: {_saved_cwd}") except Exception: pass # never let cwd-restore break a resume + finally: + if _resume_db is not None: + try: + _resume_db.close() + except Exception: + pass # xAI retirement warning — one-shot, non-blocking, never fails startup try: @@ -11472,6 +11486,7 @@ def cmd_tools(args): def cmd_insights(args): + db = None try: from hermes_state import SessionDB from agent.insights import InsightsEngine @@ -11480,9 +11495,14 @@ def cmd_insights(args): engine = InsightsEngine(db) report = engine.generate(days=args.days, source=args.source) print(engine.format_terminal(report)) - db.close() except Exception as e: print(f"Error generating insights: {e}") + finally: + if db is not None: + try: + db.close() + except Exception: + pass def cmd_monitoring(args): diff --git a/hermes_state.py b/hermes_state.py index 2b5a66de7c..0e7a094c4c 100644 --- a/hermes_state.py +++ b/hermes_state.py @@ -2836,6 +2836,16 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) data["system_prompt"] = resolved return data + @staticmethod + def _close_connection_quietly(conn: Optional[sqlite3.Connection]) -> None: + """Close a partially initialized connection without masking its error.""" + if conn is None: + return + try: + conn.close() + except Exception: + logger.debug("Could not close a SessionDB connection", exc_info=True) + def __init__(self, db_path: Path = None, read_only: bool = False): self.db_path = db_path or _default_db_path() # Fail hard (before any connection/pragma/mkdir) if a pytest-context @@ -2933,6 +2943,7 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) self._token_writer_thread: Optional[threading.Thread] = None self._token_writer_stop = False self._token_writer_busy = False + initialization_complete = False try: if read_only: # Read-only attach for cross-profile aggregation: SELECT-only, @@ -2957,8 +2968,9 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) # only so read-only search keeps its FTS and trigram paths. # Close the connection on ANY probe failure (e.g. malformed # schema raises DatabaseError, not the OperationalError the - # probe handles): the outer except re-raises without cleanup, - # and a leaked tracked connection blocks _backup_db_file's + # probe handles). The constructor's outer finally also covers + # failures before this probe and BaseException paths, so a + # leaked tracked connection cannot block _backup_db_file's # raw-copy for the rest of the process — the writable heal # that follows would then repair WITHOUT its forensic backup. try: @@ -2982,6 +2994,7 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) except Exception: pass raise + initialization_complete = True return self.db_path.parent.mkdir(parents=True, exist_ok=True) @@ -3115,6 +3128,7 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) # racing session lifecycle and the surprise disk/latency cost on # an unattended open. (An interrupted optimize resumes when the # user re-runs the command.) + initialization_complete = True except Exception as exc: # Capture the cause so /resume and friends can surface WHY the # session DB is unavailable instead of a bare "Session database @@ -3130,6 +3144,10 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) # ``hermes_state._set_last_init_error(None)`` explicitly. _set_last_init_error(f"{type(exc).__name__}: {exc}") raise + finally: + if not initialization_complete: + conn, self._conn = self._conn, None + self._close_connection_quietly(conn) # ── Read-path split ── @@ -4001,8 +4019,8 @@ class SessionDB(SessionSearchMixin, SessionSchemaMixin, SessionPortabilityMixin) "WAL checkpoint (PASSIVE) at close failed: %s", exc, ) - self._conn.close() - self._conn = None + conn, self._conn = self._conn, None + self._close_connection_quietly(conn) # ── Chunked FTS rebuild engine (v23 opt-in optimize) ── # diff --git a/mcp_serve.py b/mcp_serve.py index 23f248164b..76fcf800b8 100644 --- a/mcp_serve.py +++ b/mcp_serve.py @@ -79,6 +79,22 @@ def _get_session_db(): return None +def _load_session_messages(session_id: str): + """Read one session and close the temporary database handle.""" + db = _get_session_db() + if db is None: + return None, "Session database unavailable" + try: + return db.get_messages(session_id), None + except Exception as e: + return None, f"Failed to read messages: {e}" + finally: + try: + db.close() + except Exception: + logger.debug("Failed to close MCP SessionDB", exc_info=True) + + def _load_sessions_index() -> dict: """Load the gateway session routing index. @@ -448,6 +464,18 @@ class EventBridge: self._new_event.set() def _establish_baseline(self) -> None: + db = _get_session_db() + if not db: + return + try: + self._establish_baseline_with_db(db) + finally: + try: + db.close() + except Exception: + logger.debug("Failed to close MCP baseline SessionDB", exc_info=True) + + def _establish_baseline_with_db(self, db) -> None: """Record the latest per-session message timestamp and the current state.db mtime WITHOUT emitting events, so startup does not replay history (#13414). @@ -457,9 +485,6 @@ class EventBridge: last_seen=0.0 in _poll_once, so a brand-new conversation's first message is still delivered on its state.db-change tick. """ - db = _get_session_db() - if not db: - return try: from hermes_constants import get_hermes_home db_file = get_hermes_home() / "state.db" @@ -486,7 +511,6 @@ class EventBridge: latest = max(all_ts) if latest > 0.0: self._last_poll_timestamps[session_key] = latest - def _poll_loop(self): """Background loop: poll SessionDB for new messages.""" db = _get_session_db() @@ -494,12 +518,18 @@ class EventBridge: logger.warning("EventBridge: SessionDB unavailable, event polling disabled") return - while self._running: + try: + while self._running: + try: + self._poll_once(db) + except Exception as e: + logger.debug("EventBridge poll error: %s", e) + time.sleep(POLL_INTERVAL) + finally: try: - self._poll_once(db) - except Exception as e: - logger.debug("EventBridge poll error: %s", e) - time.sleep(POLL_INTERVAL) + db.close() + except Exception: + logger.debug("Failed to close MCP polling SessionDB", exc_info=True) def _poll_once(self, db): """Check for new messages across all sessions. @@ -722,14 +752,9 @@ def create_mcp_server(event_bridge: Optional[EventBridge] = None) -> "FastMCP": if not session_id: return json.dumps({"error": "No session ID for this conversation"}) - db = _get_session_db() - if not db: - return json.dumps({"error": "Session database unavailable"}) - - try: - all_messages = db.get_messages(session_id) - except Exception as e: - return json.dumps({"error": f"Failed to read messages: {e}"}) + all_messages, error = _load_session_messages(session_id) + if error: + return json.dumps({"error": error}) filtered = [] for msg in all_messages: @@ -778,14 +803,9 @@ def create_mcp_server(event_bridge: Optional[EventBridge] = None) -> "FastMCP": if not session_id: return json.dumps({"error": "No session ID for this conversation"}) - db = _get_session_db() - if not db: - return json.dumps({"error": "Session database unavailable"}) - - try: - all_messages = db.get_messages(session_id) - except Exception as e: - return json.dumps({"error": f"Failed to read messages: {e}"}) + all_messages, error = _load_session_messages(session_id) + if error: + return json.dumps({"error": error}) # Find the target message target_msg = None diff --git a/plugins/memory/retaindb/__init__.py b/plugins/memory/retaindb/__init__.py index f65d5018be..b4ca845249 100644 --- a/plugins/memory/retaindb/__init__.py +++ b/plugins/memory/retaindb/__init__.py @@ -364,6 +364,10 @@ class _WriteQueue: self._db_path.parent.mkdir(parents=True, exist_ok=True) # Thread-local connection cache — one connection per thread, reused. self._local = threading.local() + self._connections: set[sqlite3.Connection] = set() + self._connections_lock = threading.Lock() + self._shutdown_lock = threading.Lock() + self._shutdown = False self._init_db() self._thread.start() # Replay any rows left from a previous crash @@ -374,11 +378,38 @@ class _WriteQueue: """Return a cached connection for the current thread.""" conn = getattr(self._local, "conn", None) if conn is None: - conn = sqlite3.connect(str(self._db_path), timeout=30) + conn = sqlite3.connect( + str(self._db_path), timeout=30, check_same_thread=False + ) conn.row_factory = sqlite3.Row self._local.conn = conn + with self._connections_lock: + self._connections.add(conn) return conn + def _close_thread_conn(self) -> None: + conn = getattr(self._local, "conn", None) + if conn is None: + return + self._local.conn = None + with self._connections_lock: + self._connections.discard(conn) + try: + conn.close() + except Exception: + pass + + def _close_all_connections(self) -> None: + """Close tracked connections left by short-lived worker threads.""" + with self._connections_lock: + connections = list(self._connections) + self._connections.clear() + for conn in connections: + try: + conn.close() + except Exception: + pass + def _init_db(self) -> None: conn = self._get_conn() conn.execute("""CREATE TABLE IF NOT EXISTS pending ( @@ -394,14 +425,17 @@ class _WriteQueue: def enqueue(self, user_id: str, session_id: str, messages: list) -> None: now = datetime.now(timezone.utc).isoformat() - conn = self._get_conn() - cur = conn.execute( - "INSERT INTO pending (user_id, session_id, messages_json, created_at) VALUES (?,?,?,?)", - (user_id, session_id, json.dumps(messages, ensure_ascii=False), now), - ) - row_id = cur.lastrowid - conn.commit() - self._q.put((row_id, user_id, session_id, messages)) + with self._shutdown_lock: + if self._shutdown: + return + conn = self._get_conn() + cur = conn.execute( + "INSERT INTO pending (user_id, session_id, messages_json, created_at) VALUES (?,?,?,?)", + (user_id, session_id, json.dumps(messages, ensure_ascii=False), now), + ) + row_id = cur.lastrowid + conn.commit() + self._q.put((row_id, user_id, session_id, messages)) def _flush_row(self, row_id: int, user_id: str, session_id: str, messages: list) -> None: try: @@ -417,20 +451,35 @@ class _WriteQueue: time.sleep(2) def _loop(self) -> None: - while True: - try: - item = self._q.get(timeout=5) - if item is _ASYNC_SHUTDOWN: - break - self._flush_row(*item) - except queue.Empty: - continue - except Exception as exc: - logger.error("RetainDB writer error: %s", exc) + try: + while True: + try: + item = self._q.get(timeout=5) + if item is _ASYNC_SHUTDOWN: + break + self._flush_row(*item) + except queue.Empty: + continue + except Exception as exc: + logger.error("RetainDB writer error: %s", exc) + finally: + # sqlite3 connections must close on their owning thread. + self._close_thread_conn() def shutdown(self) -> None: - self._q.put(_ASYNC_SHUTDOWN) + with self._shutdown_lock: + if self._shutdown: + return + self._shutdown = True + self._q.put(_ASYNC_SHUTDOWN) + # Caller thread owns connection opened by _init_db/_pending_rows. + self._close_thread_conn() self._thread.join(timeout=10) + if not self._thread.is_alive(): + # MemoryManager's executor may have opened a connection on a + # worker that has already exited; check_same_thread=False lets + # shutdown close that tracked handle deterministically. + self._close_all_connections() # --------------------------------------------------------------------------- @@ -581,6 +630,9 @@ class RetainDBMemoryProvider(MemoryProvider): # Prevents thread accumulation if turns fire faster than prefetches complete. for t in self._prefetch_threads: t.join(timeout=2.0) + if any(t.is_alive() for t in self._prefetch_threads): + logger.debug("RetainDB prefetch still running; skipping new batch") + return threads = [ threading.Thread(target=self._prefetch_context, args=(query,), name="retaindb-ctx", daemon=True), threading.Thread(target=self._prefetch_dialectic, args=(query,), name="retaindb-dialectic", daemon=True), @@ -795,8 +847,12 @@ class RetainDBMemoryProvider(MemoryProvider): def shutdown(self) -> None: for t in self._prefetch_threads: t.join(timeout=3.0) - if self._queue: - self._queue.shutdown() + self._prefetch_threads = [] + queue_obj = self._queue + self._queue = None + if queue_obj: + queue_obj.shutdown() + self._client = None def register(ctx) -> None: diff --git a/run_agent.py b/run_agent.py index bc344b2cf0..583a3881c9 100644 --- a/run_agent.py +++ b/run_agent.py @@ -4208,13 +4208,14 @@ class AIAgent: ) def shutdown_memory_provider(self, messages: list = None) -> None: - """Shut down the memory provider and context engine — call at actual session boundaries. + """Shut down the memory provider and context engine at session end. - This calls on_session_end() then shutdown_all() on the memory - manager, and on_session_end() on the context engine. - NOT called per-turn — only at CLI exit, /reset, gateway - session expiry, etc. + Idempotent: gateway cleanup and AIAgent.close() may share this + ownership boundary. """ + if getattr(self, "_memory_provider_shutdown", False): + return + self._memory_provider_shutdown = True if self._memory_manager: try: self._memory_manager.on_session_end(messages or []) @@ -4399,6 +4400,17 @@ class AIAgent: Safe to call multiple times (idempotent). Each cleanup step is independently guarded so a failure in one does not prevent the rest. """ + # AIAgent.close() is the hard owner boundary. Gateway cleanup may + # call shutdown_memory_provider() first; its idempotence prevents + # duplicate extraction while direct callers cannot skip provider close. + try: + session_messages = getattr(self, "_session_messages", None) + self.shutdown_memory_provider( + session_messages if isinstance(session_messages, list) else None + ) + except Exception: + pass + task_id = getattr(self, "session_id", None) or "" # 1. Kill background processes for this task diff --git a/tests/agent/test_trace_upload.py b/tests/agent/test_trace_upload.py index c77ecba5e0..a4f4d4ebd6 100644 --- a/tests/agent/test_trace_upload.py +++ b/tests/agent/test_trace_upload.py @@ -13,6 +13,7 @@ import pytest from agent import trace_upload from agent.trace_upload import ( build_trace_jsonl, + load_session_messages, upload_session_trace, _resolve_hf_token, _do_upload, @@ -97,6 +98,17 @@ def test_converter_keeps_secrets_when_redact_disabled(): assert secret in jsonl +def test_load_session_messages_closes_database_on_failure(monkeypatch): + db = MagicMock() + db.resolve_session_id.side_effect = RuntimeError("read failed") + monkeypatch.setattr("hermes_state.SessionDB", lambda: db) + + with pytest.raises(RuntimeError, match="read failed"): + load_session_messages("s1") + + db.close.assert_called_once() + + diff --git a/tests/cli/test_cli_insights_command.py b/tests/cli/test_cli_insights_command.py index 66c3c73b5d..e6cd0ed2eb 100644 --- a/tests/cli/test_cli_insights_command.py +++ b/tests/cli/test_cli_insights_command.py @@ -1,6 +1,8 @@ from unittest.mock import MagicMock, patch +from types import SimpleNamespace from cli import HermesCLI +from hermes_cli.main import cmd_insights class _InsightsEngineStub: @@ -41,3 +43,13 @@ def test_cli_insights_keeps_days_flag_and_source(capsys): assert calls == [{"days": 14, "source": "discord"}] db.close.assert_called_once() assert "days=14 source=discord" in capsys.readouterr().out + + +def test_subcommand_insights_closes_database_when_generation_fails(capsys): + db = MagicMock() + with patch("hermes_state.SessionDB", return_value=db), \ + patch("agent.insights.InsightsEngine", side_effect=RuntimeError("boom")): + cmd_insights(SimpleNamespace(days=30, source=None)) + + db.close.assert_called_once() + assert "Error generating insights: boom" in capsys.readouterr().out diff --git a/tests/gateway/test_platform_reconnect_fd_leak.py b/tests/gateway/test_platform_reconnect_fd_leak.py index 8ab9343cc4..50fd14ef99 100644 --- a/tests/gateway/test_platform_reconnect_fd_leak.py +++ b/tests/gateway/test_platform_reconnect_fd_leak.py @@ -21,6 +21,7 @@ this file would have caught the regression and now pins the fix. from __future__ import annotations import asyncio +import threading import time from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -254,9 +255,43 @@ class TestAPIServerDisconnectClosesResponseStore: adapter._runner = None adapter._app = None adapter._response_store = store + adapter._session_dbs = {} + adapter._session_db_cache_lock = threading.Lock() + adapter._session_db_cache_closed = False adapter.platform = Platform.API_SERVER return adapter + @pytest.mark.asyncio + async def test_disconnect_closes_cached_session_dbs(self): + """Disconnect must release per-profile SessionDB cache handles.""" + store = MagicMock() + adapter = self._build_adapter_with_store(store) + first_db = MagicMock() + second_db = MagicMock() + adapter._session_dbs = {"default": first_db, "work": second_db} + + await adapter.disconnect() + + first_db.close.assert_called_once_with() + second_db.close.assert_called_once_with() + assert adapter._session_dbs == {} + + @pytest.mark.asyncio + async def test_disconnect_closes_cached_session_dbs_when_runner_cleanup_fails( + self, + ): + """Runner teardown errors must not strand cached DB handles.""" + adapter = self._build_adapter_with_store(MagicMock()) + cached_db = MagicMock() + adapter._session_dbs = {"default": cached_db} + adapter._runner = MagicMock() + adapter._runner.cleanup = AsyncMock(side_effect=RuntimeError("boom")) + + with pytest.raises(RuntimeError, match="boom"): + await adapter.disconnect() + + cached_db.close.assert_called_once_with() + @pytest.mark.asyncio async def test_disconnect_closes_response_store(self, tmp_path): """Closing the adapter's ResponseStore releases its SQLite connection. diff --git a/tests/gateway/test_shutdown_flush.py b/tests/gateway/test_shutdown_flush.py index 29cd633f95..fe67ab7168 100644 --- a/tests/gateway/test_shutdown_flush.py +++ b/tests/gateway/test_shutdown_flush.py @@ -97,6 +97,42 @@ def test_recover_inserts_via_append_message_and_deletes_file(tmp_path, monkeypat assert not flush_file.exists() +def test_recover_closes_owned_db_when_unexpected_exception_escapes( + tmp_path, monkeypatch +): + """Owned SessionDB must close even when recovery is interrupted.""" + flush_dir = _make_flush_dir(tmp_path) + monkeypatch.setattr( + "gateway.shutdown_flush._get_flush_dir", lambda: flush_dir + ) + (flush_dir / "pending.json").write_text( + json.dumps( + { + "session_key": "agent:main:telegram:123", + "data": {"text": "message", "session_id": "sid"}, + } + ), + encoding="utf-8", + ) + + class InterruptingDB: + closed = False + + def append_message(self, **_kwargs): + raise KeyboardInterrupt + + def close(self): + self.closed = True + + db = InterruptingDB() + monkeypatch.setattr("hermes_state.SessionDB", lambda: db) + + with pytest.raises(KeyboardInterrupt): + recover_pending_to_db() + + assert db.closed is True + + def test_serialise_object_with_text(): obj = MagicMock() obj.text = "msg" diff --git a/tests/plugins/memory/test_retaindb_provider.py b/tests/plugins/memory/test_retaindb_provider.py index 0372edaefd..ea50a7ea19 100644 --- a/tests/plugins/memory/test_retaindb_provider.py +++ b/tests/plugins/memory/test_retaindb_provider.py @@ -1,12 +1,65 @@ from __future__ import annotations +import sqlite3 from unittest.mock import MagicMock import agent.file_safety as fs +import pytest + +import plugins.memory.retaindb as retaindb from plugins.memory.retaindb import RetainDBMemoryProvider +def test_write_queue_closes_owner_connection(tmp_path): + queue = retaindb._WriteQueue(object(), tmp_path / "retaindb.db") + owner_conn = queue._local.conn + worker = retaindb.threading.Thread(target=queue._get_conn) + worker.start() + worker.join() + queue.shutdown() + assert not queue._connections + with pytest.raises(sqlite3.ProgrammingError): + owner_conn.execute("SELECT 1") + + +def test_write_queue_ignores_enqueue_after_shutdown(tmp_path): + queue = retaindb._WriteQueue(object(), tmp_path / "retaindb.db") + queue.shutdown() + + queue.enqueue("user", "session", []) + + assert not queue._connections + + +def test_prefetch_does_not_spawn_when_previous_batch_is_alive(monkeypatch): + provider = RetainDBMemoryProvider() + provider._client = object() + + class _RunningThread: + def join(self, timeout): + pass + + def is_alive(self): + return True + + previous = _RunningThread() + provider._prefetch_threads = [previous] + created = [] + + class _Thread: + def __init__(self, *args, **kwargs): + created.append((args, kwargs)) + + def start(self): + pass + + monkeypatch.setattr(retaindb.threading, "Thread", _Thread) + provider.queue_prefetch("query") + assert provider._prefetch_threads == [previous] + assert not created + + def test_upload_file_rejects_hermes_credential_store(tmp_path, monkeypatch): hermes_home = tmp_path / "hermes_home" hermes_home.mkdir() diff --git a/tests/run_agent/test_async_httpx_del_neuter.py b/tests/run_agent/test_async_httpx_del_neuter.py index 7070ff59c9..1583d94fbf 100644 --- a/tests/run_agent/test_async_httpx_del_neuter.py +++ b/tests/run_agent/test_async_httpx_del_neuter.py @@ -100,6 +100,7 @@ class TestCleanupStaleAsyncClients: try: cleanup_stale_async_clients() + mock_client.close.assert_called_once() with _client_cache_lock: assert key not in _client_cache, "Stale entry should be removed" finally: @@ -107,6 +108,102 @@ class TestCleanupStaleAsyncClients: with _client_cache_lock: _client_cache.pop(key, None) + def test_awaits_async_close_for_closed_loop(self): + from agent.auxiliary_client import ( + _client_cache, + _client_cache_lock, + cleanup_stale_async_clients, + ) + + class AsyncClient: + def __init__(self): + self._client = MagicMock() + self._client.is_closed = False + self.closed = False + + async def close(self): + self.closed = True + + loop = asyncio.new_event_loop() + loop.close() + client = AsyncClient() + key = ("test_async_close", True, "", "", "", (), False) + with _client_cache_lock: + _client_cache[key] = (client, "test-model", loop) + + try: + cleanup_stale_async_clients() + assert client.closed + finally: + with _client_cache_lock: + _client_cache.pop(key, None) + + + def test_shutdown_closes_outside_cache_lock(self): + from agent.auxiliary_client import ( + _client_cache, + _client_cache_lock, + shutdown_cached_clients, + ) + + lock_observations = [] + + class Client: + _client = None + + def close(self): + acquired = _client_cache_lock.acquire(blocking=False) + lock_observations.append(acquired) + if acquired: + _client_cache_lock.release() + + key = ("test_shutdown_lock", False, "", "", "", (), False) + with _client_cache_lock: + previous = dict(_client_cache) + _client_cache.clear() + _client_cache[key] = (Client(), "test-model", None) + + try: + shutdown_cached_clients() + finally: + with _client_cache_lock: + _client_cache.clear() + _client_cache.update(previous) + + assert lock_observations == [True] + + def test_shutdown_does_not_await_live_foreign_loop_client(self): + from agent.auxiliary_client import ( + _client_cache, + _client_cache_lock, + shutdown_cached_clients, + ) + + owner_loop = asyncio.new_event_loop() + + class Client: + def __init__(self): + self.awaited = False + + async def close(self): + self.awaited = True + + client = Client() + key = ("test_shutdown_foreign_loop", True, "", "", "", (), False) + with _client_cache_lock: + previous = dict(_client_cache) + _client_cache.clear() + _client_cache[key] = (client, "test-model", owner_loop) + + try: + shutdown_cached_clients() + assert client.awaited is False + finally: + owner_loop.close() + with _client_cache_lock: + _client_cache.clear() + _client_cache.update(previous) + def test_keeps_live_entries(self): """Entries with an open loop should be preserved.""" from agent.auxiliary_client import ( diff --git a/tests/run_agent/test_memory_provider_init.py b/tests/run_agent/test_memory_provider_init.py index ff647a3ae6..05edbad509 100644 --- a/tests/run_agent/test_memory_provider_init.py +++ b/tests/run_agent/test_memory_provider_init.py @@ -25,6 +25,24 @@ class RecordingMemoryProvider: pass +def test_shutdown_memory_provider_is_idempotent(): + from unittest.mock import MagicMock + + from run_agent import AIAgent + + manager = MagicMock() + agent = object.__new__(AIAgent) + agent._memory_manager = manager + agent.context_compressor = None + agent.session_id = "session-1" + + agent.shutdown_memory_provider([{"role": "user", "content": "one"}]) + agent.shutdown_memory_provider([{"role": "user", "content": "two"}]) + + manager.on_session_end.assert_called_once() + manager.shutdown_all.assert_called_once() + + def test_blank_memory_provider_does_not_auto_enable_honcho(): """Blank memory.provider should remain opt-out even if Honcho fallback looks configured.""" cfg = {"memory": {"provider": ""}, "agent": {}} @@ -59,6 +77,22 @@ def test_blank_memory_provider_does_not_auto_enable_honcho(): save_config.assert_not_called() +def test_close_shuts_down_memory_provider(): + from unittest.mock import MagicMock + + from run_agent import AIAgent + + agent = object.__new__(AIAgent) + agent._memory_manager = MagicMock() + agent.context_compressor = None + agent.session_id = "" + agent._session_messages = [] + + agent.close() + + agent._memory_manager.shutdown_all.assert_called_once() + + def test_aiagent_forwards_user_id_alt_to_memory_provider(): provider = RecordingMemoryProvider() cfg = {"memory": {"provider": "recording"}, "agent": {}} diff --git a/tests/test_hermes_state.py b/tests/test_hermes_state.py index f7721090e3..07236504ae 100644 --- a/tests/test_hermes_state.py +++ b/tests/test_hermes_state.py @@ -3,6 +3,7 @@ import sqlite3 import time import json +import threading from pathlib import Path from unittest import mock @@ -103,6 +104,114 @@ def _no_fts_rebuild_throttle(monkeypatch): class TestConnectionLifecycle: + def test_failed_writable_open_does_not_leak_tracked_connection( + self, tmp_path, monkeypatch + ): + """A failed schema init must close the connection opened before it.""" + from hermes_cli.sqlite_safe_read import has_live_connection + + db_path = tmp_path / "state.db" + opened = [] + real_connect = hermes_state._connect_tracked_db + + def capture_connect(*args, **kwargs): + conn = real_connect(*args, **kwargs) + opened.append(conn) + return conn + + monkeypatch.setattr(hermes_state, "_connect_tracked_db", capture_connect) + monkeypatch.setattr( + SessionDB, + "_init_schema", + mock.Mock(side_effect=RuntimeError("schema init failed")), + ) + + try: + with pytest.raises(RuntimeError, match="schema init failed"): + SessionDB(db_path=db_path) + assert has_live_connection(db_path) is False + finally: + for conn in opened: + try: + conn.close() + except Exception: + pass + + def test_failed_wal_read_open_does_not_leak_tracked_connection( + self, tmp_path, monkeypatch + ): + """A post-open read setup failure must close its unregistered conn.""" + from hermes_cli import sqlite_safe_read + + db_path = tmp_path / "state.db" + db = SessionDB(db_path=db_path) + opened = [] + real_connect = hermes_state._connect_tracked_db + real_pragmas = hermes_state.apply_database_pragmas + + def capture_connect(*args, **kwargs): + conn = real_connect(*args, **kwargs) + opened.append(conn) + return conn + + def fail_pragmas(*args, **kwargs): + raise RuntimeError("read setup failed") + + monkeypatch.setattr(hermes_state, "_connect_tracked_db", capture_connect) + monkeypatch.setattr(hermes_state, "apply_database_pragmas", fail_pragmas) + before = dict(sqlite_safe_read._live_connections) + db._wal_active = True + + try: + with pytest.raises(RuntimeError, match="read setup failed"): + db._get_read_conn() + assert sqlite_safe_read._live_connections == before + finally: + monkeypatch.setattr( + hermes_state, "apply_database_pragmas", real_pragmas + ) + for conn in opened: + try: + conn.close() + except Exception: + pass + db.close() + + def test_close_closes_wal_read_connection_created_on_worker_thread( + self, tmp_path + ): + """SessionDB.close() must drain read conns created by other threads.""" + from hermes_cli.sqlite_safe_read import has_live_connection + + db_path = tmp_path / "state.db" + db = SessionDB(db_path=db_path) + db._wal_active = True + opened = threading.Event() + release = threading.Event() + errors = [] + + def open_read_connection(): + try: + assert db._get_read_conn() is not None + opened.set() + release.wait(timeout=10) + except BaseException as exc: + errors.append(exc) + opened.set() + + worker = threading.Thread(target=open_read_connection) + worker.start() + assert opened.wait(timeout=10) + assert not errors + + db.close() + assert has_live_connection(db_path) is False + + release.set() + worker.join(timeout=10) + assert not worker.is_alive() + assert not errors + def test_read_only_close_never_requests_wal_checkpoint(self, tmp_path): db_path = tmp_path / "state.db" writable = SessionDB(db_path=db_path) diff --git a/tests/test_mcp_serve.py b/tests/test_mcp_serve.py index 6c58796577..d17aa4754e 100644 --- a/tests/test_mcp_serve.py +++ b/tests/test_mcp_serve.py @@ -272,6 +272,19 @@ class TestImports: class TestHelpers: + def test_load_session_messages_closes_database_on_error(self, monkeypatch): + import mcp_serve + + db = MagicMock() + db.get_messages.side_effect = RuntimeError("read failed") + monkeypatch.setattr(mcp_serve, "_get_session_db", lambda: db) + + messages, error = mcp_serve._load_session_messages("s1") + + assert messages is None + assert "read failed" in error + db.close.assert_called_once() + def test_get_sessions_dir(self, tmp_path): from mcp_serve import _get_sessions_dir result = _get_sessions_dir() diff --git a/tests/test_tui_gateway_server.py b/tests/test_tui_gateway_server.py index 43a03c577a..c144f09c59 100644 --- a/tests/test_tui_gateway_server.py +++ b/tests/test_tui_gateway_server.py @@ -10348,6 +10348,95 @@ def test_prompt_submit_merges_on_model_switch_marker(monkeypatch): server._sessions.pop("sid", None) +def test_prompt_submit_merges_on_personality_pivot_marker(monkeypatch): + """A personality pivot injected mid-turn must merge like a model switch. + + `/personality` applies immediately — there is no deferred queue for it the + way `pending_model_switch` defers a mid-turn model change — so choosing a + personality while a turn is running bumps `history_version` from the RPC + thread. The mid-turn reconciliation only recognized the model-switch + marker, so the pivot read as a genuine desync and the finished turn was + dropped from session history: the user saw the reply and it was never + stored (#82756). + """ + session_ref: dict[str, dict | None] = {"s": None} + + class _PivotAgent: + def run_conversation( + self, prompt, conversation_history=None, stream_callback=None, **_kwargs + ): + # Real injection point, mid-turn, exactly as the personality RPC + # would reach it from the other thread. + server._apply_personality_to_session( + "sid", session_ref["s"], "Answer tersely.", "terse" + ) + return { + "final_response": "agent reply", + "messages": list(conversation_history) + + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "agent reply"}, + ], + } + + class _ImmediateThread: + def __init__(self, target=None, daemon=None): + self._target = target + + def start(self): + self._target() + + server._sessions["sid"] = _session( + agent=_PivotAgent(), + history=[{"role": "user", "content": "hello"}], + ) + session_ref["s"] = server._sessions["sid"] + emits: list[tuple] = [] + try: + monkeypatch.setattr(server.threading, "Thread", _ImmediateThread) + monkeypatch.setattr(server, "_get_usage", lambda _a: {}) + monkeypatch.setattr(server, "render_message", lambda _t, _c: "") + monkeypatch.setattr(server, "_session_info", lambda *a, **k: {}) + monkeypatch.setattr(server, "_emit", lambda *a: emits.append(a)) + + resp = server.handle_request( + { + "id": "1", + "method": "prompt.submit", + "params": {"session_id": "sid", "text": "hi"}, + } + ) + assert resp.get("result"), f"got error: {resp.get('error')}" + + final_history = server._sessions["sid"]["history"] + + assistant_msgs = [ + e + for e in final_history + if isinstance(e, dict) + and e.get("role") == "assistant" + and e.get("content") == "agent reply" + ] + assert len(assistant_msgs) == 1, ( + "the personality pivot discarded the finished turn instead of " + f"merging it (got {len(assistant_msgs)} assistant replies)" + ) + + pivots = [ + e + for e in final_history + if isinstance(e, dict) and e.get("display_kind") == "personality_switch" + ] + assert len(pivots) == 1, f"expected exactly 1 pivot, got {len(pivots)}" + + complete_calls = [a for a in emits if a[0] == "message.complete"] + assert len(complete_calls) == 1 + _, _, payload = complete_calls[0] + assert "warning" not in payload, "merge path should not surface a warning" + finally: + server._sessions.pop("sid", None) + + def test_prompt_submit_sanitizes_bracketed_paste_before_agent(monkeypatch): """prompt.submit must sanitize corrupted user text before run_conversation.""" captured: dict[str, str] = {} diff --git a/tests/tools/test_react_to_message_tool.py b/tests/tools/test_react_to_message_tool.py new file mode 100644 index 0000000000..f7db04af6b --- /dev/null +++ b/tests/tools/test_react_to_message_tool.py @@ -0,0 +1,22 @@ +"""Ownership tests for desktop message reactions.""" + +from unittest.mock import MagicMock + +from tools import react_to_message_tool as reactions + + +def test_reaction_database_closes_when_write_fails(monkeypatch): + db = MagicMock() + db.latest_message_row_id.return_value = 42 + db.set_message_reaction.side_effect = RuntimeError("write failed") + monkeypatch.setattr(reactions, "_open_session_db", lambda: db) + monkeypatch.setattr( + reactions, + "get_session_env", + lambda _name, _default="": "session-1", + ) + + result = reactions.react_to_message_tool("👍") + + assert "write failed" in result + db.close.assert_called_once() diff --git a/tests/tools/test_session_search.py b/tests/tools/test_session_search.py index 8407c3eda7..5098127948 100644 --- a/tests/tools/test_session_search.py +++ b/tests/tools/test_session_search.py @@ -94,6 +94,50 @@ class TestFormatTimestamp: # ========================================================================= class TestBrowseShape: + def test_lazy_database_is_closed_after_search(self, monkeypatch): + class _DB: + closed = 0 + + def list_sessions_rich(self, **_kwargs): + return [] + + def close(self): + self.closed += 1 + + db = _DB() + monkeypatch.setattr("hermes_state.SessionDB", lambda: db) + + result = json.loads(session_search()) + + assert result["success"] is True + assert db.closed == 1 + + def test_cross_profile_database_is_closed_but_shared_database_is_not( + self, monkeypatch + ): + class _DB: + def __init__(self): + self.closed = 0 + + def list_sessions_rich(self, **_kwargs): + return [] + + def close(self): + self.closed += 1 + + shared_db = _DB() + profile_db = _DB() + monkeypatch.setattr( + "tools.session_search_tool._resolve_profile_db", + lambda _profile: profile_db, + ) + + result = json.loads(session_search(db=shared_db, profile="work")) + + assert result["success"] is True + assert profile_db.closed == 1 + assert shared_db.closed == 0 + def test_no_args_returns_recent_sessions(self, db): _seed_modpack_sessions(db) result = json.loads(session_search(db=db)) diff --git a/tools/react_to_message_tool.py b/tools/react_to_message_tool.py index a79bb25508..4fb759cdf0 100644 --- a/tools/react_to_message_tool.py +++ b/tools/react_to_message_tool.py @@ -31,20 +31,18 @@ def _open_session_db(): return None -def react_to_message_tool(emoji: str, message_row_id=None, messages_back=None) -> str: +def _react_to_message_with_db( + emoji: str, + message_row_id=None, + messages_back=None, + *, + db, + session_key: str, +) -> str: """Attach (or with an empty ``emoji`` retract) the agent's reaction.""" - emoji = (emoji or "").strip() - session_key = get_session_env("HERMES_SESSION_KEY", "") or get_session_env( - "HERMES_SESSION_ID", "" - ) - if not session_key: return tool_error("No active session — reactions need a persisted conversation.") - db = _open_session_db() - if db is None: - return tool_error("Session storage is unavailable.") - row_id = message_row_id target_role = "user" if row_id is None: @@ -89,6 +87,35 @@ def react_to_message_tool(emoji: str, message_row_id=None, messages_back=None) - ) +def react_to_message_tool(emoji: str, message_row_id=None, messages_back=None) -> str: + """Attach (or with an empty ``emoji`` retract) the agent's reaction.""" + emoji = (emoji or "").strip() + session_key = get_session_env("HERMES_SESSION_KEY", "") or get_session_env( + "HERMES_SESSION_ID", "" + ) + + if not session_key: + return tool_error("No active session — reactions need a persisted conversation.") + + db = _open_session_db() + if db is None: + return tool_error("Session storage is unavailable.") + + try: + return _react_to_message_with_db( + emoji, + message_row_id, + messages_back, + db=db, + session_key=session_key, + ) + finally: + try: + db.close() + except Exception: + pass + + def check_react_requirements() -> bool: """Opt-in feature flag — surface eligibility is the toolset's job. diff --git a/tools/session_search_tool.py b/tools/session_search_tool.py index ca3ebc970e..f69ecba32b 100644 --- a/tools/session_search_tool.py +++ b/tools/session_search_tool.py @@ -916,7 +916,7 @@ def _discover( return json.dumps(_final_payload, ensure_ascii=False) -def session_search( +def _session_search_impl( query: str = "", role_filter: str = None, limit: int = 3, @@ -930,6 +930,8 @@ def session_search( sort: str = None, # Cross-profile (any shape) profile: str = None, + *, + _owned_dbs: Optional[List[Any]] = None, ) -> str: """Single-shape tool. Mode inferred from which args are set. @@ -942,15 +944,6 @@ def session_search( ``@session:/`` link). Scroll wins over read/discovery when an anchor is set — the agent has asked for a specific slice. """ - if db is None: - try: - from hermes_state import SessionDB - db = SessionDB() - except Exception: - logging.debug("SessionDB unavailable for session_search", exc_info=True) - from hermes_state import format_session_db_unavailable - return tool_error(format_session_db_unavailable(), success=False) - # Normalise a raw `@session:/` link value passed as session_id. # Session ids never contain "/", so a slash unambiguously means profile/id — # always strip the prefix off the id, and adopt the embedded profile only @@ -973,6 +966,8 @@ def session_search( return tool_error(f"profile '{profile}': {e}", success=False) if profile_db is not None: db = profile_db + if _owned_dbs is not None: + _owned_dbs.append(profile_db) current_session_id = None # Scroll shape takes precedence — explicit anchor beats any query. @@ -1041,6 +1036,57 @@ def session_search( ) +def session_search( + query: str = "", + role_filter: str = None, + limit: int = 3, + db=None, + current_session_id: str = None, + # Scroll shape + session_id: str = None, + around_message_id: int = None, + window: int = 5, + # Discovery shape + sort: str = None, + # Cross-profile (any shape) + profile: str = None, +) -> str: + """Run session search and close databases opened by this invocation.""" + owned_dbs: List[Any] = [] + if db is None: + try: + from hermes_state import SessionDB + + db = SessionDB() + owned_dbs.append(db) + except Exception: + logging.debug("SessionDB unavailable for session_search", exc_info=True) + from hermes_state import format_session_db_unavailable + + return tool_error(format_session_db_unavailable(), success=False) + + try: + return _session_search_impl( + query=query, + role_filter=role_filter, + limit=limit, + db=db, + current_session_id=current_session_id, + session_id=session_id, + around_message_id=around_message_id, + window=window, + sort=sort, + profile=profile, + _owned_dbs=owned_dbs, + ) + finally: + for owned_db in reversed(owned_dbs): + try: + owned_db.close() + except Exception: + logging.debug("Failed to close session_search SessionDB", exc_info=True) + + def check_session_search_requirements() -> bool: """Requires the SQLite state database.""" try: diff --git a/tui_gateway/server.py b/tui_gateway/server.py index 19a78fdbe5..6be948282d 100644 --- a/tui_gateway/server.py +++ b/tui_gateway/server.py @@ -4128,6 +4128,20 @@ def _is_model_switch_marker(entry: Any) -> bool: return isinstance(content, str) and content.startswith(_MODEL_SWITCH_MARKER_PREFIX) +def _is_pivot_marker(entry: Any) -> bool: + """Whether a history entry is a marker the gateway splices in mid-turn. + + Model switches and personality changes both inject a ``role=user`` pivot + into the live history from the RPC thread while a turn may be running, so + either one can be the sole reason turn-start and current history differ. + Only the model-switch marker is self-replacing, which is why the dedup in + :func:`_append_model_switch_marker` stays narrower than this. + """ + if _is_model_switch_marker(entry): + return True + return isinstance(entry, dict) and entry.get("display_kind") == "personality_switch" + + def _append_model_switch_marker(session: dict | None, *, model: str, provider: str) -> None: """Record a real system-history pivot after a live model switch. @@ -10632,10 +10646,15 @@ def _run_prompt_submit( session["history_version"] = history_version + 1 else: # History mutated externally during the turn. - # Check if the only mutation was a model-switch - # marker inserted mid-turn (#76870). If so the - # agent output is still valid — merge it into the - # current history that now contains the marker. + # Check if the only mutation was a pivot marker + # the gateway itself inserted mid-turn (#76870). + # If so the agent output is still valid — merge it + # into the current history that now contains the + # marker. A personality change counts here too: + # unlike a model switch it has no pending queue, so + # `/personality` during a running turn lands + # immediately and used to read as a genuine desync, + # dropping the finished turn (#82756). # # _append_model_switch_marker strips prior markers # in-place then appends a new one, so the delta @@ -10643,19 +10662,19 @@ def _run_prompt_submit( # content, not indices. current_history = list(session["history"]) history_no_markers = [ - e for e in history if not _is_model_switch_marker(e) + e for e in history if not _is_pivot_marker(e) ] current_no_markers = [ - e for e in current_history if not _is_model_switch_marker(e) + e for e in current_history if not _is_pivot_marker(e) ] - model_switch_only = ( + pivot_only = ( current_no_markers == history_no_markers and any( - _is_model_switch_marker(e) + _is_pivot_marker(e) for e in current_history ) ) - if model_switch_only: + if pivot_only: # The agent's new messages start after the # turn-start history. Guard against # auto-compression making result["messages"]