refactor(v2): remove redundant type casts

Author: PsiACE

src/acp/experimental/negotiation.py

@@ -2,7 +2,7 @@
 
 import asyncio
 from collections.abc import Callable
-from typing import Any, cast
+from typing import Any
 
 from pydantic import BaseModel
 
@@ -97,15 +97,28 @@ async def _initialize(self, method: str, params: Any, is_notification: bool) ->
         if is_notification or method != V2_AGENT_METHODS["initialize"]:
             raise RequestError.invalid_request({"details": "The first ACP request must be initialize"})
         requested = _read_protocol_version(params)
-        selected = self._select(requested)
         connection = self._connection
         if connection is None:
             raise RuntimeError("Protocol router is not connected")
 
-        if selected == v2.PROTOCOL_VERSION:
-            endpoint, handler = V2AgentSideConnection._attach(cast(Any, self._v2_agent), connection)
+        if self._v2_agent is not None and requested >= v2.PROTOCOL_VERSION:
+            selected = v2.PROTOCOL_VERSION
+            endpoint, 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)
         else:
-            endpoint, handler = V1AgentSideConnection._attach(cast(Any, self._v1_agent), connection)
+            supported = [
+                version
+                for version, implementation in (
+                    (v1_meta.PROTOCOL_VERSION, self._v1_agent),
+                    (v2.PROTOCOL_VERSION, self._v2_agent),
+                )
+                if implementation is not None
+            ]
+            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)
@@ -117,23 +130,6 @@ async def _initialize(self, method: str, params: Any, is_notification: bool) ->
             })
         return response
 
-    def _select(self, requested: int) -> int:
-        if self._v2_agent is not None and requested >= v2.PROTOCOL_VERSION:
-            return v2.PROTOCOL_VERSION
-        if self._v1_agent is not None and requested >= v1_meta.PROTOCOL_VERSION:
-            return v1_meta.PROTOCOL_VERSION
-        supported = [
-            version
-            for version, implementation in (
-                (v1_meta.PROTOCOL_VERSION, self._v1_agent),
-                (v2.PROTOCOL_VERSION, self._v2_agent),
-            )
-            if implementation is not None
-        ]
-        raise RequestError.invalid_request({
-            "details": f"Unsupported ACP protocol {requested}; configured versions are {supported}"
-        })
-
 
 class AgentProtocolConnection:
     def __init__(self, connection: Connection) -> None:

src/acp/experimental/v2/agent.py

@@ -2,7 +2,7 @@
 
 import asyncio
 from collections.abc import Callable
-from typing import Any, cast
+from typing import Any
 
 from pydantic import BaseModel
 
@@ -37,10 +37,10 @@ def __init__(self, agent: Agent, state: InitializationState) -> None:
     async def __call__(self, method: str, params: Any | None, is_notification: bool) -> Any:
         initialize = self._router.request_spec("initialize")
         if not is_notification and initialize is not None and method == initialize.method:
-            request = cast(schema.InitializeRequest, initialize.request.validate_python(params))
+            request: schema.InitializeRequest = initialize.request.validate_python(params)
             self._state.begin(request)
             try:
-                response = cast(schema.InitializeResponse, await self._router.handle_request(initialize, params))
+                response: schema.InitializeResponse = await self._router.handle_request(initialize, params)
                 self._state.complete(response)
             except BaseException as error:
                 self._state.fail(error)

tests/test_protocol_negotiation.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from typing import Any, cast
+from typing import Any
 
 import pytest
 
@@ -16,6 +16,10 @@ class Client:
     pass
 
 
+class V1Client(acp.Client):
+    pass
+
+
 class V1Agent:
     def __init__(self) -> None:
         self.initialize_calls = 0
@@ -103,7 +107,7 @@ async def test_agent_protocol_router_selects_v1() -> None:
     v2_agent = V2Agent()
     wire: list[StreamEvent] = []
     agent_connection = AgentProtocolRouter(v1=lambda _: v1_agent, v2=lambda _: v2_agent).connect(agent_transport)
-    client_connection = acp.connect_to_agent(cast(acp.Client, Client()), client_transport, observers=[wire.append])
+    client_connection = acp.connect_to_agent(V1Client(), client_transport, observers=[wire.append])
 
     try:
         initialized = await client_connection.initialize(