From e829728e138cbea42524b1645b9709d047f96633 Mon Sep 17 00:00:00 2001 From: VectorPeak Date: Fri, 3 Jul 2026 20:13:19 +0800 Subject: [PATCH 1/2] Python: Accept AG-UI state data URI parameters --- .../ag-ui/agent_framework_ag_ui/_client.py | 5 ++-- .../ag-ui/tests/ag_ui/test_ag_ui_client.py | 24 +++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_client.py b/python/packages/ag-ui/agent_framework_ag_ui/_client.py index ae81162b6d5..93b9ce7b9aa 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_client.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_client.py @@ -294,10 +294,11 @@ def _extract_state_from_messages(self, messages: Sequence[Message]) -> tuple[lis if isinstance(content, Content) and content.type == "data" and content.media_type == "application/json": try: uri = content.uri - if uri.startswith("data:application/json;base64,"): # type: ignore[union-attr] + 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 - encoded_data = uri.split(",", 1)[1] # type: ignore[union-attr] decoded_bytes = base64.b64decode(encoded_data) state = json.loads(decoded_bytes.decode("utf-8")) diff --git a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py index a04d13ab1c1..0a9ffc5dfa9 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py +++ b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py @@ -104,6 +104,30 @@ async def test_extract_state_from_messages_with_state(self) -> None: assert result_messages[0].text == "Hello" assert state == state_data + async def test_extract_state_from_messages_with_parameterized_data_uri(self) -> None: + """Test state extraction from JSON data URIs with media type parameters.""" + import base64 + + client = StubAGUIChatClient(endpoint="http://localhost:8888/") + + state_data = {"key": "value", "count": 42} + state_json = json.dumps(state_data) + state_b64 = base64.b64encode(state_json.encode("utf-8")).decode("utf-8") + + messages = [ + Message(role="user", contents=["Hello"]), + Message( + role="user", + contents=[Content.from_uri(uri=f"data:application/json;charset=utf-8;base64,{state_b64}")], + ), + ] + + result_messages, state = client.extract_state_from_messages(messages) + + assert len(result_messages) == 1 + assert result_messages[0].text == "Hello" + assert state == state_data + async def test_extract_state_invalid_json(self) -> None: """Test state extraction with invalid JSON.""" import base64 From 829a628f141518bf32975368853260806c9bd58e Mon Sep 17 00:00:00 2001 From: VectorPeak Date: Fri, 3 Jul 2026 20:36:57 +0800 Subject: [PATCH 2/2] Python: Handle invalid AG-UI state base64 --- .../ag-ui/agent_framework_ag_ui/_client.py | 5 +++-- .../ag-ui/tests/ag_ui/test_ag_ui_client.py | 16 ++++++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_client.py b/python/packages/ag-ui/agent_framework_ag_ui/_client.py index 93b9ce7b9aa..a6b68cfa2d8 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_client.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_client.py @@ -8,6 +8,7 @@ import logging import sys import uuid +from binascii import Error as BinasciiError from collections.abc import AsyncIterable, Awaitable, Mapping, MutableSequence, Sequence from functools import wraps from typing import TYPE_CHECKING, Any, Generic, TypedDict, cast @@ -299,12 +300,12 @@ def _extract_state_from_messages(self, messages: Sequence[Message]) -> tuple[lis if prefix.startswith("data:") and media_type == "application/json" and "base64" in parameters: import base64 - decoded_bytes = base64.b64decode(encoded_data) + 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 (json.JSONDecodeError, ValueError, KeyError) as e: + except (BinasciiError, json.JSONDecodeError, ValueError, KeyError) as e: logger.warning(f"Failed to extract state from message: {e}") return list(messages), None diff --git a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py index 0a9ffc5dfa9..e564f21228f 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py +++ b/python/packages/ag-ui/tests/ag_ui/test_ag_ui_client.py @@ -149,6 +149,22 @@ async def test_extract_state_invalid_json(self) -> None: assert result_messages == messages assert state is None + async def test_extract_state_invalid_base64(self) -> None: + """Test state extraction with invalid base64.""" + client = StubAGUIChatClient(endpoint="http://localhost:8888/") + + messages = [ + Message( + role="user", + contents=[Content.from_uri(uri="data:application/json;base64,not-valid-base64!")], + ), + ] + + result_messages, state = client.extract_state_from_messages(messages) + + assert result_messages == messages + assert state is None + async def test_convert_messages_to_agui_format(self) -> None: """Test message conversion to AG-UI format.""" client = StubAGUIChatClient(endpoint="http://localhost:8888/")