diff --git a/gateway/run.py b/gateway/run.py index 59f3c8a44ca..714949e030e 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -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 ) diff --git a/gateway/slash_commands.py b/gateway/slash_commands.py index aab3b733ae2..4c68ac5607f 100644 --- a/gateway/slash_commands.py +++ b/gateway/slash_commands.py @@ -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 @@ -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")]