refactor(protocol): streamline connection attachment

Author: PsiACE

src/acp/agent/connection.py

@@ -102,13 +102,14 @@ def __init__(
         self._notify_connected(agent)
 
     @classmethod
-    def _attach(
+    def attach(
         cls,
         to_agent: Callable[[Client], Agent] | Agent,
         connection: Connection,
         *,
         use_unstable_protocol: bool = False,
     ) -> tuple[AgentSideConnection, MethodHandler]:
+        """Attach an agent-side wrapper to an existing connection."""
         self = cls.__new__(cls)
         agent, handler = self._prepare(to_agent, use_unstable_protocol=use_unstable_protocol)
         self._conn = connection

src/acp/client/connection.py

@@ -6,7 +6,7 @@
 from typing import Any, cast, final
 
 from .._transport import Transport
-from ..connection import Connection, MethodHandler
+from ..connection import Connection
 from ..exceptions import RequestError
 from ..interfaces import Agent, Client
 from ..meta import AGENT_METHODS, CLIENT_METHODS
@@ -122,7 +122,9 @@ def __init__(
         use_unstable_protocol: bool = False,
         **connection_kwargs: Any,
     ) -> None:
-        client, handler = self._prepare(to_client, use_unstable_protocol=use_unstable_protocol)
+        client = to_client(self) if callable(to_client) else to_client
+        self._session_updates = _SessionUpdateTracker(cast(Client, client))
+        handler = build_client_router(cast(Client, self._session_updates), use_unstable_protocol=use_unstable_protocol)
 
         if isinstance(input_stream, Transport):
             if output_stream is not None:
@@ -134,34 +136,6 @@ def __init__(
             ):
                 raise TypeError(_CLIENT_CONNECTION_ERROR)
             self._conn = Connection(handler, input_stream, output_stream, **connection_kwargs)
-        self._notify_connected(client)
-
-    @classmethod
-    def _attach(
-        cls,
-        to_client: Callable[[Agent], Client] | Client,
-        connection: Connection,
-        *,
-        use_unstable_protocol: bool = False,
-    ) -> tuple[ClientSideConnection, MethodHandler]:
-        self = cls.__new__(cls)
-        client, handler = self._prepare(to_client, use_unstable_protocol=use_unstable_protocol)
-        self._conn = connection
-        self._notify_connected(client)
-        return self, handler
-
-    def _prepare(
-        self,
-        to_client: Callable[[Agent], Client] | Client,
-        *,
-        use_unstable_protocol: bool,
-    ) -> tuple[Client, MethodHandler]:
-        client = cast(Client, to_client(self) if callable(to_client) else to_client)
-        self._session_updates = _SessionUpdateTracker(client)
-        handler = build_client_router(cast(Client, self._session_updates), use_unstable_protocol=use_unstable_protocol)
-        return client, handler
-
-    def _notify_connected(self, client: Client) -> None:
         if on_connect := getattr(client, "on_connect", None):
             on_connect(self)
 

src/acp/experimental/negotiation.py

@@ -78,7 +78,6 @@ def __init__(
         self._v2_agent = v2_agent
         self._connection: Connection | None = None
         self._selected: MethodHandler | None = None
-        self._endpoint: V1AgentSideConnection | V2AgentSideConnection | None = None
         self._lock = asyncio.Lock()
 
     def bind_connection(self, connection: Connection) -> None:
@@ -103,10 +102,10 @@ async def _initialize(self, method: str, params: Any, is_notification: bool) ->
 
         if self._v2_agent is not None and requested >= v2.PROTOCOL_VERSION:
             selected = v2.PROTOCOL_VERSION
-            endpoint, handler = V2AgentSideConnection._attach(self._v2_agent, connection)
+            _, handler = V2AgentSideConnection.attach(self._v2_agent, connection)
         elif self._v1_agent is not None and requested >= v1_meta.PROTOCOL_VERSION:
             selected = v1_meta.PROTOCOL_VERSION
-            endpoint, handler = V1AgentSideConnection._attach(self._v1_agent, connection)
+            _, handler = V1AgentSideConnection.attach(self._v1_agent, connection)
         else:
             supported = [
                 version
@@ -119,7 +118,6 @@ async def _initialize(self, method: str, params: Any, is_notification: bool) ->
             raise RequestError.invalid_request({
                 "details": f"Unsupported ACP protocol {requested}; configured versions are {supported}"
             })
-        self._endpoint = endpoint
         self._selected = handler
         normalized = _normalize_initialize(params, selected)
         response = await handler(method, normalized, False)

src/acp/experimental/v2/agent.py

@@ -6,7 +6,7 @@
 
 from pydantic import BaseModel
 
-from acp.connection import Connection
+from acp.connection import Connection, MethodHandler
 
 from . import schema
 from ._connection import open_connection
@@ -75,11 +75,12 @@ def __init__(
             on_connect(self)
 
     @classmethod
-    def _attach(
+    def attach(
         cls,
         agent_factory: Callable[[AgentSideConnection], object],
         connection: Connection,
-    ) -> tuple[AgentSideConnection, _AgentRouter]:
+    ) -> tuple[AgentSideConnection, MethodHandler]:
+        """Attach an agent-side wrapper to an existing connection."""
         self = cls.__new__(cls)
         self._state = InitializationState()
         self._conn = connection