From bd68e1af3dd4c4556f286d39d40e0463995220ba Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Mon, 16 Mar 2026 11:01:04 -0700 Subject: [PATCH] fix: S3 upload and rename ordering in cached download path (#5115) --- skyvern/services/script_service.py | 79 ++-- tests/unit/test_cached_download_s3_upload.py | 375 +++++++++++++++++++ 2 files changed, 433 insertions(+), 21 deletions(-) create mode 100644 tests/unit/test_cached_download_s3_upload.py diff --git a/skyvern/services/script_service.py b/skyvern/services/script_service.py index 4eb904e64..aa2937b67 100644 --- a/skyvern/services/script_service.py +++ b/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( diff --git a/tests/unit/test_cached_download_s3_upload.py b/tests/unit/test_cached_download_s3_upload.py new file mode 100644 index 000000000..633d53856 --- /dev/null +++ b/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)