feat(python): add overloads for CopilotClient.on()

python/copilot/client.py

@@ -21,7 +21,7 @@
 import threading
 from dataclasses import asdict, is_dataclass
 from pathlib import Path
-from typing import Any, Callable, Optional, cast
+from typing import Any, Callable, Optional, cast, overload
 
 from .generated.rpc import ServerRpc
 from .generated.session_events import session_event_from_dict
@@ -50,6 +50,8 @@
     ToolResult,
 )
 
+HandlerUnsubcribe = Callable[[], None]
+
 
 def _get_bundled_cli_path() -> Optional[str]:
     """Get the path to the bundled CLI binary, if available."""
@@ -1006,11 +1008,20 @@ async def set_foreground_session_id(self, session_id: str) -> None:
             error = response.get("error", "Unknown error")
             raise RuntimeError(f"Failed to set foreground session: {error}")
 
+    @overload
+    def on(self, handler: SessionLifecycleHandler, /) -> HandlerUnsubcribe: ...
+
+    @overload
+    def on(
+        self, event_type: SessionLifecycleEventType, /, handler: SessionLifecycleHandler
+    ) -> HandlerUnsubcribe: ...
+
     def on(
         self,
         event_type_or_handler: SessionLifecycleEventType | SessionLifecycleHandler,
+        /,
         handler: Optional[SessionLifecycleHandler] = None,
-    ) -> Callable[[], None]:
+    ) -> HandlerUnsubcribe:
         """
         Subscribe to session lifecycle events.
 
@@ -1567,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