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 c988eae1e6a..69c9d5455bd 100644 --- a/dify-agent/src/dify_agent/agent_stub/cli/_drive.py +++ b/dify-agent/src/dify_agent/agent_stub/cli/_drive.py @@ -113,8 +113,8 @@ def pull_drive_from_environment( When omitted, the historical Agent Stub drive base is used. Returns: - A structured JSON-ready result with downloaded drive keys and their - written local paths. + A structured JSON-ready result with requested drive targets/prefixes + that matched at least one manifest item and their local paths. Observable behavior: Requests a manifest with ``include_download_url=True``, requires every @@ -129,8 +129,9 @@ def pull_drive_from_environment( archive validates successfully. Successfully extracted skill archives are deleted from disk. - Extracted files are materialized on disk but are not added to the - returned item list. + Downloaded files and extracted files are materialized on disk but are + not enumerated in the returned item list; prefix pulls return the local + path corresponding to the requested prefix. Raises: AgentStubValidationError: if a manifest item omits ``download_url``, a @@ -157,6 +158,17 @@ def pull_drive_from_environment( 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() + result_items: list[DrivePullResult.Item] = [] + seen_result_targets: set[str] = set() + for target, response in zip(manifest_targets, responses, strict=True): + if not response.items or target in seen_result_targets: + continue + seen_result_targets.add(target) + try: + local_path = resolve_drive_destination(resolved_base_path, target) + except DriveMaterializationValidationError as exc: + raise AgentStubValidationError(str(exc)) from exc + result_items.append(DrivePullResult.Item(key=target, local_path=str(local_path))) 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 @@ -170,7 +182,7 @@ def pull_drive_from_environment( downloads.append(DriveDownloadPayload(key=item.key, payload=payload, size=item.size)) try: - written_paths = materialize_drive_downloads( + _ = materialize_drive_downloads( base_path=resolved_base_path, downloads=downloads, ) @@ -179,12 +191,7 @@ def pull_drive_from_environment( 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) - ] - ) + return DrivePullResult(items=result_items) def push_drive_from_environment( 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 eef9c0b001f..23ff4216fb8 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 @@ -133,9 +133,9 @@ def test_pull_drive_from_environment_writes_files_under_drive_base( result = pull_drive_from_environment(targets=["skills/"], local_base=str(tmp_path)) assert result.model_dump() == { - "items": [{"key": "skills/example/SKILL.md", "local_path": str(tmp_path / "skills" / "example" / "SKILL.md")}] + "items": [{"key": "skills/", "local_path": str(tmp_path / "skills")}] } - assert Path(result.items[0].local_path).read_bytes() == b"hello world" + assert (tmp_path / "skills" / "example" / "SKILL.md").read_bytes() == b"hello world" assert captured["prefix"] == "skills/" assert captured["include_download_url"] is True @@ -177,7 +177,7 @@ def test_pull_drive_from_environment_auto_extracts_skill_archive( archive_path = tmp_path / "skills" / "foo" / ".DIFY-SKILL-FULL.zip" assert result.model_dump() == { - "items": [{"key": "skills/foo/.DIFY-SKILL-FULL.zip", "local_path": str(archive_path)}] + "items": [{"key": "skills/foo", "local_path": str(tmp_path / "skills" / "foo")}] } assert not archive_path.exists() assert (tmp_path / "skills" / "foo" / "SKILL.md").read_text(encoding="utf-8") == "# Example\n" @@ -475,7 +475,7 @@ def test_pull_drive_from_environment_requests_multiple_targets_and_deduplicates_ 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")), + ("skills/foo", str(tmp_path / "skills" / "foo")), } assert set(downloaded_urls) == {"https://files.example.com/a-txt", "https://files.example.com/skill-md"} assert len(downloaded_urls) == 2