diff --git a/python/copilot/client.py b/python/copilot/client.py index 88b3d97a7..5a5ce5874 100644 --- a/python/copilot/client.py +++ b/python/copilot/client.py @@ -579,6 +579,9 @@ async def create_session(self, config: SessionConfig) -> CopilotSession: session_id = response["sessionId"] workspace_path = response.get("workspacePath") session = CopilotSession(session_id, self._client, workspace_path) + usage_info = response.get("usageInfo") + if usage_info: + session.usage_info = usage_info session._register_tools(tools) session._register_permission_handler(on_permission_request) if on_user_input_request: @@ -761,6 +764,9 @@ async def resume_session(self, session_id: str, config: ResumeSessionConfig) -> resumed_session_id = response["sessionId"] workspace_path = response.get("workspacePath") session = CopilotSession(resumed_session_id, self._client, workspace_path) + usage_info = response.get("usageInfo") + if usage_info: + session.usage_info = usage_info session._register_tools(cfg.get("tools")) session._register_permission_handler(on_permission_request) if on_user_input_request: diff --git a/python/copilot/session.py b/python/copilot/session.py index 658d2902a..591e4c2aa 100644 --- a/python/copilot/session.py +++ b/python/copilot/session.py @@ -82,6 +82,7 @@ def __init__(self, session_id: str, client: Any, workspace_path: Optional[str] = self._user_input_handler_lock = threading.Lock() self._hooks: Optional[SessionHooks] = None self._hooks_lock = threading.Lock() + self.usage_info = None self._rpc: Optional[SessionRpc] = None @property @@ -243,6 +244,9 @@ def _dispatch_event(self, event: SessionEvent) -> None: Args: event: The session event to dispatch to all handlers. """ + if hasattr(event, "usage_info") and event.usage_info: + self.usage_info = event.usage_info + with self._event_handlers_lock: handlers = list(self._event_handlers)