Python: preserve multimodal messages in AG-UI chat client requests

python/packages/ag-ui/README.md

@@ -239,7 +239,7 @@ Review focus: whether these names are the right stable contract for Python users
 
 | Surface | Public exports |
 | --- | --- |
-| `agent_framework.ag_ui` facade | `AgentFrameworkAgent`, `AgentFrameworkWorkflow`, `AGUIChatClient`, `AGUIEventConverter`, `AGUIHttpService`, `AGUIThreadSnapshot`, `AGUIThreadSnapshotStore`, `InMemoryAGUIThreadSnapshotStore`, `SnapshotScopeResolver`, `add_agent_framework_fastapi_endpoint`, `state_update`, `__version__` |
+| `agent_framework.ag_ui` facade | `AgentFrameworkAgent`, `AgentFrameworkWorkflow`, `AGUIChatClient`, `AGUIEventConverter`, `AGUIHttpService`, `AGUIThreadSnapshot`, `AGUIThreadSnapshotStore`, `InMemoryAGUIThreadSnapshotStore`, `SnapshotScopeResolver`, `add_agent_framework_fastapi_endpoint`, `state_carrier`, `state_update`, `__version__` |
 | Direct `agent_framework_ag_ui` package | Facade exports plus `AGUIChatOptions`, `AGUIRequest`, `AGUIThreadID`, `AgentState`, `DEFAULT_MAX_THREAD_SNAPSHOTS`, `DEFAULT_TAGS`, `PredictStateConfig`, `RunMetadata`, `SnapshotScope`, `WorkflowFactory` |
 | AG-UI protocol package (`ag_ui.core`) | `Interrupt`, `ResumeEntry`, `RunFinishedInterruptOutcome`, and related run outcome models |
 

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

@@ -19,7 +19,7 @@
     SnapshotScope,
     SnapshotScopeResolver,
 )
-from ._state import state_update
+from ._state import state_carrier, state_update
 from ._types import AgentState, AGUIChatOptions, AGUIRequest, PredictStateConfig, RunMetadata
 from ._workflow import AgentFrameworkWorkflow, WorkflowFactory
 
@@ -55,6 +55,7 @@
     "SnapshotScopeResolver",
     "DEFAULT_MAX_THREAD_SNAPSHOTS",
     "DEFAULT_TAGS",
+    "state_carrier",
     "state_update",
     "__version__",
     # A2UI (lazy — require ag-ui-a2ui-toolkit)

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

@@ -32,6 +32,7 @@
 from ._feature_usage import FeatureIndex
 from ._http_service import AGUIHttpService, _serialize_available_interrupts, _serialize_resume
 from ._message_adapters import agent_framework_messages_to_agui
+from ._state import STATE_CARRIER_KEY
 from ._utils import convert_tools_to_agui_format
 
 if sys.version_info >= (3, 13):
