machineconfig 5.15__py3-none-any.whl → 7.66__py3-none-any.whl

machineconfig/cluster/sessions_managers/wt_local_manager.py

@@ -1,28 +1,30 @@
 #!/usr/bin/env python3
 from datetime import datetime
-import json
-import uuid
 import logging
 import subprocess
 from pathlib import Path
-from typing import TypedDict, Optional, Dict, List, Any
+from typing import Optional, Any
+from rich.console import Console
 from machineconfig.utils.scheduler import Scheduler
 from machineconfig.cluster.sessions_managers.wt_local import WTLayoutGenerator
+from machineconfig.cluster.sessions_managers.wt_utils.wt_helpers import check_wt_session_status
 from machineconfig.utils.schemas.layouts.layout_types import LayoutConfig
-from machineconfig.cluster.sessions_managers.zellij_utils.monitoring_types import (
-    StartResult, GlobalSummary, ActiveSessionInfo, CommandStatus, CommandSummary
+from machineconfig.cluster.sessions_managers.zellij_utils.monitoring_types import StartResult, ActiveSessionInfo
+from machineconfig.cluster.sessions_managers.wt_utils.manager_persistence import (
+    generate_session_id, save_json_file, load_json_file, list_saved_sessions_in_dir, delete_session_dir, ensure_session_dir_exists
 )
+from machineconfig.cluster.sessions_managers.wt_utils.status_reporting import (
+    print_global_summary, print_session_health_status, print_commands_status, calculate_session_summary, calculate_global_summary_from_status
+)
+
 
 
-class WTSessionReport(TypedDict):
-    session_status: Dict[str, Any]  # WT-specific session status
-    commands_status: Dict[str, CommandStatus]
-    summary: CommandSummary
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
+console = Console()
+TMP_SERIALIZATION_DIR = Path.home() / "tmp_results" / "wt_sessions" / "serialized"
 
-TMP_SERIALIZATION_DIR = Path.home().joinpath("tmp_results", "session_manager", "wt", "local_manager")
 
 
 class WTLocalManager:
@@ -39,21 +41,22 @@ class WTLocalManager:
         """
         self.session_name_prefix = session_name_prefix
         self.session_layouts = session_layouts  # Store the original config
-        self.managers: List[WTLayoutGenerator] = []
+        self.managers: list[WTLayoutGenerator] = []
 
         # Create a WTLayoutGenerator for each session
         for layout_config in session_layouts:
-            manager = WTLayoutGenerator()
-            manager.create_wt_layout(layout_config=layout_config, output_dir=None)
+            session_name = layout_config["layoutName"]
+            manager = WTLayoutGenerator(layout_config=layout_config, session_name=session_name)
+            manager.create_layout_file()
             self.managers.append(manager)
 
         logger.info(f"Initialized WTLocalManager with {len(self.managers)} sessions")
 
-    def get_all_session_names(self) -> List[str]:
+    def get_all_session_names(self) -> list[str]:
         """Get all managed session names."""
-        return [manager.session_name for manager in self.managers if manager.session_name is not None]
+        return [manager.session_name for manager in self.managers]
 
-    def start_all_sessions(self) -> Dict[str, StartResult]:
+    def start_all_sessions(self) -> dict[str, StartResult]:
         """Start all Windows Terminal sessions with their layouts."""
         results = {}
         for manager in self.managers:
@@ -84,7 +87,7 @@ class WTLocalManager:
 
         return results
 
-    def kill_all_sessions(self) -> Dict[str, StartResult]:
+    def kill_all_sessions(self) -> dict[str, StartResult]:
         """Kill all managed Windows Terminal sessions."""
         results = {}
         for manager in self.managers:
@@ -128,105 +131,30 @@ class WTLocalManager:
                 commands.append("")
             return "\n".join(commands)
 
-    def check_all_sessions_status(self) -> Dict[str, WTSessionReport]:
-        """Check the status of all sessions and their commands."""
+    def check_all_sessions_status(self) -> dict[str, dict[str, Any]]:
         status_report = {}
-
         for manager in self.managers:
             session_name = manager.session_name or "default"
-
-            # Get session status
-            session_status = WTLayoutGenerator.check_wt_session_status(session_name)
-
-            # Get commands status for this session
+            session_status = check_wt_session_status(session_name)
             commands_status = manager.check_all_commands_status()
-
-            # Calculate summary for this session
-            running_count = sum(1 for status in commands_status.values() if status.get("running", False))
-            total_count = len(commands_status)
-
-            status_report[session_name] = {
-                "session_status": session_status,
-                "commands_status": commands_status,
-                "summary": {"total_commands": total_count, "running_commands": running_count, "stopped_commands": total_count - running_count, "session_healthy": session_status.get("session_exists", False)},
-            }
-
+            summary = calculate_session_summary(commands_status, session_status.get("session_exists", False))
+            status_report[session_name] = {"session_status": session_status, "commands_status": commands_status, "summary": summary}
         return status_report
 
-    def get_global_summary(self) -> GlobalSummary:
-        """Get a global summary across all sessions."""
+    def get_global_summary(self) -> dict[str, Any]:
         all_status = self.check_all_sessions_status()
-
-        total_sessions = len(all_status)
-        healthy_sessions = sum(1 for status in all_status.values() if status["summary"]["session_healthy"])
-        total_commands = sum(status["summary"]["total_commands"] for status in all_status.values())
-        total_running = sum(status["summary"]["running_commands"] for status in all_status.values())
-
-        return {
-            "total_sessions": total_sessions,
-            "healthy_sessions": healthy_sessions,
-            "unhealthy_sessions": total_sessions - healthy_sessions,
-            "total_commands": total_commands,
-            "running_commands": total_running,
-            "stopped_commands": total_commands - total_running,
-            "all_sessions_healthy": healthy_sessions == total_sessions,
-            "all_commands_running": total_running == total_commands,
-        }
+        return calculate_global_summary_from_status(all_status, include_remote_machines=False)
 
     def print_status_report(self) -> None:
-        """Print a comprehensive status report for all sessions."""
         all_status = self.check_all_sessions_status()
         global_summary = self.get_global_summary()
-
-        print("=" * 80)
-        print("🖥️  WINDOWS TERMINAL LOCAL MANAGER STATUS REPORT")
-        print("=" * 80)
-
-        # Global summary
-        print("🌐 GLOBAL SUMMARY:")
-        print(f"   Total sessions: {global_summary['total_sessions']}")
-        print(f"   Healthy sessions: {global_summary['healthy_sessions']}")
-        print(f"   Total commands: {global_summary['total_commands']}")
-        print(f"   Running commands: {global_summary['running_commands']}")
-        print(f"   All healthy: {'✅' if global_summary['all_sessions_healthy'] else '❌'}")
-        print()
-
-        # Per-session details
+        print_global_summary(global_summary, "WINDOWS TERMINAL LOCAL MANAGER STATUS REPORT")
         for session_name, status in all_status.items():
-            session_status = status["session_status"]
-            commands_status = status["commands_status"]
-            summary = status["summary"]
-
             print(f"🪟 SESSION: {session_name}")
             print("-" * 60)
-
-            # Session health
-            if session_status.get("wt_running", False):
-                if session_status.get("session_exists", False):
-                    session_windows = session_status.get("session_windows", [])
-                    all_windows = session_status.get("all_windows", [])
-                    print("✅ Windows Terminal is running")
-                    print(f"   Session windows: {len(session_windows)}")
-                    print(f"   Total WT windows: {len(all_windows)}")
-                else:
-                    print("⚠️  Windows Terminal is running but no session windows found")
-            else:
-                print(f"❌ Windows Terminal session issue: {session_status.get('error', 'Unknown error')}")
-
-            # Commands in this session
-            print(f"   Commands ({summary['running_commands']}/{summary['total_commands']} running):")
-            for tab_name, cmd_status in commands_status.items():
-                status_icon = "✅" if cmd_status.get("running", False) else "❌"
-                cmd_text = cmd_status.get("command", "Unknown")[:50]
-                if len(cmd_status.get("command", "")) > 50:
-                    cmd_text += "..."
-                print(f"     {status_icon} {tab_name}: {cmd_text}")
-
-                if cmd_status.get("processes"):
-                    for proc in cmd_status["processes"][:2]:  # Show first 2 processes
-                        print(f"        └─ PID {proc['pid']}: {proc['name']}")
+            print_session_health_status(status["session_status"], remote_name=None)
+            print_commands_status(status["commands_status"], status["summary"])
             print()
-
         print("=" * 80)
 
     def run_monitoring_routine(self, wait_ms: int = 30000) -> None:
@@ -288,128 +216,59 @@ class WTLocalManager:
         sched.run(max_cycles=None)
 
     def save(self, session_id: Optional[str] = None) -> str:
-        """Save the manager state to disk."""
         if session_id is None:
-            session_id = str(uuid.uuid4())[:8]
-
-        # Create session directory
+            session_id = generate_session_id()
         session_dir = TMP_SERIALIZATION_DIR / session_id
-        session_dir.mkdir(parents=True, exist_ok=True)
-
-        # Save the session2wt_tabs configuration
-        config_file = session_dir / "session_layouts.json"
-        text = json.dumps(self.session_layouts, indent=2, ensure_ascii=False)
-        config_file.write_text(text, encoding="utf-8")
-
-        # Save metadata
+        ensure_session_dir_exists(session_dir)
+        save_json_file(session_dir / "session_layouts.json", self.session_layouts, "session layouts")
         metadata = {"session_name_prefix": self.session_name_prefix, "created_at": str(datetime.now()), "num_managers": len(self.managers), "sessions": [item["layoutName"] for item in self.session_layouts], "manager_type": "WTLocalManager"}
-        metadata_file = session_dir / "metadata.json"
-        text = json.dumps(metadata, indent=2, ensure_ascii=False)
-        metadata_file.write_text(text, encoding="utf-8")
-
-        # Save each manager's state
+        save_json_file(session_dir / "metadata.json", metadata, "metadata")
         managers_dir = session_dir / "managers"
         managers_dir.mkdir(exist_ok=True)
-
         for i, manager in enumerate(self.managers):
             manager_data = {"session_name": manager.session_name, "layout_config": manager.layout_config, "script_path": manager.script_path}
-            manager_file = managers_dir / f"manager_{i}_{manager.session_name}.json"
-            text = json.dumps(manager_data, indent=2, ensure_ascii=False)
-            manager_file.write_text(text, encoding="utf-8")
-
+            save_json_file(managers_dir / f"manager_{i}_{manager.session_name}.json", manager_data, f"manager {i}")
         logger.info(f"✅ Saved WTLocalManager session to: {session_dir}")
         return session_id
 
     @classmethod
     def load(cls, session_id: str) -> "WTLocalManager":
-        """Load a saved manager state from disk."""
         session_dir = TMP_SERIALIZATION_DIR / session_id
-
         if not session_dir.exists():
             raise FileNotFoundError(f"Session directory not found: {session_dir}")
-
-        # Load configuration
-        config_file = session_dir / "session_layouts.json"
-        if not config_file.exists():
-            raise FileNotFoundError(f"Configuration file not found: {config_file}")
-
-        text = config_file.read_text(encoding="utf-8")
-        session_layouts = json.loads(text)
-
-        # Load metadata
-        metadata_file = session_dir / "metadata.json"
-        session_name_prefix = "LocalWTMgr"  # default fallback
-        if metadata_file.exists():
-            text = metadata_file.read_text(encoding="utf-8")
-            metadata = json.loads(text)
-            session_name_prefix = metadata.get("session_name_prefix", "LocalWTMgr")
-
-        # Create new instance
+        loaded_data = load_json_file(session_dir / "session_layouts.json", "Configuration file")
+        session_layouts = loaded_data if isinstance(loaded_data, list) else []  # type: ignore[arg-type]
+        metadata_data = load_json_file(session_dir / "metadata.json", "Metadata file") if (session_dir / "metadata.json").exists() else {}
+        metadata = metadata_data if isinstance(metadata_data, dict) else {}  # type: ignore[arg-type]
+        session_name_prefix = metadata.get("session_name_prefix", "LocalWTMgr")  # type: ignore[union-attr]
         instance = cls(session_layouts=session_layouts, session_name_prefix=session_name_prefix)
-
-        # Load saved manager states
         managers_dir = session_dir / "managers"
         if managers_dir.exists():
             instance.managers = []
-            manager_files = sorted(managers_dir.glob("manager_*.json"))
-
-            for manager_file in manager_files:
+            for manager_file in sorted(managers_dir.glob("manager_*.json")):
                 try:
-                    text = manager_file.read_text(encoding="utf-8")
-                    manager_data = json.loads(text)
-
-                    # Recreate the manager
-                    manager = WTLayoutGenerator()
-                    manager.session_name = manager_data["session_name"]
-                    manager.layout_config = manager_data["layout_config"]
-                    manager.script_path = manager_data["script_path"]
-
+                    loaded_manager_data = load_json_file(manager_file, "Manager data")
+                    manager_data = loaded_manager_data if isinstance(loaded_manager_data, dict) else {}  # type: ignore[arg-type]
+                    manager = WTLayoutGenerator(layout_config=manager_data["layout_config"], session_name=manager_data["session_name"])  # type: ignore[typeddict-item]
+                    manager.script_path = manager_data["script_path"]  # type: ignore[typeddict-item]
                     instance.managers.append(manager)
-
                 except Exception as e:
                     logger.warning(f"Failed to load manager from {manager_file}: {e}")
-
         logger.info(f"✅ Loaded WTLocalManager session from: {session_dir}")
         return instance
 
     @staticmethod
-    def list_saved_sessions() -> List[str]:
-        """List all saved session IDs."""
-        if not TMP_SERIALIZATION_DIR.exists():
-            return []
-
-        sessions = []
-        for item in TMP_SERIALIZATION_DIR.iterdir():
-            if item.is_dir() and (item / "metadata.json").exists():
-                sessions.append(item.name)
-
-        return sorted(sessions)
+    def list_saved_sessions() -> list[str]:
+        return list_saved_sessions_in_dir(TMP_SERIALIZATION_DIR)
 
     @staticmethod
     def delete_session(session_id: str) -> bool:
-        """Delete a saved session."""
-        session_dir = TMP_SERIALIZATION_DIR / session_id
-
-        if not session_dir.exists():
-            logger.warning(f"Session directory not found: {session_dir}")
-            return False
-
-        try:
-            import shutil
-
-            shutil.rmtree(session_dir)
-            logger.info(f"✅ Deleted session: {session_id}")
-            return True
-        except Exception as e:
-            logger.error(f"Failed to delete session {session_id}: {e}")
-            return False
+        return delete_session_dir(TMP_SERIALIZATION_DIR / session_id, session_id)
 
-    def list_active_sessions(self) -> List[ActiveSessionInfo]:
-        """List currently active Windows Terminal sessions managed by this instance."""
-        active_sessions = []
+    def list_active_sessions(self) -> list[ActiveSessionInfo]:
+        active_sessions: list[ActiveSessionInfo] = []
 
         try:
-            # Get all running Windows Terminal processes
             result = subprocess.run(
                 ["powershell", "-Command", 'Get-Process -Name "WindowsTerminal" -ErrorAction SilentlyContinue | Select-Object Id, ProcessName, StartTime, MainWindowTitle | ConvertTo-Json -Depth 2'], capture_output=True, text=True, timeout=10
             )
@@ -421,7 +280,6 @@ class WTLocalManager:
                 if not isinstance(all_processes, list):
                     all_processes = [all_processes]
 
-                # Filter to only our managed sessions
                 for manager in self.managers:
                     session_name = manager.session_name
                     session_windows = []
@@ -437,7 +295,6 @@ class WTLocalManager:
                             "is_active": len(session_windows) > 0,
                             "tab_count": len(manager.layout_config["layoutTabs"]) if manager.layout_config else 0,
                             "tabs": [tab["tabName"] for tab in manager.layout_config["layoutTabs"]] if manager.layout_config else [],
-                            "windows": session_windows,
                         }
                     )
 
@@ -446,7 +303,7 @@ class WTLocalManager:
 
         return active_sessions
 
-    def get_wt_overview(self) -> Dict[str, Any]:
+    def get_wt_overview(self) -> dict[str, Any]:
         """Get overview of all Windows Terminal windows and processes."""
         try:
             result = subprocess.run(

machineconfig/cluster/sessions_managers/wt_remote.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-from typing import Dict, Optional, List, Any
+from typing import Optional, Any
 from pathlib import Path
 import logging
 import json
@@ -11,18 +11,18 @@ from machineconfig.cluster.sessions_managers.wt_utils.layout_generator import WT
 from machineconfig.cluster.sessions_managers.wt_utils.process_monitor import WTProcessMonitor
 from machineconfig.cluster.sessions_managers.wt_utils.session_manager import WTSessionManager
 from machineconfig.cluster.sessions_managers.wt_utils.status_reporter import WTStatusReporter
-from machineconfig.utils.schemas.layouts.layout_types import TabConfig
+from machineconfig.utils.schemas.layouts.layout_types import TabConfig, LayoutConfig
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-TMP_LAYOUT_DIR = Path.home().joinpath("tmp_results", "wt_layouts", "layout_manager")
+TMP_LAYOUT_DIR = Path.home() / "tmp_results" / "wt_layouts"
 
 
 class WTRemoteLayoutGenerator:
-    def __init__(self, remote_name: str, session_name_prefix: str):
-        self.remote_name = remote_name
-        self.session_name = session_name_prefix + "_" + WTLayoutGenerator.generate_random_suffix()
-        self.tabs: List[TabConfig] = []
+    def __init__(self, layout_config: LayoutConfig, remote_name: str, session_name: str):
+        self.remote_name: str = remote_name
+        self.session_name: str = session_name
+        self.layout_config: LayoutConfig = layout_config.copy()
         self.script_path: Optional[str] = None
 
         # Initialize modular components
@@ -32,22 +32,30 @@ class WTRemoteLayoutGenerator:
         self.session_manager = WTSessionManager(self.remote_executor, self.session_name, TMP_LAYOUT_DIR)
         self.status_reporter = WTStatusReporter(self.process_monitor, self.session_manager)
 
-    # Tabs are stored and used as List[TabConfig]; no legacy dict compatibility
-
-    def create_wt_layout(self, tabs: List[TabConfig], output_dir: Optional[str]) -> str:
-        logger.info(f"Creating Windows Terminal layout with {len(tabs)} tabs for remote '{self.remote_name}'")
-        self.tabs = tabs
-        if output_dir:
-            output_path = Path(output_dir)
-        else:
-            output_path = TMP_LAYOUT_DIR
-        self.script_path = self.layout_generator.create_wt_script(self.tabs, output_path, self.session_name)
-        return self.script_path
+    def create_layout_file(self) -> bool:
+        """Create Windows Terminal layout file and return success status."""
+        tab_count = len(self.layout_config["layoutTabs"])
+        logger.info(f"Creating Windows Terminal layout with {tab_count} tabs for remote '{self.remote_name}'")
+        
+        # Extract tabs from layout_config
+        tabs: list[TabConfig] = self.layout_config["layoutTabs"]
+        script_content = self.layout_generator.create_wt_script(tabs, self.session_name, window_name=None)
+        
+        # Write to file
+        tmp_layout_dir = Path.home() / "tmp_results" / "wt_layouts" / "remote"
+        tmp_layout_dir.mkdir(parents=True, exist_ok=True)
+        random_suffix = WTLayoutGenerator.generate_random_suffix(8)
+        script_file = tmp_layout_dir / f"wt_layout_{self.session_name}_{random_suffix}.ps1"
+        script_file.write_text(script_content, encoding="utf-8")
+        self.script_path = str(script_file.absolute())
+        
+        logger.info(f"✅ Remote layout created: {self.script_path}")
+        return True
 
     # Legacy methods for backward compatibility
 
-    def to_dict(self) -> Dict[str, Any]:
-        return {"remote_name": self.remote_name, "session_name": self.session_name, "tabs": self.tabs, "script_path": self.script_path, "created_at": datetime.now().isoformat(), "class_name": self.__class__.__name__}
+    def to_dict(self) -> dict[str, Any]:
+        return {"remote_name": self.remote_name, "session_name": self.session_name, "layout_config": self.layout_config, "script_path": self.script_path, "created_at": datetime.now().isoformat(), "class_name": self.__class__.__name__}
 
     def to_json(self, file_path: Optional[str]) -> str:
         # Generate file path if not provided
@@ -95,29 +103,18 @@ class WTRemoteLayoutGenerator:
             logger.warning(f"Class name mismatch: expected {cls.__name__}, got {data.get('class_name')}")
 
         # Create new instance
-        # Extract session name prefix by removing the suffix
-        session_name = data["session_name"]
-        if "_" in session_name:
-            session_name_prefix = "_".join(session_name.split("_")[:-1])
-        else:
-            session_name_prefix = session_name
-
-        instance = cls(remote_name=data["remote_name"], session_name_prefix=session_name_prefix)
-
-        # Restore state
-        instance.session_name = data["session_name"]
-        # New schema only
-        if "tabs" in data:
-            instance.tabs = data["tabs"]
-        else:
-            instance.tabs = []
+        instance = cls(
+            layout_config=data["layout_config"],
+            remote_name=data["remote_name"],
+            session_name=data["session_name"]
+        )
         instance.script_path = data["script_path"]
 
         logger.info(f"✅ Loaded WTRemoteLayoutGenerator from: {file_path}")
         return instance
 
     @staticmethod
-    def list_saved_sessions(directory_path: Optional[str]) -> List[str]:
+    def list_saved_sessions(directory_path: Optional[str]) -> list[str]:
         if directory_path is None:
             dir_path = Path.home() / "tmp_results" / "wt_sessions" / "serialized"
         else:
@@ -132,7 +129,7 @@ class WTRemoteLayoutGenerator:
 
 if __name__ == "__main__":
     # Example usage
-    sample_tabs: List[TabConfig] = [
+    sample_tabs: list[TabConfig] = [
         {"tabName": "🤖Bot1", "startDir": "~/code/bytesense/bithence", "command": "python bot1.py --create_new_bot True"},
         {"tabName": "🤖Bot2", "startDir": "~/code/bytesense/bithence", "command": "python bot2.py --create_new_bot True"},
         {"tabName": "📊Monitor", "startDir": "~", "command": "Get-Process | Sort-Object CPU -Descending | Select-Object -First 10"},
@@ -144,10 +141,21 @@ if __name__ == "__main__":
     session_name = "test_remote_session"
 
     try:
+        # Create layout config from tabs
+        sample_layout: LayoutConfig = {
+            "layoutName": "RemoteBots",
+            "layoutTabs": sample_tabs
+        }
+        
         # Create layout using the remote generator
-        generator = WTRemoteLayoutGenerator(remote_name=remote_name, session_name_prefix=session_name)
-        script_path = generator.create_wt_layout(sample_tabs, None)
-        print(f"✅ Remote layout created successfully: {script_path}")
+        generator = WTRemoteLayoutGenerator(
+            layout_config=sample_layout,
+            remote_name=remote_name,
+            session_name=session_name
+        )
+        generator.create_layout_file()
+        
+        print(f"✅ Remote layout created successfully: {generator.script_path}")
 
         # Check if Windows Terminal is available on remote
         wt_available = generator.remote_executor.check_wt_available()
@@ -170,7 +178,7 @@ if __name__ == "__main__":
         # Demonstrate loading (using the full path)
         loaded_generator = WTRemoteLayoutGenerator.from_json(saved_path)
         print(f"✅ Session loaded successfully: {loaded_generator.session_name}")
-        print(f"📊 Loaded tabs: {[tab['tabName'] for tab in loaded_generator.tabs]}")
+        print(f"📊 Loaded tabs: {[tab['tabName'] for tab in loaded_generator.layout_config['layoutTabs']]}")
 
         # Show command preview
         preview = generator.layout_generator.generate_wt_command(sample_tabs)
@@ -178,7 +186,7 @@ if __name__ == "__main__":
 
         # Demonstrate status checking
         print(f"\n🔍 Checking command status on remote '{remote_name}':")
-        generator.status_reporter.print_status_report(generator.tabs)
+        generator.status_reporter.print_status_report(sample_tabs)
 
         # Show Windows Terminal overview
         print("\n🖥️  Windows Terminal Overview:")