diff --git a/dify-agent/src/dify_agent/agent_stub/_drive_materialization.py b/dify-agent/src/dify_agent/agent_stub/_drive_materialization.py new file mode 100644 index 00000000000..74181fc1edc --- /dev/null +++ b/dify-agent/src/dify_agent/agent_stub/_drive_materialization.py @@ -0,0 +1,167 @@ +"""Shared drive download materialization helpers. + +This module centralizes the safety-critical filesystem logic used by both the +sandbox-visible CLI and the runtime drive layer. It owns path resolution under +one local drive base, overwrite-via-temp-file semantics, payload size checks, +and safe extraction of downloaded skill archives so those invariants cannot +drift between the two call sites. +""" + +from __future__ import annotations + +import stat +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from tempfile import TemporaryDirectory +from typing import Collection, Final, Mapping +from uuid import uuid4 +from zipfile import BadZipFile, ZipFile, ZipInfo + + +SKILL_ARCHIVE_FILENAME: Final[str] = ".DIFY-SKILL-FULL.zip" + + +@dataclass(frozen=True, slots=True) +class DriveDownloadPayload: + """One downloaded drive payload ready to materialize under a local base.""" + + key: str + payload: bytes + size: int | None = None + + +class DriveMaterializationValidationError(ValueError): + """Raised when one drive key or archive entry is structurally unsafe.""" + + +class DriveMaterializationTransferError(RuntimeError): + """Raised when one downloaded payload cannot be safely materialized.""" + + +def materialize_drive_downloads( + *, + base_path: Path, + downloads: list[DriveDownloadPayload], + archive_skip_entry_names_by_dir: Mapping[str, Collection[str]] | None = None, +) -> list[Path]: + """Write downloaded drive payloads under one local base and extract skills. + + The helper preserves caller-provided order in the returned list of written + paths. Skill archives are extracted only after every payload has been + written successfully so partial extraction cannot outlive a later failure in + the same batch. + """ + + resolved_base_path = base_path.expanduser().resolve() + try: + _ = resolved_base_path.mkdir(parents=True, exist_ok=True) + except OSError as exc: + raise DriveMaterializationTransferError(f"failed to prepare drive base {resolved_base_path}") from exc + + written_paths: list[Path] = [] + archive_paths: list[Path] = [] + skip_entry_names_by_dir = archive_skip_entry_names_by_dir or {} + for download in downloads: + if download.size is not None and len(download.payload) != download.size: + raise DriveMaterializationTransferError(f"downloaded drive file size mismatch for {download.key}") + destination = resolve_drive_destination(resolved_base_path, download.key) + try: + destination.parent.mkdir(parents=True, exist_ok=True) + temp_path = destination.with_name(f"{destination.name}.tmp-{uuid4().hex}") + _ = temp_path.write_bytes(download.payload) + _ = temp_path.replace(destination) + except OSError as exc: + raise DriveMaterializationTransferError(f"failed to materialize drive file {download.key}") from exc + written_paths.append(destination) + if destination.name == SKILL_ARCHIVE_FILENAME: + archive_paths.append(destination) + + for archive_path in sorted(archive_paths): + archive_skill_dir = archive_path.parent.relative_to(resolved_base_path).as_posix() + extract_skill_archive( + archive_path, + skip_entry_names=frozenset(skip_entry_names_by_dir.get(archive_skill_dir, ())), + ) + return written_paths + + +def resolve_drive_destination(base_path: Path, drive_key: str) -> Path: + """Resolve one drive key under a local base and reject path traversal.""" + + destination = (base_path / Path(drive_key)).resolve() + try: + destination.relative_to(base_path) + except ValueError as exc: + raise DriveMaterializationValidationError(f"drive key resolves outside the drive base: {drive_key}") from exc + return destination + + +def extract_skill_archive(archive_path: Path, *, skip_entry_names: Collection[str] = ()) -> None: + """Safely extract one downloaded skill archive into its containing directory.""" + + target_dir = archive_path.parent.resolve() + normalized_skip_entry_names = {entry_name.replace("\\", "/").rstrip("/") for entry_name in skip_entry_names} + try: + with TemporaryDirectory(dir=target_dir, prefix=".dify-skill-extract-") as staging_dir_name: + staging_dir = Path(staging_dir_name).resolve() + with ZipFile(archive_path) as archive: + for zip_info in archive.infolist(): + if zip_info.filename.replace("\\", "/").rstrip("/") in normalized_skip_entry_names: + continue + destination = _resolve_zip_entry_destination(staging_dir, zip_info.filename) + if _is_zip_symlink(zip_info): + raise DriveMaterializationValidationError( + f"skill archive contains unsupported symlink entry: {zip_info.filename}" + ) + if zip_info.is_dir(): + destination.mkdir(parents=True, exist_ok=True) + continue + destination.parent.mkdir(parents=True, exist_ok=True) + with archive.open(zip_info) as source_file: + temp_path = destination.with_name(f"{destination.name}.tmp-{uuid4().hex}") + _ = temp_path.write_bytes(source_file.read()) + _ = temp_path.replace(destination) + for staged_path in sorted(staging_dir.rglob("*")): + if staged_path.is_dir(): + continue + relative_path = staged_path.relative_to(staging_dir) + destination = (target_dir / relative_path).resolve() + destination.parent.mkdir(parents=True, exist_ok=True) + _ = staged_path.replace(destination) + except DriveMaterializationValidationError: + raise + except (BadZipFile, OSError) as exc: + raise DriveMaterializationTransferError(f"downloaded skill archive is invalid: {archive_path.name}") from exc + + +def _resolve_zip_entry_destination(target_dir: Path, entry_name: str) -> Path: + normalized_name = entry_name.replace("\\", "/") + pure_path = PurePosixPath(normalized_name) + if not normalized_name or normalized_name.startswith("/") or pure_path.is_absolute(): + raise DriveMaterializationValidationError(f"skill archive contains unsafe absolute path: {entry_name}") + if any(part in {"", ".", ".."} for part in pure_path.parts): + raise DriveMaterializationValidationError(f"skill archive contains unsafe path traversal entry: {entry_name}") + destination = (target_dir / Path(*pure_path.parts)).resolve() + try: + destination.relative_to(target_dir) + except ValueError as exc: + raise DriveMaterializationValidationError( + f"skill archive entry resolves outside the skill directory: {entry_name}" + ) from exc + return destination + + +def _is_zip_symlink(zip_info: ZipInfo) -> bool: + file_mode = zip_info.external_attr >> 16 + return stat.S_ISLNK(file_mode) + + +__all__ = [ + "DriveDownloadPayload", + "DriveMaterializationTransferError", + "DriveMaterializationValidationError", + "SKILL_ARCHIVE_FILENAME", + "extract_skill_archive", + "materialize_drive_downloads", + "resolve_drive_destination", +] diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_drive.py b/dify-agent/src/dify_agent/agent_stub/cli/_drive.py index d3d6241449b..b4bba338fbe 100644 --- a/dify-agent/src/dify_agent/agent_stub/cli/_drive.py +++ b/dify-agent/src/dify_agent/agent_stub/cli/_drive.py @@ -10,14 +10,23 @@ ToolFile ids back into the drive. from __future__ import annotations -import stat from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -from pathlib import Path, PurePosixPath +from pathlib import Path from tempfile import TemporaryDirectory -from uuid import uuid4 -from zipfile import BadZipFile, ZIP_DEFLATED, ZipFile, ZipInfo +from typing import ClassVar, Literal +from zipfile import ZIP_DEFLATED, ZipFile +from pydantic import BaseModel, ConfigDict + +from dify_agent.agent_stub._drive_materialization import ( + DriveDownloadPayload, + DriveMaterializationTransferError, + DriveMaterializationValidationError, + SKILL_ARCHIVE_FILENAME, + materialize_drive_downloads, + resolve_drive_destination, +) from dify_agent.agent_stub.cli._env import read_agent_stub_environment from dify_agent.agent_stub.cli._files import upload_tool_file_resource_from_environment from dify_agent.agent_stub.client._agent_stub import ( @@ -37,11 +46,11 @@ from dify_agent.agent_stub.protocol.agent_stub import ( ) _SKILL_MD_FILENAME = "SKILL.md" -_SKILL_ARCHIVE_FILENAME = ".DIFY-SKILL-FULL.zip" _SKIP_DIR_NAMES = frozenset( {".git", "__pycache__", ".pytest_cache", ".mypy_cache", ".ruff_cache", ".venv", "node_modules"} ) -_SKIP_FILE_NAMES = frozenset({".DS_Store", _SKILL_ARCHIVE_FILENAME}) +_SKIP_FILE_NAMES = frozenset({".DS_Store", SKILL_ARCHIVE_FILENAME}) +DrivePushKind = Literal["file", "skill", "dir"] @dataclass(frozen=True, slots=True) @@ -52,18 +61,28 @@ class _DriveUploadItem: drive_key: str -def list_drive_from_environment(prefix: str, json_output: bool) -> str | AgentStubDriveManifestResponse: +class DrivePullResult(BaseModel): + """Structured JSON result for ``dify-agent drive pull --json``.""" + + class Item(BaseModel): + key: str + local_path: str + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + items: list[Item] + + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") + + +def list_drive_manifest_from_environment(prefix: str) -> AgentStubDriveManifestResponse: """List drive items through the Agent Stub using the current environment. Args: prefix: Optional drive-key prefix forwarded to the manifest request. - json_output: When ``True``, return the validated manifest response model. - When ``False``, return one human-readable tab-separated line per item - containing size, mime type, hash, and key. Returns: - Either ``AgentStubDriveManifestResponse`` for JSON callers or a formatted - string for human-facing CLI output. + The validated manifest response model. Side effects: Calls the Agent Stub drive manifest control-plane endpoint with @@ -78,24 +97,24 @@ def list_drive_from_environment(prefix: str, json_output: bool) -> str | AgentSt prefix=prefix, include_download_url=False, ) - if json_output: - return response - return _format_manifest(response) + return response def pull_drive_from_environment( targets: list[str] | None = None, - drive_base: str = DEFAULT_AGENT_STUB_DRIVE_BASE, -) -> list[Path]: + local_base: str | None = None, +) -> DrivePullResult: """Pull drive files into one local drive base via signed download URLs. Args: targets: Optional drive-key targets or prefixes. An empty list preserves the historical whole-drive pull by using ``[""]``. - drive_base: Local base directory that receives downloaded drive files. + local_base: Local base directory that receives downloaded drive files. + When omitted, the historical Agent Stub drive base is used. Returns: - A list of written local paths under ``drive_base``. + A structured JSON-ready result with downloaded drive keys and their + written local paths. Observable behavior: Requests a manifest with ``include_download_url=True``, requires every @@ -109,8 +128,8 @@ def pull_drive_from_environment( under a temporary directory and only moved into place after the full archive validates successfully. - The return value remains the list of downloaded paths only; extracted - files are materialized on disk but are not added to the returned list. + Extracted files are materialized on disk but are not added to the + returned item list. Raises: AgentStubValidationError: if a manifest item omits ``download_url``, a @@ -124,48 +143,62 @@ def pull_drive_from_environment( environment = read_agent_stub_environment() manifest_targets = targets or [""] - with ThreadPoolExecutor(max_workers=min(len(manifest_targets), 4)) as executor: - responses = list( - executor.map( - lambda target: request_agent_stub_drive_manifest_sync( - url=environment.url, - auth_jwe=environment.auth_jwe, - prefix=target, - include_download_url=True, - ), - manifest_targets, - ) + + def _fetch_manifest(target: str) -> AgentStubDriveManifestResponse: + return request_agent_stub_drive_manifest_sync( + url=environment.url, + auth_jwe=environment.auth_jwe, + prefix=target, + include_download_url=True, ) - base_path = Path(drive_base).expanduser().resolve() - base_path.mkdir(parents=True, exist_ok=True) - written_paths: list[Path] = [] + + with ThreadPoolExecutor(max_workers=min(len(manifest_targets), 4)) as executor: + responses = list(executor.map(_fetch_manifest, manifest_targets)) + downloads: list[DriveDownloadPayload] = [] + resolved_base_path = Path(local_base or DEFAULT_AGENT_STUB_DRIVE_BASE).expanduser().resolve() deduplicated_items = {item.key: item for response in responses for item in response.items} for item in [deduplicated_items[key] for key in sorted(deduplicated_items)]: download_url = item.download_url if not isinstance(download_url, str) or not download_url: raise AgentStubValidationError(f"drive manifest item is missing download_url: {item.key}") - destination = _resolve_drive_destination(base_path, item.key) + try: + _ = resolve_drive_destination(resolved_base_path, item.key) + except DriveMaterializationValidationError as exc: + raise AgentStubValidationError(str(exc)) from exc payload = download_file_bytes_from_signed_url_sync(download_url=download_url) - if item.size is not None and len(payload) != item.size: - raise AgentStubTransferError(f"downloaded drive file size mismatch for {item.key}") - destination.parent.mkdir(parents=True, exist_ok=True) - temp_path = destination.with_name(f"{destination.name}.tmp-{uuid4().hex}") - temp_path.write_bytes(payload) - temp_path.replace(destination) - written_paths.append(destination) - if destination.name == _SKILL_ARCHIVE_FILENAME: - _extract_skill_archive(destination) - return written_paths + downloads.append(DriveDownloadPayload(key=item.key, payload=payload, size=item.size)) + + try: + written_paths = materialize_drive_downloads( + base_path=resolved_base_path, + downloads=downloads, + ) + except DriveMaterializationValidationError as exc: + raise AgentStubValidationError(str(exc)) from exc + except DriveMaterializationTransferError as exc: + raise AgentStubTransferError(str(exc)) from exc + + return DrivePullResult( + items=[ + DrivePullResult.Item(key=download.key, local_path=str(path)) + for download, path in zip(downloads, written_paths, strict=True) + ] + ) -def push_drive_from_environment(local_path: str, drive_path: str, recursive: bool) -> AgentStubDriveCommitResponse: +def push_drive_from_environment( + local_path: str, + drive_path: str, + *, + kind: DrivePushKind | None, +) -> AgentStubDriveCommitResponse: """Upload local files through the Agent Stub and commit them into the drive. Args: local_path: Source file or directory in the sandbox filesystem. drive_path: Destination drive key or drive-key prefix. - recursive: Select directory mode. ``False`` standardizes skill - directories, while ``True`` uploads raw directory contents. + kind: Optional public upload mode. Files infer file mode when omitted, + while directories require explicit ``skill`` or ``dir`` selection. Returns: The validated drive commit response returned by the Agent Stub. @@ -173,25 +206,40 @@ def push_drive_from_environment(local_path: str, drive_path: str, recursive: boo Mode split: * If ``local_path`` is a file, upload that file and commit exactly one ``tool_file`` binding to ``drive_path``. - * If ``local_path`` is a directory and ``recursive`` is ``False``, + * If ``local_path`` is a directory and ``kind`` is ``"skill"``, require ``SKILL.md`` and standardize the upload into ``/SKILL.md`` plus ``/.DIFY-SKILL-FULL.zip``. - * If ``local_path`` is a directory and ``recursive`` is ``True``, upload + * If ``local_path`` is a directory and ``kind`` is ``"dir"``, upload each regular file under ``drive_path/`` without skill standardization. Observable safety behavior: - Rejects missing local paths, rejects recursive directory pushes with no - regular files, and rejects symlinked or escaping paths while preparing - directory uploads or skill archives. + Rejects missing local paths, rejects directory pushes without an + explicit mode, rejects raw directory pushes with no regular files, and + rejects symlinked or escaping paths, including symlinked top-level + ``local_path`` roots, while preparing directory uploads or skill + archives. """ - source_path = Path(local_path).expanduser().resolve() + source_path = Path(local_path).expanduser() + if kind not in {None, "file", "skill", "dir"}: + raise AgentStubValidationError(f"invalid drive push kind: {kind}") + if source_path.is_symlink(): + raise AgentStubValidationError(f"drive push does not support symlinked local paths: {source_path}") + source_path = source_path.resolve() if source_path.is_file(): + if kind == "skill": + raise AgentStubValidationError("--kind skill requires a directory containing SKILL.md") + if kind == "dir": + raise AgentStubValidationError("--kind dir requires a directory") return _commit_uploaded_items([_prepare_uploaded_file(source_path, drive_path)]) if not source_path.is_dir(): raise AgentStubValidationError(f"local path not found: {source_path}") - if recursive: + if kind is None: + raise AgentStubValidationError("directory drive push requires --kind skill or --kind dir") + if kind == "file": + raise AgentStubValidationError("--kind file requires a file") + if kind == "dir": upload_items = [ _prepare_uploaded_file(path, _join_drive_key(drive_path, relative_path)) for path, relative_path in _iter_regular_files(source_path) @@ -205,14 +253,14 @@ def push_drive_from_environment(local_path: str, drive_path: str, recursive: boo def _push_skill_directory(source_path: Path, drive_path: str) -> AgentStubDriveCommitResponse: skill_md_path = source_path / _SKILL_MD_FILENAME if not skill_md_path.is_file(): - raise AgentStubValidationError(f"non-recursive drive push requires {_SKILL_MD_FILENAME}: {source_path}") + raise AgentStubValidationError("--kind skill requires a directory containing SKILL.md") with TemporaryDirectory() as temp_dir: - archive_path = Path(temp_dir) / _SKILL_ARCHIVE_FILENAME + archive_path = Path(temp_dir) / SKILL_ARCHIVE_FILENAME _build_skill_archive(source_path, archive_path) return _commit_uploaded_items( [ _prepare_uploaded_file(skill_md_path.resolve(), _join_drive_key(drive_path, _SKILL_MD_FILENAME)), - _prepare_uploaded_file(archive_path, _join_drive_key(drive_path, _SKILL_ARCHIVE_FILENAME)), + _prepare_uploaded_file(archive_path, _join_drive_key(drive_path, SKILL_ARCHIVE_FILENAME)), ] ) @@ -239,7 +287,7 @@ def _commit_uploaded_items(items: list[_DriveUploadItem]) -> AgentStubDriveCommi ) -def _format_manifest(response: AgentStubDriveManifestResponse) -> str: +def format_drive_manifest(response: AgentStubDriveManifestResponse) -> str: return "\n".join(_format_manifest_item(item) for item in response.items) @@ -250,15 +298,6 @@ def _format_manifest_item(item: AgentStubDriveItem) -> str: return f"{size}\t{mime_type}\t{item_hash}\t{item.key}" -def _resolve_drive_destination(base_path: Path, drive_key: str) -> Path: - destination = (base_path / Path(drive_key)).resolve() - try: - destination.relative_to(base_path) - except ValueError as exc: - raise AgentStubValidationError(f"drive key resolves outside the drive base: {drive_key}") from exc - return destination - - def _iter_regular_files(root_path: Path) -> list[tuple[Path, str]]: """Return all regular files under one directory, rejecting unsafe symlinks.""" @@ -305,82 +344,6 @@ def _build_skill_archive(source_path: Path, archive_path: Path) -> None: archive.write(file_path, arcname=relative_path) -def _extract_skill_archive(archive_path: Path) -> None: - """Safely extract one downloaded skill archive into its containing directory. - - Extraction is staged under a temporary directory created inside the target - skill directory. Every entry is validated and materialized into staging - first, and only after the full archive succeeds are staged files moved into - their final locations under the skill directory. Existing files at those - final locations are overwritten in place by the extracted archive content. - - Error mapping is intentionally stable for CLI callers: unsafe archive entry - names raise ``AgentStubValidationError``, while malformed archives and zip - parsing / archive I/O failures are translated into ``AgentStubTransferError``. - """ - - target_dir = archive_path.parent.resolve() - try: - with TemporaryDirectory(dir=target_dir, prefix=".dify-skill-extract-") as staging_dir_name: - staging_dir = Path(staging_dir_name).resolve() - with ZipFile(archive_path) as archive: - for zip_info in archive.infolist(): - destination = _resolve_zip_entry_destination(staging_dir, zip_info.filename) - if _is_zip_symlink(zip_info): - raise AgentStubValidationError( - f"skill archive contains unsupported symlink entry: {zip_info.filename}" - ) - if zip_info.is_dir(): - destination.mkdir(parents=True, exist_ok=True) - continue - destination.parent.mkdir(parents=True, exist_ok=True) - with archive.open(zip_info) as source_file: - temp_path = destination.with_name(f"{destination.name}.tmp-{uuid4().hex}") - temp_path.write_bytes(source_file.read()) - temp_path.replace(destination) - for staged_path in sorted(staging_dir.rglob("*")): - if staged_path.is_dir(): - continue - relative_path = staged_path.relative_to(staging_dir) - destination = (target_dir / relative_path).resolve() - destination.parent.mkdir(parents=True, exist_ok=True) - staged_path.replace(destination) - except AgentStubValidationError: - raise - except (BadZipFile, OSError) as exc: - raise AgentStubTransferError(f"downloaded skill archive is invalid: {archive_path.name}") from exc - - -def _resolve_zip_entry_destination(target_dir: Path, entry_name: str) -> Path: - """Resolve one zip entry path under a target skill directory. - - Zip metadata may contain POSIX or backslash-separated names, so entry names - are normalized to forward slashes before validation. The resolved entry must - not be absolute, empty, ``.`` / ``..`` based, or otherwise escape the target - skill directory after resolution. - """ - - normalized_name = entry_name.replace("\\", "/") - pure_path = PurePosixPath(normalized_name) - if not normalized_name or normalized_name.startswith("/") or pure_path.is_absolute(): - raise AgentStubValidationError(f"skill archive contains unsafe absolute path: {entry_name}") - if any(part in {"", ".", ".."} for part in pure_path.parts): - raise AgentStubValidationError(f"skill archive contains unsafe path traversal entry: {entry_name}") - destination = (target_dir / Path(*pure_path.parts)).resolve() - try: - destination.relative_to(target_dir) - except ValueError as exc: - raise AgentStubValidationError( - f"skill archive entry resolves outside the skill directory: {entry_name}" - ) from exc - return destination - - -def _is_zip_symlink(zip_info: ZipInfo) -> bool: - file_mode = zip_info.external_attr >> 16 - return stat.S_ISLNK(file_mode) - - def _join_drive_key(base_key: str, child_key: str) -> str: stripped_base = base_key.rstrip("/") stripped_child = child_key.lstrip("/") @@ -388,7 +351,10 @@ def _join_drive_key(base_key: str, child_key: str) -> str: __all__ = [ - "list_drive_from_environment", + "DrivePullResult", + "DrivePushKind", + "format_drive_manifest", + "list_drive_manifest_from_environment", "pull_drive_from_environment", "push_drive_from_environment", ] diff --git a/dify-agent/src/dify_agent/agent_stub/cli/_files.py b/dify-agent/src/dify_agent/agent_stub/cli/_files.py index dccc6d36c44..38388cf097e 100644 --- a/dify-agent/src/dify_agent/agent_stub/cli/_files.py +++ b/dify-agent/src/dify_agent/agent_stub/cli/_files.py @@ -5,7 +5,7 @@ from __future__ import annotations import mimetypes from dataclasses import dataclass from pathlib import Path -from typing import Literal, cast +from typing import ClassVar, Literal, cast from pydantic import BaseModel, ConfigDict, ValidationError @@ -26,7 +26,7 @@ class UploadedToolFileMapping(BaseModel): transfer_method: Literal["tool_file"] = "tool_file" reference: str - model_config = ConfigDict(extra="forbid") + model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") @dataclass(frozen=True, slots=True) @@ -76,12 +76,14 @@ def upload_tool_file_resource_from_environment(*, path: str) -> UploadedToolFile environment = read_agent_stub_environment() filename = source_path.name mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream" - upload_request = request_agent_stub_file_upload_sync( + upload_request: object = request_agent_stub_file_upload_sync( url=environment.url, auth_jwe=environment.auth_jwe, filename=filename, mimetype=mime_type, ) + if not hasattr(upload_request, "upload_url") or not isinstance(upload_request.upload_url, str): + raise AgentStubTransferError("signed file upload response is missing upload_url") with source_path.open("rb") as file_obj: payload = upload_file_to_signed_url_sync( upload_url=upload_request.upload_url, @@ -94,19 +96,69 @@ def upload_tool_file_resource_from_environment(*, path: str) -> UploadedToolFile def download_file_from_environment( *, - transfer_method: str, - reference_or_url: str, - directory: str | None = None, + transfer_method: str | None = None, + reference_or_url: str | None = None, + mapping: str | None = None, + local_dir: str | None = None, ) -> DownloadedFileResult: - """Download one workflow file mapping into the sandbox filesystem.""" + """Download one workflow file mapping into the sandbox filesystem. + Callers may provide either the public positional pair + ``TRANSFER_METHOD REFERENCE_OR_URL`` or one JSON ``--mapping`` payload. + The helper normalizes both forms into ``AgentStubFileMapping`` before + requesting a signed download URL from the Agent Stub. + """ + + file_mapping = _build_download_mapping( + transfer_method=transfer_method, + reference_or_url=reference_or_url, + mapping=mapping, + ) environment = read_agent_stub_environment() + + download_request: object = request_agent_stub_file_download_sync( + url=environment.url, + auth_jwe=environment.auth_jwe, + file=file_mapping, + ) + if not hasattr(download_request, "filename") or not isinstance(download_request.filename, str): + raise AgentStubTransferError("signed file download response is missing filename") + if not hasattr(download_request, "download_url") or not isinstance(download_request.download_url, str): + raise AgentStubTransferError("signed file download response is missing download_url") + target_dir = Path(local_dir).expanduser().resolve() if local_dir else Path.cwd() + target_dir.mkdir(parents=True, exist_ok=True) + destination = _deduplicate_destination_path(target_dir / _sanitize_download_filename(download_request.filename)) + _ = destination.write_bytes(download_file_bytes_from_signed_url_sync(download_url=download_request.download_url)) + return DownloadedFileResult(path=destination) + + +def _build_download_mapping( + *, + transfer_method: str | None, + reference_or_url: str | None, + mapping: str | None, +) -> AgentStubFileMapping: + if mapping is not None: + if transfer_method is not None or reference_or_url is not None: + raise AgentStubValidationError( + "--mapping cannot be combined with TRANSFER_METHOD or REFERENCE_OR_URL" + ) + try: + return AgentStubFileMapping.model_validate_json(mapping) + except ValidationError as exc: + raise AgentStubValidationError("invalid file download mapping") from exc + + if transfer_method is None or reference_or_url is None: + raise AgentStubValidationError( + "file download requires either --mapping or TRANSFER_METHOD REFERENCE_OR_URL" + ) + normalized_transfer_method = cast( Literal["local_file", "tool_file", "datasource_file", "remote_url"], transfer_method, ) try: - file_mapping = AgentStubFileMapping( + return AgentStubFileMapping( transfer_method=normalized_transfer_method, url=reference_or_url if normalized_transfer_method == "remote_url" else None, reference=reference_or_url if normalized_transfer_method != "remote_url" else None, @@ -114,17 +166,6 @@ def download_file_from_environment( except ValidationError as exc: raise AgentStubValidationError("invalid file download arguments") from exc - download_request = request_agent_stub_file_download_sync( - url=environment.url, - auth_jwe=environment.auth_jwe, - file=file_mapping, - ) - target_dir = Path(directory).expanduser().resolve() if directory else Path.cwd() - target_dir.mkdir(parents=True, exist_ok=True) - destination = _deduplicate_destination_path(target_dir / _sanitize_download_filename(download_request.filename)) - destination.write_bytes(download_file_bytes_from_signed_url_sync(download_url=download_request.download_url)) - return DownloadedFileResult(path=destination) - def _normalize_uploaded_tool_file_resource(payload: dict[str, object]) -> UploadedToolFileResource: reference = payload.get("reference") diff --git a/dify-agent/src/dify_agent/agent_stub/cli/main.py b/dify-agent/src/dify_agent/agent_stub/cli/main.py index c5c722d7e03..3d34b96d62e 100644 --- a/dify-agent/src/dify_agent/agent_stub/cli/main.py +++ b/dify-agent/src/dify_agent/agent_stub/cli/main.py @@ -11,13 +11,16 @@ does not pull in FastAPI, Redis, shellctl, or JWE runtime dependencies. from __future__ import annotations import sys +from typing import cast import typer from typer.main import get_command from dify_agent.agent_stub.cli._agent_stub import connect_from_environment from dify_agent.agent_stub.cli._drive import ( - list_drive_from_environment, + DrivePushKind, + format_drive_manifest, + list_drive_manifest_from_environment, pull_drive_from_environment, push_drive_from_environment, ) @@ -60,21 +63,23 @@ def upload(path: str = typer.Argument(..., metavar="PATH")) -> None: @file_app.command("download") def download( - transfer_method: str = typer.Argument(..., metavar="TRANSFER_METHOD"), - reference_or_url: str = typer.Argument(..., metavar="REFERENCE_OR_URL"), - directory: str | None = typer.Argument(default=None, metavar="DIR"), + transfer_method: str | None = typer.Argument(None, metavar="TRANSFER_METHOD"), + reference_or_url: str | None = typer.Argument(None, metavar="REFERENCE_OR_URL"), + mapping: str | None = typer.Option(None, "--mapping", help="Download one file from a mapping JSON object."), + local_dir: str | None = typer.Option(None, "--to", help="Local directory for the downloaded file."), ) -> None: """Download one workflow file mapping into the local sandbox directory.""" _run_file_download( transfer_method=transfer_method, reference_or_url=reference_or_url, - directory=directory, + mapping=mapping, + local_dir=local_dir, ) @drive_app.command("list") def drive_list( - path_prefix: str = typer.Argument("", metavar="PATH_PREFIX"), + path_prefix: str = typer.Argument("", metavar="REMOTE_PREFIX"), json_output: bool = typer.Option(False, "--json", help="Emit the drive manifest as JSON."), ) -> None: """List drive files visible to the current sandbox execution.""" @@ -83,32 +88,39 @@ def drive_list( @drive_app.command("pull") def drive_pull( - targets: list[str] = typer.Argument(None, metavar="TARGET"), - drive_base: str | None = typer.Option( + targets: list[str] = typer.Argument(None, metavar="REMOTE"), + local_base: str | None = typer.Option( None, - "--drive-base", + "--to", help=( f"Local base directory for pulled drive files. Defaults to ${AGENT_STUB_DRIVE_BASE_ENV_VAR} " f"or {DEFAULT_AGENT_STUB_DRIVE_BASE}." ), ), + json_output: bool = typer.Option(False, "--json", help="Emit the pull result as JSON."), ) -> None: """Pull one or more drive keys/prefixes into one local directory tree. Passing no ``TARGET`` preserves the historical whole-drive behavior by pulling from the empty prefix. """ - _run_drive_pull(targets=targets, drive_base=drive_base) + _run_drive_pull(targets=targets or None, local_base=local_base, json_output=json_output) @drive_app.command("push") def drive_push( local_path: str = typer.Argument(..., metavar="LOCAL_PATH"), - drive_path: str = typer.Argument(..., metavar="DRIVE_PATH"), - recursive: bool = typer.Option(False, "-r", "--recursive", help="Recursively upload directory contents."), + drive_path: str = typer.Argument(..., metavar="REMOTE_PATH"), + kind: str | None = typer.Option(None, "--kind", help="Directory upload kind: skill or dir."), + json_output: bool = typer.Option( + False, + "--json", + help="Accepted for consistency; drive push output is already emitted as JSON.", + ), ) -> None: """Upload one local file or directory into the agent drive.""" - _run_drive_push(local_path=local_path, drive_path=drive_path, recursive=recursive) + del json_output + _run_drive_push(local_path=local_path, drive_path=drive_path, kind=kind) def main(argv: list[str] | None = None) -> None: @@ -190,12 +202,19 @@ def _run_file_upload(*, path: str) -> None: typer.echo(response.model_dump_json()) -def _run_file_download(*, transfer_method: str, reference_or_url: str, directory: str | None) -> None: +def _run_file_download( + *, + transfer_method: str | None, + reference_or_url: str | None, + mapping: str | None, + local_dir: str | None, +) -> None: try: response = download_file_from_environment( transfer_method=transfer_method, reference_or_url=reference_or_url, - directory=directory, + mapping=mapping, + local_dir=local_dir, ) except MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) @@ -208,7 +227,7 @@ def _run_file_download(*, transfer_method: str, reference_or_url: str, directory def _run_drive_list(*, path_prefix: str, json_output: bool) -> None: try: - response = list_drive_from_environment(prefix=path_prefix, json_output=json_output) + response = list_drive_manifest_from_environment(prefix=path_prefix) except MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc @@ -216,29 +235,34 @@ def _run_drive_list(*, path_prefix: str, json_output: bool) -> None: typer.echo(str(exc), err=True) raise SystemExit(1) from exc if json_output: - if isinstance(response, str): - raise RuntimeError("drive list JSON output expected a manifest response") typer.echo(response.model_dump_json()) return - typer.echo(response) + typer.echo(format_drive_manifest(response)) -def _run_drive_pull(*, targets: list[str] | None, drive_base: str | None) -> None: +def _run_drive_pull(*, targets: list[str] | None, local_base: str | None, json_output: bool) -> None: try: - response = pull_drive_from_environment(targets=targets, drive_base=drive_base or read_agent_stub_drive_base()) + response = pull_drive_from_environment(targets=targets, local_base=local_base or read_agent_stub_drive_base()) except MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc except AgentStubClientError as exc: typer.echo(str(exc), err=True) raise SystemExit(1) from exc - for path in response: - typer.echo(str(path)) + if json_output: + typer.echo(response.model_dump_json()) + return + for item in response.items: + typer.echo(item.local_path) -def _run_drive_push(*, local_path: str, drive_path: str, recursive: bool) -> None: +def _run_drive_push(*, local_path: str, drive_path: str, kind: str | None) -> None: try: - response = push_drive_from_environment(local_path=local_path, drive_path=drive_path, recursive=recursive) + response = push_drive_from_environment( + local_path=local_path, + drive_path=drive_path, + kind=cast(DrivePushKind | None, kind), + ) except MissingAgentStubEnvironmentError as exc: typer.echo(str(exc), err=True) raise SystemExit(2) from exc diff --git a/dify-agent/src/dify_agent/layers/drive/layer.py b/dify-agent/src/dify_agent/layers/drive/layer.py index 29ec996c402..6812bf4900b 100644 --- a/dify-agent/src/dify_agent/layers/drive/layer.py +++ b/dify-agent/src/dify_agent/layers/drive/layer.py @@ -15,41 +15,44 @@ from __future__ import annotations import asyncio from dataclasses import dataclass, field -from pathlib import Path, PurePosixPath -from tempfile import TemporaryDirectory +from pathlib import Path from typing import Any, ClassVar, cast -from uuid import uuid4 -from zipfile import BadZipFile, ZipFile, ZipInfo import httpx from typing_extensions import Self, override from agenton.layers import EmptyRuntimeState, Layer, LayerDeps, PlainLayer +from dify_agent.agent_stub._drive_materialization import ( + DriveDownloadPayload, + DriveMaterializationTransferError, + DriveMaterializationValidationError, + materialize_drive_downloads, +) from dify_agent.agent_stub.protocol import agent_stub_drive_base_for_ref +from dify_agent.agent_stub.protocol.agent_stub import AgentStubDriveItem, AgentStubDriveManifestResponse from dify_agent.layers.drive.configs import DIFY_DRIVE_LAYER_TYPE_ID, DifyDriveLayerConfig -_SKILL_ARCHIVE_FILENAME = ".DIFY-SKILL-FULL.zip" _DOWNLOAD_CONCURRENCY = 4 _AGENT_STUB_CLI_USAGE_PROMPT = """Agent Stub CLI usage is available inside shell jobs: -Drive commands: +Drive assets are Agent Soul versioned assets: -- List drive items: `dify-agent drive list [PATH_PREFIX]` -- Emit the drive manifest as JSON: `dify-agent drive list [PATH_PREFIX] --json` -- Pull drive keys or prefixes: `dify-agent drive pull TARGET ...` - Pulled files are written under `$DIFY_AGENT_STUB_DRIVE_BASE` by default. - Use `--drive-base .` to materialize pulled files under the current working directory. -- Upload a local file or directory: `dify-agent drive push LOCAL_PATH DRIVE_PATH` - Add `--recursive` to upload raw directory contents. Without `--recursive`, a directory must contain `SKILL.md` - and is uploaded as a standardized skill. +- List drive assets: `dify-agent drive list [REMOTE_PREFIX]` +- Pull drive assets: `dify-agent drive pull [REMOTE ...] [--to LOCAL_DIR]` + With no remote, pulls the whole visible drive. Pull overwrites local files. + Defaults to `$DIFY_AGENT_STUB_DRIVE_BASE`; use `--to .` for cwd. + `--to` is a local root; remote keys keep their path under it. + Skill archives are automatically extracted after pull. +- Push one file: `dify-agent drive push LOCAL_FILE REMOTE_PATH` +- Push a skill package: `dify-agent drive push LOCAL_DIR REMOTE_PATH --kind skill` +- Push a raw directory: `dify-agent drive push LOCAL_DIR REMOTE_PATH --kind dir` -File commands: +Workflow file mappings: -- Download one workflow file mapping: `dify-agent file download TRANSFER_METHOD REFERENCE_OR_URL [DIR]` - `TRANSFER_METHOD` is one of `local_file`, `tool_file`, `datasource_file`, or `remote_url`. - If `DIR` is omitted, the file is saved in the current working directory. -- Upload one sandbox-local output file: `dify-agent file upload PATH` - The command prints a JSON file mapping such as `{"transfer_method":"tool_file","reference":"..."}`.""" +- Download a mapping: `dify-agent file download TRANSFER_METHOD REFERENCE_OR_URL [--to LOCAL_DIR]` +- Or pass a mapping object: `dify-agent file download --mapping '{"transfer_method":"tool_file","reference":"..."}'` +- Upload an output file: `dify-agent file upload PATH` + Prints JSON like `{"transfer_method":"tool_file","reference":"..."}`.""" class DifyDriveLayerError(RuntimeError): @@ -59,14 +62,6 @@ class DifyDriveLayerError(RuntimeError): class DifyDriveDeps(LayerDeps): execution_context: Layer[Any, Any, Any, Any, Any, Any] # pyright: ignore[reportUninitializedInstanceVariable] - -@dataclass(frozen=True, slots=True) -class _DriveManifestItem: - key: str - download_url: str - size: int | None = None - - @dataclass(slots=True) class DifyDriveLayer(PlainLayer[DifyDriveDeps, DifyDriveLayerConfig, EmptyRuntimeState]): """Drive runtime layer that eagerly materializes prompt-mentioned drive targets.""" @@ -182,12 +177,12 @@ class DifyDriveLayer(PlainLayer[DifyDriveDeps, DifyDriveLayerConfig, EmptyRuntim *, tenant_id: str, targets: list[tuple[str, bool]], - ) -> list[_DriveManifestItem]: + ) -> list[AgentStubDriveItem]: semaphore = asyncio.Semaphore(_DOWNLOAD_CONCURRENCY) async with httpx.AsyncClient(timeout=30.0, follow_redirects=True, trust_env=False) as client: - async def fetch_one(target: tuple[str, bool]) -> list[_DriveManifestItem]: + async def fetch_one(target: tuple[str, bool]) -> list[AgentStubDriveItem]: prefix, exact = target try: async with semaphore: @@ -205,74 +200,58 @@ class DifyDriveLayer(PlainLayer[DifyDriveDeps, DifyDriveLayerConfig, EmptyRuntim if response.is_error: raise DifyDriveLayerError(f"drive manifest request failed for {prefix}: {response.status_code}") try: - payload = response.json() - except ValueError as exc: + payload = AgentStubDriveManifestResponse.model_validate(response.json()) + except (ValueError, TypeError) as exc: raise DifyDriveLayerError(f"drive manifest response is invalid for {prefix}") from exc - items = payload.get("items") if isinstance(payload, dict) else None - if not isinstance(items, list): - raise DifyDriveLayerError(f"drive manifest response is invalid for {prefix}") - manifest_items: list[_DriveManifestItem] = [] - for item in items: - if not isinstance(item, dict): - continue - key = item.get("key") - download_url = item.get("download_url") - if not isinstance(key, str) or not isinstance(download_url, str) or not download_url: + manifest_items: list[AgentStubDriveItem] = [] + for item in payload.items: + if not item.download_url: raise DifyDriveLayerError(f"drive manifest item is missing download_url for {prefix}") - if exact and key != prefix: + if exact and item.key != prefix: continue - manifest_items.append(_DriveManifestItem(key=key, download_url=download_url, size=item.get("size"))) + manifest_items.append(item) return manifest_items grouped_items = await asyncio.gather(*(fetch_one(target) for target in targets)) - deduplicated: dict[str, _DriveManifestItem] = {} + deduplicated: dict[str, AgentStubDriveItem] = {} for items in grouped_items: for item in items: deduplicated.setdefault(item.key, item) return [deduplicated[key] for key in sorted(deduplicated)] - async def _download_items(self, items: list[_DriveManifestItem]) -> dict[str, str]: + async def _download_items(self, items: list[AgentStubDriveItem]) -> dict[str, str]: base_path = Path(agent_stub_drive_base_for_ref(self.config.drive_ref)) - try: - base_path.mkdir(parents=True, exist_ok=True) - except OSError as exc: - raise DifyDriveLayerError(f"failed to prepare drive base {base_path}") from exc semaphore = asyncio.Semaphore(_DOWNLOAD_CONCURRENCY) - archive_paths: list[Path] = [] canonical_skill_dirs = {item.key.rsplit("/", 1)[0] for item in items if item.key.endswith("/SKILL.md")} async with httpx.AsyncClient(timeout=30.0, follow_redirects=True, trust_env=False) as client: - async def download_one(item: _DriveManifestItem) -> tuple[str, str]: + async def download_one(item: AgentStubDriveItem) -> DriveDownloadPayload: + download_url = item.download_url + if not download_url: + raise DifyDriveLayerError(f"drive manifest item is missing download_url for {item.key}") try: async with semaphore: - response = await client.get(item.download_url) + response = await client.get(download_url) except (httpx.InvalidURL, httpx.TimeoutException, httpx.RequestError) as exc: raise DifyDriveLayerError(f"drive download failed for {item.key}") from exc if response.is_error: raise DifyDriveLayerError(f"drive download failed for {item.key}: {response.status_code}") - payload = response.content - if item.size is not None and len(payload) != item.size: - raise DifyDriveLayerError(f"downloaded drive file size mismatch for {item.key}") - try: - destination = _resolve_drive_destination(base_path, item.key) - destination.parent.mkdir(parents=True, exist_ok=True) - temp_path = destination.with_name(f"{destination.name}.tmp-{uuid4().hex}") - temp_path.write_bytes(payload) - temp_path.replace(destination) - except OSError as exc: - raise DifyDriveLayerError(f"failed to materialize drive file {item.key}") from exc - if destination.name == _SKILL_ARCHIVE_FILENAME: - archive_paths.append(destination) - return item.key, str(destination) + return DriveDownloadPayload(key=item.key, payload=response.content, size=item.size) - pairs = await asyncio.gather(*(download_one(item) for item in items)) - for archive_path in sorted(archive_paths): - archive_skill_dir = archive_path.parent.relative_to(base_path).as_posix() - skip_entry_names = {"SKILL.md"} if archive_skill_dir in canonical_skill_dirs else set() - _extract_skill_archive(archive_path, skip_entry_names=skip_entry_names) - return {key: path for key, path in pairs} + downloads = await asyncio.gather(*(download_one(item) for item in items)) + + try: + written_paths = materialize_drive_downloads( + base_path=base_path, + downloads=downloads, + archive_skip_entry_names_by_dir={skill_dir: {"SKILL.md"} for skill_dir in canonical_skill_dirs}, + ) + except (DriveMaterializationValidationError, DriveMaterializationTransferError) as exc: + raise DifyDriveLayerError(str(exc)) from exc + + return {download.key: str(path) for download, path in zip(downloads, written_paths, strict=True)} def _require_tenant_id(self) -> str: execution_context = self.deps.execution_context.config @@ -286,68 +265,4 @@ class DifyDriveLayer(PlainLayer[DifyDriveDeps, DifyDriveLayerConfig, EmptyRuntim return f"{skill_key.rsplit('/', 1)[0]}/" -def _resolve_drive_destination(base_path: Path, drive_key: str) -> Path: - destination = (base_path / Path(drive_key)).resolve() - try: - destination.relative_to(base_path) - except ValueError as exc: - raise DifyDriveLayerError(f"drive key resolves outside the drive base: {drive_key}") from exc - return destination - - -def _extract_skill_archive(archive_path: Path, *, skip_entry_names: set[str]) -> None: - target_dir = archive_path.parent.resolve() - try: - with TemporaryDirectory(dir=target_dir, prefix=".dify-skill-extract-") as staging_dir_name: - staging_dir = Path(staging_dir_name).resolve() - with ZipFile(archive_path) as archive: - for zip_info in archive.infolist(): - if zip_info.filename.replace("\\", "/").rstrip("/") in skip_entry_names: - continue - destination = _resolve_zip_entry_destination(staging_dir, zip_info.filename) - if _is_zip_symlink(zip_info): - raise DifyDriveLayerError( - f"skill archive contains unsupported symlink entry: {zip_info.filename}" - ) - if zip_info.is_dir(): - destination.mkdir(parents=True, exist_ok=True) - continue - destination.parent.mkdir(parents=True, exist_ok=True) - with archive.open(zip_info) as source_file: - temp_path = destination.with_name(f"{destination.name}.tmp-{uuid4().hex}") - temp_path.write_bytes(source_file.read()) - temp_path.replace(destination) - for staged_path in sorted(staging_dir.rglob("*")): - if staged_path.is_dir(): - continue - relative_path = staged_path.relative_to(staging_dir) - destination = (target_dir / relative_path).resolve() - destination.parent.mkdir(parents=True, exist_ok=True) - staged_path.replace(destination) - except DifyDriveLayerError: - raise - except (BadZipFile, OSError) as exc: - raise DifyDriveLayerError(f"downloaded skill archive is invalid: {archive_path.name}") from exc - - -def _resolve_zip_entry_destination(target_dir: Path, entry_name: str) -> Path: - normalized_name = entry_name.replace("\\", "/") - pure_path = PurePosixPath(normalized_name) - if not normalized_name or normalized_name.startswith("/") or pure_path.is_absolute(): - raise DifyDriveLayerError(f"skill archive contains unsafe absolute path: {entry_name}") - if any(part in {"", ".", ".."} for part in pure_path.parts): - raise DifyDriveLayerError(f"skill archive contains unsafe path traversal entry: {entry_name}") - destination = (target_dir / Path(*pure_path.parts)).resolve() - try: - destination.relative_to(target_dir) - except ValueError as exc: - raise DifyDriveLayerError(f"skill archive entry resolves outside the skill directory: {entry_name}") from exc - return destination - - -def _is_zip_symlink(zip_info: ZipInfo) -> bool: - file_mode = zip_info.external_attr >> 16 - return (file_mode & 0o170000) == 0o120000 - - __all__ = ["DifyDriveLayer", "DifyDriveLayerError"] diff --git a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py index 6e45e14a034..c1773165bc1 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py @@ -8,7 +8,9 @@ from zipfile import ZipFile, ZipInfo import pytest from dify_agent.agent_stub.cli._drive import ( - list_drive_from_environment, + DrivePullResult, + format_drive_manifest, + list_drive_manifest_from_environment, pull_drive_from_environment, push_drive_from_environment, ) @@ -22,7 +24,7 @@ from dify_agent.agent_stub.protocol.agent_stub import ( ) -def test_list_drive_from_environment_returns_manifest_json_model(monkeypatch: pytest.MonkeyPatch) -> None: +def test_list_drive_manifest_from_environment_returns_manifest_model(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") captured: dict[str, object] = {} @@ -47,7 +49,7 @@ def test_list_drive_from_environment_returns_manifest_json_model(monkeypatch: py fake_manifest, ) - result = list_drive_from_environment(prefix="skills/", json_output=True) + result = list_drive_manifest_from_environment(prefix="skills/") assert isinstance(result, AgentStubDriveManifestResponse) assert result.items[0].key == "skills/example/SKILL.md" @@ -55,7 +57,7 @@ def test_list_drive_from_environment_returns_manifest_json_model(monkeypatch: py assert captured["include_download_url"] is False -def test_list_drive_from_environment_returns_human_readable_listing(monkeypatch: pytest.MonkeyPatch) -> None: +def test_format_drive_manifest_returns_human_readable_listing(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") captured: dict[str, object] = {} @@ -88,7 +90,9 @@ def test_list_drive_from_environment_returns_human_readable_listing(monkeypatch: fake_manifest, ) - result = list_drive_from_environment(prefix="skills/", json_output=False) + result = format_drive_manifest( + list_drive_manifest_from_environment(prefix="skills/") + ) assert result == ("12\ttext/markdown\t-\tskills/example/SKILL.md\n-\t-\tsha256:abc\tskills/example/helper.py") assert captured["prefix"] == "skills/" @@ -128,10 +132,12 @@ def test_pull_drive_from_environment_writes_files_under_drive_base( lambda **_kwargs: b"hello world", ) - results = pull_drive_from_environment(targets=["skills/"], drive_base=str(tmp_path)) + result = pull_drive_from_environment(targets=["skills/"], local_base=str(tmp_path)) - assert results == [tmp_path / "skills" / "example" / "SKILL.md"] - assert results[0].read_bytes() == b"hello world" + assert result.model_dump() == { + "items": [{"key": "skills/example/SKILL.md", "local_path": str(tmp_path / "skills" / "example" / "SKILL.md")}] + } + assert Path(result.items[0].local_path).read_bytes() == b"hello world" assert captured["prefix"] == "skills/" assert captured["include_download_url"] is True @@ -169,10 +175,12 @@ def test_pull_drive_from_environment_auto_extracts_skill_archive( lambda **_kwargs: archive_bytes, ) - results = pull_drive_from_environment(targets=["skills/foo"], drive_base=str(tmp_path)) + result = pull_drive_from_environment(targets=["skills/foo"], local_base=str(tmp_path)) archive_path = tmp_path / "skills" / "foo" / ".DIFY-SKILL-FULL.zip" - assert results == [archive_path] + assert result.model_dump() == { + "items": [{"key": "skills/foo/.DIFY-SKILL-FULL.zip", "local_path": str(archive_path)}] + } assert archive_path.read_bytes() == archive_bytes assert (tmp_path / "skills" / "foo" / "SKILL.md").read_text(encoding="utf-8") == "# Example\n" assert (tmp_path / "skills" / "foo" / "nested" / "helper.py").read_text(encoding="utf-8") == "print('x')\n" @@ -202,7 +210,7 @@ def test_pull_drive_from_environment_rejects_traversal_keys( ) with pytest.raises(AgentStubValidationError, match="outside the drive base"): - _ = pull_drive_from_environment(targets=[""], drive_base=str(tmp_path)) + _ = pull_drive_from_environment(targets=[""], local_base=str(tmp_path)) def test_pull_drive_from_environment_rejects_skill_archive_path_traversal( @@ -239,7 +247,7 @@ def test_pull_drive_from_environment_rejects_skill_archive_path_traversal( ) with pytest.raises(AgentStubValidationError, match="path traversal"): - _ = pull_drive_from_environment(targets=["skills/foo"], drive_base=str(tmp_path)) + _ = pull_drive_from_environment(targets=["skills/foo"], local_base=str(tmp_path)) assert not (tmp_path / "skills" / "foo" / "SKILL.md").exists() @@ -276,7 +284,7 @@ def test_pull_drive_from_environment_rejects_skill_archive_absolute_entry( ) with pytest.raises(AgentStubValidationError, match="absolute path"): - _ = pull_drive_from_environment(targets=["skills/foo"], drive_base=str(tmp_path)) + _ = pull_drive_from_environment(targets=["skills/foo"], local_base=str(tmp_path)) def test_pull_drive_from_environment_rejects_skill_archive_symlink_entry( @@ -314,7 +322,7 @@ def test_pull_drive_from_environment_rejects_skill_archive_symlink_entry( ) with pytest.raises(AgentStubValidationError, match="symlink entry"): - _ = pull_drive_from_environment(targets=["skills/foo"], drive_base=str(tmp_path)) + _ = pull_drive_from_environment(targets=["skills/foo"], local_base=str(tmp_path)) def test_pull_drive_from_environment_rejects_invalid_skill_archive( @@ -347,7 +355,7 @@ def test_pull_drive_from_environment_rejects_invalid_skill_archive( ) with pytest.raises(AgentStubTransferError, match="downloaded skill archive is invalid"): - _ = pull_drive_from_environment(targets=["skills/foo"], drive_base=str(tmp_path)) + _ = pull_drive_from_environment(targets=["skills/foo"], local_base=str(tmp_path)) def test_pull_drive_from_environment_rejects_missing_download_url( @@ -373,7 +381,7 @@ def test_pull_drive_from_environment_rejects_missing_download_url( ) with pytest.raises(AgentStubValidationError, match="missing download_url"): - _ = pull_drive_from_environment(targets=["skills/"], drive_base=str(tmp_path)) + _ = pull_drive_from_environment(targets=["skills/"], local_base=str(tmp_path)) def test_pull_drive_from_environment_rejects_size_mismatch( @@ -404,7 +412,7 @@ def test_pull_drive_from_environment_rejects_size_mismatch( ) with pytest.raises(AgentStubTransferError, match="size mismatch"): - _ = pull_drive_from_environment(targets=["skills/"], drive_base=str(tmp_path)) + _ = pull_drive_from_environment(targets=["skills/"], local_base=str(tmp_path)) def test_pull_drive_from_environment_requests_multiple_targets_and_deduplicates_overlaps( @@ -463,11 +471,16 @@ def test_pull_drive_from_environment_requests_multiple_targets_and_deduplicates_ ), ) - results = pull_drive_from_environment(targets=["skills/foo", "files/a.txt"], drive_base=str(tmp_path)) + result = pull_drive_from_environment(targets=["skills/foo", "files/a.txt"], local_base=str(tmp_path)) - assert captured_prefixes == ["skills/foo", "files/a.txt"] - assert results == [tmp_path / "files" / "a.txt", tmp_path / "skills" / "foo" / "SKILL.md"] - assert downloaded_urls == ["https://files.example.com/a-txt", "https://files.example.com/skill-md"] + assert set(captured_prefixes) == {"skills/foo", "files/a.txt"} + assert len(captured_prefixes) == 2 + assert {(item.key, item.local_path) for item in result.items} == { + ("files/a.txt", str(tmp_path / "files" / "a.txt")), + ("skills/foo/SKILL.md", str(tmp_path / "skills" / "foo" / "SKILL.md")), + } + assert set(downloaded_urls) == {"https://files.example.com/a-txt", "https://files.example.com/skill-md"} + assert len(downloaded_urls) == 2 def test_pull_drive_from_environment_without_targets_preserves_whole_drive_pull( @@ -482,10 +495,46 @@ def test_pull_drive_from_environment_without_targets_preserves_whole_drive_pull( lambda **kwargs: captured_prefixes.append(kwargs["prefix"]) or AgentStubDriveManifestResponse(items=[]), ) - assert pull_drive_from_environment(drive_base=str(tmp_path)) == [] + assert pull_drive_from_environment(local_base=str(tmp_path)).model_dump() == {"items": []} assert captured_prefixes == [""] +def test_pull_drive_from_environment_returns_json_result( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") + monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") + monkeypatch.setattr( + "dify_agent.agent_stub.cli._drive.request_agent_stub_drive_manifest_sync", + lambda **_kwargs: AgentStubDriveManifestResponse( + items=[ + AgentStubDriveItem( + key="files/a.txt", + size=1, + hash=None, + mime_type="text/plain", + file_kind="tool_file", + file_id="tool-file-1", + download_url="https://files.example.com/a-txt", + ) + ] + ), + ) + monkeypatch.setattr( + "dify_agent.agent_stub.cli._drive.download_file_bytes_from_signed_url_sync", + lambda **_kwargs: b"a", + ) + + result = pull_drive_from_environment(targets=["files/a.txt"], local_base=str(tmp_path)) + + assert isinstance(result, DrivePullResult) + assert result.model_dump() == { + "items": [{"key": "files/a.txt", "local_path": str(tmp_path / "files" / "a.txt")}] + } + assert (tmp_path / "files" / "a.txt").read_bytes() == b"a" + + def test_push_drive_from_environment_commits_single_file(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: source = tmp_path / "report.pdf" source.write_bytes(b"report") @@ -518,21 +567,46 @@ def test_push_drive_from_environment_commits_single_file(monkeypatch: pytest.Mon monkeypatch.setattr("dify_agent.agent_stub.cli._drive.request_agent_stub_drive_commit_sync", fake_commit) - response = push_drive_from_environment(local_path=str(source), drive_path="files/report.pdf", recursive=False) + response = push_drive_from_environment(local_path=str(source), drive_path="files/report.pdf", kind=None) assert response.items[0].key == "files/report.pdf" request = captured["request"] assert isinstance(request, AgentStubDriveCommitRequest) - assert request.items[0].model_dump(mode="json") == { - "key": "files/report.pdf", - "file_ref": {"kind": "tool_file", "id": "tool-file-1"}, - "value_owned_by_drive": True, - "is_skill": False, - "skill_metadata": None, - } + assert request.items[0].key == "files/report.pdf" + assert request.items[0].file_ref is not None + assert request.items[0].file_ref.kind == "tool_file" + assert request.items[0].file_ref.id == "tool-file-1" -def test_push_drive_from_environment_requires_skill_md_for_non_recursive_directory( +def test_push_drive_from_environment_rejects_file_with_kind_skill( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + source = tmp_path / "report.pdf" + source.write_bytes(b"report") + monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") + monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") + + with pytest.raises(AgentStubValidationError, match="--kind skill requires a directory containing SKILL.md"): + _ = push_drive_from_environment(local_path=str(source), drive_path="files/report.pdf", kind="skill") + + +def test_push_drive_from_environment_rejects_symlinked_file_root( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + source = tmp_path / "report.pdf" + source.write_bytes(b"report") + symlink_path = tmp_path / "report-link.pdf" + symlink_path.symlink_to(source) + monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") + monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") + + with pytest.raises(AgentStubValidationError, match="symlink"): + _ = push_drive_from_environment(local_path=str(symlink_path), drive_path="files/report.pdf", kind=None) + + +def test_push_drive_from_environment_requires_kind_for_directory( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -541,11 +615,24 @@ def test_push_drive_from_environment_requires_skill_md_for_non_recursive_directo monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") - with pytest.raises(AgentStubValidationError, match="SKILL.md"): - _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", recursive=False) + with pytest.raises(AgentStubValidationError, match="requires --kind skill or --kind dir"): + _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind=None) -def test_push_drive_from_environment_standardizes_non_recursive_skill_directory( +def test_push_drive_from_environment_kind_skill_requires_skill_md( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + skill_dir = tmp_path / "skill" + skill_dir.mkdir() + monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") + monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") + + with pytest.raises(AgentStubValidationError, match="requires a directory containing SKILL.md"): + _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind="skill") + + +def test_push_drive_from_environment_kind_skill_standardizes_skill_directory( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -584,7 +671,7 @@ def test_push_drive_from_environment_standardizes_non_recursive_skill_directory( ), ) - response = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", recursive=False) + response = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind="skill") assert set(uploaded_paths) == {"SKILL.md", ".DIFY-SKILL-FULL.zip"} assert {item.key for item in response.items} == { @@ -593,7 +680,7 @@ def test_push_drive_from_environment_standardizes_non_recursive_skill_directory( } -def test_push_drive_from_environment_non_recursive_archive_excludes_transient_entries( +def test_push_drive_from_environment_kind_skill_archive_excludes_transient_entries( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -641,7 +728,7 @@ def test_push_drive_from_environment_non_recursive_archive_excludes_transient_en ), ) - _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", recursive=False) + _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind="skill") assert {"SKILL.md", "helper.py"}.issubset(archive_entries) assert ".git/config" not in archive_entries @@ -649,7 +736,7 @@ def test_push_drive_from_environment_non_recursive_archive_excludes_transient_en assert ".DIFY-SKILL-FULL.zip" not in archive_entries -def test_push_drive_from_environment_non_recursive_rejects_symlinked_archive_entries( +def test_push_drive_from_environment_kind_skill_rejects_symlinked_archive_entries( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -663,10 +750,52 @@ def test_push_drive_from_environment_non_recursive_rejects_symlinked_archive_ent monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") with pytest.raises(AgentStubValidationError, match="symlink"): - _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", recursive=False) + _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind="skill") -def test_push_drive_from_environment_rejects_symlinked_recursive_files( +def test_push_drive_from_environment_kind_dir_requires_directory( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + source = tmp_path / "report.pdf" + source.write_bytes(b"report") + monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") + monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") + + with pytest.raises(AgentStubValidationError, match="--kind dir requires a directory"): + _ = push_drive_from_environment(local_path=str(source), drive_path="files/report.pdf", kind="dir") + + +def test_push_drive_from_environment_kind_file_requires_file( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + skill_dir = tmp_path / "skill" + skill_dir.mkdir() + monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") + monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") + + with pytest.raises(AgentStubValidationError, match="--kind file requires a file"): + _ = push_drive_from_environment(local_path=str(skill_dir), drive_path="skills/example", kind="file") + + +def test_push_drive_from_environment_rejects_symlinked_directory_root( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + source_dir = tmp_path / "skill" + source_dir.mkdir() + (source_dir / "SKILL.md").write_text("# Example\n", encoding="utf-8") + symlink_path = tmp_path / "skill-link" + symlink_path.symlink_to(source_dir, target_is_directory=True) + monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") + monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") + + with pytest.raises(AgentStubValidationError, match="symlink"): + _ = push_drive_from_environment(local_path=str(symlink_path), drive_path="skills/example", kind="skill") + + +def test_push_drive_from_environment_kind_dir_rejects_symlinked_files( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -679,10 +808,10 @@ def test_push_drive_from_environment_rejects_symlinked_recursive_files( monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") with pytest.raises(AgentStubValidationError, match="symlink"): - _ = push_drive_from_environment(local_path=str(root), drive_path="skills/example", recursive=True) + _ = push_drive_from_environment(local_path=str(root), drive_path="skills/example", kind="dir") -def test_push_drive_from_environment_recursive_keeps_user_files_that_skill_packaging_skips( +def test_push_drive_from_environment_kind_dir_keeps_user_files_that_skill_packaging_skips( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -723,7 +852,7 @@ def test_push_drive_from_environment_recursive_keeps_user_files_that_skill_packa ), ) - response = push_drive_from_environment(local_path=str(root), drive_path="skills/example", recursive=True) + response = push_drive_from_environment(local_path=str(root), drive_path="skills/example", kind="dir") assert set(uploaded_paths) == {".DIFY-SKILL-FULL.zip", "node_modules/module.js"} assert {item.key for item in response.items} == { diff --git a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py index f2de6e8edc5..c2f2d44afb4 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py @@ -11,7 +11,7 @@ from dify_agent.agent_stub.cli._files import ( upload_file_from_environment, upload_tool_file_resource_from_environment, ) -from dify_agent.agent_stub.client._errors import AgentStubTransferError +from dify_agent.agent_stub.client._errors import AgentStubTransferError, AgentStubValidationError def _reference(record_id: str) -> str: @@ -120,7 +120,7 @@ def test_download_file_from_environment_saves_bytes_and_renames_on_collision( result = download_file_from_environment( transfer_method="tool_file", reference_or_url=_reference("tool-file-1"), - directory=str(target_dir), + local_dir=str(target_dir), ) assert result.path.name == "report (1).pdf" @@ -157,7 +157,7 @@ def test_download_file_from_environment_sanitizes_server_filename( result = download_file_from_environment( transfer_method="tool_file", reference_or_url=_reference("tool-file-1"), - directory=str(target_dir), + local_dir=str(target_dir), ) assert result.path.parent == target_dir @@ -187,6 +187,62 @@ def test_upload_file_from_environment_rejects_non_canonical_reference( _ = upload_file_from_environment(path=str(source)) +def test_download_file_from_environment_supports_mapping_json( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + target_dir = tmp_path / "inputs" + monkeypatch.setenv("DIFY_AGENT_STUB_API_BASE_URL", "https://agent.example.com/agent-stub") + monkeypatch.setenv("DIFY_AGENT_STUB_AUTH_JWE", "test-jwe") + captured: dict[str, object] = {} + + def fake_request_download(**kwargs): + captured["file"] = kwargs["file"] + return type( + "Response", + (), + { + "filename": "report.pdf", + "mime_type": "application/pdf", + "size": 12, + "download_url": "https://files.example.com/download", + }, + )() + + monkeypatch.setattr("dify_agent.agent_stub.cli._files.request_agent_stub_file_download_sync", fake_request_download) + monkeypatch.setattr( + "dify_agent.agent_stub.cli._files.download_file_bytes_from_signed_url_sync", + lambda **_kwargs: b"downloaded", + ) + + result = download_file_from_environment( + mapping=json.dumps({"transfer_method": "tool_file", "reference": _reference("tool-file-1")}), + local_dir=str(target_dir), + ) + + assert captured["file"].model_dump() == { + "transfer_method": "tool_file", + "reference": _reference("tool-file-1"), + "url": None, + } + assert result.path == target_dir / "report.pdf" + assert result.path.read_bytes() == b"downloaded" + + +def test_download_file_from_environment_requires_mapping_or_positional_pair() -> None: + with pytest.raises(AgentStubValidationError, match="requires either --mapping or TRANSFER_METHOD REFERENCE_OR_URL"): + _ = download_file_from_environment() + + +def test_download_file_from_environment_rejects_mapping_mixed_with_positionals() -> None: + with pytest.raises(AgentStubValidationError, match="cannot be combined"): + _ = download_file_from_environment( + transfer_method="tool_file", + reference_or_url=_reference("tool-file-1"), + mapping=json.dumps({"transfer_method": "tool_file", "reference": _reference("tool-file-1")}), + ) + + def test_upload_tool_file_resource_from_environment_rejects_missing_id( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py index d66b801778c..9044ba336a6 100644 --- a/dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py +++ b/dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py @@ -6,6 +6,7 @@ from pathlib import Path import pytest +from dify_agent.agent_stub.cli._drive import DrivePullResult from dify_agent.agent_stub.cli.main import main from dify_agent.agent_stub.protocol.agent_stub import ( AgentStubDriveCommitResponse, @@ -194,20 +195,77 @@ def test_cli_file_download_prints_saved_path( ) with pytest.raises(SystemExit) as exc_info: - main(["file", "download", "tool_file", _reference("tool-file-1"), "/tmp"]) + main(["file", "download", "tool_file", _reference("tool-file-1"), "--to", "/tmp"]) captured = capsys.readouterr() assert exc_info.value.code == 0 assert captured.out.strip() == "/tmp/report.pdf" +def test_cli_file_download_supports_mapping_json( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + captured_kwargs: dict[str, object] = {} + + def fake_download_file_from_environment(**kwargs): + captured_kwargs.update(kwargs) + return type("Response", (), {"path": Path("/tmp/inputs/report.pdf")})() + + monkeypatch.setattr("dify_agent.agent_stub.cli.main.download_file_from_environment", fake_download_file_from_environment) + + with pytest.raises(SystemExit) as exc_info: + main( + [ + "file", + "download", + "--mapping", + json.dumps({"transfer_method": "tool_file", "reference": _reference("tool-file-1")}), + "--to", + "/tmp/inputs", + ] + ) + + captured = capsys.readouterr() + assert exc_info.value.code == 0 + assert captured_kwargs == { + "transfer_method": None, + "reference_or_url": None, + "mapping": json.dumps({"transfer_method": "tool_file", "reference": _reference("tool-file-1")}), + "local_dir": "/tmp/inputs", + } + assert captured.out.strip() == "/tmp/inputs/report.pdf" + + +def test_cli_file_download_rejects_legacy_positional_directory( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + called = False + + def fake_download_file_from_environment(**_kwargs): + nonlocal called + called = True + return type("Response", (), {"path": Path("/tmp/report.pdf")})() + + monkeypatch.setattr("dify_agent.agent_stub.cli.main.download_file_from_environment", fake_download_file_from_environment) + + with pytest.raises(SystemExit) as exc_info: + main(["file", "download", "tool_file", _reference("tool-file-1"), "/tmp"]) + + captured = capsys.readouterr() + assert exc_info.value.code == 2 + assert called is False + assert "/tmp" in captured.err + + def test_cli_drive_list_prints_manifest_json( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.list_drive_from_environment", - lambda *, prefix, json_output: AgentStubDriveManifestResponse( + "dify_agent.agent_stub.cli.main.list_drive_manifest_from_environment", + lambda *, prefix: AgentStubDriveManifestResponse( items=[ AgentStubDriveItem( key=prefix + "example/SKILL.md", @@ -234,8 +292,19 @@ def test_cli_drive_list_prints_human_readable_listing( capsys: pytest.CaptureFixture[str], ) -> None: monkeypatch.setattr( - "dify_agent.agent_stub.cli.main.list_drive_from_environment", - lambda *, prefix, json_output: f"12\ttext/markdown\t-\t{prefix}example/SKILL.md", + "dify_agent.agent_stub.cli.main.list_drive_manifest_from_environment", + lambda *, prefix: AgentStubDriveManifestResponse( + items=[ + AgentStubDriveItem( + key=f"{prefix}example/SKILL.md", + size=12, + hash=None, + mime_type="text/markdown", + file_kind="tool_file", + file_id="tool-file-1", + ) + ] + ), ) with pytest.raises(SystemExit) as exc_info: @@ -252,14 +321,16 @@ def test_cli_drive_pull_prints_downloaded_paths( ) -> None: monkeypatch.setattr( "dify_agent.agent_stub.cli.main.pull_drive_from_environment", - lambda *, targets, drive_base: [ - Path(drive_base) / targets[0] / "SKILL.md", - Path(drive_base) / targets[0] / "helper.py", - ], + lambda *, targets, local_base: DrivePullResult( + items=[ + DrivePullResult.Item(key=f"{targets[0]}/SKILL.md", local_path=str(Path(local_base) / targets[0] / "SKILL.md")), + DrivePullResult.Item(key=f"{targets[0]}/helper.py", local_path=str(Path(local_base) / targets[0] / "helper.py")), + ] + ), ) with pytest.raises(SystemExit) as exc_info: - main(["drive", "pull", "skills/example", "--drive-base", "/tmp/drive"]) + main(["drive", "pull", "skills/example", "--to", "/tmp/drive"]) captured = capsys.readouterr() assert exc_info.value.code == 0 @@ -269,16 +340,45 @@ def test_cli_drive_pull_prints_downloaded_paths( ] +def test_cli_drive_pull_prints_json_result( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + monkeypatch.setattr( + "dify_agent.agent_stub.cli.main.pull_drive_from_environment", + lambda *, targets, local_base: DrivePullResult( + items=[ + DrivePullResult.Item(key="files/a.txt", local_path=f"{local_base}/files/a.txt"), + DrivePullResult.Item(key="skills/foo/SKILL.md", local_path=f"{local_base}/skills/foo/SKILL.md"), + ] + ), + ) + + with pytest.raises(SystemExit) as exc_info: + main(["drive", "pull", "files/a.txt", "--to", "/tmp/drive", "--json"]) + + captured = capsys.readouterr() + assert exc_info.value.code == 0 + assert json.loads(captured.out) == { + "items": [ + {"key": "files/a.txt", "local_path": "/tmp/drive/files/a.txt"}, + {"key": "skills/foo/SKILL.md", "local_path": "/tmp/drive/skills/foo/SKILL.md"}, + ] + } + + def test_cli_drive_pull_forwards_multiple_targets( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: captured_kwargs: dict[str, object] = {} - def fake_pull_drive_from_environment(*, targets, drive_base): + def fake_pull_drive_from_environment(*, targets, local_base): captured_kwargs["targets"] = targets - captured_kwargs["drive_base"] = drive_base - return [Path(drive_base) / "skills" / "foo" / "SKILL.md"] + captured_kwargs["local_base"] = local_base + return DrivePullResult( + items=[DrivePullResult.Item(key="skills/foo/SKILL.md", local_path=str(Path(local_base) / "skills" / "foo" / "SKILL.md"))] + ) monkeypatch.setattr( "dify_agent.agent_stub.cli.main.pull_drive_from_environment", @@ -286,11 +386,11 @@ def test_cli_drive_pull_forwards_multiple_targets( ) with pytest.raises(SystemExit) as exc_info: - main(["drive", "pull", "skills/foo", "files/a.txt", "--drive-base", "/tmp/drive"]) + main(["drive", "pull", "skills/foo", "files/a.txt", "--to", "/tmp/drive"]) captured = capsys.readouterr() assert exc_info.value.code == 0 - assert captured_kwargs == {"targets": ["skills/foo", "files/a.txt"], "drive_base": "/tmp/drive"} + assert captured_kwargs == {"targets": ["skills/foo", "files/a.txt"], "local_base": "/tmp/drive"} assert captured.out.strip() == "/tmp/drive/skills/foo/SKILL.md" @@ -301,10 +401,12 @@ def test_cli_drive_pull_uses_environment_drive_base_default( monkeypatch.setenv("DIFY_AGENT_STUB_DRIVE_BASE", "/env/drive") captured_kwargs: dict[str, object] = {} - def fake_pull_drive_from_environment(*, targets, drive_base): + def fake_pull_drive_from_environment(*, targets, local_base): captured_kwargs["targets"] = targets - captured_kwargs["drive_base"] = drive_base - return [Path(drive_base) / "skills" / "foo" / "SKILL.md"] + captured_kwargs["local_base"] = local_base + return DrivePullResult( + items=[DrivePullResult.Item(key="skills/foo/SKILL.md", local_path=str(Path(local_base) / "skills" / "foo" / "SKILL.md"))] + ) monkeypatch.setattr( "dify_agent.agent_stub.cli.main.pull_drive_from_environment", @@ -316,7 +418,7 @@ def test_cli_drive_pull_uses_environment_drive_base_default( captured = capsys.readouterr() assert exc_info.value.code == 0 - assert captured_kwargs == {"targets": ["skills/foo"], "drive_base": "/env/drive"} + assert captured_kwargs == {"targets": ["skills/foo"], "local_base": "/env/drive"} assert captured.out.strip() == "/env/drive/skills/foo/SKILL.md" @@ -327,10 +429,12 @@ def test_cli_drive_pull_keeps_historical_drive_base_when_env_is_missing( monkeypatch.delenv("DIFY_AGENT_STUB_DRIVE_BASE", raising=False) captured_kwargs: dict[str, object] = {} - def fake_pull_drive_from_environment(*, targets, drive_base): + def fake_pull_drive_from_environment(*, targets, local_base): captured_kwargs["targets"] = targets - captured_kwargs["drive_base"] = drive_base - return [Path(drive_base) / "skills" / "foo" / "SKILL.md"] + captured_kwargs["local_base"] = local_base + return DrivePullResult( + items=[DrivePullResult.Item(key="skills/foo/SKILL.md", local_path=str(Path(local_base) / "skills" / "foo" / "SKILL.md"))] + ) monkeypatch.setattr( "dify_agent.agent_stub.cli.main.pull_drive_from_environment", @@ -342,17 +446,42 @@ def test_cli_drive_pull_keeps_historical_drive_base_when_env_is_missing( captured = capsys.readouterr() assert exc_info.value.code == 0 - assert captured_kwargs == {"targets": ["skills/foo"], "drive_base": "/mnt/drive"} + assert captured_kwargs == {"targets": ["skills/foo"], "local_base": "/mnt/drive"} assert captured.out.strip() == "/mnt/drive/skills/foo/SKILL.md" +def test_cli_drive_pull_without_targets_pulls_whole_visible_drive( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + captured_kwargs: dict[str, object] = {} + + def fake_pull_drive_from_environment(*, targets, local_base): + captured_kwargs["targets"] = targets + captured_kwargs["local_base"] = local_base + return DrivePullResult(items=[DrivePullResult.Item(key="files/a.txt", local_path=str(Path(local_base) / "files" / "a.txt"))]) + + monkeypatch.setattr( + "dify_agent.agent_stub.cli.main.pull_drive_from_environment", + fake_pull_drive_from_environment, + ) + + with pytest.raises(SystemExit) as exc_info: + main(["drive", "pull", "--to", "/tmp/drive"]) + + captured = capsys.readouterr() + assert exc_info.value.code == 0 + assert captured_kwargs == {"targets": None, "local_base": "/tmp/drive"} + assert captured.out.strip() == "/tmp/drive/files/a.txt" + + def test_cli_drive_push_prints_commit_json( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], ) -> None: monkeypatch.setattr( "dify_agent.agent_stub.cli.main.push_drive_from_environment", - lambda *, local_path, drive_path, recursive: AgentStubDriveCommitResponse( + lambda *, local_path, drive_path, kind: AgentStubDriveCommitResponse( items=[ AgentStubDriveItem( key=drive_path, @@ -361,7 +490,7 @@ def test_cli_drive_push_prints_commit_json( mime_type="text/markdown", file_kind="tool_file", file_id=Path(local_path).name, - value_owned_by_drive=recursive is False, + value_owned_by_drive=kind != "dir", ) ] ), @@ -373,3 +502,87 @@ def test_cli_drive_push_prints_commit_json( captured = capsys.readouterr() assert exc_info.value.code == 0 assert json.loads(captured.out)["items"][0]["key"] == "skills/example/SKILL.md" + + +def test_cli_drive_push_forwards_kind( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + captured_kwargs: dict[str, object] = {} + + def fake_push_drive_from_environment(*, local_path, drive_path, kind): + captured_kwargs["local_path"] = local_path + captured_kwargs["drive_path"] = drive_path + captured_kwargs["kind"] = kind + return AgentStubDriveCommitResponse(items=[]) + + monkeypatch.setattr( + "dify_agent.agent_stub.cli.main.push_drive_from_environment", + fake_push_drive_from_environment, + ) + + with pytest.raises(SystemExit) as exc_info: + main(["drive", "push", "/tmp/skill", "skills/example", "--kind", "skill"]) + + capsys.readouterr() + assert exc_info.value.code == 0 + assert captured_kwargs == { + "local_path": "/tmp/skill", + "drive_path": "skills/example", + "kind": "skill", + } + + +def test_cli_drive_push_accepts_json_flag( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + captured_kwargs: dict[str, object] = {} + + def fake_push_drive_from_environment(*, local_path, drive_path, kind): + captured_kwargs["local_path"] = local_path + captured_kwargs["drive_path"] = drive_path + captured_kwargs["kind"] = kind + return AgentStubDriveCommitResponse(items=[]) + + monkeypatch.setattr( + "dify_agent.agent_stub.cli.main.push_drive_from_environment", + fake_push_drive_from_environment, + ) + + with pytest.raises(SystemExit) as exc_info: + main(["drive", "push", "/tmp/report.md", "files/report.md", "--json"]) + + captured = capsys.readouterr() + assert exc_info.value.code == 0 + assert json.loads(captured.out) == {"items": []} + assert captured_kwargs == { + "local_path": "/tmp/report.md", + "drive_path": "files/report.md", + "kind": None, + } + + +def test_cli_drive_push_rejects_recursive_option( + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + called = False + + def fake_push_drive_from_environment(*, local_path, drive_path, kind): + nonlocal called + called = True + return AgentStubDriveCommitResponse(items=[]) + + monkeypatch.setattr( + "dify_agent.agent_stub.cli.main.push_drive_from_environment", + fake_push_drive_from_environment, + ) + + with pytest.raises(SystemExit) as exc_info: + main(["drive", "push", "/tmp/dir", "files/dir", "--recursive"]) + + captured = capsys.readouterr() + assert exc_info.value.code == 2 + assert called is False + assert "--recursive" in captured.err diff --git a/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py b/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py index 93283f539fe..8825c466399 100644 --- a/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py +++ b/dify-agent/tests/local/dify_agent/layers/drive/test_layer.py @@ -4,11 +4,14 @@ from __future__ import annotations from pathlib import Path from typing import ClassVar + import pytest from agenton.layers import EmptyRuntimeState, LayerConfig, NoLayerDeps, PlainLayer +from dify_agent.agent_stub._drive_materialization import DriveDownloadPayload +from dify_agent.agent_stub.protocol.agent_stub import AgentStubDriveItem from dify_agent.layers.drive import DifyDriveLayerConfig, DifyDriveSkillConfig -from dify_agent.layers.drive.layer import DifyDriveLayer, DifyDriveLayerError, _DriveManifestItem +from dify_agent.layers.drive.layer import DifyDriveLayer, DifyDriveLayerError class _FakeExecutionContextConfig(LayerConfig): @@ -57,18 +60,55 @@ def _build_layer(tmp_path: Path) -> DifyDriveLayer: return layer +class _FakeAsyncResponse: + def __init__(self, *, status_code: int = 200, json_data: object | None = None, content: bytes = b"") -> None: + self.status_code = status_code + self._json_data = json_data + self.content = content + + @property + def is_error(self) -> bool: + return self.status_code >= 400 + + def json(self) -> object: + if isinstance(self._json_data, Exception): + raise self._json_data + return self._json_data + + +class _FakeAsyncClient: + def __init__(self, responses: dict[str, _FakeAsyncResponse]) -> None: + self._responses = responses + + async def __aenter__(self) -> _FakeAsyncClient: + return self + + async def __aexit__(self, exc_type, exc, tb) -> None: + del exc_type, exc, tb + + async def get(self, url: str, **kwargs) -> _FakeAsyncResponse: + prefix = kwargs.get("params", {}).get("prefix") + key = f"manifest:{prefix}" if prefix is not None else f"download:{url}" + return self._responses[key] + + def test_drive_layer_exposes_agent_stub_cli_usage_suffix_prompt(tmp_path: Path) -> None: layer = _build_layer(tmp_path) assert len(layer.suffix_prompts) == 1 prompt = layer.suffix_prompts[0] - assert "dify-agent drive list [PATH_PREFIX]" in prompt - assert "dify-agent drive pull TARGET ..." in prompt - assert "--drive-base ." in prompt - assert "dify-agent drive push LOCAL_PATH DRIVE_PATH" in prompt - assert "dify-agent file download TRANSFER_METHOD REFERENCE_OR_URL [DIR]" in prompt + assert "dify-agent drive list [REMOTE_PREFIX]" in prompt + assert "dify-agent drive pull [REMOTE ...] [--to LOCAL_DIR]" in prompt + assert "--to ." in prompt + assert "dify-agent drive push LOCAL_FILE REMOTE_PATH" in prompt + assert "dify-agent drive push LOCAL_DIR REMOTE_PATH --kind skill" in prompt + assert "dify-agent drive push LOCAL_DIR REMOTE_PATH --kind dir" in prompt + assert "dify-agent file download TRANSFER_METHOD REFERENCE_OR_URL [--to LOCAL_DIR]" in prompt + assert "dify-agent file download --mapping" in prompt assert "dify-agent file upload PATH" in prompt assert '{"transfer_method":"tool_file","reference":"..."}' in prompt + assert "--recursive" not in prompt + assert "--drive-base" not in prompt @pytest.mark.anyio @@ -77,15 +117,15 @@ async def test_on_context_create_loads_mentioned_targets_into_prompt( ) -> None: layer = _build_layer(tmp_path) - async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[_DriveManifestItem]: + async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[AgentStubDriveItem]: assert tenant_id == "tenant-1" assert targets == [("tender-analyzer/", False), ("files/report.pdf", True)] return [ - _DriveManifestItem(key="tender-analyzer/SKILL.md", download_url="https://files/skill-md"), - _DriveManifestItem(key="files/report.pdf", download_url="https://files/report"), + AgentStubDriveItem(key="tender-analyzer/SKILL.md", download_url="https://files/skill-md"), + AgentStubDriveItem(key="files/report.pdf", download_url="https://files/report"), ] - async def _download_items(items: list[_DriveManifestItem]) -> dict[str, str]: + async def _download_items(items: list[AgentStubDriveItem]) -> dict[str, str]: assert {item.key for item in items} == {"files/report.pdf", "tender-analyzer/SKILL.md"} skill_path = tmp_path / "tender-analyzer" / "SKILL.md" skill_path.parent.mkdir(parents=True, exist_ok=True) @@ -117,15 +157,15 @@ async def test_on_context_resume_loads_mentioned_targets_into_prompt( ) -> None: layer = _build_layer(tmp_path) - async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[_DriveManifestItem]: + async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[AgentStubDriveItem]: assert tenant_id == "tenant-1" assert targets == [("tender-analyzer/", False), ("files/report.pdf", True)] return [ - _DriveManifestItem(key="tender-analyzer/SKILL.md", download_url="https://files/skill-md"), - _DriveManifestItem(key="files/report.pdf", download_url="https://files/report"), + AgentStubDriveItem(key="tender-analyzer/SKILL.md", download_url="https://files/skill-md"), + AgentStubDriveItem(key="files/report.pdf", download_url="https://files/report"), ] - async def _download_items(items: list[_DriveManifestItem]) -> dict[str, str]: + async def _download_items(items: list[AgentStubDriveItem]) -> dict[str, str]: assert {item.key for item in items} == {"files/report.pdf", "tender-analyzer/SKILL.md"} skill_path = tmp_path / "tender-analyzer" / "SKILL.md" skill_path.parent.mkdir(parents=True, exist_ok=True) @@ -158,11 +198,11 @@ async def test_on_context_create_raises_when_mentioned_file_is_missing( ) -> None: layer = _build_layer(tmp_path) - async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[_DriveManifestItem]: + async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[AgentStubDriveItem]: del tenant_id, targets - return [_DriveManifestItem(key="tender-analyzer/SKILL.md", download_url="https://files/skill-md")] + return [AgentStubDriveItem(key="tender-analyzer/SKILL.md", download_url="https://files/skill-md")] - async def _download_items(items: list[_DriveManifestItem]) -> dict[str, str]: + async def _download_items(items: list[AgentStubDriveItem]) -> dict[str, str]: del items skill_path = tmp_path / "tender-analyzer" / "SKILL.md" skill_path.parent.mkdir(parents=True, exist_ok=True) @@ -176,6 +216,130 @@ async def test_on_context_create_raises_when_mentioned_file_is_missing( await layer.on_context_create() +@pytest.mark.anyio +async def test_fetch_manifest_items_validates_payload_filters_exact_targets_and_deduplicates( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + layer = _build_layer(tmp_path) + responses = { + "manifest:tender-analyzer/": _FakeAsyncResponse( + json_data={ + "items": [ + {"key": "tender-analyzer/SKILL.md", "download_url": "https://files/skill-md", "size": 7}, + {"key": "files/report.pdf", "download_url": "https://files/report", "size": 3}, + ] + } + ), + "manifest:files/report.pdf": _FakeAsyncResponse( + json_data={ + "items": [ + {"key": "files/report.pdf", "download_url": "https://files/report", "size": 3}, + {"key": "files/other.pdf", "download_url": "https://files/other", "size": 4}, + ] + } + ), + } + monkeypatch.setattr( + "dify_agent.layers.drive.layer.httpx.AsyncClient", + lambda **_kwargs: _FakeAsyncClient(responses), + ) + + result = await layer._fetch_manifest_items( + tenant_id="tenant-1", + targets=[("tender-analyzer/", False), ("files/report.pdf", True)], + ) + + assert {(item.key, item.download_url, item.size) for item in result} == { + ("files/report.pdf", "https://files/report", 3), + ("tender-analyzer/SKILL.md", "https://files/skill-md", 7), + } + assert [item.key for item in result].count("files/report.pdf") == 1 + + +@pytest.mark.anyio +async def test_fetch_manifest_items_rejects_invalid_manifest_payload( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + layer = _build_layer(tmp_path) + responses = {"manifest:tender-analyzer/": _FakeAsyncResponse(json_data={"items": "bad"})} + monkeypatch.setattr( + "dify_agent.layers.drive.layer.httpx.AsyncClient", + lambda **_kwargs: _FakeAsyncClient(responses), + ) + + with pytest.raises(DifyDriveLayerError, match="drive manifest response is invalid"): + await layer._fetch_manifest_items(tenant_id="tenant-1", targets=[("tender-analyzer/", False)]) + + +@pytest.mark.anyio +async def test_fetch_manifest_items_rejects_missing_download_url( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + layer = _build_layer(tmp_path) + responses = { + "manifest:tender-analyzer/": _FakeAsyncResponse( + json_data={"items": [{"key": "tender-analyzer/SKILL.md", "download_url": None, "size": 7}]} + ) + } + monkeypatch.setattr( + "dify_agent.layers.drive.layer.httpx.AsyncClient", + lambda **_kwargs: _FakeAsyncClient(responses), + ) + + with pytest.raises(DifyDriveLayerError, match="missing download_url"): + await layer._fetch_manifest_items(tenant_id="tenant-1", targets=[("tender-analyzer/", False)]) + + +@pytest.mark.anyio +async def test_download_items_hands_validated_downloads_to_materialization( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + layer = _build_layer(tmp_path) + responses = { + "download:https://files/skill-md": _FakeAsyncResponse(content=b"skill-md"), + "download:https://files/report": _FakeAsyncResponse(content=b"pdf"), + } + monkeypatch.setattr( + "dify_agent.layers.drive.layer.httpx.AsyncClient", + lambda **_kwargs: _FakeAsyncClient(responses), + ) + captured: dict[str, object] = {} + + def fake_materialize_drive_downloads(*, base_path: Path, downloads: list[DriveDownloadPayload], archive_skip_entry_names_by_dir): + captured["base_path"] = base_path + captured["downloads"] = downloads + captured["archive_skip_entry_names_by_dir"] = archive_skip_entry_names_by_dir + return [tmp_path / "tender-analyzer" / "SKILL.md", tmp_path / "files" / "report.pdf"] + + monkeypatch.setattr( + "dify_agent.layers.drive.layer.materialize_drive_downloads", + fake_materialize_drive_downloads, + ) + + result = await layer._download_items( + [ + AgentStubDriveItem(key="tender-analyzer/SKILL.md", download_url="https://files/skill-md", size=8), + AgentStubDriveItem(key="files/report.pdf", download_url="https://files/report", size=3), + ] + ) + + downloads = captured["downloads"] + assert isinstance(downloads, list) + assert downloads == [ + DriveDownloadPayload(key="tender-analyzer/SKILL.md", payload=b"skill-md", size=8), + DriveDownloadPayload(key="files/report.pdf", payload=b"pdf", size=3), + ] + assert captured["archive_skip_entry_names_by_dir"] == {"tender-analyzer": {"SKILL.md"}} + assert result == { + "tender-analyzer/SKILL.md": str(tmp_path / "tender-analyzer" / "SKILL.md"), + "files/report.pdf": str(tmp_path / "files" / "report.pdf"), + } + + @pytest.mark.anyio async def test_on_context_resume_raises_when_mentioned_targets_are_missing( monkeypatch: pytest.MonkeyPatch, @@ -183,11 +347,11 @@ async def test_on_context_resume_raises_when_mentioned_targets_are_missing( ) -> None: layer = _build_layer(tmp_path) - async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[_DriveManifestItem]: + async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[AgentStubDriveItem]: del tenant_id, targets return [] - async def _download_items(items: list[_DriveManifestItem]) -> dict[str, str]: + async def _download_items(items: list[AgentStubDriveItem]) -> dict[str, str]: assert items == [] return {} @@ -205,11 +369,11 @@ async def test_on_context_create_raises_when_manifest_is_empty_for_mentioned_tar ) -> None: layer = _build_layer(tmp_path) - async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[_DriveManifestItem]: + async def _fetch_manifest_items(*, tenant_id: str, targets: list[tuple[str, bool]]) -> list[AgentStubDriveItem]: del tenant_id, targets return [] - async def _download_items(items: list[_DriveManifestItem]) -> dict[str, str]: + async def _download_items(items: list[AgentStubDriveItem]) -> dict[str, str]: assert items == [] return {}