Python: preserve Responses API citations through AG-UI

python/packages/ag-ui/AGENTS.md

@@ -27,6 +27,10 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard.
 
 - Outbound custom events are emitted as AG-UI `CUSTOM`.
 - Usage metadata from `Content(type="usage")` is surfaced as `CUSTOM` events with `name="usage"`.
+- Text annotations are surfaced as `CUSTOM` events with `name="annotations"` and
+  `value={"messageId": ..., "annotations": [...]}`. Annotation-only updates announce a text message if needed
+  without emitting empty text deltas. Provider `raw_representation` is omitted; `AGUIChatClient` restores the
+  annotations on text content. These live custom events are not replayed by `MESSAGES_SNAPSHOT`.
 - Inbound custom event aliases are accepted: `CUSTOM`, `CUSTOM_EVENT`, and `custom_event`.
 - Multimodal user inputs support both legacy (`text`, `binary`) and draft-style (`image`, `audio`, `video`, `document`) shapes.
 - Interrupted runs complete with `RUN_FINISHED.outcome.type == "interrupt"` and canonical `outcome.interrupts`; do not document or add new flows that depend on the legacy top-level `RUN_FINISHED.interrupt` field.

python/packages/ag-ui/README.md

@@ -119,6 +119,42 @@ and response-cookie handling, and must be closed by the caller. Scope a cookie-b
 client to a single authenticated principal; do not share it across users. AG-UI thread IDs
 are correlation identifiers, not authentication boundaries.
 
+## Citations and annotations
+
+Text content carrying `Content.annotations` emits a message-linked `CUSTOM` event named `annotations`.
+This includes annotation-only updates received after the response text, such as SharePoint grounding
+citations from the Responses API. The event arrives before `RUN_FINISHED` and does not repeat response text:
+
+```json
+{
+  "type": "CUSTOM",
+  "name": "annotations",
+  "value": {
+    "messageId": "assistant-message-id",
+    "annotations": [
+      {
+        "type": "citation",
+        "title": "Document",
+        "url": "https://example.sharepoint.com/document.pdf",
+        "annotated_regions": [
+          {"type": "text_span", "start_index": 0, "end_index": 6}
+        ]
+      }
+    ]
+  }
+}
+```
+
+Each event contains a batch of newly emitted annotations for the indicated text message. Frontends can
+append those annotations to that message and use the citation URL, title, or file ID to render sources.
+Framework annotation fields and additional properties are retained, but provider `raw_representation`
+objects are omitted.
+
+`AGUIChatClient` restores these batches as `Content.annotations` in streaming updates and aggregated
+responses, while retaining the custom-event payload in `update.additional_properties["ag_ui_custom_event"]`.
+Custom events are live run metadata: citation rendering and persistence alongside frontend message history
+remain application responsibilities; `MESSAGES_SNAPSHOT` does not restore these custom events.
+
 ## Tool Return Helpers
 
 Use `state_update` when a backend tool needs to send different payloads to the model, the UI, and shared state. The `text` value remains the LLM-bound tool result, `tool_result` becomes the AG-UI `ToolCallResultEvent.content` for frontend rendering, and `state` is merged into durable shared state.

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

@@ -30,7 +30,7 @@
 from agent_framework._tools import FunctionInvocationConfiguration, FunctionInvocationLayer
 from agent_framework.observability import ChatTelemetryLayer
 
-from ._event_converters import AGUIEventConverter
+from ._event_converters import AGUIEventConverter, _finalize_agui_response
 from ._feature_usage import FeatureIndex
 from ._http_service import AGUIHttpService, _serialize_available_interrupts, _serialize_resume
 from ._message_adapters import agent_framework_messages_to_agui
@@ -430,17 +430,19 @@ def _inner_get_response(
                     options=options,
                     **kwargs,
                 ),
-                finalizer=ChatResponse.from_updates,
+                finalizer=_finalize_agui_response,
             )
 
         async def _get_response() -> ChatResponse:
