agentpool 2.1.9__py3-none-any.whl

acp/client/connection.py

@@ -0,0 +1,251 @@
+"""Client ACP Connection."""
+
+from __future__ import annotations
+
+from functools import partial
+from typing import TYPE_CHECKING, Any, Self
+
+import logfire
+
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+    from anyio.abc import ByteReceiveStream, ByteSendStream
+
+    from acp.client.protocol import Client
+    from acp.connection import StreamObserver
+    from acp.schema import (
+        AuthenticateRequest,
+        CancelNotification,
+        ClientMethod,
+        CreateTerminalResponse,
+        ForkSessionRequest,
+        InitializeRequest,
+        KillTerminalCommandResponse,
+        ListSessionsRequest,
+        LoadSessionRequest,
+        NewSessionRequest,
+        PromptRequest,
+        ReadTextFileResponse,
+        ReleaseTerminalResponse,
+        RequestPermissionResponse,
+        ResumeSessionRequest,
+        SetSessionModelRequest,
+        SetSessionModeRequest,
+        TerminalOutputResponse,
+        WaitForTerminalExitResponse,
+        WriteTextFileResponse,
+    )
+
+from acp.agent.protocol import Agent
+from acp.connection import Connection
+from acp.exceptions import RequestError
+from acp.schema import (
+    AuthenticateResponse,
+    CreateTerminalRequest,
+    ForkSessionResponse,
+    InitializeResponse,
+    KillTerminalCommandRequest,
+    ListSessionsResponse,
+    LoadSessionResponse,
+    NewSessionResponse,
+    PromptResponse,
+    ReadTextFileRequest,
+    ReleaseTerminalRequest,
+    RequestPermissionRequest,
+    ResumeSessionResponse,
+    SessionNotification,
+    SetSessionModelResponse,
+    SetSessionModeResponse,
+    TerminalOutputRequest,
+    WaitForTerminalExitRequest,
+    WriteTextFileRequest,
+)
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+
+class ClientSideConnection(Agent):
+    """Client-side connection.
+
+    Use when you implement the Client and need to talk to an Agent.
+
+    Args:
+      to_client: factory that receives this connection and returns your Client
+      input: ByteSendStream (local -> peer)
+      output: ByteReceiveStream (peer -> local)
+    """
+
+    def __init__(
+        self,
+        to_client: Callable[[Agent], Client],
+        input_stream: ByteSendStream,
+        output_stream: ByteReceiveStream,
+        observers: list[StreamObserver] | None = None,
+    ) -> None:
+        # Build client first so handler can delegate
+        client = to_client(self)
+        handler = partial(_handle_client_method, client)
+        self._conn = Connection(handler, input_stream, output_stream, observers=observers)
+
+    # agent-bound methods (client -> agent)
+    async def initialize(self, params: InitializeRequest) -> InitializeResponse:
+        dct = params.model_dump(
+            mode="json", by_alias=True, exclude_none=True, exclude_defaults=True
+        )
+        resp = await self._conn.send_request("initialize", dct)
+        return InitializeResponse.model_validate(resp)
+
+    async def new_session(self, params: NewSessionRequest) -> NewSessionResponse:
+        # Don't exclude_defaults - MCP server 'type' field is required for discriminated unions
+        dct = params.model_dump(mode="json", by_alias=True, exclude_none=True)
+        resp = await self._conn.send_request("session/new", dct)
+        return NewSessionResponse.model_validate(resp)
+
+    async def load_session(self, params: LoadSessionRequest) -> LoadSessionResponse:
+        dct = params.model_dump(
+            mode="json", by_alias=True, exclude_none=True, exclude_defaults=True
+        )
+        resp = await self._conn.send_request("session/load", dct)
+        payload = resp if isinstance(resp, dict) else {}
+        return LoadSessionResponse.model_validate(payload)
+
+    async def list_sessions(self, params: ListSessionsRequest) -> ListSessionsResponse:
+        dct = params.model_dump(
+            mode="json", by_alias=True, exclude_none=True, exclude_defaults=True
+        )
+        resp = await self._conn.send_request("session/list", dct)
+        return ListSessionsResponse.model_validate(resp)
+
+    async def fork_session(self, params: ForkSessionRequest) -> ForkSessionResponse:
+        dct = params.model_dump(
+            mode="json", by_alias=True, exclude_none=True, exclude_defaults=True
+        )
+        resp = await self._conn.send_request("session/fork", dct)
+        return ForkSessionResponse.model_validate(resp)
+
+    async def resume_session(self, params: ResumeSessionRequest) -> ResumeSessionResponse:
+        dct = params.model_dump(
+            mode="json", by_alias=True, exclude_none=True, exclude_defaults=True
+        )
+        resp = await self._conn.send_request("session/resume", dct)
+        payload = resp if isinstance(resp, dict) else {}
+        return ResumeSessionResponse.model_validate(payload)
+
+    async def set_session_mode(self, params: SetSessionModeRequest) -> SetSessionModeResponse:
+        dct = params.model_dump(
+            mode="json", by_alias=True, exclude_none=True, exclude_defaults=True
+        )
+        resp = await self._conn.send_request("session/set_mode", dct)
+        payload = resp if isinstance(resp, dict) else {}
+        return SetSessionModeResponse.model_validate(payload)
+
+    async def set_session_model(self, params: SetSessionModelRequest) -> SetSessionModelResponse:
+        dct = params.model_dump(
+            mode="json", by_alias=True, exclude_none=True, exclude_defaults=True
+        )
+        resp = await self._conn.send_request("session/set_model", dct)
+        payload = resp if isinstance(resp, dict) else {}
+        return SetSessionModelResponse.model_validate(payload)
+
+    async def authenticate(self, params: AuthenticateRequest) -> AuthenticateResponse:
+        dct = params.model_dump(
+            mode="json", by_alias=True, exclude_none=True, exclude_defaults=True
+        )
+        resp = await self._conn.send_request("authenticate", dct)
+        payload = resp if isinstance(resp, dict) else {}
+        return AuthenticateResponse.model_validate(payload)
+
+    async def prompt(self, params: PromptRequest) -> PromptResponse:
+        # Don't exclude_defaults here - the 'type' field in content blocks has a default
+        # value but is required for discriminated unions to work
+        dct = params.model_dump(mode="json", by_alias=True, exclude_none=True)
+        resp = await self._conn.send_request("session/prompt", dct)
+        return PromptResponse.model_validate(resp)
+
+    async def cancel(self, params: CancelNotification) -> None:
+        dct = params.model_dump(
+            mode="json", by_alias=True, exclude_none=True, exclude_defaults=True
+        )
+        await self._conn.send_notification("session/cancel", dct)
+
+    async def ext_method(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
+        return await self._conn.send_request(f"_{method}", params)  # type: ignore[no-any-return]
+
+    async def ext_notification(self, method: str, params: dict[str, Any]) -> None:
+        await self._conn.send_notification(f"_{method}", params)
+
+    async def close(self) -> None:
+        await self._conn.close()
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self.close()
+
+
+@logfire.instrument(r"Handle Client Method {method}")
+async def _handle_client_method(  # noqa: PLR0911
+    client: Client,
+    method: ClientMethod | str,
+    params: dict[str, Any] | None,
+    is_notification: bool,
+) -> (
+    WriteTextFileResponse
+    | ReadTextFileResponse
+    | RequestPermissionResponse
+    | SessionNotification
+    | CreateTerminalResponse
+    | TerminalOutputResponse
+    | WaitForTerminalExitResponse
+    | ReleaseTerminalResponse
+    | KillTerminalCommandResponse
+    | dict[str, Any]
+    | None
+):
+    """Handle client method calls."""
+    match method:
+        case "fs/write_text_file":
+            write_file_request = WriteTextFileRequest.model_validate(params)
+            return await client.write_text_file(write_file_request)
+        case "fs/read_text_file":
+            read_file_request = ReadTextFileRequest.model_validate(params)
+            return await client.read_text_file(read_file_request)
+        case "session/request_permission":
+            permission_request = RequestPermissionRequest.model_validate(params)
+            return await client.request_permission(permission_request)
+        case "session/update":
+            notification = SessionNotification.model_validate(params)
+            await client.session_update(notification)
+            return None
+        case "terminal/create":
+            create_request = CreateTerminalRequest.model_validate(params)
+            return await client.create_terminal(create_request)
+        case "terminal/output":
+            output_request = TerminalOutputRequest.model_validate(params)
+            return await client.terminal_output(output_request)
+        case "terminal/release":
+            release_request = ReleaseTerminalRequest.model_validate(params)
+            return await client.release_terminal(release_request)
+        case "terminal/wait_for_exit":
+            wait_request = WaitForTerminalExitRequest.model_validate(params)
+            return await client.wait_for_terminal_exit(wait_request)
+        case "terminal/kill":
+            kill_request = KillTerminalCommandRequest.model_validate(params)
+            return await client.kill_terminal(kill_request)
+        case str() if method.startswith("_") and is_notification:
+            await client.ext_notification(method[1:], params or {})
+            return None
+        case str() if method.startswith("_") and is_notification:
+            return await client.ext_method(method[1:], params or {})
+        case _:
+            raise RequestError.method_not_found(method)

acp/client/implementations/__init__.py

@@ -0,0 +1,7 @@
+"""ACP Client implementations."""
+
+from acp.client.implementations.default_client import DefaultACPClient
+from acp.client.implementations.headless_client import HeadlessACPClient
+from acp.client.implementations.noop_client import NoOpClient
+
+__all__ = ["DefaultACPClient", "HeadlessACPClient", "NoOpClient"]

acp/client/implementations/default_client.py

@@ -0,0 +1,185 @@
+"""Default ACP client implementation.
+
+This module provides a basic client implementation for the Agent Client Protocol (ACP)
+that can be used for testing or as a base for more sophisticated client implementations.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from acp.client import Client
+from acp.schema import ReadTextFileResponse, RequestPermissionResponse, WriteTextFileResponse
+from agentpool import log
+
+
+if TYPE_CHECKING:
+    from acp.schema import (
+        AllowedOutcome,
+        CreateTerminalRequest,
+        CreateTerminalResponse,
+        DeniedOutcome,
+        KillTerminalCommandRequest,
+        KillTerminalCommandResponse,
+        ReadTextFileRequest,
+        ReleaseTerminalRequest,
+        ReleaseTerminalResponse,
+        RequestPermissionRequest,
+        SessionNotification,
+        TerminalOutputRequest,
+        TerminalOutputResponse,
+        WaitForTerminalExitRequest,
+        WaitForTerminalExitResponse,
+        WriteTextFileRequest,
+    )
+
+logger = log.get_logger(__name__)
+
+
+class DefaultACPClient(Client):
+    """Default implementation of ACP Client interface for basic operations.
+
+    This provides a basic client implementation that can be used for testing
+    or as a base for more sophisticated client implementations.
+    """
+
+    def __init__(
+        self,
+        *,
+        allow_file_operations: bool = False,
+        permission_outcomes: list[dict[str, DeniedOutcome | AllowedOutcome]] | None = None,
+        use_real_files: bool = True,
+    ) -> None:
+        """Initialize default ACP client.
+
+        Args:
+            allow_file_operations: Whether to allow file read/write operations
+            permission_outcomes: Queue of permission outcomes for testing
+            use_real_files: Whether to use real filesystem or in-memory storage
+        """
+        self.allow_file_operations = allow_file_operations
+        self.use_real_files = use_real_files
+        self.permission_outcomes = permission_outcomes or []
+        self.files: dict[str, str] = {}  # In-memory file storage for testing
+        self.ext_calls: list[tuple[str, dict[str, Any]]] = []
+        self.ext_notes: list[tuple[str, dict[str, Any]]] = []
+        self.notifications: list[SessionNotification] = []
+
+    async def request_permission(
+        self, params: RequestPermissionRequest
+    ) -> RequestPermissionResponse:
+        """Default permission handler - grants all permissions or uses test queue."""
+        title = params.tool_call.title or "operation"
+        logger.info("Permission requested", title=title)
+
+        # If we have test outcomes queued, use them
+        if self.permission_outcomes:
+            outcome = self.permission_outcomes.pop(0)
+            return RequestPermissionResponse.model_validate({"outcome": outcome})
+
+        # Default: grant permission for the first option
+        if params.options:
+            id_ = params.options[0].option_id
+            return RequestPermissionResponse.allowed(id_)
+
+        # No options - deny
+        return RequestPermissionResponse.denied()
+
+    async def session_update(self, params: SessionNotification) -> None:
+        """Handle session update notifications."""
+        msg = "Session update for %s: %s"
+        logger.debug(msg, params.session_id, params.update.session_update)
+        self.notifications.append(params)
+
+    async def write_text_file(self, params: WriteTextFileRequest) -> WriteTextFileResponse:
+        """Write text to file (if allowed)."""
+        if not self.allow_file_operations:
+            raise RuntimeError("File operations not allowed")
+
+        if self.use_real_files:
+            try:
+                path = Path(params.path)
+                path.write_text(params.content, encoding="utf-8")
+                logger.info("Wrote file", path=params.path)
+            except Exception:
+                logger.exception("Failed to write file", path=params.path)
+                raise
+        else:
+            # In-memory storage for testing
+            self.files[str(params.path)] = params.content
+
+        return WriteTextFileResponse()
+
+    async def read_text_file(self, params: ReadTextFileRequest) -> ReadTextFileResponse:
+        """Read text from file (if allowed)."""
+        if not self.allow_file_operations:
+            raise RuntimeError("File operations not allowed")
+
+        if self.use_real_files:
+            try:
+                path = Path(params.path)
+
+                if not path.exists():
+                    raise FileNotFoundError(f"File not found: {params.path}")  # noqa: TRY301
+
+                content = path.read_text(encoding="utf-8")
+
+                # Apply line filtering if requested
+                if params.line is not None or params.limit is not None:
+                    lines = content.splitlines()
+                    start_line = (params.line - 1) if params.line else 0
+                    end_line = start_line + params.limit if params.limit else len(lines)
+                    content = "\n".join(lines[start_line:end_line])
+
+                logger.info("Read file", path=params.path)
+                return ReadTextFileResponse(content=content)
+
+            except Exception:
+                logger.exception("Failed to read file", path=params.path)
+                raise
+        else:
+            # In-memory storage for testing
+            content = self.files.get(str(params.path), "default content")
+            return ReadTextFileResponse(content=content)
+
+    async def create_terminal(self, params: CreateTerminalRequest) -> CreateTerminalResponse:
+        """Create terminal (not implemented)."""
+        raise NotImplementedError("Terminal operations not implemented")
+
+    async def terminal_output(self, params: TerminalOutputRequest) -> TerminalOutputResponse:
+        """Get terminal output (not implemented)."""
+        raise NotImplementedError("Terminal operations not implemented")
+
+    async def release_terminal(
+        self, params: ReleaseTerminalRequest
+    ) -> ReleaseTerminalResponse | None:
+        """Release terminal (not implemented)."""
+        raise NotImplementedError("Terminal operations not implemented")
+
+    async def wait_for_terminal_exit(
+        self, params: WaitForTerminalExitRequest
+    ) -> WaitForTerminalExitResponse:
+        """Wait for terminal exit (not implemented)."""
+        raise NotImplementedError("Terminal operations not implemented")
+
+    async def kill_terminal(
+        self, params: KillTerminalCommandRequest
+    ) -> KillTerminalCommandResponse | None:
+        """Kill terminal (not implemented)."""
+        raise NotImplementedError("Terminal operations not implemented")
+
+    def get_session_updates(self) -> list[SessionNotification]:
+        """Get all received session updates."""
+        return self.notifications.copy()
+
+    def clear_session_updates(self) -> None:
+        """Clear all stored session updates."""
+        self.notifications.clear()
+
+    async def ext_method(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
+        self.ext_calls.append((method, params))
+        return {"ok": True, "method": method}
+
+    async def ext_notification(self, method: str, params: dict[str, Any]) -> None:
+        self.ext_notes.append((method, params))