fix(api_server): offload synchronous SessionDB calls off the event loop

This commit is contained in:
dsad 2026-07-08 18:16:07 +03:00 committed by Brooklyn Nicholson
parent c7035ef252
commit 7ba944d054
2 changed files with 94 additions and 34 deletions

View file

@ -2165,21 +2165,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"]]:
async def _get_existing_session_or_404(self, session_id: str) -> tuple[Optional[Dict[str, Any]], Optional["web.Response"]]:
db = self._ensure_session_db()
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]]:
async def _conversation_history_for_session(self, session_id: str) -> List[Dict[str, Any]]:
db = self._ensure_session_db()
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 []
@ -2198,7 +2202,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,
@ -2233,22 +2237,22 @@ 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):
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)
db.create_session(session_id, "api_server", model=str(model) if model else None, system_prompt=system_prompt)
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:
db.set_session_title(session_id, str(title))
await asyncio.to_thread(db.set_session_title, session_id, str(title))
except ValueError as exc:
db.delete_session(session_id)
await asyncio.to_thread(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}
session = await asyncio.to_thread(db.get_session, session_id) or {"id": session_id, "source": "api_server", "model": model, "title": title}
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 +2260,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 +2271,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)
@ -2281,12 +2285,12 @@ class APIServerAdapter(BasePlatformAdapter):
db = self._ensure_session_db()
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 +2299,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)
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 +2312,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)
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,7 +2330,7 @@ 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)
@ -2336,35 +2340,35 @@ class APIServerAdapter(BasePlatformAdapter):
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 +2378,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 +2390,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 +2420,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 +2477,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,
@ -2657,7 +2661,7 @@ class APIServerAdapter(BasePlatformAdapter):
try:
db = self._ensure_session_db()
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,59 @@ 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()