mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-20 19:52:23 +02:00
103 lines
5.1 KiB
Python
103 lines
5.1 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 resolves 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 (
|
|
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."""
|
|
marked = 0
|
|
for sid, data in load_all_session_data():
|
|
dirty = False
|
|
if data.get("status") in ("running", "waiting_approval"):
|
|
data["status"] = "stopped"
|
|
dirty = True
|
|
marked += 1
|
|
# 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)
|
|
if marked:
|
|
logger.info(f"Marked {marked} stale session(s) as stopped")
|
|
|
|
@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 = []
|
|
doc_data = session.model_dump(mode="json")
|
|
doc_data["search_text"] = self.build_search_text(session)
|
|
save_session(session_id, doc_data)
|
|
if self.sessions:
|
|
logger.info(f"Persisted {len(self.sessions)} session(s) 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.
|
|
"""
|
|
restored = 0
|
|
deferred = 0
|
|
for sid, data in load_all_session_data():
|
|
if not isinstance(data, dict):
|
|
continue
|
|
# Only mid-turn sessions need boot hydration (their status finalize below). Everything
|
|
# else loads on demand: history reads disk, dashboard lists promote by layout card ids,
|
|
# and resume/send_message both lazy-load. Eagerly validating thousands of settled
|
|
# sessions was the 2.7s (and hundreds of MB) of every boot.
|
|
if data.get("closed_at") is not None:
|
|
continue
|
|
if data.get("status") not in ("running", "waiting_approval"):
|
|
deferred += 1
|
|
continue
|
|
try:
|
|
session = AgentSession(**data)
|
|
except Exception as e:
|
|
logger.warning(f"Skipping corrupt session file {sid}: {e}")
|
|
continue
|
|
if session.status in ("running", "waiting_approval"):
|
|
# The app died mid-turn. If the last message in the active branch is already an assistant reply, the turn finished streaming and only the status finalize was lost (-> completed, no spurious "Resume" button); otherwise the agent was genuinely cut off owing a response (-> stopped, resumable).
|
|
branch = session.active_branch_id or "main"
|
|
p_branch_msgs = [m for m in session.messages if (m.branch_id or "main") == branch]
|
|
p_last = p_branch_msgs[-1] if p_branch_msgs else None
|
|
session.status = "completed" if (p_last is not None and p_last.role == "assistant") else "stopped"
|
|
session.pending_approvals = []
|
|
apply_context_window(session)
|
|
# The file stays on disk: unlinking here made RAM the only copy, so any non-graceful shutdown (updater SIGKILL, crash) destroyed the chat.
|
|
self.sessions[session.id] = session
|
|
restored += 1
|
|
# One summary line, not one per session (startups with hundreds of sessions flooded the console).
|
|
if restored or deferred:
|
|
logger.info(f"Restored {restored} mid-turn session(s); {deferred} settled session(s) stay on disk for lazy load")
|