emdash-core 0.1.7__py3-none-any.whl

emdash_core/swarm/swarm_runner.py

@@ -0,0 +1,226 @@
+"""Main swarm orchestrator combining all components."""
+
+import re
+from pathlib import Path
+from typing import Callable, Optional
+from uuid import uuid4
+
+from .task_definition import SwarmTask, SwarmState, TaskStatus
+from .worktree_manager import WorktreeManager
+from .worker_spawner import WorkerSpawner, WorkerResult
+from .merge_agent import MergeAgent, MergeResult
+from ..utils.logger import log
+
+
+class SwarmRunner:
+    """Orchestrates the full swarm lifecycle.
+
+    1. Creates worktrees for each task
+    2. Spawns parallel agent workers
+    3. Collects results
+    4. Merges completed work
+
+    Example:
+        runner = SwarmRunner(repo_root=Path("."))
+
+        tasks = [
+            "Add user authentication",
+            "Fix the login page CSS",
+            "Add unit tests for auth module",
+        ]
+
+        state = runner.run(tasks)
+        runner.merge_completed()
+    """
+
+    def __init__(
+        self,
+        repo_root: Path,
+        model: str = "gpt-4o-mini",
+        max_workers: int = 3,
+        timeout_per_task: int = 600,
+        base_branch: str = "main",
+        include_graph_tools: bool = False,
+    ):
+        self.repo_root = repo_root.resolve()
+        self.model = model
+        self.max_workers = max_workers
+        self.timeout = timeout_per_task
+        self.base_branch = base_branch
+        self.include_graph_tools = include_graph_tools
+
+        self.worktree_manager = WorktreeManager(repo_root)
+        self.spawner = WorkerSpawner(
+            repo_root=repo_root,
+            model=model,
+            include_graph_tools=include_graph_tools,
+            timeout_per_task=timeout_per_task,
+        )
+        self.merge_agent = MergeAgent(repo_root, model)
+
+        self.state: Optional[SwarmState] = None
+
+    def slugify(self, text: str) -> str:
+        """Create URL-safe slug from text."""
+        text = text.lower().strip()
+        text = re.sub(r"[^\w\s-]", "", text)
+        text = re.sub(r"[-\s]+", "-", text)
+        return text[:50]
+
+    def run(
+        self,
+        task_descriptions: list[str],
+        progress_callback: Optional[Callable[[SwarmTask, WorkerResult], None]] = None,
+    ) -> SwarmState:
+        """Execute the full swarm workflow.
+
+        Args:
+            task_descriptions: List of task descriptions
+            progress_callback: Called when each task completes
+
+        Returns:
+            SwarmState with all task results
+        """
+        # Create swarm state
+        swarm_id = str(uuid4())[:8]
+        tasks = []
+
+        for i, desc in enumerate(task_descriptions):
+            slug = self.slugify(desc)
+            # Add index suffix if slug is duplicate
+            existing_slugs = [t.slug for t in tasks]
+            if slug in existing_slugs:
+                slug = f"{slug}-{i+1}"
+
+            task = SwarmTask(
+                id=str(uuid4()),
+                slug=slug,
+                title=desc[:100],
+                description=desc,
+                base_branch=self.base_branch,
+            )
+            tasks.append(task)
+
+        self.state = SwarmState(
+            id=swarm_id,
+            tasks=tasks,
+            base_branch=self.base_branch,
+        )
+
+        # Save initial state
+        self.state.save(self.repo_root)
+
+        log.info(f"Starting swarm {swarm_id} with {len(tasks)} tasks")
+
+        # Run parallel workers
+        self.spawner.run_parallel(
+            tasks=tasks,
+            max_workers=self.max_workers,
+            progress_callback=progress_callback,
+        )
+
+        # Update and save final state
+        self.state.save(self.repo_root)
+
+        return self.state
+
+    def merge_completed(
+        self,
+        use_llm: bool = False,
+        target_branch: Optional[str] = None,
+        cleanup_worktrees: bool = True,
+    ) -> list[MergeResult]:
+        """Merge all completed task branches.
+
+        Args:
+            use_llm: Whether to use LLM assistance for merges
+            target_branch: Target branch (defaults to base_branch)
+            cleanup_worktrees: Whether to delete worktrees after successful merge
+
+        Returns:
+            List of merge results
+        """
+        if not self.state:
+            self.state = SwarmState.load(self.repo_root)
+
+        if not self.state:
+            raise ValueError("No swarm state found")
+
+        target = target_branch or self.state.base_branch
+
+        results = self.merge_agent.merge_all_completed(
+            self.state.tasks,
+            target=target,
+            use_llm=use_llm,
+            cleanup_worktrees=cleanup_worktrees,
+        )
+
+        # Update state
+        self.state.save(self.repo_root)
+
+        return results
+
+    def cleanup(self) -> int:
+        """Clean up all worktrees from this swarm.
+
+        Returns:
+            Number of worktrees removed
+        """
+        return self.worktree_manager.cleanup_all()
+
+    def get_status(self) -> dict:
+        """Get current swarm status.
+
+        Returns:
+            Dict with status summary
+        """
+        if not self.state:
+            self.state = SwarmState.load(self.repo_root)
+
+        if not self.state:
+            return {"active": False, "message": "No active swarm"}
+
+        summary = {
+            "active": True,
+            "swarm_id": self.state.id,
+            "base_branch": self.state.base_branch,
+            "created_at": self.state.created_at,
+            "total_tasks": len(self.state.tasks),
+            "pending": sum(1 for t in self.state.tasks if t.status == TaskStatus.PENDING),
+            "running": sum(1 for t in self.state.tasks if t.status == TaskStatus.RUNNING),
+            "completed": sum(1 for t in self.state.tasks if t.status == TaskStatus.COMPLETED),
+            "failed": sum(1 for t in self.state.tasks if t.status == TaskStatus.FAILED),
+            "merged": sum(1 for t in self.state.tasks if t.status == TaskStatus.MERGED),
+            "tasks": [
+                {
+                    "slug": t.slug,
+                    "title": t.title,
+                    "status": t.status.value,
+                    "branch": t.branch,
+                    "files_modified": len(t.files_modified),
+                }
+                for t in self.state.tasks
+            ],
+        }
+        return summary
+
+    @classmethod
+    def load(cls, repo_root: Path) -> Optional["SwarmRunner"]:
+        """Load an existing swarm runner from state.
+
+        Args:
+            repo_root: Repository root path
+
+        Returns:
+            SwarmRunner if state exists, None otherwise
+        """
+        state = SwarmState.load(repo_root)
+        if not state:
+            return None
+
+        runner = cls(
+            repo_root=repo_root,
+            base_branch=state.base_branch,
+        )
+        runner.state = state
+        return runner

