dify/api/tasks/agent_backend_session_cleanup_task.py

"""Celery tasks that execute Agent backend lifecycle-only session cleanup."""

from __future__ import annotations

import logging

from celery import shared_task

from clients.agent_backend.factory import create_agent_backend_run_client
from clients.agent_backend.request_builder import AgentBackendRunRequestBuilder
from clients.agent_backend.session_cleanup import (
    AgentBackendSessionCleanupPayload,
    cleanup_agent_backend_session,
)
from configs import dify_config

logger = logging.getLogger(__name__)


def _create_agent_backend_client():
    if not (dify_config.AGENT_BACKEND_USE_FAKE or dify_config.AGENT_BACKEND_BASE_URL):
        return None
    return create_agent_backend_run_client(
        base_url=dify_config.AGENT_BACKEND_BASE_URL,
        use_fake=dify_config.AGENT_BACKEND_USE_FAKE,
        fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO,
    )


def _run_cleanup_task(payload_dict: dict[str, object]) -> None:
    payload = AgentBackendSessionCleanupPayload.model_validate(payload_dict)
    result = cleanup_agent_backend_session(
        payload=payload,
        client=_create_agent_backend_client(),
        request_builder=AgentBackendRunRequestBuilder(),
    )
    if result.status == "succeeded":
        return

    log_fields = {
        "tenant_id": payload.metadata.get("tenant_id"),
        "app_id": payload.metadata.get("app_id"),
        "workflow_run_id": payload.metadata.get("workflow_run_id"),
        "node_id": payload.metadata.get("node_id"),
        "conversation_id": payload.metadata.get("conversation_id"),
        "agent_id": payload.metadata.get("agent_id"),
        "previous_agent_backend_run_id": payload.metadata.get("previous_agent_backend_run_id"),
        "failed_agent_backend_run_id": payload.metadata.get("failed_agent_backend_run_id"),
        "cleanup_run_id": result.cleanup_run_id,
        "reason": result.reason,
    }
    if result.status == "skipped":
        logger.info("Agent backend session cleanup skipped: %s", log_fields)
        return

    logger.warning("Agent backend session cleanup failed: %s", log_fields)


@shared_task(queue="workflow_storage")
def cleanup_workflow_agent_runtime_session(payload_dict: dict[str, object]) -> None:
    """Run one workflow-owned Agent backend cleanup payload."""
    _run_cleanup_task(payload_dict)


@shared_task(queue="conversation")
def cleanup_conversation_agent_runtime_session(payload_dict: dict[str, object]) -> None:
    """Run one conversation-owned Agent backend cleanup payload."""
    _run_cleanup_task(payload_dict)