...

herolib/infra/tmuxrunner/process_monitor.py

@@ -0,0 +1,106 @@
+import os
+import time
+import re
+import sys
+import toml
+import libtmux
+from libtmux.pane import Pane
+from libtmux.window import Window
+from libtmux.session import Session
+import psutil
+from typing import Dict, List, Optional, Any, Set, Tuple
+from dataclasses import dataclass, field, asdict
+from datetime import datetime
+import uuid
+from pathlib import Path
+import asyncio
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import threading
+
+
+class ProcessMonitor:
+    """Monitor processes running in tmux panes using psutil."""
+    
+    @staticmethod
+    def get_pane_process_tree(pane: Pane) -> Tuple[Optional[psutil.Process], List[psutil.Process]]:
+        """Get the main process and all child processes for a tmux pane."""
+        try:
+            pane_pid = pane.pane_pid
+            if pane_pid is None:
+                return None, []
+            
+            # Get the main process
+            try:
+                main_process = psutil.Process(int(pane_pid))
+            except (psutil.NoSuchProcess, ValueError):
+                return None, []
+            
+            # Get all children recursively
+            children = []
+            try:
+                children = main_process.children(recursive=True)
+            except psutil.NoSuchProcess:
+                pass
+            
+            return main_process, children
+        except Exception as e:
+            print(f"Error getting process tree: {e}")
+            return None, []
+    
+    @staticmethod
+    def get_process_metrics(pane: Pane) -> ProcessMetrics:
+        """Get CPU and memory metrics for all processes in a pane."""
+        metrics = ProcessMetrics()
+        metrics.last_updated = datetime.now().isoformat()
+        
+        main_proc, children = ProcessMonitor.get_pane_process_tree(pane)
+        
+        if main_proc is None:
+            return metrics
+        
+        try:
+            # Get main process metrics
+            if main_proc.is_running():
+                metrics.cpu_percent = main_proc.cpu_percent(interval=0.1)
+                mem_info = main_proc.memory_info()
+                metrics.memory_rss = mem_info.rss
+                metrics.memory_vms = mem_info.vms
+                metrics.memory_percent = main_proc.memory_percent()
+                metrics.num_threads = main_proc.num_threads()
+            
+            # Get children metrics
+            metrics.num_children = len(children)
+            for child in children:
+                try:
+                    if child.is_running():
+                        metrics.children_cpu_percent += child.cpu_percent(interval=0.1)
+                        child_mem = child.memory_info()
+                        metrics.children_memory_rss += child_mem.rss
+                except (psutil.NoSuchProcess, psutil.AccessDenied):
+                    pass
+                    
+        except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
+            print(f"Error getting process metrics: {e}")
+            
+        return metrics
+    
+    @staticmethod
+    def is_process_running_command(pane: Pane, command_pattern: str) -> bool:
+        """Check if a specific command is running in the pane."""
+        main_proc, children = ProcessMonitor.get_pane_process_tree(pane)
+        
+        all_processes = [main_proc] + children if main_proc else children
+        
+        for proc in all_processes:
+            try:
+                if proc and proc.is_running():
+                    cmdline = " ".join(proc.cmdline())
+                    if command_pattern in cmdline:
+                        return True
+            except (psutil.NoSuchProcess, psutil.AccessDenied):
+                pass
+        
+        return False