emdash_core/swarm/task_definition.py

@@ -0,0 +1,137 @@
+"""Task definitions and state for swarm execution."""
+
+from dataclasses import dataclass, field, asdict
+from enum import Enum
+from pathlib import Path
+from typing import Optional, Any
+from datetime import datetime
+import json
+
+
+class TaskStatus(Enum):
+    """Status of a swarm task."""
+    PENDING = "pending"
+    WORKTREE_CREATED = "worktree_created"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    MERGED = "merged"
+
+
+@dataclass
+class SwarmTask:
+    """A task to be executed by a swarm worker.
+
+    Each task gets its own worktree and subprocess.
+    State is persisted to {worktree}/.emdash-task/task.json
+    """
+    # Core identity
+    id: str
+    slug: str
+    title: str
+    description: str
+
+    # Git context
+    base_branch: str = "main"
+    branch: str = ""
+    worktree_path: Optional[str] = None
+
+    # Status tracking
+    status: TaskStatus = TaskStatus.PENDING
+    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
+    started_at: Optional[str] = None
+    completed_at: Optional[str] = None
+
+    # Results
+    files_modified: list[str] = field(default_factory=list)
+    completion_summary: Optional[str] = None
+    error_message: Optional[str] = None
+    agent_output: Optional[str] = None
+
+    # Merge status
+    merge_status: Optional[str] = None
+    merge_conflicts: list[str] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        d = asdict(self)
+        d["status"] = self.status.value
+        return d
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "SwarmTask":
+        """Create from dictionary."""
+        data = data.copy()
+        data["status"] = TaskStatus(data["status"])
+        return cls(**data)
+
+    def save(self, worktree_path: Path) -> None:
+        """Save task state to worktree."""
+        task_dir = worktree_path / ".emdash-task"
+        task_dir.mkdir(parents=True, exist_ok=True)
+
+        task_file = task_dir / "task.json"
+        with open(task_file, "w") as f:
+            json.dump(self.to_dict(), f, indent=2)
+
+    @classmethod
+    def load(cls, worktree_path: Path) -> Optional["SwarmTask"]:
+        """Load task state from worktree."""
+        task_file = worktree_path / ".emdash-task" / "task.json"
+        if not task_file.exists():
+            return None
+
+        with open(task_file) as f:
+            return cls.from_dict(json.load(f))
+
+
+@dataclass
+class SwarmState:
+    """Overall state of a swarm execution.
+
+    Persisted to {repo}/.emdash-worktrees/.swarm-state.json
+    """
+    id: str
+    tasks: list[SwarmTask]
+    base_branch: str = "main"
+    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
+    completed_at: Optional[str] = None
+
+    def save(self, repo_root: Path) -> None:
+        """Save swarm state to repo root."""
+        state_file = repo_root / ".emdash-worktrees" / ".swarm-state.json"
+        state_file.parent.mkdir(parents=True, exist_ok=True)
+
+        data = {
+            "id": self.id,
+            "tasks": [t.to_dict() for t in self.tasks],
+            "base_branch": self.base_branch,
+            "created_at": self.created_at,
+            "completed_at": self.completed_at,
+            "summary": {
+                "total": len(self.tasks),
+                "completed": sum(1 for t in self.tasks if t.status == TaskStatus.COMPLETED),
+                "failed": sum(1 for t in self.tasks if t.status == TaskStatus.FAILED),
+                "merged": sum(1 for t in self.tasks if t.status == TaskStatus.MERGED),
+            },
+        }
+        with open(state_file, "w") as f:
+            json.dump(data, f, indent=2)
+
+    @classmethod
+    def load(cls, repo_root: Path) -> Optional["SwarmState"]:
+        """Load swarm state from repo root."""
+        state_file = repo_root / ".emdash-worktrees" / ".swarm-state.json"
+        if not state_file.exists():
+            return None
+
+        with open(state_file) as f:
+            data = json.load(f)
+            tasks = [SwarmTask.from_dict(t) for t in data.get("tasks", [])]
+            return cls(
+                id=data["id"],
+                tasks=tasks,
+                base_branch=data.get("base_branch", "main"),
+                created_at=data.get("created_at", ""),
+                completed_at=data.get("completed_at"),
+            )

