dify/api/tests/unit_tests/services/test_agent_config_service.py

"""Focused tests for the Agent Soul-backed config service."""

from __future__ import annotations

import io
import zipfile
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest

from models.agent_config_entities import (
    AgentConfigFileRefConfig,
    AgentConfigSkillRefConfig,
    AgentEnvVariableConfig,
    AgentSoulConfig,
)
from services.agent.skill_package_service import SkillPackageError
from services.agent_config_service import (
    AgentConfigService,
    AgentConfigServiceError,
    AgentConfigTarget,
    AgentConfigVersionKind,
    ConfigPushPayload,
    ConfigPushSkillItem,
)

MODULE = "services.agent_config_service"
TENANT = "tenant-1"
AGENT = "agent-1"
USER = "user-1"


def _session_cm(session: MagicMock) -> MagicMock:
    context_manager = MagicMock()
    context_manager.__enter__.return_value = session
    context_manager.__exit__.return_value = None
    return context_manager


def _soul(**updates) -> AgentSoulConfig:
    payload = AgentSoulConfig().model_dump(mode="json")
    payload.update(updates)
    return AgentSoulConfig.model_validate(payload)


def _version(*, version_id: str = "version-1", snapshot: AgentSoulConfig | None = None) -> SimpleNamespace:
    agent_soul = snapshot or _soul()
    return SimpleNamespace(
        id=version_id,
        config_snapshot_dict=agent_soul.model_dump(mode="json"),
        config_snapshot=agent_soul,
    )


def _target(
    *,
    kind: AgentConfigVersionKind,
    writable: bool,
    version_id: str = "version-1",
    soul: AgentSoulConfig | None = None,
) -> AgentConfigTarget:
    agent_soul = soul or _soul()
    return AgentConfigTarget(
        agent_id=AGENT,
        version_id=version_id,
        kind=kind,
        writable=writable,
        version=_version(version_id=version_id, snapshot=agent_soul),
        agent_soul=agent_soul,
    )


def _zip_bytes(members: dict[str, bytes]) -> bytes:
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, "w") as archive:
        for name, payload in members.items():
            zip_info = zipfile.ZipInfo(filename=name)
            zip_info.date_time = (1980, 1, 1, 0, 0, 0)
            archive.writestr(zip_info, payload)
    return buffer.getvalue()


@pytest.mark.parametrize(
    ("kind", "user_id", "version_row", "expected_writable"),
    [
        (AgentConfigVersionKind.SNAPSHOT, None, _version(version_id="snapshot-1"), False),
        (AgentConfigVersionKind.DRAFT, USER, _version(version_id="draft-1"), False),
        (AgentConfigVersionKind.BUILD_DRAFT, USER, _version(version_id="build-draft-1"), True),
    ],
)
def test_resolve_target_supports_snapshot_draft_and_build_draft(
    kind: AgentConfigVersionKind,
    user_id: str | None,
    version_row: SimpleNamespace,
    expected_writable: bool,
) -> None:
    session = MagicMock()
    session.scalar.side_effect = [AGENT, version_row]
    service = AgentConfigService()

    with patch(f"{MODULE}.session_factory.create_session", return_value=_session_cm(session)):
        target = service.resolve_target(
            tenant_id=TENANT,
            agent_id=AGENT,
            config_version_id=version_row.id,
            config_version_kind=kind,
            user_id=user_id,
        )

    assert target.agent_id == AGENT
    assert target.version_id == version_row.id
    assert target.kind == kind
    assert target.writable is expected_writable


def test_resolve_target_requires_user_for_build_draft() -> None:
    session = MagicMock()
    session.scalar.side_effect = [AGENT]
    service = AgentConfigService()

    with patch(f"{MODULE}.session_factory.create_session", return_value=_session_cm(session)):
        with pytest.raises(AgentConfigServiceError, match="user_id is required") as exc_info:
            service.resolve_target(
                tenant_id=TENANT,
                agent_id=AGENT,
                config_version_id="build-draft-1",
                config_version_kind=AgentConfigVersionKind.BUILD_DRAFT,
            )

    assert exc_info.value.code == "missing_user_id"


@pytest.mark.parametrize(
    ("first_scalar", "expected_code"),
    [
        (None, "agent_not_found"),
        (AGENT, "config_version_not_found"),
    ],
)
def test_resolve_target_maps_missing_agent_and_version(first_scalar: str | None, expected_code: str) -> None:
    session = MagicMock()
    if first_scalar is None:
        session.scalar.return_value = None
    else:
        session.scalar.side_effect = [first_scalar, None]
    service = AgentConfigService()

    with patch(f"{MODULE}.session_factory.create_session", return_value=_session_cm(session)):
        with pytest.raises(AgentConfigServiceError) as exc_info:
            service.resolve_target(
                tenant_id=TENANT,
                agent_id=AGENT,
                config_version_id="missing",
                config_version_kind=AgentConfigVersionKind.SNAPSHOT,
                user_id=USER,
            )

    assert exc_info.value.code == expected_code


