feat(plugins): add file size validation and Files API routing for SaveFilesAsArtifactsPlugin with robust error handling and improved test mocking (#3781)

Author: AakashSuresh2003

src/google/adk/plugins/save_files_as_artifacts_plugin.py

@@ -16,6 +16,7 @@
 
 import copy
 import logging
+import mimetypes
 import os
 import tempfile
 from typing import Optional
@@ -92,7 +93,7 @@ async def on_user_message_callback(
       try:
         # Check file size before processing
         inline_data = part.inline_data
-        file_size = len(inline_data.data) if inline_data.data else 0
+        file_size = len(inline_data.data or b'')
 
         # Use display_name if available, otherwise generate a filename
         file_name = inline_data.display_name
@@ -110,7 +111,10 @@ async def on_user_message_callback(
           file_size_gb = file_size / (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' maximumFile {display_name} ({file_size_gb:.2f} GB) exceeds the'
+              ' maximum supported size of'
+              f' {_MAX_FILES_API_SIZE_BYTES / (1024*1024*1024):.0f}GB. Please'
+              ' upload a smaller file.'
           )
           logger.warning(error_message)
           new_parts.append(types.Part(text=f'[Upload Error: {error_message}]'))
@@ -121,8 +125,8 @@ async def on_user_message_callback(
         if file_size > _MAX_INLINE_DATA_SIZE_BYTES:
           file_size_mb = file_size / (1024 * 1024)
           logger.info(
-              f'File {display_name} ({file_size_mb:.2f} MB) exceeds'
-              ' inline_data limit. Uploading via Files API...'
+              f'File {display_name} ({file_size_mb:.2f} MB) exceeds inline_data'
+              ' limit. Uploading via Files API...'
           )
 
           # Upload to Files API and convert to file_data
@@ -214,16 +218,8 @@ async def _upload_to_files_api(
       if inline_data.display_name and '.' in inline_data.display_name:
         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, '')
+        # Use mimetypes for robust mime type to extension mapping
+        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

@@ -321,23 +321,22 @@ async def test_file_size_exceeds_limit(self):
     user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
 
     # Mock the Files API upload
-    with (
-        patch.object(Client, "__init__", return_value=None),
-        patch.object(Client, "files") as mock_files,
-    ):
+    with patch(
+        "google.adk.plugins.save_files_as_artifacts_plugin.Client"
+    ) as mock_client:
       # Mock uploaded file response
       mock_uploaded_file = MagicMock()
       mock_uploaded_file.uri = (
           "https://generativelanguage.googleapis.com/v1beta/files/test-file-id"
       )
-      mock_files.upload.return_value = mock_uploaded_file
+      mock_client.return_value.files.upload.return_value = mock_uploaded_file
 
       result = await self.plugin.on_user_message_callback(
           invocation_context=self.mock_context, user_message=user_message
       )
 
       # Should upload via Files API
-      mock_files.upload.assert_called_once()
+      mock_client.return_value.files.upload.assert_called_once()
 
       # Should save the artifact with file_data
       self.mock_context.artifact_service.save_artifact.assert_called_once()
@@ -387,7 +386,7 @@ async def test_file_size_just_over_limit(self):
     user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
 
     # Mock the Files API upload
-    with patch.object(Client, "files", create=True) as mock_files:
+    with patch.object(Client, "files") as mock_files:
       mock_uploaded_file = MagicMock()
       mock_uploaded_file.uri = (
           "https://generativelanguage.googleapis.com/v1beta/files/test-file-id"
@@ -434,7 +433,7 @@ async def test_mixed_file_sizes(self):
     )
 
     # Mock the Files API upload for large file
-    with patch.object(Client, "files", create=True) as mock_files:
+    with patch.object(Client, "files") as mock_files:
       mock_uploaded_file = MagicMock()
       mock_uploaded_file.uri = (
           "https://generativelanguage.googleapis.com/v1beta/files/test-file-id"
@@ -475,7 +474,7 @@ async def test_files_api_upload_failure(self):
     user_message = types.Content(parts=[types.Part(inline_data=inline_data)])
 
     # Mock the Files API to raise an exception
-    with patch.object(Client, "files", create=True) as mock_files:
+    with patch.object(Client, "files") as mock_files:
       mock_files.upload.side_effect = Exception("API quota exceeded")
 
       result = await self.plugin.on_user_message_callback(
@@ -498,57 +497,31 @@ async def test_files_api_upload_failure(self):
   @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
+    # 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 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}]")],
-            )
-      return await original_callback(
-          invocation_context=invocation_context, user_message=user_message
+    # 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
       )
 
-    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 return error message about 2GB limit
+    # 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
-    assert "2.00 GB" 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