From 6117e992b7f4c114c8d4a0a33ae9bcac49f7a7c6 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Fri, 8 May 2026 01:23:52 -0700 Subject: [PATCH] feat: expose file size in artifact file info (#5894) --- ...78a8db531e69_add_file_size_to_artifacts.py | 29 ++++++++++ .../browser-automations/work-with-files.mdx | 1 + fern/running-tasks/api-spec.mdx | 4 +- skyvern/forge/sdk/api/real_azure.py | 3 +- skyvern/forge/sdk/artifact/manager.py | 29 ++++++++++ skyvern/forge/sdk/artifact/models.py | 1 + skyvern/forge/sdk/artifact/storage/azure.py | 58 ++++++++++++++----- skyvern/forge/sdk/artifact/storage/base.py | 1 + skyvern/forge/sdk/artifact/storage/local.py | 23 ++++++-- skyvern/forge/sdk/artifact/storage/s3.py | 30 +++++++++- skyvern/forge/sdk/db/models.py | 1 + .../forge/sdk/db/repositories/artifacts.py | 2 + skyvern/forge/sdk/db/utils.py | 1 + skyvern/forge/sdk/schemas/files.py | 1 + ...test_browser_session_download_artifacts.py | 16 ++++- ...est_browser_session_recording_artifacts.py | 15 ++++- .../test_downloaded_files_artifact_urls.py | 35 +++++++++-- 17 files changed, 216 insertions(+), 34 deletions(-) create mode 100644 alembic/versions/2026_05_08_0815-78a8db531e69_add_file_size_to_artifacts.py diff --git a/alembic/versions/2026_05_08_0815-78a8db531e69_add_file_size_to_artifacts.py b/alembic/versions/2026_05_08_0815-78a8db531e69_add_file_size_to_artifacts.py new file mode 100644 index 000000000..26c227384 --- /dev/null +++ b/alembic/versions/2026_05_08_0815-78a8db531e69_add_file_size_to_artifacts.py @@ -0,0 +1,29 @@ +"""add file size to artifacts + +Revision ID: 78a8db531e69 +Revises: 289f2d72b0c2 +Create Date: 2026-05-08T08:15:20.233046+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "78a8db531e69" +down_revision: Union[str, None] = "289f2d72b0c2" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute("SET lock_timeout = '5s'") + op.add_column("artifacts", sa.Column("file_size", sa.BigInteger(), nullable=True)) + + +def downgrade() -> None: + op.execute("SET lock_timeout = '5s'") + op.drop_column("artifacts", "file_size") diff --git a/docs/developers/browser-automations/work-with-files.mdx b/docs/developers/browser-automations/work-with-files.mdx index b8fd4ac86..11ef5c280 100644 --- a/docs/developers/browser-automations/work-with-files.mdx +++ b/docs/developers/browser-automations/work-with-files.mdx @@ -111,6 +111,7 @@ Each file in `downloaded_files` has: | `url` | `str` | Presigned URL to download the file. Valid for a limited time. | | `filename` | `str \| None` | Saved filename. Reflects `download_suffix` if you set one, otherwise system-generated. | | `checksum` | `str \| None` | SHA-256 checksum for integrity verification | +| `file_size` | `int \| None` | Size of the file in bytes | | `modified_at` | `datetime \| None` | Last-modified timestamp if the source provides one | --- diff --git a/fern/running-tasks/api-spec.mdx b/fern/running-tasks/api-spec.mdx index 65a04572e..9c6e68ff0 100644 --- a/fern/running-tasks/api-spec.mdx +++ b/fern/running-tasks/api-spec.mdx @@ -66,7 +66,7 @@ Each task run has an associated `run_id` -- a unique identifier you can use to l | run_id | String | yes | "t_123456" or "oc_12345" | Unique identifier for this run | | status | String | yes | "running" | Current status of the run. Possible values: "created", "queued", "running", "timed_out", "failed", "terminated", "completed", "canceled" | | output | Object/Array/String/null | no | \{ "success": true } | Output data from the run, if any. Format depends on the schema in the input | -| downloaded_files | Array/null | no | [\{ "url": "...", "checksum": "...", "filename": "..." }] | List of files downloaded during the run | +| downloaded_files | Array/null | no | [\{ "url": "...", "checksum": "...", "filename": "...", "file_size": 1024 }] | List of files downloaded during the run | | recording_url | String | no | "https://..." | URL to the recording of the run | | failure_reason | String | no | "Max steps exceeded" | Reason for failure if the run failed | | created_at | String | yes | "2024-03-20T10:00:00Z" | Timestamp when this run was created | @@ -107,7 +107,7 @@ The response follows the same schema as the task creation response, containing: | run_id | String | "t_123456" | Unique identifier for this run | | status | String | "created" | Current status of the run. Possible values: "created", "queued", "running", "timed_out", "failed", "terminated", "completed", "canceled" | | output | Object | \{ "additionalProp1": \{ } } | Output data from the run | -| downloaded_files | Array | [\{ "url": "string", "checksum": "string", "filename": "string" }] | List of files downloaded during the run | +| downloaded_files | Array | [\{ "url": "string", "checksum": "string", "filename": "string", "file_size": 1024 }] | List of files downloaded during the run | | recording_url | String | "https://..." | URL to the recording of the run | | failure_reason | String | "Max steps exceeded" | Reason for failure if the run failed | | created_at | String | "2025-04-03T14:08:29.505Z" | Timestamp when this run was created | diff --git a/skyvern/forge/sdk/api/real_azure.py b/skyvern/forge/sdk/api/real_azure.py index 88957212d..0108723bb 100644 --- a/skyvern/forge/sdk/api/real_azure.py +++ b/skyvern/forge/sdk/api/real_azure.py @@ -308,13 +308,14 @@ class RealAsyncAzureStorageClient(AsyncAzureStorageClient): return await self.list_blobs(parsed.container, parsed.blob_path) async def get_object_info(self, uri: str) -> dict | None: - """Get object info including metadata. Returns dict with Metadata and LastModified keys.""" + """Get object info including metadata, last-modified time, and size.""" props = await self.get_blob_properties(uri) if props is None: return None return { "Metadata": props.get("metadata", {}), "LastModified": props.get("last_modified"), + "ContentLength": props.get("size"), } async def delete_file(self, uri: str) -> None: diff --git a/skyvern/forge/sdk/artifact/manager.py b/skyvern/forge/sdk/artifact/manager.py index 8d6d1c765..410210109 100644 --- a/skyvern/forge/sdk/artifact/manager.py +++ b/skyvern/forge/sdk/artifact/manager.py @@ -1,5 +1,6 @@ import asyncio import io +import os import time import zipfile from collections import defaultdict @@ -35,6 +36,16 @@ _SCREENSHOT_PREFIX_MAP: dict[ArtifactType, str] = { } +def _safe_file_size_from_path(path: str | None) -> int | None: + if path is None: + return None + try: + return os.path.getsize(path) + except OSError: + LOG.warning("Failed to get artifact file size", path=path, exc_info=True) + return None + + @dataclass class ArtifactBatchData: """ @@ -114,6 +125,7 @@ class ArtifactManager: uri: str, organization_id: str, bundle_key: str | None = None, + file_size: int | None = None, step_id: str | None = None, task_id: str | None = None, workflow_run_id: str | None = None, @@ -148,6 +160,7 @@ class ArtifactManager: artifact_type=artifact_type, uri=uri, bundle_key=bundle_key, + file_size=file_size, organization_id=organization_id, task_id=task_id, step_id=step_id, @@ -174,6 +187,7 @@ class ArtifactManager: task_v2_id: str | None = None, run_id: str | None = None, ai_suggestion_id: str | None = None, + file_size: int | None = None, data: bytes | None = None, path: str | None = None, ) -> str: @@ -194,6 +208,9 @@ class ArtifactManager: if not workflow_run_block_id and context: workflow_run_block_id = context.parent_workflow_run_block_id + if file_size is None: + file_size = _safe_file_size_from_path(path) + artifact = await app.DATABASE.artifacts.create_artifact( artifact_id, artifact_type, @@ -207,6 +224,7 @@ class ArtifactManager: run_id=run_id, organization_id=organization_id, ai_suggestion_id=ai_suggestion_id, + file_size=file_size, ) if data: # Fire and forget @@ -230,6 +248,7 @@ class ArtifactManager: uri = app.STORAGE.build_uri( organization_id=step.organization_id, artifact_id=artifact_id, step=step, artifact_type=artifact_type ) + file_size = len(data) if data is not None else _safe_file_size_from_path(path) return await self._create_artifact( aio_task_primary_key=step.task_id, artifact_id=artifact_id, @@ -238,6 +257,7 @@ class ArtifactManager: step_id=step.step_id, task_id=step.task_id, organization_id=step.organization_id, + file_size=file_size, data=data, path=path, ) @@ -286,6 +306,7 @@ class ArtifactManager: filename: str, workflow_run_id: str | None = None, checksum: str | None = None, + file_size: int | None = None, ) -> str: """Register a downloaded file as an Artifact row without re-uploading. @@ -318,6 +339,7 @@ class ArtifactManager: run_id=run_id, workflow_run_id=workflow_run_id, checksum=checksum, + file_size=file_size, ) LOG.debug( "Registered downloaded file as artifact", @@ -335,6 +357,7 @@ class ArtifactManager: uri: str, filename: str, checksum: str | None = None, + file_size: int | None = None, ) -> str: """Register a session-scoped downloaded file as an Artifact row. @@ -357,6 +380,7 @@ class ArtifactManager: filename=filename, artifact_type=ArtifactType.DOWNLOAD, checksum=checksum, + file_size=file_size, ) async def create_browser_session_recording_artifact( @@ -367,6 +391,7 @@ class ArtifactManager: uri: str, filename: str, checksum: str | None = None, + file_size: int | None = None, ) -> str: """Register a session-scoped recording (video) as a RECORDING Artifact row. @@ -383,6 +408,7 @@ class ArtifactManager: filename=filename, artifact_type=ArtifactType.RECORDING, checksum=checksum, + file_size=file_size, ) async def _create_browser_session_artifact( @@ -394,6 +420,7 @@ class ArtifactManager: filename: str, artifact_type: ArtifactType, checksum: str | None = None, + file_size: int | None = None, ) -> str: """Shared idempotent insert keyed on ``(browser_session_id, uri, artifact_type)``.""" existing = await app.DATABASE.artifacts.find_artifact_for_browser_session( @@ -413,6 +440,7 @@ class ArtifactManager: organization_id=organization_id, browser_session_id=browser_session_id, checksum=checksum, + file_size=file_size, ) LOG.debug( "Registered session-scoped artifact", @@ -631,6 +659,7 @@ class ArtifactManager: artifact_type=ArtifactType.SCRIPT_FILE, uri=uri, organization_id=organization_id, + file_size=len(data), data=data, ) diff --git a/skyvern/forge/sdk/artifact/models.py b/skyvern/forge/sdk/artifact/models.py index a954db0d3..d594084ae 100644 --- a/skyvern/forge/sdk/artifact/models.py +++ b/skyvern/forge/sdk/artifact/models.py @@ -86,6 +86,7 @@ class Artifact(BaseModel): uri: str bundle_key: str | None = None checksum: str | None = None + file_size: int | None = None task_id: str | None = None step_id: str | None = None workflow_run_id: str | None = None diff --git a/skyvern/forge/sdk/artifact/storage/azure.py b/skyvern/forge/sdk/artifact/storage/azure.py index 39b5593c5..472889b27 100644 --- a/skyvern/forge/sdk/artifact/storage/azure.py +++ b/skyvern/forge/sdk/artifact/storage/azure.py @@ -35,6 +35,22 @@ from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock LOG = structlog.get_logger() +def _safe_get_file_size(path: str) -> int | None: + try: + return os.path.getsize(path) + except OSError: + LOG.warning("Failed to get file size", path=path, exc_info=True) + return None + + +async def _get_object_info_safe(storage: "AzureStorage", uri: str) -> dict | None: + try: + return await storage.async_client.get_object_info(uri) + except Exception: + LOG.warning("Object info retrieval failed", uri=uri, exc_info=True) + return None + + class AzureStorage(BaseStorage): _PATH_VERSION = "v1" @@ -332,14 +348,13 @@ class AzureStorage(BaseStorage): for key in object_keys: metadata = {} modified_at: datetime | None = None + content_length: int | None = None # Get metadata (including checksum) - try: - object_info = await self.async_client.get_object_info(key) - if object_info: - metadata = object_info.get("Metadata", {}) - modified_at = object_info.get("LastModified") - except Exception: - LOG.exception("Object info retrieval failed", uri=key) + object_info = await _get_object_info_safe(self, key) + if object_info: + metadata = object_info.get("Metadata", {}) + modified_at = object_info.get("LastModified") + content_length = object_info.get("ContentLength") # Create FileInfo object filename = os.path.basename(key) @@ -354,6 +369,7 @@ class AzureStorage(BaseStorage): url=sas_urls[0], checksum=checksum, filename=metadata.get("original_filename", filename) if metadata else filename, + file_size=content_length, modified_at=modified_at, ) file_infos.append(file_info) @@ -443,14 +459,11 @@ class AzureStorage(BaseStorage): modified_at: datetime | None = None content_length: int | None = None # Get metadata (including checksum) - try: - object_info = await self.async_client.get_object_info(key) - if object_info: - metadata = object_info.get("Metadata", {}) - modified_at = object_info.get("LastModified") - content_length = object_info.get("ContentLength") or object_info.get("Size") - except Exception: - LOG.exception("Recording object info retrieval failed", uri=key) + object_info = await _get_object_info_safe(self, key) + if object_info: + metadata = object_info.get("Metadata", {}) + modified_at = object_info.get("LastModified") + content_length = object_info.get("ContentLength") # Skip zero-byte objects (if any incomplete uploads) if content_length == 0: @@ -469,6 +482,7 @@ class AzureStorage(BaseStorage): url=sas_urls[0], checksum=checksum, filename=metadata.get("original_filename", filename) if metadata else filename, + file_size=content_length, modified_at=modified_at, ) file_infos.append(file_info) @@ -512,6 +526,7 @@ class AzureStorage(BaseStorage): continue uri = f"{base_uri}/{file}" checksum = calculate_sha256_for_file(fpath) + file_size = _safe_get_file_size(fpath) # Azure Blob metadata values must be ASCII; preserve the full # filename via the blob path / Artifact URI instead. metadata: dict[str, str] = {"sha256_checksum": checksum} @@ -549,6 +564,7 @@ class AzureStorage(BaseStorage): uri=uri, filename=file, checksum=checksum, + file_size=file_size, ) except Exception: LOG.warning( @@ -599,7 +615,12 @@ class AzureStorage(BaseStorage): for key in object_keys: object_uri = f"azure://{settings.AZURE_STORAGE_CONTAINER_UPLOADS}/{key}" - metadata = await self.async_client.get_file_metadata(object_uri, log_exception=False) + metadata = {} + content_length: int | None = None + object_info = await _get_object_info_safe(self, object_uri) + if object_info: + metadata = object_info.get("Metadata", {}) + content_length = object_info.get("ContentLength") filename = os.path.basename(key) checksum = metadata.get("sha256_checksum") if metadata else None display_name = metadata.get("original_filename", filename) if metadata else filename @@ -613,6 +634,7 @@ class AzureStorage(BaseStorage): url=sas_urls[0], checksum=checksum, filename=display_name, + file_size=content_length, ) ) @@ -709,12 +731,14 @@ class AzureStorage(BaseStorage): # transient DB outage — both ops are idempotent. is_partial = remote_path.endswith(BROWSER_DOWNLOADING_SUFFIX) checksum = None if is_partial else calculate_sha256_for_file(local_file_path) + file_size = None if is_partial else _safe_get_file_size(local_file_path) await app.ARTIFACT_MANAGER.create_browser_session_download_artifact( organization_id=organization_id, browser_session_id=browser_session_id, uri=uri, filename=os.path.basename(remote_path), checksum=checksum, + file_size=file_size, ) elif artifact_type == "videos": # Recording uploaded once at session close — see s3.py. Artifact- @@ -725,12 +749,14 @@ class AzureStorage(BaseStorage): # we fall through to the Azure LIST path, so a row-less recording # still surfaces via the legacy SAS URL). checksum = calculate_sha256_for_file(local_file_path) + file_size = _safe_get_file_size(local_file_path) await app.ARTIFACT_MANAGER.create_browser_session_recording_artifact( organization_id=organization_id, browser_session_id=browser_session_id, uri=uri, filename=os.path.basename(remote_path), checksum=checksum, + file_size=file_size, ) return uri diff --git a/skyvern/forge/sdk/artifact/storage/base.py b/skyvern/forge/sdk/artifact/storage/base.py index 7f9ca10db..8ef68115b 100644 --- a/skyvern/forge/sdk/artifact/storage/base.py +++ b/skyvern/forge/sdk/artifact/storage/base.py @@ -44,6 +44,7 @@ async def _file_infos_from_artifacts(artifacts: list[Artifact], *, artifact_type url=url, checksum=artifact.checksum, filename=filename, + file_size=artifact.file_size, modified_at=artifact.created_at, artifact_id=artifact.artifact_id, ) diff --git a/skyvern/forge/sdk/artifact/storage/local.py b/skyvern/forge/sdk/artifact/storage/local.py index cc8ecba76..1761f4033 100644 --- a/skyvern/forge/sdk/artifact/storage/local.py +++ b/skyvern/forge/sdk/artifact/storage/local.py @@ -359,13 +359,18 @@ class LocalStorage(BaseStorage): if not path_obj.exists(): continue - file_size = path_obj.stat().st_size + try: + stat_result = path_obj.stat() + except OSError: + LOG.warning("Failed to stat local recording file", path=file_path, exc_info=True) + continue + file_size = stat_result.st_size if file_size == 0: continue # Return UTC-aware so consumers can safely compare against S3 LastModified # (also UTC-aware) without hitting naive-vs-aware TypeErrors. - modified_at = datetime.fromtimestamp(path_obj.stat().st_mtime, tz=UTC) + modified_at = datetime.fromtimestamp(stat_result.st_mtime, tz=UTC) checksum = calculate_sha256_for_file(file_path) filename = path_obj.name @@ -373,6 +378,7 @@ class LocalStorage(BaseStorage): url=uri, checksum=checksum, filename=filename, + file_size=file_size, modified_at=modified_at, ) file_infos.append(file_info) @@ -387,9 +393,18 @@ class LocalStorage(BaseStorage): for file_or_folder in files_and_folders: path = os.path.join(download_dir, file_or_folder) if os.path.isfile(path): - # Calculate checksum for the file checksum = calculate_sha256_for_file(path) - file_info = FileInfo(url=f"file://{path}", checksum=checksum, filename=file_or_folder) + try: + file_size = os.path.getsize(path) + except OSError: + LOG.warning("Failed to get local downloaded file size", path=path, exc_info=True) + file_size = None + file_info = FileInfo( + url=f"file://{path}", + checksum=checksum, + filename=file_or_folder, + file_size=file_size, + ) file_infos.append(file_info) return file_infos diff --git a/skyvern/forge/sdk/artifact/storage/s3.py b/skyvern/forge/sdk/artifact/storage/s3.py index 9b6c347a7..8c63e8b1f 100644 --- a/skyvern/forge/sdk/artifact/storage/s3.py +++ b/skyvern/forge/sdk/artifact/storage/s3.py @@ -41,6 +41,14 @@ LOG = structlog.get_logger() S3_ZSTD_COMPRESSED_SUFFIX = ".zst" +def _safe_get_file_size(path: str) -> int | None: + try: + return os.path.getsize(path) + except OSError: + LOG.warning("Failed to get file size", path=path, exc_info=True) + return None + + class S3Storage(BaseStorage): _PATH_VERSION = "v1" # Cap concurrent head_object fan-out in the legacy listing fallback. @@ -385,7 +393,7 @@ class S3Storage(BaseStorage): try: return await self.async_client.get_object_info(uri) except Exception: - LOG.exception("Object info retrieval failed", uri=uri) + LOG.warning("Object info retrieval failed", uri=uri, exc_info=True) return None async def _bounded_get_object_infos(self, keys: list[str]) -> list[dict | None]: @@ -457,6 +465,7 @@ class S3Storage(BaseStorage): continue metadata = (object_info or {}).get("Metadata") or {} modified_at: datetime | None = (object_info or {}).get("LastModified") + content_length: int | None = (object_info or {}).get("ContentLength") filename = os.path.basename(key) checksum = metadata.get("sha256_checksum") if metadata else None @@ -466,6 +475,7 @@ class S3Storage(BaseStorage): url=url, checksum=checksum, filename=metadata.get("original_filename", filename) if metadata else filename, + file_size=content_length, modified_at=modified_at, ) ) @@ -583,6 +593,7 @@ class S3Storage(BaseStorage): continue metadata = (object_info or {}).get("Metadata") or {} modified_at: datetime | None = (object_info or {}).get("LastModified") + content_length: int | None = (object_info or {}).get("ContentLength") filename = os.path.basename(key) checksum = metadata.get("sha256_checksum") if metadata else None @@ -592,6 +603,7 @@ class S3Storage(BaseStorage): url=url, checksum=checksum, filename=metadata.get("original_filename", filename) if metadata else filename, + file_size=content_length, modified_at=modified_at, ) ) @@ -634,6 +646,7 @@ class S3Storage(BaseStorage): continue uri = f"{base_uri}/{file}" checksum = calculate_sha256_for_file(fpath) + file_size = _safe_get_file_size(fpath) # S3 object metadata only allows ASCII; non-ASCII filenames (CJK, # emoji) would otherwise raise ParamValidationError at upload time. # The full filename is still preserved in the S3 key and on the @@ -674,6 +687,7 @@ class S3Storage(BaseStorage): uri=uri, filename=file, checksum=checksum, + file_size=file_size, ) except Exception: LOG.warning( @@ -729,7 +743,14 @@ class S3Storage(BaseStorage): for key in object_keys: object_uri = f"s3://{bucket}/{key}" - metadata = await self.async_client.get_file_metadata(object_uri, log_exception=False) + metadata = {} + content_length: int | None = None + try: + object_info = await self.async_client.get_object_info(object_uri) + metadata = object_info.get("Metadata", {}) + content_length = object_info.get("ContentLength") + except Exception: + LOG.warning("Object info retrieval failed", uri=object_uri, exc_info=True) filename = os.path.basename(key) checksum = metadata.get("sha256_checksum") if metadata else None display_name = metadata.get("original_filename", filename) if metadata else filename @@ -743,6 +764,7 @@ class S3Storage(BaseStorage): url=presigned_urls[0], checksum=checksum, filename=display_name, + file_size=content_length, ) ) return file_infos @@ -844,12 +866,14 @@ class S3Storage(BaseStorage): # (idempotent on ``(session, uri)``) are safe to retry. is_partial = remote_path.endswith(BROWSER_DOWNLOADING_SUFFIX) checksum = None if is_partial else calculate_sha256_for_file(local_file_path) + file_size = None if is_partial else _safe_get_file_size(local_file_path) await app.ARTIFACT_MANAGER.create_browser_session_download_artifact( organization_id=organization_id, browser_session_id=browser_session_id, uri=uri, filename=os.path.basename(remote_path), checksum=checksum, + file_size=file_size, ) elif artifact_type == "videos": # Register a RECORDING artifact row so @@ -870,12 +894,14 @@ class S3Storage(BaseStorage): # a row-less recording still surfaces via the legacy presigned # URL until a subsequent close writes the row. checksum = calculate_sha256_for_file(local_file_path) + file_size = _safe_get_file_size(local_file_path) await app.ARTIFACT_MANAGER.create_browser_session_recording_artifact( organization_id=organization_id, browser_session_id=browser_session_id, uri=uri, filename=os.path.basename(remote_path), checksum=checksum, + file_size=file_size, ) return uri diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 8627ccd04..d8bab6e52 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -259,6 +259,7 @@ class ArtifactModel(Base): run_id = Column(String, nullable=True) browser_session_id = Column(String, nullable=True) checksum = Column(String, nullable=True) + file_size = Column(BigInteger, nullable=True) created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False) modified_at = Column( DateTime, diff --git a/skyvern/forge/sdk/db/repositories/artifacts.py b/skyvern/forge/sdk/db/repositories/artifacts.py index 1825b8e85..4d1e7d6fa 100644 --- a/skyvern/forge/sdk/db/repositories/artifacts.py +++ b/skyvern/forge/sdk/db/repositories/artifacts.py @@ -49,6 +49,7 @@ class ArtifactsRepository(BaseRepository): thought_id: str | None = None, ai_suggestion_id: str | None = None, checksum: str | None = None, + file_size: int | None = None, browser_session_id: str | None = None, ) -> Artifact: async with self.Session() as session: @@ -66,6 +67,7 @@ class ArtifactsRepository(BaseRepository): ai_suggestion_id=ai_suggestion_id, organization_id=organization_id, checksum=checksum, + file_size=file_size, browser_session_id=browser_session_id, ) session.add(new_artifact) diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index be64c2927..f3af302ad 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/skyvern/forge/sdk/db/utils.py @@ -388,6 +388,7 @@ def convert_to_artifact(artifact_model: ArtifactModel, debug_enabled: bool = Fal uri=artifact_model.uri, bundle_key=artifact_model.bundle_key, checksum=artifact_model.checksum, + file_size=artifact_model.file_size, task_id=artifact_model.task_id, step_id=artifact_model.step_id, workflow_run_id=artifact_model.workflow_run_id, diff --git a/skyvern/forge/sdk/schemas/files.py b/skyvern/forge/sdk/schemas/files.py index 9877da836..6b3bf9bb8 100644 --- a/skyvern/forge/sdk/schemas/files.py +++ b/skyvern/forge/sdk/schemas/files.py @@ -9,6 +9,7 @@ class FileInfo(BaseModel): url: str = Field(..., description="URL to access the file") checksum: str | None = Field(None, description="SHA-256 checksum of the file") filename: str | None = Field(None, description="Original filename") + file_size: int | None = Field(None, description="Size of the file in bytes") modified_at: datetime | None = Field(None, description="Modified time of the file") # Optional: when the FileInfo is built from a DOWNLOAD Artifact row, the # row's id is carried through so persisted snapshots (e.g. block outputs) diff --git a/tests/unit/test_browser_session_download_artifacts.py b/tests/unit/test_browser_session_download_artifacts.py index c00354c1c..768301945 100644 --- a/tests/unit/test_browser_session_download_artifacts.py +++ b/tests/unit/test_browser_session_download_artifacts.py @@ -44,6 +44,7 @@ def _make_artifact( browser_session_id: str = "pbs_1", run_id: str | None = None, checksum: str | None = None, + file_size: int | None = None, created_at: str = "2026-04-25T00:00:00Z", ) -> Artifact: return Artifact( @@ -54,6 +55,7 @@ def _make_artifact( run_id=run_id, browser_session_id=browser_session_id, checksum=checksum, + file_size=file_size, created_at=created_at, modified_at=created_at, ) @@ -94,6 +96,7 @@ async def test_create_browser_session_download_artifact_inserts_when_no_existing uri="s3://skyvern-artifacts/v1/local/o_1/browser_sessions/pbs_1/downloads/file.pdf", filename="file.pdf", checksum="sha-xyz", + file_size=1234, ) assert artifact_id.startswith("a_") @@ -103,6 +106,7 @@ async def test_create_browser_session_download_artifact_inserts_when_no_existing assert kwargs["browser_session_id"] == "pbs_1" assert kwargs["organization_id"] == "o_1" assert kwargs["checksum"] == "sha-xyz" + assert kwargs["file_size"] == 1234 # No run_id at write time — claim happens at run finalization. assert kwargs.get("run_id") is None @@ -160,6 +164,7 @@ async def test_sync_browser_session_file_registers_download_artifact(): with ( patch.object(storage, "_get_storage_class_for_org", new=AsyncMock(return_value=MagicMock())), patch("skyvern.forge.sdk.artifact.storage.s3.calculate_sha256_for_file", return_value="sha-1"), + patch("skyvern.forge.sdk.artifact.storage.s3._safe_get_file_size", return_value=4321), patch("skyvern.forge.sdk.artifact.storage.s3.app") as app_module, ): app_module.ARTIFACT_MANAGER = mock_artifact_manager @@ -177,6 +182,7 @@ async def test_sync_browser_session_file_registers_download_artifact(): assert kwargs["browser_session_id"] == "pbs_1" assert kwargs["filename"] == "file.pdf" assert kwargs["checksum"] == "sha-1" + assert kwargs["file_size"] == 4321 assert kwargs["uri"].startswith("s3://") and "browser_sessions/pbs_1/downloads/" in kwargs["uri"] @@ -261,6 +267,7 @@ async def test_sync_browser_session_file_creates_partial_artifact_with_null_chec with ( patch.object(storage, "_get_storage_class_for_org", new=AsyncMock(return_value=MagicMock())), patch("skyvern.forge.sdk.artifact.storage.s3.calculate_sha256_for_file", mock_checksum), + patch("skyvern.forge.sdk.artifact.storage.s3._safe_get_file_size") as mock_file_size, patch("skyvern.forge.sdk.artifact.storage.s3.app") as app_module, ): app_module.ARTIFACT_MANAGER = mock_artifact_manager @@ -275,9 +282,11 @@ async def test_sync_browser_session_file_creates_partial_artifact_with_null_chec mock_create.assert_awaited_once() _, kwargs = mock_create.call_args assert kwargs["checksum"] is None + assert kwargs["file_size"] is None assert kwargs["uri"].endswith(BROWSER_DOWNLOADING_SUFFIX) # Partial files: checksum computation skipped — file is mid-write. mock_checksum.assert_not_called() + mock_file_size.assert_not_called() # --------------------------------------------------------------------------- @@ -296,6 +305,7 @@ async def test_get_shared_downloaded_files_in_browser_session_uses_artifact_urls "a_42", "s3://skyvern-artifacts/v1/local/o_1/browser_sessions/pbs_1/downloads/invoice.pdf", checksum="sha-from-db", + file_size=2048, ) mock_list = AsyncMock(return_value=[artifact]) build_url = MagicMock(return_value="https://api.skyvern.com/v1/artifacts/a_42/content?expiry=x&kid=y&sig=z") @@ -312,6 +322,7 @@ async def test_get_shared_downloaded_files_in_browser_session_uses_artifact_urls assert len(result) == 1 assert result[0].url.startswith("https://api.skyvern.com/v1/artifacts/a_42/content") assert result[0].checksum == "sha-from-db" + assert result[0].file_size == 2048 storage.async_client.list_files.assert_not_awaited() storage.async_client.create_presigned_urls.assert_not_awaited() @@ -327,6 +338,7 @@ async def test_get_shared_downloaded_files_in_browser_session_falls_back_to_pres return_value={ "Metadata": {"sha256_checksum": "sha-old", "original_filename": "legacy.pdf"}, "LastModified": None, + "ContentLength": 1024, } ) storage.async_client.create_presigned_urls = AsyncMock( @@ -346,6 +358,7 @@ async def test_get_shared_downloaded_files_in_browser_session_falls_back_to_pres assert len(result) == 1 assert _is_amazonaws_s3_url(result[0].url) + assert result[0].file_size == 1024 build_url.assert_not_called() @@ -474,7 +487,7 @@ async def test_get_shared_downloaded_files_in_browser_session_keyring_unset_skip object_uri = "s3://skyvern-artifacts/v1/local/o_1/browser_sessions/pbs_1/downloads/legacy.pdf" storage.async_client.list_files = AsyncMock(return_value=[object_uri.split("/", 3)[-1]]) storage.async_client.get_object_info = AsyncMock( - return_value={"Metadata": {"sha256_checksum": "sha-x"}, "LastModified": None} + return_value={"Metadata": {"sha256_checksum": "sha-x"}, "LastModified": None, "ContentLength": 2048} ) storage.async_client.create_presigned_urls = AsyncMock( return_value=["https://skyvern-artifacts.s3.amazonaws.com/...?sig=fallback"] @@ -493,6 +506,7 @@ async def test_get_shared_downloaded_files_in_browser_session_keyring_unset_skip assert len(result) == 1 assert _is_amazonaws_s3_url(result[0].url) + assert result[0].file_size == 2048 mock_list.assert_not_awaited() diff --git a/tests/unit/test_browser_session_recording_artifacts.py b/tests/unit/test_browser_session_recording_artifacts.py index a37c51d69..ea9178e30 100644 --- a/tests/unit/test_browser_session_recording_artifacts.py +++ b/tests/unit/test_browser_session_recording_artifacts.py @@ -38,6 +38,7 @@ def _make_recording_artifact( *, browser_session_id: str = "pbs_1", checksum: str | None = "sha-r", + file_size: int | None = None, created_at: str = "2026-04-26T00:00:00Z", ) -> Artifact: return Artifact( @@ -47,6 +48,7 @@ def _make_recording_artifact( organization_id="o_1", browser_session_id=browser_session_id, checksum=checksum, + file_size=file_size, created_at=created_at, modified_at=created_at, ) @@ -87,6 +89,7 @@ async def test_create_browser_session_recording_artifact_inserts_when_no_existin uri="s3://skyvern-artifacts/v1/local/o_1/browser_sessions/pbs_1/videos/2026-04-26/recording.webm", filename="recording.webm", checksum="sha-r", + file_size=4096, ) assert artifact_id.startswith("a_") @@ -96,6 +99,7 @@ async def test_create_browser_session_recording_artifact_inserts_when_no_existin assert kwargs["browser_session_id"] == "pbs_1" assert kwargs["organization_id"] == "o_1" assert kwargs["checksum"] == "sha-r" + assert kwargs["file_size"] == 4096 # Recordings are session-scoped, not run-scoped. assert kwargs.get("run_id") is None @@ -139,11 +143,13 @@ async def test_create_browser_session_recording_artifact_is_idempotent(): @pytest.mark.asyncio -async def test_sync_browser_session_file_registers_recording_artifact(): +async def test_sync_browser_session_file_registers_recording_artifact(tmp_path): """A successful 'videos' sync must create a RECORDING artifact row.""" storage = S3Storage() storage.async_client = MagicMock() storage.async_client.upload_file_from_path = AsyncMock() + recording_file = tmp_path / "recording.webm" + recording_file.write_bytes(b"fake-video") mock_create = AsyncMock(return_value="a_new") mock_artifact_manager = MagicMock() @@ -159,7 +165,7 @@ async def test_sync_browser_session_file_registers_recording_artifact(): organization_id="o_1", browser_session_id="pbs_1", artifact_type="videos", - local_file_path="/tmp/recording.webm", + local_file_path=str(recording_file), remote_path="recording.webm", date="2026-04-26", ) @@ -170,6 +176,7 @@ async def test_sync_browser_session_file_registers_recording_artifact(): assert kwargs["browser_session_id"] == "pbs_1" assert kwargs["filename"] == "recording.webm" assert kwargs["checksum"] == "sha-r" + assert kwargs["file_size"] == recording_file.stat().st_size assert kwargs["uri"].startswith("s3://") and "/browser_sessions/pbs_1/videos/" in kwargs["uri"] @@ -226,6 +233,7 @@ async def test_get_shared_recordings_returns_short_signed_urls(keyring_configure artifact = _make_recording_artifact( "a_42", "s3://skyvern-artifacts/v1/local/o_1/browser_sessions/pbs_1/videos/2026-04-26/recording.webm", + file_size=16384, ) mock_list = AsyncMock(return_value=[artifact]) build_url = MagicMock(return_value="https://api.skyvern.com/v1/artifacts/a_42/content?expiry=x&kid=y&sig=z") @@ -244,6 +252,7 @@ async def test_get_shared_recordings_returns_short_signed_urls(keyring_configure assert len(result) == 1 assert result[0].url.startswith("https://api.skyvern.com/v1/artifacts/a_42/content") assert result[0].artifact_id == "a_42" + assert result[0].file_size == 16384 storage.async_client.list_files.assert_not_awaited() storage.async_client.create_presigned_urls.assert_not_awaited() @@ -340,6 +349,7 @@ async def test_get_shared_recordings_falls_back_to_presigned_for_legacy_session( assert len(result) == 1 assert _is_amazonaws_s3_url(result[0].url) + assert result[0].file_size == 1024 build_url.assert_not_called() @@ -373,4 +383,5 @@ async def test_get_shared_recordings_keyring_unset_skips_artifact_lookup(): assert len(result) == 1 assert _is_amazonaws_s3_url(result[0].url) + assert result[0].file_size == 1024 mock_list.assert_not_awaited() diff --git a/tests/unit/test_downloaded_files_artifact_urls.py b/tests/unit/test_downloaded_files_artifact_urls.py index 53009cc7c..b99a47ef8 100644 --- a/tests/unit/test_downloaded_files_artifact_urls.py +++ b/tests/unit/test_downloaded_files_artifact_urls.py @@ -102,6 +102,7 @@ async def test_create_download_artifact_inserts_row_without_uploading(): workflow_run_id="wr_1", uri="s3://skyvern-uploads/download/prod/o_1/wr_1/file.pdf", filename="file.pdf", + file_size=123, ) assert artifact_id.startswith("a_") @@ -112,6 +113,7 @@ async def test_create_download_artifact_inserts_row_without_uploading(): assert kwargs["organization_id"] == "o_1" assert kwargs["run_id"] == "wr_1" assert kwargs["workflow_run_id"] == "wr_1" + assert kwargs["file_size"] == 123 mock_store.assert_not_awaited() @@ -144,7 +146,12 @@ async def test_save_downloaded_files_registers_artifact_per_file(tmp_path): assert mock_create_download.await_count == 2 uris = {call.kwargs["uri"] for call in mock_create_download.await_args_list} filenames = {call.kwargs["filename"] for call in mock_create_download.await_args_list} + file_sizes = {call.kwargs["filename"]: call.kwargs["file_size"] for call in mock_create_download.await_args_list} assert filenames == {"invoice.pdf", "report.csv"} + assert file_sizes == { + "invoice.pdf": (download_dir / "invoice.pdf").stat().st_size, + "report.csv": (download_dir / "report.csv").stat().st_size, + } assert all(u.startswith("s3://") and "/downloads/" in u and "/o_1/wr_1/" in u for u in uris) for call in mock_create_download.await_args_list: assert call.kwargs["organization_id"] == "o_1" @@ -157,6 +164,7 @@ def _make_artifact( run_id: str = "wr_1", *, checksum: str | None = None, + file_size: int | None = None, created_at: str = "2026-04-23T00:00:00Z", ) -> Artifact: return Artifact( @@ -167,6 +175,7 @@ def _make_artifact( run_id=run_id, workflow_run_id=run_id if run_id.startswith("wr_") else None, checksum=checksum, + file_size=file_size, created_at=created_at, modified_at=created_at, ) @@ -200,6 +209,7 @@ async def test_get_downloaded_files_uses_artifact_urls_when_rows_exist(keyring_c "a_42", "s3://skyvern-uploads/downloads/local/o_1/wr_1/invoice.pdf", checksum="sha-from-db", + file_size=4096, ) mock_list = AsyncMock(return_value=[artifact]) build_url = MagicMock(return_value="https://api.skyvern.com/v1/artifacts/a_42/content?expiry=x&kid=y&sig=z") @@ -215,6 +225,7 @@ async def test_get_downloaded_files_uses_artifact_urls_when_rows_exist(keyring_c assert result[0].url.startswith("https://api.skyvern.com/v1/artifacts/a_42/content") assert result[0].filename == "invoice.pdf" assert result[0].checksum == "sha-from-db" + assert result[0].file_size == 4096 assert result[0].modified_at is not None storage.async_client.list_files.assert_not_awaited() storage.async_client.get_file_metadata.assert_not_awaited() @@ -263,8 +274,11 @@ async def test_get_downloaded_files_falls_back_to_presigned_for_legacy_runs(keyr storage.async_client = MagicMock() s3_key = "downloads/local/o_1/wr_old/legacy.pdf" storage.async_client.list_files = AsyncMock(return_value=[s3_key]) - storage.async_client.get_file_metadata = AsyncMock( - return_value={"sha256_checksum": "sha-old", "original_filename": "legacy.pdf"} + storage.async_client.get_object_info = AsyncMock( + return_value={ + "Metadata": {"sha256_checksum": "sha-old", "original_filename": "legacy.pdf"}, + "ContentLength": 2048, + } ) storage.async_client.create_presigned_urls = AsyncMock( return_value=["https://skyvern-uploads.s3.amazonaws.com/...?sig=old"] @@ -282,6 +296,7 @@ async def test_get_downloaded_files_falls_back_to_presigned_for_legacy_runs(keyr assert len(result) == 1 assert result[0].filename == "legacy.pdf" assert result[0].checksum == "sha-old" + assert result[0].file_size == 2048 assert _is_amazonaws_s3_url(result[0].url) build_url.assert_not_called() storage.async_client.list_files.assert_awaited_once() @@ -299,8 +314,11 @@ async def test_get_downloaded_files_falls_back_to_presigned_when_keyring_unset(t storage.async_client = MagicMock() s3_key = "downloads/local/o_1/wr_1/invoice.pdf" storage.async_client.list_files = AsyncMock(return_value=[s3_key]) - storage.async_client.get_file_metadata = AsyncMock( - return_value={"sha256_checksum": "sha-abc", "original_filename": "invoice.pdf"} + storage.async_client.get_object_info = AsyncMock( + return_value={ + "Metadata": {"sha256_checksum": "sha-abc", "original_filename": "invoice.pdf"}, + "ContentLength": 1024, + } ) storage.async_client.create_presigned_urls = AsyncMock( return_value=["https://skyvern-uploads.s3.amazonaws.com/...?sig=fallback"] @@ -320,6 +338,7 @@ async def test_get_downloaded_files_falls_back_to_presigned_when_keyring_unset(t assert len(result) == 1 assert _is_amazonaws_s3_url(result[0].url) + assert result[0].file_size == 1024 build_url.assert_not_called() @@ -332,8 +351,11 @@ async def test_get_downloaded_files_artifact_lookup_failure_falls_back_to_listin storage.async_client = MagicMock() s3_key = "downloads/local/o_1/wr_1/recoverable.pdf" storage.async_client.list_files = AsyncMock(return_value=[s3_key]) - storage.async_client.get_file_metadata = AsyncMock( - return_value={"sha256_checksum": "sha-recover", "original_filename": "recoverable.pdf"} + storage.async_client.get_object_info = AsyncMock( + return_value={ + "Metadata": {"sha256_checksum": "sha-recover", "original_filename": "recoverable.pdf"}, + "ContentLength": 512, + } ) storage.async_client.create_presigned_urls = AsyncMock( return_value=["https://skyvern-uploads.s3.amazonaws.com/...?sig=fallback"] @@ -347,6 +369,7 @@ async def test_get_downloaded_files_artifact_lookup_failure_falls_back_to_listin assert len(result) == 1 assert _is_amazonaws_s3_url(result[0].url) + assert result[0].file_size == 512 storage.async_client.list_files.assert_awaited_once()