fix(file): clean up failed cloud downloads

src/backend/tests/unit/components/files_and_knowledge/test_file_component.py

@@ -862,7 +862,7 @@ def test_output_has_tool_mode_true(self):
 
     @patch("lfx.base.data.cloud_storage_utils.create_s3_client")
     @patch("lfx.base.data.cloud_storage_utils.validate_aws_credentials")
-    def test_s3_temp_file_cleanup_on_download_failure(self, mock_validate, mock_create_client):  # noqa: ARG002
+    def test_s3_temp_file_cleanup_on_download_failure(self, mock_validate, mock_create_client, tmp_path):  # noqa: ARG002
         """Test that temp file is cleaned up when S3 download fails."""
         from pathlib import Path
 
@@ -877,24 +877,37 @@ def test_s3_temp_file_cleanup_on_download_failure(self, mock_validate, mock_crea
             }
         )
 
-        # Mock S3 client to raise an exception during download
         mock_s3_client = MagicMock()
         mock_s3_client.download_fileobj.side_effect = Exception("S3 download failed")
         mock_create_client.return_value = mock_s3_client
 
-        # Track temp files created
-        temp_dir = Path(tempfile.gettempdir())
-        temp_files_before = set(temp_dir.glob("tmp*.txt")) if temp_dir.exists() else set()
-
-        # Attempt to read from S3 - should fail and clean up temp file
-        with pytest.raises(RuntimeError, match="Failed to download file from S3"):
+        temp_state = {"handle": None, "path": None}
+        original_named_temporary_file = tempfile.NamedTemporaryFile
+        original_unlink = Path.unlink
+
+        def tracked_named_temporary_file(*args, **kwargs):
+            kwargs.setdefault("dir", tmp_path)
+            kwargs["delete"] = False
+            handle = original_named_temporary_file(*args, **kwargs)
+            temp_state["handle"] = handle
+            temp_state["path"] = Path(handle.name)
+            return handle
+
+        def guarded_unlink(self, *args, **kwargs):
+            if temp_state["path"] == self and temp_state["handle"] is not None and not temp_state["handle"].closed:
+                msg = "temp file still open"
+                raise PermissionError(msg)
+            return original_unlink(self, *args, **kwargs)
+
+        with (
+            patch("tempfile.NamedTemporaryFile", side_effect=tracked_named_temporary_file),
+            patch.object(Path, "unlink", autospec=True, side_effect=guarded_unlink),
+            pytest.raises(RuntimeError, match="Failed to download file from S3"),
+        ):
             component._read_from_aws_s3()
 
-        # Verify no new temp files are left behind
-        with tempfile.TemporaryDirectory() as temp_dir:
-            temp_files_after = set(Path(temp_dir).glob("tmp*.txt"))
-        new_temp_files = temp_files_after - temp_files_before
-        assert len(new_temp_files) == 0, f"Temp files not cleaned up: {new_temp_files}"
+        assert temp_state["path"] is not None
+        assert not temp_state["path"].exists()
 
     @patch("lfx.base.data.cloud_storage_utils.create_s3_client")
     @patch("lfx.base.data.cloud_storage_utils.validate_aws_credentials")
@@ -923,7 +936,7 @@ def test_s3_download_uses_explicit_local_cleanup(self, mock_validate, mock_creat
 
     @patch("lfx.base.data.cloud_storage_utils.create_google_drive_service")
     @pytest.mark.usefixtures("fake_googleapiclient")
-    def test_google_drive_temp_file_cleanup_on_download_failure(self, mock_create_service):
+    def test_google_drive_temp_file_cleanup_on_download_failure(self, mock_create_service, tmp_path):
         """Test that temp file is cleaned up when Google Drive download fails."""
         from pathlib import Path
 
@@ -936,29 +949,38 @@ def test_google_drive_temp_file_cleanup_on_download_failure(self, mock_create_se
             }
         )
 
-        # Mock Google Drive service
         mock_drive_service = MagicMock()
-        # Metadata call succeeds
         mock_drive_service.files().get().execute.return_value = {"name": "test-file.txt"}
-        # Media download fails
         mock_drive_service.files().get_media.side_effect = Exception("Drive download failed")
         mock_create_service.return_value = mock_drive_service
 
