fix(session): Implement session shelving to persist resume state across account switches

Author: salacoste

src/codex_account_manager/config/manager.py

@@ -205,6 +205,29 @@ def delete_account(self, name: str):
         """Alias for remove_account (Legacy compatibility)."""
         self.remove_account(name)
 
+    def _restore_target_session(self, target_account: Account, vault_slug: str):
+        """Restores sessions from storage to ~/.codex/sessions."""
+        legacy_auth_file = self._get_legacy_auth_file()
+        live_sessions_dir = legacy_auth_file.parent / "sessions"
+        
+        # Ensure it's clean (should be clean from shelving, but just in case)
+        if live_sessions_dir.exists():
+            shutil.rmtree(live_sessions_dir)
+            
+        # Derive canonical key for retrieval
+        if vault_slug == "personal":
+            storage_key = target_account.name
+        else:
+            storage_key = f"{vault_slug}/{target_account.name}"
+            
+        storage_dir = self._get_session_storage_dir(storage_key)
+        
+        if storage_dir.exists():
+             shutil.copytree(storage_dir, live_sessions_dir)
+        else:
+            # If no stored session, just ensure the dir exists empty
+            live_sessions_dir.mkdir(parents=True, exist_ok=True)
+
     def switch_account(self, name: str):
         """Switches active account. Validates existence first."""
         # 1. Validate & Normalize
@@ -216,9 +239,15 @@ def switch_account(self, name: str):
         # Check existence in vault (slugifies internally)
         acc = vault.get_account(acc_slug)
 
-        # 2. Sync Legacy Auth
+        # 2a. Shelve Current Session (Save state of the OLD account)
+        self._shelve_current_session()
+
+        # 2b. Sync Legacy Auth (Update credentials)
         self.sync_legacy_auth(acc)
 
+        # 2c. Restore Target Session (Load state of the NEW account)
+        self._restore_target_session(acc, vault_slug)
+        
         # 3. Update Config with Canonical Reference
         cfg = self.load_config()
 
@@ -231,6 +260,42 @@ def switch_account(self, name: str):
 
         self.save_config(cfg)
 
+    def _get_session_storage_dir(self, account_name: str) -> Path:
+        """Returns the directory where sessions for a specific account are stored."""
+        # We store shelved sessions in the root config dir separated by account
+        return self.root / "shelved_sessions" / slugify(account_name)
+
+    def _shelve_current_session(self):
+        """Moves current ~/.codex/sessions to storage."""
+        cfg = self.load_config()
+        if not cfg.active_account:
+            return # No active account to save session for
+
+        legacy_auth_file = self._get_legacy_auth_file()
+        live_sessions_dir = legacy_auth_file.parent / "sessions"
+        
+        if not live_sessions_dir.exists():
+            return
+
+        # Identify where to save it
+        # We handle 'personal/slug' vs 'slug' normalisation via the config
+        # We just need a unique folder name.
+        storage_dir = self._get_session_storage_dir(cfg.active_account)
+        storage_dir.parent.mkdir(parents=True, exist_ok=True)
+        
+        # If storage dir exists, we merge or replace? 
+        # rmtree+copytree is safest to ensure exact mirroring of state
+        if storage_dir.exists():
+            shutil.rmtree(storage_dir)
+            
+        # Move mechanism: We copytree then rmtree orig, or move. 
+        # shutil.move can fail across filesystems, copytree is robust.
+        # But we want to LEAVE the directory empty for the next guy? No, we restore next.
+        shutil.copytree(live_sessions_dir, storage_dir)
+        
+        # Clean up the live directory to be ready for the new account (or just to leave it clean)
+        shutil.rmtree(live_sessions_dir) 
+
     def sync_legacy_auth(self, account: Account):
         """Writes credentials to ~/.codex/auth.json (Legacy Support)."""
         # This logic used to be implicit in `list` or `status` checks.
@@ -260,11 +325,14 @@ def sync_legacy_auth(self, account: Account):
             f.write(json.dumps(data, indent=2))
 
 
