opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-06-26 23:01:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Upgrade dify-agent CLI experience

Browse Source
This commit is contained in:
Yanli 盐粒 2026-06-25 21:29:42 +08:00
parent 0d7d1704f7
commit 4c01f1b2f5
9 changed files with 1094 additions and 419 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

167
dify-agent/src/dify_agent/agent_stub/_drive_materialization.py Normal file
View File

@ -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",
]

258
dify-agent/src/dify_agent/agent_stub/cli/_drive.py
View File

@ -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
``<drive_path>/SKILL.md`` plus ``<drive_path>/.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/<relative_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",
]

79
dify-agent/src/dify_agent/agent_stub/cli/_files.py
View File

@ -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")

74
dify-agent/src/dify_agent/agent_stub/cli/main.py
View File

@ -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

189
dify-agent/src/dify_agent/layers/drive/layer.py
View File

@ -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"]

215
dify-agent/tests/local/dify_agent/agent_stub/cli/test_drive.py
View File

@ -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} == {

62
dify-agent/tests/local/dify_agent/agent_stub/cli/test_files.py
View File

@ -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,

263
dify-agent/tests/local/dify_agent/agent_stub/cli/test_main.py
View File

@ -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

206
dify-agent/tests/local/dify_agent/layers/drive/test_layer.py
View File

@ -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 {}