From 8c0a5010e175a968e9d14e39bcdeaef52a467ad1 Mon Sep 17 00:00:00 2001 From: ciregenz Date: Tue, 23 Jun 2026 14:40:20 -0700 Subject: [PATCH] [eric] agents: split bulk startup/shutdown persistence into SessionPersistenceMixin (one-session vs all-sessions) --- backend/apps/agents/agent_manager.py | 3 +- .../manager/session/SessionLifecycleMixin.py | 62 -------------- .../session/SessionPersistenceMixin.py | 83 +++++++++++++++++++ 3 files changed, 85 insertions(+), 63 deletions(-) create mode 100644 backend/apps/agents/manager/session/SessionPersistenceMixin.py diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 6a262c6e..9a8409a1 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -50,6 +50,7 @@ from backend.apps.agents.manager.permissions.effective_tools import build_effect from backend.apps.agents.manager.builtin_mcp_servers import register_builtin_mcp_servers from backend.apps.agents.manager.provider_env import configure_provider_env from backend.apps.agents.manager.session.SessionLifecycleMixin import SessionLifecycleMixin +from backend.apps.agents.manager.session.SessionPersistenceMixin import SessionPersistenceMixin from backend.apps.agents.manager.MessagingMixin import MessagingMixin from backend.apps.agents.manager.SessionControlMixin import SessionControlMixin from backend.apps.agents.manager.AgentLaunchMixin import AgentLaunchMixin @@ -72,7 +73,7 @@ logger = logging.getLogger(__name__) os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000") -class AgentManager(SessionLifecycleMixin, MessagingMixin, SessionControlMixin, AgentLaunchMixin, MockAgentMixin, RunSupportMixin): +class AgentManager(SessionLifecycleMixin, SessionPersistenceMixin, MessagingMixin, SessionControlMixin, AgentLaunchMixin, MockAgentMixin, RunSupportMixin): @typechecked def __init__(self): self.sessions: Dict[str, AgentSession] = {} diff --git a/backend/apps/agents/manager/session/SessionLifecycleMixin.py b/backend/apps/agents/manager/session/SessionLifecycleMixin.py index 79fd1338..d69378ae 100644 --- a/backend/apps/agents/manager/session/SessionLifecycleMixin.py +++ b/backend/apps/agents/manager/session/SessionLifecycleMixin.py @@ -186,68 +186,6 @@ class SessionLifecycleMixin: "has_more": offset + limit < total, } - @typechecked - async def reconcile_on_startup(self) -> None: - """Mark any stale running sessions as stopped.""" - for sid, data in load_all_session_data(): - dirty = False - if data.get("status") in ("running", "waiting_approval"): - data["status"] = "stopped" - dirty = True - logger.info(f"Marked stale session {sid} as stopped") - # Mode migration: Chat was merged into Ask. Rewrite mode="chat" - # so old sessions keep loading after the chat.json file is gone. - if data.get("mode") == "chat": - data["mode"] = "ask" - dirty = True - if dirty: - save_session(sid, data) - - @typechecked - async def persist_all_sessions(self) -> None: - """Flush every in-memory session to JSON files (for graceful shutdown).""" - for session_id, session in list(self.sessions.items()): - if session.status in ("running", "waiting_approval"): - session.status = "stopped" - session.closed_at = None - for req in list(session.pending_approvals): - ws_manager.resolve_approval(req.id, {"behavior": "deny", "message": "Server shutting down"}) - session.pending_approvals = [] - # Tag this close as "shutdown" so the cloud can tell it apart - # from a user-initiated close. The desktop doesn't care; the - # tag rides along in the dump for whoever consumes it. - self.p_sync_session_close(session, close_reason="shutdown") - doc_data = session.model_dump(mode="json") - doc_data["search_text"] = self.p_build_search_text(session) - save_session(session_id, doc_data) - logger.info(f"Persisted session {session_id} on shutdown") - self.sessions.clear() - self.tasks.clear() - - @typechecked - async def restore_all_sessions(self) -> None: - """On startup, reload all persisted sessions from JSON files back into memory. - - Only sessions without closed_at are restored (they were active at - shutdown). Sessions with closed_at were explicitly closed by the user - and stay on disk so the history endpoint can still serve them. - """ - for sid, data in load_all_session_data(): - try: - session = AgentSession(**data) - except Exception as e: - logger.warning(f"Skipping corrupt session file {sid}: {e}") - continue - if session.closed_at is not None: - continue - if session.status in ("running", "waiting_approval"): - session.status = "stopped" - session.pending_approvals = [] - apply_context_window(session) - self.sessions[session.id] = session - delete_session_file(sid) - logger.info(f"Restored session {session.id}") - @typechecked async def duplicate_session(self, session_id: str, dashboard_id: Optional[str] = None, up_to_message_id: Optional[str] = None) -> AgentSession: new_session = lifecycle.build_duplicate_session(self.sessions.get(session_id), session_id, dashboard_id, up_to_message_id) diff --git a/backend/apps/agents/manager/session/SessionPersistenceMixin.py b/backend/apps/agents/manager/session/SessionPersistenceMixin.py new file mode 100644 index 00000000..7001029c --- /dev/null +++ b/backend/apps/agents/manager/session/SessionPersistenceMixin.py @@ -0,0 +1,83 @@ +"""Bulk session persistence across the WHOLE store, the startup/shutdown orchestration that +operates on every session at once (reconcile stale-running, flush-all on shutdown, restore-all +on boot). Split from SessionLifecycleMixin (which handles ONE session at a time) so each file is +one concern. self.sessions / self.p_sync_session_close resolve across the MRO as before.""" + +import logging + +from typeguard import typechecked + +from backend.apps.agents.core.models import AgentSession +from backend.apps.agents.core.ws_manager import ws_manager +from backend.apps.agents.manager.session.session_store import ( + delete_session_file, + load_all_session_data, + save_session, +) +from backend.apps.agents.manager.session.apply_context_window import apply_context_window + +logger = logging.getLogger(__name__) + + +class SessionPersistenceMixin: + @typechecked + async def reconcile_on_startup(self) -> None: + """Mark any stale running sessions as stopped.""" + for sid, data in load_all_session_data(): + dirty = False + if data.get("status") in ("running", "waiting_approval"): + data["status"] = "stopped" + dirty = True + logger.info(f"Marked stale session {sid} as stopped") + # Mode migration: Chat was merged into Ask. Rewrite mode="chat" + # so old sessions keep loading after the chat.json file is gone. + if data.get("mode") == "chat": + data["mode"] = "ask" + dirty = True + if dirty: + save_session(sid, data) + + @typechecked + async def persist_all_sessions(self) -> None: + """Flush every in-memory session to JSON files (for graceful shutdown).""" + for session_id, session in list(self.sessions.items()): + if session.status in ("running", "waiting_approval"): + session.status = "stopped" + session.closed_at = None + for req in list(session.pending_approvals): + ws_manager.resolve_approval(req.id, {"behavior": "deny", "message": "Server shutting down"}) + session.pending_approvals = [] + # Tag this close as "shutdown" so the cloud can tell it apart + # from a user-initiated close. The desktop doesn't care; the + # tag rides along in the dump for whoever consumes it. + self.p_sync_session_close(session, close_reason="shutdown") + doc_data = session.model_dump(mode="json") + doc_data["search_text"] = self.p_build_search_text(session) + save_session(session_id, doc_data) + logger.info(f"Persisted session {session_id} on shutdown") + self.sessions.clear() + self.tasks.clear() + + @typechecked + async def restore_all_sessions(self) -> None: + """On startup, reload all persisted sessions from JSON files back into memory. + + Only sessions without closed_at are restored (they were active at + shutdown). Sessions with closed_at were explicitly closed by the user + and stay on disk so the history endpoint can still serve them. + """ + for sid, data in load_all_session_data(): + try: + session = AgentSession(**data) + except Exception as e: + logger.warning(f"Skipping corrupt session file {sid}: {e}") + continue + if session.closed_at is not None: + continue + if session.status in ("running", "waiting_approval"): + session.status = "stopped" + session.pending_approvals = [] + apply_context_window(session) + self.sessions[session.id] = session + delete_session_file(sid) + logger.info(f"Restored session {session.id}")