from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import MagicMock import pytest from sqlalchemy.orm import Session from models.agent import ( Agent, AgentConfigVersionKind, AgentHomeSnapshot, AgentKind, AgentScope, AgentSource, AgentStatus, AgentWorkingResourceStatus, AgentWorkspace, AgentWorkspaceBinding, AgentWorkspaceOwnerType, WorkflowAgentBindingType, WorkflowAgentNodeBinding, ) from models.enums import AppStatus from models.model import App, AppMode from models.workflow import Workflow, WorkflowType from services.agent.home_snapshot_service import AgentHomeSnapshotService from services.agent.retirement_service import WorkflowAgentRetirementService from services.agent.workspace_service import AgentWorkspaceService def test_retire_unowned_commits_resource_retirement(monkeypatch: pytest.MonkeyPatch) -> None: context = MagicMock() session = context.__enter__.return_value session.scalars.return_value.all.return_value = [SimpleNamespace(id="binding-1")] monkeypatch.setattr( "services.agent.retirement_service.session_factory.create_session", lambda: context, ) monkeypatch.setattr( WorkflowAgentRetirementService, "archive_unowned", MagicMock(return_value=["agent-1"]), ) monkeypatch.setattr( AgentWorkspaceService, "retire_binding", MagicMock(return_value="binding-1"), ) monkeypatch.setattr( AgentHomeSnapshotService, "retire_all_for_agent", MagicMock(return_value=["home-1"]), ) result = WorkflowAgentRetirementService.retire_unowned( tenant_id="tenant-1", agent_ids=["agent-1"], account_id="account-1", ) assert result == (["binding-1"], ["home-1"]) session.commit.assert_called_once_with() def _workflow_only_agent() -> Agent: return Agent( id="agent-1", tenant_id="tenant-1", name="Inline Agent", description="", role="", agent_kind=AgentKind.DIFY_AGENT, scope=AgentScope.WORKFLOW_ONLY, source=AgentSource.WORKFLOW, status=AgentStatus.ACTIVE, ) @pytest.mark.parametrize( "sqlite_session", [(Agent, App, Workflow, WorkflowAgentNodeBinding)], indirect=True, ) def test_retire_unowned_keeps_effectively_owned_agent_active( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, ) -> None: agent = _workflow_only_agent() app = App( id="app-1", tenant_id="tenant-1", name="Workflow", mode=AppMode.WORKFLOW, status=AppStatus.NORMAL, enable_site=True, enable_api=True, ) workflow = Workflow.new( tenant_id="tenant-1", app_id=app.id, type=WorkflowType.WORKFLOW.value, version=Workflow.VERSION_DRAFT, graph="{}", features="{}", created_by="account-1", environment_variables=[], conversation_variables=[], rag_pipeline_variables=[], ) binding = WorkflowAgentNodeBinding( tenant_id="tenant-1", app_id=app.id, workflow_id=workflow.id, workflow_version=workflow.version, node_id="agent-node", binding_type=WorkflowAgentBindingType.INLINE_AGENT, agent_id=agent.id, current_snapshot_id="config-1", node_job_config={}, ) sqlite_session.add_all([agent, app, workflow, binding]) sqlite_session.commit() monkeypatch.setattr( "services.agent.retirement_service.session_factory.create_session", lambda: nullcontext(sqlite_session), ) result = WorkflowAgentRetirementService.retire_unowned( tenant_id="tenant-1", agent_ids=[agent.id], account_id="account-1", ) assert result == ([], []) stored_agent = sqlite_session.get(Agent, agent.id) assert stored_agent is not None assert stored_agent.status is AgentStatus.ACTIVE @pytest.mark.parametrize( "sqlite_session", [(Agent, App, Workflow, WorkflowAgentNodeBinding, AgentHomeSnapshot, AgentWorkspace, AgentWorkspaceBinding)], indirect=True, ) def test_retire_unowned_archives_orphan_and_retires_resources( monkeypatch: pytest.MonkeyPatch, sqlite_session: Session, ) -> None: agent = _workflow_only_agent() home = AgentHomeSnapshot( id="home-1", tenant_id="tenant-1", agent_id=agent.id, snapshot_ref="home-ref", status=AgentWorkingResourceStatus.ACTIVE, ) workspace = AgentWorkspace( id="workspace-1", tenant_id="tenant-1", app_id="app-1", owner_type=AgentWorkspaceOwnerType.CONVERSATION, owner_id="conversation-1", owner_scope_key="root", backend_workspace_ref="workspace-ref", status=AgentWorkingResourceStatus.ACTIVE, active_guard=1, ) binding = AgentWorkspaceBinding( id="binding-1", tenant_id="tenant-1", app_id="app-1", workspace_id=workspace.id, agent_id=agent.id, base_home_snapshot_id=home.id, agent_config_version_id="config-1", agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, backend_binding_ref="binding-ref", status=AgentWorkingResourceStatus.ACTIVE, ) sqlite_session.add_all([agent, home, workspace, binding]) sqlite_session.commit() monkeypatch.setattr( "services.agent.retirement_service.session_factory.create_session", lambda: nullcontext(sqlite_session), ) result = WorkflowAgentRetirementService.retire_unowned( tenant_id="tenant-1", agent_ids=[agent.id], account_id="account-1", ) assert result == ([binding.id], [home.id]) stored_agent = sqlite_session.get(Agent, agent.id) stored_binding = sqlite_session.get(AgentWorkspaceBinding, binding.id) stored_workspace = sqlite_session.get(AgentWorkspace, workspace.id) stored_home = sqlite_session.get(AgentHomeSnapshot, home.id) assert stored_agent is not None assert stored_binding is not None assert stored_workspace is not None assert stored_home is not None assert stored_agent.status is AgentStatus.ARCHIVED assert stored_binding.status is AgentWorkingResourceStatus.RETIRED assert stored_workspace.status is AgentWorkingResourceStatus.RETIRED assert stored_home.status is AgentWorkingResourceStatus.RETIRED