refactor: split into files

Signed-off-by: Chojan Shang <chojan.shang@vesoft.com>

Author: Chojan Shang

src/acp/agent/__init__.py

@@ -0,0 +1,3 @@
+from .connection import AgentSideConnection
+
+__all__ = ["AgentSideConnection"]

src/acp/agent/connection.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+from typing import Any
+
+from ..connection import Connection, MethodHandler
+from ..interfaces import Agent
+from ..meta import CLIENT_METHODS
+from ..schema import (
+    CreateTerminalRequest,
+    CreateTerminalResponse,
+    KillTerminalCommandRequest,
+    KillTerminalCommandResponse,
+    ReadTextFileRequest,
+    ReadTextFileResponse,
+    ReleaseTerminalRequest,
+    ReleaseTerminalResponse,
+    RequestPermissionRequest,
+    RequestPermissionResponse,
+    SessionNotification,
+    TerminalOutputRequest,
+    TerminalOutputResponse,
+    WaitForTerminalExitRequest,
+    WaitForTerminalExitResponse,
+    WriteTextFileRequest,
+    WriteTextFileResponse,
+)
+from ..terminal import TerminalHandle
+from ..utils import notify_model, request_model, request_optional_model
+from .handlers import dispatch_agent_method
+
+__all__ = ["AgentSideConnection"]
+
+_AGENT_CONNECTION_ERROR = "AgentSideConnection requires asyncio StreamWriter/StreamReader"
+
+
+class AgentSideConnection:
+    """Agent-side connection wrapper that dispatches JSON-RPC messages to a Client implementation."""
+
+    def __init__(
+        self,
+        to_agent: Callable[[AgentSideConnection], Agent],
+        input_stream: Any,
+        output_stream: Any,
+    ) -> None:
+        agent = to_agent(self)
+        handler = self._create_handler(agent)
+
+        if not isinstance(input_stream, asyncio.StreamWriter) or not isinstance(output_stream, asyncio.StreamReader):
+            raise TypeError(_AGENT_CONNECTION_ERROR)
+        self._conn = Connection(handler, input_stream, output_stream)
+
+    def _create_handler(self, agent: Agent) -> MethodHandler:
+        async def handler(method: str, params: Any | None, is_notification: bool) -> Any:
+            return await dispatch_agent_method(agent, method, params, is_notification)
+
+        return handler
+
+    async def sessionUpdate(self, params: SessionNotification) -> None:
+        await notify_model(self._conn, CLIENT_METHODS["session_update"], params)
+
+    async def requestPermission(self, params: RequestPermissionRequest) -> RequestPermissionResponse:
+        return await request_model(
+            self._conn,
+            CLIENT_METHODS["session_request_permission"],
+            params,
+            RequestPermissionResponse,
+        )
+
+    async def readTextFile(self, params: ReadTextFileRequest) -> ReadTextFileResponse:
+        return await request_model(
+            self._conn,
+            CLIENT_METHODS["fs_read_text_file"],
+            params,
+            ReadTextFileResponse,
+        )
+
+    async def writeTextFile(self, params: WriteTextFileRequest) -> WriteTextFileResponse | None:
+        return await request_optional_model(
+            self._conn,
+            CLIENT_METHODS["fs_write_text_file"],
+            params,
+            WriteTextFileResponse,
+        )
+
+    async def createTerminal(self, params: CreateTerminalRequest) -> TerminalHandle:
+        create_response = await request_model(
+            self._conn,
+            CLIENT_METHODS["terminal_create"],
+            params,
+            CreateTerminalResponse,
+        )
+        return TerminalHandle(create_response.terminalId, params.sessionId, self._conn)
+
+    async def terminalOutput(self, params: TerminalOutputRequest) -> TerminalOutputResponse:
+        return await request_model(
+            self._conn,
+            CLIENT_METHODS["terminal_output"],
+            params,
+            TerminalOutputResponse,
+        )
+
+    async def releaseTerminal(self, params: ReleaseTerminalRequest) -> ReleaseTerminalResponse | None:
+        return await request_optional_model(
+            self._conn,
+            CLIENT_METHODS["terminal_release"],
+            params,
+            ReleaseTerminalResponse,
+        )
+
+    async def waitForTerminalExit(self, params: WaitForTerminalExitRequest) -> WaitForTerminalExitResponse:
+        return await request_model(
+            self._conn,
+            CLIENT_METHODS["terminal_wait_for_exit"],
+            params,
+            WaitForTerminalExitResponse,
+        )
+
+    async def killTerminal(self, params: KillTerminalCommandRequest) -> KillTerminalCommandResponse | None:
+        return await request_optional_model(
+            self._conn,
+            CLIENT_METHODS["terminal_kill"],
+            params,
+            KillTerminalCommandResponse,
+        )
+
+    async def extMethod(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
+        return await self._conn.send_request(f"_{method}", params)
+
+    async def extNotification(self, method: str, params: dict[str, Any]) -> None:
+        await self._conn.send_notification(f"_{method}", params)

src/acp/agent/handlers.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from typing import Any
+
+from ..exceptions import RequestError
+from ..interfaces import Agent
+from ..meta import AGENT_METHODS
+from ..schema import (
+    AuthenticateRequest,
+    CancelNotification,
+    InitializeRequest,
+    LoadSessionRequest,
+    NewSessionRequest,
+    PromptRequest,
+    SetSessionModelRequest,
+    SetSessionModeRequest,
+)
+from ..utils import normalize_result
+
+__all__ = [
+    "NO_MATCH",
+    "dispatch_agent_method",
+]
+
+
+class _NoMatch:
+    """Sentinel returned by routing helpers when no handler matches."""
+
+
+NO_MATCH = _NoMatch()
+
+
+async def _handle_agent_init(agent: Agent, method: str, params: Any | None) -> Any:
+    if method == AGENT_METHODS["initialize"]:
+        request = InitializeRequest.model_validate(params)
+        return await agent.initialize(request)
+    if method == AGENT_METHODS["session_new"]:
+        request = NewSessionRequest.model_validate(params)
+        return await agent.newSession(request)
+    return NO_MATCH
+
+
+async def _handle_agent_session(agent: Agent, method: str, params: Any | None) -> Any:
+    if method == AGENT_METHODS["session_load"]:
+        if not hasattr(agent, "loadSession"):
+            raise RequestError.method_not_found(method)
+        request = LoadSessionRequest.model_validate(params)
+        result = await agent.loadSession(request)
+        return normalize_result(result)
+    if method == AGENT_METHODS["session_set_mode"]:
+        if not hasattr(agent, "setSessionMode"):
+            raise RequestError.method_not_found(method)
+        request = SetSessionModeRequest.model_validate(params)
+        result = await agent.setSessionMode(request)
+        return normalize_result(result)
+    if method == AGENT_METHODS["session_prompt"]:
+        request = PromptRequest.model_validate(params)
+        return await agent.prompt(request)
+    if method == AGENT_METHODS["session_set_model"]:
+        if not hasattr(agent, "setSessionModel"):
+            raise RequestError.method_not_found(method)
+        request = SetSessionModelRequest.model_validate(params)
+        result = await agent.setSessionModel(request)
+        return normalize_result(result)
+    if method == AGENT_METHODS["session_cancel"]:
+        request = CancelNotification.model_validate(params)
+        return await agent.cancel(request)
+    return NO_MATCH
+
+
+async def _handle_agent_auth(agent: Agent, method: str, params: Any | None) -> Any:
+    if method == AGENT_METHODS["authenticate"]:
+        if not hasattr(agent, "authenticate"):
+            raise RequestError.method_not_found(method)
+        request = AuthenticateRequest.model_validate(params)
+        result = await agent.authenticate(request)
+        return normalize_result(result)
+    return NO_MATCH
+
+
+async def _handle_agent_extensions(agent: Agent, method: str, params: Any | None, is_notification: bool) -> Any:
+    if isinstance(method, str) and method.startswith("_"):
+        ext_name = method[1:]
+        if is_notification:
+            if hasattr(agent, "extNotification"):
+                await agent.extNotification(ext_name, params or {})  # type: ignore[arg-type]
+                return None
+            return None
+        if hasattr(agent, "extMethod"):
+            return await agent.extMethod(ext_name, params or {})  # type: ignore[arg-type]
+        return NO_MATCH
+    return NO_MATCH
+
+
+async def dispatch_agent_method(agent: Agent, method: str, params: Any | None, is_notification: bool) -> Any:
+    """Dispatch agent-bound methods, mirroring the upstream ACP routing."""
+    for resolver in (_handle_agent_init, _handle_agent_session, _handle_agent_auth):
+        result = await resolver(agent, method, params)
+        if result is not NO_MATCH:
+            return result
+    extension_result = await _handle_agent_extensions(agent, method, params, is_notification)
+    if extension_result is not NO_MATCH:
+        return extension_result
+    raise RequestError.method_not_found(method)

src/acp/client/__init__.py

@@ -0,0 +1,3 @@
+from .connection import ClientSideConnection
+
+__all__ = ["ClientSideConnection"]

src/acp/client/connection.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+from typing import Any
+
+from ..connection import Connection, MethodHandler
+from ..interfaces import Agent, Client
+from ..meta import AGENT_METHODS
+from ..schema import (
+    AuthenticateRequest,
+    AuthenticateResponse,
+    CancelNotification,
+    InitializeRequest,
+    InitializeResponse,
+    LoadSessionRequest,
+    LoadSessionResponse,
+    NewSessionRequest,
+    NewSessionResponse,
+    PromptRequest,
+    PromptResponse,
+    SetSessionModelRequest,
+    SetSessionModelResponse,
+    SetSessionModeRequest,
+    SetSessionModeResponse,
+)
+from ..utils import (
+    notify_model,
+    request_model,
+    request_model_from_dict,
+)
+from .handlers import dispatch_client_method
+
+__all__ = ["ClientSideConnection"]
+
+_CLIENT_CONNECTION_ERROR = "ClientSideConnection requires asyncio StreamWriter/StreamReader"
+
+
+class ClientSideConnection:
+    """Client-side connection wrapper that dispatches JSON-RPC messages to an Agent implementation."""
+
+    def __init__(
+        self,
+        to_client: Callable[[Agent], Client],
+        input_stream: Any,
+        output_stream: Any,
+    ) -> None:
+        if not isinstance(input_stream, asyncio.StreamWriter) or not isinstance(output_stream, asyncio.StreamReader):
+            raise TypeError(_CLIENT_CONNECTION_ERROR)
+
+        client = to_client(self)  # type: ignore[arg-type]
+        handler = self._create_handler(client)
+        self._conn = Connection(handler, input_stream, output_stream)
+
+    def _create_handler(self, client: Client) -> MethodHandler:
+        async def handler(method: str, params: Any | None, is_notification: bool) -> Any:
+            return await dispatch_client_method(client, method, params, is_notification)
+
+        return handler
+
+    async def initialize(self, params: InitializeRequest) -> InitializeResponse:
+        return await request_model(
+            self._conn,
+            AGENT_METHODS["initialize"],
+            params,
+            InitializeResponse,
+        )
+
+    async def newSession(self, params: NewSessionRequest) -> NewSessionResponse:
+        return await request_model(
+            self._conn,
+            AGENT_METHODS["session_new"],
+            params,
+            NewSessionResponse,
+        )
+
+    async def loadSession(self, params: LoadSessionRequest) -> LoadSessionResponse:
+        return await request_model_from_dict(
+            self._conn,
+            AGENT_METHODS["session_load"],
+            params,
+            LoadSessionResponse,
+        )
+
+    async def setSessionMode(self, params: SetSessionModeRequest) -> SetSessionModeResponse:
+        return await request_model_from_dict(
+            self._conn,
+            AGENT_METHODS["session_set_mode"],
+            params,
+            SetSessionModeResponse,
+        )
+
+    async def setSessionModel(self, params: SetSessionModelRequest) -> SetSessionModelResponse:
+        return await request_model_from_dict(
+            self._conn,
+            AGENT_METHODS["session_set_model"],
+            params,
+            SetSessionModelResponse,
+        )
+
+    async def authenticate(self, params: AuthenticateRequest) -> AuthenticateResponse:
+        return await request_model_from_dict(
+            self._conn,
+            AGENT_METHODS["authenticate"],
+            params,
+            AuthenticateResponse,
+        )
+
+    async def prompt(self, params: PromptRequest) -> PromptResponse:
+        return await request_model(
+            self._conn,
+            AGENT_METHODS["session_prompt"],
+            params,
+            PromptResponse,
+        )
+
+    async def cancel(self, params: CancelNotification) -> None:
+        await notify_model(self._conn, AGENT_METHODS["session_cancel"], params)
+
+    async def extMethod(self, method: str, params: dict[str, Any]) -> dict[str, Any]:
+        return await self._conn.send_request(f"_{method}", params)
+
+    async def extNotification(self, method: str, params: dict[str, Any]) -> None:
+        await self._conn.send_notification(f"_{method}", params)