def test_push_rejects_non_build_draft_writes() -> None:
    session = MagicMock()
    service = AgentConfigService()

    with (
        patch(f"{MODULE}.session_factory.create_session", return_value=_session_cm(session)),
        patch.object(
            service,
            "_resolve_target_in_session",
            return_value=_target(kind=AgentConfigVersionKind.DRAFT, writable=False),
        ),
    ):
        with pytest.raises(AgentConfigServiceError, match="build drafts") as exc_info:
            service.push(
                tenant_id=TENANT,
                agent_id=AGENT,
                user_id=USER,
                config_version_id="draft-1",
                config_version_kind=AgentConfigVersionKind.DRAFT,
                payload=ConfigPushPayload(note="ignored"),
            )

    assert exc_info.value.code == "config_not_writable"
    session.commit.assert_not_called()


def test_push_for_console_allows_shared_draft_mutations() -> None:
    session = MagicMock()
    service = AgentConfigService()
    target = _target(kind=AgentConfigVersionKind.DRAFT, writable=False, soul=_soul(config_note="before"))

    with (
        patch(f"{MODULE}.session_factory.create_session", return_value=_session_cm(session)),
        patch.object(service, "_resolve_target_in_session", return_value=target),
    ):
        manifest = service.push_for_console(
            tenant_id=TENANT,
            agent_id=AGENT,
            user_id=USER,
            config_version_id="draft-1",
            config_version_kind=AgentConfigVersionKind.DRAFT,
            payload=ConfigPushPayload(note="after"),
        )

    assert manifest["note"] == "after"
    assert target.version.config_snapshot.config_note == "after"
    session.commit.assert_called_once()


def test_push_accepts_tenant_scoped_tool_file_sources_from_different_upload_owner() -> None:
    session = MagicMock()
    service = AgentConfigService()
    target = _target(kind=AgentConfigVersionKind.BUILD_DRAFT, writable=True)
    file_source = SimpleNamespace(
        id="tool-file-file",
        tenant_id=TENANT,
        user_id="end-user-1",
        size=7,
        mimetype="text/plain",
        file_key="file-key",
        name="guide.txt",
    )
    skill_source = SimpleNamespace(
        id="tool-file-skill",
        tenant_id=TENANT,
        user_id="end-user-1",
        size=123,
        mimetype="application/zip",
        file_key="skill-key",
        name="alpha.zip",
    )
    skill_ref = AgentConfigSkillRefConfig(
        name="alpha",
        description="Alpha skill",
        file_id="normalized-skill-file",
        size=321,
        mime_type="application/zip",
    )

    with (
        patch(f"{MODULE}.session_factory.create_session", return_value=_session_cm(session)),
        patch.object(service, "_resolve_target_in_session", return_value=target),
        patch.object(service, "_require_tool_file_source", side_effect=[file_source, skill_source]) as require_source,
        patch(f"{MODULE}.storage.load_once", return_value=b"skill-archive"),
        patch.object(service._skill_normalizer, "normalize", return_value=(skill_ref, object())),
    ):
        manifest = service.push(
            tenant_id=TENANT,
            agent_id=AGENT,
            user_id=USER,
            config_version_id="build-draft-1",
            config_version_kind=AgentConfigVersionKind.BUILD_DRAFT,
            payload=ConfigPushPayload.model_validate(
                {
                    "files": [{"name": "guide.txt", "file_ref": {"kind": "tool_file", "id": "tool-file-file"}}],
                    "skills": [{"name": "alpha", "file_ref": {"kind": "tool_file", "id": "tool-file-skill"}}],
                }
            ),
        )

    assert [call.args for call in require_source.call_args_list] == [(session,), (session,)]
    assert [call.kwargs for call in require_source.call_args_list] == [
        {"tenant_id": TENANT, "file_id": "tool-file-file"},
        {"tenant_id": TENANT, "file_id": "tool-file-skill"},
    ]
    files = manifest["files"]
    skills = manifest["skills"]
    assert isinstance(files, dict)
    assert isinstance(skills, dict)
    assert files["items"][0]["file_id"] == "tool-file-file"
    assert skills["items"][0]["file_id"] == "normalized-skill-file"
    session.commit.assert_called_once()


