dify/api/services/agent/retirement_service.py

"""Workflow-only Agent ownership retirement after product transactions commit."""

from __future__ import annotations

import logging
from collections.abc import Iterable

from sqlalchemy import delete, select
from sqlalchemy.orm import Session

from core.db.session_factory import session_factory
from libs.datetime_utils import naive_utc_now
from models.agent import (
    Agent,
    AgentScope,
    AgentStatus,
    AgentWorkingResourceStatus,
    AgentWorkspace,
    AgentWorkspaceBinding,
    WorkflowAgentNodeBinding,
)
from models.model import App, AppMode
from models.workflow import Workflow
from services.agent.home_snapshot_service import AgentHomeSnapshotService
from services.agent.workspace_service import AgentWorkspaceService
from tasks.collect_agent_resources_task import enqueue_agent_resource_collection
from tasks.remove_app_and_related_data_task import remove_app_and_related_data_task

logger = logging.getLogger(__name__)


class WorkflowAgentRetirementService:
    """Delete workflow-only Agent aggregates after their last Workflow owner is gone."""

    @classmethod
    def retire_unowned(
        cls,
        *,
        tenant_id: str,
        agent_ids: Iterable[str],
        account_id: str | None,
    ) -> None:
        """Retire unowned workflow-only Agents in an independent transaction.

        This method returns ``None``. It archives orphan Agents, retires their
        working resources, and deletes their hidden Apps before committing. It
        then publishes every hidden-App cleanup before publishing the Agent
        resource collector; database and task-publication errors propagate.

        Archived Agents, missing hidden App rows, and already-retired resources
        remain cleanup candidates. A retry can therefore publish duplicate
        cleanup tasks, which are expected to be idempotent.
        """

        candidates = tuple(sorted({agent_id for agent_id in agent_ids if agent_id}))
        if not candidates:
            return
        backing_app_ids: list[str] = []
        retired_bindings: list[str] = []
        retired_workspaces: list[str] = []
        retired_snapshots: list[str] = []
        try:
            with session_factory.create_session() as session:
                retired_agent_ids = cls.archive_unowned(
                    session=session,
                    tenant_id=tenant_id,
                    agent_ids=candidates,
                    account_id=account_id,
                )
                retired_agents = session.scalars(
                    select(Agent).where(
                        Agent.tenant_id == tenant_id,
                        Agent.id.in_(retired_agent_ids),
                    )
                ).all()
                backing_app_ids = sorted({agent.backing_app_id for agent in retired_agents if agent.backing_app_id})
                for app_id in backing_app_ids:
                    AgentWorkspaceService.retire_all_for_app(
                        session=session,
                        tenant_id=tenant_id,
                        app_id=app_id,
                    )
                    retired_workspaces.extend(
                        session.scalars(
                            select(AgentWorkspace.id).where(
                                AgentWorkspace.tenant_id == tenant_id,
                                AgentWorkspace.app_id == app_id,
                                AgentWorkspace.status == AgentWorkingResourceStatus.RETIRED,
                            )
                        ).all()
                    )
                for agent_id in retired_agent_ids:
                    bindings = session.scalars(
                        select(AgentWorkspaceBinding).where(
                            AgentWorkspaceBinding.tenant_id == tenant_id,
                            AgentWorkspaceBinding.agent_id == agent_id,
                        )
                    ).all()
                    for binding in bindings:
                        if binding.status == AgentWorkingResourceStatus.ACTIVE:
                            AgentWorkspaceService.retire_binding(
                                session=session,
                                tenant_id=tenant_id,
                                binding_id=binding.id,
                            )
                        retired_bindings.append(binding.id)
                    retired_snapshots.extend(
                        AgentHomeSnapshotService.retire_all_for_agent(
                            session=session,
                            tenant_id=tenant_id,
                            agent_id=agent_id,
                        )
                    )
                if backing_app_ids:
                    session.execute(
                        delete(App).where(
                            App.tenant_id == tenant_id,
                            App.id.in_(backing_app_ids),
                            App.mode == AppMode.AGENT,
                        )
                    )
                session.commit()
        except Exception:
            logger.exception(
                "Failed to retire unowned Workflow Agents",
                extra={
                    "tenant_id": tenant_id,
                    "agent_ids": candidates,
                },
            )
            raise

        for app_id in backing_app_ids:
            try:
                remove_app_and_related_data_task.delay(tenant_id=tenant_id, app_id=app_id)
            except Exception:
                logger.exception(
                    "Failed to enqueue hidden Agent App cleanup",
                    extra={"tenant_id": tenant_id, "app_id": app_id},
                )
                raise
        enqueue_agent_resource_collection(
            tenant_id=tenant_id,
            workspace_ids=retired_workspaces,
            binding_ids=retired_bindings,
            home_snapshot_ids=retired_snapshots,
            purge_agent_ids=retired_agent_ids,
        )

    @classmethod
    def archive_unowned(
        cls,
        *,
        session: Session,
        tenant_id: str,
        agent_ids: Iterable[str],
        account_id: str | None,
    ) -> list[str]:
        """Archive active orphans and return complete aggregate purge candidates."""
        candidates = tuple(sorted({agent_id for agent_id in agent_ids if agent_id}))
        if not candidates:
            return []
        agents = session.scalars(
            select(Agent).where(
                Agent.tenant_id == tenant_id,
                Agent.id.in_(candidates),
                Agent.scope == AgentScope.WORKFLOW_ONLY,
                Agent.status.in_((AgentStatus.ACTIVE, AgentStatus.ARCHIVED)),
            )
        ).all()
        retained_agent_ids = cls.retained_agent_ids(
            session=session,
            tenant_id=tenant_id,
            agent_ids=[agent.id for agent in agents],
        )
        now = naive_utc_now()
        cleanup_candidates: list[str] = []
        for agent in agents:
            if agent.id in retained_agent_ids:
                continue
            if agent.status == AgentStatus.ACTIVE:
                agent.status = AgentStatus.ARCHIVED
                agent.archived_by = account_id
                agent.archived_at = now
                agent.updated_by = account_id or agent.updated_by
                agent.updated_at = now
            cleanup_candidates.append(agent.id)
        session.flush()
        return cleanup_candidates

    @staticmethod
    def retained_agent_ids(
        *,
        session: Session,
        tenant_id: str,
        agent_ids: list[str],
    ) -> set[str]:
        """Return Agents that still have an exact persisted Workflow owner.

        The owner key is tenant, App, Workflow, and Workflow version. Draft and
        every published version, whether current or historical, count equally;
        the App's current-Workflow pointer is not part of ownership.
        """
        if not agent_ids:
            return set()
        values = session.scalars(
            select(WorkflowAgentNodeBinding.agent_id)
            .join(
                Workflow,
                (Workflow.tenant_id == WorkflowAgentNodeBinding.tenant_id)
                & (Workflow.app_id == WorkflowAgentNodeBinding.app_id)
                & (Workflow.id == WorkflowAgentNodeBinding.workflow_id)
                & (Workflow.version == WorkflowAgentNodeBinding.workflow_version),
            )
            .where(
                WorkflowAgentNodeBinding.tenant_id == tenant_id,
                WorkflowAgentNodeBinding.agent_id.in_(agent_ids),
            )
            .distinct()
        ).all()
        return {agent_id for agent_id in values if agent_id}


__all__ = ["WorkflowAgentRetirementService"]