diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index 37ab94dd6253..65d5b324cc45 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -990,6 +990,7 @@ 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_db_lock: Optional[asyncio.Lock] = None # Single-flight for lazy init # Concurrency cap shared across all agent-serving endpoints # (/v1/chat/completions, /v1/responses, /v1/runs). Read from # config.yaml gateway.api_server.max_concurrent_runs; 0 disables @@ -1591,6 +1592,28 @@ class APIServerAdapter(BasePlatformAdapter): # Session DB helper # ------------------------------------------------------------------ + def _open_and_cache_session_db(self, home) -> Optional[Any]: + """Sync core: return the cached SessionDB for ``home``, opening it once. + + Shared by the sync (``_ensure_session_db``) and async + (``_ensure_session_db_async``) entry points so both honor the same + per-profile cache. Deliberately does NOT write into ``self._session_db`` + — that stays reserved for an explicit test/manual override, so the first + profile served can't pin every later request to its DB. + """ + 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 + def _ensure_session_db(self): """Lazily initialise and return the SessionDB for the active profile home. @@ -1599,29 +1622,48 @@ class APIServerAdapter(BasePlatformAdapter): Under multiplex ``/p//`` requests the profile runtime scope redirects ``get_hermes_home()``, so each profile gets its own DB — - never the default profile's file. + never the default profile's file. Synchronous: used by ``_create_agent`` + (itself sync, and run in both loop and worker contexts). Request + handlers use ``_ensure_session_db_async`` to keep the SQLite open off + the event loop. + """ + # Explicit override (tests / manual wiring) wins. + if self._session_db is not None: + return self._session_db + try: + from hermes_constants import get_hermes_home + + return self._open_and_cache_session_db(get_hermes_home()) + except Exception as e: + logger.debug("SessionDB unavailable for API server: %s", e) + return None + + async def _ensure_session_db_async(self): + """Async variant for request handlers: offload the SQLite open/schema + init off the single aiohttp event-loop thread. + + The active profile home is captured on the loop thread (its runtime + scope is not visible inside ``asyncio.to_thread``); only the blocking + construction runs in the worker. A single-flight lock prevents duplicate + concurrent construction for the same home. """ - # Explicit override (tests / manual wiring) wins. Production never sets - # this externally, so the per-home cache below is the live path — and - # we deliberately do NOT write back into ``self._session_db`` there, or - # the first profile served would pin every later request to its DB. if self._session_db is not None: return self._session_db try: from hermes_constants import get_hermes_home - from hermes_state import SessionDB home = get_hermes_home() - cache = getattr(self, "_session_dbs", None) - if cache is None: - cache = {} - self._session_dbs = cache key = str(home) - db = cache.get(key) - if db is None: - db = SessionDB(db_path=home / "state.db") - cache[key] = db - return db + cache = getattr(self, "_session_dbs", None) + if cache is not None and cache.get(key) is not None: + return cache[key] + 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] + 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) return None @@ -2165,21 +2207,25 @@ class APIServerAdapter(BasePlatformAdapter): return {}, web.json_response(_openai_error("Request body must be a JSON object"), status=400) return body, None - def _get_existing_session_or_404(self, session_id: str) -> tuple[Optional[Dict[str, Any]], Optional["web.Response"]]: - db = self._ensure_session_db() + async def _get_existing_session_or_404(self, session_id: str) -> tuple[Optional[Dict[str, Any]], Optional["web.Response"]]: + db = await self._ensure_session_db_async() if db is None: return None, web.json_response(_openai_error("Session database unavailable", code="session_db_unavailable"), status=503) - session = db.get_session(session_id) + # Offload the blocking SQLite read off the event loop (CWE/perf: the + # API server is single-threaded aiohttp; a sync SessionDB call here + # freezes every in-flight request, see PR discussion on event-loop + # blocking SQLite in the gateway surface). + session = await asyncio.to_thread(db.get_session, session_id) if not session: return None, web.json_response(_openai_error(f"Session not found: {session_id}", code="session_not_found"), status=404) return session, None - def _conversation_history_for_session(self, session_id: str) -> List[Dict[str, Any]]: - db = self._ensure_session_db() + async def _conversation_history_for_session(self, session_id: str) -> List[Dict[str, Any]]: + db = await self._ensure_session_db_async() if db is None: return [] try: - return db.get_messages_as_conversation(session_id) + return await asyncio.to_thread(db.get_messages_as_conversation, session_id) except Exception as exc: logger.warning("Failed to load session history for %s: %s", session_id, exc) return [] @@ -2190,7 +2236,7 @@ class APIServerAdapter(BasePlatformAdapter): if auth_err: return auth_err - db = self._ensure_session_db() + db = await self._ensure_session_db_async() if db is None: return web.json_response(_openai_error("Session database unavailable", code="session_db_unavailable"), status=503) @@ -2198,7 +2244,7 @@ class APIServerAdapter(BasePlatformAdapter): offset = self._parse_nonnegative_int(request.query.get("offset"), default=0, maximum=1_000_000) source = request.query.get("source") or None include_children = _coerce_request_bool(request.query.get("include_children"), default=False) - sessions = db.list_sessions_rich( + sessions = await asyncio.to_thread(db.list_sessions_rich, source=source, limit=limit, offset=offset, @@ -2214,7 +2260,14 @@ class APIServerAdapter(BasePlatformAdapter): }) async def _handle_create_session(self, request: "web.Request") -> "web.Response": - """POST /api/sessions — create an empty Hermes session row.""" + """POST /api/sessions -- create an empty Hermes session row. + + The existence check, insert, title handling, and invalid-title + rollback run as a single off-loop operation to avoid a TOCTOU + window between the duplicate check and the insert (concurrent + same-ID creates could otherwise both pass the check and both + return 201 via the ON CONFLICT enrichment upsert). + """ auth_err = self._check_auth(request) if auth_err: return auth_err @@ -2222,7 +2275,7 @@ class APIServerAdapter(BasePlatformAdapter): if err: return err - db = self._ensure_session_db() + db = await self._ensure_session_db_async() if db is None: return web.json_response(_openai_error("Session database unavailable", code="session_db_unavailable"), status=503) @@ -2233,22 +2286,68 @@ class APIServerAdapter(BasePlatformAdapter): return web.json_response(_openai_error("Invalid session ID", code="invalid_session_id"), status=400) if len(session_id) > self._MAX_SESSION_HEADER_LEN: return web.json_response(_openai_error("Session ID too long", code="invalid_session_id"), status=400) - if db.get_session(session_id): - return web.json_response(_openai_error(f"Session already exists: {session_id}", code="session_exists"), status=409) model = body.get("model") or self._model_name system_prompt = body.get("system_prompt") if system_prompt is not None and not isinstance(system_prompt, str): return web.json_response(_openai_error("system_prompt must be a string", code="invalid_system_prompt"), status=400) - db.create_session(session_id, "api_server", model=str(model) if model else None, system_prompt=system_prompt) title = body.get("title") - if title is not None: - try: - db.set_session_title(session_id, str(title)) - except ValueError as exc: - db.delete_session(session_id) - return web.json_response(_openai_error(str(exc), code="invalid_title"), status=400) - session = db.get_session(session_id) or {"id": session_id, "source": "api_server", "model": model, "title": title} + + # Run the entire check-insert-title sequence inside a single + # _execute_write call (BEGIN IMMEDIATE + commit) so the existence + # check and the insert are atomic at the SQLite level. Two + # concurrent requests for the same ID serialize here: the second + # one blocks on the write lock and sees the row the first inserted. + def _do_create(): + def _atomic(conn): + row = conn.execute( + "SELECT id FROM sessions WHERE id = ?", (session_id,) + ).fetchone() + if row: + return None, "exists" + import time as _time + conn.execute( + """INSERT INTO sessions ( + id, source, model, system_prompt, started_at + ) VALUES (?, ?, ?, ?, ?)""", + ( + session_id, + "api_server", + str(model) if model else None, + system_prompt, + _time.time(), + ), + ) + if title is not None: + clean_title = db.sanitize_title(str(title)) + if clean_title: + conflict = conn.execute( + "SELECT id FROM sessions WHERE title = ? AND id != ?", + (clean_title, session_id), + ).fetchone() + if conflict: + conn.execute( + "DELETE FROM sessions WHERE id = ?", (session_id,) + ) + return None, f"title:Title already in use by session {conflict['id']}" + conn.execute( + "UPDATE sessions SET title = ? WHERE id = ?", + (clean_title, session_id), + ) + session_row = conn.execute( + "SELECT * FROM sessions WHERE id = ?", (session_id,) + ).fetchone() + return (dict(session_row) if session_row else { + "id": session_id, "source": "api_server", + "model": model, "title": title, + }), None + return db._execute_write(_atomic) + + session, err = await asyncio.to_thread(_do_create) + if err == "exists": + return web.json_response(_openai_error(f"Session already exists: {session_id}", code="session_exists"), status=409) + if err and err.startswith("title:"): + return web.json_response(_openai_error(err[len("title:"):], code="invalid_title"), status=400) return web.json_response({"object": "hermes.session", "session": self._session_response(session)}, status=201) async def _handle_get_session(self, request: "web.Request") -> "web.Response": @@ -2256,7 +2355,7 @@ class APIServerAdapter(BasePlatformAdapter): auth_err = self._check_auth(request) if auth_err: return auth_err - session, err = self._get_existing_session_or_404(request.match_info["session_id"]) + session, err = await self._get_existing_session_or_404(request.match_info["session_id"]) if err: return err return web.json_response({"object": "hermes.session", "session": self._session_response(session)}) @@ -2267,7 +2366,7 @@ class APIServerAdapter(BasePlatformAdapter): if auth_err: return auth_err session_id = request.match_info["session_id"] - session, err = self._get_existing_session_or_404(session_id) + session, err = await self._get_existing_session_or_404(session_id) if err: return err body, err = await self._read_json_body(request) @@ -2278,15 +2377,15 @@ class APIServerAdapter(BasePlatformAdapter): if unknown: return web.json_response(_openai_error(f"Unsupported session fields: {', '.join(unknown)}", code="unsupported_session_field"), status=400) - db = self._ensure_session_db() + db = await self._ensure_session_db_async() if "title" in body: try: - db.set_session_title(session_id, "" if body["title"] is None else str(body["title"])) + await asyncio.to_thread(db.set_session_title, session_id, "" if body["title"] is None else str(body["title"])) except ValueError as exc: return web.json_response(_openai_error(str(exc), code="invalid_title"), status=400) if body.get("end_reason"): - db.end_session(session_id, str(body["end_reason"])) - session = db.get_session(session_id) or session + await asyncio.to_thread(db.end_session, session_id, str(body["end_reason"])) + session = await asyncio.to_thread(db.get_session, session_id) or session return web.json_response({"object": "hermes.session", "session": self._session_response(session)}) async def _handle_delete_session(self, request: "web.Request") -> "web.Response": @@ -2295,11 +2394,11 @@ class APIServerAdapter(BasePlatformAdapter): if auth_err: return auth_err session_id = request.match_info["session_id"] - session, err = self._get_existing_session_or_404(session_id) + session, err = await self._get_existing_session_or_404(session_id) if err: return err - db = self._ensure_session_db() - deleted = db.delete_session(session_id) + db = await self._ensure_session_db_async() + deleted = await asyncio.to_thread(db.delete_session, session_id) return web.json_response({"object": "hermes.session.deleted", "id": session_id, "deleted": bool(deleted)}) async def _handle_session_messages(self, request: "web.Request") -> "web.Response": @@ -2308,12 +2407,12 @@ class APIServerAdapter(BasePlatformAdapter): if auth_err: return auth_err session_id = request.match_info["session_id"] - _, err = self._get_existing_session_or_404(session_id) + _, err = await self._get_existing_session_or_404(session_id) if err: return err - db = self._ensure_session_db() - resolved_id = db.resolve_resume_session_id(session_id) - messages = db.get_messages(resolved_id) + db = await self._ensure_session_db_async() + resolved_id = await asyncio.to_thread(db.resolve_resume_session_id, session_id) + messages = await asyncio.to_thread(db.get_messages, resolved_id) return web.json_response({ "object": "list", "session_id": resolved_id, @@ -2326,45 +2425,45 @@ class APIServerAdapter(BasePlatformAdapter): if auth_err: return auth_err source_id = request.match_info["session_id"] - source, err = self._get_existing_session_or_404(source_id) + source, err = await self._get_existing_session_or_404(source_id) if err: return err body, err = await self._read_json_body(request) if err: return err - db = self._ensure_session_db() + db = await self._ensure_session_db_async() fork_id = str(body.get("id") or body.get("session_id") or f"api_{int(time.time())}_{uuid.uuid4().hex[:8]}").strip() if not fork_id or re.search(r'[\r\n\x00]', fork_id): return web.json_response(_openai_error("Invalid session ID", code="invalid_session_id"), status=400) - if db.get_session(fork_id): + if await asyncio.to_thread(db.get_session, fork_id): return web.json_response(_openai_error(f"Session already exists: {fork_id}", code="session_exists"), status=409) # Match the CLI /branch semantics: mark the original as branched, then # create a child session that carries the transcript forward. This uses # SessionDB's native parent_session_id/end_reason visibility model rather # than inventing a parallel fork store. - db.end_session(source_id, "branched") - db.create_session( + await asyncio.to_thread(db.end_session, source_id, "branched") + await asyncio.to_thread(db.create_session, fork_id, "api_server", model=source.get("model"), system_prompt=source.get("system_prompt"), parent_session_id=source_id, ) - messages = db.get_messages(source_id) - db.replace_messages(fork_id, messages) + messages = await asyncio.to_thread(db.get_messages, source_id) + await asyncio.to_thread(db.replace_messages, fork_id, messages) title = body.get("title") if title is None: base = source.get("title") or "fork" try: - title = db.get_next_title_in_lineage(base) + title = await asyncio.to_thread(db.get_next_title_in_lineage, base) except Exception: title = f"{base} fork" try: - db.set_session_title(fork_id, str(title)) + await asyncio.to_thread(db.set_session_title, fork_id, str(title)) except ValueError as exc: return web.json_response(_openai_error(str(exc), code="invalid_title"), status=400) - fork = db.get_session(fork_id) or {"id": fork_id, "parent_session_id": source_id} + fork = await asyncio.to_thread(db.get_session, fork_id) or {"id": fork_id, "parent_session_id": source_id} return web.json_response({"object": "hermes.session", "session": self._session_response(fork)}, status=201) @_admit_api_agent_request @@ -2374,7 +2473,7 @@ class APIServerAdapter(BasePlatformAdapter): if key_err is not None: return key_err session_id = request.match_info["session_id"] - _, err = self._get_existing_session_or_404(session_id) + _, err = await self._get_existing_session_or_404(session_id) if err: return err body, err = await self._read_json_body(request) @@ -2386,7 +2485,7 @@ class APIServerAdapter(BasePlatformAdapter): system_prompt = body.get("system_message") or body.get("instructions") if system_prompt is not None and not isinstance(system_prompt, str): return web.json_response(_openai_error("system_message must be a string", code="invalid_system_message"), status=400) - history = self._conversation_history_for_session(session_id) + history = await self._conversation_history_for_session(session_id) result, usage = await self._run_agent( user_message=user_message, conversation_history=history, @@ -2416,7 +2515,7 @@ class APIServerAdapter(BasePlatformAdapter): if key_err is not None: return key_err session_id = request.match_info["session_id"] - _, err = self._get_existing_session_or_404(session_id) + _, err = await self._get_existing_session_or_404(session_id) if err: return err body, err = await self._read_json_body(request) @@ -2473,7 +2572,7 @@ class APIServerAdapter(BasePlatformAdapter): try: await queue.put(_event_payload("run.started", {"user_message": {"role": "user", "content": user_message}})) await queue.put(_event_payload("message.started", {"message": {"id": message_id, "role": "assistant"}})) - history = self._conversation_history_for_session(session_id) + history = await self._conversation_history_for_session(session_id) result, usage = await self._run_agent( user_message=user_message, conversation_history=history, @@ -2655,9 +2754,9 @@ class APIServerAdapter(BasePlatformAdapter): ) session_id = provided_session_id try: - db = self._ensure_session_db() + db = await self._ensure_session_db_async() if db is not None: - history = db.get_messages_as_conversation(session_id) + history = await asyncio.to_thread(db.get_messages_as_conversation, session_id) except Exception as e: logger.warning("Failed to load session history for %s: %s", session_id, e) history = [] diff --git a/tests/gateway/test_api_server.py b/tests/gateway/test_api_server.py index aa48e5323daa..a9621924b577 100644 --- a/tests/gateway/test_api_server.py +++ b/tests/gateway/test_api_server.py @@ -4249,3 +4249,129 @@ class TestModelRoutesAgentCreation: assert adapter._session_model_override_for("chan-1") == {"model": "user/model"} assert adapter._session_model_override_for("chan-2") is None assert adapter._session_model_override_for(None) is None + + +# --------------------------------------------------------------------------- +# Event-loop offloading for synchronous SessionDB calls (P1) +# --------------------------------------------------------------------------- + + +class TestSessionDbOffEventLoop: + """Regression: synchronous SessionDB calls in the OpenAI-compatible API + server must run OFF the aiohttp event loop. A blocking SQLite read/write on + the loop freezes every in-flight request under load (same class of bug as + gateway build_channel_directory, #60794 / #60810), so each call is wrapped + in asyncio.to_thread. + """ + + @pytest.mark.asyncio + async def test_get_existing_session_or_404_offloads(self, auth_adapter): + import threading + + captured = {} + + class FakeDB: + def get_session(self, session_id): + captured["thread"] = threading.current_thread() + return {"id": session_id, "source": "api_server"} + + auth_adapter._session_db = FakeDB() + session, err = await auth_adapter._get_existing_session_or_404("sess-x") + assert err is None + assert session["id"] == "sess-x" + # The blocking DB call must NOT execute on the event-loop thread. + assert captured["thread"] is not None + assert captured["thread"] != threading.current_thread() + + @pytest.mark.asyncio + async def test_list_sessions_offloads_db_off_event_loop(self, auth_adapter): + import threading + + captured = {} + + class FakeDB: + def list_sessions_rich(self, **kwargs): + captured["thread"] = threading.current_thread() + return [] + + auth_adapter._session_db = FakeDB() + app = _create_app(auth_adapter) + app.router.add_get("/api/sessions", auth_adapter._handle_list_sessions) + async with TestClient(TestServer(app)) as cli: + resp = await cli.get( + "/api/sessions", + headers={"Authorization": "Bearer sk-secret"}, + ) + assert resp.status == 200 + assert captured["thread"] is not None + assert captured["thread"] != threading.current_thread() + + @pytest.mark.asyncio + async def test_concurrent_same_id_create_one_201_one_409(self, auth_adapter): + """Two concurrent creates for the same ID must yield one 201 and one 409. + + The create sequence (existence check + insert + title) runs as a + single off-loop call, so concurrent same-ID requests serialize at + the DB level. Before the fix the TOCTOU window between the check + and the insert let both requests pass the existence guard and both + return 201 via the ON CONFLICT enrichment upsert. + """ + import asyncio + + app = _create_app(auth_adapter) + app.router.add_post("/api/sessions", auth_adapter._handle_create_session) + + async with TestClient(TestServer(app)) as cli: + # Fire both requests concurrently through the same server. + resp_a, resp_b = await asyncio.gather( + cli.post( + "/api/sessions", + json={"id": "race-same-id"}, + headers={"Authorization": "Bearer sk-secret"}, + ), + cli.post( + "/api/sessions", + json={"id": "race-same-id"}, + headers={"Authorization": "Bearer sk-secret"}, + ), + ) + assert sorted([resp_a.status, resp_b.status]) == [201, 409] + + @pytest.mark.asyncio + async def test_ensure_session_db_first_request_path(self, auth_adapter): + """First /api/sessions request initializes SessionDB off the event loop.""" + import threading + + captured = {} + + class FakeDB: + def __init__(self, db_path=None): + captured["init_thread"] = threading.current_thread() + + def list_sessions_rich(self, **kwargs): + return [] + + # Simulate cold start -- no DB yet. + auth_adapter._session_db = None + auth_adapter._session_db_lock = None + + original_class = None + import hermes_state + original_class = hermes_state.SessionDB + hermes_state.SessionDB = FakeDB + try: + app = _create_app(auth_adapter) + app.router.add_get("/api/sessions", auth_adapter._handle_list_sessions) + async with TestClient(TestServer(app)) as cli: + resp = await cli.get( + "/api/sessions", + headers={"Authorization": "Bearer sk-secret"}, + ) + assert resp.status == 200 + # SessionDB() was constructed -- the init must NOT be on the event-loop thread. + assert "init_thread" in captured + assert captured["init_thread"] != threading.current_thread() + finally: + hermes_state.SessionDB = original_class + auth_adapter._session_db = None + auth_adapter._session_db_lock = None