def test_push_file_for_console_rejects_snapshot_writes() -> None:
    session = MagicMock()
    service = AgentConfigService()

    with (
        patch(f"{MODULE}.session_factory.create_session", return_value=_session_cm(session)),
        patch.object(
            service,
            "_resolve_target_in_session",
            return_value=_target(kind=AgentConfigVersionKind.SNAPSHOT, writable=False),
        ),
    ):
        with pytest.raises(AgentConfigServiceError, match="editable drafts") as exc_info:
            service.push_file_for_console(
                tenant_id=TENANT,
                agent_id=AGENT,
                user_id=USER,
                config_version_id="snapshot-1",
                config_version_kind=AgentConfigVersionKind.SNAPSHOT,
                upload_file_id="upload-1",
            )

    assert exc_info.value.code == "config_not_writable"


def test_push_file_for_console_uses_service_owned_upload_lookup_and_naming() -> None:
    session = MagicMock()
    service = AgentConfigService()
    target = _target(kind=AgentConfigVersionKind.DRAFT, writable=False)
    upload_file = SimpleNamespace(
        id="upload-1",
        name="guide.txt",
        size=7,
        hash="sha256:abc",
        mime_type="text/plain",
    )

    with (
        patch(f"{MODULE}.session_factory.create_session", return_value=_session_cm(session)),
        patch.object(service, "_resolve_target_in_session", return_value=target),
        patch.object(service, "_require_console_upload_file_source", return_value=upload_file),
    ):
        response = service.push_file_for_console(
            tenant_id=TENANT,
            agent_id=AGENT,
            user_id=USER,
            config_version_id="draft-1",
            config_version_kind=AgentConfigVersionKind.DRAFT,
            upload_file_id="upload-1",
        )

    assert response == {
        "file": {
            "id": "guide.txt",
            "name": "guide.txt",
            "file_id": "upload-1",
            "size": 7,
            "hash": "sha256:abc",
            "mime_type": "text/plain",
        },
        "config_version": {
            "id": "version-1",
            "kind": "draft",
            "writable": True,
        },
    }
    session.commit.assert_called_once()


def test_upload_skill_for_console_maps_package_validation_failures() -> None:
    session = MagicMock()
    service = AgentConfigService()
    target = _target(kind=AgentConfigVersionKind.DRAFT, writable=False)
    message = "skill package must contain exactly one skill; multiple skill folders in one archive are not supported"

    with (
        patch(f"{MODULE}.session_factory.create_session", return_value=_session_cm(session)),
        patch.object(service, "_resolve_target_in_session", return_value=target),
        patch.object(
            service._skill_normalizer,
            "normalize",
            side_effect=SkillPackageError("files_outside_skill_root", message, status_code=400),
        ),
    ):
        with pytest.raises(AgentConfigServiceError, match="exactly one skill") as exc_info:
            service.upload_skill_for_console(
                tenant_id=TENANT,
                agent_id=AGENT,
                user_id=USER,
                config_version_id="draft-1",
                config_version_kind=AgentConfigVersionKind.DRAFT,
                content=b"bad-archive",
                filename="skills.zip",
            )

    assert exc_info.value.code == "files_outside_skill_root"
    assert exc_info.value.message == message
    assert exc_info.value.status_code == 400
    session.commit.assert_not_called()


def test_apply_skill_updates_rejects_non_tool_file_refs() -> None:
    service = AgentConfigService()

    with pytest.raises(AgentConfigServiceError, match="tool files") as exc_info:
        service._apply_skill_updates(
            MagicMock(),
            tenant_id=TENANT,
            user_id=USER,
            current=[],
            updates=[
                ConfigPushSkillItem.model_validate(
                    {"name": "alpha", "file_ref": {"kind": "upload_file", "id": "upload-1"}}
                )
            ],
        )

    assert exc_info.value.code == "invalid_skill_file_ref"


@pytest.mark.parametrize(
    ("error_code", "message"),
    [
        ("skill_name_mismatch", "skill name does not match requested config key"),
        ("invalid_archive", "stored tool file is not a valid skill archive"),
    ],
)
def test_apply_skill_updates_maps_normalizer_failures(error_code: str, message: str) -> None:
    service = AgentConfigService()
    tool_file = SimpleNamespace(name="alpha.zip", file_key="tool-files/alpha.zip")

    with (
        patch.object(service, "_require_tool_file_source", return_value=tool_file),
        patch(f"{MODULE}.storage.load_once", return_value=b"bad-archive"),
        patch.object(
            service._skill_normalizer,
            "normalize",
            side_effect=SkillPackageError(error_code, message, status_code=400),
        ),
    ):
        with pytest.raises(AgentConfigServiceError, match=message) as exc_info:
            service._apply_skill_updates(
                MagicMock(),
                tenant_id=TENANT,
                user_id=USER,
                current=[],
                updates=[
                    ConfigPushSkillItem.model_validate(
                        {"name": "alpha", "file_ref": {"kind": "tool_file", "id": "tool-file-1"}}
                    )
                ],
            )

    assert exc_info.value.code == error_code


