llmcode-cli 1.0.0__py3-none-any.whl

llm_code/swarm/coordinator.py

@@ -0,0 +1,236 @@
+"""Coordinator — auto-decompose tasks and dispatch to swarm workers."""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import re
+from typing import Any
+
+from llm_code.api.provider import LLMProvider
+from llm_code.api.types import Message, MessageRequest, TextBlock
+from llm_code.swarm.manager import SwarmManager
+
+logger = logging.getLogger(__name__)
+
+_DECOMPOSE_PROMPT = """\
+Break the following task into independent subtasks that can be executed in parallel by specialized agents.
+
+Output ONLY a JSON array (no explanation, no markdown fences). Each element must have:
+  - "role": a short role label (e.g. "coder", "tester", "reviewer", "researcher")
+  - "task": a clear, self-contained task description
+
+Example output:
+[
+  {{"role": "coder", "task": "Implement the binary search function in utils.py"}},
+  {{"role": "tester", "task": "Write unit tests for the binary search function"}}
+]
+
+Task to decompose:
+{task}
+"""
+
+_AGGREGATE_PROMPT = """\
+You are a coordinator agent summarizing the results of parallel worker agents.
+
+Original task: {original_task}
+
+Worker results:
+{results}
+
+Provide a concise summary of what was accomplished, any issues encountered, and the combined outcome.
+"""
+
+
+class Coordinator:
+    """Orchestrate task decomposition and parallel worker dispatch.
+
+    Sends the original task to the LLM for decomposition into subtasks,
+    creates swarm members per subtask, monitors completion via the mailbox,
+    then aggregates results with a final LLM summary.
+    """
+
+    POLL_INTERVAL: float = 5.0
+    TIMEOUT: float = 300.0
+    COORDINATOR_ID: str = "coordinator"
+
+    def __init__(
+        self,
+        manager: SwarmManager,
+        provider: LLMProvider,
+        config: Any,
+    ) -> None:
+        self._manager = manager
+        self._provider = provider
+        self._config = config
+
+    async def orchestrate(self, task: str) -> str:
+        """Decompose task, dispatch workers, wait for completion, return summary.
+
+        Args:
+            task: High-level task description to decompose and delegate.
+
+        Returns:
+            Aggregated summary string from all worker results.
+        """
+        subtasks = await self._decompose(task)
+        if not subtasks:
+            return f"No subtasks generated for: {task}"
+
+        max_members = getattr(
+            getattr(self._config, "swarm", None), "max_members", 5
+        )
+        subtasks = subtasks[:max_members]
+
+        members = []
+        for subtask in subtasks:
+            role = subtask.get("role", "worker")
+            subtask_desc = subtask.get("task", "")
+            if not subtask_desc:
+                continue
+            try:
+                member = await self._manager.create_member(role=role, task=subtask_desc)
+                members.append(member)
+                logger.info("Spawned swarm member %s (%s)", member.id, role)
+            except ValueError as exc:
+                logger.warning("Could not create member for role=%s: %s", role, exc)
+                break
+
+        if not members:
+            return "Failed to create any swarm members."
+
+        results = await self._wait_for_completion(
+            member_ids=[m.id for m in members],
+            timeout=self.TIMEOUT,
+            poll_interval=self.POLL_INTERVAL,
+        )
+
+        summary = await self._aggregate(task, members, results)
+        return summary
+
+    async def _decompose(self, task: str) -> list[dict]:
+        """Ask the LLM to decompose task into subtasks. Returns list of dicts."""
+        prompt = _DECOMPOSE_PROMPT.format(task=task)
+        model = getattr(self._config, "model", None) or "default"
+        request = MessageRequest(
+            model=model,
+            messages=(
+                Message(
+                    role="user",
+                    content=(TextBlock(text=prompt),),
+                ),
+            ),
+            max_tokens=1024,
+            stream=False,
+        )
+        try:
+            response = await self._provider.send_message(request)
+            text = ""
+            for block in response.content:
+                if isinstance(block, TextBlock):
+                    text += block.text
+            return self._parse_json_list(text)
+        except Exception as exc:
+            logger.error("Decomposition failed: %s", exc)
+            return []
+
+    def _parse_json_list(self, text: str) -> list[dict]:
+        """Extract a JSON array from LLM output (strips markdown fences)."""
+        # Strip markdown code fences if present
+        cleaned = re.sub(r"```(?:json)?\s*", "", text).strip().rstrip("`").strip()
+        # Find first '[' and last ']'
+        start = cleaned.find("[")
+        end = cleaned.rfind("]")
+        if start == -1 or end == -1:
+            logger.warning("Could not find JSON array in decomposition output: %r", text[:200])
+            return []
+        try:
+            data = json.loads(cleaned[start : end + 1])
+            if isinstance(data, list):
+                return [item for item in data if isinstance(item, dict)]
+            return []
+        except json.JSONDecodeError as exc:
+            logger.warning("JSON parse error in decomposition: %s", exc)
+            return []
+
+    async def _wait_for_completion(
+        self,
+        member_ids: list[str],
+        timeout: float,
+        poll_interval: float,
+    ) -> dict[str, list[str]]:
+        """Poll the mailbox until all members report completion or timeout.
+
+        Each member is expected to send a message to COORDINATOR_ID containing
+        "DONE" or "COMPLETE" when finished.
+
+        Returns:
+            Mapping of member_id -> list of message texts received.
+        """
+        results: dict[str, list[str]] = {mid: [] for mid in member_ids}
+        completed: set[str] = set()
+        elapsed = 0.0
+
+        while len(completed) < len(member_ids) and elapsed < timeout:
+            for mid in member_ids:
+                if mid in completed:
+                    continue
+                msgs = self._manager.mailbox.receive_and_clear(
+                    from_id=mid, to_id=self.COORDINATOR_ID
+                )
+                for msg in msgs:
+                    results[mid].append(msg.text)
+                    upper = msg.text.upper()
+                    if "DONE" in upper or "COMPLETE" in upper or "FINISHED" in upper:
+                        completed.add(mid)
+                        logger.info("Member %s reported completion", mid)
+
+            if len(completed) < len(member_ids):
+                await asyncio.sleep(poll_interval)
+                elapsed += poll_interval
+
+        if elapsed >= timeout:
+            logger.warning(
+                "Coordinator timed out waiting for members: %s",
+                [mid for mid in member_ids if mid not in completed],
+            )
+
+        return results
+
+    async def _aggregate(self, original_task: str, members: list, results: dict[str, list[str]]) -> str:
+        """Ask the LLM to aggregate all worker results into a summary."""
+        result_lines = []
+        for member in members:
+            texts = results.get(member.id, [])
+            result_lines.append(
+                f"[{member.role}] {member.task}\n"
+                + (("\n".join(texts)) if texts else "(no output received)")
+            )
+
+        results_text = "\n\n".join(result_lines)
+        prompt = _AGGREGATE_PROMPT.format(
+            original_task=original_task,
+            results=results_text,
+        )
+        model = getattr(self._config, "model", None) or "default"
+        request = MessageRequest(
+            model=model,
+            messages=(
+                Message(
+                    role="user",
+                    content=(TextBlock(text=prompt),),
+                ),
+            ),
+            max_tokens=2048,
+            stream=False,
+        )
+        try:
+            response = await self._provider.send_message(request)
+            text = ""
+            for block in response.content:
+                if isinstance(block, TextBlock):
+                    text += block.text
+            return text.strip() or results_text
+        except Exception as exc:
+            logger.error("Aggregation LLM call failed: %s", exc)
+            return results_text

