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

machineconfig/cluster/sessions_managers/zellij_local.py

@@ -1,31 +1,22 @@
 #!/usr/bin/env python3
-import shlex
-import subprocess
-from machineconfig.cluster.sessions_managers.zellij_utils.monitoring_types import CommandStatus, ZellijSessionStatus, ComprehensiveStatus, ProcessInfo
-import psutil
-import random
-import string
-from typing import List, Optional
+from typing import Optional
 from pathlib import Path
 import logging
 
 from rich.console import Console
 
-from machineconfig.utils.schemas.layouts.layout_types import LayoutConfig, TabConfig
+from machineconfig.utils.schemas.layouts.layout_types import LayoutConfig
+from machineconfig.cluster.sessions_managers.zellij_utils.monitoring_types import ComprehensiveStatus, CommandStatus
+from machineconfig.cluster.sessions_managers.zellij_utils.zellij_local_helper import validate_layout_config, create_tab_section, check_command_status, check_zellij_session_status
+from machineconfig.cluster.sessions_managers.zellij_utils.zellij_local_helper_restart import restart_tab_process
 
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 console = Console()
-TMP_LAYOUT_DIR = Path.home().joinpath("tmp_results", "session_manager", "zellij", "layout_manager")
 
 
-class ZellijLayoutGenerator:
-    def __init__(self):
-        self.session_name: Optional[str] = None
-        self.layout_config: Optional[LayoutConfig] = None  # Store the complete layout config
-        self.layout_path: Optional[str] = None  # Store the full path to the layout file
-        self.layout_template = """layout {
+DEFAULT_LAYOUT_TEMPLATE = """layout {
     default_tab_template {
         // the default zellij tab-bar and status bar plugins
         pane size=1 borderless=true {
@@ -35,303 +26,47 @@ class ZellijLayoutGenerator:
     }
 """
 
-    @staticmethod
-    def _generate_random_suffix(length: int) -> str:
-        """Generate a random string suffix for unique layout file names."""
-        return "".join(random.choices(string.ascii_lowercase + string.digits, k=length))
-
-    @staticmethod
-    def _parse_command(command: str) -> tuple[str, List[str]]:
-        try:
-            parts = shlex.split(command)
-            if not parts:
-                raise ValueError("Empty command provided")
-            return parts[0], parts[1:] if len(parts) > 1 else []
-        except ValueError as e:
-            logger.error(f"Error parsing command '{command}': {e}")
-            parts = command.split()
-            return parts[0] if parts else "", parts[1:] if len(parts) > 1 else []
-
-    @staticmethod
-    def _format_args_for_kdl(args: List[str]) -> str:
-        if not args:
-            return ""
-        formatted_args = []
-        for arg in args:
-            if " " in arg or '"' in arg or "'" in arg:
-                escaped_arg = arg.replace('"', '\\"')
-                formatted_args.append(f'"{escaped_arg}"')
-            else:
-                formatted_args.append(f'"{arg}"')
-        return " ".join(formatted_args)
-
-    @staticmethod
-    def _create_tab_section(tab_config: TabConfig) -> str:
-        tab_name = tab_config["tabName"]
-        cwd = tab_config["startDir"]
-        command = tab_config["command"]
-        cmd, args = ZellijLayoutGenerator._parse_command(command)
-        args_str = ZellijLayoutGenerator._format_args_for_kdl(args)
-        tab_cwd = cwd or "~"
-        escaped_tab_name = tab_name.replace('"', '\\"')
-        tab_section = f'  tab name="{escaped_tab_name}" cwd="{tab_cwd}" {{\n'
-        tab_section += f'    pane command="{cmd}" {{\n'
-        if args_str:
-            tab_section += f"      args {args_str}\n"
-        tab_section += "    }\n  }\n"
-        return tab_section
-
-    @staticmethod
-    def _validate_layout_config(layout_config: LayoutConfig) -> None:
-        if not layout_config["layoutTabs"]:
-            raise ValueError("Layout must contain at least one tab")
-        for tab in layout_config["layoutTabs"]:
-            if not tab["tabName"].strip():
-                raise ValueError(f"Invalid tab name: {tab['tabName']}")
-            if not tab["command"].strip():
-                raise ValueError(f"Invalid command for tab '{tab['tabName']}': {tab['command']}")
-            if not tab["startDir"].strip():
-                raise ValueError(f"Invalid startDir for tab '{tab['tabName']}': {tab['startDir']}")
-
-    def create_zellij_layout(self, layout_config: LayoutConfig, output_dir: Optional[str], session_name: Optional[str]) -> str:
-        ZellijLayoutGenerator._validate_layout_config(layout_config)
-
-        # Enhanced Rich logging
-        tab_count = len(layout_config["layoutTabs"])
-        layout_name = layout_config["layoutName"]
-        console.print(f"[bold cyan]📋 Creating Zellij layout[/bold cyan] [bright_green]'{layout_name}' with {tab_count} tabs[/bright_green]")
-
-        # Display tab summary with emojis and colors
-        for tab in layout_config["layoutTabs"]:
-            console.print(f"  [yellow]→[/yellow] [bold]{tab['tabName']}[/bold] [dim]in[/dim] [blue]{tab['startDir']}[/blue]")
-
-        # Store session name and layout config for status checking
-        self.session_name = session_name or layout_name
-        self.layout_config = layout_config.copy()
-
-        layout_content = self.layout_template
-        for tab in layout_config["layoutTabs"]:
-            layout_content += "\n" + ZellijLayoutGenerator._create_tab_section(tab)
-        layout_content += "\n}\n"
-
-        try:
-            random_suffix = ZellijLayoutGenerator._generate_random_suffix(8)
-            if output_dir:
-                output_path = Path(output_dir)
-                output_path.mkdir(parents=True, exist_ok=True)
-                layout_file = output_path / f"zellij_layout_{random_suffix}.kdl"
-                layout_file.write_text(layout_content, encoding="utf-8")
-                self.layout_path = str(layout_file.absolute())
-            else:
-                # Use the predefined TMP_LAYOUT_DIR for temporary files
-                TMP_LAYOUT_DIR.mkdir(parents=True, exist_ok=True)
-                layout_file = TMP_LAYOUT_DIR / f"zellij_layout_{self.session_name or 'default'}_{random_suffix}.kdl"
-                layout_file.write_text(layout_content, encoding="utf-8")
-                self.layout_path = str(layout_file.absolute())
-
-            # Enhanced Rich logging for file creation
-            console.print(f"[bold green]✅ Zellij layout file created:[/bold green] [cyan]{self.layout_path}[/cyan]")
-            return self.layout_path
-        except OSError as e:
-            logger.error(f"Failed to create layout file: {e}")
-            raise
-
-    @staticmethod
-    def get_layout_preview(layout_config: LayoutConfig, layout_template: str | None) -> str:
-        if layout_template is None:
-            layout_template = """layout {
-    default_tab_template {
-        // the default zellij tab-bar and status bar plugins
-        pane size=1 borderless=true {
-            plugin location="zellij:compact-bar"
-        }
-        children
-    }
-"""
-        ZellijLayoutGenerator._validate_layout_config(layout_config)
-        layout_content = layout_template
-        for tab in layout_config["layoutTabs"]:
-            layout_content += "\n" + ZellijLayoutGenerator._create_tab_section(tab)
-        return layout_content + "\n}\n"
-
-    @staticmethod
-    def check_command_status(tab_name: str, layout_config: LayoutConfig) -> CommandStatus:
-        # Find the tab with the given name
-        tab_config = None
-        for tab in layout_config["layoutTabs"]:
-            if tab["tabName"] == tab_name:
-                tab_config = tab
-                break
-
-        if tab_config is None:
-            return {"status": "unknown", "error": f"Tab '{tab_name}' not found in layout config", "running": False, "command": "", "cwd": "", "tab_name": tab_name, "processes": []}
-
-        command = tab_config["command"]
-        cwd = tab_config["startDir"]
-        cmd, args = ZellijLayoutGenerator._parse_command(command)
-        try:
-            shells = {"bash", "sh", "zsh", "fish"}
-            matching_processes: list[ProcessInfo] = []
-            for proc in psutil.process_iter(["pid", "name", "cmdline", "status", "ppid", "create_time", "memory_info"]):
-                try:
-                    info = proc.info
-                    proc_cmdline: list[str] | None = info.get("cmdline")  # type: ignore[assignment]
-                    if not proc_cmdline:
-                        continue
-                    if info.get("status") in ["zombie", "dead", "stopped"]:
-                        continue
-                    proc_name = info.get("name", "")
-                    is_match = False
-                    joined_cmdline = " ".join(proc_cmdline)
-                    # Primary matching heuristics - more precise matching
-                    if proc_name == cmd and cmd not in shells:
-                        # For non-shell commands, match if args appear in cmdline
-                        if not args or any(arg in joined_cmdline for arg in args):
-                            is_match = True
-                    elif proc_name == cmd and cmd in shells:
-                        # For shell commands, require more precise matching to avoid false positives
-                        if args:
-                            # Check if all args appear as separate cmdline arguments (not just substrings)
-                            args_found = 0
-                            for arg in args:
-                                for cmdline_arg in proc_cmdline[1:]:  # Skip shell name
-                                    if arg == cmdline_arg or (len(arg) > 3 and arg in cmdline_arg):
-                                        args_found += 1
-                                        break
-                            # Require at least as many args found as we're looking for
-                            if args_found >= len(args):
-                                is_match = True
-                    elif cmd in proc_cmdline[0] and cmd not in shells:
-                        # Non-shell command in first argument
-                        is_match = True
-
-                    # Additional shell wrapper filter - be more restrictive for shells
-                    if is_match and proc_name in shells and args:
-                        # For shell processes, ensure the match is actually meaningful
-                        # Don't match generic shell sessions just because they contain common paths
-                        meaningful_match = False
-                        for arg in args:
-                            # Only consider it meaningful if the arg is substantial (not just a common path)
-                            if len(arg) > 10 and any(arg == cmdline_arg for cmdline_arg in proc_cmdline[1:]):
-                                meaningful_match = True
-                                break
-                            # Or if it's an exact script name match
-                            elif arg.endswith(".py") or arg.endswith(".sh") or arg.endswith(".rb"):
-                                if any(arg in cmdline_arg for cmdline_arg in proc_cmdline[1:]):
-                                    meaningful_match = True
-                                    break
-                        if not meaningful_match:
-                            is_match = False
-                    if not is_match:
-                        continue
-                    try:
-                        proc_obj = psutil.Process(info["pid"])  # type: ignore[index]
-                        if proc_obj.status() not in ["running", "sleeping"]:
-                            continue
-                        mem_info = None
-                        try:
-                            mem = proc_obj.memory_info()
-                            mem_info = mem.rss / (1024 * 1024)
-                        except (psutil.NoSuchProcess, psutil.AccessDenied):
-                            pass
-                        matching_processes.append(
-                            {
-                                "pid": info["pid"],  # type: ignore[index]
-                                "name": proc_name,
-                                "cmdline": proc_cmdline,
-                                "status": info.get("status", "unknown"),
-                                "cmdline_str": joined_cmdline,
-                                "create_time": info.get("create_time", 0.0),
-                                **({"memory_mb": float(mem_info)} if mem_info is not None else {}),
-                            }
-                        )
-                    except (psutil.NoSuchProcess, psutil.AccessDenied):
-                        continue
-                except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
-                    continue
-
-            # Second-pass filtering: remove idle wrapper shells that have no meaningful (non-shell) descendants
-            filtered_active: list[ProcessInfo] = []
-            for proc_info in matching_processes:
-                try:
-                    proc_obj = psutil.Process(proc_info["pid"])  # type: ignore[index]
-                    if not proc_obj.is_running():
-                        continue
-                    status_val = proc_obj.status()
-                    if status_val not in ["running", "sleeping"]:
-                        continue
-                    proc_name = proc_info.get("name", "")
-                    if proc_name in shells:
-                        descendants = proc_obj.children(recursive=True)
-                        # Keep shell only if there exists a non-shell alive descendant OR descendant cmdline still includes our command token
-                        meaningful = False
-                        for child in descendants:
-                            try:
-                                if not child.is_running():
-                                    continue
-                                child_name = child.name()
-                                child_cmdline = " ".join(child.cmdline())
-                                if child_name not in shells:
-                                    meaningful = True
-                                    break
-                                if cmd in child_cmdline or any(arg in child_cmdline for arg in args):
-                                    meaningful = True
-                                    break
-                            except (psutil.NoSuchProcess, psutil.AccessDenied):
-                                continue
-                        if not meaningful:
-                            continue  # discard idle wrapper shell
-                    filtered_active.append(proc_info)
-                except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
-                    continue
-
-            if filtered_active:
-                # Heuristic: if the only remaining processes are wrapper shells invoking a script that already completed, mark as not running.
-                # Case: layout launches 'bash <script.sh>' where script finishes and leaves an idle shell whose cmdline still shows the script path.
-                try:
-                    if all(p.get("name") in shells for p in filtered_active):
-                        script_paths = [arg for arg in args if arg.endswith(".sh")]
-                        shell_only = True
-                        stale_script_overall = False
-                        for p in filtered_active:
-                            proc_shell = psutil.Process(p["pid"])  # type: ignore[index]
-                            create_time = getattr(proc_shell, "create_time", lambda: None)()
-                            cmdline_joined = " ".join(p.get("cmdline", []))
-                            stale_script = False
-                            for spath in script_paths:
-                                script_file = Path(spath)
-                                if script_file.exists():
-                                    try:
-                                        # If script mtime older than process start AND no non-shell descendants -> likely finished
-                                        if create_time and script_file.stat().st_mtime < create_time:
-                                            stale_script = True
-                                    except OSError:
-                                        pass
-                                if spath not in cmdline_joined:
-                                    stale_script = False
-                            # If shell has any alive non-shell descendants, treat as running
-                            descendants = proc_shell.children(recursive=True)
-                            for d in descendants:
-                                try:
-                                    if d.is_running() and d.name() not in shells:
-                                        shell_only = False
-                                        break
-                                except (psutil.NoSuchProcess, psutil.AccessDenied):
-                                    continue
-                            if not shell_only:
-                                break
-                            if stale_script:
-                                stale_script_overall = True
-                        if shell_only and stale_script_overall:
-                            return {"status": "not_running", "running": False, "processes": [], "command": command, "cwd": cwd, "tab_name": tab_name}
-                except Exception:
-                    pass
-                return {"status": "running", "running": True, "processes": filtered_active, "command": command, "cwd": cwd, "tab_name": tab_name}
-            return {"status": "not_running", "running": False, "processes": [], "command": command, "cwd": cwd, "tab_name": tab_name}
-
-        except Exception as e:
-            logger.error(f"Error checking command status for tab '{tab_name}': {e}")
-            return {"status": "error", "error": str(e), "running": False, "command": command, "cwd": cwd, "tab_name": tab_name, "processes": []}
+
+def create_zellij_layout(layout_config: LayoutConfig, layout_template: str) -> str:
+    """Standalone function to create Zellij layout content from config."""
+    validate_layout_config(layout_config)
+    # Enhanced Rich logging
+    tab_count = len(layout_config["layoutTabs"])
+    layout_name = layout_config["layoutName"]
+    console.print(f"[bold cyan]📋 Creating Zellij layout[/bold cyan] [bright_green]'{layout_name}' with {tab_count} tabs[/bright_green]")
+
+    # Display tab summary with emojis and colors
+    for tab in layout_config["layoutTabs"]:
+        console.print(f"  [yellow]→[/yellow] [bold]{tab['tabName']}[/bold] [dim]in[/dim] [blue]{tab['startDir']}[/blue]")
+    from copy import deepcopy
+    layout_config = deepcopy(layout_config)  # Avoid mutating the original
+    layout_content = layout_template
+    for tab in layout_config["layoutTabs"]:
+        layout_content += "\n" + create_tab_section(tab)
+    layout_content += "\n}\n"
+
+    console.print("[bold green]✅ Zellij layout content generated[/bold green]")
+    return layout_content
+
+
+class ZellijLayoutGenerator:
+    def __init__(self, layout_config: LayoutConfig, session_name: str):
+        self.session_name: str = session_name
+        self.layout_config: LayoutConfig = layout_config.copy()
+        self.layout_path: Optional[str] = None
+    def create_layout_file(self) -> bool:
+        """Create zellij layout file and return the path."""
+        layout_content = create_zellij_layout(self.layout_config, layout_template=DEFAULT_LAYOUT_TEMPLATE)
+        # Write to file
+        tmp_dir = Path.home() / "tmp_results" / "sessions" / "zellij_layouts"
+        tmp_dir.mkdir(parents=True, exist_ok=True)
+        import tempfile
+        layout_file = Path(tempfile.mkstemp(suffix="_layout.kdl", dir=tmp_dir)[1])
+        layout_file.write_text(layout_content, encoding="utf-8")
+        self.layout_path = str(layout_file.absolute())
+
+        console.print(f"[bold green]✅ Layout created successfully:[/bold green] [cyan]{self.layout_path}[/cyan]")
+        return True
 
     def check_all_commands_status(self) -> dict[str, CommandStatus]:
         if not self.layout_config:
@@ -341,41 +76,16 @@ class ZellijLayoutGenerator:
         status_report: dict[str, CommandStatus] = {}
         for tab in self.layout_config["layoutTabs"]:
             tab_name = tab["tabName"]
-            status_report[tab_name] = ZellijLayoutGenerator.check_command_status(tab_name, self.layout_config)
+            status_report[tab_name] = check_command_status(tab_name, self.layout_config)
 
         return status_report
 
-    @staticmethod
-    def check_zellij_session_status(session_name: str) -> ZellijSessionStatus:
-        try:
-            # Run zellij list-sessions command
-            result = subprocess.run(["zellij", "list-sessions"], capture_output=True, text=True, timeout=10)
-
-            if result.returncode == 0:
-                sessions = result.stdout.strip().split("\n") if result.stdout.strip() else []
-                session_running = any(session_name in session for session in sessions)
-
-                return {"zellij_running": True, "session_exists": session_running, "session_name": session_name, "all_sessions": sessions}
-            else:
-                return {"zellij_running": False, "session_exists": False, "session_name": session_name, "all_sessions": [], "error": result.stderr}
-
-        except subprocess.TimeoutExpired:
-            return {"zellij_running": False, "session_exists": False, "session_name": session_name, "all_sessions": [], "error": "Timeout while checking Zellij sessions"}
-        except FileNotFoundError:
-            return {"zellij_running": False, "session_exists": False, "session_name": session_name, "all_sessions": [], "error": "Zellij not found in PATH"}
-        except Exception as e:
-            return {"zellij_running": False, "session_exists": False, "session_name": session_name, "all_sessions": [], "error": str(e)}
-
     def get_comprehensive_status(self) -> ComprehensiveStatus:
-        zellij_status = ZellijLayoutGenerator.check_zellij_session_status(self.session_name or "default")
+        zellij_status = check_zellij_session_status(self.session_name or "default")
         commands_status = self.check_all_commands_status()
         running_count = sum(1 for status in commands_status.values() if status.get("running", False))
         total_count = len(commands_status)
-        return {
-            "zellij_session": zellij_status,
-            "commands": commands_status,
-            "summary": {"total_commands": total_count, "running_commands": running_count, "stopped_commands": total_count - running_count, "session_healthy": zellij_status.get("session_exists", False)},
-        }
+        return {"zellij_session": zellij_status, "commands": commands_status, "summary": {"total_commands": total_count, "running_commands": running_count, "stopped_commands": total_count - running_count, "session_healthy": zellij_status.get("session_exists", False)}}
 
     def print_status_report(self) -> None:
         from rich.panel import Panel
@@ -446,26 +156,29 @@ class ZellijLayoutGenerator:
 
         console.print(Panel(summary_text, title="📊 Summary", style="blue"))
 
-
-def created_zellij_layout(layout_config: LayoutConfig, output_dir: Optional[str]) -> str:
-    generator = ZellijLayoutGenerator()
-    return generator.create_zellij_layout(layout_config, output_dir, None)
-
-
-def run_zellij_layout(layout_config: LayoutConfig):
-    layout_path = created_zellij_layout(layout_config, None)
+    def restart_tab(self, tab_name: str) -> bool:
+        """Restart the process running in a specific tab without changing the layout.
+        
+        This method will:
+        1. Navigate to the specified tab
+        2. Send Ctrl+C to stop the running process
+        3. Clear the screen
+        4. Re-execute the original command
+        
+        The tab layout and configuration remain unchanged - only the process is restarted.
+        """
+        return restart_tab_process(tab_name=tab_name, layout_config=self.layout_config, session_name=self.session_name)
+
+    def run(self):
+        from machineconfig.cluster.sessions_managers.helpers.enhanced_command_runner import enhanced_zellij_session_start
+        enhanced_zellij_session_start(session_name=self.session_name, layout_path=str(self.layout_path))
+
+
+def run_zellij_layout(layout_config: LayoutConfig) -> None:
     session_name = layout_config["layoutName"]
-    try:
-        from machineconfig.cluster.sessions_managers.utils.enhanced_command_runner import enhanced_zellij_session_start
-
-        enhanced_zellij_session_start(session_name, layout_path)
-    except ImportError:
-        # Fallback to original implementation
-        cmd = f"zellij delete-session --force {session_name}; zellij --layout {layout_path} a -b {session_name}"
-        import subprocess
-
-        subprocess.run(cmd, shell=True, check=True)
-        console.print(f"[bold green]🚀 Zellij layout is running[/bold green] [yellow]@[/yellow] [bold cyan]{session_name}[/bold cyan]")
+    generator = ZellijLayoutGenerator(layout_config, session_name)
+    generator.create_layout_file()
+    generator.run()
 
 
 def run_command_in_zellij_tab(command: str, tab_name: str, cwd: Optional[str]) -> str:
@@ -493,30 +206,23 @@ if __name__ == "__main__":
     # Example usage with new schema
     sample_layout: LayoutConfig = {
         "layoutName": "SampleBots",
-        "layoutTabs": [
-            {"tabName": "🤖Bot1", "startDir": "~/code/bytesense/bithence", "command": "~/scripts/fire -mO go1.py bot1 -- --create_new_bot=True"},
-            {"tabName": "🤖Bot2", "startDir": "~/code/bytesense/bithence", "command": "~/scripts/fire -mO go2.py bot2 -- --create_new_bot=True"},
-            {"tabName": "📊Monitor", "startDir": "~", "command": "htop"},
-            {"tabName": "📝Logs", "startDir": "/var/log", "command": "tail -f /var/log/app.log"},
-        ],
+        "layoutTabs": [{"tabName": "Explorer", "startDir": "~/code", "command": "lf"}, {"tabName": "🤖Bot2", "startDir": "~", "command": "cmatrix"}, {"tabName": "📊Monitor", "startDir": "~", "command": "htop"}, {"tabName": "📝Logs", "startDir": "/var/log", "command": "tail -f /var/log/app.log"}],
     }
-
     try:
-        # Create layout using the generator directly to access status methods
-        generator = ZellijLayoutGenerator()
-        layout_path = generator.create_zellij_layout(sample_layout, None, "test_session")
-        print(f"✅ Layout created successfully: {layout_path}")
+        # Create layout using the generator with new design
+        generator = ZellijLayoutGenerator(layout_config=sample_layout, session_name="test_session")
+        generator.create_layout_file()
 
+        generator.run()
         # Demonstrate status checking
         print("\n🔍 Checking command status (this is just a demo - commands aren't actually running):")
         generator.print_status_report()
 
         # Individual command status check
-        print("\n🔎 Individual command status for Bot1:")
-        if generator.layout_config:
-            bot1_status = ZellijLayoutGenerator.check_command_status("🤖Bot1", generator.layout_config)
-            print(f"Status: {bot1_status['status']}")
-            print(f"Running: {bot1_status['running']}")
+        print("\n🔎 Individual command status for 🤖Bot2:")
+        bot2_status = check_command_status("🤖Bot2", generator.layout_config)
+        print(f"Status: {bot2_status['status']}")
+        print(f"Running: {bot2_status['running']}")
 
     except Exception as e:
         print(f"❌ Error: {e}")