beads/integrations/beads-mcp/src/beads_mcp/bd_daemon_client.py

"""Client for interacting with bd daemon via RPC over Unix socket."""

import asyncio
import json
import os
import socket
from pathlib import Path
from typing import Any, Dict, List, Optional

from .bd_client import BdClientBase, BdError
from .models import (
    AddDependencyParams,
    BlockedIssue,
    CloseIssueParams,
    CreateIssueParams,
    InitParams,
    Issue,
    ListIssuesParams,
    ReadyWorkParams,
    ReopenIssueParams,
    ShowIssueParams,
    Stats,
    UpdateIssueParams,
)


class DaemonError(Exception):
    """Base exception for daemon client errors."""

    pass


class DaemonNotRunningError(DaemonError):
    """Raised when daemon is not running."""

    pass


class DaemonConnectionError(DaemonError):
    """Raised when connection to daemon fails."""

    pass


class BdDaemonClient(BdClientBase):
    """Client for calling bd daemon via RPC over Unix socket."""

    socket_path: str | None
    working_dir: str | None
    actor: str | None
    timeout: float

    def __init__(
        self,
        socket_path: str | None = None,
        working_dir: str | None = None,
        actor: str | None = None,
        timeout: float = 30.0,
    ):
        """Initialize daemon client.

        Args:
            socket_path: Path to daemon Unix socket (optional, will auto-discover if not provided)
            working_dir: Working directory for database discovery (optional)
            actor: Actor name for audit trail (optional)
            timeout: Socket timeout in seconds (default: 30.0)
        """
        self.socket_path = socket_path
        self.working_dir = working_dir or os.getcwd()
        self.actor = actor
        self.timeout = timeout

    async def _find_socket_path(self) -> str:
        """Find daemon socket path by searching for .beads directory.

        Returns:
            Path to bd.sock file

        Raises:
            DaemonNotRunningError: If no .beads directory or socket file found
        """
        if self.socket_path:
            return self.socket_path

        # Walk up from working_dir to find .beads/bd.sock
        current = Path(self.working_dir).resolve()
        while True:
            beads_dir = current / ".beads"
            if beads_dir.is_dir():
                sock_path = beads_dir / "bd.sock"
                if sock_path.exists():
                    return str(sock_path)
                # Found .beads but no socket - daemon not running
                raise DaemonNotRunningError(
                    f"Daemon socket not found at {sock_path}. Is the daemon running? Try: bd daemon"
                )

            # Move up one directory
            parent = current.parent
            if parent == current:
                # Reached filesystem root
                raise DaemonNotRunningError(
                    "No .beads directory found. Initialize with: bd init"
                )
            current = parent

    async def _send_request(self, operation: str, args: Dict[str, Any]) -> Dict[str, Any]:
        """Send RPC request to daemon and get response.

        Args:
            operation: RPC operation name (e.g., "create", "list")
            args: Operation-specific arguments

        Returns:
            Parsed response data

        Raises:
            DaemonNotRunningError: If daemon is not running
            DaemonConnectionError: If connection fails
            DaemonError: If request fails
        """
        sock_path = await self._find_socket_path()

        # Build request
        request = {
            "operation": operation,
            "args": args,
            "cwd": self.working_dir,
        }
        if self.actor:
            request["actor"] = self.actor

        # Connect to socket and send request
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_unix_connection(sock_path),
                timeout=self.timeout,
            )
        except FileNotFoundError:
            raise DaemonNotRunningError(
                f"Daemon socket not found: {sock_path}. Is the daemon running?"
            )
        except asyncio.TimeoutError:
            raise DaemonConnectionError(
                f"Timeout connecting to daemon at {sock_path}"
            )
        except Exception as e:
            raise DaemonConnectionError(
                f"Failed to connect to daemon at {sock_path}: {e}"
            )

        try:
            # Send request as newline-delimited JSON
            request_json = json.dumps(request) + "\n"
            writer.write(request_json.encode())
            await writer.drain()

            # Read response (also newline-delimited JSON)
            try:
                response_line = await asyncio.wait_for(
                    reader.readline(),
                    timeout=self.timeout,
                )
            except asyncio.TimeoutError:
                raise DaemonError(
                    f"Timeout waiting for response from daemon (operation: {operation})"
                )

            if not response_line:
                raise DaemonError("Daemon closed connection without responding")

            response = json.loads(response_line.decode())

            # Check for errors
            if not response.get("success"):
                error = response.get("error", "Unknown error")
                raise DaemonError(f"Daemon returned error: {error}")

            # Return data
            return response.get("data", {})

        finally:
            writer.close()
            await writer.wait_closed()

    async def ping(self) -> Dict[str, str]:
        """Ping daemon to check if it's running.

        Returns:
            Dict with "message" and "version" fields

        Raises:
            DaemonNotRunningError: If daemon is not running
        """
        data = await self._send_request("ping", {})
        return json.loads(data) if isinstance(data, str) else data

    async def init(self, params: Optional[InitParams] = None) -> str:
        """Initialize new beads database (not typically used via daemon).

        Args:
            params: Initialization parameters (optional)

        Returns:
            Success message

        Note:
            This command is typically run via CLI, not daemon
        """
        params = params or InitParams()
        args: Dict[str, Any] = {}
        if params.prefix:
            args["prefix"] = params.prefix
        result = await self._send_request("init", args)
        return str(result) if result else "Initialized"

    async def create(self, params: CreateIssueParams) -> Issue:
        """Create a new issue.

        Args:
            params: Issue creation parameters

        Returns:
            Created issue
        """
        args = {
            "title": params.title,
            "issue_type": params.issue_type or "task",
            "priority": params.priority if params.priority is not None else 2,
        }
        if params.id:
            args["id"] = params.id
        if params.description:
            args["description"] = params.description
        if params.design:
            args["design"] = params.design
        if params.acceptance:
            args["acceptance_criteria"] = params.acceptance
        if params.assignee:
            args["assignee"] = params.assignee
        if params.labels:
            args["labels"] = params.labels
        if params.deps:
            args["dependencies"] = params.deps

        data = await self._send_request("create", args)
        return Issue(**(json.loads(data) if isinstance(data, str) else data))

    async def update(self, params: UpdateIssueParams) -> Issue:
        """Update an existing issue.

        Args:
            params: Issue update parameters

        Returns:
            Updated issue
        """
        args: Dict[str, Any] = {"id": params.issue_id}
        if params.status:
            args["status"] = params.status
        if params.priority is not None:
            args["priority"] = params.priority
        if params.design is not None:
            args["design"] = params.design
        if params.acceptance_criteria is not None:
            args["acceptance_criteria"] = params.acceptance_criteria
        if params.notes is not None:
            args["notes"] = params.notes
        if params.assignee is not None:
            args["assignee"] = params.assignee
        if params.title is not None:
            args["title"] = params.title

        data = await self._send_request("update", args)
        return Issue(**(json.loads(data) if isinstance(data, str) else data))

    async def close(self, params: CloseIssueParams) -> List[Issue]:
        """Close an issue.

        Args:
            params: Close parameters

        Returns:
            List containing the closed issue
        """
        args = {"id": params.issue_id}
        if params.reason:
            args["reason"] = params.reason

        data = await self._send_request("close", args)
        issue = Issue(**(json.loads(data) if isinstance(data, str) else data))
        return [issue]

    async def reopen(self, params: ReopenIssueParams) -> List[Issue]:
        """Reopen one or more closed issues.

        Args:
            params: Reopen parameters with issue IDs

        Returns:
            List of reopened issues

        Note:
            Reopen operation may not be implemented in daemon RPC yet
        """
        # Note: reopen operation may not be in RPC protocol yet
        # This is a placeholder for when it's added
        raise NotImplementedError("Reopen operation not yet supported via daemon")

    async def list_issues(self, params: Optional[ListIssuesParams] = None) -> List[Issue]:
        """List issues with optional filters.

        Args:
            params: List filter parameters (optional)

        Returns:
            List of matching issues
        """
        params = params or ListIssuesParams()
        args: Dict[str, Any] = {}
        if params.status:
            args["status"] = params.status
        if params.priority is not None:
            args["priority"] = params.priority
        if params.issue_type:
            args["issue_type"] = params.issue_type
        if params.assignee:
            args["assignee"] = params.assignee
        if params.limit:
            args["limit"] = params.limit

        data = await self._send_request("list", args)
        issues_data = json.loads(data) if isinstance(data, str) else data
        return [Issue(**issue) for issue in issues_data]

    async def show(self, params: ShowIssueParams) -> Issue:
        """Show detailed issue information.

        Args:
            params: Show parameters with issue_id

        Returns:
            Issue details
        """
        args = {"id": params.issue_id}
        data = await self._send_request("show", args)
        return Issue(**(json.loads(data) if isinstance(data, str) else data))

    async def ready(self, params: Optional[ReadyWorkParams] = None) -> List[Issue]:
        """Get ready work (issues with no blockers).

        Args:
            params: Ready work filter parameters (optional)

        Returns:
            List of ready issues
        """
        params = params or ReadyWorkParams()
        args: Dict[str, Any] = {}
        if params.assignee:
            args["assignee"] = params.assignee
        if params.priority is not None:
            args["priority"] = params.priority
        if params.limit:
            args["limit"] = params.limit

        data = await self._send_request("ready", args)
        issues_data = json.loads(data) if isinstance(data, str) else data
        return [Issue(**issue) for issue in issues_data]

    async def stats(self) -> Stats:
        """Get repository statistics.

        Returns:
            Statistics object
        """
        data = await self._send_request("stats", {})
        stats_data = json.loads(data) if isinstance(data, str) else data
        return Stats(**stats_data)

    async def blocked(self) -> List[BlockedIssue]:
        """Get blocked issues.

        Returns:
            List of blocked issues with their blockers

        Note:
            This operation may not be implemented in daemon RPC yet
        """
        # Note: blocked operation may not be in RPC protocol yet
        # This is a placeholder for when it's added
        raise NotImplementedError("Blocked operation not yet supported via daemon")

    async def add_dependency(self, params: AddDependencyParams) -> None:
        """Add a dependency between issues.

        Args:
            params: Dependency parameters
        """
        args = {
            "from_id": params.from_id,
            "to_id": params.to_id,
            "dep_type": params.dep_type or "blocks",
        }
        await self._send_request("dep_add", args)

    async def is_daemon_running(self) -> bool:
        """Check if daemon is running.

        Returns:
            True if daemon is running and responsive
        """
        try:
            await self.ping()
            return True
        except (DaemonNotRunningError, DaemonConnectionError, DaemonError):
            return False