llm_code/swarm/mailbox.py

@@ -0,0 +1,88 @@
+"""File-based JSONL mailbox for inter-agent communication."""
+from __future__ import annotations
+
+import fcntl
+import json
+from datetime import datetime, timezone
+from pathlib import Path
+
+from llm_code.swarm.types import SwarmMessage
+
+
+class Mailbox:
+    """JSONL-based message passing between swarm members.
+
+    Messages stored at: <base_dir>/<sender>_to_<receiver>.jsonl
+    Uses file locking to prevent concurrent write corruption.
+    """
+
+    def __init__(self, base_dir: Path) -> None:
+        self._dir = Path(base_dir)
+        self._dir.mkdir(parents=True, exist_ok=True)
+
+    def send(self, from_id: str, to_id: str, text: str) -> SwarmMessage:
+        """Append a message to the sender->receiver JSONL file (with file lock)."""
+        ts = datetime.now(timezone.utc).isoformat()
+        msg = SwarmMessage(from_id=from_id, to_id=to_id, text=text, timestamp=ts)
+        path = self._msg_path(from_id, to_id)
+        with open(path, "a", encoding="utf-8") as f:
+            fcntl.flock(f, fcntl.LOCK_EX)
+            try:
+                f.write(json.dumps({
+                    "from_id": msg.from_id,
+                    "to_id": msg.to_id,
+                    "text": msg.text,
+                    "timestamp": msg.timestamp,
+                }) + "\n")
+            finally:
+                fcntl.flock(f, fcntl.LOCK_UN)
+        return msg
+
+    def receive(self, from_id: str, to_id: str) -> list[SwarmMessage]:
+        """Read all messages from sender->receiver."""
+        path = self._msg_path(from_id, to_id)
+        if not path.exists():
+            return []
+        messages: list[SwarmMessage] = []
+        for line in path.read_text(encoding="utf-8").strip().splitlines():
+            if not line:
+                continue
+            data = json.loads(line)
+            messages.append(SwarmMessage(
+                from_id=data["from_id"],
+                to_id=data["to_id"],
+                text=data["text"],
+                timestamp=data["timestamp"],
+            ))
+        return messages
+
+    def receive_and_clear(self, from_id: str, to_id: str) -> list[SwarmMessage]:
+        """Read all messages then delete the file."""
+        msgs = self.receive(from_id, to_id)
+        path = self._msg_path(from_id, to_id)
+        if path.exists():
+            path.unlink()
+        return msgs
+
+    def broadcast(self, from_id: str, to_ids: list[str], text: str) -> list[SwarmMessage]:
+        """Send the same message to multiple receivers."""
+        return [self.send(from_id, to_id, text) for to_id in to_ids]
+
+    def pending_for(self, to_id: str) -> list[SwarmMessage]:
+        """Return all unread messages addressed to a given member."""
+        messages: list[SwarmMessage] = []
+        for path in self._dir.glob(f"*_to_{to_id}.jsonl"):
+            for line in path.read_text(encoding="utf-8").strip().splitlines():
+                if not line:
+                    continue
+                data = json.loads(line)
+                messages.append(SwarmMessage(
+                    from_id=data["from_id"],
+                    to_id=data["to_id"],
+                    text=data["text"],
+                    timestamp=data["timestamp"],
+                ))
+        return messages
+
+    def _msg_path(self, from_id: str, to_id: str) -> Path:
+        return self._dir / f"{from_id}_to_{to_id}.jsonl"

