Patch: Fix file size limit, MIME type uses library (#3781)

Author: AakashSuresh2003

src/google/adk/plugins/save_files_as_artifacts_plugin.py

@@ -26,6 +26,7 @@
 
 from ..agents.invocation_context import InvocationContext
 from .base_plugin import BasePlugin
+import mimetypes
 
 logger = logging.getLogger('google_adk.' + __name__)
 
@@ -108,9 +109,10 @@ async def on_user_message_callback(
         # Check if file exceeds Files API limit (2GB)
         if file_size > _MAX_FILES_API_SIZE_BYTES:
           file_size_gb = file_size / (1024 * 1024 * 1024)
+          max_size_gb = _MAX_FILES_API_SIZE_BYTES / (1024 * 1024 * 1024)
           error_message = (
               f'File {display_name} ({file_size_gb:.2f} GB) exceeds the'
-              ' maximum supported size of 2GB. Please upload a smaller file.'
+              f' maximum supported size of {max_size_gb:.2f} GB. Please upload a smaller file.'
           )
           logger.warning(error_message)
           new_parts.append(types.Part(text=f'[Upload Error: {error_message}]'))
@@ -215,15 +217,7 @@ async def _upload_to_files_api(
         file_extension = os.path.splitext(inline_data.display_name)[1]
       elif inline_data.mime_type:
         # Simple mime type to extension mapping
-        mime_to_ext = {
-            'application/pdf': '.pdf',
-            'image/png': '.png',
-            'image/jpeg': '.jpg',
-            'image/gif': '.gif',
-            'text/plain': '.txt',
-            'application/json': '.json',
-        }
-        file_extension = mime_to_ext.get(inline_data.mime_type, '')
+        file_extension = mimetypes.guess_extension(inline_data.mime_type) or ''
 
       # Create temporary file
       with tempfile.NamedTemporaryFile(

tests/unittests/plugins/test_save_files_as_artifacts.py

@@ -495,60 +495,33 @@ async def test_files_api_upload_failure(self):
       assert "huge_file.pdf" in result.parts[0].text
       assert "API quota exceeded" in result.parts[0].text
 
-  @pytest.mark.asyncio
-  async def test_file_exceeds_files_api_limit(self):
-    """Test that files exceeding 2GB limit are rejected with clear error."""
-    # Create a file larger than 2GB (simulated with a descriptor that reports large size)
-    # Create a mock object that behaves like bytes but reports 2GB+ size
-    large_data = b"x" * 1000  # Small actual data for testing
-
-    # Create inline_data with the small data
-    inline_data = types.Blob(
-        display_name="huge_video.mp4",
-        data=large_data,
-        mime_type="video/mp4",
-    )
-
-    user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
-
-    # Patch the file size check to simulate a 2GB+ file
-    original_callback = self.plugin.on_user_message_callback
-
-    async def patched_callback(*, invocation_context, user_message):
-      # Temporarily replace the data length check
-      for part in user_message.parts:
-        if part.inline_data:
-          # Simulate 2GB + 1 byte size
-          file_size_over_limit = (2 * 1024 * 1024 * 1024) + 1
-          # Manually inject the check that would happen in the real code
-          if file_size_over_limit > (2 * 1024 * 1024 * 1024):
-            file_size_gb = file_size_over_limit / (1024 * 1024 * 1024)
-            display_name = part.inline_data.display_name or "unknown"
-            error_message = (
-                f"File {display_name} ({file_size_gb:.2f} GB) exceeds the"
-                " maximum supported size of 2GB. Please upload a smaller file."
-            )
-            return types.Content(
-                role="user",
-                parts=[types.Part(text=f"[Upload Error: {error_message}]")],
+    @pytest.mark.asyncio
+    async def test_file_exceeds_files_api_limit(self):
+        """Test that files exceeding 2GB limit are rejected with clear error."""
+        # Use a small file for the test
+        large_data = b"x" * 1000
+        inline_data = types.Blob(
+            display_name="huge_video.mp4",
+            data=large_data,
+            mime_type="video/mp4",
+        )
+        user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
+
+        # Patch the size limit to be smaller than the file data
+        with patch(
+            "google.adk.plugins.save_files_as_artifacts_plugin._MAX_FILES_API_SIZE_BYTES", 500
+        ):
+            result = await self.plugin.on_user_message_callback(
+                invocation_context=self.mock_context, user_message=user_message
             )
-      return await original_callback(
-          invocation_context=invocation_context, user_message=user_message
-      )
-
-    self.plugin.on_user_message_callback = patched_callback
-
-    result = await self.plugin.on_user_message_callback(
-        invocation_context=self.mock_context, user_message=user_message
-    )
 
-    # Should not attempt any upload
-    self.mock_context.artifact_service.save_artifact.assert_not_called()
+        # Should not attempt any upload
+        self.mock_context.artifact_service.save_artifact.assert_not_called()
 
-    # Should return error message about 2GB limit
-    assert result is not None
-    assert len(result.parts) == 1
-    assert "[Upload Error:" in result.parts[0].text
-    assert "huge_video.mp4" in result.parts[0].text
-    assert "2.00 GB" in result.parts[0].text
-    assert "exceeds the maximum supported size" in result.parts[0].text
+        # Should return error message about the limit
+        assert result is not None
+        assert len(result.parts) == 1
+        assert "[Upload Error:" in result.parts[0].text
+        assert "huge_video.mp4" in result.parts[0].text
+        # Note: This assertion will depend on fixing the hardcoded "2GB" in the error message.
+        assert "exceeds the maximum supported size" in result.parts[0].text