zwarm 2.3.5__py3-none-any.whl

zwarm/__init__.py

@@ -0,0 +1,38 @@
+"""
+zwarm: Multi-Agent CLI Orchestration Research Platform
+
+A framework for orchestrating multiple CLI coding agents (codex, claude-code, gemini)
+with support for sync (conversational) and async (fire-and-forget) delegation.
+"""
+
+from zwarm.core.config import ZwarmConfig, load_config
+from zwarm.core.models import (
+    ConversationSession,
+    Event,
+    Message,
+    SessionMode,
+    SessionStatus,
+    Task,
+    TaskStatus,
+)
+from zwarm.core.state import StateManager
+from zwarm.orchestrator import Orchestrator, build_orchestrator
+
+__all__ = [
+    # Config
+    "ZwarmConfig",
+    "load_config",
+    # Models
+    "ConversationSession",
+    "Event",
+    "Message",
+    "SessionMode",
+    "SessionStatus",
+    "Task",
+    "TaskStatus",
+    # State
+    "StateManager",
+    # Orchestrator
+    "Orchestrator",
+    "build_orchestrator",
+]

zwarm/adapters/__init__.py

@@ -0,0 +1,21 @@
+"""
+Adapters: Executor wrappers for CLI coding agents.
+
+Adapters provide a unified interface to different coding CLIs (Codex, Claude Code).
+Use the registry to discover and instantiate adapters by name.
+"""
+
+from zwarm.adapters.base import ExecutorAdapter
+from zwarm.adapters.registry import register_adapter, get_adapter, list_adapters, adapter_exists
+
+# Import built-in adapters to register them
+from zwarm.adapters import codex_mcp as _codex_mcp  # noqa: F401
+from zwarm.adapters import claude_code as _claude_code  # noqa: F401
+
+__all__ = [
+    "ExecutorAdapter",
+    "register_adapter",
+    "get_adapter",
+    "list_adapters",
+    "adapter_exists",
+]

zwarm/adapters/base.py

