fix: TOCTOU race in session create + offload SessionDB init
- Make create sequence (check + insert + title) atomic via single _execute_write call with BEGIN IMMEDIATE, closing the TOCTOU window where two concurrent same-ID creates could both return 201. - Offload _ensure_session_db() to asyncio.to_thread with single-flight lock so first-request SQLite init doesn't block the event loop. - Add concurrent same-ID create test (one 201, one 409) and first-request path test covering the initialization.
This commit is contained in:
dsad
Brooklyn Nicholson
parent
7ba944d054
commit
5529175084
|
|
@ -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/<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. 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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue