webber/webber-cli/webber_cli/client.py

"""
Webber API client.

Communicates with the Webber API backend for agent execution.
Supports permission modes for controlling agent tool access.
"""
import json
import httpx
from collections.abc import AsyncIterator
from dataclasses import dataclass
from enum import Enum
from typing import Any


class PermissionMode(str, Enum):
    """
    Permission modes controlling agent tool access.

    - default: All tools available (approval may be required)
    - plan: Read-only tools only
    - auto_accept: All tools, no approval prompts
    """
    default = "default"
    plan = "plan"
    auto_accept = "auto_accept"


@dataclass
class AgentResponse:
    """Response from agent execution."""
    response: str
    agent_type: str
    success: bool
    mode: PermissionMode = PermissionMode.default
    error: str | None = None


@dataclass
class AgentInfo:
    """Information about an available agent."""
    name: str
    description: str


class WebberClient:
    """
    Client for the Webber API.

    Usage:
        client = WebberClient("http://localhost:8086")
        response = await client.run_agent("explore", "find python files", "/path/to/project")
    """

    def __init__(
        self,
        base_url: str = "http://localhost:8086",
        api_key: str | None = None,
        timeout: float = 120.0,
    ):
        """
        Initialize the Webber client.

        Args:
            base_url: Webber API URL
            api_key: Optional API key for authentication
            timeout: Request timeout in seconds
        """
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.timeout = timeout
        self._client: httpx.AsyncClient | None = None

    async def _get_client(self) -> httpx.AsyncClient:
        """Get or create the HTTP client."""
        if self._client is None or self._client.is_closed:
            headers = {}
            if self.api_key:
                headers["X-API-Key"] = self.api_key
            self._client = httpx.AsyncClient(
                base_url=self.base_url,
                headers=headers,
                timeout=self.timeout,
            )
        return self._client

    async def close(self) -> None:
        """Close the HTTP client."""
        if self._client and not self._client.is_closed:
            await self._client.aclose()
            self._client = None

    async def health_check(self) -> bool:
        """Check if the API is healthy."""
        try:
            client = await self._get_client()
            response = await client.get("/health")
            return response.status_code == 200
        except httpx.RequestError:
            return False

    async def list_agents(self) -> list[AgentInfo]:
        """List available agents."""
        client = await self._get_client()
        response = await client.get("/agents/")
        response.raise_for_status()
        data = response.json()
        return [AgentInfo(**a) for a in data.get("agents", [])]

    async def get_agent(self, agent_type: str) -> AgentInfo | None:
        """Get information about a specific agent."""
        client = await self._get_client()
        response = await client.get(f"/agents/{agent_type}")
        if response.status_code == 404:
            return None
        response.raise_for_status()
        return AgentInfo(**response.json())

    async def run_agent(
        self,
        agent_type: str,
        prompt: str,
        working_dir: str = ".",
        mode: PermissionMode = PermissionMode.default,
    ) -> AgentResponse:
        """
        Run an agent with the given prompt.

        Args:
            agent_type: Type of agent (e.g., "task")
            prompt: User prompt/query
            working_dir: Working directory for the agent
            mode: Permission mode controlling tool access

        Returns:
            AgentResponse with the result
        """
        client = await self._get_client()
        response = await client.post(
            "/agents/run",
            json={
                "agent_type": agent_type,
                "prompt": prompt,
                "working_dir": working_dir,
                "mode": mode.value,
            },
        )
        response.raise_for_status()
        data = response.json()
        return AgentResponse(
            response=data.get("response", ""),
            agent_type=data.get("agent_type", agent_type),
            success=data.get("success", True),
            mode=PermissionMode(data.get("mode", "default")),
            error=data.get("error"),
        )

    async def run_agent_stream(
        self,
        agent_type: str,
        prompt: str,
        working_dir: str = ".",
        mode: PermissionMode = PermissionMode.default,
    ) -> AsyncIterator[str]:
        """
        Run an agent with streaming response.

        Args:
            agent_type: Type of agent (e.g., "task")
            prompt: User prompt/query
            working_dir: Working directory for the agent
            mode: Permission mode controlling tool access

        Yields:
            Text chunks as they arrive
        """
        # Use a fresh client for streaming with longer timeout
        async with httpx.AsyncClient(
            base_url=self.base_url,
            timeout=httpx.Timeout(300.0, connect=10.0),
        ) as client:
            async with client.stream(
                "POST",
                "/agents/stream",
                json={
                    "agent_type": agent_type,
                    "prompt": prompt,
                    "working_dir": working_dir,
                    "mode": mode.value,
                },
            ) as response:
                response.raise_for_status()
                async for line in response.aiter_lines():
                    if line.startswith("data: "):
                        try:
                            data = json.loads(line[6:])
                            event = data.get("event")
                            if event == "chunk":
                                yield data.get("data", "")
                            elif event == "error":
                                raise Exception(data.get("data", "Unknown error"))
                            elif event == "done":
                                break
                        except json.JSONDecodeError:
                            continue

    async def __aenter__(self) -> "WebberClient":
        """Async context manager entry."""
        return self

    async def __aexit__(self, *args: Any) -> None:
        """Async context manager exit."""
        await self.close()