Python: fix(copilot): propagate usage details, finish_reason, and model from SDK events

python/packages/github_copilot/agent_framework_github_copilot/_agent.py

@@ -25,10 +25,13 @@
     BaseAgent,
     Content,
     ContextProvider,
+    FinishReason,
     HistoryProvider,
     Message,
     ResponseStream,
     SessionContext,
+    UsageDetails,
+    add_usage_details,
     normalize_messages,
 )
 from agent_framework._settings import load_settings
@@ -131,6 +134,51 @@ def _deny_all_permissions(
     return PermissionDecisionUserNotAvailable()
 
 
+def _parse_copilot_usage(data: Any) -> UsageDetails | None:
+    """Convert the data payload of an ``ASSISTANT_USAGE`` event to :class:`UsageDetails`.
+
+    The Copilot SDK emits an ``assistant.usage`` event after every model turn.
+    In SDK v1.0.2 the payload is a single ``Data`` object whose optional fields
+    include ``input_tokens``, ``output_tokens``, ``cache_read_tokens``, and
+    ``cache_write_tokens``.  Newer SDK releases use a dedicated
+    ``AssistantUsageData`` dataclass with the same attribute names, so this
+    helper accesses all fields via :func:`getattr` for forward compatibility.
+
+    Returns ``None`` when the payload carries no token counts at all so callers
+    can skip emitting an empty usage update.
+
+    Args:
+        data: The ``.data`` attribute of a ``SessionEvent`` whose type is
+            ``SessionEventType.ASSISTANT_USAGE``.
+
+    Returns:
+        A :class:`UsageDetails` dict if any token data is present, else ``None``.
+    """
+    input_tokens = getattr(data, "input_tokens", None)
+    output_tokens = getattr(data, "output_tokens", None)
+    if input_tokens is None and output_tokens is None:
+        return None
+    details = UsageDetails()
+    if input_tokens is not None:
+        details["input_token_count"] = int(input_tokens)
+    if output_tokens is not None:
+        details["output_token_count"] = int(output_tokens)
+    input_count = details.get("input_token_count") or 0
+    output_count = details.get("output_token_count") or 0
+    if input_count or output_count:
+        details["total_token_count"] = input_count + output_count
+    cache_read = getattr(data, "cache_read_tokens", None)
+    if cache_read is not None:
+        details["cache_read_input_token_count"] = int(cache_read)
+    cache_write = getattr(data, "cache_write_tokens", None)
+    if cache_write is not None:
+        details["cache_creation_input_token_count"] = int(cache_write)
+    reasoning = getattr(data, "reasoning_tokens", None)
+    if reasoning is not None:
+        details["reasoning_output_token_count"] = int(reasoning)
+    return details
+
+
 class GitHubCopilotSettings(TypedDict, total=False):
     """GitHub Copilot model settings.
 
@@ -594,10 +642,25 @@ async def _run_impl(
         if session_context.instructions:
             prompt = "\n".join(session_context.instructions) + "\n" + prompt
 
+        # Subscribe to events before send_and_wait so we can capture every
+        # ASSISTANT_USAGE event.  Multi-turn runs (e.g. with tool calls) emit
+        # one ASSISTANT_USAGE event per model turn; we collect them all and sum
+        # into a per-run total.  The handler coexists with send_and_wait's own
+        # internal handler — both receive every event the session emits.
+        usage_events: list[Any] = []
+
+        def _collect_usage(event: SessionEvent) -> None:
+            if event.type == SessionEventType.ASSISTANT_USAGE:
+                usage_events.append(event.data)
+
+        unsubscribe_usage = copilot_session.on(_collect_usage)
         try:
-            response_event = await copilot_session.send_and_wait(prompt, timeout=timeout)
-        except Exception as ex:
-            raise AgentException(f"GitHub Copilot request failed: {ex}") from ex
+            try:
+                response_event = await copilot_session.send_and_wait(prompt, timeout=timeout)
+            except Exception as ex:
+                raise AgentException(f"GitHub Copilot request failed: {ex}") from ex
+        finally:
+            unsubscribe_usage()
 
         response_messages: list[Message] = []
         response_id: str | None = None
@@ -619,7 +682,38 @@ async def _run_impl(
                 )
             response_id = message_id
 
-        response = AgentResponse(messages=response_messages, response_id=response_id)
+        # Build usage_details and finish_reason by aggregating all ASSISTANT_USAGE
+        # events collected during this run.  Summing per-turn counts gives the
+        # same per-run total as the streaming path (which relies on _process_update
+        # to accumulate Content.from_usage() updates).
+        # finish_reason and model are taken from the *last* event so they reflect
+        # the final model turn.
+        # Fall back to output_tokens on the ASSISTANT_MESSAGE data when no
+        # dedicated usage event arrived (older CLI versions).
+        response_usage: UsageDetails | None = None
+        response_finish_reason: FinishReason | None = None
+        response_model: str | None = None
+        if usage_events:
+            for evt_data in usage_events:
+                per_turn = _parse_copilot_usage(evt_data)
+                if per_turn:
+                    response_usage = add_usage_details(response_usage, per_turn)
+            last_evt = usage_events[-1]
+            _raw_reason = getattr(last_evt, "reason", None) or getattr(last_evt, "finish_reason", None)
+            response_finish_reason = FinishReason(_raw_reason) if _raw_reason else None
+            response_model = getattr(last_evt, "model", None)
+        elif response_event and response_event.type == SessionEventType.ASSISTANT_MESSAGE:
+            output_tokens = getattr(response_event.data, "output_tokens", None)
+            if output_tokens is not None:
+                response_usage = UsageDetails(output_token_count=int(output_tokens))
+
+        response = AgentResponse(
+            messages=response_messages,
+            response_id=response_id,
+            usage_details=response_usage,
+            finish_reason=response_finish_reason,
+            additional_properties={"model": response_model} if response_model else None,
+        )
         session_context._response = response  # type: ignore[assignment]
         await self._run_after_providers(session=session, context=session_context)
         return response
@@ -742,6 +836,32 @@ def event_handler(event: SessionEvent) -> None:
                     raw_representation=event,
                 )
                 queue.put_nowait(update)
+            elif event.type == SessionEventType.ASSISTANT_USAGE:
+                # Capture per-turn token usage and emit it as a usage update so that
+                # AgentResponse.usage_details and AgentResponse.finish_reason are
+                # populated by the standard _process_update() aggregation path.
+                # Emit whenever usage OR finish_reason is present so finish_reason
+                # is never dropped for events that carry only non-token fields
+                # (e.g. cache-only or reasoning-only turns).
+                usage_data: Any = event.data
+                streamed_usage = _parse_copilot_usage(usage_data)
+                _raw_stream_reason = getattr(usage_data, "reason", None) or getattr(
+                    usage_data, "finish_reason", None
+                )
+                stream_finish_reason = FinishReason(_raw_stream_reason) if _raw_stream_reason else None
+                if streamed_usage or stream_finish_reason:
+                    usage_contents = (
+                        [Content.from_usage(usage_details=streamed_usage, raw_representation=event)]
+                        if streamed_usage
+                        else []
+                    )
+                    usage_update = AgentResponseUpdate(
+                        role="assistant",
+                        contents=usage_contents,
+                        raw_representation=event,
+                        finish_reason=stream_finish_reason,
+                    )
+                    queue.put_nowait(usage_update)
             elif event.type == SessionEventType.SESSION_IDLE:
                 queue.put_nowait(None)
             elif event.type == SessionEventType.SESSION_ERROR: