diff --git a/python/packages/a2a/agent_framework_a2a/_utils.py b/python/packages/a2a/agent_framework_a2a/_utils.py index 64a3e6a7fa7..08f5b018880 100644 --- a/python/packages/a2a/agent_framework_a2a/_utils.py +++ b/python/packages/a2a/agent_framework_a2a/_utils.py @@ -1,6 +1,9 @@ # Copyright (c) Microsoft. All rights reserved. +from __future__ import annotations + import re +from urllib.parse import unquote # Key in ``SendMessageRequest.metadata`` that carries forwarded ``function_invocation_kwargs`` and # ``client_kwargs``. @@ -21,7 +24,11 @@ def get_uri_data(uri: str) -> str: Raises: ValueError: If the URI format is invalid. """ - match = URI_PATTERN.match(uri) + metadata, separator, payload = uri.partition(",") + if not separator or not URI_PATTERN.match(f"{metadata},"): + raise ValueError(f"Invalid data URI format: {uri}") + + match = URI_PATTERN.match(f"{metadata},{unquote(payload)}") if not match: raise ValueError(f"Invalid data URI format: {uri}") diff --git a/python/packages/a2a/tests/test_a2a_agent.py b/python/packages/a2a/tests/test_a2a_agent.py index 81f3e271a40..1de77c7b382 100644 --- a/python/packages/a2a/tests/test_a2a_agent.py +++ b/python/packages/a2a/tests/test_a2a_agent.py @@ -512,11 +512,12 @@ def test_prepare_message_for_a2a_with_uri_content(a2a_agent: A2AAgent) -> None: assert a2a_message.parts[0].media_type == "application/pdf" -def test_prepare_message_for_a2a_with_data_content(a2a_agent: A2AAgent) -> None: +@mark.parametrize("payload", ["SGVsbG8gV29ybGQ=", "SGVsbG8gV29ybGQ%3D"]) +def test_prepare_message_for_a2a_with_data_content(a2a_agent: A2AAgent, payload: str) -> None: """Test _prepare_message_for_a2a with DataContent.""" # Create Message with DataContent (base64 data URI) - data_content = Content.from_uri(uri="data:text/plain;base64,SGVsbG8gV29ybGQ=", media_type="text/plain") + data_content = Content.from_uri(uri=f"data:text/plain;base64,{payload}", media_type="text/plain") message = Message(role="user", contents=[data_content]) # Convert to A2A message diff --git a/python/packages/a2a/tests/test_a2a_executor.py b/python/packages/a2a/tests/test_a2a_executor.py index a215a853de5..f0266ee9a65 100644 --- a/python/packages/a2a/tests/test_a2a_executor.py +++ b/python/packages/a2a/tests/test_a2a_executor.py @@ -648,14 +648,18 @@ async def test_handle_multiple_text_contents(self, executor: A2AExecutor, mock_u mock_updater.update_status.assert_called_once() assert mock_updater.new_agent_message.called + @mark.parametrize("escape_payload", [False, True]) @mark.parametrize("data", [b"test file data", b""]) - async def test_handle_data_content(self, executor: A2AExecutor, mock_updater: MagicMock, data: bytes) -> None: + async def test_handle_data_content( + self, executor: A2AExecutor, mock_updater: MagicMock, data: bytes, escape_payload: bool + ) -> None: """Test handling messages with data content, including empty attachments.""" # Arrange - message = Message( - contents=[Content.from_data(data=data, media_type="application/octet-stream")], - role="assistant", - ) + content = Content.from_data(data=data, media_type="application/octet-stream") + if escape_payload: + assert content.uri is not None + content = Content.from_uri(content.uri.replace("=", "%3D"), media_type="application/octet-stream") + message = Message(contents=[content], role="assistant") # Act await executor.handle_events(message, mock_updater) diff --git a/python/packages/a2a/tests/test_utils.py b/python/packages/a2a/tests/test_utils.py index 61b06ab88a6..d2e33dd56c5 100644 --- a/python/packages/a2a/tests/test_utils.py +++ b/python/packages/a2a/tests/test_utils.py @@ -58,3 +58,21 @@ def test_get_uri_data_empty() -> None: """Test get_uri_data with empty string.""" with pytest.raises(ValueError, match="Invalid data URI format"): get_uri_data("") + + +@pytest.mark.parametrize( + ("uri", "expected"), + [ + ("data:application/octet-stream;base64,+/8=", "+/8="), + ("data:application/octet-stream;base64,%2B%2F8%3D", "+/8="), + ("data:application/octet-stream;base64,+%2F8=", "+/8="), + ("data:text/plain;name=keep%2Fraw;base64,SGVsbG8%3D", "SGVsbG8="), + ], +) +def test_get_uri_data_percent_encoded_payload(uri: str, expected: str) -> None: + assert get_uri_data(uri) == expected + + +def test_get_uri_data_does_not_decode_percent_escapes_twice() -> None: + with pytest.raises(ValueError, match="Invalid data URI format"): + get_uri_data("data:application/octet-stream;base64,%252B%252F8%253D") diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 0d9aed13f58..1bd922814e5 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -38,6 +38,7 @@ get_args, overload, ) +from urllib.parse import unquote from typing_extensions import Required, TypedDict @@ -196,7 +197,7 @@ def detect_media_type_from_base64( raise ValueError(str(exc)) from exc if not parsed.is_base64: raise ValueError("Data URI must use base64 encoding.") - data_str = parsed.payload + data_str = unquote(parsed.payload) if data_str is not None: if data is not None: raise ValueError("Provide exactly one of data_bytes, data_str, or data_uri.") @@ -277,7 +278,7 @@ def _get_data_bytes_as_str(content: Content) -> str | None: parsed = _parse_data_uri(uri) if not parsed.is_base64: raise ContentError("Data URI must use base64 encoding") - return parsed.payload + return unquote(parsed.payload) def _get_data_bytes(content: Content) -> bytes | None: # pyright: ignore[reportUnusedFunction] diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 771e94fb799..24c27b63c69 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -12,6 +12,7 @@ from datetime import datetime, timezone from functools import partial from typing import Any, Literal, cast +from urllib.parse import unquote import pytest from pydantic import BaseModel, Field, ValidationError @@ -7689,3 +7690,52 @@ def test_data_uri_readers_agree_on_the_base64_marker(): payload_marker = f"data:image/png,x;base64,{png_b64}" with pytest.raises(ValueError, match="Data URI must use base64 encoding."): detect_media_type_from_base64(data_uri=payload_marker) + + +@mark.parametrize("payload", ["+/8=", "%2B%2F8%3D", "+%2F8="]) +def test_data_uri_percent_encoded_payload(payload: str) -> None: + uri = f"data:application/octet-stream;name=keep%2Fraw;base64,{payload}" + content = Content.from_uri(uri) + assert _get_data_bytes_as_str(content) == "+/8=" + assert _get_data_bytes(content) == b"\xfb\xff" + assert content.uri == uri + + +@mark.parametrize("escaped", [False, True]) +def test_detect_media_type_from_percent_encoded_data_uri(escaped: bool) -> None: + payload = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGP4z8AAAAMBAQDJ/pLvAAAAAElFTkSuQmCC" + uri = "data:image/png;base64," + (payload.replace("/", "%2F") if escaped else payload) + assert detect_media_type_from_base64(data_uri=uri) == "image/png" + assert _get_data_bytes(Content.from_uri(uri)) == base64.b64decode(payload, validate=True) + + +def test_data_uri_metadata_and_serialization_do_not_decode_payload(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[str] = [] + + def tracked_unquote(payload: str) -> str: + calls.append(payload) + return unquote(payload) + + monkeypatch.setattr("agent_framework._types.unquote", tracked_unquote, raising=False) + uri = "data:application/octet-stream;name=keep%2Fraw;base64,%2B%2F8%3D" + content = Content.from_uri(uri) + message = Message("user", [content]) + restored = Message.from_dict(json.loads(json.dumps(message.to_dict()))) + assert restored.contents[0].uri == uri + assert calls == [] + + assert _get_data_bytes(restored.contents[0]) == b"\xfb\xff" + assert calls == ["%2B%2F8%3D"] + assert restored.contents[0].uri == uri + + +def test_data_uri_percent_escapes_are_decoded_only_once() -> None: + uri = "data:application/octet-stream;base64,%252B%252F8%253D" + assert _get_data_bytes_as_str(Content.from_uri(uri)) == "%2B%2F8%3D" + with raises(ValueError, match="Invalid base64"): + detect_media_type_from_base64(data_uri=uri) + + +def test_raw_base64_is_not_percent_decoded() -> None: + with raises(ValueError, match="Invalid base64"): + detect_media_type_from_base64(data_str="iVBORw0KGgo%3D") diff --git a/python/packages/hosting-a2a/agent_framework_hosting_a2a/_conversion.py b/python/packages/hosting-a2a/agent_framework_hosting_a2a/_conversion.py index 881d14414d4..cb9f6b25cf1 100644 --- a/python/packages/hosting-a2a/agent_framework_hosting_a2a/_conversion.py +++ b/python/packages/hosting-a2a/agent_framework_hosting_a2a/_conversion.py @@ -9,6 +9,7 @@ import logging from collections.abc import Collection, Mapping, Sequence from typing import Any, cast +from urllib.parse import unquote from a2a.types import Message as A2AMessage from a2a.types import Part @@ -313,7 +314,7 @@ def a2a_from_run( if not separator or not prefix.startswith("data:") or ";base64" not in prefix: raise ValueError("Agent Framework data content must contain a base64 data URI.") try: - raw = base64.b64decode(encoded, validate=True) + raw = base64.b64decode(unquote(encoded), validate=True) except ValueError as exc: raise ValueError("Agent Framework data content contains invalid base64 data.") from exc parts.append( diff --git a/python/packages/hosting-a2a/tests/hosting_a2a/test_conversion.py b/python/packages/hosting-a2a/tests/hosting_a2a/test_conversion.py index 8f5d143158a..2c7137241d8 100644 --- a/python/packages/hosting-a2a/tests/hosting_a2a/test_conversion.py +++ b/python/packages/hosting-a2a/tests/hosting_a2a/test_conversion.py @@ -17,7 +17,7 @@ ) from agent_framework.security import ContentLabel, IntegrityLabel, LabelTrackingFunctionMiddleware from google.protobuf.json_format import MessageToDict, ParseDict -from pytest import raises +from pytest import mark, raises from agent_framework_hosting_a2a import ( a2a_from_run, @@ -378,3 +378,28 @@ def test_a2a_from_workflow_run_rejects_pending_input() -> None: with raises(ValueError, match="requires external input"): a2a_from_workflow_run(result) + + +@mark.parametrize("streaming", [False, True]) +@mark.parametrize("payload", ["+/8=", "%2B%2F8%3D", "+%2F8="]) +def test_a2a_from_run_percent_encoded_data_uri(streaming: bool, payload: str) -> None: + uri = f"data:application/octet-stream;name=keep%2Fraw;base64,{payload}" + content = Content.from_uri(uri, additional_properties={"source": "binary"}) + result = ( + AgentResponseUpdate(role="assistant", contents=[content]) + if streaming + else AgentResponse(messages=[Message("assistant", [content])]) + ) + parts = a2a_from_run(result) + assert len(parts) == 1 + assert parts[0].raw == b"\xfb\xff" + assert parts[0].media_type == "application/octet-stream" + assert MessageToDict(parts[0].metadata) == {"source": "binary"} + assert content.uri == uri + + +@mark.parametrize("payload", ["%252B%252F8%253D", "+/8=\n"]) +def test_a2a_from_run_rejects_invalid_encoded_payload(payload: str) -> None: + content = Content.from_uri(f"data:application/octet-stream;base64,{payload}") + with raises(ValueError, match="invalid base64"): + a2a_from_run(Message("assistant", [content])) diff --git a/python/packages/hosting-mcp/agent_framework_hosting_mcp/_conversion.py b/python/packages/hosting-mcp/agent_framework_hosting_mcp/_conversion.py index 91d8381e8c2..677e254b370 100644 --- a/python/packages/hosting-mcp/agent_framework_hosting_mcp/_conversion.py +++ b/python/packages/hosting-mcp/agent_framework_hosting_mcp/_conversion.py @@ -8,6 +8,7 @@ import logging from collections.abc import Collection, Mapping, Sequence from typing import Any, cast +from urllib.parse import unquote from agent_framework import AgentResponse, ChatOptions, Content, Message from agent_framework_hosting import AgentRunArgs @@ -107,6 +108,7 @@ def mcp_from_run( prefix, separator, encoded = content.uri.partition(",") if not separator or not prefix.startswith("data:") or ";base64" not in prefix: raise ValueError("Agent Framework data content must contain a base64 data URI.") + encoded = unquote(encoded) try: base64.b64decode(encoded, validate=True) except ValueError as exc: diff --git a/python/packages/hosting-mcp/tests/hosting_mcp/test_conversion.py b/python/packages/hosting-mcp/tests/hosting_mcp/test_conversion.py index af133c7ec28..65cb18bf423 100644 --- a/python/packages/hosting-mcp/tests/hosting_mcp/test_conversion.py +++ b/python/packages/hosting-mcp/tests/hosting_mcp/test_conversion.py @@ -5,7 +5,7 @@ from agent_framework import AgentResponse, Content, Message from mcp import types from mcp.shared.exceptions import MCPError -from pytest import raises +from pytest import mark, raises from agent_framework_hosting_mcp import mcp_from_run, mcp_to_run @@ -153,3 +153,31 @@ def test_mcp_from_run_rejects_invalid_base64_data() -> None: with raises(ValueError, match="invalid base64"): mcp_from_run(Message("assistant", [content])) + + +@mark.parametrize("media_type", ["image/png", "audio/wav", "application/octet-stream"]) +@mark.parametrize("payload", ["+/8=", "%2B%2F8%3D", "+%2F8="]) +def test_mcp_from_run_percent_encoded_data_uri(media_type: str, payload: str) -> None: + uri = f"data:{media_type};name=keep%2Fraw;base64,{payload}" + content = Content.from_uri(uri, additional_properties={"source": "binary", "uri": "af://documents/data"}) + blocks = mcp_from_run(Message("assistant", [content])) + assert len(blocks) == 1 + block = blocks[0] + if isinstance(block, types.EmbeddedResource): + assert isinstance(block.resource, types.BlobResourceContents) + assert block.resource.blob == "+/8=" + assert str(block.resource.uri) == "af://documents/data" + assert block.resource.mime_type == media_type + else: + assert isinstance(block, (types.ImageContent, types.AudioContent)) + assert block.data == "+/8=" + assert block.mime_type == media_type + assert block.meta == {"source": "binary", "uri": "af://documents/data"} + assert content.uri == uri + + +@mark.parametrize("payload", ["%252B%252F8%253D", "+/8=\n"]) +def test_mcp_from_run_rejects_invalid_encoded_payload(payload: str) -> None: + content = Content.from_uri(f"data:image/png;base64,{payload}") + with raises(ValueError, match="invalid base64"): + mcp_from_run(Message("assistant", [content])) diff --git a/python/packages/ollama/agent_framework_ollama/_chat_client.py b/python/packages/ollama/agent_framework_ollama/_chat_client.py index 970840ed6e4..e1157bbef69 100644 --- a/python/packages/ollama/agent_framework_ollama/_chat_client.py +++ b/python/packages/ollama/agent_framework_ollama/_chat_client.py @@ -18,6 +18,7 @@ ) from itertools import chain from typing import Any, ClassVar, Generic, Literal, TypedDict +from urllib.parse import unquote from agent_framework import ( BaseChatClient, @@ -558,7 +559,7 @@ def _decode_image_data_uri(content: Content) -> bytes: raise ChatClientInvalidRequestException("Ollama image content must be a base64-encoded data:image URI.") try: - return base64.b64decode(payload, validate=True) + return base64.b64decode(unquote(payload), validate=True) except (binascii.Error, ValueError) as exc: raise ChatClientInvalidRequestException("Ollama image content contains invalid base64 data.") from exc diff --git a/python/packages/ollama/tests/test_ollama_chat_client.py b/python/packages/ollama/tests/test_ollama_chat_client.py index b404bf2701c..fd7709314f0 100644 --- a/python/packages/ollama/tests/test_ollama_chat_client.py +++ b/python/packages/ollama/tests/test_ollama_chat_client.py @@ -630,20 +630,25 @@ async def test_cmc_with_dict_tool_passthrough( assert call_kwargs["tools"] == [{"type": "function", "function": {"name": "custom_tool", "parameters": {}}}] +@pytest.mark.parametrize("escape_payload", [False, True]) @patch.object(AsyncClient, "chat", new_callable=AsyncMock) async def test_cmc_with_data_content_type( mock_chat: AsyncMock, ollama_unit_test_env: dict[str, str], chat_history: list[Message], mock_chat_completion_response: OllamaChatResponse, + escape_payload: bool, ) -> None: image_bytes = b"\x89PNG\r\n\x1a\n" + payload = base64.b64encode(image_bytes).decode() + if escape_payload: + payload = payload.replace("=", "%3D") mock_chat.return_value = mock_chat_completion_response chat_history.append( Message( contents=[ Content.from_uri( - uri=f"data:image/png;base64,{base64.b64encode(image_bytes).decode()}", + uri=f"data:image/png;base64,{payload}", media_type="image/png", ) ],