diff --git a/api/controllers/console/agent/roster.py b/api/controllers/console/agent/roster.py index dc3923177f0e64..3b819d7d5ab998 100644 --- a/api/controllers/console/agent/roster.py +++ b/api/controllers/console/agent/roster.py @@ -676,6 +676,7 @@ def put( return _serialize_agent_app_detail(session, updated, current_user=current_user) @console_ns.response(204, "Agent app deleted successfully") + @console_ns.response(409, "Agent is still referenced by an active workflow") @console_ns.response(403, "Insufficient permissions") @setup_required @login_required diff --git a/api/controllers/console/app/workflow.py b/api/controllers/console/app/workflow.py index d506a443a3af64..6698883957cf84 100644 --- a/api/controllers/console/app/workflow.py +++ b/api/controllers/console/app/workflow.py @@ -1320,7 +1320,7 @@ def post(self, current_user: Account, app_model: App): workflow_created_at = TimestampField().format(workflow.created_at) - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + binding_ids, home_snapshot_ids, purge_agent_ids = WorkflowAgentRetirementService.retire_unowned( tenant_id=app_model.tenant_id, agent_ids=retirement_candidates, account_id=current_user.id, @@ -1329,6 +1329,7 @@ def post(self, current_user: Account, app_model: App): tenant_id=app_model.tenant_id, binding_ids=binding_ids, home_snapshot_ids=home_snapshot_ids, + purge_agent_ids=purge_agent_ids, ) return { "result": "success", diff --git a/api/controllers/console/snippets/snippet_workflow.py b/api/controllers/console/snippets/snippet_workflow.py index d1399801a06eef..bf29e9a50f52e0 100644 --- a/api/controllers/console/snippets/snippet_workflow.py +++ b/api/controllers/console/snippets/snippet_workflow.py @@ -310,7 +310,7 @@ def post(self, current_user: Account, snippet: CustomizedSnippet): except ValueError as e: return {"message": str(e)}, 400 - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + binding_ids, home_snapshot_ids, purge_agent_ids = WorkflowAgentRetirementService.retire_unowned( tenant_id=tenant_id, agent_ids=retirement_candidates, account_id=current_user.id, @@ -319,6 +319,7 @@ def post(self, current_user: Account, snippet: CustomizedSnippet): tenant_id=tenant_id, binding_ids=binding_ids, home_snapshot_ids=home_snapshot_ids, + purge_agent_ids=purge_agent_ids, ) return { "result": "success", diff --git a/api/openapi/markdown/console-openapi.md b/api/openapi/markdown/console-openapi.md index b403b6713daf5e..bcf801dc137b7d 100644 --- a/api/openapi/markdown/console-openapi.md +++ b/api/openapi/markdown/console-openapi.md @@ -362,6 +362,7 @@ Check if activation token is valid | ---- | ----------- | | 204 | Agent app deleted successfully | | 403 | Insufficient permissions | +| 409 | Agent is still referenced by an active workflow | ### [GET] /agent/{agent_id} #### Parameters diff --git a/api/services/agent/composer_service.py b/api/services/agent/composer_service.py index d29d8ed6b41d3f..3449af0a629206 100644 --- a/api/services/agent/composer_service.py +++ b/api/services/agent/composer_service.py @@ -286,7 +286,7 @@ def save_workflow_composer( agent_id=binding.agent_id, ) session.commit() - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + binding_ids, home_snapshot_ids, purge_agent_ids = WorkflowAgentRetirementService.retire_unowned( tenant_id=tenant_id, agent_ids=retirement_candidates, account_id=account_id, @@ -295,6 +295,7 @@ def save_workflow_composer( tenant_id=tenant_id, binding_ids=binding_ids, home_snapshot_ids=home_snapshot_ids, + purge_agent_ids=purge_agent_ids, ) return state diff --git a/api/services/agent/deletion_service.py b/api/services/agent/deletion_service.py new file mode 100644 index 00000000000000..846daaf5009d27 --- /dev/null +++ b/api/services/agent/deletion_service.py @@ -0,0 +1,119 @@ +"""Hard-delete archived Agent aggregates after external resources are collected.""" + +from __future__ import annotations + +from collections.abc import Iterable + +from sqlalchemy import delete, select + +from core.db.session_factory import session_factory +from models.agent import ( + Agent, + AgentConfigDraft, + AgentConfigRevision, + AgentConfigSnapshot, + AgentDebugConversation, + AgentHomeSnapshot, + AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspaceBinding, + WorkflowAgentNodeBinding, +) +from services.agent.retirement_service import WorkflowAgentRetirementService + + +class AgentDeletionInvariantError(RuntimeError): + """An archived Agent no longer satisfies the hard-deletion contract.""" + + +class AgentDeletionService: + """Delete Agent aggregates and stale Workflow-owned binding soft references. + + The aggregate includes Agent-owned configuration, debug, Home, and Workspace + Binding rows. Purge also removes invalid or historical + ``WorkflowAgentNodeBinding`` soft-reference rows that are owned by Workflows + rather than by the Agent aggregate. + """ + + @classmethod + def purge_archived_agents(cls, *, tenant_id: str, agent_ids: Iterable[str]) -> None: + """Idempotently hard-delete eligible archived Agent aggregates. + + Missing targets are a no-op. Every stored target must be ``ARCHIVED``, + have no ACTIVE Workspace Binding or Home Snapshot, and have no effective + Workflow ownership. All dependent rows and Agents are deleted and + committed in one transaction; an exception before commit leaves the + transaction to roll back without a partial aggregate deletion. + """ + candidates = tuple(sorted({agent_id for agent_id in agent_ids if agent_id})) + if not candidates: + return + + with session_factory.create_session() as session: + agents = session.scalars( + select(Agent).where(Agent.tenant_id == tenant_id, Agent.id.in_(candidates)).with_for_update() + ).all() + if not agents: + return + + stored_ids = [agent.id for agent in agents] + non_archived_ids = [agent.id for agent in agents if agent.status != AgentStatus.ARCHIVED] + if non_archived_ids: + raise AgentDeletionInvariantError( + f"Agents must be ARCHIVED before deletion: {', '.join(non_archived_ids)}" + ) + + active_binding_id = session.scalar( + select(AgentWorkspaceBinding.id) + .where( + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.agent_id.in_(stored_ids), + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, + ) + .limit(1) + ) + if active_binding_id is not None: + raise AgentDeletionInvariantError(f"Agent aggregate still has ACTIVE Binding {active_binding_id}") + + active_home_id = session.scalar( + select(AgentHomeSnapshot.id) + .where( + AgentHomeSnapshot.tenant_id == tenant_id, + AgentHomeSnapshot.agent_id.in_(stored_ids), + AgentHomeSnapshot.status == AgentWorkingResourceStatus.ACTIVE, + ) + .limit(1) + ) + if active_home_id is not None: + raise AgentDeletionInvariantError(f"Agent aggregate still has ACTIVE Home Snapshot {active_home_id}") + + effective_ids = WorkflowAgentRetirementService.effective_agent_ids( + session=session, + tenant_id=tenant_id, + agent_ids=stored_ids, + ) + if effective_ids: + raise AgentDeletionInvariantError( + f"Agents regained effective Workflow ownership: {', '.join(sorted(effective_ids))}" + ) + + for model in ( + WorkflowAgentNodeBinding, + AgentDebugConversation, + AgentConfigRevision, + AgentConfigDraft, + AgentConfigSnapshot, + AgentHomeSnapshot, + AgentWorkspaceBinding, + ): + session.execute( + delete(model).where( + model.tenant_id == tenant_id, + model.agent_id.in_(stored_ids), + ) + ) + session.execute(delete(Agent).where(Agent.tenant_id == tenant_id, Agent.id.in_(stored_ids))) + session.commit() + + +__all__ = ["AgentDeletionInvariantError", "AgentDeletionService"] diff --git a/api/services/agent/errors.py b/api/services/agent/errors.py index 45f10031491646..01a025f310f129 100644 --- a/api/services/agent/errors.py +++ b/api/services/agent/errors.py @@ -19,6 +19,15 @@ class AgentArchivedError(Conflict): description = "Archived agent cannot be modified." +class AgentWorkflowReferenceConflictError(BaseHTTPException): + error_code = "agent_workflow_reference_conflict" + description = "Agent is still referenced by an active workflow." + code = 409 + + def __init__(self, reference_count: int): + super().__init__(description=f"Agent is still referenced by {reference_count} active workflow app(s).") + + class AgentVersionConflictError(Conflict): description = "Agent config version changed. Please reload and try again." diff --git a/api/services/agent/home_snapshot_service.py b/api/services/agent/home_snapshot_service.py index 54b0eb2591ede8..9e2113c56a65d0 100644 --- a/api/services/agent/home_snapshot_service.py +++ b/api/services/agent/home_snapshot_service.py @@ -15,7 +15,6 @@ from models.agent import ( Agent, AgentConfigDraft, - AgentConfigSnapshot, AgentConfigVersionKind, AgentHomeSnapshot, AgentStatus, @@ -108,13 +107,13 @@ def retire_all_for_agent(cls, *, session: Session, tenant_id: str, agent_id: str select(AgentHomeSnapshot).where( AgentHomeSnapshot.tenant_id == tenant_id, AgentHomeSnapshot.agent_id == agent_id, - AgentHomeSnapshot.status == AgentWorkingResourceStatus.ACTIVE, ) ).all() now = naive_utc_now() for row in rows: - row.status = AgentWorkingResourceStatus.RETIRED - row.retired_at = now + if row.status == AgentWorkingResourceStatus.ACTIVE: + row.status = AgentWorkingResourceStatus.RETIRED + row.retired_at = now return [row.id for row in rows] @classmethod @@ -129,13 +128,6 @@ def collect_retired_home_snapshot(cls, *, tenant_id: str, home_snapshot_id: str) ) if snapshot is None: return - referenced = session.scalar( - select(AgentConfigDraft.id).where(AgentConfigDraft.home_snapshot_id == home_snapshot_id).limit(1) - ) or session.scalar( - select(AgentConfigSnapshot.id).where(AgentConfigSnapshot.home_snapshot_id == home_snapshot_id).limit(1) - ) - if referenced is not None: - return snapshot_ref = snapshot.snapshot_ref cls.delete(snapshot_ref=snapshot_ref) with session_factory.create_session() as session: diff --git a/api/services/agent/retirement_service.py b/api/services/agent/retirement_service.py index 25b133e14f9dfb..8ad361d5df0fae 100644 --- a/api/services/agent/retirement_service.py +++ b/api/services/agent/retirement_service.py @@ -37,12 +37,17 @@ def retire_unowned( tenant_id: str, agent_ids: Iterable[str], account_id: str | None, - ) -> tuple[list[str], list[str]]: - """Re-check ownership, archive orphans, and commit their resource retirement.""" + ) -> tuple[list[str], list[str], list[str]]: + """Re-check ownership, archive orphans, and commit resource retirement. + + Returns ``(binding_ids, home_snapshot_ids, purge_agent_ids)`` for the + resources and complete Agent aggregates made eligible by the committed + retirement transaction. + """ candidates = tuple(sorted({agent_id for agent_id in agent_ids if agent_id})) if not candidates: - return [], [] + return [], [], [] retired_bindings: list[str] = [] retired_snapshots: list[str] = [] try: @@ -58,17 +63,16 @@ def retire_unowned( select(AgentWorkspaceBinding).where( AgentWorkspaceBinding.tenant_id == tenant_id, AgentWorkspaceBinding.agent_id == agent_id, - AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, ) ).all() for binding in bindings: - binding_id = AgentWorkspaceService.retire_binding( - session=session, - tenant_id=tenant_id, - binding_id=binding.id, - ) - if binding_id is not None: - retired_bindings.append(binding_id) + 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, @@ -85,8 +89,8 @@ def retire_unowned( "agent_ids": candidates, }, ) - return [], [] - return retired_bindings, retired_snapshots + raise + return retired_bindings, retired_snapshots, retired_agent_ids @classmethod def archive_unowned( @@ -97,7 +101,7 @@ def archive_unowned( agent_ids: Iterable[str], account_id: str | None, ) -> list[str]: - """Archive active orphans and return every orphan eligible for Home cleanup.""" + """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 [] @@ -109,7 +113,7 @@ def archive_unowned( Agent.status.in_((AgentStatus.ACTIVE, AgentStatus.ARCHIVED)), ) ).all() - effective_agent_ids = cls._effective_agent_ids( + effective_agent_ids = cls.effective_agent_ids( session=session, tenant_id=tenant_id, agent_ids=[agent.id for agent in agents], @@ -130,7 +134,7 @@ def archive_unowned( return cleanup_candidates @staticmethod - def _effective_agent_ids( + def effective_agent_ids( *, session: Session, tenant_id: str, diff --git a/api/services/agent/roster_service.py b/api/services/agent/roster_service.py index af7bed03879e96..81766327880272 100644 --- a/api/services/agent/roster_service.py +++ b/api/services/agent/roster_service.py @@ -41,6 +41,7 @@ AgentNameConflictError, AgentNotFoundError, AgentVersionNotFoundError, + AgentWorkflowReferenceConflictError, ) from services.agent.home_snapshot_service import AgentHomeSnapshotService from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService, WorkspaceOwnerScope @@ -1262,7 +1263,20 @@ def update_roster_agent( return self.get_roster_agent_detail(tenant_id=tenant_id, agent_id=agent_id) def archive_roster_agent(self, *, tenant_id: str, agent_id: str, account_id: str) -> None: - agent = self._get_agent(tenant_id=tenant_id, agent_id=agent_id, roster_only=True) + agent = self._session.scalar( + select(Agent) + .where( + Agent.tenant_id == tenant_id, + Agent.id == agent_id, + Agent.scope == AgentScope.ROSTER, + ) + .with_for_update() + ) + if agent is None: + raise AgentNotFoundError() + reference_count = self.count_effective_workflow_references(tenant_id=tenant_id, agent_id=agent_id) + if reference_count: + raise AgentWorkflowReferenceConflictError(reference_count) retired_binding_ids: list[str] = [] if agent.status != AgentStatus.ARCHIVED: agent.status = AgentStatus.ARCHIVED @@ -1273,17 +1287,16 @@ def archive_roster_agent(self, *, tenant_id: str, agent_id: str, account_id: str select(AgentWorkspaceBinding).where( AgentWorkspaceBinding.tenant_id == tenant_id, AgentWorkspaceBinding.agent_id == agent_id, - AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, ) ).all() for binding in bindings: - retired_id = AgentWorkspaceService.retire_binding( - session=self._session, - tenant_id=tenant_id, - binding_id=binding.id, - ) - if retired_id is not None: - retired_binding_ids.append(retired_id) + if binding.status == AgentWorkingResourceStatus.ACTIVE: + AgentWorkspaceService.retire_binding( + session=self._session, + tenant_id=tenant_id, + binding_id=binding.id, + ) + retired_binding_ids.append(binding.id) retired_snapshot_ids = AgentHomeSnapshotService.retire_all_for_agent( session=self._session, tenant_id=tenant_id, @@ -1294,6 +1307,7 @@ def archive_roster_agent(self, *, tenant_id: str, agent_id: str, account_id: str tenant_id=tenant_id, binding_ids=retired_binding_ids, home_snapshot_ids=retired_snapshot_ids, + purge_agent_ids=[agent_id], ) @staticmethod @@ -1655,6 +1669,34 @@ def _load_reference_counts_by_agent_id(self, *, tenant_id: str, agent_ids: list[ return {agent_id: len(app_ids) for agent_id, app_ids in referenced_app_ids_by_agent_id.items()} + def count_effective_workflow_references(self, *, tenant_id: str, agent_id: str) -> int: + """Count normal Apps whose current draft or published Workflow references an Agent.""" + + return self._load_reference_counts_by_agent_id(tenant_id=tenant_id, agent_ids=[agent_id]).get(agent_id, 0) + + @staticmethod + def lock_workflow_bindable_roster_agent(*, session: Any, tenant_id: str, agent_id: str) -> Agent | None: + """Lock one callable roster Agent under the binding-write protocol. + + Every write that establishes or replaces a binding to an existing roster + Agent must first acquire this row lock and retain it until its transaction + ends. This serializes binding writes with direct and App-backed Agent + deletion. Operations touching multiple Agents must acquire these locks in + sorted Agent-ID order to avoid inverse-order deadlocks. + """ + + return session.scalar( + select(Agent) + .where( + Agent.tenant_id == tenant_id, + Agent.id == agent_id, + Agent.scope == AgentScope.ROSTER, + Agent.status == AgentStatus.ACTIVE, + workflow_callable_active_snapshot_filter(), + ) + .with_for_update() + ) + def _load_versions_by_id(self, version_ids: list[str]) -> dict[str, AgentConfigSnapshot]: if not version_ids: return {} diff --git a/api/services/agent/workflow_publish_service.py b/api/services/agent/workflow_publish_service.py index b43f3091f889f7..28e87d6298b75a 100644 --- a/api/services/agent/workflow_publish_service.py +++ b/api/services/agent/workflow_publish_service.py @@ -8,7 +8,6 @@ from sqlalchemy import select from sqlalchemy.orm import Session -from core.agent.publish_visibility import workflow_callable_active_snapshot_filter from core.workflow.nodes.agent_v2.validators import WorkflowAgentNodeValidationError, WorkflowAgentNodeValidator from models.agent import ( Agent, @@ -31,6 +30,7 @@ extract_workflow_node_output_selectors, workflow_previous_node_output_refs_from_selectors, ) +from services.agent.roster_service import AgentRosterService from services.entities.agent_entities import ( ComposerSavePayload, ComposerSaveStrategy, @@ -227,6 +227,23 @@ def sync_agent_bindings_for_draft( account_id: str, ) -> set[str]: agent_nodes = dict(WorkflowAgentNodeValidator.iter_agent_v2_nodes(draft_workflow.graph_dict)) + roster_agent_ids: set[str] = set() + for node_data in agent_nodes.values(): + binding_payload = node_data.get(cls._AGENT_BINDING_KEY) + if not isinstance(binding_payload, Mapping): + continue + agent_id = binding_payload.get("agent_id") + if ( + binding_payload.get("binding_type") == WorkflowAgentBindingType.ROSTER_AGENT.value + and isinstance(agent_id, str) + and agent_id + ): + roster_agent_ids.add(agent_id) + locked_roster_agents = cls._lock_workflow_bindable_roster_agents( + session=session, + tenant_id=draft_workflow.tenant_id, + agent_ids=roster_agent_ids, + ) existing_bindings = list( session.scalars( select(WorkflowAgentNodeBinding).where( @@ -272,6 +289,7 @@ def sync_agent_bindings_for_draft( node_binding=binding_payload, existing_binding=existing_binding, account_id=account_id, + locked_roster_agents=locked_roster_agents, ) if ( replaced_inline_agent_id @@ -310,6 +328,7 @@ def _sync_agent_binding_for_node( node_binding: Mapping[str, Any], existing_binding: WorkflowAgentNodeBinding | None, account_id: str, + locked_roster_agents: Mapping[str, Agent], ) -> None: binding_type = node_binding.get("binding_type") agent_id = node_binding.get("agent_id") @@ -323,10 +342,9 @@ def _sync_agent_binding_for_node( if binding_type == WorkflowAgentBindingType.ROSTER_AGENT.value: agent, current_snapshot_id = cls._resolve_roster_agent_graph_binding( - session=session, - draft_workflow=draft_workflow, node_id=node_id, agent_id=agent_id, + locked_roster_agents=locked_roster_agents, ) resolved_binding_type = WorkflowAgentBindingType.ROSTER_AGENT elif binding_type == WorkflowAgentBindingType.INLINE_AGENT.value: @@ -465,24 +483,13 @@ def _clone_inline_graph_binding_for_node( def _resolve_roster_agent_graph_binding( cls, *, - session: Session, - draft_workflow: Workflow, node_id: str, agent_id: str, + locked_roster_agents: Mapping[str, Agent], ) -> tuple[Agent, str]: """Resolve an active roster Agent whose published snapshot is callable.""" - agent = session.scalar( - select(Agent) - .where( - Agent.tenant_id == draft_workflow.tenant_id, - Agent.id == agent_id, - Agent.scope == AgentScope.ROSTER, - Agent.status == AgentStatus.ACTIVE, - workflow_callable_active_snapshot_filter(), - ) - .limit(1) - ) + agent = locked_roster_agents.get(agent_id) if agent is None: raise ValueError(f"Workflow Agent node {node_id} references an unavailable or unpublished roster agent.") if agent.scope != AgentScope.ROSTER: @@ -491,6 +498,25 @@ def _resolve_roster_agent_graph_binding( raise ValueError(f"Workflow Agent node {node_id} roster agent has no active config snapshot.") return agent, agent.active_config_snapshot_id + @staticmethod + def _lock_workflow_bindable_roster_agents( + *, + session: Session, + tenant_id: str, + agent_ids: set[str], + ) -> dict[str, Agent]: + """Acquire and retain roster binding-write locks in sorted Agent-ID order.""" + locked_agents: dict[str, Agent] = {} + for agent_id in sorted(agent_ids): + agent = AgentRosterService.lock_workflow_bindable_roster_agent( + session=session, + tenant_id=tenant_id, + agent_id=agent_id, + ) + if agent is not None: + locked_agents[agent_id] = agent + return locked_agents + @classmethod def _resolve_inline_agent_graph_binding( cls, @@ -622,6 +648,21 @@ def copy_agent_node_bindings_to_published( if not bindings: return retirement_candidates + roster_agent_ids = { + binding.agent_id + for binding in bindings + if binding.binding_type == WorkflowAgentBindingType.ROSTER_AGENT and binding.agent_id + } + locked_roster_agents = cls._lock_workflow_bindable_roster_agents( + session=session, + tenant_id=draft_workflow.tenant_id, + agent_ids=roster_agent_ids, + ) + unavailable_roster_agent_ids = roster_agent_ids - locked_roster_agents.keys() + if unavailable_roster_agent_ids: + agent_id = min(unavailable_roster_agent_ids) + raise ValueError(f"Published Workflow references unavailable roster Agent {agent_id}.") + agents_by_id = { agent.id: agent for agent in session.scalars( @@ -631,6 +672,7 @@ def copy_agent_node_bindings_to_published( ) ).all() } + agents_by_id.update(locked_roster_agents) for binding in bindings: agent = agents_by_id.get(binding.agent_id) if binding.agent_id else None @@ -679,9 +721,6 @@ def restore_agent_node_bindings_to_draft( for binding in existing if binding.binding_type == WorkflowAgentBindingType.INLINE_AGENT and binding.agent_id } - for binding in existing: - session.delete(binding) - source_bindings = session.scalars( select(WorkflowAgentNodeBinding).where( WorkflowAgentNodeBinding.tenant_id == source_workflow.tenant_id, @@ -690,6 +729,25 @@ def restore_agent_node_bindings_to_draft( WorkflowAgentNodeBinding.workflow_version == source_workflow.version, ) ).all() + roster_agent_ids = { + source.agent_id + for source in source_bindings + if source.binding_type == WorkflowAgentBindingType.ROSTER_AGENT and source.agent_id + } + locked_roster_agents = cls._lock_workflow_bindable_roster_agents( + session=session, + tenant_id=draft_workflow.tenant_id, + agent_ids=roster_agent_ids, + ) + unavailable_roster_agent_ids = roster_agent_ids - locked_roster_agents.keys() + if unavailable_roster_agent_ids: + agent_id = min(unavailable_roster_agent_ids) + raise ValueError(f"Published Workflow references unavailable roster Agent {agent_id}.") + + for binding in existing: + session.delete(binding) + session.flush() + for source in source_bindings: agent_id = source.agent_id snapshot_id = source.current_snapshot_id diff --git a/api/services/agent/workspace_service.py b/api/services/agent/workspace_service.py index 2cc6c34e881c49..cdf2736b523aff 100644 --- a/api/services/agent/workspace_service.py +++ b/api/services/agent/workspace_service.py @@ -13,7 +13,7 @@ from dify_agent.client import Client from dify_agent.protocol import CreateExecutionBindingRequest, DestroyExecutionBindingRequest -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.orm import Session from clients.agent_backend.factory import create_agent_backend_client @@ -364,46 +364,68 @@ def collect_retired_workspace(cls, *, tenant_id: str, workspace_id: str) -> None .order_by(AgentWorkspaceBinding.created_at) ).all() if not bindings: - logger.error( - "RETIRED Workspace has no Binding available for physical collection", - extra={"tenant_id": tenant_id, "workspace_id": workspace_id}, + raise AgentWorkspaceError( + f"RETIRED Workspace has no RETIRED Binding: tenant_id={tenant_id}, workspace_id={workspace_id}" ) - return anchor = bindings[0] - remaining_ids = [binding.id for binding in bindings[1:]] + remaining = [(binding.id, binding.backend_binding_ref) for binding in bindings[1:]] workspace_ref = workspace.backend_workspace_ref binding_ref = anchor.backend_binding_ref anchor_id = anchor.id + + failures: list[str] = [] + first_error: Exception | None = None with cls._client() as client: - client.destroy_execution_binding_sync( - DestroyExecutionBindingRequest( - binding_ref=binding_ref, - workspace_ref=workspace_ref, - destroy_workspace=True, + targets = [(anchor_id, binding_ref, workspace_ref, True)] + [ + (binding_id, backend_binding_ref, None, False) for binding_id, backend_binding_ref in remaining + ] + for binding_id, backend_binding_ref, target_workspace_ref, destroy_workspace in targets: + try: + client.destroy_execution_binding_sync( + DestroyExecutionBindingRequest( + binding_ref=backend_binding_ref, + workspace_ref=target_workspace_ref, + destroy_workspace=destroy_workspace, + ) + ) + except Exception as exc: + failures.append(binding_id) + if first_error is None: + first_error = exc + logger.exception( + "Failed to destroy retired Agent Workspace Binding", + extra={ + "tenant_id": tenant_id, + "workspace_id": workspace_id, + "binding_id": binding_id, + "destroy_workspace": destroy_workspace, + }, + ) + if failures: + if len(failures) == 1 and first_error is not None: + raise first_error + raise AgentWorkspaceError( + f"Failed to destroy {len(failures)} RETIRED Workspace Binding(s): {', '.join(failures)}" + ) from first_error + + binding_ids = [anchor_id, *(binding_id for binding_id, _binding_ref in remaining)] + with session_factory.create_session() as session: + session.execute( + delete(AgentWorkspaceBinding).where( + AgentWorkspaceBinding.id.in_(binding_ids), + AgentWorkspaceBinding.tenant_id == tenant_id, + AgentWorkspaceBinding.workspace_id == workspace_id, + AgentWorkspaceBinding.status == AgentWorkingResourceStatus.RETIRED, ) ) - with session_factory.create_session() as session: - stored_workspace = session.scalar( - select(AgentWorkspace).where( + session.execute( + delete(AgentWorkspace).where( AgentWorkspace.id == workspace_id, AgentWorkspace.tenant_id == tenant_id, AgentWorkspace.status == AgentWorkingResourceStatus.RETIRED, ) ) - stored_anchor = session.scalar( - select(AgentWorkspaceBinding).where( - AgentWorkspaceBinding.id == anchor_id, - AgentWorkspaceBinding.tenant_id == tenant_id, - AgentWorkspaceBinding.status == AgentWorkingResourceStatus.RETIRED, - ) - ) - if stored_workspace is not None: - session.delete(stored_workspace) - if stored_anchor is not None: - session.delete(stored_anchor) session.commit() - for remaining_id in remaining_ids: - cls.collect_retired_binding(tenant_id=tenant_id, binding_id=remaining_id) @staticmethod def validate_binding_generation( diff --git a/api/services/app_dsl_service.py b/api/services/app_dsl_service.py index 98fe5f9d75843f..86cb1984c23602 100644 --- a/api/services/app_dsl_service.py +++ b/api/services/app_dsl_service.py @@ -583,7 +583,7 @@ def _create_or_update_app( draft_workflow=draft_workflow, ) self._session.commit() - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + binding_ids, home_snapshot_ids, purge_agent_ids = WorkflowAgentRetirementService.retire_unowned( tenant_id=app.tenant_id, agent_ids=retirement_candidates, account_id=account.id, @@ -592,6 +592,7 @@ def _create_or_update_app( tenant_id=app.tenant_id, binding_ids=binding_ids, home_snapshot_ids=home_snapshot_ids, + purge_agent_ids=purge_agent_ids, ) case AppMode.CHAT | AppMode.AGENT_CHAT | AppMode.COMPLETION: # Initialize model config diff --git a/api/services/app_service.py b/api/services/app_service.py index d7788ef6bf45a8..7a209299f9def7 100644 --- a/api/services/app_service.py +++ b/api/services/app_service.py @@ -39,7 +39,7 @@ ) from models.model import App, AppMode, AppModelConfig, IconType, Site, load_annotation_reply_config from models.workflow import Workflow -from services.agent.errors import AgentAccessNotReadyError, AgentNameConflictError +from services.agent.errors import AgentAccessNotReadyError, AgentNameConflictError, AgentWorkflowReferenceConflictError from services.agent.home_snapshot_service import AgentHomeSnapshotService from services.agent.retirement_service import WorkflowAgentRetirementService from services.agent.workspace_service import AgentWorkspaceService @@ -745,13 +745,15 @@ def _get_backing_agent_for_update(app: App, *, session: Session) -> Agent | None if app.mode != AppMode.AGENT: return None return session.scalar( - select(Agent).where( + select(Agent) + .where( Agent.tenant_id == app.tenant_id, Agent.app_id == app.id, Agent.scope == AgentScope.ROSTER, Agent.source.in_(APP_BACKED_AGENT_SOURCES), Agent.status == AgentStatus.ACTIVE, ) + .with_for_update() ) @staticmethod @@ -992,9 +994,19 @@ def delete_app(self, app: App, *, session: Session) -> None: Delete app :param app: App instance """ + backing_agent = self._get_backing_agent_for_update(app, session=session) + if backing_agent is not None: + from services.agent.roster_service import AgentRosterService + + reference_count = AgentRosterService(session).count_effective_workflow_references( + tenant_id=app.tenant_id, + agent_id=backing_agent.id, + ) + if reference_count: + raise AgentWorkflowReferenceConflictError(reference_count) + app_was_deleted.send(app) - backing_agent = self._get_backing_agent_for_update(app, session=session) workflow_agent_ids = session.scalars( select(Agent.id).where( Agent.tenant_id == app.tenant_id, @@ -1019,17 +1031,16 @@ def delete_app(self, app: App, *, session: Session) -> None: select(AgentWorkspaceBinding).where( AgentWorkspaceBinding.tenant_id == app.tenant_id, AgentWorkspaceBinding.agent_id == backing_agent.id, - AgentWorkspaceBinding.status == AgentWorkingResourceStatus.ACTIVE, ) ).all() for binding in bindings: - binding_id = AgentWorkspaceService.retire_binding( - session=session, - tenant_id=app.tenant_id, - binding_id=binding.id, - ) - if binding_id is not None: - retired_binding_ids.append(binding_id) + if binding.status == AgentWorkingResourceStatus.ACTIVE: + AgentWorkspaceService.retire_binding( + session=session, + tenant_id=app.tenant_id, + binding_id=binding.id, + ) + retired_binding_ids.append(binding.id) retired_snapshot_ids = AgentHomeSnapshotService.retire_all_for_agent( session=session, tenant_id=app.tenant_id, @@ -1044,16 +1055,22 @@ def delete_app(self, app: App, *, session: Session) -> None: session.delete(app) session.commit() - workflow_binding_ids, workflow_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( - tenant_id=app.tenant_id, - agent_ids=workflow_agent_ids, - account_id=account_id, + workflow_binding_ids, workflow_snapshot_ids, workflow_purge_agent_ids = ( + WorkflowAgentRetirementService.retire_unowned( + tenant_id=app.tenant_id, + agent_ids=workflow_agent_ids, + account_id=account_id, + ) ) enqueue_agent_resource_collection( tenant_id=app.tenant_id, workspace_ids=retired_workspace_ids, binding_ids=[*retired_binding_ids, *workflow_binding_ids], home_snapshot_ids=[*retired_snapshot_ids, *workflow_snapshot_ids], + purge_agent_ids=[ + *([backing_agent.id] if backing_agent is not None else []), + *workflow_purge_agent_ids, + ], ) # clean up web app settings diff --git a/api/services/data_migration/import_service.py b/api/services/data_migration/import_service.py index 48ff50ac6558c4..c3a2508c277cee 100644 --- a/api/services/data_migration/import_service.py +++ b/api/services/data_migration/import_service.py @@ -723,7 +723,7 @@ def _ensure_workflow_app_is_published( app_in_session.workflow_id = workflow.id app_in_session.updated_by = account.id app_in_session.updated_at = naive_utc_now() - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + binding_ids, home_snapshot_ids, purge_agent_ids = WorkflowAgentRetirementService.retire_unowned( tenant_id=target.tenant_id, agent_ids=retirement_candidates, account_id=account.id, @@ -732,6 +732,7 @@ def _ensure_workflow_app_is_published( tenant_id=target.tenant_id, binding_ids=binding_ids, home_snapshot_ids=home_snapshot_ids, + purge_agent_ids=purge_agent_ids, ) def _import_mcp_tools( diff --git a/api/services/snippet_dsl_service.py b/api/services/snippet_dsl_service.py index 22f495a2370131..296bbf5869e5a1 100644 --- a/api/services/snippet_dsl_service.py +++ b/api/services/snippet_dsl_service.py @@ -491,7 +491,7 @@ def _create_or_update_snippet( self._session.commit() if workflow_data: - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + binding_ids, home_snapshot_ids, purge_agent_ids = WorkflowAgentRetirementService.retire_unowned( tenant_id=snippet.tenant_id, agent_ids=retirement_candidates, account_id=account.id, @@ -500,6 +500,7 @@ def _create_or_update_snippet( tenant_id=snippet.tenant_id, binding_ids=binding_ids, home_snapshot_ids=home_snapshot_ids, + purge_agent_ids=purge_agent_ids, ) return snippet diff --git a/api/services/snippet_service.py b/api/services/snippet_service.py index f34f9789fa575c..00c425938216f3 100644 --- a/api/services/snippet_service.py +++ b/api/services/snippet_service.py @@ -9,6 +9,7 @@ from sqlalchemy.orm import Session, sessionmaker from configs import dify_config +from core.db.session_factory import session_factory from core.workflow.node_factory import LATEST_VERSION, NODE_TYPE_CLASSES_MAPPING from enums import DeploymentEdition from graphon.enums import BuiltinNodeTypes, NodeType @@ -377,34 +378,52 @@ def delete_snippet( Agent.status == AgentStatus.ACTIVE, ) ).all() - now = datetime.now(UTC).replace(tzinfo=None) - backing_app_ids = {agent.backing_app_id for agent in owned_agents if agent.backing_app_id} - for agent in owned_agents: - 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 - - if backing_app_ids: - session.execute( - delete(App) - .where( - App.tenant_id == snippet.tenant_id, - App.id.in_(backing_app_ids), - App.mode == AppMode.AGENT, - ) - .execution_options(synchronize_session=False) - ) + if owned_agents: tenant_id = snippet.tenant_id + candidate_agent_ids = [agent.id for agent in owned_agents] - def cleanup_backing_apps(_session: Session) -> None: - from tasks.remove_app_and_related_data_task import remove_app_and_related_data_task + def collect_agent_resources(_session: Session) -> None: + binding_ids, home_snapshot_ids, purge_agent_ids = WorkflowAgentRetirementService.retire_unowned( + tenant_id=tenant_id, + agent_ids=candidate_agent_ids, + account_id=account_id, + ) + if not purge_agent_ids: + return + with session_factory.create_session() as cleanup_session: + backing_app_ids = { + app_id + for app_id in cleanup_session.scalars( + select(Agent.backing_app_id).where( + Agent.tenant_id == tenant_id, + Agent.id.in_(purge_agent_ids), + Agent.backing_app_id.is_not(None), + ) + ).all() + if app_id + } + if backing_app_ids: + cleanup_session.execute( + delete(App).where( + App.tenant_id == tenant_id, + App.id.in_(backing_app_ids), + App.mode == AppMode.AGENT, + ) + ) + cleanup_session.commit() + enqueue_agent_resource_collection( + tenant_id=tenant_id, + binding_ids=binding_ids, + home_snapshot_ids=home_snapshot_ids, + purge_agent_ids=purge_agent_ids, + ) + if backing_app_ids: + from tasks.remove_app_and_related_data_task import remove_app_and_related_data_task - for app_id in backing_app_ids: - remove_app_and_related_data_task.delay(tenant_id=tenant_id, app_id=app_id) + for app_id in backing_app_ids: + remove_app_and_related_data_task.delay(tenant_id=tenant_id, app_id=app_id) - event.listen(session, "after_commit", cleanup_backing_apps, once=True) + event.listen(session, "after_commit", collect_agent_resources, once=True) session.execute( delete(WorkflowAgentNodeBinding) @@ -620,7 +639,7 @@ def sync_draft_workflow( ) self._commit_if_owned(session) if self._session is None: - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + binding_ids, home_snapshot_ids, purge_agent_ids = WorkflowAgentRetirementService.retire_unowned( tenant_id=snippet.tenant_id, agent_ids=retirement_candidates, account_id=account.id, @@ -629,6 +648,7 @@ def sync_draft_workflow( tenant_id=snippet.tenant_id, binding_ids=binding_ids, home_snapshot_ids=home_snapshot_ids, + purge_agent_ids=purge_agent_ids, ) return workflow @@ -679,7 +699,7 @@ def restore_published_workflow_to_draft( ) self._commit_if_owned(session) if self._session is None: - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + binding_ids, home_snapshot_ids, purge_agent_ids = WorkflowAgentRetirementService.retire_unowned( tenant_id=snippet.tenant_id, agent_ids=retirement_candidates, account_id=account.id, @@ -688,6 +708,7 @@ def restore_published_workflow_to_draft( tenant_id=snippet.tenant_id, binding_ids=binding_ids, home_snapshot_ids=home_snapshot_ids, + purge_agent_ids=purge_agent_ids, ) return draft_workflow diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index c821e7b7c807c4..cbc7c9967a2b94 100644 --- a/api/services/workflow_service.py +++ b/api/services/workflow_service.py @@ -491,7 +491,7 @@ def sync_draft_workflow( # commit db session changes if commit: session.commit() - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + binding_ids, home_snapshot_ids, purge_agent_ids = WorkflowAgentRetirementService.retire_unowned( tenant_id=app_model.tenant_id, agent_ids=retirement_candidates, account_id=account.id, @@ -500,6 +500,7 @@ def sync_draft_workflow( tenant_id=app_model.tenant_id, binding_ids=binding_ids, home_snapshot_ids=home_snapshot_ids, + purge_agent_ids=purge_agent_ids, ) # trigger app workflow events @@ -656,7 +657,7 @@ def restore_published_workflow_to_draft( ) session.commit() - binding_ids, home_snapshot_ids = WorkflowAgentRetirementService.retire_unowned( + binding_ids, home_snapshot_ids, purge_agent_ids = WorkflowAgentRetirementService.retire_unowned( tenant_id=app_model.tenant_id, agent_ids=retirement_candidates, account_id=account.id, @@ -665,6 +666,7 @@ def restore_published_workflow_to_draft( tenant_id=app_model.tenant_id, binding_ids=binding_ids, home_snapshot_ids=home_snapshot_ids, + purge_agent_ids=purge_agent_ids, ) app_draft_workflow_was_synced.send(app_model, synced_draft_workflow=draft_workflow) diff --git a/api/tasks/collect_agent_resources_task.py b/api/tasks/collect_agent_resources_task.py index bfec87b7bf3c9c..54addc2f45a6fb 100644 --- a/api/tasks/collect_agent_resources_task.py +++ b/api/tasks/collect_agent_resources_task.py @@ -1,4 +1,10 @@ -"""Asynchronously collect retired Agent working resources.""" +"""Collect retired Agent data under a two-phase task contract. + +Phase one attempts every explicitly identified RETIRED working resource. Phase +two purges the requested archived Agent aggregates only after the whole first +phase succeeds. Any collection failure skips aggregate purge and is included in +the error raised after all explicit resources have been attempted. +""" from __future__ import annotations @@ -7,6 +13,7 @@ from celery import shared_task +from services.agent.deletion_service import AgentDeletionService from services.agent.home_snapshot_service import AgentHomeSnapshotService from services.agent.workspace_service import AgentWorkspaceService @@ -20,8 +27,14 @@ def collect_agent_resources( binding_ids: list[str], workspace_ids: list[str], home_snapshot_ids: list[str], + purge_agent_ids: list[str] | None = None, ) -> None: - """Collect only the explicitly identified RETIRED resources.""" + """Collect the explicit RETIRED batch, then purge Agents only on full success. + + Collection is best-effort across the complete explicit batch so one failed + resource does not hide later failures. If any resource fails, aggregate + purge is skipped and one summary error is raised after all attempts. + """ collectors = ( (workspace_ids, "workspace_id", AgentWorkspaceService.collect_retired_workspace), @@ -32,20 +45,30 @@ def collect_agent_resources( AgentHomeSnapshotService.collect_retired_home_snapshot, ), ) + failures: list[str] = [] + first_error: Exception | None = None for resource_ids, argument_name, collector in collectors: for resource_id in resource_ids: try: collector(tenant_id=tenant_id, **{argument_name: resource_id}) - except Exception: + except Exception as exc: + resource_type = argument_name.removesuffix("_id") + failures.append(f"{resource_type}:{resource_id}") + if first_error is None: + first_error = exc logger.exception( "Failed to collect retired Agent resource", extra={ "tenant_id": tenant_id, - "resource_type": argument_name.removesuffix("_id"), + "resource_type": resource_type, "resource_id": resource_id, }, ) - raise + if failures: + raise RuntimeError( + f"Failed to collect {len(failures)} retired Agent resource(s): {', '.join(failures)}" + ) from first_error + AgentDeletionService.purge_archived_agents(tenant_id=tenant_id, agent_ids=purge_agent_ids or ()) def enqueue_agent_resource_collection( @@ -54,13 +77,15 @@ def enqueue_agent_resource_collection( binding_ids: Iterable[str] = (), workspace_ids: Iterable[str] = (), home_snapshot_ids: Iterable[str] = (), + purge_agent_ids: Iterable[str] = (), ) -> None: - """Best-effort enqueue of physical collection after retire has committed.""" + """Enqueue physical collection after retirement has committed.""" payload = { "binding_ids": sorted({resource_id for resource_id in binding_ids if resource_id}), "workspace_ids": sorted({resource_id for resource_id in workspace_ids if resource_id}), "home_snapshot_ids": sorted({resource_id for resource_id in home_snapshot_ids if resource_id}), + "purge_agent_ids": sorted({agent_id for agent_id in purge_agent_ids if agent_id}), } if not any(payload.values()): return @@ -71,6 +96,7 @@ def enqueue_agent_resource_collection( "Failed to enqueue retired Agent resource collection", extra={"tenant_id": tenant_id, **payload}, ) + raise __all__ = ["collect_agent_resources", "enqueue_agent_resource_collection"] diff --git a/api/tasks/remove_app_and_related_data_task.py b/api/tasks/remove_app_and_related_data_task.py index 4562b7d1d90355..f8af48e09908e3 100644 --- a/api/tasks/remove_app_and_related_data_task.py +++ b/api/tasks/remove_app_and_related_data_task.py @@ -40,6 +40,7 @@ TraceAppConfig, WorkflowSchedulePlan, ) +from models.agent import WorkflowAgentNodeBinding from models.tools import WorkflowToolProvider from models.trigger import WorkflowPluginTrigger, WorkflowTriggerLog, WorkflowWebhookTrigger from models.web import PinnedConversation, SavedMessage @@ -70,6 +71,7 @@ def remove_app_and_related_data_task(self, tenant_id: str, app_id: str): _delete_recommended_apps(tenant_id, app_id) _delete_app_annotation_data(tenant_id, app_id) _delete_app_dataset_joins(tenant_id, app_id) + _delete_workflow_agent_node_bindings(tenant_id, app_id) _delete_app_workflows(tenant_id, app_id) _delete_app_workflow_runs(tenant_id, app_id) _delete_app_workflow_node_executions(tenant_id, app_id) @@ -262,6 +264,17 @@ def del_workflow(session, workflow_id: str): ) +def _delete_workflow_agent_node_bindings(tenant_id: str, app_id: str) -> None: + with session_factory.create_session() as session: + session.execute( + delete(WorkflowAgentNodeBinding).where( + WorkflowAgentNodeBinding.tenant_id == tenant_id, + WorkflowAgentNodeBinding.app_id == app_id, + ) + ) + session.commit() + + def _delete_app_workflow_runs(tenant_id: str, app_id: str): """Delete all workflow runs for an app using the service repository.""" session_maker = sessionmaker(bind=db.engine) diff --git a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py index d91005fa377319..6ec88087b6a7b6 100644 --- a/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py +++ b/api/tests/unit_tests/controllers/console/agent/test_agent_controllers.py @@ -3,12 +3,14 @@ from types import SimpleNamespace from typing import Any, cast from unittest.mock import MagicMock, Mock, call +from uuid import UUID import pytest from flask import Flask from sqlalchemy.orm import Session from werkzeug.exceptions import InternalServerError, NotFound +from controllers.console import bp as console_bp from controllers.console import console_ns from controllers.console.agent import composer as composer_controller from controllers.console.agent import roster as roster_controller @@ -68,6 +70,7 @@ from models.agent import Agent, AgentConfigDraftType, AgentScope, AgentSource, AgentStatus from models.enums import ApiTokenType, ConversationFromSource from models.model import ApiToken, App, AppMode, Conversation, IconType, Message +from services.agent.errors import AgentWorkflowReferenceConflictError from services.entities.agent_entities import ( ComposerSavePayload, ComposerSaveStrategy, @@ -622,6 +625,41 @@ def delete_app(self, app_obj: object, *, session: object) -> None: assert captured["delete"] is app_model +def test_agent_app_delete_reference_conflict_returns_http_409(monkeypatch: pytest.MonkeyPatch) -> None: + http_app = Flask(__name__) + http_app.config["TESTING"] = True + http_app.register_blueprint(console_bp) + agent_id = "00000000-0000-0000-0000-000000000001" + app_model = _app_detail_obj(id="app-1", tenant_id="tenant-1") + monkeypatch.setattr(roster_controller, "_resolve_agent_app_model", lambda *_args, **_kwargs: app_model) + + class FakeAppService: + def delete_app(self, _app_model: object, *, session: object) -> None: + raise AgentWorkflowReferenceConflictError(2) + + monkeypatch.setattr(roster_controller, "AppService", FakeAppService) + delete_agent = unwrap(AgentAppApi.delete) + + def delete_through_registered_resource(self: AgentAppApi, agent_id: UUID) -> tuple[str, int]: + return delete_agent( + self, + MagicMock(), + "tenant-1", + agent_id, + ) + + monkeypatch.setattr(AgentAppApi, "delete", delete_through_registered_resource) + + response = http_app.test_client().delete(f"/console/api/agent/{agent_id}") + + assert response.status_code == 409 + assert response.get_json() == { + "code": "agent_workflow_reference_conflict", + "message": "Agent is still referenced by 2 active workflow app(s).", + "status": 409, + } + + def test_agent_app_copy_uses_agent_id_and_returns_agent_detail( app: Flask, monkeypatch: pytest.MonkeyPatch, account_id: str ) -> None: diff --git a/api/tests/unit_tests/controllers/console/app/test_workflow.py b/api/tests/unit_tests/controllers/console/app/test_workflow.py index 3677b0cbb80bb2..841ad252316ed1 100644 --- a/api/tests/unit_tests/controllers/console/app/test_workflow.py +++ b/api/tests/unit_tests/controllers/console/app/test_workflow.py @@ -2,6 +2,7 @@ import inspect import json +from contextlib import nullcontext from datetime import datetime from types import SimpleNamespace from typing import cast @@ -79,6 +80,50 @@ def _make_workflow(**overrides): return workflow +def test_publish_workflow_forwards_retired_agent_purge_ids( + app: Flask, + monkeypatch: pytest.MonkeyPatch, +) -> None: + current_user = SimpleNamespace(id="account-1") + app_model = SimpleNamespace(id="app-1", tenant_id="tenant-1") + workflow = SimpleNamespace(id="published-workflow", created_at=datetime(2026, 8, 17, 12, 0, 0)) + session = Mock() + session.get.return_value = app_model + monkeypatch.setattr( + workflow_module, + "WorkflowService", + Mock(return_value=SimpleNamespace(publish_workflow=Mock(return_value=(workflow, {"retired-agent"})))), + ) + monkeypatch.setattr( + workflow_module, + "sessionmaker", + lambda _engine: SimpleNamespace(begin=lambda: nullcontext(session)), + ) + monkeypatch.setattr(workflow_module, "db", SimpleNamespace(engine=object())) + monkeypatch.setattr( + workflow_module.WorkflowAgentRetirementService, + "retire_unowned", + Mock(return_value=(["binding-1"], ["home-1"], ["retired-agent"])), + ) + enqueue_collection = Mock() + monkeypatch.setattr(workflow_module, "enqueue_agent_resource_collection", enqueue_collection) + + with app.test_request_context("/apps/app-1/workflows/publish", method="POST", json={}): + response = inspect.unwrap(workflow_module.PublishedWorkflowApi.post)( + workflow_module.PublishedWorkflowApi(), + current_user, + app_model, + ) + + assert response["result"] == "success" + enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + purge_agent_ids=["retired-agent"], + ) + + def test_parse_file_no_config(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(workflow_module.FileUploadConfigManager, "convert", lambda *_args, **_kwargs: None) workflow = SimpleNamespace(features_dict={}, tenant_id="t1") diff --git a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py index b5b2d79411b0f9..7f7aa47220ad9c 100644 --- a/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py +++ b/api/tests/unit_tests/controllers/console/snippets/test_snippet_workflow.py @@ -197,6 +197,46 @@ def fail_publish(*, session: Session, snippet: CustomizedSnippet, account: Accou assert snippet.name == "Snippet" +@pytest.mark.parametrize("sqlite_session", [(CustomizedSnippet,)], indirect=True) +def test_published_workflow_post_forwards_retired_agent_purge_ids( + app: Flask, + monkeypatch: pytest.MonkeyPatch, + sqlite_engine: Engine, + sqlite_session: Session, +) -> None: + user = _account("account-1") + snippet = _snippet() + sqlite_session.add(snippet) + sqlite_session.commit() + workflow = SimpleNamespace(created_at=datetime(2026, 8, 17, 12, 0, 0)) + monkeypatch.setattr(snippet_workflow_module, "db", SimpleNamespace(engine=sqlite_engine)) + monkeypatch.setattr( + snippet_workflow_module, + "_snippet_service", + lambda: SimpleNamespace(publish_workflow=Mock(return_value=(workflow, {"retired-agent"}))), + ) + monkeypatch.setattr( + snippet_workflow_module.WorkflowAgentRetirementService, + "retire_unowned", + Mock(return_value=(["binding-1"], ["home-1"], ["retired-agent"])), + ) + enqueue_collection = Mock() + monkeypatch.setattr(snippet_workflow_module, "enqueue_agent_resource_collection", enqueue_collection) + + api = snippet_workflow_module.SnippetPublishedWorkflowApi() + handler = unwrap(api.post) + with app.test_request_context("/snippets/snippet-1/workflows/publish", method="POST", json={}): + response = handler(api, user, snippet) + + assert response["result"] == "success" + enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + purge_agent_ids=["retired-agent"], + ) + + def test_default_block_configs_delegates_to_service(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: get_default_block_configs = Mock(return_value=[{"type": "llm"}]) monkeypatch.setattr( diff --git a/api/tests/unit_tests/services/agent/test_agent_services.py b/api/tests/unit_tests/services/agent/test_agent_services.py index 5dbce1c2549473..561b456e59f5a5 100644 --- a/api/tests/unit_tests/services/agent/test_agent_services.py +++ b/api/tests/unit_tests/services/agent/test_agent_services.py @@ -27,6 +27,7 @@ AgentScope, AgentSource, AgentStatus, + AgentWorkingResourceStatus, AgentWorkspaceBinding, AgentWorkspaceOwnerType, WorkflowAgentBindingType, @@ -53,6 +54,7 @@ AgentNotFoundError, AgentVersionConflictError, AgentVersionNotFoundError, + AgentWorkflowReferenceConflictError, InvalidComposerConfigError, ) from services.agent.home_snapshot_service import AgentHomeSnapshotService @@ -564,14 +566,11 @@ def test_save_workflow_composer_commits_before_retiring_replaced_inline_agent( def retire_unowned(**kwargs): assert kwargs["agent_ids"] == {"old-inline-agent"} events.append("retire") - return ["binding-1"], ["home-1"] + return ["binding-1"], ["home-1"], ["old-inline-agent"] monkeypatch.setattr(composer_service.WorkflowAgentRetirementService, "retire_unowned", retire_unowned) - monkeypatch.setattr( - composer_service, - "enqueue_agent_resource_collection", - MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")), - ) + enqueue_collection = MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")) + monkeypatch.setattr(composer_service, "enqueue_agent_resource_collection", enqueue_collection) payload = ComposerSavePayload.model_validate( { "variant": ComposerVariant.WORKFLOW, @@ -592,6 +591,12 @@ def retire_unowned(**kwargs): ) assert events == ["commit", "retire", "enqueue"] + enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + purge_agent_ids=["old-inline-agent"], + ) def test_save_workflow_composer_rejects_agent_app_variant(sqlite_session: Session): @@ -4089,7 +4094,7 @@ def test_reference_counts_include_draft_and_published_bindings_once_per_app(sqli assert result == {"agent-1": 1} -def test_roster_update_archive_versions_and_detail(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): +def test_roster_update_versions_detail_and_archive(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session): session = sqlite_session listed_version = AgentConfigSnapshot( id="version-4", @@ -4146,6 +4151,7 @@ def test_roster_update_archive_versions_and_detail(monkeypatch: pytest.MonkeyPat service = AgentRosterService(session) retire_snapshots = MagicMock(return_value=[]) monkeypatch.setattr(AgentHomeSnapshotService, "retire_all_for_agent", retire_snapshots) + monkeypatch.setattr(roster_service, "enqueue_agent_resource_collection", MagicMock()) monkeypatch.setattr( service, "get_roster_agent_detail", @@ -4158,9 +4164,9 @@ def test_roster_update_archive_versions_and_detail(monkeypatch: pytest.MonkeyPat account_id="account-1", payload=roster_service.RosterAgentUpdatePayload(description="new"), ) - service.archive_roster_agent(tenant_id="tenant-1", agent_id="agent-1", account_id="account-1") versions = service.list_agent_versions(tenant_id="tenant-1", agent_id="agent-1") detail = service.get_agent_version_detail(tenant_id="tenant-1", agent_id="agent-1", version_id="version-2") + service.archive_roster_agent(tenant_id="tenant-1", agent_id="agent-1", account_id="account-1") assert updated["description"] == "new" assert agent.status == AgentStatus.ARCHIVED @@ -4211,15 +4217,89 @@ def test_roster_archive_retires_then_commits_before_enqueue( MagicMock(side_effect=lambda **_kwargs: events.append("retire-home") or ["home-1"]), ) event.listen(session, "after_commit", lambda _session: events.append("commit")) - monkeypatch.setattr( - roster_service, - "enqueue_agent_resource_collection", - MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")), - ) + enqueue_collection = MagicMock(side_effect=lambda **_kwargs: events.append("enqueue")) + monkeypatch.setattr(roster_service, "enqueue_agent_resource_collection", enqueue_collection) service.archive_roster_agent(tenant_id="tenant-1", agent_id="agent-1", account_id="account-1") assert events == ["retire-binding", "retire-home", "commit", "enqueue"] + enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + purge_agent_ids=["agent-1"], + ) + + +def test_roster_archive_lookup_locks_agent_row() -> None: + session = MagicMock() + session.scalar.return_value = None + + with pytest.raises(AgentNotFoundError): + AgentRosterService(session).archive_roster_agent( + tenant_id="tenant-1", + agent_id="agent-1", + account_id="account-1", + ) + + statement = session.scalar.call_args.args[0] + assert statement._for_update_arg is not None + + +def test_roster_archive_rejects_effective_workflow_reference( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + agent = _agent() + app = _app(app_id="app-1", mode=AppMode.WORKFLOW) + workflow = _workflow(workflow_id="workflow-draft", app_id=app.id) + binding = WorkflowAgentNodeBinding( + tenant_id=agent.tenant_id, + agent_id=agent.id, + app_id=app.id, + workflow_id=workflow.id, + workflow_version=Workflow.VERSION_DRAFT, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + current_snapshot_id="version-1", + node_job_config=WorkflowNodeJobConfig(), + ) + workspace_binding = AgentWorkspaceBinding( + id="workspace-binding-1", + tenant_id=agent.tenant_id, + app_id="runtime-app-1", + workspace_id="workspace-1", + agent_id=agent.id, + agent_config_version_id="version-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="backend-binding-1", + status=AgentWorkingResourceStatus.ACTIVE, + ) + home = AgentHomeSnapshot( + id="home-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + snapshot_ref="home-ref-1", + status=AgentWorkingResourceStatus.ACTIVE, + ) + sqlite_session.add_all([agent, app, workflow, binding, workspace_binding, home]) + sqlite_session.commit() + retire_homes = MagicMock() + enqueue_collection = MagicMock() + monkeypatch.setattr(AgentHomeSnapshotService, "retire_all_for_agent", retire_homes) + monkeypatch.setattr(roster_service, "enqueue_agent_resource_collection", enqueue_collection) + + with pytest.raises(AgentWorkflowReferenceConflictError, match="1 active workflow app"): + AgentRosterService(sqlite_session).archive_roster_agent( + tenant_id=agent.tenant_id, + agent_id=agent.id, + account_id="account-1", + ) + + assert sqlite_session.get(Agent, agent.id).status == AgentStatus.ACTIVE # type: ignore[union-attr] + assert sqlite_session.get(AgentWorkspaceBinding, workspace_binding.id).status == AgentWorkingResourceStatus.ACTIVE # type: ignore[union-attr] + assert sqlite_session.get(AgentHomeSnapshot, home.id).status == AgentWorkingResourceStatus.ACTIVE # type: ignore[union-attr] + retire_homes.assert_not_called() + enqueue_collection.assert_not_called() def test_roster_archive_commit_failure_does_not_enqueue( diff --git a/api/tests/unit_tests/services/agent/test_deletion_service.py b/api/tests/unit_tests/services/agent/test_deletion_service.py new file mode 100644 index 00000000000000..96929014ae8abd --- /dev/null +++ b/api/tests/unit_tests/services/agent/test_deletion_service.py @@ -0,0 +1,436 @@ +from collections.abc import Iterator +from contextlib import contextmanager, nullcontext +from decimal import Decimal +from unittest.mock import MagicMock + +import pytest +from sqlalchemy.orm import Session, sessionmaker +from sqlalchemy.sql.dml import Delete + +from models.agent import ( + Agent, + AgentConfigDraft, + AgentConfigDraftType, + AgentConfigRevision, + AgentConfigRevisionOperation, + AgentConfigSnapshot, + AgentConfigVersionKind, + AgentDebugConversation, + AgentHomeSnapshot, + AgentKind, + AgentScope, + AgentSource, + AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspace, + AgentWorkspaceBinding, + AgentWorkspaceOwnerType, + WorkflowAgentBindingType, + WorkflowAgentNodeBinding, +) +from models.agent_config_entities import AgentSoulConfig +from models.enums import AppStatus, ConversationFromSource, ConversationStatus +from models.model import App, AppMode, Conversation, Message +from models.workflow import Workflow, WorkflowType +from services.agent.deletion_service import AgentDeletionInvariantError, AgentDeletionService + + +def _archived_agent( + *, + agent_id: str = "agent-1", + tenant_id: str = "tenant-1", + status: AgentStatus = AgentStatus.ARCHIVED, +) -> Agent: + return Agent( + id=agent_id, + tenant_id=tenant_id, + name="Agent", + description="", + role="", + agent_kind=AgentKind.DIFY_AGENT, + scope=AgentScope.ROSTER, + source=AgentSource.AGENT_APP, + status=status, + ) + + +def test_purge_archived_agent_deletes_complete_aggregate( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], +) -> None: + agent = _archived_agent() + snapshot = AgentConfigSnapshot( + id="snapshot-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + version=1, + config_snapshot=AgentSoulConfig(), + ) + rows = [ + agent, + snapshot, + AgentConfigDraft( + id="draft-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + config_snapshot=AgentSoulConfig(), + ), + AgentConfigDraft( + id="build-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + draft_type=AgentConfigDraftType.DEBUG_BUILD, + account_id="account-1", + draft_owner_key="account-1", + config_snapshot=AgentSoulConfig(), + ), + AgentConfigRevision( + id="revision-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + current_snapshot_id=snapshot.id, + revision=1, + operation=AgentConfigRevisionOperation.CREATE_VERSION, + ), + AgentDebugConversation( + id="debug-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + app_id="app-1", + account_id="account-1", + draft_type=AgentConfigDraftType.DRAFT, + conversation_id="conversation-1", + ), + AgentHomeSnapshot( + id="home-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + snapshot_ref="home-ref", + status=AgentWorkingResourceStatus.RETIRED, + ), + AgentWorkspaceBinding( + id="binding-1", + tenant_id=agent.tenant_id, + app_id="app-1", + workspace_id="workspace-1", + agent_id=agent.id, + agent_config_version_id=snapshot.id, + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="binding-ref", + status=AgentWorkingResourceStatus.RETIRED, + ), + WorkflowAgentNodeBinding( + id="workflow-binding-1", + tenant_id=agent.tenant_id, + app_id="missing-app", + workflow_id="missing-workflow", + workflow_version="old-version", + node_id="agent-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id=agent.id, + current_snapshot_id=snapshot.id, + node_job_config={}, + ), + ] + sibling = _archived_agent(agent_id="agent-2") + other_tenant = _archived_agent(agent_id="agent-3", tenant_id="tenant-2") + unrelated_app = App( + id="unrelated-app", + tenant_id=agent.tenant_id, + name="Unrelated", + mode=AppMode.WORKFLOW, + status=AppStatus.NORMAL, + enable_site=True, + enable_api=True, + ) + conversation = Conversation( + id="conversation-1", + app_id=unrelated_app.id, + mode=AppMode.AGENT_CHAT, + name="Preserved conversation", + _inputs={}, + status=ConversationStatus.NORMAL, + from_source=ConversationFromSource.CONSOLE, + from_account_id="account-1", + ) + preserved_rows = [ + sibling, + other_tenant, + AgentConfigDraft( + id="sibling-draft", + tenant_id=sibling.tenant_id, + agent_id=sibling.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + config_snapshot=AgentSoulConfig(), + ), + AgentHomeSnapshot( + id="other-home", + tenant_id=other_tenant.tenant_id, + agent_id=other_tenant.id, + snapshot_ref="other-home-ref", + status=AgentWorkingResourceStatus.RETIRED, + ), + unrelated_app, + AgentWorkspace( + id="workspace-1", + tenant_id=agent.tenant_id, + app_id=unrelated_app.id, + owner_type=AgentWorkspaceOwnerType.CONVERSATION, + owner_id=conversation.id, + owner_scope_key="root", + backend_workspace_ref="workspace-ref", + status=AgentWorkingResourceStatus.ACTIVE, + active_guard=1, + ), + conversation, + Message( + id="message-1", + app_id=unrelated_app.id, + conversation_id=conversation.id, + _inputs={}, + query="hello", + message={"role": "user", "content": "hello"}, + answer="world", + message_unit_price=Decimal(0), + answer_unit_price=Decimal(0), + currency="USD", + from_source=ConversationFromSource.CONSOLE, + from_account_id="account-1", + ), + ] + sqlite_session.add_all([*rows, *preserved_rows]) + sqlite_session.commit() + monkeypatch.setattr( + "services.agent.deletion_service.session_factory.create_session", + sqlite_session_factory, + ) + + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[agent.id]) + + with sqlite_session_factory() as observer_session: + for row in rows: + assert observer_session.get(type(row), row.id) is None + for row in preserved_rows: + assert observer_session.get(type(row), row.id) is not None + + +def test_purge_bulk_deletes_aggregate_dependencies_before_agent(monkeypatch: pytest.MonkeyPatch) -> None: + context = MagicMock() + session = context.__enter__.return_value + session.scalars.return_value.all.return_value = [_archived_agent()] + session.scalar.side_effect = [None, None] + deleted_tables: list[str] = [] + + def record_bulk_delete(statement: object) -> None: + if isinstance(statement, Delete): + deleted_tables.append(statement.table.name) + + session.execute.side_effect = record_bulk_delete + monkeypatch.setattr("services.agent.deletion_service.session_factory.create_session", lambda: context) + monkeypatch.setattr( + "services.agent.deletion_service.WorkflowAgentRetirementService.effective_agent_ids", + lambda **_kwargs: set(), + ) + + AgentDeletionService.purge_archived_agents(tenant_id="tenant-1", agent_ids=["agent-1"]) + + assert deleted_tables == [ + model.__table__.name + for model in ( + WorkflowAgentNodeBinding, + AgentDebugConversation, + AgentConfigRevision, + AgentConfigDraft, + AgentConfigSnapshot, + AgentHomeSnapshot, + AgentWorkspaceBinding, + Agent, + ) + ] + + +@pytest.mark.parametrize( + ("invariant", "expected_error"), + [ + ("non_archived", "must be ARCHIVED"), + ("active_binding", "still has ACTIVE Binding"), + ("active_home", "still has ACTIVE Home Snapshot"), + ], +) +def test_purge_rejects_invalid_aggregate_invariants( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + invariant: str, + expected_error: str, +) -> None: + agent = _archived_agent(status=AgentStatus.ACTIVE if invariant == "non_archived" else AgentStatus.ARCHIVED) + related: AgentWorkspaceBinding | AgentHomeSnapshot | None = None + if invariant == "active_binding": + related = AgentWorkspaceBinding( + id="binding-1", + tenant_id=agent.tenant_id, + app_id="app-1", + workspace_id="workspace-1", + agent_id=agent.id, + agent_config_version_id="snapshot-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="binding-ref", + status=AgentWorkingResourceStatus.ACTIVE, + ) + elif invariant == "active_home": + related = AgentHomeSnapshot( + id="home-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + snapshot_ref="home-ref", + status=AgentWorkingResourceStatus.ACTIVE, + ) + sqlite_session.add(agent) + if related is not None: + sqlite_session.add(related) + sqlite_session.commit() + monkeypatch.setattr( + "services.agent.deletion_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + + with pytest.raises(AgentDeletionInvariantError, match=expected_error): + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[agent.id]) + + assert sqlite_session.get(Agent, agent.id) is not None + if related is not None: + assert sqlite_session.get(type(related), related.id) is not None + + +def test_purge_is_idempotent_for_empty_missing_and_repeated_ids( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + agent = _archived_agent() + sqlite_session.add(agent) + sqlite_session.commit() + monkeypatch.setattr( + "services.agent.deletion_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + + agent_id = agent.id + + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[]) + assert sqlite_session.get(Agent, agent_id) is not None + + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=["missing-agent"]) + assert sqlite_session.get(Agent, agent_id) is not None + + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[agent_id, agent_id]) + assert sqlite_session.get(Agent, agent_id) is None + + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[agent_id]) + assert sqlite_session.get(Agent, agent_id) is None + + +def test_purge_locks_candidate_agents_for_update(monkeypatch: pytest.MonkeyPatch) -> None: + context = MagicMock() + session = context.__enter__.return_value + session.scalars.return_value.all.return_value = [] + monkeypatch.setattr("services.agent.deletion_service.session_factory.create_session", lambda: context) + + AgentDeletionService.purge_archived_agents(tenant_id="tenant-1", agent_ids=["agent-1"]) + + statement = session.scalars.call_args.args[0] + assert statement._for_update_arg is not None + + +@pytest.mark.parametrize("failure_stage", ["delete", "commit"]) +def test_purge_failure_rolls_back_complete_aggregate( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + failure_stage: str, +) -> None: + agent = _archived_agent() + draft = AgentConfigDraft( + id="draft-1", + tenant_id=agent.tenant_id, + agent_id=agent.id, + draft_type=AgentConfigDraftType.DRAFT, + draft_owner_key="", + config_snapshot=AgentSoulConfig(), + ) + sqlite_session.add_all([agent, draft]) + sqlite_session.commit() + agent_id = agent.id + draft_id = draft.id + error = RuntimeError(f"{failure_stage} failed") + + @contextmanager + def failing_session() -> Iterator[Session]: + with sqlite_session_factory() as service_session: + failure_method = failure_stage if failure_stage == "commit" else "execute" + monkeypatch.setattr(service_session, failure_method, MagicMock(side_effect=error)) + yield service_session + + monkeypatch.setattr( + "services.agent.deletion_service.session_factory.create_session", + failing_session, + ) + + with pytest.raises(RuntimeError) as exc_info: + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[agent_id]) + + assert exc_info.value is error + with sqlite_session_factory() as observer_session: + assert observer_session.get(Agent, agent_id) is not None + assert observer_session.get(AgentConfigDraft, draft_id) is not None + + +def test_purge_rejects_agent_with_effective_workflow_reference( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + agent = _archived_agent() + app = App( + id="app-1", + tenant_id=agent.tenant_id, + name="Workflow", + mode=AppMode.WORKFLOW, + status=AppStatus.NORMAL, + enable_site=True, + enable_api=True, + ) + workflow = Workflow.new( + tenant_id=agent.tenant_id, + 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=agent.tenant_id, + app_id=app.id, + workflow_id=workflow.id, + workflow_version=workflow.version, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id=agent.id, + current_snapshot_id="snapshot-1", + node_job_config={}, + ) + sqlite_session.add_all([agent, app, workflow, binding]) + sqlite_session.commit() + monkeypatch.setattr( + "services.agent.deletion_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + + with pytest.raises(AgentDeletionInvariantError, match="regained effective Workflow ownership"): + AgentDeletionService.purge_archived_agents(tenant_id=agent.tenant_id, agent_ids=[agent.id]) + + assert sqlite_session.get(Agent, agent.id) is not None diff --git a/api/tests/unit_tests/services/agent/test_home_snapshot_service.py b/api/tests/unit_tests/services/agent/test_home_snapshot_service.py index abe565aaa942b2..b3476f1c1204a5 100644 --- a/api/tests/unit_tests/services/agent/test_home_snapshot_service.py +++ b/api/tests/unit_tests/services/agent/test_home_snapshot_service.py @@ -148,6 +148,42 @@ def test_home_snapshot_collection_database_failure_propagates(monkeypatch: pytes assert exc_info.value is error +@pytest.mark.parametrize("snapshot_state", ["missing", "active"]) +@pytest.mark.parametrize("sqlite_session", [(AgentHomeSnapshot,)], indirect=True) +def test_home_snapshot_collection_non_retired_target_is_idempotent_noop( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, + snapshot_state: str, +) -> None: + if snapshot_state == "active": + sqlite_session.add( + AgentHomeSnapshot( + id="home-1", + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_ref="snapshot-ref-1", + status=AgentWorkingResourceStatus.ACTIVE, + ) + ) + sqlite_session.commit() + delete = MagicMock() + commit = MagicMock(wraps=sqlite_session.commit) + monkeypatch.setattr( + "services.agent.home_snapshot_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + monkeypatch.setattr(AgentHomeSnapshotService, "delete", delete) + monkeypatch.setattr(sqlite_session, "commit", commit) + + AgentHomeSnapshotService.collect_retired_home_snapshot( + tenant_id="tenant-1", + home_snapshot_id="home-1", + ) + + delete.assert_not_called() + commit.assert_not_called() + + @pytest.mark.parametrize( "sqlite_session", [(AgentHomeSnapshot, AgentConfigDraft, AgentConfigSnapshot)], @@ -185,6 +221,50 @@ def test_home_snapshot_collection_backend_failure_propagates_and_preserves_retir assert stored_snapshot.status is AgentWorkingResourceStatus.RETIRED +@pytest.mark.parametrize( + "sqlite_session", + [(AgentHomeSnapshot, AgentConfigDraft, AgentConfigSnapshot)], + indirect=True, +) +def test_home_snapshot_collection_ignores_config_references( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + snapshot = AgentHomeSnapshot( + id="home-1", + tenant_id="tenant-1", + agent_id="agent-1", + snapshot_ref="snapshot-ref-1", + status=AgentWorkingResourceStatus.RETIRED, + ) + draft = _build_draft(home_snapshot_id=snapshot.id) + config_snapshot = AgentConfigSnapshot( + id="config-1", + tenant_id="tenant-1", + agent_id="agent-1", + version=1, + home_snapshot_id=snapshot.id, + config_snapshot=AgentSoulConfig(), + ) + sqlite_session.add_all([snapshot, draft, config_snapshot]) + sqlite_session.commit() + delete = MagicMock() + monkeypatch.setattr( + "services.agent.home_snapshot_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + monkeypatch.setattr(AgentHomeSnapshotService, "delete", delete) + + AgentHomeSnapshotService.collect_retired_home_snapshot( + tenant_id="tenant-1", + home_snapshot_id=snapshot.id, + ) + + delete.assert_called_once_with(snapshot_ref=snapshot.snapshot_ref) + assert sqlite_session.get(AgentHomeSnapshot, snapshot.id) is None + assert sqlite_session.get(AgentConfigDraft, draft.id) is not None + assert sqlite_session.get(AgentConfigSnapshot, config_snapshot.id) is not None + + @pytest.mark.parametrize( "sqlite_session", [(AgentHomeSnapshot, AgentConfigDraft, AgentConfigSnapshot)], diff --git a/api/tests/unit_tests/services/agent/test_retirement_service.py b/api/tests/unit_tests/services/agent/test_retirement_service.py index 71362eecfe330b..8721ea45cc9bcb 100644 --- a/api/tests/unit_tests/services/agent/test_retirement_service.py +++ b/api/tests/unit_tests/services/agent/test_retirement_service.py @@ -31,7 +31,9 @@ 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")] + session.scalars.return_value.all.return_value = [ + SimpleNamespace(id="binding-1", status=AgentWorkingResourceStatus.ACTIVE) + ] monkeypatch.setattr( "services.agent.retirement_service.session_factory.create_session", lambda: context, @@ -58,10 +60,33 @@ def test_retire_unowned_commits_resource_retirement(monkeypatch: pytest.MonkeyPa account_id="account-1", ) - assert result == (["binding-1"], ["home-1"]) + assert result == (["binding-1"], ["home-1"], ["agent-1"]) session.commit.assert_called_once_with() +def test_retire_unowned_failure_propagates(monkeypatch: pytest.MonkeyPatch) -> None: + context = MagicMock() + error = RuntimeError("retirement failed") + monkeypatch.setattr( + "services.agent.retirement_service.session_factory.create_session", + lambda: context, + ) + monkeypatch.setattr( + WorkflowAgentRetirementService, + "archive_unowned", + MagicMock(side_effect=error), + ) + + with pytest.raises(RuntimeError) as exc_info: + WorkflowAgentRetirementService.retire_unowned( + tenant_id="tenant-1", + agent_ids=["agent-1"], + account_id="account-1", + ) + + assert exc_info.value is error + + def _workflow_only_agent() -> Agent: return Agent( id="agent-1", @@ -131,7 +156,7 @@ def test_retire_unowned_keeps_effectively_owned_agent_active( account_id="account-1", ) - assert result == ([], []) + assert result == ([], [], []) stored_agent = sqlite_session.get(Agent, agent.id) assert stored_agent is not None assert stored_agent.status is AgentStatus.ACTIVE @@ -190,7 +215,7 @@ def test_retire_unowned_archives_orphan_and_retires_resources( account_id="account-1", ) - assert result == ([binding.id], [home.id]) + assert result == ([binding.id], [home.id], [agent.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) diff --git a/api/tests/unit_tests/services/agent/test_workflow_publish_service.py b/api/tests/unit_tests/services/agent/test_workflow_publish_service.py index 70cd220877ae69..79f15db700687c 100644 --- a/api/tests/unit_tests/services/agent/test_workflow_publish_service.py +++ b/api/tests/unit_tests/services/agent/test_workflow_publish_service.py @@ -5,12 +5,13 @@ from sqlalchemy import select from sqlalchemy.orm import Session -from models.agent import Agent, WorkflowAgentBindingType, WorkflowAgentNodeBinding +from models.agent import Agent, AgentScope, WorkflowAgentBindingType, WorkflowAgentNodeBinding from models.agent_config_entities import WorkflowNodeJobConfig from models.enums import AppStatus from models.model import App, AppMode from models.workflow import Workflow, WorkflowType from services.agent.dsl_service import AgentDslService +from services.agent.roster_service import AgentRosterService from services.agent.workflow_publish_service import WorkflowAgentPublishService, _InlineAgentOwnershipError @@ -57,6 +58,7 @@ def test_inline_binding_from_another_node_is_cloned(monkeypatch: pytest.MonkeyPa }, existing_binding=None, account_id="account-1", + locked_roster_agents={}, ) clone.assert_called_once() @@ -67,7 +69,40 @@ def test_inline_binding_from_another_node_is_cloned(monkeypatch: pytest.MonkeyPa assert binding.node_job_config.workflow_prompt == "Summarize the input" -def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent() -> None: +def test_draft_sync_locks_roster_agents_in_sorted_order(monkeypatch: pytest.MonkeyPatch) -> None: + draft_workflow = _workflow() + draft_workflow.graph = ( + '{"nodes":[' + '{"id":"node-b","data":{"type":"agent","version":"2","agent_node_kind":"dify_agent",' + '"agent_binding":{"binding_type":"roster_agent","agent_id":"agent-b"}}},' + '{"id":"node-a","data":{"type":"agent","version":"2","agent_node_kind":"dify_agent",' + '"agent_binding":{"binding_type":"roster_agent","agent_id":"agent-a"}}}' + '],"edges":[]}' + ) + session = Mock() + session.scalars.return_value = SimpleNamespace(all=lambda: []) + agents = { + agent_id: SimpleNamespace( + id=agent_id, + scope=AgentScope.ROSTER, + active_config_snapshot_id=f"{agent_id}-snapshot", + ) + for agent_id in ("agent-a", "agent-b") + } + lock_agent = Mock(side_effect=lambda **kwargs: agents[kwargs["agent_id"]]) + monkeypatch.setattr(AgentRosterService, "lock_workflow_bindable_roster_agent", lock_agent) + + WorkflowAgentPublishService.sync_agent_bindings_for_draft( + session=session, + draft_workflow=draft_workflow, + account_id="account-1", + ) + + assert [call.kwargs["agent_id"] for call in lock_agent.call_args_list] == ["agent-a", "agent-b"] + assert {call.args[0].agent_id for call in session.add.call_args_list} == {"agent-a", "agent-b"} + + +def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent(monkeypatch: pytest.MonkeyPatch) -> None: existing_inline = WorkflowAgentNodeBinding( tenant_id="tenant-1", app_id="app-1", @@ -109,6 +144,15 @@ def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent() -> N SimpleNamespace(all=lambda: [existing_inline, existing_roster]), SimpleNamespace(all=lambda: [source]), ] + events: list[str] = [] + session.delete.side_effect = lambda _binding: events.append("delete") + + def lock_roster_agent(**_kwargs: object) -> SimpleNamespace: + events.append("lock") + return SimpleNamespace(id="roster-agent") + + lock_agent = Mock(side_effect=lock_roster_agent) + monkeypatch.setattr(AgentRosterService, "lock_workflow_bindable_roster_agent", lock_agent) retirement_candidates = WorkflowAgentPublishService.restore_agent_node_bindings_to_draft( session=session, source_workflow=_workflow(workflow_id="published-workflow", version="2026-07-13 00:00:00"), @@ -127,8 +171,85 @@ def test_restore_replaces_bindings_and_returns_only_replaced_inline_agent() -> N assert restored.agent_id == "roster-agent" assert restored.current_snapshot_id == "published-snapshot" assert restored.node_job_config.workflow_prompt == "Use the roster agent" - session.flush.assert_called_once() + assert events.index("lock") < events.index("delete") assert retirement_candidates == {"old-inline-agent"} + lock_agent.assert_called_once_with(session=session, tenant_id="tenant-1", agent_id="roster-agent") + + +def test_restore_locks_roster_agents_in_sorted_order(monkeypatch: pytest.MonkeyPatch) -> None: + source_bindings = [ + WorkflowAgentNodeBinding( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="published-workflow", + workflow_version="published", + node_id=f"node-{suffix}", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id=f"agent-{suffix}", + current_snapshot_id=f"snapshot-{suffix}", + node_job_config={}, + created_by="account-1", + ) + for suffix in ("b", "a") + ] + session = Mock() + session.scalars.side_effect = [ + SimpleNamespace(all=lambda: []), + SimpleNamespace(all=lambda: source_bindings), + ] + lock_agent = Mock(side_effect=lambda **kwargs: SimpleNamespace(id=kwargs["agent_id"])) + monkeypatch.setattr(AgentRosterService, "lock_workflow_bindable_roster_agent", lock_agent) + + WorkflowAgentPublishService.restore_agent_node_bindings_to_draft( + session=session, + source_workflow=_workflow(workflow_id="published-workflow", version="published"), + draft_workflow=_workflow(workflow_id="draft-workflow"), + account_id="account-2", + ) + + assert [call.kwargs["agent_id"] for call in lock_agent.call_args_list] == ["agent-a", "agent-b"] + assert {call.args[0].agent_id for call in session.add.call_args_list} == {"agent-a", "agent-b"} + + +def test_publish_copy_locks_roster_agent_before_creating_binding(monkeypatch: pytest.MonkeyPatch) -> None: + draft_workflow = _workflow() + draft_workflow.graph = ( + '{"nodes":[{"id":"agent-node","data":{"type":"agent","version":"2",' + '"agent_node_kind":"dify_agent"}}],"edges":[]}' + ) + published_workflow = _workflow(workflow_id="published", version="published") + binding = WorkflowAgentNodeBinding( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_version=Workflow.VERSION_DRAFT, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id="roster-agent", + current_snapshot_id="old-snapshot", + node_job_config={}, + created_by="account-1", + ) + session = Mock() + session.scalar.return_value = None + session.scalars.side_effect = [ + SimpleNamespace(all=lambda: [binding]), + SimpleNamespace(all=lambda: []), + ] + locked_agent = SimpleNamespace(id="roster-agent", active_config_snapshot_id="active-snapshot") + lock_agent = Mock(return_value=locked_agent) + monkeypatch.setattr(AgentRosterService, "lock_workflow_bindable_roster_agent", lock_agent) + + WorkflowAgentPublishService.copy_agent_node_bindings_to_published( + session=session, + draft_workflow=draft_workflow, + published_workflow=published_workflow, + ) + + lock_agent.assert_called_once_with(session=session, tenant_id="tenant-1", agent_id="roster-agent") + copied = session.add.call_args.args[0] + assert copied.agent_id == "roster-agent" + assert copied.current_snapshot_id == "active-snapshot" @pytest.mark.parametrize( @@ -258,6 +379,7 @@ def test_inline_binding_reuses_existing_node_owned_agent(monkeypatch: pytest.Mon }, existing_binding=existing_binding, account_id="account-1", + locked_roster_agents={}, ) assert existing_binding.agent_id == "existing-agent" @@ -305,18 +427,28 @@ def test_resolve_existing_inline_binding_agent_returns_valid_agent_or_none(monke def test_resolve_roster_binding_rejects_unpublished_agent() -> None: - session = Mock() - session.scalar.return_value = None - with pytest.raises(ValueError, match="unavailable or unpublished roster agent"): WorkflowAgentPublishService._resolve_roster_agent_graph_binding( - session=session, - draft_workflow=_workflow(), node_id="agent-node", agent_id="agent-1", + locked_roster_agents={}, ) +def test_lock_workflow_bindable_roster_agent_uses_row_lock() -> None: + session = Mock() + session.scalar.return_value = None + + AgentRosterService.lock_workflow_bindable_roster_agent( + session=session, + tenant_id="tenant-1", + agent_id="agent-1", + ) + + statement = session.scalar.call_args.args[0] + assert statement._for_update_arg is not None + + def test_clone_inline_graph_binding_for_node_clones_source(monkeypatch: pytest.MonkeyPatch) -> None: session = Mock() source_agent = SimpleNamespace(id="source-agent") diff --git a/api/tests/unit_tests/services/agent/test_workspace_service.py b/api/tests/unit_tests/services/agent/test_workspace_service.py index d678a92106d4b1..8191f04cbfcc79 100644 --- a/api/tests/unit_tests/services/agent/test_workspace_service.py +++ b/api/tests/unit_tests/services/agent/test_workspace_service.py @@ -15,7 +15,12 @@ AgentWorkspaceBinding, AgentWorkspaceOwnerType, ) -from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService, WorkspaceOwnerScope +from services.agent.workspace_service import ( + AgentWorkspaceError, + AgentWorkspaceNotFoundError, + AgentWorkspaceService, + WorkspaceOwnerScope, +) def _scope() -> WorkspaceOwnerScope: @@ -412,9 +417,11 @@ def test_collect_workspace_destroys_workspace_then_remaining_bindings( sqlite_session.add_all([workspace, anchor, remaining]) sqlite_session.commit() client = MagicMock() + commit = MagicMock(wraps=sqlite_session.commit) monkeypatch.setattr( "services.agent.workspace_service.session_factory.create_session", lambda: nullcontext(sqlite_session) ) + monkeypatch.setattr(sqlite_session, "commit", commit) monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client)) AgentWorkspaceService.collect_retired_workspace(tenant_id="tenant-1", workspace_id=workspace.id) @@ -429,6 +436,73 @@ def test_collect_workspace_destroys_workspace_then_remaining_bindings( assert sqlite_session.get(AgentWorkspace, workspace.id) is None assert sqlite_session.get(AgentWorkspaceBinding, anchor.id) is None assert sqlite_session.get(AgentWorkspaceBinding, remaining.id) is None + commit.assert_called_once() + + +@pytest.mark.parametrize("sqlite_session", [(AgentWorkspace, AgentWorkspaceBinding)], indirect=True) +def test_collect_workspace_remaining_failure_preserves_ledgers_and_replay_converges( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + workspace = _workspace(status=AgentWorkingResourceStatus.RETIRED) + anchor = _binding(status=AgentWorkingResourceStatus.RETIRED) + remaining = [ + _binding( + binding_id=f"binding-{index}", + agent_id=f"agent-{index}", + status=AgentWorkingResourceStatus.RETIRED, + ) + for index in (2, 3) + ] + anchor.created_at = datetime(2026, 7, 23, 10) + for index, binding in enumerate(remaining, start=1): + binding.created_at = anchor.created_at + timedelta(minutes=index) + sqlite_session.add_all([workspace, anchor, *remaining]) + sqlite_session.commit() + workspace_id = workspace.id + binding_ids = [anchor.id, *(binding.id for binding in remaining)] + error = RuntimeError("middle Binding destroy failed") + client = MagicMock() + client.destroy_execution_binding_sync.side_effect = [None, error, None] + monkeypatch.setattr( + "services.agent.workspace_service.session_factory.create_session", lambda: nullcontext(sqlite_session) + ) + monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client)) + + with pytest.raises(RuntimeError) as exc_info: + AgentWorkspaceService.collect_retired_workspace(tenant_id="tenant-1", workspace_id=workspace_id) + + assert exc_info.value is error + assert client.destroy_execution_binding_sync.call_count == 3 + first_attempt = [call.args[0] for call in client.destroy_execution_binding_sync.call_args_list] + assert [request.destroy_workspace for request in first_attempt] == [True, False, False] + assert sqlite_session.get(AgentWorkspace, workspace_id) is not None + assert all(sqlite_session.get(AgentWorkspaceBinding, binding_id) is not None for binding_id in binding_ids) + + client.reset_mock() + client.destroy_execution_binding_sync.side_effect = None + AgentWorkspaceService.collect_retired_workspace(tenant_id="tenant-1", workspace_id=workspace_id) + + assert client.destroy_execution_binding_sync.call_count == 3 + assert sqlite_session.get(AgentWorkspace, workspace_id) is None + assert all(sqlite_session.get(AgentWorkspaceBinding, binding_id) is None for binding_id in binding_ids) + + +def test_collect_retired_workspace_without_retired_binding_raises( + monkeypatch: pytest.MonkeyPatch, sqlite_session: Session +) -> None: + workspace = _workspace(status=AgentWorkingResourceStatus.RETIRED) + sqlite_session.add(workspace) + sqlite_session.commit() + monkeypatch.setattr( + "services.agent.workspace_service.session_factory.create_session", + lambda: nullcontext(sqlite_session), + ) + + with pytest.raises(AgentWorkspaceError, match="tenant_id=tenant-1, workspace_id=workspace-1"): + AgentWorkspaceService.collect_retired_workspace( + tenant_id="tenant-1", + workspace_id=workspace.id, + ) def test_binding_collection_database_failure_propagates(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/api/tests/unit_tests/services/data_migration/test_import_service.py b/api/tests/unit_tests/services/data_migration/test_import_service.py index 02a897a8f20939..3549196154548d 100644 --- a/api/tests/unit_tests/services/data_migration/test_import_service.py +++ b/api/tests/unit_tests/services/data_migration/test_import_service.py @@ -1,5 +1,8 @@ +from contextlib import nullcontext from dataclasses import dataclass +from types import SimpleNamespace from typing import cast +from unittest.mock import Mock import pytest import yaml @@ -489,6 +492,43 @@ def create_workflow_tool(**kwargs) -> None: assert events == [("published", app_id), ("created", app_id)] +def test_ensure_workflow_app_is_published_forwards_purge_ids( + monkeypatch: pytest.MonkeyPatch, + database: Database, +) -> None: + _, account = _persist_tenant_account(database.session) + app_id = "00000000-0000-0000-0000-000000000001" + _persist_app(database.session, app_id=app_id) + publish = Mock(return_value=(SimpleNamespace(id="published-workflow"), {"retired-agent"})) + monkeypatch.setattr(import_service, "WorkflowService", Mock(return_value=SimpleNamespace(publish_workflow=publish))) + monkeypatch.setattr( + import_service, + "sessionmaker", + lambda _engine: SimpleNamespace(begin=lambda: nullcontext(database.session)), + ) + monkeypatch.setattr( + import_service.WorkflowAgentRetirementService, + "retire_unowned", + Mock(return_value=(["binding-1"], ["home-1"], ["retired-agent"])), + ) + enqueue_collection = Mock() + monkeypatch.setattr(import_service, "enqueue_agent_resource_collection", enqueue_collection) + + MigrationImportService()._ensure_workflow_app_is_published( + ImportTarget("tenant-1", "target", "account-1", "owner@example.com"), + account, + app_id, + session=database.session, + ) + + enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + purge_agent_ids=["retired-agent"], + ) + + @pytest.mark.parametrize("id_strategy", [IdStrategy.PRESERVE_ID, IdStrategy.GENERATE_NEW_ID]) def test_workflow_tool_import_id_follows_id_strategy( monkeypatch: pytest.MonkeyPatch, database: Database, id_strategy: IdStrategy diff --git a/api/tests/unit_tests/services/test_app_dsl_service.py b/api/tests/unit_tests/services/test_app_dsl_service.py index 7afbb0fa042e91..b90bbc858502f5 100644 --- a/api/tests/unit_tests/services/test_app_dsl_service.py +++ b/api/tests/unit_tests/services/test_app_dsl_service.py @@ -301,6 +301,61 @@ def record_signal(*_args: object, **_kwargs: object) -> None: assert sqlite_session.in_transaction() +def test_create_or_update_app_forwards_imported_agent_purge_ids(monkeypatch: pytest.MonkeyPatch) -> None: + session = SimpleNamespace(add=Mock(), flush=Mock(), commit=Mock(), get=Mock()) + service = AppDslService(session=session) + app = SimpleNamespace( + id="app-1", + tenant_id="tenant-1", + name="Workflow", + description="", + icon_type=IconType.EMOJI, + icon="robot", + icon_background="#FFFFFF", + ) + workflow = SimpleNamespace(id="workflow-1") + workflow_service = SimpleNamespace( + get_draft_workflow=Mock(return_value=None), + sync_draft_workflow=Mock(return_value=workflow), + ) + monkeypatch.setattr("services.app_dsl_service.WorkflowService", Mock(return_value=workflow_service)) + monkeypatch.setattr( + "services.app_dsl_service.AgentDslService.graph_without_package_bindings", + Mock(return_value={"nodes": [], "edges": []}), + ) + monkeypatch.setattr( + "services.app_dsl_service.AgentDslService.import_workflow_packages", + Mock(return_value=(workflow, [], {"retired-agent"})), + ) + monkeypatch.setattr( + "services.app_dsl_service.WorkflowAgentPublishService.validate_agent_nodes_for_draft_sync", + Mock(), + ) + monkeypatch.setattr( + "services.app_dsl_service.WorkflowAgentRetirementService.retire_unowned", + Mock(return_value=(["binding-1"], ["home-1"], ["retired-agent"])), + ) + enqueue_collection = Mock() + monkeypatch.setattr("services.app_dsl_service.enqueue_agent_resource_collection", enqueue_collection) + + service._create_or_update_app( + app=cast(App, app), + data={ + "app": {"mode": AppMode.WORKFLOW.value}, + "workflow": {"graph": {"nodes": [], "edges": []}}, + "agent_packages": {"package-1": {}}, + }, + account=Mock(id="account-1"), + ) + + enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + purge_agent_ids=["retired-agent"], + ) + + def test_export_dsl_loads_model_config_and_annotation_reply_with_request_session( monkeypatch: pytest.MonkeyPatch, sqlite_session_factory: sessionmaker[Session], diff --git a/api/tests/unit_tests/services/test_app_service.py b/api/tests/unit_tests/services/test_app_service.py index 794e146db97279..9ba0e2af564759 100644 --- a/api/tests/unit_tests/services/test_app_service.py +++ b/api/tests/unit_tests/services/test_app_service.py @@ -14,9 +14,22 @@ from graphon.model_runtime.entities.model_entities import ModelType from models import Account, Tenant from models.account import TenantAccountJoin, TenantAccountRole -from models.agent import Agent, AgentIconType, AgentScope, AgentSource, AgentStatus +from models.agent import ( + Agent, + AgentConfigVersionKind, + AgentHomeSnapshot, + AgentIconType, + AgentScope, + AgentSource, + AgentStatus, + AgentWorkingResourceStatus, + AgentWorkspaceBinding, + WorkflowAgentBindingType, + WorkflowAgentNodeBinding, +) from models.model import App, AppMode, AppModelConfig, IconType -from services.agent.errors import AgentAccessNotReadyError, AgentNameConflictError +from models.workflow import Workflow, WorkflowType +from services.agent.errors import AgentAccessNotReadyError, AgentNameConflictError, AgentWorkflowReferenceConflictError from services.app_service import AppListParams, AppService, CreateAppParams @@ -493,6 +506,15 @@ def test_bound_agent_id_is_none_for_non_agent_app(self): ) assert app.bound_agent_id is None + def test_backing_agent_delete_lookup_locks_agent_row(self) -> None: + app = SimpleNamespace(mode=AppMode.AGENT, tenant_id="tenant-1", id="app-1") + session = MagicMock() + + AppService._get_backing_agent_for_update(app, session=session) + + statement = session.scalar.call_args.args[0] + assert statement._for_update_arg is not None + def test_update_agent_app_syncs_backing_agent_identity(self, sqlite_session: Session): app, backing_agent = _persist_agent_app(sqlite_session) account_id = str(uuid4()) @@ -655,7 +677,8 @@ def test_delete_agent_app_archives_backing_agent(self, sqlite_session: Session): patch( "services.app_service.WorkflowAgentRetirementService.retire_unowned", side_effect=lambda **_kwargs: ( - events.append("retire-workflow-agents") or (["workflow-binding-1"], ["workflow-home-1"]) + events.append("retire-workflow-agents") + or (["workflow-binding-1"], ["workflow-home-1"], ["workflow-agent-1"]) ), ) as mock_workflow_retirement, patch( @@ -693,7 +716,73 @@ def test_delete_agent_app_archives_backing_agent(self, sqlite_session: Session): workspace_ids=["workspace-1"], binding_ids=["workflow-binding-1"], home_snapshot_ids=["home-1", "workflow-home-1"], + purge_agent_ids=[backing_agent.id, "workflow-agent-1"], + ) + + def test_delete_agent_app_rejects_effective_workflow_reference(self, sqlite_session: Session) -> None: + agent_app, backing_agent = _persist_agent_app(sqlite_session) + workflow_app = _persist_app(sqlite_session, tenant_id=agent_app.tenant_id, name="Workflow") + workflow_app.mode = AppMode.WORKFLOW + workflow = Workflow.new( + tenant_id=agent_app.tenant_id, + app_id=workflow_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=agent_app.tenant_id, + app_id=workflow_app.id, + workflow_id=workflow.id, + workflow_version=workflow.version, + node_id="agent-node", + binding_type=WorkflowAgentBindingType.ROSTER_AGENT, + agent_id=backing_agent.id, + current_snapshot_id="snapshot-1", + node_job_config={}, + ) + workspace_binding = AgentWorkspaceBinding( + id="workspace-binding-1", + tenant_id=agent_app.tenant_id, + app_id=agent_app.id, + workspace_id="workspace-1", + agent_id=backing_agent.id, + agent_config_version_id="snapshot-1", + agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT, + backend_binding_ref="backend-binding-1", + status=AgentWorkingResourceStatus.ACTIVE, ) + home = AgentHomeSnapshot( + id="home-1", + tenant_id=agent_app.tenant_id, + agent_id=backing_agent.id, + snapshot_ref="home-ref-1", + status=AgentWorkingResourceStatus.ACTIVE, + ) + sqlite_session.add_all([workflow, binding, workspace_binding, home]) + sqlite_session.commit() + + with ( + patch("services.app_service.app_was_deleted.send") as deleted_signal, + patch("services.app_service.enqueue_agent_resource_collection") as enqueue_collection, + pytest.raises(AgentWorkflowReferenceConflictError, match="1 active workflow app"), + ): + AppService().delete_app(agent_app, session=sqlite_session) + + assert sqlite_session.get(App, agent_app.id) is not None + assert sqlite_session.get(Agent, backing_agent.id).status == AgentStatus.ACTIVE # type: ignore[union-attr] + assert ( + sqlite_session.get(AgentWorkspaceBinding, workspace_binding.id).status # type: ignore[union-attr] + == AgentWorkingResourceStatus.ACTIVE + ) + assert sqlite_session.get(AgentHomeSnapshot, home.id).status == AgentWorkingResourceStatus.ACTIVE # type: ignore[union-attr] + deleted_signal.assert_not_called() + enqueue_collection.assert_not_called() def test_delete_app_commit_failure_does_not_retire_workflow_agents_or_enqueue(self, sqlite_session: Session): app = _persist_app(sqlite_session, tenant_id=str(uuid4())) diff --git a/api/tests/unit_tests/services/test_snippet_dsl_service.py b/api/tests/unit_tests/services/test_snippet_dsl_service.py index a9bbdc984959ca..74116569a77816 100644 --- a/api/tests/unit_tests/services/test_snippet_dsl_service.py +++ b/api/tests/unit_tests/services/test_snippet_dsl_service.py @@ -533,12 +533,18 @@ def test_create_or_update_snippet_updates_existing_snippet_and_syncs_workflow(mo monkeypatch.setattr("services.snippet_dsl_service.SnippetService", lambda *_args, **_kwargs: snippet_service) monkeypatch.setattr( "services.snippet_dsl_service.WorkflowAgentPublishService.sync_agent_bindings_for_draft", - Mock(return_value=set()), + Mock(return_value={"retired-agent"}), ) monkeypatch.setattr( "services.snippet_dsl_service.WorkflowAgentPublishService.validate_agent_nodes_for_draft_sync", Mock(), ) + monkeypatch.setattr( + "services.snippet_dsl_service.WorkflowAgentRetirementService.retire_unowned", + Mock(return_value=(["binding-1"], ["home-1"], ["retired-agent"])), + ) + enqueue_collection = Mock() + monkeypatch.setattr("services.snippet_dsl_service.enqueue_agent_resource_collection", enqueue_collection) result = service._create_or_update_snippet( snippet=snippet, @@ -561,6 +567,12 @@ def test_create_or_update_snippet_updates_existing_snippet_and_syncs_workflow(mo assert snippet.icon_info == {"icon": "x"} snippet_service.sync_draft_workflow.assert_called_once() session.commit.assert_called_once() + enqueue_collection.assert_called_once_with( + tenant_id="tenant-1", + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + purge_agent_ids=["retired-agent"], + ) def test_create_or_update_snippet_creates_new_snippet_and_flushes(monkeypatch: pytest.MonkeyPatch): diff --git a/api/tests/unit_tests/services/test_snippet_service.py b/api/tests/unit_tests/services/test_snippet_service.py index 21a9deb824e223..79bfb598a3bed1 100644 --- a/api/tests/unit_tests/services/test_snippet_service.py +++ b/api/tests/unit_tests/services/test_snippet_service.py @@ -13,9 +13,16 @@ from extensions.storage.storage_type import StorageType from graphon.variables.segments import StringSegment from graphon.variables.types import SegmentType -from models.agent import Agent, AgentScope, AgentSource, AgentStatus -from models.enums import CreatorUserRole -from models.model import UploadFile +from models.agent import ( + Agent, + AgentScope, + AgentSource, + AgentStatus, + WorkflowAgentBindingType, + WorkflowAgentNodeBinding, +) +from models.enums import AppStatus, CreatorUserRole +from models.model import App, AppMode, UploadFile from models.snippet import CustomizedSnippet, SnippetType from models.workflow import ( Workflow, @@ -207,6 +214,16 @@ def test_sync_draft_workflow_creates_draft_and_updates_input_fields( ) -> None: service = SnippetService(session_maker=sqlite_session_factory) monkeypatch.setattr(service, "get_draft_workflow", Mock(return_value=None)) + monkeypatch.setattr( + "services.agent.workflow_publish_service.WorkflowAgentPublishService.sync_agent_bindings_for_draft", + Mock(return_value={"retired-agent"}), + ) + monkeypatch.setattr( + "services.snippet_service.WorkflowAgentRetirementService.retire_unowned", + Mock(return_value=(["binding-1"], ["home-1"], ["retired-agent"])), + ) + enqueue_collection = Mock() + monkeypatch.setattr("services.snippet_service.enqueue_agent_resource_collection", enqueue_collection) snippet = _snippet() account = SimpleNamespace(id="account-1") @@ -227,6 +244,12 @@ def test_sync_draft_workflow_creates_draft_and_updates_input_fields( assert stored_workflow is not None assert stored_snippet is not None assert stored_snippet.input_fields_list == [{"variable": "query"}] + enqueue_collection.assert_called_once_with( + tenant_id=snippet.tenant_id, + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + purge_agent_ids=["retired-agent"], + ) def test_sync_draft_workflow_raises_when_hash_mismatches( @@ -394,6 +417,16 @@ def test_restore_published_snippet_workflow_to_draft_copies_source_snapshot( monkeypatch.setattr(service, "get_published_workflow_by_id", Mock(return_value=source_workflow)) monkeypatch.setattr(service, "get_draft_workflow", Mock(return_value=draft_workflow)) + monkeypatch.setattr( + "services.agent.workflow_publish_service.WorkflowAgentPublishService.restore_agent_node_bindings_to_draft", + Mock(return_value={"retired-agent"}), + ) + monkeypatch.setattr( + "services.snippet_service.WorkflowAgentRetirementService.retire_unowned", + Mock(return_value=(["binding-1"], ["home-1"], ["retired-agent"])), + ) + enqueue_collection = Mock() + monkeypatch.setattr("services.snippet_service.enqueue_agent_resource_collection", enqueue_collection) result = service.restore_published_workflow_to_draft( snippet=snippet, @@ -409,6 +442,12 @@ def test_restore_published_snippet_workflow_to_draft_copies_source_snapshot( stored = sqlite_session.get(Workflow, draft_workflow.id) assert stored is not None assert stored.graph_dict == source_graph + enqueue_collection.assert_called_once_with( + tenant_id=snippet.tenant_id, + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + purge_agent_ids=["retired-agent"], + ) def test_restore_published_snippet_workflow_to_draft_raises_when_source_missing( @@ -629,7 +668,9 @@ def test_delete_snippet_archives_owned_agents_and_schedules_backing_app_cleanup( sqlite_session.add_all([snippet, agent]) sqlite_session.flush() cleanup_delay = Mock() + enqueue_collection = Mock() monkeypatch.setattr("tasks.remove_app_and_related_data_task.remove_app_and_related_data_task.delay", cleanup_delay) + monkeypatch.setattr("services.snippet_service.enqueue_agent_resource_collection", enqueue_collection) result = SnippetService.delete_snippet( session=sqlite_session, @@ -638,13 +679,89 @@ def test_delete_snippet_archives_owned_agents_and_schedules_backing_app_cleanup( ) assert result is True - assert agent.status == "archived" - assert agent.archived_by == "account-1" - assert agent.archived_at is not None - assert agent.updated_by == "account-1" + assert agent.status == AgentStatus.ACTIVE sqlite_session.commit() - assert sqlite_session.get(Agent, agent.id).status == AgentStatus.ARCHIVED + sqlite_session.expire_all() + stored_agent = sqlite_session.get(Agent, agent.id) + assert stored_agent is not None + assert stored_agent.status == AgentStatus.ARCHIVED + assert stored_agent.archived_by == "account-1" + assert stored_agent.archived_at is not None + assert stored_agent.updated_by == "account-1" cleanup_delay.assert_called_once_with(tenant_id=snippet.tenant_id, app_id="backing-app-1") + enqueue_collection.assert_called_once_with( + tenant_id=snippet.tenant_id, + binding_ids=[], + home_snapshot_ids=[], + purge_agent_ids=[agent.id], + ) + + +def test_delete_snippet_keeps_agent_with_effective_external_owner( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +) -> None: + snippet = _snippet() + agent = Agent( + id="agent-1", + tenant_id=snippet.tenant_id, + name="Shared workflow Agent", + description="", + role="", + scope=AgentScope.WORKFLOW_ONLY, + source=AgentSource.WORKFLOW, + app_id=snippet.id, + backing_app_id="backing-app-1", + status=AgentStatus.ACTIVE, + ) + workflow_app = App( + id="app-1", + tenant_id=snippet.tenant_id, + name="Workflow", + mode=AppMode.WORKFLOW, + status=AppStatus.NORMAL, + enable_site=True, + enable_api=True, + ) + workflow = Workflow( + id="workflow-1", + tenant_id=snippet.tenant_id, + app_id=workflow_app.id, + type=WorkflowType.WORKFLOW, + version=Workflow.VERSION_DRAFT, + graph="{}", + features="{}", + created_by="account-1", + environment_variables=[], + conversation_variables=[], + rag_pipeline_variables=[], + ) + external_binding = WorkflowAgentNodeBinding( + tenant_id=snippet.tenant_id, + app_id=workflow_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="snapshot-1", + node_job_config={}, + ) + sqlite_session.add_all([snippet, agent, workflow_app, workflow, external_binding]) + sqlite_session.flush() + enqueue_collection = Mock() + cleanup_delay = Mock() + monkeypatch.setattr("services.snippet_service.enqueue_agent_resource_collection", enqueue_collection) + monkeypatch.setattr("tasks.remove_app_and_related_data_task.remove_app_and_related_data_task.delay", cleanup_delay) + + SnippetService.delete_snippet(session=sqlite_session, snippet=snippet, account_id="account-1") + sqlite_session.commit() + sqlite_session.expire_all() + + assert sqlite_session.get(Agent, agent.id).status == AgentStatus.ACTIVE # type: ignore[union-attr] + assert sqlite_session.get(WorkflowAgentNodeBinding, external_binding.id) is not None + enqueue_collection.assert_not_called() + cleanup_delay.assert_not_called() def test_delete_draft_variable_files_removes_storage_objects( diff --git a/api/tests/unit_tests/services/test_workflow_service.py b/api/tests/unit_tests/services/test_workflow_service.py index d0ff4a085e9477..b2d3446ed233f1 100644 --- a/api/tests/unit_tests/services/test_workflow_service.py +++ b/api/tests/unit_tests/services/test_workflow_service.py @@ -419,7 +419,18 @@ def test_sync_draft_workflow_creates_new_draft(self, workflow_service: WorkflowS graph = TestWorkflowAssociatedDataFactory.create_valid_workflow_graph() features = {"file_upload": {"enabled": False}} - with patch("services.workflow_service.app_draft_workflow_was_synced"): + with ( + patch("services.workflow_service.app_draft_workflow_was_synced"), + patch( + "services.agent.workflow_publish_service.WorkflowAgentPublishService.sync_agent_bindings_for_draft", + return_value={"retired-agent"}, + ), + patch( + "services.workflow_service.WorkflowAgentRetirementService.retire_unowned", + return_value=(["binding-1"], ["home-1"], ["retired-agent"]), + ), + patch("services.workflow_service.enqueue_agent_resource_collection") as enqueue_collection, + ): result = workflow_service.sync_draft_workflow( app_model=app, graph=graph, @@ -435,6 +446,12 @@ def test_sync_draft_workflow_creates_new_draft(self, workflow_service: WorkflowS assert persisted_workflow is result assert result.graph_dict == graph assert result.features_dict == features + enqueue_collection.assert_called_once_with( + tenant_id=app.tenant_id, + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + purge_agent_ids=["retired-agent"], + ) def test_sync_draft_workflow_updates_existing_draft( self, workflow_service: WorkflowService, sqlite_session: Session @@ -759,7 +776,18 @@ def test_restore_published_workflow_to_draft_keeps_source_features_unmodified( sqlite_session.add_all([source_workflow, draft_workflow]) sqlite_session.commit() - with patch("services.workflow_service.app_draft_workflow_was_synced"): + with ( + patch("services.workflow_service.app_draft_workflow_was_synced"), + patch( + "services.agent.workflow_publish_service.WorkflowAgentPublishService.restore_agent_node_bindings_to_draft", + return_value={"retired-agent"}, + ), + patch( + "services.workflow_service.WorkflowAgentRetirementService.retire_unowned", + return_value=(["binding-1"], ["home-1"], ["retired-agent"]), + ), + patch("services.workflow_service.enqueue_agent_resource_collection") as enqueue_collection, + ): result = workflow_service.restore_published_workflow_to_draft( app_model=app, workflow_id=source_workflow.id, @@ -772,6 +800,12 @@ def test_restore_published_workflow_to_draft_keeps_source_features_unmodified( assert draft_workflow.serialized_features == json.dumps(legacy_features) sqlite_session.refresh(draft_workflow) assert draft_workflow.serialized_features == json.dumps(legacy_features) + enqueue_collection.assert_called_once_with( + tenant_id=app.tenant_id, + binding_ids=["binding-1"], + home_snapshot_ids=["home-1"], + purge_agent_ids=["retired-agent"], + ) # ==================== Workflow Validation Tests ==================== # These tests verify graph structure and feature configuration validation diff --git a/api/tests/unit_tests/tasks/test_collect_agent_resources_task.py b/api/tests/unit_tests/tasks/test_collect_agent_resources_task.py index eb751061f17bf9..d660ec489fe513 100644 --- a/api/tests/unit_tests/tasks/test_collect_agent_resources_task.py +++ b/api/tests/unit_tests/tasks/test_collect_agent_resources_task.py @@ -1,8 +1,9 @@ from typing import Protocol, cast -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import pytest +from services.agent.deletion_service import AgentDeletionService from services.agent.home_snapshot_service import AgentHomeSnapshotService from services.agent.workspace_service import AgentWorkspaceService from tasks.collect_agent_resources_task import ( @@ -29,6 +30,7 @@ def test_enqueue_deduplicates_ids_and_skips_empty_input(monkeypatch: pytest.Monk tenant_id="tenant-1", binding_ids=["binding-2", "binding-1", "binding-2"], workspace_ids=["workspace-1"], + purge_agent_ids=["agent-2", "", "agent-1", "agent-2"], ) delay.assert_called_once_with( @@ -36,6 +38,7 @@ def test_enqueue_deduplicates_ids_and_skips_empty_input(monkeypatch: pytest.Monk binding_ids=["binding-1", "binding-2"], workspace_ids=["workspace-1"], home_snapshot_ids=[], + purge_agent_ids=["agent-1", "agent-2"], ) @@ -56,67 +59,126 @@ def test_collection_runs_in_workspace_binding_snapshot_order(monkeypatch: pytest "collect_retired_home_snapshot", lambda **_kwargs: calls.append("home"), ) + purge = MagicMock(side_effect=lambda **_kwargs: calls.append("purge")) + monkeypatch.setattr(AgentDeletionService, "purge_archived_agents", purge) collect_agent_resources.run( tenant_id="tenant-1", workspace_ids=["workspace-1"], binding_ids=["binding-1"], home_snapshot_ids=["home-1"], + purge_agent_ids=["agent-1"], ) - assert calls == ["workspace", "binding", "home"] + assert calls == ["workspace", "binding", "home", "purge"] + purge.assert_called_once_with(tenant_id="tenant-1", agent_ids=["agent-1"]) -def test_collection_failure_propagates_and_stops_task(monkeypatch: pytest.MonkeyPatch) -> None: +def test_collection_failure_propagates_after_attempting_remaining_resources(monkeypatch: pytest.MonkeyPatch) -> None: calls: list[str] = [] - error = RuntimeError("workspace failed") + first_error = RuntimeError("workspace-1 failed") + errors = { + "workspace-1": first_error, + "workspace-2": RuntimeError("workspace-2 failed"), + "binding-1": RuntimeError("binding-1 failed"), + "home-2": RuntimeError("home-2 failed"), + } log_exception = MagicMock() - def collect_workspace(**_kwargs: object) -> None: - calls.append("workspace") - raise error + def collect_workspace(*, workspace_id: str, **_kwargs: object) -> None: + calls.append(f"workspace:{workspace_id}") + if error := errors.get(workspace_id): + raise error + + def collect_binding(*, binding_id: str, **_kwargs: object) -> None: + calls.append(f"binding:{binding_id}") + if error := errors.get(binding_id): + raise error + + def collect_home(*, home_snapshot_id: str, **_kwargs: object) -> None: + calls.append(f"home:{home_snapshot_id}") + if error := errors.get(home_snapshot_id): + raise error monkeypatch.setattr(AgentWorkspaceService, "collect_retired_workspace", collect_workspace) - monkeypatch.setattr( - AgentWorkspaceService, - "collect_retired_binding", - lambda **_kwargs: calls.append("binding"), + monkeypatch.setattr(AgentWorkspaceService, "collect_retired_binding", collect_binding) + monkeypatch.setattr(AgentHomeSnapshotService, "collect_retired_home_snapshot", collect_home) + monkeypatch.setattr("tasks.collect_agent_resources_task.logger.exception", log_exception) + purge = MagicMock() + monkeypatch.setattr(AgentDeletionService, "purge_archived_agents", purge) + + with pytest.raises(RuntimeError) as exc_info: + collect_agent_resources.run( + tenant_id="tenant-1", + workspace_ids=["workspace-1", "workspace-2", "workspace-3"], + binding_ids=["binding-1", "binding-2"], + home_snapshot_ids=["home-1", "home-2", "home-3"], + purge_agent_ids=["agent-1"], + ) + + assert exc_info.value.__cause__ is first_error + assert str(exc_info.value) == ( + "Failed to collect 4 retired Agent resource(s): " + "workspace:workspace-1, workspace:workspace-2, binding:binding-1, home_snapshot:home-2" ) - monkeypatch.setattr( - AgentHomeSnapshotService, - "collect_retired_home_snapshot", - lambda **_kwargs: calls.append("home"), + assert calls == [ + "workspace:workspace-1", + "workspace:workspace-2", + "workspace:workspace-3", + "binding:binding-1", + "binding:binding-2", + "home:home-1", + "home:home-2", + "home:home-3", + ] + purge.assert_not_called() + log_exception.assert_has_calls( + [ + call( + "Failed to collect retired Agent resource", + extra={ + "tenant_id": "tenant-1", + "resource_type": resource_type, + "resource_id": resource_id, + }, + ) + for resource_type, resource_id in ( + ("workspace", "workspace-1"), + ("workspace", "workspace-2"), + ("binding", "binding-1"), + ("home_snapshot", "home-2"), + ) + ], + any_order=False, ) + assert log_exception.call_count == 4 + + +def test_enqueue_failure_propagates(monkeypatch: pytest.MonkeyPatch) -> None: + error = RuntimeError("queue unavailable") + delay = MagicMock(side_effect=error) + log_exception = MagicMock() + monkeypatch.setattr(collect_agent_resources, "delay", delay) monkeypatch.setattr("tasks.collect_agent_resources_task.logger.exception", log_exception) with pytest.raises(RuntimeError) as exc_info: - collect_agent_resources.run( + enqueue_agent_resource_collection( tenant_id="tenant-1", - workspace_ids=["workspace-1"], binding_ids=["binding-1"], + workspace_ids=["workspace-1"], home_snapshot_ids=["home-1"], + purge_agent_ids=["agent-2", "agent-1", "agent-2"], ) assert exc_info.value is error - assert calls == ["workspace"] + payload = { + "binding_ids": ["binding-1"], + "workspace_ids": ["workspace-1"], + "home_snapshot_ids": ["home-1"], + "purge_agent_ids": ["agent-1", "agent-2"], + } + delay.assert_called_once_with(tenant_id="tenant-1", **payload) log_exception.assert_called_once_with( - "Failed to collect retired Agent resource", - extra={ - "tenant_id": "tenant-1", - "resource_type": "workspace", - "resource_id": "workspace-1", - }, - ) - - -def test_enqueue_failure_is_best_effort(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - collect_agent_resources, - "delay", - MagicMock(side_effect=RuntimeError("queue unavailable")), - ) - - enqueue_agent_resource_collection( - tenant_id="tenant-1", - binding_ids=["binding-1"], + "Failed to enqueue retired Agent resource collection", + extra={"tenant_id": "tenant-1", **payload}, ) diff --git a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py index 3f0ce2475ef099..125e447a10e685 100644 --- a/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py +++ b/api/tests/unit_tests/tasks/test_remove_app_and_related_data_task.py @@ -6,9 +6,12 @@ import pytest from sqlalchemy.orm import Session +import tasks.remove_app_and_related_data_task as remove_app_task_module +from enums import DeploymentEdition from graphon.enums import WorkflowExecutionStatus from libs.archive_storage import ArchiveStorageNotConfiguredError from models import AppStar +from models.agent import WorkflowAgentBindingType, WorkflowAgentNodeBinding from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom from models.workflow import WorkflowArchiveLog from tasks.remove_app_and_related_data_task import ( @@ -17,10 +20,104 @@ _delete_archived_workflow_run_files, _delete_draft_variable_offload_data, _delete_draft_variables, + _delete_workflow_agent_node_bindings, delete_draft_variables_batch, ) +def test_delete_workflow_agent_node_bindings_is_scoped_to_tenant_and_app(sqlite_session: Session) -> None: + target = WorkflowAgentNodeBinding( + tenant_id="tenant-1", + app_id="app-1", + workflow_id="workflow-1", + workflow_version="draft", + node_id="node-1", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="agent-1", + current_snapshot_id="snapshot-1", + node_job_config={}, + ) + kept = WorkflowAgentNodeBinding( + tenant_id="tenant-1", + app_id="app-2", + workflow_id="workflow-2", + workflow_version="draft", + node_id="node-2", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="agent-2", + current_snapshot_id="snapshot-2", + node_job_config={}, + ) + other_tenant = WorkflowAgentNodeBinding( + tenant_id="tenant-2", + app_id="app-1", + workflow_id="workflow-3", + workflow_version="draft", + node_id="node-3", + binding_type=WorkflowAgentBindingType.INLINE_AGENT, + agent_id="agent-3", + current_snapshot_id="snapshot-3", + node_job_config={}, + ) + sqlite_session.add_all([target, kept, other_tenant]) + sqlite_session.commit() + target_id = target.id + kept_id = kept.id + other_tenant_id = other_tenant.id + + _delete_workflow_agent_node_bindings("tenant-1", "app-1") + + sqlite_session.expire_all() + assert sqlite_session.get(WorkflowAgentNodeBinding, target_id) is None + assert sqlite_session.get(WorkflowAgentNodeBinding, kept_id) is not None + assert sqlite_session.get(WorkflowAgentNodeBinding, other_tenant_id) is not None + + +def test_app_cleanup_removes_agent_bindings_before_workflows(monkeypatch: pytest.MonkeyPatch) -> None: + events: list[str] = [] + monkeypatch.setattr(remove_app_task_module.dify_config, "DEPLOYMENT_EDITION", DeploymentEdition.COMMUNITY) + other_cleanup_names = ( + "_delete_app_model_configs", + "_delete_app_site", + "_delete_app_mcp_servers", + "_delete_app_api_tokens", + "_delete_installed_apps", + "_delete_app_stars", + "_delete_recommended_apps", + "_delete_app_annotation_data", + "_delete_app_dataset_joins", + "_delete_app_workflow_runs", + "_delete_app_workflow_node_executions", + "_delete_app_workflow_app_logs", + "_delete_app_conversations", + "_delete_app_messages", + "_delete_workflow_tool_providers", + "_delete_app_tag_bindings", + "_delete_end_users", + "_delete_trace_app_configs", + "_delete_conversation_variables", + "_delete_draft_variables", + "_delete_app_triggers", + "_delete_workflow_plugin_triggers", + "_delete_workflow_webhook_triggers", + "_delete_workflow_schedule_plans", + "_delete_workflow_trigger_logs", + ) + for name in other_cleanup_names: + monkeypatch.setattr(remove_app_task_module, name, MagicMock()) + + delete_bindings = MagicMock(side_effect=lambda *_args: events.append("bindings")) + delete_workflows = MagicMock(side_effect=lambda *_args: events.append("workflows")) + monkeypatch.setattr(remove_app_task_module, "_delete_workflow_agent_node_bindings", delete_bindings) + monkeypatch.setattr(remove_app_task_module, "_delete_app_workflows", delete_workflows) + + remove_app_task_module.remove_app_and_related_data_task.run(tenant_id="tenant-1", app_id="app-1") + + assert events == ["bindings", "workflows"] + delete_bindings.assert_called_once_with("tenant-1", "app-1") + delete_workflows.assert_called_once_with("tenant-1", "app-1") + + class TestDeleteDraftVariablesBatch: def test_delete_draft_variables_batch_invalid_batch_size(self): """Test that invalid batch size raises ValueError.""" diff --git a/dify-agent/docs/dify-agent/concepts/runtime-resources/index.md b/dify-agent/docs/dify-agent/concepts/runtime-resources/index.md index c3bae6fb3c68fd..1196566198c5d8 100644 --- a/dify-agent/docs/dify-agent/concepts/runtime-resources/index.md +++ b/dify-agent/docs/dify-agent/concepts/runtime-resources/index.md @@ -118,10 +118,12 @@ product use without performing network I/O inside the caller's transaction. Product lifecycle paths commit this transition synchronously. After the transaction commits, one Celery task asks Dify Agent to destroy the physical resources. A successful collector deletes the corresponding ledger row. If a -collector raises, the task logs the tenant, resource type, and resource ID, -re-raises the exception so collection stops and Celery records the task as -failed, and leaves the RETIRED row intact. No automatic retry or reconciliation -is performed. +collector raises, the task logs the tenant, resource type, and resource ID and +continues with the other independent resources in the batch. After all resources +have been attempted, any failure makes the Celery task fail and prevents Agent +aggregate deletion. Failed RETIRED rows remain available for a later retry. A +failure to publish the Celery task is also propagated to the product caller. No +automatic retry or reconciliation is performed. The unified `collect_agent_resources` task is registered on normal Celery workers and explicitly uses the existing `retention` queue. Standard workers @@ -130,14 +132,24 @@ is required. At a Workflow terminal event, the graph layer synchronously retires and commits the run's Workspaces before enqueueing collection. When a Workflow change may orphan Workflow-only Agents, the main product transaction commits first; a fresh session then rechecks effective ownership and retires only Agents -that remain unowned. +that remain unowned. An effective reference is a binding in a normal App's +current draft or current published Workflow. Such a reference blocks direct +Agent deletion. Retiring a final Binding also retires its Workspace. Workspace collection destroys the physical Workspace through one Binding and then collects remaining materialized Homes. Home Snapshots are retired when their owning Agent is -retired and are collected only after no draft or config snapshot references -them. Celery performs physical collection only; it does not decide or perform -the initial retirement. Dify Agent itself remains stateless. +retired. `RETIRED` is the sole physical-deletion condition for a Home Snapshot; +Draft and Config Snapshot references are historical pointers and do not keep it +alive. After every external resource in a deletion batch succeeds, Dify API +hard-deletes the archived Agent together with its Drafts, Config Snapshots, +Config Revisions, debug-conversation mappings, resource ledgers, and stale +Workflow Agent bindings in one database transaction. Dify Agent itself remains +stateless. + +A `RETIRED` Workspace without a `RETIRED` Binding cannot identify a backend +participant through which to destroy the Workspace. That state is a lifecycle +invariant violation and fails collection instead of being logged as success. There is currently no age-based TTL, periodic GC, or global orphan reconciler. Backend destroy operations are idempotent where supported. Dify API does not diff --git a/packages/contracts/generated/api/console/agent/types.gen.ts b/packages/contracts/generated/api/console/agent/types.gen.ts index ddef2e3addface..ddf8feafdb4ddc 100644 --- a/packages/contracts/generated/api/console/agent/types.gen.ts +++ b/packages/contracts/generated/api/console/agent/types.gen.ts @@ -2078,6 +2078,7 @@ export type DeleteAgentByAgentIdData = { export type DeleteAgentByAgentIdErrors = { 403: unknown + 409: unknown } export type DeleteAgentByAgentIdResponses = {