Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion python/packages/ag-ui/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -239,12 +239,14 @@ 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 |

Interrupt support is protocol data rather than a separate Agent Framework Python class. Requests accept canonical `availableInterrupts`/`available_interrupts` and `resume` values; `AGUIChatClient` and `AGUIHttpService.post_run(...)` forward those fields with AG-UI wire aliases; agent approval and workflow `request_info` pauses emit `RUN_FINISHED.outcome.interrupts`; `AGUIEventConverter` preserves canonical interrupt outcome metadata on the final `ChatResponseUpdate`; and thread snapshot hydration replays the canonical interrupt outcome when a scoped snapshot stores an unresolved pause.

Use `state_carrier(...)` to mark JSON content that should be sent in the AG-UI request's `state` field rather than as a model-visible document. The client removes explicitly marked carriers from all client-controlled history and uses the most recent carrier. Ordinary `application/json` content remains a document. For migration, pass `allow_legacy_state_carrier=True` in `AGUIChatOptions` to recognize the deprecated final single-content base64 JSON convention; mixed text and document messages remain model-visible. This client-only option emits a `DeprecationWarning` when the legacy convention is used and is not sent to the remote server.

## Features

This integration supports all 7 AG-UI features:
Expand Down
3 changes: 2 additions & 1 deletion python/packages/ag-ui/agent_framework_ag_ui/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -55,6 +55,7 @@
"SnapshotScopeResolver",
"DEFAULT_MAX_THREAD_SNAPSHOTS",
"DEFAULT_TAGS",
"state_carrier",
"state_update",
"__version__",
# A2UI (lazy — require ag-ui-a2ui-toolkit)
Expand Down
111 changes: 86 additions & 25 deletions python/packages/ag-ui/agent_framework_ag_ui/_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,12 @@

from __future__ import annotations

import base64
import json
import logging
import sys
import uuid
import warnings
from binascii import Error as BinasciiError
from collections.abc import AsyncIterable, Awaitable, Mapping, MutableSequence, Sequence
from functools import wraps
Expand All @@ -32,6 +34,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):
Expand Down Expand Up @@ -72,6 +75,54 @@ def _unwrap_server_function_call_contents(contents: MutableSequence[Content | di
)


def _is_state_carrier_message(message: Message) -> bool:
"""Return whether a message is a dedicated, explicitly marked state carrier."""
if len(message.contents) != 1:
return False
content = message.contents[0]
return isinstance(content, Content) and (content.additional_properties or {}).get(STATE_CARRIER_KEY) is True


def _decode_json_state(content: Content) -> dict[str, Any] | None:
"""Decode a base64 JSON state content, returning None for invalid input."""
if content.type != "data" or content.media_type != "application/json":
return None

try:
uri = content.uri
prefix, _, encoded_data = uri.partition(",") # type: ignore[union-attr]
if not prefix.startswith("data:"):
return None

media_type, *parameters = prefix[5:].split(";")
if media_type != "application/json" or "base64" not in parameters:
return None

decoded_bytes = base64.b64decode(encoded_data, validate=True)
state = json.loads(decoded_bytes.decode("utf-8"))
if not isinstance(state, dict):
logger.warning("AG-UI state carrier JSON must decode to an object")
return None
return state
except (BinasciiError, UnicodeDecodeError, json.JSONDecodeError, ValueError, TypeError, AttributeError) as e:
logger.warning(f"Failed to extract state from message: {e}")
return None


def _extract_legacy_json_state(message: Message) -> dict[str, Any] | None:
"""Extract the historical implicit state convention from a final state-only message."""
if len(message.contents) != 1:
return None

content = message.contents[0]
if not isinstance(content, Content):
return None
if (content.additional_properties or {}).get(STATE_CARRIER_KEY) is True:
return None

return _decode_json_state(content)


def _apply_server_function_call_unwrap(client: BaseChatClientT) -> BaseChatClientT:
"""Class decorator that unwraps server-side function calls after tool handling."""

Expand Down Expand Up @@ -279,38 +330,45 @@ def _register_server_tool_placeholder(self, tool_name: str) -> None:
self._registered_server_tools = registered
logger.debug(f"[AGUIChatClient] Registered server placeholder: {tool_name}")

def _extract_state_from_messages(self, messages: Sequence[Message]) -> tuple[list[Message], dict[str, Any] | None]:
"""Extract state from last message if present.
def _extract_state_from_messages(
self,
messages: Sequence[Message],
*,
allow_legacy_state_carrier: bool = False,
) -> tuple[list[Message], dict[str, Any] | None]:
"""Extract explicitly marked state from client-controlled message history.

Args:
messages: List of chat messages
allow_legacy_state_carrier: Whether to recognize the deprecated implicit
final base64 JSON state convention.

Returns:
Tuple of (messages_without_state, state_dict)
"""
if not messages:
return list(messages), None

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}")
messages_to_send: list[Message] = []
state: dict[str, Any] | None = None

for message in messages:
if _is_state_carrier_message(message):
content = cast(Content, message.contents[0])
if (extracted_state := _decode_json_state(content)) is not None:
state = extracted_state
continue
messages_to_send.append(message)

if allow_legacy_state_carrier and messages and not _is_state_carrier_message(messages[-1]):
legacy_state = _extract_legacy_json_state(messages[-1])
if legacy_state is not None:
messages_to_send.pop()
Comment on lines +360 to +363

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Following up on the earlier mixed-JSON fix: could legacy extraction keep the dedicated state-only boundary? With allow_legacy_state_carrier=True, _extract_legacy_json_state() accepts JSON inside a final text-plus-document message, then messages_to_send.pop() drops the entire prompt and sends the document as state. Could the fallback require exactly one unmarked JSON content so opting into migration does not reintroduce the message loss fixed earlier?

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Evan Mattson (@moonbox3) I’ve addressed this in a follow-up. The legacy fallback now requires the final message to contain exactly one unmarked application/json content, so mixed text-plus-document messages remain intact and are sent through the AG-UI converter. The single-content legacy convention remains available with allow_legacy_state_carrier=True, including the deprecation warning. A request-level regression test covers the mixed-message case.

state = legacy_state
warnings.warn(
"Implicit AG-UI JSON state extraction is deprecated; use state_carrier() instead.",
DeprecationWarning,
stacklevel=3,
)

return list(messages), None
return messages_to_send, state

def _convert_messages_to_agui_format(self, messages: list[Message]) -> list[dict[str, Any]]:
"""Convert Agent Framework messages to AG-UI format.
Expand Down Expand Up @@ -401,7 +459,10 @@ async def _streaming_impl(
ChatResponseUpdate objects
"""
mark_feature_used(FeatureIndex.AG_UI)
messages_to_send, state = self._extract_state_from_messages(messages)
messages_to_send, state = self._extract_state_from_messages(
messages,
allow_legacy_state_carrier=options.get("allow_legacy_state_carrier") is True,
)

thread_id = self._get_thread_id(options)
run_id = f"run_{uuid.uuid4().hex}"
Expand Down
61 changes: 49 additions & 12 deletions python/packages/ag-ui/agent_framework_ag_ui/_message_adapters.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
{
Expand All @@ -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]]:
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
45 changes: 43 additions & 2 deletions python/packages/ag-ui/agent_framework_ag_ui/_state.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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."""
Expand All @@ -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 message carrier for ``AGUIChatClient`` request state.

Add the returned content as the only content in a user message. The client
recognizes its explicit marker anywhere in client-controlled history, moves
the most recent carrier's JSON object into the AG-UI request's ``state``
field, and does not send carriers as chat messages. 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 = "",
*,
Expand Down
Loading
Loading