Python: Preserve AG-UI citations through client and agent streaming

python/packages/ag-ui/agent_framework_ag_ui/_client.py

@@ -17,6 +17,8 @@
 
 import httpx
 from agent_framework import (
+    AgentResponse,
+    AgentResponseUpdate,
     BaseChatClient,
     ChatResponse,
     ChatResponseUpdate,
@@ -30,7 +32,7 @@
 from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer
 from agent_framework.observability import ChatTelemetryLayer
 
-from ._event_converters import AGUIEventConverter, _finalize_agui_response
+from ._event_converters import AGUIEventConverter, _finalize_agui_response, _finalize_agui_updates
 from ._feature_usage import FeatureIndex
 from ._http_service import AGUIHttpService, _serialize_available_interrupts, _serialize_resume
 from ._message_adapters import agent_framework_messages_to_agui
@@ -403,6 +405,28 @@ def _get_thread_id(self, options: Mapping[str, Any]) -> str:
 
         return thread_id
 
+    @override
+    def _finalize_function_invocation_updates(
+        self,
+        updates: Sequence[ChatResponseUpdate],
+        *,
+        response_format: Any | None = None,
+    ) -> ChatResponse[Any]:
+        """Keep message-linked annotations attached across the complete tool loop."""
+        return _finalize_agui_response(updates, response_format=response_format)
+
+    @override
+    def _finalize_agent_response_updates(
+        self,
+        updates: Sequence[AgentResponseUpdate],
+        *,
+        response_format: Any | None = None,
+    ) -> AgentResponse[Any]:
+        """Preserve message-linked annotations after chat updates are mapped to an agent."""
+        return _finalize_agui_updates(
+            updates, lambda items: AgentResponse.from_updates(items, output_format_type=response_format)
+        )
+
     @override
     def _inner_get_response(
         self,

python/packages/ag-ui/agent_framework_ag_ui/_event_converters.py

@@ -5,43 +5,74 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Sequence
-from typing import Any, cast
+from collections.abc import Callable, Sequence
+from copy import copy
+from typing import Any, TypeVar, cast
 
 from agent_framework import (
+    AgentResponse,
+    AgentResponseUpdate,
     Annotation,
     ChatResponse,
     ChatResponseUpdate,
     Content,
     Message,
 )
 
+AGUIUpdateT = TypeVar("AGUIUpdateT", ChatResponseUpdate, AgentResponseUpdate)
+AGUIResponseT = TypeVar("AGUIResponseT", bound=ChatResponse[Any] | AgentResponse[Any])
+
 logger = logging.getLogger(__name__)
 
 
-def _annotation_batch_from_update(update: ChatResponseUpdate) -> tuple[str, list[Annotation]] | None:
+def _annotation_batch_from_update(
+    update: ChatResponseUpdate | AgentResponseUpdate,
+) -> tuple[str, list[Annotation]] | None:
     """Extract a message-linked annotation batch produced by this converter."""
     custom_event = (update.additional_properties or {}).get("ag_ui_custom_event")
     if not isinstance(custom_event, dict) or custom_event.get("name") != "annotations" or not update.message_id:
         return None
+    # Buffered middleware can copy this response-level event metadata onto a rebuilt message.
+    # Only the original annotation-only update should be excluded from message aggregation.
+    value = custom_event.get("value")
+    if not isinstance(value, dict) or value.get("messageId") != update.message_id:
+        return None
+    if any(content.type != "text" or content.text for content in update.contents):
+        return None
     annotations = [
         annotation for content in update.contents if content.type == "text" for annotation in content.annotations or []
     ]
     return update.message_id, annotations
 
 
-def _finalize_agui_response(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
-    """Aggregate AG-UI updates while attaching annotation events by message ID."""
+def _finalize_agui_response(
+    updates: Sequence[ChatResponseUpdate], *, response_format: Any | None = None
+) -> ChatResponse[Any]:
+    """Aggregate chat updates while attaching annotation events by message ID."""
+    return _finalize_agui_updates(
+        updates, lambda items: ChatResponse.from_updates(items, output_format_type=response_format)
+    )
+
+
+def _finalize_agui_updates(
+    updates: Sequence[AGUIUpdateT], aggregate: Callable[[Sequence[AGUIUpdateT]], AGUIResponseT]
+) -> AGUIResponseT:
+    """Attach citations while preserving the response type and original update metadata."""
     annotation_batches: list[tuple[str, list[Annotation]]] = []
-    aggregatable_updates: list[ChatResponseUpdate] = []
+    aggregatable_updates: list[AGUIUpdateT] = []
     for update in updates:
         annotation_batch = _annotation_batch_from_update(update)
         if annotation_batch is None:
             aggregatable_updates.append(update)
         else:
             annotation_batches.append(annotation_batch)
+            # Keep response-level fields from the event without changing any message.
+            metadata_update = copy(update)
+            metadata_update.contents = []
+            metadata_update.message_id = metadata_update.role = metadata_update.author_name = None
+            aggregatable_updates.append(metadata_update)
 
-    response = ChatResponse.from_updates(aggregatable_updates)
+    response = aggregate(aggregatable_updates)
 
     # Annotation events are excluded from message aggregation, but their response metadata
     # and raw representations retain their original stream ordering.