kollabor 0.4.9__py3-none-any.whl → 0.4.15__py3-none-any.whl

plugins/agent_orchestrator/orchestrator.py

@@ -0,0 +1,403 @@
+"""Manages tmux sessions for agent sub-processes."""
+
+import subprocess
+import time
+import asyncio
+import fnmatch
+import logging
+from pathlib import Path
+from typing import List, Optional, Tuple, Dict
+
+from .models import AgentSession, AgentTask
+from .file_attacher import FileAttacher
+
+logger = logging.getLogger(__name__)
+
+
+class AgentOrchestrator:
+    """Manages tmux sessions for agent sub-processes."""
+
+    def __init__(self, project_name: str = None):
+        """Initialize orchestrator.
+
+        Args:
+            project_name: Project name for session naming. Defaults to cwd name.
+        """
+        self.project_name = project_name or Path.cwd().name
+        self.agents: Dict[str, AgentSession] = {}
+        self.file_attacher = FileAttacher()
+
+        # Timing configuration
+        self.session_init_delay = 1
+        self.kollab_init_delay = 7  # Allow time for kollab to fully initialize
+        self.message_delay = 2
+
+    # -------------------------------------------------------------------------
+    # Spawn
+    # -------------------------------------------------------------------------
+
+    async def spawn(self, name: str, task: str, files: List[str] = None) -> bool:
+        """Spawn a new agent session.
+
+        Args:
+            name: Agent name
+            task: Task description
+            files: Optional list of files to attach
+
+        Returns:
+            True if successful
+        """
+        full_name = f"{self.project_name}-{name}"
+
+        # Check if already exists
+        if self._session_exists(full_name):
+            logger.warning(f"Session already exists: {full_name}")
+            return False
+
+        # Create session
+        if not self._create_session(full_name):
+            logger.error(f"Failed to create session: {full_name}")
+            return False
+
+        # Start kollab
+        logger.info(f"[spawn] Sending 'kollab' command to {full_name}")
+        self._send_keys(full_name, "kollab")
+        logger.info(f"[spawn] Waiting {self.kollab_init_delay}s for kollab to initialize")
+        await asyncio.sleep(self.kollab_init_delay)
+
+        # Prepare message with attached files
+        message = ""
+        if files:
+            message = self.file_attacher.attach(files)
+            message += "\n\n"
+        message += task
+
+        logger.info(f"[spawn] Sending task to {full_name}: {message[:100]}...")
+
+        # Send task
+        result = self._send_keys(full_name, message)
+        logger.info(f"[spawn] Task send result: {result}")
+
+        # Track agent
+        self.agents[name] = AgentSession(
+            name=name,
+            full_name=full_name,
+            status="running",
+            start_time=time.time(),
+        )
+
+        logger.info(f"Spawned agent: {name}")
+        return True
+
+    async def spawn_clone(
+        self,
+        name: str,
+        task: str,
+        files: List[str],
+        conversation_file: str,
+    ) -> bool:
+        """Spawn agent with conversation context.
+
+        Args:
+            name: Agent name
+            task: Task description
+            files: Files to attach
+            conversation_file: Path to exported conversation JSON
+
+        Returns:
+            True if successful
+        """
+        full_name = f"{self.project_name}-{name}"
+
+        if self._session_exists(full_name):
+            logger.warning(f"Session already exists: {full_name}")
+            return False
+
+        if not self._create_session(full_name):
+            logger.error(f"Failed to create session: {full_name}")
+            return False
+
+        # Start kollab with resume
+        self._send_keys(full_name, f"kollab --resume {conversation_file}")
+        await asyncio.sleep(self.kollab_init_delay)
+
+        # Send task with files
+        message = ""
+        if files:
+            message = self.file_attacher.attach(files) + "\n\n"
+        message += task
+
+        self._send_keys(full_name, message)
+
+        self.agents[name] = AgentSession(
+            name=name,
+            full_name=full_name,
+            status="running",
+            start_time=time.time(),
+        )
+
+        logger.info(f"Spawned clone agent: {name} with conversation context")
+        return True
+
+    async def spawn_team_lead(
+        self,
+        lead_name: str,
+        max_workers: int,
+        task: AgentTask,
+    ) -> bool:
+        """Spawn a team lead agent that can spawn workers.
+
+        Args:
+            lead_name: Lead agent name
+            max_workers: Maximum number of workers the lead can spawn
+            task: Task for the lead
+
+        Returns:
+            True if successful
+        """
+        full_name = f"{self.project_name}-{lead_name}"
+
+        if self._session_exists(full_name):
+            logger.warning(f"Session already exists: {full_name}")
+            return False
+
+        if not self._create_session(full_name):
+            logger.error(f"Failed to create session: {full_name}")
+            return False
+
+        self._send_keys(full_name, "kollab")
+        await asyncio.sleep(self.kollab_init_delay)
+
+        # Inject team lead prompt
+        lead_prompt = f"""You are a team lead agent.
+You can spawn up to {max_workers} worker agents using <agent> tags.
+Coordinate their work and integrate results.
+Use <status /> to check on workers.
+Use <capture>worker-name 100</capture> to see their progress.
+
+"""
+        message = lead_prompt
+        if task.files:
+            message += self.file_attacher.attach(task.files) + "\n\n"
+        message += task.task
+
+        self._send_keys(full_name, message)
+
+        self.agents[lead_name] = AgentSession(
+            name=lead_name,
+            full_name=full_name,
+            status="running",
+            start_time=time.time(),
+        )
+
+        logger.info(f"Spawned team lead: {lead_name} with {max_workers} max workers")
+        return True
+
+    # -------------------------------------------------------------------------
+    # Message
+    # -------------------------------------------------------------------------
+
+    async def message(self, name: str, content: str) -> bool:
+        """Send message to agent.
+
+        Args:
+            name: Agent name
+            content: Message content
+
+        Returns:
+            True if successful
+        """
+        full_name = f"{self.project_name}-{name}"
+
+        if not self._session_exists(full_name):
+            logger.warning(f"Session does not exist: {full_name}")
+            return False
+
+        self._send_keys(full_name, content)
+        logger.info(f"Sent message to agent: {name}")
+        return True
+
+    # -------------------------------------------------------------------------
+    # Stop
+    # -------------------------------------------------------------------------
+
+    async def stop(self, name: str) -> Tuple[str, str]:
+        """Stop agent and return final output + duration.
+
+        Args:
+            name: Agent name
+
+        Returns:
+            Tuple of (output, duration)
+        """
+        full_name = f"{self.project_name}-{name}"
+
+        # Capture final output
+        output = self.capture_output(name, 100)
+
+        # Get duration
+        agent = self.agents.get(name)
+        duration = agent.duration if agent else "?"
+
+        # Kill session
+        self._kill_session(full_name)
+
+        # Remove from tracking
+        if name in self.agents:
+            del self.agents[name]
+
+        logger.info(f"Stopped agent: {name} @ {duration}")
+        return output, duration
+
+    # -------------------------------------------------------------------------
+    # Status / Capture
+    # -------------------------------------------------------------------------
+
+    def list_agents(self) -> List[AgentSession]:
+        """List all active agents.
+
+        Returns:
+            List of agent sessions
+        """
+        # Refresh status from tmux
+        self._refresh_agents()
+        return list(self.agents.values())
+
+    def get_agent(self, name: str) -> Optional[AgentSession]:
+        """Get specific agent.
+
+        Args:
+            name: Agent name
+
+        Returns:
+            Agent session or None
+        """
+        return self.agents.get(name)
+
+    def find_agents(self, pattern: str) -> List[str]:
+        """Find agents matching glob pattern.
+
+        Args:
+            pattern: Glob pattern (e.g., "lint-*")
+
+        Returns:
+            List of matching agent names
+        """
+        return [
+            name for name in self.agents.keys() if fnmatch.fnmatch(name, pattern)
+        ]
+
+    def capture_output(self, name: str, lines: int = 50) -> str:
+        """Capture last N lines from agent.
+
+        Args:
+            name: Agent name
+            lines: Number of lines to capture
+
+        Returns:
+            Captured output
+        """
+        full_name = f"{self.project_name}-{name}"
+
+        result = subprocess.run(
+            ["tmux", "capture-pane", "-t", full_name, "-p", "-S", f"-{lines}"],
+            capture_output=True,
+            text=True,
+        )
+        return result.stdout
+
+    # -------------------------------------------------------------------------
+    # Private Helpers
+    # -------------------------------------------------------------------------
+
+    def _session_exists(self, full_name: str) -> bool:
+        """Check if tmux session exists."""
+        result = subprocess.run(
+            ["tmux", "has-session", "-t", full_name], capture_output=True
+        )
+        return result.returncode == 0
+
+    def _create_session(self, full_name: str) -> bool:
+        """Create new tmux session in current working directory."""
+        cwd = str(Path.cwd())
+        result = subprocess.run(
+            ["tmux", "new-session", "-d", "-s", full_name, "-c", cwd], capture_output=True
+        )
+        if result.returncode != 0:
+            logger.error(f"Failed to create tmux session: {result.stderr}")
+            return False
+
+        time.sleep(self.session_init_delay)
+        return True
+
+    def _send_keys(self, full_name: str, content: str) -> bool:
+        """Send keys to tmux session."""
+        logger.debug(f"[send_keys] Sending to {full_name}: {len(content)} chars")
+
+        # Send content first
+        result = subprocess.run(
+            ["tmux", "send-keys", "-t", full_name, content],
+            capture_output=True,
+        )
+        if result.returncode != 0:
+            logger.error(f"[send_keys] Content send failed: {result.stderr}")
+            return False
+
+        logger.debug(f"[send_keys] Content sent, waiting 1s before Enter")
+        # Wait for paste to be processed
+        time.sleep(1)
+
+        # Then send Enter to submit
+        logger.debug(f"[send_keys] Sending Enter")
+        result = subprocess.run(
+            ["tmux", "send-keys", "-t", full_name, "Enter"],
+            capture_output=True,
+        )
+        if result.returncode != 0:
+            logger.error(f"[send_keys] Enter send failed: {result.stderr}")
+
+        logger.debug(f"[send_keys] Waiting {self.message_delay}s message delay")
+        time.sleep(self.message_delay)
+        return result.returncode == 0
+
+    def _kill_session(self, full_name: str) -> bool:
+        """Kill tmux session."""
+        result = subprocess.run(
+            ["tmux", "kill-session", "-t", full_name], capture_output=True
+        )
+        return result.returncode == 0
+
+    def _refresh_agents(self) -> None:
+        """Refresh agent status from tmux."""
+        result = subprocess.run(
+            ["tmux", "list-sessions", "-F", "#{session_name}"],
+            capture_output=True,
+            text=True,
+        )
+
+        if result.returncode != 0:
+            return
+
+        active_sessions = set(result.stdout.strip().split("\n"))
+        prefix = f"{self.project_name}-"
+
+        # Remove dead agents
+        dead = [
+            name
+            for name in self.agents
+            if f"{prefix}{name}" not in active_sessions
+        ]
+        for name in dead:
+            del self.agents[name]
+
+        # Discover new agents (created externally)
+        for session in active_sessions:
+            if session.startswith(prefix):
+                agent_name = session[len(prefix):]
+                if agent_name and agent_name not in self.agents:
+                    self.agents[agent_name] = AgentSession(
+                        name=agent_name,
+                        full_name=session,
+                        status="running",
+                        start_time=time.time(),  # approximate
+                    )