aiagent-runner 0.1.3__py3-none-any.whl

aiagent_runner/mcp_client.py

@@ -0,0 +1,698 @@
+# src/aiagent_runner/mcp_client.py
+# MCP client for communication with AI Agent PM server
+# Reference: docs/plan/PHASE3_PULL_ARCHITECTURE.md - Phase 3-5
+# Reference: docs/plan/PHASE4_COORDINATOR_ARCHITECTURE.md
+# Reference: docs/design/MULTI_DEVICE_IMPLEMENTATION_PLAN.md - Phase 4.3
+
+import asyncio
+import json
+import logging
+import os
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Optional
+
+# HTTP transport support (optional dependency)
+try:
+    import aiohttp
+    HAS_AIOHTTP = True
+except ImportError:
+    HAS_AIOHTTP = False
+
+logger = logging.getLogger(__name__)
+
+
+class AuthenticationError(Exception):
+    """Raised when authentication fails."""
+    pass
+
+
+class SessionExpiredError(Exception):
+    """Raised when session token has expired."""
+    pass
+
+
+class MCPError(Exception):
+    """General MCP communication error."""
+    pass
+
+
+# Phase 4: Coordinator API data classes
+
+@dataclass
+class HealthCheckResult:
+    """Result of health check."""
+    status: str
+    version: Optional[str] = None
+    timestamp: Optional[str] = None
+
+
+@dataclass
+class ProjectWithAgents:
+    """Project with its assigned agents."""
+    project_id: str
+    project_name: str
+    working_directory: str
+    agents: list[str] = field(default_factory=list)
+
+
+@dataclass
+class AgentActionResult:
+    """Result of get_agent_action check."""
+    action: str                          # "start", "hold", "stop", "restart"
+    reason: Optional[str] = None         # Reason for the action
+    provider: Optional[str] = None       # "claude", "gemini", "openai", "other"
+    model: Optional[str] = None          # "claude-sonnet-4-5", "gemini-2.0-flash", etc.
+    kick_command: Optional[str] = None   # Custom CLI command (takes priority if set)
+    task_id: Optional[str] = None        # Phase 4: タスクID（ログファイル登録用）
+
+
+# Phase 3/4: Agent API data classes
+
+@dataclass
+class AuthResult:
+    """Result of authentication."""
+    session_token: str
+    expires_in: int
+    agent_name: Optional[str] = None
+    system_prompt: Optional[str] = None
+    instruction: Optional[str] = None
+
+
+@dataclass
+class TaskInfo:
+    """Information about a pending task."""
+    task_id: str
+    project_id: str
+    title: str
+    description: str
+    priority: str
+    working_directory: Optional[str] = None
+    context: Optional[dict] = None
+    handoff: Optional[dict] = None
+
+
+@dataclass
+class ExecutionStartResult:
+    """Result of reporting execution start."""
+    execution_id: str
+    started_at: datetime
+
+
+class MCPClient:
+    """Client for MCP server communication.
+
+    Handles authentication, task retrieval, and execution reporting.
+
+    Supports two transport modes:
+    - Unix socket: For local connections (default)
+    - HTTP: For remote connections (multi-device operation)
+
+    The transport is automatically selected based on the URL:
+    - http:// or https:// → HTTP transport
+    - Everything else → Unix socket transport
+    """
+
+    def __init__(
+        self,
+        socket_path: Optional[str] = None,
+        coordinator_token: Optional[str] = None
+    ):
+        """Initialize MCP client.
+
+        Args:
+            socket_path: Path to MCP Unix socket, or HTTP URL for remote connections.
+                        Unix socket: ~/Library/Application Support/AIAgentPM/mcp.sock
+                        HTTP URL: http://hostname:port/mcp
+                        Defaults to standard Unix socket location.
+            coordinator_token: Token for Coordinator-only API calls (Phase 5).
+                              If not provided, reads from MCP_COORDINATOR_TOKEN env var.
+        """
+        # Determine transport type based on URL scheme
+        if socket_path and socket_path.startswith(("http://", "https://")):
+            self._url = socket_path
+            self._use_http = True
+            logger.info(f"Using HTTP transport: {self._url}")
+        else:
+            # Unix socket path - expand tilde
+            if socket_path:
+                self._url = os.path.expanduser(socket_path)
+            else:
+                self._url = self._default_socket_path()
+            self._use_http = False
+            logger.info(f"Using Unix socket transport: {self._url}")
+
+        # Backward compatibility
+        self.socket_path = self._url if not self._use_http else None
+
+        self._session_token: Optional[str] = None
+        # Phase 5: Coordinator token for Coordinator-only API calls
+        self._coordinator_token = coordinator_token or os.environ.get("MCP_COORDINATOR_TOKEN")
+
+    def _default_socket_path(self) -> str:
+        """Get default MCP socket path."""
+        return os.path.expanduser(
+            "~/Library/Application Support/AIAgentPM/mcp.sock"
+        )
+
+    async def _call_tool(self, tool_name: str, args: dict) -> dict:
+        """Call an MCP tool via Unix socket or HTTP.
+
+        Automatically selects the transport based on the URL scheme.
+
+        Args:
+            tool_name: Name of the tool to call
+            args: Arguments for the tool
+
+        Returns:
+            Tool result as dictionary
+
+        Raises:
+            MCPError: If communication fails
+        """
+        if self._use_http:
+            return await self._call_tool_http(tool_name, args)
+        else:
+            return await self._call_tool_unix(tool_name, args)
+
+    async def _call_tool_unix(self, tool_name: str, args: dict) -> dict:
+        """Call an MCP tool via Unix socket.
+
+        Args:
+            tool_name: Name of the tool to call
+            args: Arguments for the tool
+
+        Returns:
+            Tool result as dictionary
+
+        Raises:
+            MCPError: If communication fails
+        """
+        try:
+            reader, writer = await asyncio.open_unix_connection(self._url)
+        except (ConnectionRefusedError, FileNotFoundError) as e:
+            raise MCPError(f"Cannot connect to MCP server at {self._url}: {e}")
+
+        try:
+            request = json.dumps({
+                "jsonrpc": "2.0",
+                "method": "tools/call",
+                "params": {"name": tool_name, "arguments": args},
+                "id": 1
+            })
+            writer.write(request.encode() + b"\n")
+            await writer.drain()
+
+            response = await reader.readline()
+            data = json.loads(response)
+
+            return self._parse_response(data)
+        finally:
+            writer.close()
+            await writer.wait_closed()
+
+    async def _call_tool_http(self, tool_name: str, args: dict) -> dict:
+        """Call an MCP tool via HTTP.
+
+        Args:
+            tool_name: Name of the tool to call
+            args: Arguments for the tool
+
+        Returns:
+            Tool result as dictionary
+
+        Raises:
+            MCPError: If communication fails or aiohttp is not installed
+        """
+        if not HAS_AIOHTTP:
+            raise MCPError(
+                "HTTP transport requires aiohttp. Install with: pip install aiohttp"
+            )
+
+        request_body = {
+            "jsonrpc": "2.0",
+            "method": "tools/call",
+            "params": {"name": tool_name, "arguments": args},
+            "id": 1
+        }
+
+        headers = {"Content-Type": "application/json"}
+
+        # Add coordinator token as Authorization header if available
+        if self._coordinator_token:
+            headers["Authorization"] = f"Bearer {self._coordinator_token}"
+
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.post(
+                    self._url,
+                    json=request_body,
+                    headers=headers,
+                    timeout=aiohttp.ClientTimeout(total=30)
+                ) as response:
+                    if response.status != 200:
+                        text = await response.text()
+                        raise MCPError(f"HTTP {response.status}: {text}")
+
+                    data = await response.json()
+                    return self._parse_response(data)
+
+        except aiohttp.ClientError as e:
+            raise MCPError(f"Cannot connect to MCP server at {self._url}: {e}")
+
+    def _parse_response(self, data: dict) -> dict:
+        """Parse MCP JSON-RPC response.
+
+        Args:
+            data: Raw JSON-RPC response
+
+        Returns:
+            Parsed tool result
+
+        Raises:
+            MCPError: If response contains an error
+        """
+        if "error" in data:
+            raise MCPError(data["error"].get("message", "Unknown error"))
+
+        # Parse MCP protocol response format
+        # MCP returns: {"result": {"content": [{"type": "text", "text": "JSON"}]}}
+        result = data.get("result", {})
+        content = result.get("content", [])
+        if content and isinstance(content, list) and len(content) > 0:
+            first_content = content[0]
+            if isinstance(first_content, dict) and first_content.get("type") == "text":
+                text = first_content.get("text", "{}")
+                try:
+                    return json.loads(text)
+                except json.JSONDecodeError:
+                    return {"text": text}
+        return result
+
+    # ==========================================================================
+    # Phase 4: Coordinator API
+    # Reference: docs/plan/PHASE4_COORDINATOR_ARCHITECTURE.md
+    # ==========================================================================
+
+    async def health_check(self) -> HealthCheckResult:
+        """Check MCP server health.
+
+        The Coordinator calls this first to verify the server is available.
+        Phase 5: Requires coordinator_token for authorization.
+
+        Returns:
+            HealthCheckResult with server status
+
+        Raises:
+            MCPError: If server is not available or unauthorized
+        """
+        args = {}
+        if self._coordinator_token:
+            args["coordinator_token"] = self._coordinator_token
+        result = await self._call_tool("health_check", args)
+        return HealthCheckResult(
+            status=result.get("status", "ok"),
+            version=result.get("version"),
+            timestamp=result.get("timestamp")
+        )
+
+    async def list_active_projects_with_agents(
+        self, root_agent_id: Optional[str] = None
+    ) -> list[ProjectWithAgents]:
+        """Get all active projects with their assigned agents.
+
+        The Coordinator calls this to discover what (agent_id, project_id)
+        combinations exist and need to be monitored.
+        Phase 5: Requires coordinator_token for authorization.
+
+        Multi-device operation:
+        When root_agent_id is specified, the server uses that agent's
+        working directories instead of the project defaults.
+
+        Args:
+            root_agent_id: Optional human agent ID for working directory resolution
+
+        Returns:
+            List of ProjectWithAgents
+
+        Raises:
+            MCPError: If request fails or unauthorized
+        """
+        args = {}
+        if self._coordinator_token:
+            args["coordinator_token"] = self._coordinator_token
+            logger.debug("list_active_projects_with_agents: passing coordinator_token")
+        else:
+            logger.warning("list_active_projects_with_agents: NO coordinator_token set!")
+
+        # Multi-device: Pass root_agent_id for working directory resolution
+        if root_agent_id:
+            args["root_agent_id"] = root_agent_id
+            logger.debug(f"list_active_projects_with_agents: passing root_agent_id={root_agent_id}")
+
+        result = await self._call_tool("list_active_projects_with_agents", args)
+        logger.debug(f"list_active_projects_with_agents result: {result}")
+
+        if not result.get("success", True):
+            raise MCPError(result.get("error", "Failed to list projects"))
+
+        projects = []
+        for p in result.get("projects", []):
+            projects.append(ProjectWithAgents(
+                project_id=p.get("project_id", p.get("projectId", "")),
+                project_name=p.get("project_name", p.get("projectName", p.get("name", ""))),
+                working_directory=p.get("working_directory", p.get("workingDirectory", "")),
+                agents=p.get("agents", [])
+            ))
+        return projects
+
+    async def get_agent_action(self, agent_id: str, project_id: str) -> AgentActionResult:
+        """Get the action an Agent Instance should take.
+
+        The Coordinator calls this for each (agent_id, project_id) pair
+        to determine what action to take.
+        Phase 5: Requires coordinator_token for authorization.
+
+        Args:
+            agent_id: Agent ID
+            project_id: Project ID
+
+        Returns:
+            AgentActionResult with action (start/hold/stop/restart), reason, provider, model, and kick_command
+
+        Raises:
+            MCPError: If request fails or unauthorized
+        """
+        args = {
+            "agent_id": agent_id,
+            "project_id": project_id
+        }
+        if self._coordinator_token:
+            args["coordinator_token"] = self._coordinator_token
+        result = await self._call_tool("get_agent_action", args)
+
+        return AgentActionResult(
+            action=result.get("action", "hold"),
+            reason=result.get("reason"),
+            provider=result.get("provider"),
+            model=result.get("model"),
+            kick_command=result.get("kick_command"),
+            task_id=result.get("task_id")  # Phase 4: Coordinatorがログファイルパス登録に使用
+        )
+
+    async def register_execution_log_file(
+        self, agent_id: str, task_id: str, log_file_path: str
+    ) -> bool:
+        """Register log file path for an execution log.
+
+        Called by Coordinator after Agent Instance process completes.
+        Phase 5: Requires coordinator_token for authorization.
+
+        Args:
+            agent_id: Agent ID
+            task_id: Task ID
+            log_file_path: Absolute path to the log file
+
+        Returns:
+            True if successful, False otherwise
+
+        Raises:
+            MCPError: If request fails or unauthorized
+        """
+        args = {
+            "agent_id": agent_id,
+            "task_id": task_id,
+            "log_file_path": log_file_path
+        }
+        if self._coordinator_token:
+            args["coordinator_token"] = self._coordinator_token
+        result = await self._call_tool("register_execution_log_file", args)
+
+        return result.get("success", False)
+
+    async def invalidate_session(self, agent_id: str, project_id: str) -> bool:
+        """Invalidate session for an agent-project pair.
+
+        Called by Coordinator when Agent Instance process exits.
+        This allows shouldStart to return True again for the next instance.
+        Phase 5: Requires coordinator_token for authorization.
+
+        Args:
+            agent_id: Agent ID
+            project_id: Project ID
+
+        Returns:
+            True if successful, False otherwise
+
+        Raises:
+            MCPError: If request fails or unauthorized
+        """
+        args = {
+            "agent_id": agent_id,
+            "project_id": project_id
+        }
+        if self._coordinator_token:
+            args["coordinator_token"] = self._coordinator_token
+        result = await self._call_tool("invalidate_session", args)
+
+        return result.get("success", False)
+
+    async def report_agent_error(
+        self, agent_id: str, project_id: str, error_message: str
+    ) -> bool:
+        """Report agent error to chat.
+
+        Called by Coordinator when Agent Instance process exits with error.
+        The error message will be displayed in the chat.
+
+        Args:
+            agent_id: Agent ID
+            project_id: Project ID
+            error_message: Error message to display
+
+        Returns:
+            True if successful, False otherwise
+
+        Raises:
+            MCPError: If request fails or unauthorized
+        """
+        args = {
+            "agent_id": agent_id,
+            "project_id": project_id,
+            "error_message": error_message
+        }
+        if self._coordinator_token:
+            args["coordinator_token"] = self._coordinator_token
+        result = await self._call_tool("report_agent_error", args)
+
+        return result.get("success", False)
+
+    # ==========================================================================
+    # Phase 3/4: Agent Instance API
+    # ==========================================================================
+
+    async def authenticate(self, agent_id: str, passkey: str, project_id: str) -> AuthResult:
+        """Authenticate with the MCP server.
+
+        Args:
+            agent_id: Agent ID
+            passkey: Agent passkey
+            project_id: Project ID (Phase 4: required for session management)
+
+        Returns:
+            AuthResult with session token
+
+        Raises:
+            AuthenticationError: If authentication fails
+        """
+        result = await self._call_tool("authenticate", {
+            "agent_id": agent_id,
+            "passkey": passkey,
+            "project_id": project_id
+        })
+
+        if not result.get("success"):
+            raise AuthenticationError(result.get("error", "Authentication failed"))
+
+        self._session_token = result["session_token"]
+        return AuthResult(
+            session_token=result["session_token"],
+            expires_in=result.get("expires_in", 3600),
+            agent_name=result.get("agent_name"),
+            system_prompt=result.get("system_prompt"),
+            instruction=result.get("instruction")
+        )
+
+    async def get_pending_tasks(self) -> list[TaskInfo]:
+        """Get pending tasks for the authenticated agent.
+
+        Returns:
+            List of pending TaskInfo objects
+
+        Raises:
+            SessionExpiredError: If session has expired
+            MCPError: If request fails or not authenticated
+        """
+        if not self._session_token:
+            raise MCPError("Not authenticated. Call authenticate() first.")
+
+        result = await self._call_tool("get_pending_tasks", {
+            "session_token": self._session_token
+        })
+
+        if not result.get("success"):
+            error = result.get("error", "")
+            if "expired" in error.lower() or "invalid" in error.lower():
+                raise SessionExpiredError(error)
+            raise MCPError(error)
+
+        tasks = []
+        for t in result.get("tasks", []):
+            tasks.append(TaskInfo(
+                task_id=t.get("task_id", t.get("taskId", t.get("id", ""))),
+                project_id=t.get("project_id", t.get("projectId", "")),
+                title=t.get("title", ""),
+                description=t.get("description", ""),
+                priority=t.get("priority", "medium"),
+                working_directory=t.get("working_directory", t.get("workingDirectory")),
+                context=t.get("context"),
+                handoff=t.get("handoff")
+            ))
+        return tasks
+
+    async def report_execution_start(
+        self, task_id: str
+    ) -> ExecutionStartResult:
+        """Report that task execution has started.
+
+        Args:
+            task_id: Task ID being executed
+
+        Returns:
+            ExecutionStartResult with execution ID
+
+        Raises:
+            MCPError: If reporting fails or not authenticated
+        """
+        if not self._session_token:
+            raise MCPError("Not authenticated. Call authenticate() first.")
+
+        result = await self._call_tool("report_execution_start", {
+            "session_token": self._session_token,
+            "task_id": task_id
+        })
+
+        if not result.get("success"):
+            raise MCPError(result.get("error", "Failed to report execution start"))
+
+        started_at_str = result.get("started_at", datetime.now().isoformat())
+        if started_at_str.endswith("Z"):
+            started_at_str = started_at_str[:-1] + "+00:00"
+
+        return ExecutionStartResult(
+            execution_id=result.get("execution_log_id", result.get("execution_id", "")),
+            started_at=datetime.fromisoformat(started_at_str)
+        )
+
+    async def report_execution_complete(
+        self,
+        execution_id: str,
+        exit_code: int,
+        duration_seconds: float,
+        log_file_path: Optional[str] = None,
+        error_message: Optional[str] = None
+    ) -> None:
+        """Report that task execution has completed.
+
+        Args:
+            execution_id: Execution log ID from report_execution_start
+            exit_code: Exit code of the CLI process
+            duration_seconds: Duration of execution in seconds
+            log_file_path: Path to log file (optional)
+            error_message: Error message if execution failed (optional)
+
+        Raises:
+            MCPError: If reporting fails or not authenticated
+        """
+        if not self._session_token:
+            raise MCPError("Not authenticated. Call authenticate() first.")
+
+        args = {
+            "session_token": self._session_token,
+            "execution_log_id": execution_id,
+            "exit_code": exit_code,
+            "duration_seconds": duration_seconds
+        }
+        if log_file_path:
+            args["log_file_path"] = log_file_path
+        if error_message:
+            args["error_message"] = error_message
+
+        result = await self._call_tool("report_execution_complete", args)
+
+        if not result.get("success"):
+            raise MCPError(result.get("error", "Failed to report execution complete"))
+
+    async def update_task_status(
+        self, task_id: str, status: str, reason: Optional[str] = None
+    ) -> None:
+        """Update task status.
+
+        Args:
+            task_id: Task ID to update
+            status: New status (todo, in_progress, done, etc.)
+            reason: Reason for status change (optional)
+
+        Raises:
+            MCPError: If update fails
+        """
+        args = {
+            "task_id": task_id,
+            "status": status
+        }
+        if reason:
+            args["reason"] = reason
+
+        result = await self._call_tool("update_task_status", args)
+
+        if not result.get("success"):
+            raise MCPError(result.get("error", "Failed to update task status"))
+
+    async def save_context(
+        self,
+        task_id: str,
+        progress: Optional[str] = None,
+        findings: Optional[str] = None,
+        blockers: Optional[str] = None,
+        next_steps: Optional[str] = None,
+        agent_id: Optional[str] = None
+    ) -> None:
+        """Save task context.
+
+        Args:
+            task_id: Task ID
+            progress: Current progress description
+            findings: Findings or discoveries
+            blockers: Current blockers
+            next_steps: Recommended next steps
+            agent_id: Agent ID (optional)
+
+        Raises:
+            MCPError: If save fails
+        """
+        args = {"task_id": task_id}
+        if progress:
+            args["progress"] = progress
+        if findings:
+            args["findings"] = findings
+        if blockers:
+            args["blockers"] = blockers
+        if next_steps:
+            args["next_steps"] = next_steps
+        if agent_id:
+            args["agent_id"] = agent_id
+
+        result = await self._call_tool("save_context", args)
+
+        if not result.get("success"):
+            raise MCPError(result.get("error", "Failed to save context"))