-        # Track temp files created
-        temp_files_before = (
-            set(Path(tempfile.gettempdir()).glob("tmp*.txt")) if Path(tempfile.gettempdir()).exists() else set()
-        )
-
-        # Attempt to read from Google Drive - should fail and clean up temp file
-        with pytest.raises(RuntimeError, match="Failed to download file from Google Drive"):
+        temp_state = {"handle": None, "path": None}
+        original_named_temporary_file = tempfile.NamedTemporaryFile
+        original_unlink = Path.unlink
+
+        def tracked_named_temporary_file(*args, **kwargs):
+            kwargs.setdefault("dir", tmp_path)
+            kwargs["delete"] = False
+            handle = original_named_temporary_file(*args, **kwargs)
+            temp_state["handle"] = handle
+            temp_state["path"] = Path(handle.name)
+            return handle
+
+        def guarded_unlink(self, *args, **kwargs):
+            if temp_state["path"] == self and temp_state["handle"] is not None and not temp_state["handle"].closed:
+                msg = "temp file still open"
+                raise PermissionError(msg)
+            return original_unlink(self, *args, **kwargs)
+
+        with (
+            patch("tempfile.NamedTemporaryFile", side_effect=tracked_named_temporary_file),
+            patch.object(Path, "unlink", autospec=True, side_effect=guarded_unlink),
+            pytest.raises(RuntimeError, match="Failed to download file from Google Drive"),
+        ):
             component._read_from_google_drive()
 
-        # Verify no new temp files are left behind
-        temp_files_after = (
-            set(Path(tempfile.gettempdir()).glob("tmp*.txt")) if Path(tempfile.gettempdir()).exists() else set()
-        )
-        new_temp_files = temp_files_after - temp_files_before
-        assert len(new_temp_files) == 0, f"Temp files not cleaned up: {new_temp_files}"
+        assert temp_state["path"] is not None
+        assert not temp_state["path"].exists()
 
     @patch("googleapiclient.http.MediaIoBaseDownload")
     @patch("lfx.base.data.cloud_storage_utils.create_google_drive_service")

src/lfx/src/lfx/components/files_and_knowledge/file.py

@@ -815,17 +815,20 @@ def _read_from_aws_s3(self) -> list[BaseFileComponent.BaseFile]:
 
         # Get file extension from S3 key
         file_extension = Path(self.s3_file_key).suffix or ""
-
+        temp_file_path = None
+        download_error = None
         with tempfile.NamedTemporaryFile(mode="wb", suffix=file_extension, delete=False) as temp_file:
             temp_file_path = temp_file.name
             try:
                 s3_client.download_fileobj(self.bucket_name, self.s3_file_key, temp_file)
-            except Exception as e:
-                # Clean up temp file on failure
-                with contextlib.suppress(OSError):
-                    Path(temp_file_path).unlink()
-                msg = f"Failed to download file from S3: {e}"
-                raise RuntimeError(msg) from e
+            except Exception as e:  # noqa: BLE001
+                download_error = e
+
+        if download_error is not None:
+            with contextlib.suppress(OSError):
+                Path(temp_file_path).unlink()
+            msg = f"Failed to download file from S3: {download_error}"
+            raise RuntimeError(msg) from download_error
 
         # Create BaseFile object
         from lfx.schema.data import Data
@@ -877,6 +880,8 @@ def _read_from_google_drive(self) -> list[BaseFileComponent.BaseFile]:
 
         # Download file to temp location
         file_extension = Path(file_name).suffix or ""
+        temp_file_path = None
+        download_error = None
         with tempfile.NamedTemporaryFile(mode="wb", suffix=file_extension, delete=False) as temp_file:
             temp_file_path = temp_file.name
             try:
@@ -885,12 +890,14 @@ def _read_from_google_drive(self) -> list[BaseFileComponent.BaseFile]:
                 done = False
                 while not done:
                     _status, done = downloader.next_chunk()
-            except Exception as e:
-                # Clean up temp file on failure
-                with contextlib.suppress(OSError):
-                    Path(temp_file_path).unlink()
-                msg = f"Failed to download file from Google Drive: {e}"
-                raise RuntimeError(msg) from e
+            except Exception as e:  # noqa: BLE001
+                download_error = e
+
+        if download_error is not None:
+            with contextlib.suppress(OSError):
+                Path(temp_file_path).unlink()
+            msg = f"Failed to download file from Google Drive: {download_error}"
+            raise RuntimeError(msg) from download_error
 
         # Create BaseFile object
         from lfx.schema.data import Data