diff --git a/src/acp/agent/connection.py b/src/acp/agent/connection.py
index 37e7f46..f0c793c 100644
--- a/src/acp/agent/connection.py
+++ b/src/acp/agent/connection.py
@@ -8,7 +8,7 @@
 
 from .._transport import Transport
 from ..connection import Connection, MethodHandler
-from ..interfaces import Agent, Client
+from ..interfaces import Agent
 from ..meta import CLIENT_METHODS
 from ..schema import (
     AcceptElicitationResponse,
@@ -95,7 +95,7 @@ class AgentSideConnection:
 
     def __init__(
         self,
-        to_agent: Callable[[Client], Agent] | Agent,
+        to_agent: Callable[[AgentSideConnection], Agent] | Agent,
         input_stream: Any,
         output_stream: Any = None,
         listening: bool = True,
@@ -119,7 +119,7 @@ def __init__(
     @classmethod
     def attach(
         cls,
-        to_agent: Callable[[Client], Agent] | Agent,
+        to_agent: Callable[[AgentSideConnection], Agent] | Agent,
         connection: Connection,
         *,
         use_unstable_protocol: bool = False,
@@ -133,7 +133,7 @@ def attach(
 
     def _prepare(
         self,
-        to_agent: Callable[[Client], Agent] | Agent,
+        to_agent: Callable[[AgentSideConnection], Agent] | Agent,
         *,
         use_unstable_protocol: bool,
     ) -> tuple[Agent, MethodHandler]:
diff --git a/src/acp/stdio.py b/src/acp/stdio.py
index d58644e..99bd157 100644
--- a/src/acp/stdio.py
+++ b/src/acp/stdio.py
@@ -185,7 +185,7 @@ async def spawn_agent_process(
 
 @asynccontextmanager
 async def spawn_client_process(
-    to_agent: Callable[[Client], Agent] | Agent,
+    to_agent: Callable[[AgentSideConnection], Agent] | Agent,
     command: str,
     *args: str,
     env: Mapping[str, str] | None = None,