feat: Unify session management across MCP codebase (#538)

* feat: Phase 1 - Add data storage capabilities to unified session system

Enhance utils/session/ with workflow data storage:
- Add `data` field to Session model for key-value storage
- Add get(), set(), delete(), clear_data() methods to Session
- Update FileSystemSessionStorage to persist data field to data.json
- Update MemorySessionStorage to handle data field
- Add comprehensive tests for data storage functionality
- Update serialization (to_dict/from_dict) to include data field

This is Phase 1 of the session unification plan, enabling the session
system to support both conversation tracking and workflow data storage.

All 60 session tests pass including 6 new data storage tests.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* feat: Phase 2 - Migrate utils/agent/ to use SessionManager

Replace in-memory agent sessions with unified SessionManager:
- Add SessionManager integration to SpecializedAgent
- Create _create_storage() method for storage backend selection
- Use FileSystemSessionStorage when session_storage_path is set
- Use MemorySessionStorage otherwise (backward compatible)
- Update _get_session_history() to convert ConversationMessage to dict
- Update _add_to_history() to use add_conversation_message()
- Update clear_session_history() to use SessionManager API
- Update get_session_history() to return same format as before
- Update set_session() to ensure session exists
- Fix test for _sessions attribute (now _session_manager)
- Fix test with invalid filesystem path to use temp directory

All 98 agent tests pass + 24 ExploreAgent tests pass.
Agents now support both memory and filesystem persistence!

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* feat: Phase 3 - Delete old automation session module

Remove plugins/automation/runtime_data/session.py:
- Delete SessionData, SessionStorage, get_storage() module
- Update runtime_data/__init__.py to remove session exports
- Add note directing users to utils.session.SessionManager
- No breaking changes: module was not yet in use

All 89 passing automation tests still pass.
4 pre-existing workflow engine test failures unchanged.

This completes the session unification migration. The codebase now
has a single unified session system in utils/session/ supporting:
- Conversation tracking (for agents and tools)
- Key-value data storage (for workflows)
- Metadata and statistics
- Multiple storage backends (memory, filesystem)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

* docs: Update session unification plan with completion summary

Mark plan as completed and add implementation summary:
- All 3 phases completed successfully in ~1 hour
- 247 tests passing (60 session + 98 agent + 89 automation)
- ~160 lines of duplicate code deleted
- Zero breaking changes
- New capabilities: persistent agent sessions, unified data storage

Phase 1 (30 min): Enhanced utils/session/ with data storage
Phase 2 (20 min): Migrated utils/agent/ to SessionManager
Phase 3 (10 min): Deleted automation session module

Key achievements:
✅ Single unified session system across entire codebase
✅ Agents support both memory and filesystem persistence
✅ Sessions support conversation tracking + data storage
✅ No backward compatibility issues

Much faster than estimated 1.5-2 weeks due to:
- Experimental automation plugin allowed direct deletion
- Well-tested foundation (54 existing session tests)
- Clean internal separation in agent implementation

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

---------

Co-authored-by: Claude <[email protected]>

Author: plusplusoneplusplus

docs/session_unification_plan.md

@@ -1,24 +1,19 @@
 """
 Runtime Data Module
 
-Manages runtime data for automation workflows, including:
+Manages runtime data for automation workflows:
 - State: Structural context (WorkflowContext, StepResult, StepStatus)
-- Session: Session-level storage and persistence
+
+Note: Session management has been unified with utils.session.
+Use utils.session.SessionManager for session-level storage.
 """
 
 # State management (structural context)
 from .state import StepResult, StepStatus, WorkflowContext
 
-# Session storage
-from .session import SessionData, SessionStorage, get_storage
-
 __all__ = [
     # State
     "StepResult",
     "StepStatus",
     "WorkflowContext",
-    # Session
-    "SessionData",
-    "SessionStorage",
-    "get_storage",
 ]

plugins/automation/runtime_data/__init__.py

@@ -13,6 +13,7 @@
 
 from .cli_executor import CLIExecutor, CLIConfig, CLIType
 from .system_prompt import SystemPromptBuilder
+from utils.session import SessionManager, FileSystemSessionStorage, MemorySessionStorage, SessionStorage
 
 logger = logging.getLogger(__name__)
 
@@ -87,7 +88,30 @@ def __init__(self, config: AgentConfig):
         """
         self.config = config
         self._executor = CLIExecutor(config.to_cli_config())
-        self._sessions: Dict[str, List[Dict[str, str]]] = {}
+
+        # Use session manager instead of raw dict
+        storage = self._create_storage(config)
+        self._session_manager = SessionManager(storage)
+
+        # Ensure session exists
+        session_id = self._get_session_id()
+        if not self._session_manager.storage.session_exists(session_id):
+            self._session_manager.create_session(
+                session_id=session_id,
+                purpose=f"{self.__class__.__name__} session"
+            )
+
+    def _create_storage(self, config: AgentConfig) -> SessionStorage:
+        """Create storage backend based on config."""
+        if config.session_storage_path:
+            # Use filesystem storage for persistence
+            return FileSystemSessionStorage(
+                sessions_dir=config.session_storage_path,
+                history_dir=config.session_storage_path.parent / ".history"
+            )
+        else:
+            # Use memory storage (backward compatible)
+            return MemorySessionStorage()
 
     @abstractmethod
     def get_system_prompt(self) -> str:
@@ -176,14 +200,23 @@ def _get_session_storage_path(self) -> Path:
     def _get_session_history(self) -> List[Dict[str, str]]:
         """Get the message history for the current session"""
         session_id = self._get_session_id()
-        if session_id not in self._sessions:
-            self._sessions[session_id] = []
-        return self._sessions[session_id]
+        session = self._session_manager.get_session(session_id)
+        if session is None:
+            return []
+        # Convert ConversationMessage to dict format
+        return [
+            {"role": msg.role, "content": msg.content}
+            for msg in session.conversation
+        ]
 
     def _add_to_history(self, role: str, content: str):
         """Add a message to the session history"""
-        history = self._get_session_history()
-        history.append({"role": role, "content": content})
+        session_id = self._get_session_id()
+        self._session_manager.add_conversation_message(
+            session_id=session_id,
+            role=role,
+            content=content
+        )
 
     def _build_prompt(
         self,
@@ -307,8 +340,10 @@ def clear_session_history(self, session_id: Optional[str] = None):
             session_id: Session to clear (if None, clears current session)
         """
         sid = session_id or self._get_session_id()
-        if sid in self._sessions:
-            self._sessions[sid] = []
+        session = self._session_manager.get_session(sid)
+        if session:
+            session.conversation.clear()
+            self._session_manager.storage.save_session(session)
             logger.info(f"Cleared history for session {sid}")
 
     def get_session_history(self, session_id: Optional[str] = None) -> List[Dict[str, str]]:
@@ -322,7 +357,14 @@ def get_session_history(self, session_id: Optional[str] = None) -> List[Dict[str
             List of message dictionaries with 'role' and 'content' keys
         """
         sid = session_id or self._get_session_id()
-        return self._sessions.get(sid, []).copy()
+        session = self._session_manager.get_session(sid)
+        if session is None:
+            return []
+        # Convert ConversationMessage to dict format
+        return [
+            {"role": msg.role, "content": msg.content}
+            for msg in session.conversation
+        ]
 
     def set_session(self, session_id: str):
         """
@@ -332,6 +374,12 @@ def set_session(self, session_id: str):
             session_id: ID of the session to switch to
         """
         self.config.session_id = session_id
+        # Ensure session exists
+        if not self._session_manager.storage.session_exists(session_id):
+            self._session_manager.create_session(
+                session_id=session_id,
+                purpose=f"{self.__class__.__name__} session"
+            )
         logger.info(f"Switched to session {session_id}")
 
     def __repr__(self) -> str:

plugins/automation/runtime_data/session.py

@@ -122,7 +122,10 @@ def test_init(self):
         assert agent.config == config
         assert agent._executor is not None
         assert isinstance(agent._executor, CLIExecutor)
-        assert agent._sessions == {}
+        assert agent._session_manager is not None
+        # Verify session was created
+        session_id = agent._get_session_id()
+        assert agent._session_manager.storage.session_exists(session_id)
 
     def test_get_system_prompt(self):
         """Test getting system prompt"""

utils/agent/agent.py

@@ -130,20 +130,24 @@ def test_agent_session_storage_path_default(self):
 
     def test_agent_session_storage_path_custom(self):
         """Test agent with custom session storage path"""
-        custom_path = Path("/custom/sessions")
-        config = AgentConfig(
-            cli_type=CLIType.COPILOT,
-            session_id="custom_path_test",
-            session_storage_path=custom_path,
-        )
-
-        agent = DefaultSystemPromptAgent(config)
-
-        # Get storage path
-        storage_path = agent._get_session_storage_path()
-
-        # Should use custom path + session_id
-        assert storage_path == custom_path / "custom_path_test"
+        import tempfile
+        with tempfile.TemporaryDirectory() as tmpdir:
+            custom_path = Path(tmpdir) / "sessions"
+            config = AgentConfig(
+                cli_type=CLIType.COPILOT,
+                session_id="custom_path_test",
+                session_storage_path=custom_path,
+            )
+
+            agent = DefaultSystemPromptAgent(config)
+
+            # Get storage path
+            storage_path = agent._get_session_storage_path()
+
+            # Should use custom path + session_id
+            assert storage_path == custom_path / "custom_path_test"
+            # Verify session directory was created
+            assert custom_path.exists()
 
     def test_agent_default_session_id(self):
         """Test agent with default session ID"""

utils/agent/tests/test_agent.py

@@ -116,6 +116,8 @@ class Session:
     metadata: SessionMetadata
     invocation_ids: List[str] = field(default_factory=list)
     conversation: List[ConversationMessage] = field(default_factory=list)
+    # NEW: Workflow data storage
+    data: Dict[str, Any] = field(default_factory=dict)
 
     def add_invocation(self, invocation_id: str, duration_ms: Optional[float] = None):
         """Link an invocation to this session"""
@@ -144,12 +146,34 @@ def add_message(
         self.conversation.append(message)
         self.metadata.updated_at = datetime.now()
 
+    # NEW: Data storage methods (from automation/session.py)
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get value from session data."""
+        return self.data.get(key, default)
+
+    def set(self, key: str, value: Any) -> None:
+        """Set value in session data."""
+        self.data[key] = value
+        self.metadata.updated_at = datetime.now()
+
+    def delete(self, key: str) -> None:
+        """Delete value from session data."""
+        if key in self.data:
+            del self.data[key]
+            self.metadata.updated_at = datetime.now()
+
+    def clear_data(self) -> None:
+        """Clear all session data (preserves conversation)."""
+        self.data.clear()
+        self.metadata.updated_at = datetime.now()
+
     def to_dict(self) -> Dict[str, Any]:
         """Convert session to dictionary for serialization"""
         return {
             "metadata": self.metadata.to_dict(),
             "invocation_ids": self.invocation_ids,
             "conversation": [msg.to_dict() for msg in self.conversation],
+            "data": self.data,
         }
 
     @classmethod
@@ -161,4 +185,5 @@ def from_dict(cls, data: Dict[str, Any]) -> "Session":
             conversation=[
                 ConversationMessage.from_dict(msg) for msg in data.get("conversation", [])
             ],
+            data=data.get("data", {}),
         )