-            return await ChatResponse.from_update_generator(
-                self._streaming_impl(
+            updates = [
+                update
+                async for update in self._streaming_impl(
                     messages=messages,
                     options=options,
                     **kwargs,
                 )
-            )
+            ]
+            return _finalize_agui_response(updates)
 
         return _get_response()
 

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

@@ -5,16 +5,76 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
+from collections.abc import Sequence
+from typing import Any, cast
 
 from agent_framework import (
+    Annotation,
+    ChatResponse,
     ChatResponseUpdate,
     Content,
+    Message,
 )
 
 logger = logging.getLogger(__name__)
 
 
+def _annotation_batch_from_update(update: ChatResponseUpdate) -> 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
+    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."""
+    annotation_batches: list[tuple[str, list[Annotation]]] = []
+    aggregatable_updates: list[ChatResponseUpdate] = []
+    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)
+
+    response = ChatResponse.from_updates(aggregatable_updates)
+
+    # Annotation events are excluded from message aggregation, but their response metadata
+    # and raw representations retain their original stream ordering.
+    response.additional_properties.clear()
+    for update in updates:
+        if update.additional_properties:
+            response.additional_properties.update(update.additional_properties)
+    if updates:
+        response.raw_representation = [update.raw_representation for update in updates]
+        response.continuation_token = updates[-1].continuation_token
+
+    for message_id, batch_annotations in annotation_batches:
+        if not batch_annotations:
+            continue
+        message = next((message for message in response.messages if message.message_id == message_id), None)
+        if message is None:
+            response.messages.append(
+                Message(
+                    role="assistant",
+                    contents=[Content.from_text(text="", annotations=batch_annotations)],
+                    message_id=message_id,
+                )
+            )
+            continue
+        text_content = next((content for content in message.contents if content.type == "text"), None)
+        if text_content is None:
+            message.contents.append(Content.from_text(text="", annotations=batch_annotations))
+        else:
+            text_content.annotations = [*(text_content.annotations or []), *batch_annotations]
+
+    return response
+
+
 class AGUIEventConverter:
     """Converter for AG-UI events to Agent Framework types.
 
@@ -261,9 +321,9 @@ def _handle_run_error(self, event: dict[str, Any]) -> ChatResponseUpdate:
     def _handle_custom_event(self, event: dict[str, Any], raw_event_type: str) -> ChatResponseUpdate:
         """Handle CUSTOM/CUSTOM_EVENT events.
 
-        Custom events are surfaced as metadata so callers can inspect protocol-specific payloads.
+        Custom events remain inspectable as metadata; annotation batches also restore text annotations.
         """
-        return ChatResponseUpdate(
+        update = ChatResponseUpdate(
             role="assistant",
             contents=[],
             additional_properties={
@@ -276,3 +336,19 @@ def _handle_custom_event(self, event: dict[str, Any], raw_event_type: str) -> Ch
                 },
             },
         )
+        if event.get("name") == "annotations":
+            value = event.get("value")
+            message_id = value.get("messageId") if isinstance(value, dict) else None
+            annotations = value.get("annotations") if isinstance(value, dict) else None
+            if (
+                not isinstance(message_id, str)
+                or not message_id
+                or not isinstance(annotations, list)
+                or not all(isinstance(annotation, dict) for annotation in annotations)
+            ):
+                logger.warning("Invalid annotations custom event: expected messageId and an annotations array")
+            else:
+                update.message_id = message_id
+                if annotations:
+                    update.contents = [Content.from_text(text="", annotations=cast("list[Annotation]", annotations))]
+        return update

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

@@ -609,15 +609,35 @@ def _track_reasoning_segment(flow: FlowState, message_id: str) -> None:
     flow.snapshot_segments.append({"kind": "reasoning", "id": message_id})
 
 
+def _offset_annotation_text_spans(annotations: list[dict[str, Any]], offset: int) -> None:
+    """Rebase text spans from a delta to the accumulated message."""
+    if offset == 0:
+        return
+    for annotation in annotations:
+        regions = annotation.get("annotated_regions")
+        if not isinstance(regions, list):
+            continue
+        for region in regions:
+            if not isinstance(region, dict) or region.get("type") != "text_span":
+                continue
+            start_index = region.get("start_index")
+            end_index = region.get("end_index")
+            if isinstance(start_index, int):
+                region["start_index"] = start_index + offset
+            if isinstance(end_index, int):
+                region["end_index"] = end_index + offset
+
+
 def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> list[BaseEvent]:
-    """Emit TextMessage events for TextContent."""
-    if not content.text:
+    """Emit text deltas and message-linked annotation batches."""
+    if not content.text and not content.annotations:
         return []
 
     if skip_text or flow.waiting_for_approval:
         return []
 
     events: list[BaseEvent] = []
+    duplicate_text = False
     if not flow.message_id:
         flow.message_id = generate_event_id()
         flow.accumulated_text = ""
@@ -626,7 +646,9 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li
     elif flow.accumulated_text and content.text == flow.accumulated_text:
         # Guard against full-message replay chunks that can appear after streaming deltas.
         logger.debug("Skipping duplicate full-text delta for message_id=%s", flow.message_id)
-        return []
+        if not content.annotations:
+            return []
+        duplicate_text = True
 
     # A tool-only response may pre-open a message before its tool-call segment
     # is tracked. If that segment claims the pre-opened ID, rotate to a fresh
@@ -648,9 +670,41 @@ def _emit_text(content: Content, flow: FlowState, skip_text: bool = False) -> li
     if segment is None:
         segment = _open_text_segment(flow, flow.message_id)
 
-    events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text))
-    flow.accumulated_text += content.text
+    annotation_offset = len(flow.accumulated_text)
+    if content.text and not duplicate_text:
+        events.append(TextMessageContentEvent(message_id=flow.message_id, delta=content.text))
+        flow.accumulated_text += content.text
     segment["text"] = flow.accumulated_text