@@ -0,0 +1,109 @@
+"""
+Base adapter protocol for executor agents.
+
+All CLI coding agent adapters (codex, claude-code, gemini) implement this protocol.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Literal
+
+from zwarm.core.models import ConversationSession, SessionMode
+
+
+class ExecutorAdapter(ABC):
+    """
+    Abstract base class for CLI coding agent adapters.
+
+    Adapters handle the mechanics of:
+    - Starting sessions (sync or async)
+    - Sending messages in sync mode
+    - Checking status in async mode
+    - Stopping sessions
+    """
+
+    name: str = "base"
+
+    @abstractmethod
+    async def start_session(
+        self,
+        task: str,
+        working_dir: Path,
+        mode: Literal["sync", "async"] = "sync",
+        model: str | None = None,
+        **kwargs,
+    ) -> ConversationSession:
+        """
+        Start a new session with the executor.
+
+        Args:
+            task: The task description/prompt
+            working_dir: Directory to work in
+            mode: "sync" for conversational, "async" for fire-and-forget
+            model: Optional model override
+            **kwargs: Adapter-specific options
+
+        Returns:
+            A ConversationSession with initial response (if sync)
+        """
+        ...
+
+    @abstractmethod
+    async def send_message(
+        self,
+        session: ConversationSession,
+        message: str,
+    ) -> str:
+        """
+        Send a message to a sync session and get response.
+
+        Args:
+            session: The active session
+            message: Message to send
+
+        Returns:
+            The agent's response
+
+        Raises:
+            ValueError: If session is not in sync mode or not active
+        """
+        ...
+
+    @abstractmethod
+    async def check_status(
+        self,
+        session: ConversationSession,
+    ) -> dict:
+        """
+        Check the status of an async session.
+
+        Args:
+            session: The session to check
+
+        Returns:
+            Status dict with at least {"status": "running"|"completed"|"failed"}
+        """
+        ...
+
+    @abstractmethod
+    async def stop(
+        self,
+        session: ConversationSession,
+    ) -> None:
+        """
+        Stop/kill a session.
+
+        Args:
+            session: The session to stop
+        """
+        ...
+
+    async def cleanup(self) -> None:
+        """
+        Clean up adapter resources (e.g., MCP server).
+
+        Called when the orchestrator shuts down.
+        """
+        pass

zwarm/adapters/claude_code.py

@@ -0,0 +1,357 @@
+"""
+Claude Code adapter for sync/async execution.
+
+Uses the claude CLI for conversations:
+- claude -p --output-format json for non-interactive mode
+- claude -r <session_id> to continue conversations
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import subprocess
+from pathlib import Path
+from typing import Any, Literal
+
+import weave
+
+from zwarm.adapters.base import ExecutorAdapter
+from zwarm.adapters.registry import register_adapter
+from zwarm.core.models import (
+    ConversationSession,
+    SessionMode,
+    SessionStatus,
+)
+
+
+@register_adapter("claude_code")
+class ClaudeCodeAdapter(ExecutorAdapter):
+    """
+    Claude Code adapter using the claude CLI.
+
+    Supports both sync (conversational) and async (fire-and-forget) modes.
+    """
+    DEFAULT_MODEL = "claude-sonnet-4-5-20250514"  # Best balance of speed and capability
+
+    def __init__(self, model: str | None = None):
+        self._model = model or self.DEFAULT_MODEL
+        self._sessions: dict[str, str] = {}  # session_id -> claude session_id
+        # Cumulative token usage for cost tracking
+        self._total_usage: dict[str, int] = {
+            "input_tokens": 0,
+            "output_tokens": 0,
+            "cache_creation_input_tokens": 0,
+            "cache_read_input_tokens": 0,
+            "total_tokens": 0,
+        }
+
+    def _accumulate_usage(self, usage: dict[str, Any]) -> None:
+        """Add usage to cumulative totals."""
+        if not usage:
+            return
+        for key in self._total_usage:
+            self._total_usage[key] += usage.get(key, 0)
+
+    def _extract_usage(self, output: dict) -> dict[str, int]:
+        """Extract token usage from claude CLI JSON output."""
+        usage = {}
+        # Claude CLI may include usage in various formats
+        if "usage" in output:
+            usage = output["usage"]
+        elif "cost_usd" in output:
+            # Alternative: estimate from cost if available
+            usage["cost_usd"] = output["cost_usd"]
+        # Also check for token counts in the output
+        for key in ["input_tokens", "output_tokens", "total_tokens"]:
+            if key in output:
+                usage[key] = output[key]
+        return usage
+
+    @property
+    def total_usage(self) -> dict[str, int]:
+        """Get cumulative token usage across all calls."""
+        return self._total_usage.copy()
+
+    @weave.op()
+    async def _call_claude(
+        self,
+        task: str,
+        cwd: str,
+        model: str | None = None,
+        permission_mode: str = "bypassPermissions",
+    ) -> dict[str, Any]:
+        """
+        Call claude CLI - traced by Weave.
+
+        This wraps the actual claude call so it appears in Weave traces
+        with full input/output visibility.
+        """
+        cmd = ["claude", "-p", "--output-format", "json"]
+
+        if permission_mode:
+            cmd.extend(["--permission-mode", permission_mode])
+        if model:
+            cmd.extend(["--model", model])
+
+        cmd.extend(["--", task])
+
+        loop = asyncio.get_event_loop()
+        result = await loop.run_in_executor(
+            None,
+            lambda: subprocess.run(
+                cmd,
+                cwd=cwd,
+                capture_output=True,
+                text=True,
+                timeout=300,
+            )
+        )
+
+        response_text = self._extract_response(result.stdout, result.stderr)
+
+        # Try to get session ID and usage from JSON output
+        session_id = None
+        usage = {}
+        try:
+            output = json.loads(result.stdout)
+            session_id = output.get("session_id")
+            usage = self._extract_usage(output)
+            self._accumulate_usage(usage)
+        except (json.JSONDecodeError, TypeError):
+            pass
+
+        return {
+            "response": response_text,
+            "session_id": session_id,
+            "exit_code": result.returncode,
+            "usage": usage,
+            "total_usage": self.total_usage,
+        }
+
+    @weave.op()
+    async def _call_claude_continue(
+        self,
+        message: str,
+        cwd: str,
+        session_id: str | None = None,
+    ) -> dict[str, Any]:
+        """
+        Continue a claude conversation - traced by Weave.
+
+        This wraps the continuation call so it appears in Weave traces
+        with full input/output visibility.
+        """
+        cmd = ["claude", "-p", "--output-format", "json"]
+
+        if session_id:
+            cmd.extend(["--resume", session_id])
+        else:
+            cmd.extend(["--continue"])
+
+        cmd.extend(["--permission-mode", "bypassPermissions"])
+        cmd.extend(["--", message])
+
+        loop = asyncio.get_event_loop()
+        result = await loop.run_in_executor(
+            None,
+            lambda: subprocess.run(
+                cmd,
+                cwd=cwd,
+                capture_output=True,
+                text=True,
+                timeout=300,
+            )
+        )
+
+        response_text = self._extract_response(result.stdout, result.stderr)
+
+        # Try to get session ID and usage from JSON output
+        new_session_id = None
+        usage = {}
+        try:
+            output = json.loads(result.stdout)
+            new_session_id = output.get("session_id")
+            usage = self._extract_usage(output)
+            self._accumulate_usage(usage)
+        except (json.JSONDecodeError, TypeError):
+            pass
+
+        return {
+            "response": response_text,
+            "usage": usage,
+            "total_usage": self.total_usage,
+            "session_id": new_session_id or session_id,
+            "exit_code": result.returncode,
+        }
+
+    @weave.op()
+    async def start_session(
+        self,
+        task: str,
+        working_dir: Path,
+        mode: Literal["sync", "async"] = "sync",
+        model: str | None = None,
+        permission_mode: str = "bypassPermissions",
+        **kwargs,
+    ) -> ConversationSession:
+        """Start a Claude Code session (sync or async mode)."""
+        session = ConversationSession(
+            adapter=self.name,
+            mode=SessionMode(mode),
+            working_dir=working_dir,
+            task_description=task,
+            model=model or self._model,
+        )
+
+        if mode == "sync":
+            # Use traced claude call
+            result = await self._call_claude(
+                task=task,
+                cwd=str(working_dir),
+                model=model or self._model,
+                permission_mode=permission_mode,
+            )
+
+            # Extract session ID and response
+            if result["session_id"]:
+                session.conversation_id = result["session_id"]
+                self._sessions[session.id] = session.conversation_id
+
+            session.add_message("user", task)
+            session.add_message("assistant", result["response"])
+
+            # Track token usage on the session
+            session.add_usage(result.get("usage", {}))
+
+        else:
+            # Async mode: run in background
+            cmd = ["claude", "-p", "--output-format", "json"]
+            if permission_mode:
+                cmd.extend(["--permission-mode", permission_mode])
+            if model or self._model:
+                cmd.extend(["--model", model or self._model])
+            cmd.extend(["--", task])
+
+            proc = subprocess.Popen(
+                cmd,
+                cwd=working_dir,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+            session.process = proc
+            session.add_message("user", task)
+
+        return session
+
+    async def send_message(
+        self,
+        session: ConversationSession,
+        message: str,
+    ) -> str:
+        """Send a message to continue a sync conversation."""
+        if session.mode != SessionMode.SYNC:
+            raise ValueError("Cannot send message to async session")
+        if session.status != SessionStatus.ACTIVE:
+            raise ValueError(f"Session is not active: {session.status}")
+
+        # Use traced continuation call
+        result = await self._call_claude_continue(
+            message=message,
+            cwd=str(session.working_dir),
+            session_id=session.conversation_id,
+        )
+
+        # Update session ID if we didn't have one
+        if not session.conversation_id and result["session_id"]:
+            session.conversation_id = result["session_id"]
+            self._sessions[session.id] = session.conversation_id
+
+        response_text = result["response"]
+        session.add_message("user", message)
+        session.add_message("assistant", response_text)
+
+        # Track token usage on the session
+        session.add_usage(result.get("usage", {}))
+
+        return response_text
+
+    @weave.op()
+    async def check_status(
+        self,
+        session: ConversationSession,
+    ) -> dict:
+        """Check status of an async session."""
+        if session.mode != SessionMode.ASYNC:
+            return {"status": session.status.value}
+
+        if session.process is None:
+            return {"status": "unknown", "error": "No process handle"}
+
+        # Check if process is still running
+        poll = session.process.poll()
+        if poll is None:
+            return {"status": "running"}
+
+        # Process finished
+        stdout, stderr = session.process.communicate()
+        if poll == 0:
+            response = self._extract_response(stdout, stderr)
+            session.complete(response[:1000] if response else "Completed")
+            return {"status": "completed", "output": response}
+        else:
+            session.fail(stderr[:1000] if stderr else f"Exit code: {poll}")
+            return {"status": "failed", "error": stderr, "exit_code": poll}
+
+    async def stop(
+        self,
+        session: ConversationSession,
+    ) -> None:
+        """Stop a session."""
+        if session.process and session.process.poll() is None:
+            session.process.terminate()
+            try:
+                session.process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                session.process.kill()
+
+        session.fail("Stopped by user")
+
+        # Remove from tracking
+        if session.id in self._sessions:
+            del self._sessions[session.id]
+
+    def _extract_response(self, stdout: str, stderr: str) -> str:
+        """Extract response text from CLI output."""
+        # Try to parse as JSON
+        try:
+            output = json.loads(stdout)
+
+            # Check for result/response fields
+            if "result" in output:
+                return output["result"]
+            if "response" in output:
+                return output["response"]
+            if "content" in output:
+                return output["content"]
+            if "text" in output:
+                return output["text"]
+
+            # Handle messages array
+            if "messages" in output and isinstance(output["messages"], list):
+                for msg in reversed(output["messages"]):
+                    if isinstance(msg, dict) and msg.get("role") == "assistant":
+                        return msg.get("content", "")
+
+            # Fallback: stringify the output
+            return json.dumps(output, indent=2)
+
+        except json.JSONDecodeError:
+            # Not JSON, return raw output
+            if stdout.strip():
+                return stdout.strip()
+            if stderr.strip():
+                return f"Error: {stderr.strip()}"
+            return "(no output)"