def test_apply_env_text_supports_delete_comments_export_and_keeps_unmentioned_values() -> None:
    current = [
        AgentEnvVariableConfig(key="KEEP", name="KEEP", value="old"),
        AgentEnvVariableConfig(key="REMOVE", name="REMOVE", value="gone"),
        AgentEnvVariableConfig(key="UNTOUCHED", name="UNTOUCHED", value="still-here"),
    ]

    updated = AgentConfigService._apply_env_text(
        current,
        "# comment\nexport KEEP=new-value\nREMOVE=\nNEW='two words'\n",
    )

    values = {item.key: item.value for item in updated}
    assert values == {
        "KEEP": "new-value",
        "UNTOUCHED": "still-here",
        "NEW": "two words",
    }


@pytest.mark.parametrize(
    "archive_bytes",
    [
        pytest.param(b"not-a-zip-archive", id="not-a-zip-archive"),
        pytest.param(_zip_bytes({"README.md": b"missing skill md"}), id="missing-skill-md"),
    ],
)
def test_inspect_skill_maps_invalid_archives_to_service_errors(archive_bytes: bytes) -> None:
    service = AgentConfigService()
    target = _target(
        kind=AgentConfigVersionKind.BUILD_DRAFT,
        writable=True,
        soul=_soul(config_skills=[AgentConfigSkillRefConfig(name="alpha", file_id="tool-file-1")]),
    )

    with (
        patch.object(service, "resolve_target", return_value=target),
        patch.object(service, "_load_tool_file_bytes", return_value=(archive_bytes, "application/zip")),
    ):
        with pytest.raises(AgentConfigServiceError, match="stored config skill archive is invalid") as exc_info:
            service.inspect_skill(
                tenant_id=TENANT,
                agent_id=AGENT,
                config_version_id="build-draft-1",
                config_version_kind=AgentConfigVersionKind.BUILD_DRAFT,
                name="alpha",
                user_id=USER,
            )

    assert exc_info.value.code == "skill_archive_invalid"
    assert exc_info.value.status_code == 500


def test_manifest_uses_items_shape_without_download_urls() -> None:
    target = _target(
        kind=AgentConfigVersionKind.DRAFT,
        writable=False,
        soul=_soul(
            config_skills=[AgentConfigSkillRefConfig(name="alpha", description="Alpha skill", file_id="tool-file-1")],
            config_files=[AgentConfigFileRefConfig(name="guide.txt", file_kind="upload_file", file_id="upload-file-1")],
            config_note="Use the guide.",
        ),
    )

    manifest = AgentConfigService._manifest_for_target(target)

    assert manifest == {
        "agent_id": AGENT,
        "config_version": {
            "id": "version-1",
            "kind": "draft",
            "writable": True,
        },
        "skills": {
            "items": [
                {
                    "id": "alpha",
                    "name": "alpha",
                    "file_id": "tool-file-1",
                    "description": "Alpha skill",
                    "size": None,
                    "hash": None,
                    "mime_type": "application/zip",
                }
            ]
        },
        "files": {
            "items": [
                {
                    "id": "guide.txt",
                    "name": "guide.txt",
                    "file_id": "upload-file-1",
                    "size": None,
                    "hash": None,
                    "mime_type": None,
                }
            ]
        },
        "env_keys": [],
        "note": "Use the guide.",
    }


def test_preview_skill_file_returns_text_preview() -> None:
    service = AgentConfigService()
    target = _target(
        kind=AgentConfigVersionKind.BUILD_DRAFT,
        writable=True,
        soul=_soul(config_skills=[AgentConfigSkillRefConfig(name="alpha", file_id="tool-file-1")]),
    )
    archive_bytes = _zip_bytes(
        {
            "SKILL.md": b"# Alpha\n",
            "references/guide.md": b"hello world",
        }
    )

    with (
        patch.object(service, "resolve_target", return_value=target),
        patch.object(service, "_load_tool_file_bytes", return_value=(archive_bytes, "application/zip")),
    ):
        preview = service.preview_skill_file(
            tenant_id=TENANT,
            agent_id=AGENT,
            config_version_id="build-draft-1",
            config_version_kind=AgentConfigVersionKind.BUILD_DRAFT,
            name="alpha",
            path="references/guide.md",
            user_id=USER,
        )

    assert preview == {
        "path": "references/guide.md",
        "size": 11,
        "truncated": False,
        "binary": False,
        "text": "hello world",
    }