+    if content.annotations:
+        annotations = cast(
+            "list[dict[str, Any]]",
+            make_json_safe(
+                [
+                    {key: value for key, value in annotation.items() if key != "raw_representation"}
+                    for annotation in content.annotations
+                ]
+            ),
+        )
+        if content.text and not duplicate_text:
+            _offset_annotation_text_spans(annotations, annotation_offset)
+        new_annotations = annotations
+        if duplicate_text:
+            # Reconcile replay by occurrence count, not source URL or dictionary uniqueness.
+            remaining_annotations = list(segment.get("annotations", []))
+            new_annotations = []
+            for annotation in annotations:
+                if annotation in remaining_annotations:
+                    remaining_annotations.remove(annotation)
+                else:
+                    new_annotations.append(annotation)
+        if new_annotations:
+            segment.setdefault("annotations", []).extend(new_annotations)
+            events.append(
+                CustomEvent(
+                    name="annotations",
+                    value={"messageId": flow.message_id, "annotations": new_annotations},
+                )
+            )
     return events
 
 

python/packages/ag-ui/tests/ag_ui/test_event_converters.py

@@ -8,7 +8,7 @@
 import pytest
 from agent_framework import ChatResponse
 
-from agent_framework_ag_ui._event_converters import AGUIEventConverter
+from agent_framework_ag_ui._event_converters import AGUIEventConverter, _finalize_agui_response
 
 
 class TestAGUIEventConverter:
@@ -391,6 +391,95 @@ def test_custom_event_alias_conversion(self) -> None:
         assert updates[0].additional_properties["ag_ui_custom_event"]["raw_type"] == "CUSTOM_EVENT"
         assert updates[1].additional_properties["ag_ui_custom_event"]["raw_type"] == "custom_event"
 
+    @pytest.mark.parametrize("event_type", ["CUSTOM", "CUSTOM_EVENT", "custom_event"])
+    def test_annotations_custom_event_restores_content(self, event_type: str) -> None:
+        """Annotation custom events preserve their explicit message correlation and metadata."""
+        converter = AGUIEventConverter()
+        converter.current_message_id = "another-message"
+        annotations = [
+            {
+                "type": "citation",
+                "title": "Document",
+                "url": "https://example.sharepoint.com/document.pdf",
+                "annotated_regions": [{"type": "text_span", "start_index": 0, "end_index": 6}],
+            }
+        ]
+        event = {
+            "type": event_type,
+            "name": "annotations",
+            "value": {"messageId": "msg_citations", "annotations": annotations},
+        }
+
+        update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.message_id == "msg_citations"
+        assert len(update.contents) == 1
+        assert update.contents[0].type == "text"
+        assert update.contents[0].text == ""
+        assert update.contents[0].annotations == annotations
+        assert update.additional_properties is not None
+        assert update.additional_properties["ag_ui_custom_event"]["value"] == event["value"]
+
+    def test_annotation_finalizer_targets_an_existing_message_after_later_events(self) -> None:
+        """Late annotations attach by message ID without creating a duplicate message."""
+        converter = AGUIEventConverter()
+        annotations = [
+            {
+                "type": "citation",
+                "url": "https://example.com/first",
+                "annotated_regions": [{"type": "text_span", "start_index": 0, "end_index": 5}],
+            }
+        ]
+        events: list[dict[str, Any]] = [
+            {"type": "TEXT_MESSAGE_START", "messageId": "m1"},
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m1", "delta": "First"},
+            {"type": "TEXT_MESSAGE_START", "messageId": "m2"},
+            {"type": "TEXT_MESSAGE_CONTENT", "messageId": "m2", "delta": "Second"},
+            {"type": "TOOL_CALL_RESULT", "toolCallId": "call-1", "result": "done"},
+            {
+                "type": "CUSTOM",
+                "name": "annotations",
+                "value": {"messageId": "m1", "annotations": annotations},
+            },
+        ]
+        updates = [update for event in events if (update := converter.convert_event(event)) is not None]
+
+        response = _finalize_agui_response(updates)
+
+        assert [message.message_id for message in response.messages].count("m1") == 1
+        first_message = next(message for message in response.messages if message.message_id == "m1")
+        first_text = next(content for content in first_message.contents if content.type == "text")
+        assert first_text.text == "First"
+        assert first_text.annotations == annotations
+        assert [message.message_id for message in response.messages].count(None) == 1
+
+    @pytest.mark.parametrize(
+        "value",
+        [
+            None,
+            {"annotations": []},
+            {"messageId": 123, "annotations": []},
+            {"messageId": "msg_citations", "annotations": "invalid"},
+            {"messageId": "msg_citations", "annotations": ["invalid"]},
+        ],
+    )
+    def test_malformed_annotations_custom_event_preserves_metadata(
+        self, value: Any, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        """Malformed known custom events remain observable and log an explicit warning."""
+        converter = AGUIEventConverter()
+        event = {"type": "CUSTOM", "name": "annotations", "value": value}
+
+        with caplog.at_level(logging.WARNING):
+            update = converter.convert_event(event)
+
+        assert update is not None
+        assert update.contents == []
+        assert update.additional_properties is not None
+        assert update.additional_properties["ag_ui_custom_event"]["value"] == value
+        assert "annotations" in caplog.text
+
     def test_full_conversation_flow(self) -> None:
         """Test complete conversation flow with multiple event types."""
         converter = AGUIEventConverter()