import ast import asyncio import threading from pathlib import Path from hermes_cli import web_server TARGET_HANDLERS = { "bulk_delete_sessions_endpoint", "count_empty_sessions_endpoint", "delete_empty_sessions_endpoint", "get_session_latest_descendant", "get_session_messages", "delete_session_endpoint", "export_session_endpoint", "prune_sessions_endpoint", "get_usage_analytics", "get_models_analytics", } def _call_name(call: ast.Call) -> str | None: if isinstance(call.func, ast.Name): return call.func.id if isinstance(call.func, ast.Attribute): return call.func.attr return None def test_sessiondb_handlers_open_connections_inside_executor_helpers(): tree = ast.parse(Path(web_server.__file__).read_text(encoding="utf-8")) handlers = { node.name: node for node in tree.body if isinstance(node, ast.AsyncFunctionDef) and node.name in TARGET_HANDLERS } top_level_helpers = { node.name: node for node in tree.body if isinstance(node, ast.FunctionDef) } assert handlers.keys() == TARGET_HANDLERS for name, handler in handlers.items(): helpers = { **top_level_helpers, **{ node.name: node for node in handler.body if isinstance(node, ast.FunctionDef) }, } offloaded = { arg.id for node in ast.walk(handler) if isinstance(node, ast.Call) and _call_name(node) == "to_thread" for arg in node.args[:1] if isinstance(arg, ast.Name) } db_open_owners = { helper_name for helper_name, helper in helpers.items() if helper_name in offloaded and any( isinstance(node, ast.Call) and _call_name(node) == "_open_session_db_for_profile" for node in ast.walk(helper) ) } assert db_open_owners, f"{name} does not offload SessionDB open + work" def test_bulk_delete_sessiondb_work_runs_off_event_loop(monkeypatch): loop_thread = threading.get_ident() db_threads: list[int] = [] class _DB: def delete_sessions(self, ids): db_threads.append(threading.get_ident()) assert ids == ["one", "two"] return 2 def close(self): db_threads.append(threading.get_ident()) monkeypatch.setattr(web_server, "_open_session_db_for_profile", lambda profile=None: _DB()) result = asyncio.run( web_server.bulk_delete_sessions_endpoint( web_server.BulkDeleteSessions(ids=["one", "two"]) ) ) assert result == {"ok": True, "deleted": 2} assert db_threads assert all(thread_id != loop_thread for thread_id in db_threads)