diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 52f07344..d331304a 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -4489,9 +4489,19 @@ class AgentManager: "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) - logger.info(f"Session {session_id} closed and persisted") + self._live_partial.pop(session_id, None) + p_view_builder_render_retry_counts.pop(session_id, None) + p_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. @@ -4511,8 +4521,7 @@ class AgentManager: except asyncio.CancelledError: pass - self.sessions.pop(session_id, None) - self.tasks.pop(session_id, None) + self._purge_session_memory(session_id) _delete_session_file(session_id) logger.info(f"Session {session_id} permanently deleted") diff --git a/backend/tests/test_session_cleanup.py b/backend/tests/test_session_cleanup.py new file mode 100644 index 00000000..a31d4f1b --- /dev/null +++ b/backend/tests/test_session_cleanup.py @@ -0,0 +1,40 @@ +"""Invariant: closing or deleting a session must strand NO per-session state. + +The orchestration core keeps several maps keyed by session id (the session +record, its asyncio task, the live partial-stream mirror, and two module-level +view-builder retry/dirty structures). Removal used to pop only `sessions` + +`tasks`, leaking the rest for the life of the process, an unbounded creep over +a long-running app. `_purge_session_memory` is the single chokepoint both the +close and delete paths route through; this pins the invariant that after it +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 + + +def test_purge_session_memory_clears_every_structure(): + mgr = am.AgentManager() + mgr.sessions = {"dead": object(), "alive": object()} + mgr.tasks = {"dead": object()} + mgr._live_partial = {"dead": {"text": "half a reply"}} + am.p_view_builder_render_retry_counts["dead"] = 4 + am.p_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.p_view_builder_render_retry_counts + assert "dead" not in am.p_view_builder_dirty_sessions + # Only the target id is purged; an unrelated live session survives. + assert "alive" in mgr.sessions + + +def test_purge_is_safe_on_an_untracked_id(): + # Purging an id that was never tracked must be a quiet no-op, not a KeyError, + # so the delete/close paths can call it unconditionally. + mgr = am.AgentManager() + mgr._purge_session_memory("never-existed") + assert mgr.sessions == {}