mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-21 20:22:22 +02:00
87 lines
3.8 KiB
Python
87 lines
3.8 KiB
Python
"""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 SessionLifecycle (which handles ONE session at a time) so each file is
|
|
one concern. self.sessions / self.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__)
|
|
|
|
|
|
from backend.apps.agents.manager.AgentManagerProtocol import AgentManagerProtocol
|
|
|
|
|
|
class SessionPersistence(AgentManagerProtocol):
|
|
@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.sync_session_close(session, close_reason="shutdown")
|
|
doc_data = session.model_dump(mode="json")
|
|
doc_data["search_text"] = self.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}")
|