fix: download_file/download_files crash on binary content with UnicodeDecodeError (#257)

download_file() and download_files() crashed with UnicodeDecodeError
when the API returned binary content (PNG, JPEG, PDF) as base64-encoded
blobs, because .decode("utf-8") was called unconditionally on decoded
blob bytes.

Fix uses try/except to decode UTF-8 where possible (preserving str
return for backwards compatibility) and falls back to raw bytes for
true binary content. Return type annotations updated to
Union[str, bytes].

Add tests covering both the binary-returns-bytes and
utf8-blob-returns-str paths for download_file and download_files.

Co-authored-by: Kevin Orellana <keorel@amazon.com>

Author: kevin-orellana

src/bedrock_agentcore/tools/code_interpreter_client.py

@@ -600,14 +600,15 @@ def install_packages(
     def download_file(
         self,
         path: str,
-    ) -> str:
+    ) -> Union[str, bytes]:
         """Download/read a file from the code interpreter environment.
 
         Args:
             path: Path to the file to read.
 
         Returns:
-            File content as string.
+            File content as string, or bytes if the file contains binary content
+            (images, PDFs, etc.).
 
         Raises:
             FileNotFoundError: If the file doesn't exist.
@@ -631,21 +632,26 @@ def download_file(
                             if "text" in resource:
                                 return resource["text"]
                             elif "blob" in resource:
-                                return base64.b64decode(resource["blob"]).decode("utf-8")
+                                raw = base64.b64decode(resource["blob"])
+                                try:
+                                    return raw.decode("utf-8")
+                                except (UnicodeDecodeError, ValueError):
+                                    return raw
 
         raise FileNotFoundError(f"Could not read file: {path}")
 
     def download_files(
         self,
         paths: List[str],
-    ) -> Dict[str, str]:
+    ) -> Dict[str, Union[str, bytes]]:
         """Download/read multiple files from the code interpreter environment.
 
         Args:
             paths: List of file paths to read.
 
         Returns:
-            Dict mapping file paths to their contents.
+            Dict mapping file paths to their contents. Values are strings for
+            text files, or bytes for binary files (images, PDFs, etc.).
 
         Example:
             >>> files = client.download_files(['data.csv', 'results.json'])
@@ -667,7 +673,11 @@ def download_files(
                             if "text" in resource:
                                 files[file_path] = resource["text"]
                             elif "blob" in resource:
-                                files[file_path] = base64.b64decode(resource["blob"]).decode("utf-8")
+                                raw = base64.b64decode(resource["blob"])
+                                try:
+                                    files[file_path] = raw.decode("utf-8")
+                                except (UnicodeDecodeError, ValueError):
+                                    files[file_path] = raw
 
         return files
 

tests/bedrock_agentcore/tools/test_code_interpreter_client.py

@@ -1052,8 +1052,44 @@ def test_download_file_binary(self, mock_boto3, mock_get_data_endpoint, mock_get
         client.identifier = "test.identifier"
         client.session_id = "test-session-id"
 
-        original_content = "binary content as text"
-        encoded_content = base64.b64encode(original_content.encode("utf-8")).decode("utf-8")
+        binary_content = b"\x89PNG\r\n\x1a\n"  # PNG header bytes
+        encoded_content = base64.b64encode(binary_content).decode("utf-8")
+
+        mock_response = {
+            "stream": [
+                {
+                    "result": {
+                        "content": [
+                            {"type": "resource", "resource": {"uri": "file://image.png", "blob": encoded_content}}
+                        ]
+                    }
+                }
+            ]
+        }
+        client.data_plane_client.invoke_code_interpreter.return_value = mock_response
+
+        # Act
+        result = client.download_file("image.png")
+
+        # Assert
+        assert result == binary_content
+        assert isinstance(result, bytes)
+
+    @patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
+    @patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
+    @patch("bedrock_agentcore.tools.code_interpreter_client.boto3")
+    def test_download_file_blob_utf8_returns_str(self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint):
+        """Blob content that is valid UTF-8 should be decoded and returned as str."""
+        # Arrange
+        mock_session = MagicMock()
+        mock_session.client.return_value = MagicMock()
+        mock_boto3.Session.return_value = mock_session
+        client = CodeInterpreter("us-west-2")
+        client.identifier = "test.identifier"
+        client.session_id = "test-session-id"
+
+        text_content = "hello world"
+        encoded_content = base64.b64encode(text_content.encode("utf-8")).decode("utf-8")
 
         mock_response = {
             "stream": [
@@ -1072,7 +1108,8 @@ def test_download_file_binary(self, mock_boto3, mock_get_data_endpoint, mock_get
         result = client.download_file("data.bin")
 
         # Assert
-        assert result == original_content
+        assert result == text_content
+        assert isinstance(result, str)
 
     @patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
     @patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
@@ -1152,6 +1189,108 @@ def test_download_files_empty_result(self, mock_boto3, mock_get_data_endpoint, m
         # Assert
         assert result == {}
 
+    @patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
+    @patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
+    @patch("bedrock_agentcore.tools.code_interpreter_client.boto3")
+    def test_download_files_binary(self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint):
+        # Arrange
+        mock_session = MagicMock()
+        mock_session.client.return_value = MagicMock()
+        mock_boto3.Session.return_value = mock_session
+        client = CodeInterpreter("us-west-2")
+        client.identifier = "test.identifier"
+        client.session_id = "test-session-id"
+
+        binary_content = b"\x89PNG\r\n\x1a\n"  # PNG header bytes
+        encoded_binary = base64.b64encode(binary_content).decode("utf-8")
+
+        mock_response = {
+            "stream": [
+                {
+                    "result": {
+                        "content": [
+                            {
+                                "type": "resource",
+                                "resource": {
+                                    "uri": "file:///opt/amazon/genesis1p-tools/var/data.csv",
+                                    "text": "col1,col2\n1,2",
+                                },
+                            },
+                            {
+                                "type": "resource",
+                                "resource": {
+                                    "uri": "file:///opt/amazon/genesis1p-tools/var/chart.png",
+                                    "blob": encoded_binary,
+                                },
+                            },
+                        ]
+                    }
+                }
+            ]
+        }
+        client.data_plane_client.invoke_code_interpreter.return_value = mock_response
+
+        # Act
+        result = client.download_files(["data.csv", "chart.png"])
+
+        # Assert
+        assert result["/opt/amazon/genesis1p-tools/var/data.csv"] == "col1,col2\n1,2"
+        assert result["/opt/amazon/genesis1p-tools/var/chart.png"] == binary_content
+        assert isinstance(result["/opt/amazon/genesis1p-tools/var/chart.png"], bytes)
+
+    @patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
+    @patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
+    @patch("bedrock_agentcore.tools.code_interpreter_client.boto3")
+    def test_download_files_blob_utf8_returns_str(self, mock_boto3, mock_get_data_endpoint, mock_get_control_endpoint):
+        """Blob content that is valid UTF-8 should be decoded and returned as str in multi-file download."""
+        # Arrange
+        mock_session = MagicMock()
+        mock_session.client.return_value = MagicMock()
+        mock_boto3.Session.return_value = mock_session
+        client = CodeInterpreter("us-west-2")
+        client.identifier = "test.identifier"
+        client.session_id = "test-session-id"
+
+        text_content = "some utf-8 blob content"
+        encoded_text = base64.b64encode(text_content.encode("utf-8")).decode("utf-8")
+        binary_content = b"\x89PNG\r\n\x1a\n"
+        encoded_binary = base64.b64encode(binary_content).decode("utf-8")
+
+        mock_response = {
+            "stream": [
+                {
+                    "result": {
+                        "content": [
+                            {
+                                "type": "resource",
+                                "resource": {
+                                    "uri": "file:///opt/amazon/genesis1p-tools/var/data.bin",
+                                    "blob": encoded_text,
+                                },
+                            },
+                            {
+                                "type": "resource",
+                                "resource": {
+                                    "uri": "file:///opt/amazon/genesis1p-tools/var/chart.png",
+                                    "blob": encoded_binary,
+                                },
+                            },
+                        ]
+                    }
+                }
+            ]
+        }
+        client.data_plane_client.invoke_code_interpreter.return_value = mock_response
+
+        # Act
+        result = client.download_files(["data.bin", "chart.png"])
+
+        # Assert — UTF-8-valid blob comes back as str, invalid UTF-8 blob comes back as bytes
+        assert result["/opt/amazon/genesis1p-tools/var/data.bin"] == text_content
+        assert isinstance(result["/opt/amazon/genesis1p-tools/var/data.bin"], str)
+        assert result["/opt/amazon/genesis1p-tools/var/chart.png"] == binary_content
+        assert isinstance(result["/opt/amazon/genesis1p-tools/var/chart.png"], bytes)
+
     @patch("bedrock_agentcore.tools.code_interpreter_client.get_control_plane_endpoint")
     @patch("bedrock_agentcore.tools.code_interpreter_client.get_data_plane_endpoint")
     @patch("bedrock_agentcore.tools.code_interpreter_client.boto3")