Fix type checking

python/copilot/client.py

@@ -50,7 +50,6 @@
     ToolResult,
 )
 
-
 HandlerUnsubcribe = Callable[[], None]
 
 
@@ -1579,9 +1578,10 @@ async def _execute_tool_call(
         }
 
         try:
-            result = handler(invocation)
-            if inspect.isawaitable(result):
-                result = await result
+            raw_result = handler(invocation)
+            if inspect.isawaitable(raw_result):
+                raw_result = await raw_result
+            result: ToolResult = cast(ToolResult, raw_result)
         except Exception as exc:  # pylint: disable=broad-except
             # Don't expose detailed error information to the LLM for security reasons.
             # The actual error is stored in the 'error' field for debugging.

python/copilot/session.py

@@ -8,7 +8,7 @@
 import asyncio
 import inspect
 import threading
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, cast
 
 from .generated.rpc import SessionRpc
 from .generated.session_events import SessionEvent, SessionEventType, session_event_from_dict
@@ -336,7 +336,7 @@ async def _handle_permission_request(
             result = handler(request, {"session_id": self.session_id})
             if inspect.isawaitable(result):
                 result = await result
-            return result
+            return cast(PermissionRequestResult, result)
         except Exception:  # pylint: disable=broad-except
             # Handler failed, deny permission
             return {"kind": "denied-no-approval-rule-and-could-not-request-from-user"}
@@ -388,7 +388,7 @@ async def _handle_user_input_request(self, request: dict) -> UserInputResponse:
             )
             if inspect.isawaitable(result):
                 result = await result
-            return result
+            return cast(UserInputResponse, result)
         except Exception:
             raise