Python: Emit AG-UI events for MCP tool calls, results, and text reasoning (#4760)

* Python: Emit AG-UI events for MCP tool calls, results, and text reasoning Fixes #4213 — `_emit_content()` in the AG-UI layer only handled `text`, `function_call`, `function_result`, `function_approval_request`, `usage`, and `oauth_consent_request` content types. Foundry MCP content types (`mcp_server_tool_call`, `mcp_server_tool_result`) and `text_reasoning` fell through unhandled, producing no SSE events for AG-UI consumers. Added three new handler functions wired into `_emit_content()`: - `_emit_mcp_tool_call`: emits TOOL_CALL_START + TOOL_CALL_ARGS and tracks in FlowState for MESSAGES_SNAPSHOT inclusion - `_emit_mcp_tool_result`: emits TOOL_CALL_END + TOOL_CALL_RESULT with full FlowState cleanup mirroring `_emit_tool_result` - `_emit_text_reasoning`: emits the protocol-defined reasoning event sequence (ReasoningStart → MessageStart → MessageContent → MessageEnd → ReasoningEnd) with ReasoningEncryptedValueEvent for protected_data * Add HTTP round-trip tests for MCP tool and reasoning SSE events Exercises the full POST → SSE bytes → parse → validate pipeline for mcp_server_tool_call, mcp_server_tool_result, text_reasoning, and ReasoningEncryptedValueEvent content through FastAPI TestClient. * Fix _emit_mcp_tool_result missing predictive_handler support (#4213) - Add predictive_handler parameter to _emit_mcp_tool_result and mirror the apply_pending_updates + StateSnapshotEvent block from _emit_tool_result - Forward predictive_handler from _emit_content to _emit_mcp_tool_result - Add assertion for stored arguments in MCP tool call test - Add test for predictive handler state snapshot after MCP tool result Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Apply pre-commit auto-fixes * Refactor MCP tool emit functions and add missing tests (#4213) - Extract _emit_tool_result_common shared helper to eliminate duplication between _emit_tool_result and _emit_mcp_tool_result - Remove server_name prefix from tool_call_name in _emit_mcp_tool_call; display_name now equals tool_name directly - Add test for tool_name fallback to 'mcp_tool' when tool_name is None - Add test for output=None fallback to empty string in _emit_mcp_tool_result Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Address review feedback for #4213: review comment fixes --------- Co-authored-by: Copilot <copilot@github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-16 21:04:09 +08:00 · 2026-03-20 09:41:37 +09:00
parent 1272ec5adf
commit 4afc088f01
3 changed files with 624 additions and 12 deletions
@@ -12,6 +12,12 @@ from typing import Any, cast
 from ag_ui.core import (
    BaseEvent,
    CustomEvent,
+    ReasoningEncryptedValueEvent,
+    ReasoningEndEvent,
+    ReasoningMessageContentEvent,
+    ReasoningMessageEndEvent,
+    ReasoningMessageStartEvent,
+    ReasoningStartEvent,
    RunFinishedEvent,
    StateSnapshotEvent,
    TextMessageContentEvent,
@@ -224,27 +230,28 @@ def _emit_tool_call(
    return events


-def _emit_tool_result(
-    content: Content,
+def _emit_tool_result_common(
+    call_id: str,
+    raw_result: Any,
    flow: FlowState,
    predictive_handler: PredictiveStateHandler | None = None,
 ) -> list[BaseEvent]:
-    """Emit ToolCallResult events for function_result content."""
+    """Shared helper for emitting ToolCallEnd + ToolCallResult events and performing FlowState cleanup.
+
+    Both ``_emit_tool_result`` (standard function results) and ``_emit_mcp_tool_result``
+    (MCP server tool results) delegate to this function.
+    """
    events: list[BaseEvent] = []

-    if not content.call_id:
-        return events
+    events.append(ToolCallEndEvent(tool_call_id=call_id))
+    flow.tool_calls_ended.add(call_id)

-    events.append(ToolCallEndEvent(tool_call_id=content.call_id))
-    flow.tool_calls_ended.add(content.call_id)
-
-    raw_result = content.result if content.result is not None else ""
    result_content = raw_result if isinstance(raw_result, str) else json.dumps(make_json_safe(raw_result))
    message_id = generate_event_id()
    events.append(
        ToolCallResultEvent(
            message_id=message_id,
-            tool_call_id=content.call_id,
+            tool_call_id=call_id,
            content=result_content,
            role="tool",
        )
@@ -254,7 +261,7 @@ def _emit_tool_result(
        {
            "id": message_id,
            "role": "tool",
-            "toolCallId": content.call_id,
+            "toolCallId": call_id,
            "content": result_content,
        }
    )
@@ -268,7 +275,7 @@ def _emit_tool_result(
    flow.tool_call_name = None

    if flow.message_id:
-        logger.debug("Closing text message (issue #3568 fix): message_id=%s", flow.message_id)
+        logger.debug("Closing text message: message_id=%s", flow.message_id)
        events.append(TextMessageEndEvent(message_id=flow.message_id))
    flow.message_id = None
    flow.accumulated_text = ""
@@ -276,6 +283,18 @@ def _emit_tool_result(
    return events


+def _emit_tool_result(
+    content: Content,
+    flow: FlowState,
+    predictive_handler: PredictiveStateHandler | None = None,
+) -> list[BaseEvent]:
+    """Emit ToolCallResult events for function_result content."""
+    if not content.call_id:
+        return []
+    raw_result = content.result if content.result is not None else ""
+    return _emit_tool_result_common(content.call_id, raw_result, flow, predictive_handler)
+
+
 def _emit_approval_request(
    content: Content,
    flow: FlowState,
@@ -381,6 +400,107 @@ def _emit_oauth_consent(content: Content) -> list[BaseEvent]:
    )


+def _emit_mcp_tool_call(content: Content, flow: FlowState) -> list[BaseEvent]:
+    """Emit ToolCall start/args events for MCP server tool call content.
+
+    MCP tool calls arrive as complete items (not streamed deltas), so we emit a
+    ``ToolCallStartEvent`` (and, when arguments are present, a ``ToolCallArgsEvent``)
+    immediately. This maps MCP-specific fields (tool_name, server_name) to the
+    same AG-UI ToolCall* events used by regular function calls, making MCP tool
+    execution visible to AG-UI consumers. Completion/end events are handled
+    separately by ``_emit_mcp_tool_result``.
+    """
+    events: list[BaseEvent] = []
+
+    tool_call_id = content.call_id or generate_event_id()
+    tool_name = content.tool_name or "mcp_tool"
+
+    display_name = tool_name
+
+    events.append(
+        ToolCallStartEvent(
+            tool_call_id=tool_call_id,
+            tool_call_name=display_name,
+            parent_message_id=flow.message_id,
+        )
+    )
+
+    # Serialize arguments
+    args_str = ""
+    if content.arguments:
+        args_str = (
+            content.arguments if isinstance(content.arguments, str) else json.dumps(make_json_safe(content.arguments))
+        )
+        events.append(ToolCallArgsEvent(tool_call_id=tool_call_id, delta=args_str))
+
+    # Track in flow state for MESSAGES_SNAPSHOT
+    tool_entry = {
+        "id": tool_call_id,
+        "type": "function",
+        "function": {"name": display_name, "arguments": args_str},
+    }
+    flow.pending_tool_calls.append(tool_entry)
+    flow.tool_calls_by_id[tool_call_id] = tool_entry
+
+    return events
+
+
+def _emit_mcp_tool_result(
+    content: Content, flow: FlowState, predictive_handler: PredictiveStateHandler | None = None
+) -> list[BaseEvent]:
+    """Emit ToolCallResult events for MCP server tool result content.
+
+    Delegates to the shared _emit_tool_result_common helper using content.output
+    (the MCP-specific result field) instead of content.result.
+    """
+    if not content.call_id:
+        logger.warning("MCP tool result content missing call_id, skipping")
+        return []
+    raw_output = content.output if content.output is not None else ""
+    return _emit_tool_result_common(content.call_id, raw_output, flow, predictive_handler)
+
+
+def _emit_text_reasoning(content: Content) -> list[BaseEvent]:
+    """Emit AG-UI reasoning events for text_reasoning content.
+
+    Uses the protocol-defined reasoning event types so that AG-UI consumers
+    such as CopilotKit can render reasoning natively.
+
+    Only ``content.text`` is used for the visible reasoning message. If
+    ``content.protected_data`` is present it is emitted as a
+    ``ReasoningEncryptedValueEvent`` so that consumers can persist encrypted
+    reasoning for state continuity without conflating it with display text.
+    """
+    text = content.text or ""
+    if not text and content.protected_data is None:
+        return []
+
+    message_id = content.id or generate_event_id()
+
+    events: list[BaseEvent] = [
+        ReasoningStartEvent(message_id=message_id),
+        ReasoningMessageStartEvent(message_id=message_id, role="assistant"),
+    ]
+
+    if text:
+        events.append(ReasoningMessageContentEvent(message_id=message_id, delta=text))
+
+    events.append(ReasoningMessageEndEvent(message_id=message_id))
+
+    if content.protected_data is not None:
+        events.append(
+            ReasoningEncryptedValueEvent(
+                subtype="message",
+                entity_id=message_id,
+                encrypted_value=content.protected_data,
+            )
+        )
+
+    events.append(ReasoningEndEvent(message_id=message_id))
+
+    return events
+
+
 def _emit_content(
    content: Any,
    flow: FlowState,
@@ -402,5 +522,11 @@ def _emit_content(
        return _emit_usage(content)
    if content_type == "oauth_consent_request":
        return _emit_oauth_consent(content)
+    if content_type == "mcp_server_tool_call":
+        return _emit_mcp_tool_call(content, flow)
+    if content_type == "mcp_server_tool_result":
+        return _emit_mcp_tool_result(content, flow, predictive_handler)
+    if content_type == "text_reasoning":
+        return _emit_text_reasoning(content)
    logger.debug("Skipping unsupported content type in AG-UI emitter: %s", content_type)
    return []