mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-09-08 02:37:45 +02:00
[eric] agents: split session lifecycle + query methods into SessionLifecycleMixin
This commit is contained in:
ciregenz
@@ -43,17 +43,12 @@ from backend.apps.agents.core.error_classify import (
|
||||
redact_for_telemetry,
|
||||
)
|
||||
from backend.apps.agents.manager.session.session_store import (
|
||||
_delete_session_file,
|
||||
_load_all_session_data,
|
||||
_load_session_data,
|
||||
_save_session,
|
||||
build_search_text,
|
||||
)
|
||||
from backend.apps.agents.manager.session.cloud_sync import _sync_session_close
|
||||
from backend.apps.agents.manager import browser_dispatch
|
||||
from backend.apps.agents.manager import metadata
|
||||
from backend.apps.agents.manager.session.apply_context_window import apply_context_window
|
||||
from backend.apps.agents.manager.session import lifecycle
|
||||
from backend.apps.agents.manager.permissions import path_gate
|
||||
from backend.apps.agents.manager import context_budget
|
||||
from backend.apps.agents.manager.streaming.state import ThinkingState, TurnState
|
||||
@@ -63,11 +58,8 @@ from backend.apps.agents.manager.streaming import tool_result_hook
|
||||
from backend.apps.agents.manager.streaming import stop_hook as stop_hook_mod
|
||||
from backend.apps.agents.manager.prompt.system_prompt import compose_turn_system_prompt
|
||||
from backend.apps.agents.tools.web import should_register_web_mcp
|
||||
from backend.apps.agents.manager.session.SessionLifecycleMixin import SessionLifecycleMixin
|
||||
from backend.apps.agents.manager.permissions import gate_hooks
|
||||
from backend.apps.agents.manager.view_builder_state import (
|
||||
view_builder_render_retry_counts,
|
||||
view_builder_dirty_sessions,
|
||||
)
|
||||
from backend.apps.agents.manager.session.workspace_git import _detect_git_identity, _ensure_cwd_git_repo
|
||||
from backend.apps.agents.manager.prompt.tool_catalog import (
|
||||
FULL_TOOLS,
|
||||
@@ -96,7 +88,7 @@ logger = logging.getLogger(__name__)
|
||||
os.environ.setdefault("CLAUDE_CODE_STREAM_CLOSE_TIMEOUT", "3600000")
|
||||
|
||||
|
||||
class AgentManager:
|
||||
class AgentManager(SessionLifecycleMixin):
|
||||
def __init__(self):
|
||||
self.sessions: dict[str, AgentSession] = {}
|
||||
self.tasks: dict[str, asyncio.Task] = {}
|
||||
@@ -2847,224 +2839,6 @@ class AgentManager:
|
||||
})
|
||||
|
||||
@staticmethod
|
||||
def _build_search_text(session: AgentSession, max_len: int = 5000) -> str:
|
||||
return build_search_text(session, max_len)
|
||||
|
||||
def _sync_session_close(self, session: AgentSession, close_reason: str = "user"):
|
||||
_sync_session_close(session, close_reason)
|
||||
|
||||
async def close_session(self, session_id: str) -> None:
|
||||
"""Close a session: pause the agent if running, persist to JSON file,
|
||||
and remove from in-memory state. Also stops browser-agent children."""
|
||||
children = [
|
||||
s for s in self.sessions.values()
|
||||
if s.parent_session_id == session_id and s.mode == "browser-agent"
|
||||
]
|
||||
for child in children:
|
||||
await self.stop_agent(child.id)
|
||||
|
||||
task = self.tasks.get(session_id)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
session = self.sessions.get(session_id)
|
||||
if not session:
|
||||
raise ValueError(f"Session {session_id} not found")
|
||||
|
||||
if session.status in ("running", "waiting_approval"):
|
||||
session.status = "stopped"
|
||||
session.closed_at = datetime.now()
|
||||
|
||||
for req in list(session.pending_approvals):
|
||||
ws_manager.resolve_approval(req.id, {"behavior": "deny", "message": "Session closed"})
|
||||
session.pending_approvals = []
|
||||
|
||||
if hasattr(session, '_cancel_event'):
|
||||
session._cancel_event.set()
|
||||
|
||||
self._sync_session_close(session)
|
||||
|
||||
doc_data = session.model_dump(mode="json")
|
||||
doc_data["search_text"] = self._build_search_text(session)
|
||||
|
||||
_save_session(session_id, doc_data)
|
||||
|
||||
await ws_manager.send_to_session(session_id, "agent:closed", {
|
||||
"session_id": session_id,
|
||||
"status": session.status,
|
||||
"name": session.name,
|
||||
"model": session.model,
|
||||
"mode": session.mode,
|
||||
"created_at": session.created_at.isoformat() if session.created_at else None,
|
||||
"closed_at": session.closed_at.isoformat() if session.closed_at else None,
|
||||
"cost_usd": session.cost_usd,
|
||||
"dashboard_id": session.dashboard_id,
|
||||
})
|
||||
|
||||
self._purge_session_memory(session_id)
|
||||
logger.info(f"Session {session_id} closed and persisted")
|
||||
|
||||
def _purge_session_memory(self, session_id: str) -> None:
|
||||
"""Drop a session from EVERY in-memory structure keyed by its id, so a
|
||||
close or delete can't strand stale per-session state that lives until
|
||||
the process dies. One chokepoint on purpose: a new per-session cache
|
||||
wires its eviction in HERE and both removal paths get it for free."""
|
||||
self.sessions.pop(session_id, None)
|
||||
self.tasks.pop(session_id, None)
|
||||
self._live_partial.pop(session_id, None)
|
||||
view_builder_render_retry_counts.pop(session_id, None)
|
||||
view_builder_dirty_sessions.discard(session_id)
|
||||
|
||||
async def delete_session(self, session_id: str) -> None:
|
||||
"""Permanently delete a session: remove from memory and JSON file.
|
||||
Also stops browser-agent children first."""
|
||||
children = [
|
||||
s for s in self.sessions.values()
|
||||
if s.parent_session_id == session_id and s.mode == "browser-agent"
|
||||
]
|
||||
for child in children:
|
||||
await self.stop_agent(child.id)
|
||||
|
||||
task = self.tasks.get(session_id)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
self._purge_session_memory(session_id)
|
||||
|
||||
_delete_session_file(session_id)
|
||||
logger.info(f"Session {session_id} permanently deleted")
|
||||
|
||||
async def resume_session(self, session_id: str) -> AgentSession:
|
||||
if session_id in self.sessions:
|
||||
return self.sessions[session_id]
|
||||
session = lifecycle.load_session_for_resume(session_id)
|
||||
self.sessions[session_id] = session
|
||||
await ws_manager.send_to_session(session_id, "agent:status", {
|
||||
"session_id": session_id,
|
||||
"status": session.status,
|
||||
"session": session.model_dump(mode="json"),
|
||||
})
|
||||
logger.info(f"Session {session_id} resumed from history")
|
||||
return session
|
||||
|
||||
def get_history(
|
||||
self,
|
||||
q: str = "",
|
||||
limit: int = 20,
|
||||
offset: int = 0,
|
||||
dashboard_id: str | None = None,
|
||||
) -> dict:
|
||||
"""Return paginated, optionally filtered summaries of closed sessions."""
|
||||
all_data = _load_all_session_data()
|
||||
all_data.sort(key=lambda pair: pair[1].get("closed_at") or "", reverse=True)
|
||||
|
||||
q_lower = q.strip().lower()
|
||||
history = []
|
||||
for sid, data in all_data:
|
||||
if dashboard_id and data.get("dashboard_id") != dashboard_id:
|
||||
continue
|
||||
if q_lower:
|
||||
name = (data.get("name") or "").lower()
|
||||
search_text = (data.get("search_text") or "").lower()
|
||||
if q_lower not in name and q_lower not in search_text:
|
||||
continue
|
||||
history.append({
|
||||
"id": data.get("id", sid),
|
||||
"name": data.get("name", "Untitled"),
|
||||
"status": data.get("status", "stopped"),
|
||||
"model": data.get("model", "sonnet"),
|
||||
"mode": data.get("mode", "agent"),
|
||||
"created_at": data.get("created_at"),
|
||||
"closed_at": data.get("closed_at"),
|
||||
"cost_usd": data.get("cost_usd", 0),
|
||||
"dashboard_id": data.get("dashboard_id"),
|
||||
})
|
||||
|
||||
total = len(history)
|
||||
page = history[offset : offset + limit]
|
||||
return {
|
||||
"sessions": page,
|
||||
"total": total,
|
||||
"has_more": offset + limit < total,
|
||||
}
|
||||
|
||||
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)
|
||||
|
||||
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()
|
||||
|
||||
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}")
|
||||
|
||||
async def duplicate_session(self, session_id: str, dashboard_id: str | None = None, up_to_message_id: str | None = None) -> AgentSession:
|
||||
new_session = lifecycle.build_duplicate_session(self.sessions.get(session_id), session_id, dashboard_id, up_to_message_id)
|
||||
self.sessions[new_session.id] = new_session
|
||||
await ws_manager.send_to_session(new_session.id, "agent:status", {
|
||||
"session_id": new_session.id,
|
||||
"status": new_session.status,
|
||||
"session": new_session.model_dump(mode="json"),
|
||||
})
|
||||
return new_session
|
||||
|
||||
async def invoke_agent(
|
||||
self,
|
||||
source_session_id: str,
|
||||
@@ -3182,65 +2956,4 @@ class AgentManager:
|
||||
"cost_usd": fork.cost_usd,
|
||||
}
|
||||
|
||||
def get_all_sessions(self, dashboard_id: str | None = None) -> list[AgentSession]:
|
||||
if not dashboard_id:
|
||||
return list(self.sessions.values())
|
||||
# Memory first, then promote on-disk sessions for this dashboard, but
|
||||
# ONLY ones the dashboard's layout still has a card for. A session keeps
|
||||
# its dashboard_id when its card is deleted, so promoting by tag alone
|
||||
# resurrected deleted chats on every reopen; the layout's cards are the
|
||||
# real source of truth for what's on the board. Imported sessions ARE in
|
||||
# the layout, so they still surface, and this bounds the disk read to
|
||||
# once per session per run, like resume_session.
|
||||
result = [s for s in self.sessions.values() if s.dashboard_id == dashboard_id]
|
||||
seen = {s.id for s in result}
|
||||
card_ids = self._dashboard_card_ids(dashboard_id)
|
||||
for sid, data in _load_all_session_data():
|
||||
if sid in seen or sid not in card_ids:
|
||||
continue
|
||||
if data.get("dashboard_id") != dashboard_id:
|
||||
continue
|
||||
try:
|
||||
sess = AgentSession(**data)
|
||||
except Exception:
|
||||
logger.warning(f"get_all_sessions: skipping unloadable session {sid}", exc_info=True)
|
||||
continue
|
||||
apply_context_window(sess)
|
||||
self.sessions[sid] = sess
|
||||
result.append(sess)
|
||||
return result
|
||||
|
||||
def _dashboard_card_ids(self, dashboard_id: str) -> set[str]:
|
||||
"""Session ids the dashboard's layout currently has agent cards for.
|
||||
Read straight off disk (no dashboards-module import, avoids a cycle)."""
|
||||
try:
|
||||
import os
|
||||
import backend.config.paths as _paths
|
||||
from backend.config.json_store import read_json_or_none
|
||||
d = read_json_or_none(os.path.join(_paths.DASHBOARDS_DIR, f"{dashboard_id}.json")) or {}
|
||||
return set((d.get("layout", {}).get("cards") or {}).keys())
|
||||
except Exception:
|
||||
return set()
|
||||
|
||||
def get_session(self, session_id: str) -> Optional[AgentSession]:
|
||||
return self.sessions.get(session_id)
|
||||
|
||||
def get_browser_agent_children(self, parent_session_id: str) -> list[dict]:
|
||||
"""Return browser-agent sessions for a parent, from memory or disk."""
|
||||
results: list[dict] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
for s in self.sessions.values():
|
||||
if s.mode == "browser-agent" and s.parent_session_id == parent_session_id:
|
||||
results.append(s.model_dump(mode="json"))
|
||||
seen.add(s.id)
|
||||
|
||||
for sid, data in _load_all_session_data():
|
||||
if sid in seen:
|
||||
continue
|
||||
if data.get("mode") == "browser-agent" and data.get("parent_session_id") == parent_session_id:
|
||||
results.append(data)
|
||||
|
||||
return results
|
||||
|
||||
agent_manager = AgentManager()
|
||||
|
||||
@@ -0,0 +1,310 @@
|
||||
"""Session lifecycle + query methods for AgentManager, split out as a mixin so the manager
|
||||
file stays under the size ceiling. Pure relocation: every method reaches self.sessions /
|
||||
self.tasks / self.stop_agent across the MRO exactly as it did inline, so behavior is identical."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
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,
|
||||
build_search_text,
|
||||
)
|
||||
from backend.apps.agents.manager.session.cloud_sync import _sync_session_close
|
||||
from backend.apps.agents.manager.session.apply_context_window import apply_context_window
|
||||
from backend.apps.agents.manager.session import lifecycle
|
||||
from backend.apps.agents.manager.view_builder_state import (
|
||||
view_builder_render_retry_counts,
|
||||
view_builder_dirty_sessions,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SessionLifecycleMixin:
|
||||
@staticmethod
|
||||
def _build_search_text(session: AgentSession, max_len: int = 5000) -> str:
|
||||
return build_search_text(session, max_len)
|
||||
|
||||
def _sync_session_close(self, session: AgentSession, close_reason: str = "user"):
|
||||
_sync_session_close(session, close_reason)
|
||||
|
||||
async def close_session(self, session_id: str) -> None:
|
||||
"""Close a session: pause the agent if running, persist to JSON file,
|
||||
and remove from in-memory state. Also stops browser-agent children."""
|
||||
children = [
|
||||
s for s in self.sessions.values()
|
||||
if s.parent_session_id == session_id and s.mode == "browser-agent"
|
||||
]
|
||||
for child in children:
|
||||
await self.stop_agent(child.id)
|
||||
|
||||
task = self.tasks.get(session_id)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
session = self.sessions.get(session_id)
|
||||
if not session:
|
||||
raise ValueError(f"Session {session_id} not found")
|
||||
|
||||
if session.status in ("running", "waiting_approval"):
|
||||
session.status = "stopped"
|
||||
session.closed_at = datetime.now()
|
||||
|
||||
for req in list(session.pending_approvals):
|
||||
ws_manager.resolve_approval(req.id, {"behavior": "deny", "message": "Session closed"})
|
||||
session.pending_approvals = []
|
||||
|
||||
if hasattr(session, '_cancel_event'):
|
||||
session._cancel_event.set()
|
||||
|
||||
self._sync_session_close(session)
|
||||
|
||||
doc_data = session.model_dump(mode="json")
|
||||
doc_data["search_text"] = self._build_search_text(session)
|
||||
|
||||
_save_session(session_id, doc_data)
|
||||
|
||||
await ws_manager.send_to_session(session_id, "agent:closed", {
|
||||
"session_id": session_id,
|
||||
"status": session.status,
|
||||
"name": session.name,
|
||||
"model": session.model,
|
||||
"mode": session.mode,
|
||||
"created_at": session.created_at.isoformat() if session.created_at else None,
|
||||
"closed_at": session.closed_at.isoformat() if session.closed_at else None,
|
||||
"cost_usd": session.cost_usd,
|
||||
"dashboard_id": session.dashboard_id,
|
||||
})
|
||||
|
||||
self._purge_session_memory(session_id)
|
||||
logger.info(f"Session {session_id} closed and persisted")
|
||||
|
||||
def _purge_session_memory(self, session_id: str) -> None:
|
||||
"""Drop a session from EVERY in-memory structure keyed by its id, so a
|
||||
close or delete can't strand stale per-session state that lives until
|
||||
the process dies. One chokepoint on purpose: a new per-session cache
|
||||
wires its eviction in HERE and both removal paths get it for free."""
|
||||
self.sessions.pop(session_id, None)
|
||||
self.tasks.pop(session_id, None)
|
||||
self._live_partial.pop(session_id, None)
|
||||
view_builder_render_retry_counts.pop(session_id, None)
|
||||
view_builder_dirty_sessions.discard(session_id)
|
||||
|
||||
async def delete_session(self, session_id: str) -> None:
|
||||
"""Permanently delete a session: remove from memory and JSON file.
|
||||
Also stops browser-agent children first."""
|
||||
children = [
|
||||
s for s in self.sessions.values()
|
||||
if s.parent_session_id == session_id and s.mode == "browser-agent"
|
||||
]
|
||||
for child in children:
|
||||
await self.stop_agent(child.id)
|
||||
|
||||
task = self.tasks.get(session_id)
|
||||
if task and not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
self._purge_session_memory(session_id)
|
||||
|
||||
_delete_session_file(session_id)
|
||||
logger.info(f"Session {session_id} permanently deleted")
|
||||
|
||||
async def resume_session(self, session_id: str) -> AgentSession:
|
||||
if session_id in self.sessions:
|
||||
return self.sessions[session_id]
|
||||
session = lifecycle.load_session_for_resume(session_id)
|
||||
self.sessions[session_id] = session
|
||||
await ws_manager.send_to_session(session_id, "agent:status", {
|
||||
"session_id": session_id,
|
||||
"status": session.status,
|
||||
"session": session.model_dump(mode="json"),
|
||||
})
|
||||
logger.info(f"Session {session_id} resumed from history")
|
||||
return session
|
||||
|
||||
def get_history(
|
||||
self,
|
||||
q: str = "",
|
||||
limit: int = 20,
|
||||
offset: int = 0,
|
||||
dashboard_id: str | None = None,
|
||||
) -> dict:
|
||||
"""Return paginated, optionally filtered summaries of closed sessions."""
|
||||
all_data = _load_all_session_data()
|
||||
all_data.sort(key=lambda pair: pair[1].get("closed_at") or "", reverse=True)
|
||||
|
||||
q_lower = q.strip().lower()
|
||||
history = []
|
||||
for sid, data in all_data:
|
||||
if dashboard_id and data.get("dashboard_id") != dashboard_id:
|
||||
continue
|
||||
if q_lower:
|
||||
name = (data.get("name") or "").lower()
|
||||
search_text = (data.get("search_text") or "").lower()
|
||||
if q_lower not in name and q_lower not in search_text:
|
||||
continue
|
||||
history.append({
|
||||
"id": data.get("id", sid),
|
||||
"name": data.get("name", "Untitled"),
|
||||
"status": data.get("status", "stopped"),
|
||||
"model": data.get("model", "sonnet"),
|
||||
"mode": data.get("mode", "agent"),
|
||||
"created_at": data.get("created_at"),
|
||||
"closed_at": data.get("closed_at"),
|
||||
"cost_usd": data.get("cost_usd", 0),
|
||||
"dashboard_id": data.get("dashboard_id"),
|
||||
})
|
||||
|
||||
total = len(history)
|
||||
page = history[offset : offset + limit]
|
||||
return {
|
||||
"sessions": page,
|
||||
"total": total,
|
||||
"has_more": offset + limit < total,
|
||||
}
|
||||
|
||||
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)
|
||||
|
||||
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()
|
||||
|
||||
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}")
|
||||
|
||||
async def duplicate_session(self, session_id: str, dashboard_id: str | None = None, up_to_message_id: str | None = None) -> AgentSession:
|
||||
new_session = lifecycle.build_duplicate_session(self.sessions.get(session_id), session_id, dashboard_id, up_to_message_id)
|
||||
self.sessions[new_session.id] = new_session
|
||||
await ws_manager.send_to_session(new_session.id, "agent:status", {
|
||||
"session_id": new_session.id,
|
||||
"status": new_session.status,
|
||||
"session": new_session.model_dump(mode="json"),
|
||||
})
|
||||
return new_session
|
||||
|
||||
def get_all_sessions(self, dashboard_id: str | None = None) -> list[AgentSession]:
|
||||
if not dashboard_id:
|
||||
return list(self.sessions.values())
|
||||
# Memory first, then promote on-disk sessions for this dashboard, but
|
||||
# ONLY ones the dashboard's layout still has a card for. A session keeps
|
||||
# its dashboard_id when its card is deleted, so promoting by tag alone
|
||||
# resurrected deleted chats on every reopen; the layout's cards are the
|
||||
# real source of truth for what's on the board. Imported sessions ARE in
|
||||
# the layout, so they still surface, and this bounds the disk read to
|
||||
# once per session per run, like resume_session.
|
||||
result = [s for s in self.sessions.values() if s.dashboard_id == dashboard_id]
|
||||
seen = {s.id for s in result}
|
||||
card_ids = self._dashboard_card_ids(dashboard_id)
|
||||
for sid, data in _load_all_session_data():
|
||||
if sid in seen or sid not in card_ids:
|
||||
continue
|
||||
if data.get("dashboard_id") != dashboard_id:
|
||||
continue
|
||||
try:
|
||||
sess = AgentSession(**data)
|
||||
except Exception:
|
||||
logger.warning(f"get_all_sessions: skipping unloadable session {sid}", exc_info=True)
|
||||
continue
|
||||
apply_context_window(sess)
|
||||
self.sessions[sid] = sess
|
||||
result.append(sess)
|
||||
return result
|
||||
|
||||
def _dashboard_card_ids(self, dashboard_id: str) -> set[str]:
|
||||
"""Session ids the dashboard's layout currently has agent cards for.
|
||||
Read straight off disk (no dashboards-module import, avoids a cycle)."""
|
||||
try:
|
||||
import os
|
||||
import backend.config.paths as _paths
|
||||
from backend.config.json_store import read_json_or_none
|
||||
d = read_json_or_none(os.path.join(_paths.DASHBOARDS_DIR, f"{dashboard_id}.json")) or {}
|
||||
return set((d.get("layout", {}).get("cards") or {}).keys())
|
||||
except Exception:
|
||||
return set()
|
||||
|
||||
def get_session(self, session_id: str) -> Optional[AgentSession]:
|
||||
return self.sessions.get(session_id)
|
||||
|
||||
def get_browser_agent_children(self, parent_session_id: str) -> list[dict]:
|
||||
"""Return browser-agent sessions for a parent, from memory or disk."""
|
||||
results: list[dict] = []
|
||||
seen: set[str] = set()
|
||||
|
||||
for s in self.sessions.values():
|
||||
if s.mode == "browser-agent" and s.parent_session_id == parent_session_id:
|
||||
results.append(s.model_dump(mode="json"))
|
||||
seen.add(s.id)
|
||||
|
||||
for sid, data in _load_all_session_data():
|
||||
if sid in seen:
|
||||
continue
|
||||
if data.get("mode") == "browser-agent" and data.get("parent_session_id") == parent_session_id:
|
||||
results.append(data)
|
||||
|
||||
return results
|
||||
|
||||
@@ -77,10 +77,11 @@ def test_read_empty_file_returns_none(tmp_path):
|
||||
|
||||
def test_sessions_skip_corrupt_and_roundtrip(tmp_path, monkeypatch):
|
||||
from backend.apps.agents import agent_manager as am
|
||||
from backend.apps.agents.manager.session.session_store import _load_all_session_data
|
||||
monkeypatch.setattr(am, "SESSIONS_DIR", str(tmp_path))
|
||||
am._save_session("good", {"id": "good", "v": 1})
|
||||
(tmp_path / "bad.json").write_text("{ truncated session ,,,")
|
||||
loaded = dict(am._load_all_session_data())
|
||||
loaded = dict(_load_all_session_data())
|
||||
assert loaded == {"good": {"id": "good", "v": 1}}
|
||||
assert (tmp_path / "bad.json").exists() # corrupt file preserved, not deleted
|
||||
assert am._load_session_data("good") == {"id": "good", "v": 1}
|
||||
|
||||
@@ -11,6 +11,7 @@ runs the id is gone from EVERY structure, while a sibling session is untouched.
|
||||
Run with: backend/.venv/bin/python -m pytest backend/tests/test_session_cleanup.py
|
||||
"""
|
||||
from backend.apps.agents import agent_manager as am
|
||||
from backend.apps.agents.manager import view_builder_state as vbs
|
||||
|
||||
|
||||
def test_purge_session_memory_clears_every_structure():
|
||||
@@ -18,16 +19,16 @@ def test_purge_session_memory_clears_every_structure():
|
||||
mgr.sessions = {"dead": object(), "alive": object()}
|
||||
mgr.tasks = {"dead": object()}
|
||||
mgr._live_partial = {"dead": {"text": "half a reply"}}
|
||||
am.view_builder_render_retry_counts["dead"] = 4
|
||||
am.view_builder_dirty_sessions.add("dead")
|
||||
vbs.view_builder_render_retry_counts["dead"] = 4
|
||||
vbs.view_builder_dirty_sessions.add("dead")
|
||||
|
||||
mgr._purge_session_memory("dead")
|
||||
|
||||
assert "dead" not in mgr.sessions
|
||||
assert "dead" not in mgr.tasks
|
||||
assert "dead" not in mgr._live_partial
|
||||
assert "dead" not in am.view_builder_render_retry_counts
|
||||
assert "dead" not in am.view_builder_dirty_sessions
|
||||
assert "dead" not in vbs.view_builder_render_retry_counts
|
||||
assert "dead" not in vbs.view_builder_dirty_sessions
|
||||
# Only the target id is purged; an unrelated live session survives.
|
||||
assert "alive" in mgr.sessions
|
||||
|
||||
|
||||
Reference in New Issue
Block a user