Merge pull request #67283 from NousResearch/bb/salvage-60980-apiserver-offload

fix(api_server): offload synchronous SessionDB calls off the event loop (supersedes #60980)
This commit is contained in:
brooklyn! 2026-07-18 23:47:03 -04:00 committed by GitHub
commit e99a0f6a97
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 289 additions and 64 deletions

View file

@ -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/<profile>/`` 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 = []

View file

@ -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