aicc-py/aicc/client.py

"""AICCClient: agent side of the AICC protocol."""

from __future__ import annotations

import asyncio
import uuid
from typing import Any, Self

from aicc.errors import (
    AICCError,
    ConnectionError_,
    ProtocolError,
    ToolError,
)
from aicc.protocol import (
    PROTOCOL_VERSION,
    EventMessage,
    MessageType,
    SessionClose,
    SessionInit,
    ToolCall,
    ToolResult,
)
from aicc.transport.base import Transport


class AICCClient:
    """Agent-side AICC client.

    Wraps a Transport and exposes high-level operations:
    - handshake() -> manifest
    - call_tool(name, input) -> ToolResult
    - next_event() -> EventMessage
    - close

    A single background reader task consumes the transport and routes
    messages into an internal queue, so any number of concurrent callers
    (tool calls, event consumers) are safe on any transport, including
    WebSocket where concurrent recv() is forbidden.

    Example:
        client = AICCClient(WebSocketClientTransport("ws://localhost:8765"))
        async with client:
            manifest = await client.handshake()
            result = await client.call_tool("vision", {})
    """

    def __init__(self, transport: Transport, protocol: str = PROTOCOL_VERSION):
        self._transport = transport
        self._protocol = protocol
        self._session_id: str | None = None
        self._manifest: SessionInit | None = None
        self._inbox: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
        self._reader_task: asyncio.Task | None = None

    # ---------- Lifecycle ----------

    async def __aenter__(self) -> Self:
        await self._transport.connect()
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self.close()

    async def connect(self) -> AICCClient:
        await self._transport.connect()
        return self

    async def close(self) -> None:
        if self._reader_task is not None:
            self._reader_task.cancel()
            try:
                await self._reader_task
            except (asyncio.CancelledError, AICCError):
                pass
            self._reader_task = None
        if self._session_id is not None:
            try:
                await self._transport.send(
                    SessionClose(
                        session_id=self._session_id,
                        reason="client_close",
                    ).model_dump(by_alias=True)
                )
            except AICCError:
                pass
            self._session_id = None
        await self._transport.close()

    # ---------- Manifest ----------

    async def handshake(self) -> SessionInit:
        """Wait for the bridge to push the initial session_init manifest."""
        self._spawn_reader()
        init = await self._receive_until(SessionInit)
        self._session_id = init.session_id
        self._manifest = init
        return init

    async def fetch_manifest(self) -> SessionInit:
        if self._session_id is None:
            raise ProtocolError("handshake() must be called before fetch_manifest().")
        await self._transport.send(
            {
                "type": MessageType.MANIFEST_REQUEST.value,
                "session_id": self._session_id,
                "message_id": str(uuid.uuid4()),
            }
        )
        return await self._receive_until(SessionInit)

    # ---------- Tool calls ----------

    async def call_tool(self, tool_id: str, input: dict[str, Any]) -> ToolResult:
        if self._session_id is None:
            raise ProtocolError("handshake() must be called before call_tool().")
        call_id = f"tc_{uuid.uuid4().hex[:12]}"
        await self._transport.send(
            ToolCall(
                session_id=self._session_id,
                call_id=call_id,
                tool=tool_id,
                input=input,
            ).model_dump(by_alias=True)
        )

        while True:
            msg = await self._inbox.get()
            self._validate_protocol(msg)
            t = msg.get("type")
            if t == MessageType.ERROR.value:
                raise ProtocolError(f"bridge error: {msg.get('error')}")
            if t != MessageType.TOOL_RESULT.value:
                continue
            result = ToolResult.model_validate(msg)
            if result.call_id != call_id:
                continue
            if not result.ok and result.error:
                raise ToolError(
                    code=result.error.code,
                    message=result.error.message,
                    retryable=result.error.retryable,
                )
            return result

    # ---------- Events ----------

    async def next_event(self, timeout: float | None = None) -> EventMessage:
        if timeout is None:
            raw = await self._inbox.get()
        else:
            raw = await asyncio.wait_for(self._inbox.get(), timeout=timeout)
        return EventMessage.model_validate(raw)

    def _spawn_reader(self) -> None:
        async def read_loop() -> None:
            try:
                while True:
                    msg = await self._transport.receive()
                    await self._inbox.put(msg)
            except (ConnectionError_, asyncio.CancelledError):
                return

        self._reader_task = asyncio.create_task(read_loop())

    # ---------- Internals ----------

    async def _receive_until(self, expect_cls: type) -> SessionInit:
        while True:
            msg = await self._inbox.get()
            self._validate_protocol(msg)
            t = msg.get("type")
            if t == MessageType.ERROR.value:
                raise ProtocolError(f"bridge error: {msg.get('error')}")
            if t == expect_cls.model_fields["type"].default:
                return expect_cls.model_validate(msg)

    def _validate_protocol(self, msg: dict[str, Any]) -> None:
        proto = msg.get("protocol")
        if proto != self._protocol:
            raise ProtocolError(
                f"Protocol mismatch: client='{self._protocol}', message='{proto}'"
            )