@@ -293,22 +294,34 @@ def _extract_state_from_messages(self, messages: Sequence[Message]) -> tuple[lis
 
         last_message = messages[-1]
 
-        for content in last_message.contents:
-            if isinstance(content, Content) and content.type == "data" and content.media_type == "application/json":
-                try:
-                    uri = content.uri
-                    prefix, _, encoded_data = uri.partition(",")  # type: ignore[union-attr]
-                    media_type, *parameters = prefix[5:].split(";")
-                    if prefix.startswith("data:") and media_type == "application/json" and "base64" in parameters:
-                        import base64
-
-                        decoded_bytes = base64.b64decode(encoded_data, validate=True)
-                        state = json.loads(decoded_bytes.decode("utf-8"))
-
-                        messages_without_state = list(messages[:-1]) if len(messages) > 1 else []
-                        return messages_without_state, state
-                except (BinasciiError, json.JSONDecodeError, ValueError, KeyError) as e:
-                    logger.warning(f"Failed to extract state from message: {e}")
+        # A state carrier is an explicitly marked, dedicated final message. Mixed
+        # messages and ordinary JSON documents must reach the normal AG-UI message
+        # converter intact.
+        if len(last_message.contents) != 1:
+            return list(messages), None
+
+        content = last_message.contents[0]
+        if not isinstance(content, Content):
+            return list(messages), None
+        if (content.additional_properties or {}).get(STATE_CARRIER_KEY) is not True:
+            return list(messages), None
+        if content.type != "data" or content.media_type != "application/json":
+            return list(messages), None
+
+        try:
+            uri = content.uri
+            prefix, _, encoded_data = uri.partition(",")  # type: ignore[union-attr]
+            media_type, *parameters = prefix[5:].split(";")
+            if prefix.startswith("data:") and media_type == "application/json" and "base64" in parameters:
+                import base64
+
+                decoded_bytes = base64.b64decode(encoded_data, validate=True)
+                state = json.loads(decoded_bytes.decode("utf-8"))
+
+                messages_without_state = list(messages[:-1]) if len(messages) > 1 else []
+                return messages_without_state, state
+        except (BinasciiError, json.JSONDecodeError, ValueError, KeyError) as e:
+            logger.warning(f"Failed to extract state from message: {e}")
 
         return list(messages), None
 

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

@@ -958,19 +958,55 @@ def _filter_modified_args(
     return result
 
 
-def _encode_agui_segment(contents: list[Content]) -> tuple[str, list[dict[str, Any]]]:
-    """Encode assistant contents into an AG-UI ``(content, tool_calls)`` pair.
+def _convert_framework_content_to_agui(content: Content) -> dict[str, Any] | None:
+    """Convert Agent Framework media content to an AG-UI input part."""
+    if content.type not in {"uri", "data"} or not content.uri:
+        return None
+
+    media_type = content.media_type
+    media_type_prefix = media_type.lower().split("/", 1)[0] if media_type else ""
+    part_type = media_type_prefix if media_type_prefix in {"image", "audio", "video"} else "document"
+
+    if content.type == "data":
+        data_uri_prefix, separator, encoded_data = content.uri.partition(",")
+        is_base64_data_uri = bool(separator) and any(
+            parameter.lower() == "base64" for parameter in data_uri_prefix.split(";")[1:]
+        )
+        if is_base64_data_uri:
+            source: dict[str, Any] = {"type": "data", "value": encoded_data}
+        else:
+            source = {"type": "url", "value": content.uri}
+    else:
+        source = {"type": "url", "value": content.uri}
+
+    if media_type is not None:
+        source["mimeType"] = media_type
+    return {"type": part_type, "source": source}
+
+
+def _encode_agui_segment(
+    contents: list[Content], role: str
+) -> tuple[str | list[dict[str, Any]], list[dict[str, Any]]]:
+    """Encode a framework content segment into AG-UI message content and tool calls.
 
-    Shared by both the single-message path (``agent_framework_messages_to_agui``) and the
-    split path (``_split_mixed_message_to_agui``) so the text / function_call
-    serialization lives in one place. A future argument-format or supported-content
-    change then updates both paths at once instead of drifting between them.
+    The shared encoder preserves ordered user text and media parts for both the
+    single-message and function-result split paths. Non-user messages retain AG-UI's
+    string-content shape.
     """
     text = ""
+    input_content_parts: list[dict[str, Any]] = []
+    has_multimodal_content = False
     tool_calls: list[dict[str, Any]] = []
     for content in contents:
         if content.type == "text":
-            text += content.text or ""
+            text_content = content.text or ""
+            text += text_content
+            if role == "user":
+                input_content_parts.append({"type": "text", "text": text_content})
+        elif role == "user" and content.type in {"uri", "data"}:
+            if input_part := _convert_framework_content_to_agui(content):
+                input_content_parts.append(input_part)
+                has_multimodal_content = True
         elif content.type == "function_call":
             tool_calls.append(
                 {
@@ -982,7 +1018,8 @@ def _encode_agui_segment(contents: list[Content]) -> tuple[str, list[dict[str, A
                     },
                 }
             )
-    return text, tool_calls
+    message_content: str | list[dict[str, Any]] = input_content_parts if has_multimodal_content else text
+    return message_content, tool_calls
 
 
 def _split_mixed_message_to_agui(msg: Message, role: str, unresolved_call_ids: set[str]) -> list[dict[str, Any]]:
@@ -1045,7 +1082,7 @@ def flush_segment() -> None:
         nonlocal seg_contents, seg_has_call
         if not seg_contents:
             return
-        seg_text, seg_tool_calls = _encode_agui_segment(seg_contents)
+        seg_text, seg_tool_calls = _encode_agui_segment(seg_contents, role)
         seg_contents = []
         seg_has_call = False
         seg_call_ids.clear()
@@ -1079,7 +1116,7 @@ def drain_queued() -> None:
         queued_results.clear()
 
     for content in msg.contents:
-        if content.type in ("text", "function_call"):
+        if content.type in ("text", "function_call") or (role == "user" and content.type in {"uri", "data"}):
             seg_contents.append(content)
             if content.type == "function_call":
                 seg_has_call = True
@@ -1184,12 +1221,12 @@ def track_emitted(
             result.extend(_split_mixed_message_to_agui(msg, role, unresolved_call_ids))
             continue
 
-        content_text, tool_calls = _encode_agui_segment(msg.contents)
+        message_content, tool_calls = _encode_agui_segment(msg.contents, role)
 
         agui_msg: dict[str, Any] = {
             "id": msg.message_id if msg.message_id else generate_event_id(),  # Always include id
             "role": role,
-            "content": content_text,
+            "content": message_content,
         }
 
         if tool_calls:

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

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-"""Deterministic tool-driven AG-UI state updates and display payloads.
+"""AG-UI state carrier and deterministic tool-result state helpers.
 
 Tools wired into the :mod:`agent_framework_ag_ui` endpoint can push a
 deterministic state update or a per-call tool result display payload by
@@ -23,9 +23,12 @@
 
 from ._utils import make_json_safe
 
-__all__ = ["TOOL_RESULT_DISPLAY_KEY", "TOOL_RESULT_STATE_KEY", "state_update"]
+__all__ = ["STATE_CARRIER_KEY", "TOOL_RESULT_DISPLAY_KEY", "TOOL_RESULT_STATE_KEY", "state_carrier", "state_update"]
 
 
+STATE_CARRIER_KEY = "__ag_ui_state_carrier__"
+"""Reserved ``Content.additional_properties`` key marking an AG-UI request state carrier."""
+
 TOOL_RESULT_STATE_KEY = "__ag_ui_tool_result_state__"
 """Reserved ``Content.additional_properties`` key used to carry a tool-driven
 state snapshot from a tool return value through to the AG-UI emitter."""
@@ -40,6 +43,44 @@ def _serialize_tool_result(value: Any) -> str:  # noqa: ANN401
     return value if isinstance(value, str) else json.dumps(make_json_safe(value))
 
 
+def state_carrier(state: Mapping[str, Any]) -> Content:
+    """Build a dedicated final-message carrier for ``AGUIChatClient`` request state.
+
+    Add the returned content as the only content in a final user message. The
+    client recognizes its explicit marker, moves the JSON object into the AG-UI
+    request's ``state`` field, and does not send the carrier as a chat message.
+    Ordinary ``application/json`` content without this marker remains a document
+    input.
+
+    Example:
+        .. code-block:: python
+
+            from agent_framework import Message
+            from agent_framework_ag_ui import state_carrier
+
+            messages = [
+                Message(role="user", contents=["Update the dashboard"]),
+                Message(role="user", contents=[state_carrier({"selected_tab": "sales"})]),
+            ]
+
+    Args:
+        state: JSON-compatible mapping to send as AG-UI shared state.
+
+    Returns:
+        A JSON ``Content`` marked as an AG-UI request state carrier.
+
+    Raises:
+        TypeError: If ``state`` is not a mapping.
+    """
+    if not isinstance(state, Mapping):
+        raise TypeError(f"state_carrier() 'state' must be a Mapping, got {type(state).__name__}")
+    return Content.from_data(
+        json.dumps(make_json_safe(dict(state))).encode("utf-8"),
+        media_type="application/json",
+        additional_properties={STATE_CARRIER_KEY: True},
+    )
+
+
 def state_update(
     text: str = "",
     *,