def test_preview_skill_file_marks_binary_and_truncated_payloads() -> None:
    service = AgentConfigService()
    target = _target(
        kind=AgentConfigVersionKind.BUILD_DRAFT,
        writable=True,
        soul=_soul(config_skills=[AgentConfigSkillRefConfig(name="alpha", file_id="tool-file-1")]),
    )
    archive_bytes = _zip_bytes(
        {
            "SKILL.md": b"# Alpha\n",
            "bin/data.bin": b"\x00" + (b"x" * (AgentConfigService.PREVIEW_MAX_BYTES + 10)),
        }
    )

    with (
        patch.object(service, "resolve_target", return_value=target),
        patch.object(service, "_load_tool_file_bytes", return_value=(archive_bytes, "application/zip")),
    ):
        preview = service.preview_skill_file(
            tenant_id=TENANT,
            agent_id=AGENT,
            config_version_id="build-draft-1",
            config_version_kind=AgentConfigVersionKind.BUILD_DRAFT,
            name="alpha",
            path="bin/data.bin",
            user_id=USER,
        )

    assert preview == {
        "path": "bin/data.bin",
        "size": AgentConfigService.PREVIEW_MAX_BYTES + 11,
        "truncated": True,
        "binary": True,
        "text": None,
    }


def test_resolve_skill_file_member_path_requires_existing_member() -> None:
    service = AgentConfigService()
    target = _target(
        kind=AgentConfigVersionKind.BUILD_DRAFT,
        writable=True,
        soul=_soul(config_skills=[AgentConfigSkillRefConfig(name="alpha", file_id="tool-file-1")]),
    )
    archive_bytes = _zip_bytes(
        {
            "SKILL.md": b"# Alpha\n",
            "references/guide.md": b"hello world",
        }
    )

    with (
        patch.object(service, "resolve_target", return_value=target),
        patch.object(service, "_load_tool_file_bytes", return_value=(archive_bytes, "application/zip")),
    ):
        assert (
            service.resolve_skill_file_member_path(
                tenant_id=TENANT,
                agent_id=AGENT,
                config_version_id="build-draft-1",
                config_version_kind=AgentConfigVersionKind.BUILD_DRAFT,
                name="alpha",
                path="references/guide.md",
                user_id=USER,
            )
            == "references/guide.md"
        )

        with pytest.raises(AgentConfigServiceError, match="config skill file not found") as exc_info:
            service.resolve_skill_file_member_path(
                tenant_id=TENANT,
                agent_id=AGENT,
                config_version_id="build-draft-1",
                config_version_kind=AgentConfigVersionKind.BUILD_DRAFT,
                name="alpha",
                path="references/missing.md",
                user_id=USER,
            )

    assert exc_info.value.code == "config_skill_file_not_found"
    assert exc_info.value.status_code == 404


def test_download_url_helpers_use_shared_url_resolution() -> None:
    service = AgentConfigService()
    target = _target(
        kind=AgentConfigVersionKind.BUILD_DRAFT,
        writable=True,
        soul=_soul(
            config_skills=[AgentConfigSkillRefConfig(name="alpha", file_id="tool-file-1")],
            config_files=[AgentConfigFileRefConfig(name="guide.txt", file_kind="upload_file", file_id="upload-file-1")],
        ),
    )

    with (
        patch.object(service, "resolve_target", return_value=target),
        patch.object(
            service,
            "_resolve_download_url",
            side_effect=["https://example.com/alpha.zip", "https://example.com/guide.txt"],
        ),
    ):
        assert (
            service.download_skill_url(
                tenant_id=TENANT,
                agent_id=AGENT,
                config_version_id="build-draft-1",
                config_version_kind=AgentConfigVersionKind.BUILD_DRAFT,
                name="alpha",
                user_id=USER,
            )
            == "https://example.com/alpha.zip"
        )
        assert (
            service.download_file_url(
                tenant_id=TENANT,
                agent_id=AGENT,
                config_version_id="build-draft-1",
                config_version_kind=AgentConfigVersionKind.BUILD_DRAFT,
                name="guide.txt",
                user_id=USER,
            )
            == "https://example.com/guide.txt"
        )