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, field
from datetime import datetime
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"


class StreamEventType(str, Enum):
    """Event types for structured agent streaming."""
    tool_start = "tool_start"
    tool_done = "tool_done"
    thinking = "thinking"
    response = "response"
    error = "error"
    done = "done"
    chunk = "chunk"  # Legacy text chunk


@dataclass
class StreamEvent:
    """
    Structured streaming event from agent execution.

    Different event types carry different data:
    - tool_start: tool, args
    - tool_done: tool, result_summary
    - thinking: message
    - response: text
    - error: error_message
    - done: mode
    - chunk: text (legacy)
    """
    event: StreamEventType
    tool: str | None = None
    args: dict | None = None
    result_summary: str | None = None
    message: str | None = None
    text: str | None = None
    error_message: str | None = None
    mode: str | None = None


@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


@dataclass
class Message:
    """A message in a conversation."""
    id: str
    role: str
    content: str
    token_count: int
    is_summary: bool
    created_at: datetime


@dataclass
class Conversation:
    """A conversation session."""
    id: str
    agent_type: str
    title: str | None
    working_dir: str
    total_tokens: int
    created_at: datetime
    updated_at: datetime | None
    messages: list[Message] = field(default_factory=list)


@dataclass
class AddMessageResult:
    """Result of adding a message to a conversation."""
    user_message: Message
    assistant_message: Message
    total_tokens: int
    summarized: bool


@dataclass
class SaveMessagesResult:
    """Result of saving a message pair without agent execution."""
    user_message: Message
    assistant_message: Message
    total_tokens: int


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[StreamEvent]:
        """
        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:
            StreamEvent objects 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_type = data.get("event")

                            # Parse event type
                            try:
                                evt_type = StreamEventType(event_type)
                            except ValueError:
                                continue  # Unknown event type

                            # Build StreamEvent from response data
                            yield StreamEvent(
                                event=evt_type,
                                tool=data.get("tool"),
                                args=data.get("args"),
                                result_summary=data.get("result_summary"),
                                message=data.get("message"),
                                text=data.get("text") or data.get("data"),  # 'data' for legacy chunk
                                error_message=data.get("error_message") or data.get("data"),
                                mode=data.get("mode"),
                            )

                            # Stop on done or error
                            if evt_type in (StreamEventType.done, StreamEventType.error):
                                break

                        except json.JSONDecodeError:
                            continue

    # === Conversation API ===

    def _parse_message(self, data: dict) -> Message:
        """Parse a Message from API response data."""
        return Message(
            id=data["id"],
            role=data["role"],
            content=data["content"],
            token_count=data["token_count"],
            is_summary=data["is_summary"],
            created_at=datetime.fromisoformat(data["created_at"].replace("Z", "+00:00")),
        )

    def _parse_conversation(self, data: dict, with_messages: bool = False) -> Conversation:
        """Parse a Conversation from API response data."""
        messages = []
        if with_messages and "messages" in data:
            messages = [self._parse_message(m) for m in data["messages"]]

        updated_at = None
        if data.get("updated_at"):
            updated_at = datetime.fromisoformat(data["updated_at"].replace("Z", "+00:00"))

        return Conversation(
            id=data["id"],
            agent_type=data["agent_type"],
            title=data.get("title"),
            working_dir=data["working_dir"],
            total_tokens=data["total_tokens"],
            created_at=datetime.fromisoformat(data["created_at"].replace("Z", "+00:00")),
            updated_at=updated_at,
            messages=messages,
        )

    async def list_conversations(
        self,
        limit: int = 50,
        offset: int = 0,
    ) -> tuple[list[Conversation], int]:
        """
        List user's conversations.

        Args:
            limit: Maximum number of conversations to return
            offset: Offset for pagination

        Returns:
            Tuple of (conversations, total_count)
        """
        client = await self._get_client()
        response = await client.get(
            "/conversations/",
            params={"limit": limit, "offset": offset},
        )
        response.raise_for_status()
        data = response.json()
        conversations = [self._parse_conversation(c) for c in data["conversations"]]
        return conversations, data["total"]

    async def get_conversation(self, conversation_id: str) -> Conversation | None:
        """
        Get a conversation with all messages.

        Args:
            conversation_id: UUID of the conversation

        Returns:
            Conversation with messages, or None if not found
        """
        client = await self._get_client()
        response = await client.get(f"/conversations/{conversation_id}")
        if response.status_code == 404:
            return None
        response.raise_for_status()
        return self._parse_conversation(response.json(), with_messages=True)

    async def create_conversation(
        self,
        agent_type: str = "task",
        working_dir: str = ".",
        title: str | None = None,
    ) -> Conversation:
        """
        Create a new conversation.

        Args:
            agent_type: Type of agent to use
            working_dir: Working directory for the agent
            title: Optional title for the conversation

        Returns:
            The created conversation
        """
        client = await self._get_client()
        response = await client.post(
            "/conversations/",
            json={
                "agent_type": agent_type,
                "working_dir": working_dir,
                "title": title,
            },
        )
        response.raise_for_status()
        return self._parse_conversation(response.json())

    async def add_message(
        self,
        conversation_id: str,
        content: str,
    ) -> AddMessageResult:
        """
        Add a message to a conversation and get agent response.

        Args:
            conversation_id: UUID of the conversation
            content: Message content

        Returns:
            AddMessageResult with user and assistant messages
        """
        client = await self._get_client()
        response = await client.post(
            f"/conversations/{conversation_id}/messages",
            json={"content": content},
        )
        response.raise_for_status()
        data = response.json()
        return AddMessageResult(
            user_message=self._parse_message(data["user_message"]),
            assistant_message=self._parse_message(data["assistant_message"]),
            total_tokens=data["total_tokens"],
            summarized=data.get("summarized", False),
        )

    async def save_messages(
        self,
        conversation_id: str,
        user_content: str,
        assistant_content: str,
    ) -> SaveMessagesResult:
        """
        Save a user/assistant message pair without triggering agent execution.

        Used when streaming responses separately via run_agent_stream().
        Allows persisting the exchange after streaming completes.

        Args:
            conversation_id: UUID of the conversation
            user_content: User message content
            assistant_content: Assistant response content

        Returns:
            SaveMessagesResult with both messages
        """
        client = await self._get_client()
        response = await client.post(
            f"/conversations/{conversation_id}/save",
            json={
                "user_content": user_content,
                "assistant_content": assistant_content,
            },
        )
        response.raise_for_status()
        data = response.json()
        return SaveMessagesResult(
            user_message=self._parse_message(data["user_message"]),
            assistant_message=self._parse_message(data["assistant_message"]),
            total_tokens=data["total_tokens"],
        )

    async def delete_conversation(self, conversation_id: str) -> bool:
        """
        Delete a conversation.

        Args:
            conversation_id: UUID of the conversation

        Returns:
            True if deleted, False if not found
        """
        client = await self._get_client()
        response = await client.delete(f"/conversations/{conversation_id}")
        if response.status_code == 404:
            return False
        response.raise_for_status()
        return True

    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()