llm_code/swarm/manager.py

@@ -0,0 +1,202 @@
+"""SwarmManager — orchestrate creation, lifecycle, and teardown of swarm members."""
+from __future__ import annotations
+
+import re
+import subprocess as sp
+import uuid
+from pathlib import Path
+
+from llm_code.runtime.config import RuntimeConfig
+from llm_code.swarm.backend_subprocess import SubprocessBackend
+from llm_code.swarm.backend_tmux import TmuxBackend, is_tmux_available
+from llm_code.swarm.backend_worktree import WorktreeBackend
+from llm_code.swarm.mailbox import Mailbox
+from llm_code.swarm.memory_sync import SharedMemory
+from llm_code.swarm.types import SwarmMember, SwarmStatus
+
+
+class SwarmManager:
+    """Manage the lifecycle of swarm worker agents.
+
+    Auto-detects tmux (if available and inside a session), otherwise falls
+    back to subprocess.  Each member is a llm-code --lite process with a
+    role prompt injected at startup.
+    """
+
+    def __init__(
+        self,
+        swarm_dir: Path,
+        max_members: int = 5,
+        backend_preference: str = "auto",
+        config: RuntimeConfig | None = None,
+    ) -> None:
+        self._swarm_dir = Path(swarm_dir)
+        self._swarm_dir.mkdir(parents=True, exist_ok=True)
+        self._max_members = max_members
+        self._backend_preference = backend_preference
+        self._members: dict[str, SwarmMember] = {}
+        self._config = config or RuntimeConfig()
+
+        # Backends (lazily used)
+        self._subprocess_backend = SubprocessBackend(swarm_dir=self._swarm_dir)
+        self._tmux_backend = TmuxBackend()
+        # WorktreeBackend is initialised on demand (requires git + project dir)
+        self._worktree_backend: WorktreeBackend | None = None
+
+        # Shared resources
+        self.mailbox = Mailbox(self._swarm_dir / "mailbox")
+        self.shared_memory = SharedMemory(self._swarm_dir / "memory.json")
+
+    async def create_member(
+        self,
+        role: str,
+        task: str,
+        backend: str = "auto",
+        model: str | None = None,
+    ) -> SwarmMember:
+        """Spawn a new swarm worker.
+
+        Args:
+            role: Role description (e.g. 'security reviewer').
+            task: The task this member should perform.
+            backend: 'tmux', 'subprocess', or 'auto' (default).
+            model: Override the model for this specific member. When None,
+                the effective model is resolved via the 4-level fallback chain.
+
+        Returns:
+            The created SwarmMember.
+
+        Raises:
+            ValueError: If max_members limit is reached.
+        """
+        if len(self._members) >= self._max_members:
+            raise ValueError(
+                f"Cannot create member: max {self._max_members} members reached"
+            )
+
+        member_id = uuid.uuid4().hex[:8]
+        effective_backend = self._resolve_backend(backend)
+        effective_model = self._resolve_model(role, model)
+
+        pid: int | str | None = None
+        if effective_backend == "tmux":
+            pid = self._tmux_backend.spawn(
+                member_id=member_id, role=role, task=task, model=effective_model,
+            )
+        elif effective_backend == "worktree":
+            if self._worktree_backend is None:
+                self._worktree_backend = WorktreeBackend(
+                    project_dir=self._swarm_dir.parent,
+                    config=self._config.swarm.worktree,
+                )
+            pid = await self._worktree_backend.spawn(
+                member_id=member_id, role=role, task=task, model=effective_model,
+            )
+        else:
+            pid = await self._subprocess_backend.spawn(
+                member_id=member_id, role=role, task=task, model=effective_model,
+            )
+
+        member = SwarmMember(
+            id=member_id,
+            role=role,
+            task=task,
+            backend=effective_backend,
+            pid=pid if isinstance(pid, int) else None,
+            status=SwarmStatus.RUNNING,
+            model=effective_model,
+        )
+        self._members[member_id] = member
+        return member
+
+    def _resolve_model(self, role: str, explicit: str | None) -> str:
+        """Determine the effective model using a 4-level fallback chain.
+
+        Priority (highest to lowest):
+          1. explicit argument
+          2. config.swarm.role_models[role]
+          3. config.model_routing.sub_agent
+          4. config.model
+
+        The resolved value is then looked up in config.model_aliases.
+        """
+        if explicit:
+            model = explicit
+        elif role in self._config.swarm.role_models:
+            model = self._config.swarm.role_models[role]
+        elif self._config.model_routing.sub_agent:
+            model = self._config.model_routing.sub_agent
+        else:
+            model = self._config.model
+        return self._config.model_aliases.get(model, model)
+
+    def list_members(self) -> list[SwarmMember]:
+        """Return all current swarm members."""
+        return list(self._members.values())
+
+    async def stop_member(self, member_id: str) -> None:
+        """Stop and remove a swarm member.
+
+        Raises:
+            KeyError: If member_id is not found.
+        """
+        member = self._members.get(member_id)
+        if member is None:
+            raise KeyError(f"No swarm member with id '{member_id}'")
+
+        if member.backend == "tmux":
+            self._tmux_backend.stop(member_id)
+        else:
+            await self._subprocess_backend.stop(member_id)
+
+        del self._members[member_id]
+
+    async def stop_all(self) -> None:
+        """Stop all swarm members."""
+        await self._subprocess_backend.stop_all()
+        self._tmux_backend.stop_all()
+        self._members.clear()
+
+    def _resolve_backend(self, requested: str) -> str:
+        """Determine which backend to use.
+
+        Priority for explicit requests: worktree > tmux > subprocess.
+        In auto mode: worktree (if git available) > tmux (if available) > subprocess.
+        """
+        if requested == "worktree":
+            return "worktree"
+        if requested == "tmux":
+            return "tmux"
+        if requested == "subprocess":
+            return "subprocess"
+        # auto path — honour backend_preference first
+        pref = self._backend_preference
+        if pref == "worktree":
+            return "worktree"
+        if pref == "tmux":
+            return "tmux"
+        # pref == "auto": try worktree > tmux > subprocess
+        if self._is_git_repo() and self._git_supports_worktree():
+            return "worktree"
+        return "tmux" if is_tmux_available() else "subprocess"
+
+    def _is_git_repo(self) -> bool:
+        """Return True if the project directory is inside a git repository."""
+        result = sp.run(
+            ["git", "rev-parse", "--is-inside-work-tree"],
+            cwd=str(self._swarm_dir.parent),
+            capture_output=True,
+            text=True,
+        )
+        return result.returncode == 0
+
+    def _git_supports_worktree(self) -> bool:
+        """Return True if the installed git version supports worktrees (>= 2.15)."""
+        result = sp.run(["git", "--version"], capture_output=True, text=True)
+        if result.returncode != 0:
+            return False
+        match = re.search(r"(\d+)\.(\d+)", result.stdout)
+        if not match:
+            return False
+        major, minor = int(match.group(1)), int(match.group(2))
+        return (major, minor) >= (2, 15)