-        # Clear sessions to prevent "Token data not available" error
-        sessions_dir = legacy_dir / "sessions"
-        if sessions_dir.exists():
-            shutil.rmtree(sessions_dir)
-            sessions_dir.mkdir()  # Recreate empty
+        # STOP DELETING SESSIONS
+        # sessions_dir = legacy_dir / "sessions"
+        # if sessions_dir.exists():
+        #     shutil.rmtree(sessions_dir)
+        #     sessions_dir.mkdir()  # Recreate empty
+        
+        # Just ensure the directory exists so Codex doesn't crash if it expects it
+        (legacy_dir / "sessions").mkdir(exist_ok=True)
 
         os.chmod(legacy_auth_file, 0o600)
 

tests/test_session_shelving.py

@@ -0,0 +1,75 @@
+import shutil
+import pytest
+from pathlib import Path
+from codex_account_manager.config.manager import ConfigManager
+from codex_account_manager.config.models import Account, AccountType
+
+@pytest.fixture
+def mock_storage_root(tmp_path):
+    """Creates a temporary home directory for testing."""
+    root = tmp_path / ".codex-accounts"
+    legacy_auth = tmp_path / ".codex" / "auth.json"
+    legacy_auth.parent.mkdir(parents=True, exist_ok=True)
+    return root, legacy_auth
+
+def test_session_persistence(mock_storage_root, monkeypatch):
+    """
+    Verifies that sessions are shelved and restored correctly when switching accounts.
+    """
+    root, legacy_auth = mock_storage_root
+    
+    # Mock environment to map to our temp paths
+    monkeypatch.setenv("HOME", str(root.parent))
+    monkeypatch.setattr("codex_account_manager.config.manager.DEFAULT_CONFIG_ROOT", root)
+    monkeypatch.setattr("codex_account_manager.config.manager.LEGACY_AUTH_FILE", legacy_auth)
+    
+    mgr = ConfigManager(root_path=root)
+    
+    # 1. Create two accounts
+    acc1 = Account(name="personal", api_key="sk-1", type=AccountType.API_KEY)
+    acc2 = Account(name="work", api_key="sk-2", type=AccountType.API_KEY)
+    mgr.save_account(acc1)
+    mgr.save_account(acc2)
+    
+    # 2. Start as 'personal'
+    mgr.switch_account("personal")
+    
+    # 3. Simulate a Codex session file being created
+    live_session_dir = legacy_auth.parent / "sessions"
+    live_session_dir.mkdir(parents=True, exist_ok=True)
+    session_file = live_session_dir / "session-123.json"
+    session_file.write_text('{"id": "session-123", "context": "personal"}')
+    
+    assert session_file.exists(), "Session file should simulate existing in ~/.codex/sessions"
+    
+    # 4. Switch to 'work'
+    mgr.switch_account("work")
+    
+    # ASSERTIONS:
+    # - Live session dir should be cleaned or empty (since 'work' has no previous session)
+    # - Personal session should be in shelving
+    
+    shelved_personal = root / "shelved_sessions" / "personal" / "session-123.json"
+    assert shelved_personal.exists(), "Personal session should be shelved"
+    assert shelved_personal.read_text() == '{"id": "session-123", "context": "personal"}'
+    
+    assert not session_file.exists(), "Live session should be cleared after switch"
+    
+    # 5. Simulate 'work' creating a session
+    live_session_dir.mkdir(parents=True, exist_ok=True) # Ensure dir exists if logic cleared it
+    work_session = live_session_dir / "work-session.json"
+    work_session.write_text('{"id": "work-1", "context": "work"}')
+    
+    # 6. Switch back to 'personal'
+    mgr.switch_account("personal")
+    
+    # ASSERTIONS
+    # - Live session dir should now contain 'session-123.json' (Restored)
+    # - 'work-session.json' should be in shelving for 'work'
+    
+    restored_personal = live_session_dir / "session-123.json"
+    assert restored_personal.exists(), "Personal session should be restored"
+    assert restored_personal.read_text() == '{"id": "session-123", "context": "personal"}'
+    
+    shelved_work = root / "shelved_sessions" / "work" / "work-session.json"
+    assert shelved_work.exists(), "Work session should be shelved"