Python: emit oauth_consent_request in hosting _to_outputs() (fixes #5535)

python/packages/core/agent_framework/__init__.py

@@ -153,6 +153,7 @@
     SessionContext,
     register_state_type,
 )
+ChatMessageStore = HistoryProvider
 from ._settings import SecretString, load_settings
 from ._skills import (
     AggregatingSkillsSource,
@@ -371,6 +372,7 @@
     "ChatMiddleware",
     "ChatMiddlewareLayer",
     "ChatMiddlewareTypes",
+    "ChatMessageStore",
     "ChatOptions",
     "ChatResponse",
     "ChatResponseUpdate",

python/packages/core/agent_framework/observability.py

@@ -200,6 +200,7 @@ class OtelAttr(str, Enum):
     # Usage attributes
     INPUT_TOKENS = "gen_ai.usage.input_tokens"
     OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
+    TOTAL_TOKENS = "gen_ai.usage.total_tokens"
     # Tool attributes
     TOOL_CALL_ID = "gen_ai.tool.call.id"
     TOOL_DESCRIPTION = "gen_ai.tool.description"
@@ -1650,8 +1651,7 @@ async def get_embeddings(
             duration = perf_counter() - start_time_stamp
             response_attributes: dict[str, Any] = {**attributes}
             usage = result.usage or {}
-            if (input_tokens := usage.get("input_token_count")) is not None:
-                response_attributes[OtelAttr.INPUT_TOKENS] = input_tokens
+            _apply_usage_details(response_attributes, usage)
             _capture_response(
                 span=span,
                 attributes=response_attributes,
@@ -2259,19 +2259,31 @@ def _mark_inner_response_telemetry_captured(response: ChatResponse | AgentRespon
             INNER_ACCUMULATED_USAGE.set(add_usage_details(accumulated, response.usage_details))
 
 
+_USAGE_ATTR_MAP: dict[str, str] = {
+    "input_token_count": OtelAttr.INPUT_TOKENS,
+    "output_token_count": OtelAttr.OUTPUT_TOKENS,
+    "total_token_count": OtelAttr.TOTAL_TOKENS,
+}
+
+
+def _apply_usage_details(attributes: dict[str, Any], usage: Mapping[str, Any]) -> None:
+    for usage_key, otel_attr in _USAGE_ATTR_MAP.items():
+        value = usage.get(usage_key)
+        if value is not None:
+            attributes[otel_attr] = value
+    for usage_key, value in usage.items():
+        if usage_key not in _USAGE_ATTR_MAP and value is not None:
+            attributes[f"gen_ai.usage.{usage_key}"] = value
+
+
 def _apply_accumulated_usage(attributes: dict[str, Any], captured_fields: set[str]) -> None:
     """Apply accumulated usage from inner chat spans to the invoke_agent span attributes."""
     if INNER_USAGE_CAPTURED_FIELD not in captured_fields:
         return
     accumulated = INNER_ACCUMULATED_USAGE.get()
     if not accumulated:
         return
-    input_tokens = accumulated.get("input_token_count")
-    if input_tokens:
-        attributes[OtelAttr.INPUT_TOKENS] = input_tokens
-    output_tokens = accumulated.get("output_token_count")
-    if output_tokens:
-        attributes[OtelAttr.OUTPUT_TOKENS] = output_tokens
+    _apply_usage_details(attributes, accumulated)
 
 
 def _get_response_attributes(
@@ -2294,12 +2306,7 @@ def _get_response_attributes(
     if model := getattr(response, "model", None):
         attributes[OtelAttr.RESPONSE_MODEL] = model
     if capture_usage and (usage := response.usage_details):
-        input_tokens = usage.get("input_token_count")
-        if input_tokens:
-            attributes[OtelAttr.INPUT_TOKENS] = input_tokens
-        output_tokens = usage.get("output_token_count")
-        if output_tokens:
-            attributes[OtelAttr.OUTPUT_TOKENS] = output_tokens
+        _apply_usage_details(attributes, usage)
     return attributes
 
 

python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py

@@ -1627,6 +1627,14 @@ async def _to_outputs(
             max_output_length=content.max_output_length,
         ):
             yield event
+    elif content.type == "oauth_consent_request":
+        server_label = (content.additional_properties or {}).get("server_label", "")
+        text = (
+            f"OAuth consent required for MCP server '{server_label}'. "
+            f"Please visit:\n{content.consent_link}"
+        )
+        async for event in stream.aoutput_item_message(text):
+            yield event
     elif content.type == "function_approval_request":
         function_call: Content = content.function_call  # type: ignore
         server_label = function_call.additional_properties.get("server_label", "agent_framework")