From 67f9aa9fe864d8ced8dc34c4ee0659f6e038cb1c Mon Sep 17 00:00:00 2001 From: VectorPeak Date: Sun, 5 Jul 2026 10:13:14 +0800 Subject: [PATCH 1/2] Python: raise ValueError for malformed data URIs --- python/packages/core/agent_framework/_types.py | 2 ++ python/packages/core/tests/core/test_types.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 23926658fb5..71d9bdb2e54 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -122,6 +122,8 @@ def detect_media_type_from_base64( if data is not None: raise ValueError("Provide exactly one of data_bytes, data_str, or data_uri.") # Remove data URI prefix if present + if ";base64," not in data_uri: + raise ValueError("Invalid data URI format.") data_str = data_uri.split(";base64,", 1)[1] if data_str is not None: if data is not None: diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 19c77585e48..3be67692657 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -208,6 +208,20 @@ def test_data_content_detect_image_format_from_base64(): detect_media_type_from_base64(data_str="data", data_uri="data:application/octet-stream;base64,AAA") +@pytest.mark.parametrize( + "data_uri", + [ + "data:text/plain,hello", + "data:image/png;base64", + "not-a-data-uri", + ], +) +def test_detect_media_type_from_base64_rejects_malformed_data_uri(data_uri: str): + """Test malformed data URI inputs raise the documented ValueError.""" + with pytest.raises(ValueError, match="Invalid data URI format."): + detect_media_type_from_base64(data_uri=data_uri) + + def test_data_content_create_data_uri_from_base64(): """Test the create_data_uri_from_base64 class method.""" # Test with PNG data From 630243a82b9f0dc4e2b7a5819f41a95dfdaffd79 Mon Sep 17 00:00:00 2001 From: VectorPeak Date: Sun, 5 Jul 2026 10:26:22 +0800 Subject: [PATCH 2/2] Address malformed data URI review feedback --- python/packages/core/agent_framework/_types.py | 6 ++++-- python/packages/core/tests/core/test_types.py | 14 +++++++------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 71d9bdb2e54..d5f7d14d925 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -122,9 +122,11 @@ def detect_media_type_from_base64( if data is not None: raise ValueError("Provide exactly one of data_bytes, data_str, or data_uri.") # Remove data URI prefix if present - if ";base64," not in data_uri: + if not data_uri.startswith("data:") or "," not in data_uri: raise ValueError("Invalid data URI format.") - data_str = data_uri.split(";base64,", 1)[1] + prefix, data_str = data_uri.split(",", 1) + if not prefix.endswith(";base64"): + raise ValueError("Data URI must use base64 encoding.") if data_str is not None: if data is not None: raise ValueError("Provide exactly one of data_bytes, data_str, or data_uri.") diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 3be67692657..9b31dee2385 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -208,17 +208,17 @@ def test_data_content_detect_image_format_from_base64(): detect_media_type_from_base64(data_str="data", data_uri="data:application/octet-stream;base64,AAA") -@pytest.mark.parametrize( - "data_uri", +@mark.parametrize( + ("data_uri", "error_message"), [ - "data:text/plain,hello", - "data:image/png;base64", - "not-a-data-uri", + ("data:text/plain,hello", "Data URI must use base64 encoding."), + ("data:image/png;base64", "Invalid data URI format."), + ("not-a-data-uri", "Invalid data URI format."), ], ) -def test_detect_media_type_from_base64_rejects_malformed_data_uri(data_uri: str): +def test_detect_media_type_from_base64_rejects_malformed_data_uri(data_uri: str, error_message: str): """Test malformed data URI inputs raise the documented ValueError.""" - with pytest.raises(ValueError, match="Invalid data URI format."): + with raises(ValueError, match=error_message): detect_media_type_from_base64(data_uri=data_uri)