fix(gateway): offload session store calls off the event loop via asyncio.to_thread

Every inbound message calls get_or_create_session which synchronously
executes _is_session_ended_in_db → db.get_session → conn.execute on
the asyncio event loop. On a ~1.4GB state.db, this blocks the loop
for seconds to minutes, starving Discord heartbeats.

Upstream #55159 fixed the same pattern for self._session_db in
gateway/run.py but missed SessionStore._db in gateway/session.py.

This follows the exact same approach as #55159:
- session.py internals stay fully synchronous (zero changes)
- Threading.Lock contract is preserved
- All hot-path callers in run.py and slash_commands.py wrap calls
  with await asyncio.to_thread(self.session_store.method, ...)

Affected: get_or_create_session, switch_session, update_session,
load_transcript, rewrite_transcript, rewind_session, reset_session,
set_model_override, _save — ~24 call sites across 2 files.

# Conflicts:
#	gateway/slash_commands.py
This commit is contained in:
kenyonxu 2026-07-10 11:21:14 +08:00 committed by kshitij
parent 04ca34b5f3
commit 24ea21993f
2 changed files with 47 additions and 43 deletions

View file

@ -10275,7 +10275,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
# on error. Let the user drive the next turn.
if _final_text.strip():
try:
session_entry = self.session_store.get_or_create_session(source)
session_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
except Exception:
session_entry = None
if session_entry is not None:
@ -10699,7 +10699,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
except Exception:
pass
session_entry = self.session_store.get_or_create_session(source)
session_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
session_key = session_entry.session_key
pinned_session_id = str(
(getattr(event, "metadata", None) or {}).get("gateway_session_id") or ""
@ -10781,7 +10781,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
# lane session is ended cleanly. Mutating session_entry in
# place here created a split-brain state where the JSON
# index pointed at one id but code downstream used another.
switched = self.session_store.switch_session(session_key, bound_session_id)
switched = await asyncio.to_thread(self.session_store.switch_session, session_key, bound_session_id)
if switched is not None:
session_entry = switched
# If the stored binding pointed at a parent, rewrite it to the
@ -12016,7 +12016,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
# Token counts and model are now persisted by the agent directly.
# Keep only last_prompt_tokens here for context-window tracking and
# compression decisions.
self.session_store.update_session(
await asyncio.to_thread(self.session_store.update_session,
session_entry.session_key,
last_prompt_tokens=agent_result.get("last_prompt_tokens", 0),
)
@ -12599,7 +12599,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
except Exception:
return 20
def _get_goal_manager_for_event(self, event: "MessageEvent"):
async def _get_goal_manager_for_event(self, event: "MessageEvent"):
"""Return a GoalManager bound to the session for this gateway event.
Returns ``(manager, session_entry)`` or ``(None, None)`` if the
@ -12611,7 +12611,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
logger.debug("goal manager unavailable: %s", exc)
return None, None
try:
session_entry = self.session_store.get_or_create_session(event.source)
session_entry = await asyncio.to_thread(self.session_store.get_or_create_session, event.source)
except Exception as exc:
logger.debug("goal manager: session lookup failed: %s", exc)
return None, None
@ -14054,7 +14054,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
"content": f"[IMPORTANT: MCP servers have been reloaded. {change_detail}{tool_summary}. The tool list for this conversation has been updated accordingly.]",
}
try:
session_entry = self.session_store.get_or_create_session(event.source)
session_entry = await asyncio.to_thread(self.session_store.get_or_create_session, event.source)
self.session_store.append_to_transcript(
session_entry.session_id, reload_msg
)

View file

@ -198,7 +198,7 @@ class GatewaySlashCommandsMixin:
pass
# Reset the session
new_entry = self.session_store.reset_session(session_key)
new_entry = await asyncio.to_thread(self.session_store.reset_session, session_key)
# Clear any session-scoped model/reasoning overrides so the next agent
# picks up configured defaults instead of previous session switches.
@ -263,7 +263,7 @@ class GatewaySlashCommandsMixin:
header = await asyncio.to_thread(self._telegram_topic_new_header, source) or t("gateway.reset.header_default")
else:
# No existing session, just create one
new_entry = self.session_store.get_or_create_session(source, force_new=True)
new_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source, force_new=True)
header = await asyncio.to_thread(self._telegram_topic_new_header, source) or t("gateway.reset.header_new")
# Set session title if provided with /new <title>
@ -495,7 +495,7 @@ class GatewaySlashCommandsMixin:
from gateway.run import _AGENT_PENDING_SENTINEL, _load_gateway_config, _resolve_gateway_model
source = event.source
session_entry = self.session_store.get_or_create_session(source)
session_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
connected_platforms = [p.value for p in self.adapters.keys()]
@ -1061,7 +1061,7 @@ class GatewaySlashCommandsMixin:
"""
from gateway.run import _AGENT_PENDING_SENTINEL, _INTERRUPT_REASON_STOP
source = event.source
session_entry = self.session_store.get_or_create_session(source)
session_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
session_key = session_entry.session_key
agent = self._running_agents.get(session_key)
@ -1598,8 +1598,9 @@ class GatewaySlashCommandsMixin:
_sess_db = getattr(_self, "_session_db", None)
if _sess_db is not None:
try:
_sess_entry = _self.session_store.get_or_create_session(
event.source
_sess_entry = await asyncio.to_thread(
_self.session_store.get_or_create_session,
event.source,
)
await _sess_db.update_session_model(
_sess_entry.session_id, result.new_model
@ -1629,7 +1630,8 @@ class GatewaySlashCommandsMixin:
# store so the picked model survives a gateway restart
# (api_key is never persisted).
try:
_self.session_store.set_model_override(
await asyncio.to_thread(
_self.session_store.set_model_override,
_session_key,
_self._session_model_overrides[_session_key],
)
@ -1840,7 +1842,7 @@ class GatewaySlashCommandsMixin:
_sess_db = getattr(self, "_session_db", None)
if _sess_db is not None:
try:
_sess_entry = self.session_store.get_or_create_session(source)
_sess_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
# If this session was auto-reset, consume the flag so the
# next regular message's cleanup does not wipe the model
# override just stored below (Closes #48031).
@ -1878,8 +1880,9 @@ class GatewaySlashCommandsMixin:
# api_key/api_mode are never persisted — they are re-resolved via
# runtime provider resolution on rehydration.
try:
self.session_store.set_model_override(
session_key, self._session_model_overrides[session_key]
await asyncio.to_thread(
self.session_store.set_model_override,
session_key, self._session_model_overrides[session_key],
)
except Exception:
logger.debug(
@ -2142,8 +2145,8 @@ class GatewaySlashCommandsMixin:
async def _handle_retry_command(self, event: MessageEvent) -> str:
"""Handle /retry command - re-send the last user message."""
source = event.source
session_entry = self.session_store.get_or_create_session(source)
history = self.session_store.load_transcript(session_entry.session_id)
session_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
history = await asyncio.to_thread(self.session_store.load_transcript, session_entry.session_id)
# Find the last user message
last_user_msg = None
@ -2159,10 +2162,10 @@ class GatewaySlashCommandsMixin:
# Truncate history to before the last user message and persist
truncated = history[:last_user_idx]
self.session_store.rewrite_transcript(session_entry.session_id, truncated)
await asyncio.to_thread(self.session_store.rewrite_transcript, session_entry.session_id, truncated)
# Reset stored token count — transcript was truncated
session_entry.last_prompt_tokens = 0
# Re-send by creating a fake text event with the old message
retry_event = MessageEvent(
text=last_user_msg,
@ -2189,7 +2192,7 @@ class GatewaySlashCommandsMixin:
args = (event.get_command_args() or "").strip()
lower = args.lower()
mgr, session_entry = self._get_goal_manager_for_event(event)
mgr, session_entry = await self._get_goal_manager_for_event(event)
if mgr is None:
return t("gateway.goal.unavailable")
@ -2323,7 +2326,7 @@ class GatewaySlashCommandsMixin:
to invoke while the agent is running.
"""
args = (event.get_command_args() or "").strip()
mgr, _session_entry = self._get_goal_manager_for_event(event)
mgr, _session_entry = await self._get_goal_manager_for_event(event)
if mgr is None:
return t("gateway.goal.unavailable")
if not mgr.has_goal():
@ -2390,8 +2393,8 @@ class GatewaySlashCommandsMixin:
if n < 1:
n = 1
session_entry = self.session_store.get_or_create_session(source)
result = self.session_store.rewind_session(session_entry.session_id, n)
session_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
result = await asyncio.to_thread(self.session_store.rewind_session, session_entry.session_id, n)
if result is None:
return t("gateway.undo.nothing")
@ -3090,8 +3093,8 @@ class GatewaySlashCommandsMixin:
https://code.claude.com/docs/en/whats-new/2026-w20).
"""
source = event.source
session_entry = self.session_store.get_or_create_session(source)
history = self.session_store.load_transcript(session_entry.session_id)
session_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
history = await asyncio.to_thread(self.session_store.load_transcript, session_entry.session_id)
if not history or len(history) < 4:
return t("gateway.compress.not_enough")
@ -3279,15 +3282,16 @@ class GatewaySlashCommandsMixin:
# original messages and replace them with only the compressed
# summary (permanent data loss #44794, #39704).
if rotated:
if not self.session_store.rewrite_transcript(
new_session_id, compressed
if not await asyncio.to_thread(
self.session_store.rewrite_transcript,
new_session_id, compressed,
):
raise RuntimeError(
f"failed to persist compressed transcript for "
f"session {new_session_id}"
)
session_entry.session_id = new_session_id
self.session_store._save()
await asyncio.to_thread(self.session_store._save)
await asyncio.to_thread(
self._sync_telegram_topic_binding,
source, session_entry, reason="compress-command",
@ -3304,8 +3308,8 @@ class GatewaySlashCommandsMixin:
"it (#44794)."
)
# Reset stored token count — transcript changed, old value is stale
self.session_store.update_session(
session_entry.session_key, last_prompt_tokens=0
await asyncio.to_thread(self.session_store.update_session,
session_entry.session_key, last_prompt_tokens=0,
)
new_tokens = estimate_request_tokens_rough(
compressed, system_prompt=_sys_prompt, tools=_tools
@ -3453,7 +3457,7 @@ class GatewaySlashCommandsMixin:
async def _handle_title_command(self, event: MessageEvent) -> str:
"""Handle /title command — set or show the current session's title."""
source = event.source
session_entry = self.session_store.get_or_create_session(source)
session_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
session_id = session_entry.session_id
if not self._session_db:
@ -3636,7 +3640,7 @@ class GatewaySlashCommandsMixin:
return t("gateway.resume.blocked_not_owner", name=name)
# Check if already on that session
current_entry = self.session_store.get_or_create_session(source)
current_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
if current_entry.session_id == target_id:
return t("gateway.resume.already_on", name=name)
@ -3644,7 +3648,7 @@ class GatewaySlashCommandsMixin:
self._release_running_agent_state(session_key)
# Switch the session entry to point at the old session
new_entry = self.session_store.switch_session(session_key, target_id)
new_entry = await asyncio.to_thread(self.session_store.switch_session, session_key, target_id)
if not new_entry:
return t("gateway.resume.switch_failed")
self._clear_session_boundary_security_state(session_key)
@ -3681,7 +3685,7 @@ class GatewaySlashCommandsMixin:
title = await self._session_db.get_session_title(target_id) or name
# Count messages for context
history = self.session_store.load_transcript(target_id)
history = await asyncio.to_thread(self.session_store.load_transcript, target_id)
msg_count = len([m for m in history if m.get("role") == "user"]) if history else 0
msg_part = f" ({msg_count} message{'s' if msg_count != 1 else ''})" if msg_count else ""
@ -3732,7 +3736,7 @@ class GatewaySlashCommandsMixin:
# `/sessions all` and enumerate other origins' session ids / titles /
# previews / sources — the enumeration half of the /resume IDOR.
cross_origin = include_all and self._resume_caller_is_admin(source)
current_entry = self.session_store.get_or_create_session(source)
current_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
rows = await asyncio.to_thread(
query_session_listing,
getattr(self._session_db, "_db", self._session_db),
@ -3781,8 +3785,8 @@ class GatewaySlashCommandsMixin:
session_key = self._session_key_for_source(source)
# Load the current session and its transcript
current_entry = self.session_store.get_or_create_session(source)
history = self.session_store.load_transcript(current_entry.session_id)
current_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
history = await asyncio.to_thread(self.session_store.load_transcript, current_entry.session_id)
if not history:
return t("gateway.branch.no_conversation")
@ -3849,7 +3853,7 @@ class GatewaySlashCommandsMixin:
pass
# Switch the session store entry to the new session
new_entry = self.session_store.switch_session(session_key, new_session_id)
new_entry = await asyncio.to_thread(self.session_store.switch_session, session_key, new_session_id)
if not new_entry:
return t("gateway.branch.switch_failed")
self._clear_session_boundary_security_state(session_key)
@ -3967,7 +3971,7 @@ class GatewaySlashCommandsMixin:
api_key = getattr(agent, "api_key", None) if agent and agent is not _AGENT_PENDING_SENTINEL else None
if not provider and getattr(self, "_session_db", None) is not None:
try:
_entry_for_billing = self.session_store.get_or_create_session(source)
_entry_for_billing = await asyncio.to_thread(self.session_store.get_or_create_session, source)
persisted = await self._session_db.get_session(_entry_for_billing.session_id) or {}
except Exception:
persisted = {}
@ -4055,8 +4059,8 @@ class GatewaySlashCommandsMixin:
return "\n".join(lines)
# No agent at all -- check session history for a rough count
session_entry = self.session_store.get_or_create_session(source)
history = self.session_store.load_transcript(session_entry.session_id)
session_entry = await asyncio.to_thread(self.session_store.get_or_create_session, source)
history = await asyncio.to_thread(self.session_store.load_transcript, session_entry.session_id)
if history:
from agent.model_metadata import estimate_messages_tokens_rough
msgs = [m for m in history if m.get("role") in {"user", "assistant"} and m.get("content")]