fix: S3 upload and rename ordering in cached download path (#5115)

skyvern/services/script_service.py

@@ -18,7 +18,7 @@ from fastapi import BackgroundTasks, HTTPException
 from jinja2.sandbox import SandboxedEnvironment
 
 from skyvern.config import settings
-from skyvern.constants import GET_DOWNLOADED_FILES_TIMEOUT
+from skyvern.constants import GET_DOWNLOADED_FILES_TIMEOUT, SAVE_DOWNLOADED_FILES_TIMEOUT
 from skyvern.core.script_generations.constants import SCRIPT_TASK_BLOCKS
 from skyvern.core.script_generations.generate_script import _build_block_fn, create_or_update_script_block
 from skyvern.core.script_generations.script_skyvern_page import script_run_context_manager
@@ -1837,7 +1837,11 @@ async def download(
                     )
                 files_before_ok = True
             except asyncio.TimeoutError:
-                LOG.warning("Timeout getting downloaded files before cached download")
+                LOG.warning(
+                    "Timeout getting downloaded files before cached download",
+                    organization_id=org_id,
+                    workflow_run_id=run_id,
+                )
 
             # Track local files before download for renaming with download_suffix
             local_download_dir = get_path_for_workflow_download_directory(run_id)
@@ -1845,44 +1849,77 @@ async def download(
 
             await _run_cached_function(cached_fn)
 
-            # Rename newly downloaded files using download_suffix if provided
+            # Rename newly downloaded files using download_suffix if provided.
+            # Rename runs BEFORE S3 upload so that remote storage receives the
+            # correctly-named file and subsequent blocks get the right URLs.
+            # This matches the agent path ordering in agent.py.
             if download_suffix and local_download_dir.exists():
                 local_files_after = list_files_in_directory(local_download_dir)
                 new_files = list(set(local_files_after) - set(local_files_before))
-                for file in new_files:
-                    file_extension = Path(file).suffix
+                for file_path in new_files:
+                    file_extension = Path(file_path).suffix
                     # Skip incomplete downloads
                     if file_extension == ".crdownload":
                         continue
-                    target_path = local_download_dir / (download_suffix + file_extension)
-                    counter = 1
                     final_file_name = download_suffix
+                    target_path = local_download_dir / (final_file_name + file_extension)
+                    counter = 1
                     while target_path.exists():
                         final_file_name = f"{download_suffix}_{counter}"
                         target_path = local_download_dir / (final_file_name + file_extension)
                         counter += 1
-                    rename_file(str(local_download_dir / file), final_file_name + file_extension)
+                    rename_file(file_path, final_file_name + file_extension)
+
+            # Upload downloaded files from local filesystem to remote storage
+            # so that get_downloaded_files() can find them for verification.
+            save_ok = False
+            try:
+                async with asyncio.timeout(SAVE_DOWNLOADED_FILES_TIMEOUT):
+                    await app.STORAGE.save_downloaded_files(
+                        organization_id=org_id,
+                        run_id=run_id,
+                    )
+                save_ok = True
+            except asyncio.TimeoutError:
+                LOG.warning(
+                    "Timeout saving downloaded files after cached download, skipping verification",
+                    organization_id=org_id,
+                    workflow_run_id=run_id,
+                )
+            except Exception:
+                LOG.warning(
+                    "Failed to save downloaded files after cached download, skipping verification",
+                    exc_info=True,
+                    organization_id=org_id,
+                    workflow_run_id=run_id,
+                )
 
             # Verify a new file was actually downloaded.
             # Retry briefly — file may not be visible in storage immediately after the click.
+            # Skip entirely if save timed out — verification would fail and waste ~6s retrying.
             files_after: list = []
             files_after_ok = False
-            for _attempt in range(3):
-                try:
-                    async with asyncio.timeout(GET_DOWNLOADED_FILES_TIMEOUT):
-                        files_after = await app.STORAGE.get_downloaded_files(
+            if save_ok:
+                for _attempt in range(3):
+                    try:
+                        async with asyncio.timeout(GET_DOWNLOADED_FILES_TIMEOUT):
+                            files_after = await app.STORAGE.get_downloaded_files(
+                                organization_id=org_id,
+                                run_id=run_id,
+                            )
+                        files_after_ok = True
+                    except asyncio.TimeoutError:
+                        LOG.warning(
+                            "Timeout getting downloaded files after cached download",
                             organization_id=org_id,
-                            run_id=run_id,
+                            workflow_run_id=run_id,
                         )
-                    files_after_ok = True
-                except asyncio.TimeoutError:
-                    LOG.warning("Timeout getting downloaded files after cached download")
-                if len(files_after) > len(files_before):
-                    break
-                if _attempt < 2:
-                    await asyncio.sleep(2)
+                    if len(files_after) > len(files_before):
+                        break
+                    if _attempt < 2:
+                        await asyncio.sleep(2)
 
-            # Only raise if both calls succeeded — if either timed out, skip
+            # Only raise if all storage calls succeeded — if any timed out, skip
             # the check to avoid spurious AI fallbacks under degraded storage.
             if files_before_ok and files_after_ok and len(files_after) <= len(files_before):
                 raise Exception(

tests/unit/test_cached_download_s3_upload.py

@@ -0,0 +1,375 @@
+"""Tests for the cached download path in script_service.download().
+
+Validates that the cached download flow:
+1. Uploads files to remote storage (save_downloaded_files) so verification works
+2. Renames files with download_suffix BEFORE the S3 upload
+3. Verifies the download produced new files via get_downloaded_files
+4. Falls back to AI on verification failure
+5. Handles timeouts gracefully
+"""
+
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+MODULE = "skyvern.services.script_service"
+
+
+def _make_mock_app(storage):
+    """Create a mock that replaces the `app` module-level reference in script_service."""
+    mock_app = MagicMock()
+    mock_app.STORAGE = storage
+    return mock_app
+
+
+def _make_storage(get_side_effect=None):
+    """Create a mock storage with async save/get methods."""
+    s = MagicMock()
+    s.save_downloaded_files = AsyncMock()
+    s.get_downloaded_files = AsyncMock(side_effect=get_side_effect or [[], ["file.pdf"]])
+    return s
+
+
+@pytest.fixture()
+def mock_context():
+    ctx = MagicMock()
+    ctx.organization_id = "o_test_org"
+    ctx.workflow_run_id = "wr_test_run"
+    ctx.prompt = None
+    return ctx
+
+
+@pytest.fixture()
+def setup(mock_context, tmp_path):
+    """Provide a helper that configures all mocks for the download() function.
+
+    Returns a callable that accepts optional overrides for storage/rename/list_files.
+    """
+    download_dir = tmp_path / "downloads"
+    download_dir.mkdir()
+
+    def _setup(
+        get_side_effect=None,
+        save_side_effect=None,
+        list_files_side_effect=None,
+        rename_mock=None,
+    ):
+        storage = _make_storage(get_side_effect)
+        if save_side_effect is not None:
+            storage.save_downloaded_files = AsyncMock(side_effect=save_side_effect)
+        mock_app = _make_mock_app(storage)
+
+        list_side = list_files_side_effect if list_files_side_effect is not None else [[]]
+        rename = rename_mock or MagicMock()
+
+        fallback_mock = AsyncMock()
+        update_block_mock = AsyncMock()
+
+        all_patches = [
+            patch(f"{MODULE}.app", mock_app),
+            patch(f"{MODULE}.script_run_context_manager.get_cached_fn", return_value=AsyncMock()),
+            patch(
+                f"{MODULE}._create_workflow_block_run_and_task",
+                new_callable=AsyncMock,
+                return_value=("wrb_1", "tsk_1", "stp_1"),
+            ),
+            patch(f"{MODULE}._render_template_with_label", side_effect=lambda p, _: p),
+            patch(f"{MODULE}.skyvern_context.ensure_context", return_value=mock_context),
+            patch(f"{MODULE}._prepare_cached_block_inputs", new_callable=AsyncMock),
+            patch(f"{MODULE}._run_cached_function", new_callable=AsyncMock),
+            patch(f"{MODULE}._update_workflow_block", update_block_mock),
+            patch(f"{MODULE}._fallback_to_ai_run", fallback_mock),
+            patch(f"{MODULE}._clear_cached_block_overrides"),
+            patch(f"{MODULE}.get_path_for_workflow_download_directory", return_value=download_dir),
+            patch(f"{MODULE}.list_files_in_directory", side_effect=list_side),
+            patch(f"{MODULE}.rename_file", rename),
+        ]
+
+        for p in all_patches:
+            p.start()
+
+        return {
+            "storage": storage,
+            "app": mock_app,
+            "patches": all_patches,
+            "download_dir": download_dir,
+            "fallback": fallback_mock,
+            "update_block": update_block_mock,
+            "rename": rename,
+        }
+
+    return _setup
+
+
+def _cleanup(refs):
+    for p in refs["patches"]:
+        p.stop()
+
+
+@pytest.mark.asyncio
+async def test_cached_download_calls_save_downloaded_files(setup):
+    """save_downloaded_files must be called so get_downloaded_files can find the file."""
+    refs = setup(get_side_effect=[[], ["file1.pdf"]])
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", label="test_block")
+
+        refs["storage"].save_downloaded_files.assert_called_once_with(
+            organization_id="o_test_org",
+            run_id="wr_test_run",
+        )
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_rename_happens_before_s3_upload(setup, tmp_path):
+    """download_suffix rename must happen BEFORE save_downloaded_files.
+
+    This ensures remote storage receives the correctly-named file and subsequent
+    blocks get the right URLs. Matches the agent path ordering in agent.py.
+    """
+    download_dir = tmp_path / "downloads"
+    call_order: list[str] = []
+
+    rename_mock = MagicMock(side_effect=lambda path, name: call_order.append("rename"))
+
+    async def track_save(**kwargs):
+        call_order.append("save")
+
+    fake_file = str(download_dir / "uuid-random.pdf")
+
+    refs = setup(
+        get_side_effect=[[], ["invoice.pdf"]],
+        save_side_effect=track_save,
+        list_files_side_effect=[[], [fake_file]],
+        rename_mock=rename_mock,
+    )
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", download_suffix="invoice", label="test_block")
+
+        assert call_order == ["rename", "save"], f"Expected rename before save, got: {call_order}"
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_download_suffix_rename_uses_file_path_directly(setup, tmp_path):
+    """rename_file should receive the absolute path from list_files_in_directory,
+    not a reconstructed path via Path joining."""
+    download_dir = tmp_path / "downloads"
+    abs_path = str(download_dir / "abc123.pdf")
+
+    rename_mock = MagicMock(return_value=str(download_dir / "invoice.pdf"))
+    refs = setup(
+        get_side_effect=[[], ["invoice.pdf"]],
+        list_files_side_effect=[[], [abs_path]],
+        rename_mock=rename_mock,
+    )
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", download_suffix="invoice", label="test_block")
+
+        rename_mock.assert_called_once_with(abs_path, "invoice.pdf")
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_verification_raises_when_no_new_file(setup):
+    """When get_downloaded_files shows no increase, the cached path should raise
+    and fall back to AI."""
+    refs = setup(get_side_effect=[[], [], [], []])
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", label="test_block")
+
+        refs["fallback"].assert_called_once()
+        error_arg = refs["fallback"].call_args.kwargs.get("error")
+        assert "Files before: 0, after: 0" in str(error_arg)
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_verification_retries_before_failing(setup):
+    """get_downloaded_files retries up to 3 times before declaring failure."""
+    # Before: 0 files. After: [], [], ["file.pdf"] (succeeds on 3rd attempt)
+    refs = setup(get_side_effect=[[], [], [], ["file.pdf"]])
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", label="test_block")
+
+        refs["fallback"].assert_not_called()
+        assert refs["storage"].get_downloaded_files.call_count == 4  # 1 before + 3 retries
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_save_timeout_skips_verification(setup):
+    """TimeoutError on save_downloaded_files should skip verification entirely.
+    No point retrying get_downloaded_files when we know S3 is degraded."""
+    refs = setup(
+        save_side_effect=asyncio.TimeoutError(),
+        get_side_effect=[[]],  # only the before-check runs; after-check is skipped
+    )
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", label="test_block")
+
+        # Verification skipped → block completes, no AI fallback
+        refs["fallback"].assert_not_called()
+        refs["update_block"].assert_called_once()
+        # get_downloaded_files called only once (before-check), not 3 more times for after
+        assert refs["storage"].get_downloaded_files.call_count == 1
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_save_generic_exception_skips_verification(setup):
+    """Non-timeout S3 failure (e.g., permission error) should also skip verification.
+    Matches agent.py which catches both TimeoutError and generic Exception."""
+    refs = setup(
+        save_side_effect=RuntimeError("S3 permission denied"),
+        get_side_effect=[[]],  # only the before-check runs; after-check is skipped
+    )
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", label="test_block")
+
+        # Verification skipped → block completes, no AI fallback
+        refs["fallback"].assert_not_called()
+        refs["update_block"].assert_called_once()
+        # get_downloaded_files called only once (before-check), not 3 more times for after
+        assert refs["storage"].get_downloaded_files.call_count == 1
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_get_before_timeout_skips_verification(setup):
+    """If the before-check times out, verification should be skipped entirely
+    to avoid spurious AI fallbacks under degraded storage."""
+    refs = setup(get_side_effect=asyncio.TimeoutError())
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", label="test_block")
+
+        # Should NOT fall back — verification was skipped
+        refs["fallback"].assert_not_called()
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_get_after_timeout_skips_verification(setup):
+    """If the after-check times out, verification should be skipped."""
+    call_count = {"n": 0}
+
+    async def get_side_effect(**kwargs):
+        call_count["n"] += 1
+        if call_count["n"] == 1:
+            return []  # before-check succeeds
+        raise asyncio.TimeoutError()
+
+    refs = setup(get_side_effect=get_side_effect)
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", label="test_block")
+
+        refs["fallback"].assert_not_called()
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_no_rename_without_download_suffix(setup):
+    """When download_suffix is not provided, rename should not be called."""
+    rename_mock = MagicMock()
+    refs = setup(get_side_effect=[[], ["file.pdf"]], rename_mock=rename_mock)
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", label="test_block")
+
+        rename_mock.assert_not_called()
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_rename_skips_crdownload_files(setup, tmp_path):
+    """Files with .crdownload extension (incomplete downloads) should be skipped."""
+    download_dir = tmp_path / "downloads"
+    incomplete = str(download_dir / "file.crdownload")
+    complete = str(download_dir / "invoice.pdf")
+
+    rename_mock = MagicMock(return_value=str(download_dir / "renamed.pdf"))
+    refs = setup(
+        get_side_effect=[[], ["invoice.pdf"]],
+        list_files_side_effect=[[], [incomplete, complete]],
+        rename_mock=rename_mock,
+    )
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", download_suffix="renamed", label="test_block")
+
+        rename_mock.assert_called_once_with(complete, "renamed.pdf")
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_rename_handles_name_collision(setup, tmp_path):
+    """When target filename already exists, a counter suffix should be added."""
+    download_dir = tmp_path / "downloads"
+    download_dir.mkdir(exist_ok=True)
+    new_file = str(download_dir / "uuid.pdf")
+    # Create collision: "invoice.pdf" already exists
+    (download_dir / "invoice.pdf").touch()
+
+    rename_mock = MagicMock(return_value=str(download_dir / "invoice_1.pdf"))
+    refs = setup(
+        get_side_effect=[[], ["invoice_1.pdf"]],
+        list_files_side_effect=[[], [new_file]],
+        rename_mock=rename_mock,
+    )
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", download_suffix="invoice", label="test_block")
+
+        rename_mock.assert_called_once_with(new_file, "invoice_1.pdf")
+    finally:
+        _cleanup(refs)
+
+
+@pytest.mark.asyncio
+async def test_block_marked_completed_on_success(setup):
+    """On successful download + verification, block should be marked completed."""
+    refs = setup(get_side_effect=[[], ["file.pdf"]])
+    try:
+        from skyvern.services.script_service import download
+
+        await download(prompt="Download invoice", label="test_block")
+
+        refs["update_block"].assert_called_once()
+        refs["fallback"].assert_not_called()
+    finally:
+        _cleanup(refs)