llm_code/swarm/memory_sync.py

@@ -0,0 +1,80 @@
+"""Shared memory store with file locking for swarm members."""
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+
+def _lock_file(f) -> None:  # type: ignore[no-untyped-def]
+    """Acquire an exclusive lock on a file handle (platform-aware)."""
+    if sys.platform == "win32":
+        import msvcrt
+        msvcrt.locking(f.fileno(), msvcrt.LK_LOCK, 1)
+    else:
+        import fcntl
+        fcntl.flock(f, fcntl.LOCK_EX)
+
+
+def _unlock_file(f) -> None:  # type: ignore[no-untyped-def]
+    """Release a file lock (platform-aware)."""
+    if sys.platform == "win32":
+        import msvcrt
+        msvcrt.locking(f.fileno(), msvcrt.LK_UNLCK, 1)
+    else:
+        import fcntl
+        fcntl.flock(f, fcntl.LOCK_UN)
+
+
+class SharedMemory:
+    """JSON-backed shared key-value store with file locking.
+
+    Multiple swarm members can safely read/write to the same file.
+    """
+
+    def __init__(self, path: Path) -> None:
+        self._path = Path(path)
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+
+    def write(self, key: str, value: str) -> None:
+        """Set a key-value pair (locked)."""
+        data = self._locked_read()
+        data[key] = value
+        self._locked_write(data)
+
+    def read(self, key: str) -> str | None:
+        """Get a value by key, or None if missing."""
+        data = self._locked_read()
+        return data.get(key)
+
+    def read_all(self) -> dict[str, str]:
+        """Return the entire shared memory dict."""
+        return self._locked_read()
+
+    def delete(self, key: str) -> None:
+        """Remove a key (no-op if missing)."""
+        data = self._locked_read()
+        data.pop(key, None)
+        self._locked_write(data)
+
+    def _locked_read(self) -> dict[str, str]:
+        if not self._path.exists():
+            return {}
+        try:
+            with open(self._path, "r", encoding="utf-8") as f:
+                _lock_file(f)
+                try:
+                    content = f.read()
+                    return json.loads(content) if content.strip() else {}
+                finally:
+                    _unlock_file(f)
+        except (json.JSONDecodeError, OSError):
+            return {}
+
+    def _locked_write(self, data: dict[str, str]) -> None:
+        with open(self._path, "w", encoding="utf-8") as f:
+            _lock_file(f)
+            try:
+                json.dump(data, f, indent=2)
+            finally:
+                _unlock_file(f)

