WEBDIENSTE/openswarm
1
0
Fork 0
mirror of https://github.com/openswarm-ai/openswarm.git synced 2026-09-11 12:17:45 +02:00
Code Issues Packages Projects Releases Wiki Activity

[eric] agents: split bulk startup/shutdown persistence into SessionPersistenceMixin (one-session vs all-sessions)

Browse Source
This commit is contained in:
ciregenz
2026-06-23 14:40:20 -07:00
parent 5c725ed52d
commit 8c0a5010e1
3 changed files with 85 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/apps/agents/agent_manager.py
+2 -1
View File
@@ -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] = {}
backend/apps/agents/manager/session/SessionLifecycleMixin.py
-62
View File
@@ -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)
backend/apps/agents/manager/session/SessionPersistenceMixin.py
+83
View File
@@ -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}")