emdash_core/swarm/worker_spawner.py

@@ -0,0 +1,319 @@
+"""Subprocess spawning and management for parallel agent workers."""
+
+import json
+import os
+import subprocess
+import sys
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Callable, Optional
+
+from .task_definition import SwarmTask, TaskStatus
+from .worktree_manager import WorktreeManager, WorktreeInfo
+from ..utils.logger import log
+
+
+@dataclass
+class WorkerResult:
+    """Result from a worker subprocess."""
+    task_id: str
+    success: bool
+    files_modified: list[str]
+    completion_summary: Optional[str]
+    error_message: Optional[str]
+    stdout: str
+    stderr: str
+    return_code: int
+
+
+def run_agent_in_worktree(
+    worktree_path: str,
+    task_description: str,
+    model: str = "gpt-4o-mini",
+    include_graph_tools: bool = False,
+    timeout: int = 600,
+) -> WorkerResult:
+    """Run a coding agent in a worktree subprocess.
+
+    This function is designed to be called in a separate process.
+    It invokes `emdash agent code` within the worktree directory.
+
+    Args:
+        worktree_path: Path to the worktree
+        task_description: Task for the agent to complete
+        model: LLM model to use
+        include_graph_tools: Whether to include graph tools
+        timeout: Maximum seconds to run
+
+    Returns:
+        WorkerResult with execution details
+    """
+    worktree = Path(worktree_path)
+    task_file = worktree / ".emdash-task" / "task.json"
+
+    # Load task for ID
+    task_id = "unknown"
+    if task_file.exists():
+        with open(task_file) as f:
+            task_data = json.load(f)
+            task_id = task_data.get("id", "unknown")
+
+    # Build command - run emdash agent code in non-interactive mode
+    cmd = [
+        sys.executable, "-m", "emdash",
+        "agent", "code",
+        task_description,
+        "--model", model,
+        "--mode", "code",
+        "-q",  # Quiet mode
+    ]
+
+    if not include_graph_tools:
+        cmd.append("--no-graph-tools")
+
+    # Set environment for isolated execution
+    env = os.environ.copy()
+    # Each worktree gets its own Kuzu DB to avoid lock contention
+    env["KUZU_DATABASE_PATH"] = str(worktree / ".emdash-task" / "kuzu_db")
+
+    try:
+        result = subprocess.run(
+            cmd,
+            cwd=str(worktree),
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+            env=env,
+        )
+
+        # Parse output to extract modified files
+        files_modified = _extract_modified_files(result.stdout)
+        completion_summary = _extract_summary(result.stdout)
+
+        return WorkerResult(
+            task_id=task_id,
+            success=result.returncode == 0,
+            files_modified=files_modified,
+            completion_summary=completion_summary,
+            error_message=result.stderr if result.returncode != 0 else None,
+            stdout=result.stdout,
+            stderr=result.stderr,
+            return_code=result.returncode,
+        )
+
+    except subprocess.TimeoutExpired:
+        return WorkerResult(
+            task_id=task_id,
+            success=False,
+            files_modified=[],
+            completion_summary=None,
+            error_message=f"Task timed out after {timeout} seconds",
+            stdout="",
+            stderr="",
+            return_code=-1,
+        )
+    except Exception as e:
+        return WorkerResult(
+            task_id=task_id,
+            success=False,
+            files_modified=[],
+            completion_summary=None,
+            error_message=str(e),
+            stdout="",
+            stderr="",
+            return_code=-1,
+        )
+
+
+def _extract_modified_files(output: str) -> list[str]:
+    """Extract list of modified files from agent output."""
+    files = []
+    for line in output.split("\n"):
+        # Look for common patterns indicating file modifications
+        lower = line.lower()
+        if any(x in lower for x in ["modified:", "created:", "applied diff to", "wrote to"]):
+            # Try to extract file path
+            parts = line.split()
+            for part in parts:
+                if "/" in part or part.endswith(".py") or part.endswith(".ts"):
+                    # Clean up the path
+                    path = part.strip(",:\"'`")
+                    if path and not path.startswith("-"):
+                        files.append(path)
+    return list(set(files))  # Dedupe
+
+
+def _extract_summary(output: str) -> Optional[str]:
+    """Extract completion summary from agent output."""
+    # Look for completion markers
+    lines = output.split("\n")
+    for i, line in enumerate(lines):
+        if "completed" in line.lower() or "summary" in line.lower():
+            # Return this line and a few following
+            summary_lines = lines[i:i+5]
+            return "\n".join(summary_lines).strip()
+
+    # Fall back to last non-empty lines
+    non_empty = [l.strip() for l in lines if l.strip()]
+    if non_empty:
+        return "\n".join(non_empty[-3:])
+    return None
+
+
+class WorkerSpawner:
+    """Orchestrates parallel worker subprocesses.
+
+    Creates worktrees for each task and runs agents in parallel,
+    collecting results as they complete.
+
+    Example:
+        spawner = WorkerSpawner(repo_root=Path("."))
+
+        tasks = [
+            SwarmTask(id="1", slug="add-auth", ...),
+            SwarmTask(id="2", slug="fix-bug", ...),
+        ]
+
+        results = spawner.run_parallel(tasks, max_workers=3)
+    """
+
+    def __init__(
+        self,
+        repo_root: Path,
+        model: str = "gpt-4o-mini",
+        include_graph_tools: bool = False,
+        timeout_per_task: int = 600,
+    ):
+        self.repo_root = repo_root.resolve()
+        self.worktree_manager = WorktreeManager(repo_root)
+        self.model = model
+        self.include_graph_tools = include_graph_tools
+        self.timeout = timeout_per_task
+
+    def prepare_worktree(self, task: SwarmTask) -> WorktreeInfo:
+        """Create worktree and initialize task state."""
+        info = self.worktree_manager.create_worktree(
+            task_name=task.slug,
+            base_branch=task.base_branch,
+            force=True,
+        )
+
+        # Update task with worktree info
+        task.worktree_path = str(info.path)
+        task.branch = info.branch
+        task.status = TaskStatus.WORKTREE_CREATED
+
+        # Save task state to worktree
+        task.save(info.path)
+
+        return info
+
+    def run_parallel(
+        self,
+        tasks: list[SwarmTask],
+        max_workers: int = 3,
+        progress_callback: Optional[Callable[[SwarmTask, WorkerResult], None]] = None,
+    ) -> dict[str, WorkerResult]:
+        """Run all tasks in parallel.
+
+        Args:
+            tasks: List of tasks to execute
+            max_workers: Maximum concurrent workers
+            progress_callback: Called when each task completes
+
+        Returns:
+            Dict mapping task_id to WorkerResult
+        """
+        results: dict[str, WorkerResult] = {}
+
+        # Create all worktrees first
+        log.info(f"Creating {len(tasks)} worktrees...")
+        for task in tasks:
+            self.prepare_worktree(task)
+            task.status = TaskStatus.RUNNING
+            task.started_at = datetime.now().isoformat()
+            if task.worktree_path:
+                task.save(Path(task.worktree_path))
+
+        # Run agents in parallel using ProcessPoolExecutor
+        log.info(f"Starting {len(tasks)} agents with max {max_workers} workers...")
+
+        with ProcessPoolExecutor(max_workers=max_workers) as executor:
+            # Submit all tasks
+            future_to_task = {
+                executor.submit(
+                    run_agent_in_worktree,
+                    task.worktree_path,
+                    task.description,
+                    self.model,
+                    self.include_graph_tools,
+                    self.timeout,
+                ): task
+                for task in tasks
+            }
+
+            # Collect results as they complete
+            for future in as_completed(future_to_task):
+                task = future_to_task[future]
+
+                try:
+                    result = future.result()
+                    results[task.id] = result
+
+                    # Update task status
+                    task.status = TaskStatus.COMPLETED if result.success else TaskStatus.FAILED
+                    task.files_modified = result.files_modified
+                    task.completion_summary = result.completion_summary
+                    task.error_message = result.error_message
+                    task.agent_output = result.stdout[:5000] if result.stdout else None
+                    task.completed_at = datetime.now().isoformat()
+
+                    # Save updated state
+                    if task.worktree_path:
+                        task.save(Path(task.worktree_path))
+
+                    if progress_callback:
+                        progress_callback(task, result)
+
+                    status = "SUCCESS" if result.success else "FAILED"
+                    log.info(f"Task {task.slug} completed: {status}")
+
+                except Exception as e:
+                    log.exception(f"Task {task.slug} raised exception")
+                    task.status = TaskStatus.FAILED
+                    task.error_message = str(e)
+                    task.completed_at = datetime.now().isoformat()
+
+                    if task.worktree_path:
+                        task.save(Path(task.worktree_path))
+
+                    results[task.id] = WorkerResult(
+                        task_id=task.id,
+                        success=False,
+                        files_modified=[],
+                        completion_summary=None,
+                        error_message=str(e),
+                        stdout="",
+                        stderr="",
+                        return_code=-1,
+                    )
+
+        return results
+
+    def run_sequential(
+        self,
+        tasks: list[SwarmTask],
+        progress_callback: Optional[Callable[[SwarmTask, WorkerResult], None]] = None,
+    ) -> dict[str, WorkerResult]:
+        """Run tasks sequentially (for debugging or dependent tasks).
+
+        Args:
+            tasks: List of tasks to execute
+            progress_callback: Called when each task completes
+
+        Returns:
+            Dict mapping task_id to WorkerResult
+        """
+        return self.run_parallel(tasks, max_workers=1, progress_callback=progress_callback)