llm_code/swarm/recovery.py

@@ -0,0 +1,21 @@
+"""Recovery policy for agent teams — retry logic and failure handling."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class RecoveryPolicy:
+    max_retries: int = 2
+    retry_delay_sec: int = 5
+    on_all_failed: str = "abort"
+
+
+class RecoveryAction:
+    @staticmethod
+    def should_retry(policy: RecoveryPolicy, attempt: int) -> bool:
+        return attempt <= policy.max_retries
+
+    @staticmethod
+    def resolve_all_failed(policy: RecoveryPolicy) -> str:
+        return policy.on_all_failed

llm_code/swarm/team.py

@@ -0,0 +1,67 @@
+"""Team template — save/load reusable agent team configurations."""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class TeamMemberTemplate:
+    """A template for a single team member."""
+    role: str
+    model: str = ""
+    backend: str = ""
+    system_prompt: str = ""
+
+
+@dataclass(frozen=True)
+class TeamTemplate:
+    """A reusable team configuration."""
+    name: str
+    description: str
+    members: tuple[TeamMemberTemplate, ...]
+    coordinator_model: str = ""
+    max_timeout: int = 600
+
+
+def save_team(team: TeamTemplate, teams_dir: Path) -> Path:
+    teams_dir.mkdir(parents=True, exist_ok=True)
+    path = teams_dir / f"{team.name}.json"
+    data = {
+        "name": team.name,
+        "description": team.description,
+        "members": [
+            {"role": m.role, "model": m.model, "backend": m.backend, "system_prompt": m.system_prompt}
+            for m in team.members
+        ],
+        "coordinator_model": team.coordinator_model,
+        "max_timeout": team.max_timeout,
+    }
+    path.write_text(json.dumps(data, indent=2), encoding="utf-8")
+    return path
+
+
+def load_team(name: str, teams_dir: Path) -> TeamTemplate:
+    path = teams_dir / f"{name}.json"
+    if not path.exists():
+        raise FileNotFoundError(f"Team template not found: {path}")
+    data = json.loads(path.read_text(encoding="utf-8"))
+    members = tuple(
+        TeamMemberTemplate(
+            role=m["role"], model=m.get("model", ""),
+            backend=m.get("backend", ""), system_prompt=m.get("system_prompt", ""),
+        )
+        for m in data.get("members", [])
+    )
+    return TeamTemplate(
+        name=data["name"], description=data.get("description", ""),
+        members=members, coordinator_model=data.get("coordinator_model", ""),
+        max_timeout=data.get("max_timeout", 600),
+    )
+
+
+def list_teams(teams_dir: Path) -> list[str]:
+    if not teams_dir.is_dir():
+        return []
+    return sorted(p.stem for p in teams_dir.glob("*.json"))