diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index 30f401873c22d..52236713b81fe 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,7 +1592,7 @@ class APIServerAdapter(BasePlatformAdapter): # Session DB helper # ------------------------------------------------------------------ - def _ensure_session_db(self): + async def _ensure_session_db(self): """Lazily initialise and return the SessionDB for the active profile home. Sessions are persisted to ``state.db`` so that ``hermes sessions list`` @@ -1599,7 +1600,9 @@ 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. The first request per home pays the + SQLite open + schema-init cost; a single-flight lock prevents duplicate + concurrent construction, and the open itself runs off the event loop. """ # Explicit override (tests / manual wiring) wins. Production never sets # this externally, so the per-home cache below is the live path — and @@ -1612,16 +1615,25 @@ class APIServerAdapter(BasePlatformAdapter): from hermes_state import SessionDB home = get_hermes_home() + key = str(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 + if db is not None: + return db + if self._session_db_lock is None: + self._session_db_lock = asyncio.Lock() + async with self._session_db_lock: + # Double-check after acquiring the lock. + db = cache.get(key) + if db is None: + # Offload the blocking SQLite open + schema-init off the + # single aiohttp event-loop thread. + db = await asyncio.to_thread(SessionDB, db_path=home / "state.db") + cache[key] = db + return db except Exception as e: logger.debug("SessionDB unavailable for API server: %s", e) return None @@ -1827,7 +1839,7 @@ class APIServerAdapter(BasePlatformAdapter): tool_progress_callback=tool_progress_callback, tool_start_callback=tool_start_callback, tool_complete_callback=tool_complete_callback, - session_db=self._ensure_session_db(), + session_db=self._session_db, fallback_model=fallback_model, reasoning_config=reasoning_config, gateway_session_key=gateway_session_key, @@ -2166,7 +2178,7 @@ class APIServerAdapter(BasePlatformAdapter): return body, None async def _get_existing_session_or_404(self, session_id: str) -> tuple[Optional[Dict[str, Any]], Optional["web.Response"]]: - db = self._ensure_session_db() + db = await self._ensure_session_db() if db is None: return None, web.json_response(_openai_error("Session database unavailable", code="session_db_unavailable"), status=503) # Offload the blocking SQLite read off the event loop (CWE/perf: the @@ -2179,7 +2191,7 @@ class APIServerAdapter(BasePlatformAdapter): return session, None async def _conversation_history_for_session(self, session_id: str) -> List[Dict[str, Any]]: - db = self._ensure_session_db() + db = await self._ensure_session_db() if db is None: return [] try: @@ -2194,7 +2206,7 @@ class APIServerAdapter(BasePlatformAdapter): if auth_err: return auth_err - db = self._ensure_session_db() + db = await self._ensure_session_db() if db is None: return web.json_response(_openai_error("Session database unavailable", code="session_db_unavailable"), status=503) @@ -2218,7 +2230,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 @@ -2226,7 +2245,7 @@ class APIServerAdapter(BasePlatformAdapter): if err: return err - db = self._ensure_session_db() + db = await self._ensure_session_db() if db is None: return web.json_response(_openai_error("Session database unavailable", code="session_db_unavailable"), status=503) @@ -2237,22 +2256,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 await asyncio.to_thread(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) - await asyncio.to_thread(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: - await asyncio.to_thread(db.set_session_title, session_id, str(title)) - except ValueError as exc: - await asyncio.to_thread(db.delete_session, session_id) - return web.json_response(_openai_error(str(exc), code="invalid_title"), status=400) - session = await asyncio.to_thread(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": @@ -2282,7 +2347,7 @@ 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() if "title" in body: try: await asyncio.to_thread(db.set_session_title, session_id, "" if body["title"] is None else str(body["title"])) @@ -2302,7 +2367,7 @@ class APIServerAdapter(BasePlatformAdapter): session, err = await self._get_existing_session_or_404(session_id) if err: return err - db = self._ensure_session_db() + db = await self._ensure_session_db() deleted = await asyncio.to_thread(db.delete_session, session_id) return web.json_response({"object": "hermes.session.deleted", "id": session_id, "deleted": bool(deleted)}) @@ -2315,7 +2380,7 @@ class APIServerAdapter(BasePlatformAdapter): _, err = await self._get_existing_session_or_404(session_id) if err: return err - db = self._ensure_session_db() + db = await self._ensure_session_db() 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({ @@ -2336,7 +2401,7 @@ class APIServerAdapter(BasePlatformAdapter): body, err = await self._read_json_body(request) if err: return err - db = self._ensure_session_db() + db = await self._ensure_session_db() 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) @@ -2659,7 +2724,7 @@ class APIServerAdapter(BasePlatformAdapter): ) session_id = provided_session_id try: - db = self._ensure_session_db() + db = await self._ensure_session_db() if db is not None: history = await asyncio.to_thread(db.get_messages_as_conversation, session_id) except Exception as e: diff --git a/tests/gateway/test_api_server.py b/tests/gateway/test_api_server.py index 4330d1515961c..d79a251eb06ca 100644 --- a/tests/gateway/test_api_server.py +++ b/tests/gateway/test_api_server.py @@ -4305,3 +4305,73 @@ class TestSessionDbOffEventLoop: 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): + 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