from __future__ import annotations import logging from typing import override from clients.agent_backend import AgentBackendError, AgentBackendRunClient, AgentBackendRunRequestBuilder from clients.agent_backend.factory import create_agent_backend_run_client from configs import dify_config from core.workflow.system_variables import SystemVariableKey, get_system_text from graphon.graph_engine.layers import GraphEngineLayer from graphon.graph_events import ( GraphEngineEvent, GraphRunAbortedEvent, GraphRunFailedEvent, GraphRunPartialSucceededEvent, GraphRunSucceededEvent, ) from .session_store import StoredWorkflowAgentSession, WorkflowAgentRuntimeSessionStore logger = logging.getLogger(__name__) # Upper bound on how long a cleanup-only run is allowed to settle before the # layer gives up and leaves the row ACTIVE so it can be retried later. Cleanup # work is mostly local agent-backend bookkeeping (no LLM inference), so 30s is # generous; a hung backend should never block workflow termination beyond this. _CLEANUP_WAIT_TIMEOUT_SECONDS = 30.0 class WorkflowAgentSessionCleanupLayer(GraphEngineLayer): """Retires workflow Agent session snapshots when a workflow reaches a terminal state. Implementation notes — there are two failure modes the cleanup path has to avoid simultaneously: 1. The agenton compositor on the agent-backend side validates the cleanup request's session snapshot against the replayed composition before running any lifecycle hook. If the snapshot's layer names diverge from the composition, the run fails asynchronously with ``run_failed`` — but the initial ``POST /runs`` already returned 202, so the API side has no visibility of the failure unless it waits for terminal status. The ``composition_layer_specs`` persistence in A.1–A.4 plus the ``_filter_snapshot_to_specs`` shape in ``build_cleanup_request`` keeps the two name lists in sync. 2. The current agent backend's ``runner.py::_run_agent`` always invokes ``run.get_layer("llm")`` and the structured-output / history validators before exiting any slot — there is no ``purpose: "cleanup"`` branch yet. A truly cleanup-only request (no LLM layer) therefore still crashes inside the runner with ``Layer 'llm' is not defined in this compositor run.``. Until the backend grows a cleanup-only purpose, this layer **does not issue an HTTP cleanup run**: it simply retires the local snapshot row so stale state cannot be re-resumed, and lets the agent backend's own retention TTL release the suspended layers. The HTTP-cleanup machinery (``build_cleanup_request`` + ``wait_run``) is intentionally still wired into the request builder + integration tests so that when the agent backend supports cleanup runs we can flip the switch here with a one-line change (see ``_HTTP_CLEANUP_SUPPORTED``). """ # Flip to True once dify-agent's runner has a ``purpose=cleanup`` branch # that skips the LLM/output/user-prompt invariants. Until then we only # update the local row; the spec list is still persisted so the future # HTTP cleanup path has everything it needs. _HTTP_CLEANUP_SUPPORTED: bool = False _TERMINAL_EVENTS = ( GraphRunSucceededEvent, GraphRunPartialSucceededEvent, GraphRunFailedEvent, GraphRunAbortedEvent, ) def __init__( self, *, session_store: WorkflowAgentRuntimeSessionStore, request_builder: AgentBackendRunRequestBuilder, agent_backend_client: AgentBackendRunClient | None, cleanup_wait_timeout_seconds: float = _CLEANUP_WAIT_TIMEOUT_SECONDS, ) -> None: super().__init__() self._session_store = session_store self._request_builder = request_builder self._agent_backend_client = agent_backend_client self._cleanup_wait_timeout_seconds = cleanup_wait_timeout_seconds @override def on_graph_start(self) -> None: return @override def on_event(self, event: GraphEngineEvent) -> None: if not isinstance(event, self._TERMINAL_EVENTS): return workflow_run_id = get_system_text( self.graph_runtime_state.variable_pool, SystemVariableKey.WORKFLOW_EXECUTION_ID, ) if not workflow_run_id: logger.warning("Skipping workflow Agent session cleanup: workflow_run_id is missing.") return for stored_session in self._session_store.list_active_sessions(workflow_run_id=workflow_run_id): self._cleanup_session(stored_session) @override def on_graph_end(self, error: Exception | None) -> None: return def _cleanup_session(self, stored_session: StoredWorkflowAgentSession) -> None: scope = stored_session.scope if not self._HTTP_CLEANUP_SUPPORTED: # Agent backend has no cleanup-only run mode yet (see class # docstring). Retire the local row so future re-entries do not # resume from stale state, and let the backend's retention TTL # release the suspended layers on its own schedule. logger.info( "Workflow Agent session retired locally; HTTP cleanup is disabled " "until the agent backend supports a cleanup-only run mode. " "workflow_run_id=%s node_id=%s binding_id=%s agent_id=%s previous_run_id=%s", scope.workflow_run_id, scope.node_id, scope.binding_id, scope.agent_id, stored_session.backend_run_id, ) self._session_store.mark_cleaned(scope=scope, backend_run_id=stored_session.backend_run_id) return if self._agent_backend_client is None: # HTTP cleanup was enabled by the caller but no client was wired # in (e.g. the API runs without AGENT_BACKEND_BASE_URL configured). # Leave the row ACTIVE so an operator restart with proper config # can drive the cleanup; do not silently retire it. logger.warning( "Skipping Agent backend cleanup: HTTP cleanup is enabled but no agent " "backend client is wired in. workflow_run_id=%s node_id=%s agent_id=%s", scope.workflow_run_id, scope.node_id, scope.agent_id, ) return if not stored_session.composition_layer_specs: # Sessions persisted before A.1 landed do not carry the spec list, # so we cannot replay a valid cleanup composition. Leave the row # ACTIVE and warn so the absence shows up in observability rather # than being silently swallowed by a doomed cleanup run. logger.warning( "Skipping Agent backend cleanup: no composition_layer_specs persisted. " "workflow_run_id=%s node_id=%s agent_id=%s", scope.workflow_run_id, scope.node_id, scope.agent_id, ) return request = self._request_builder.build_cleanup_request( session_snapshot=stored_session.session_snapshot, composition_layer_specs=stored_session.composition_layer_specs, idempotency_key=f"{scope.workflow_run_id}:{scope.node_id}:{scope.binding_id}:agent-session-cleanup", metadata={ "tenant_id": scope.tenant_id, "app_id": scope.app_id, "workflow_id": scope.workflow_id, "workflow_run_id": scope.workflow_run_id, "node_id": scope.node_id, "node_execution_id": scope.node_execution_id, "binding_id": scope.binding_id, "agent_id": scope.agent_id, "agent_config_snapshot_id": scope.agent_config_snapshot_id, "previous_agent_backend_run_id": stored_session.backend_run_id, }, ) try: response = self._agent_backend_client.create_run(request) except AgentBackendError: logger.warning( "Agent backend session cleanup request failed: workflow_run_id=%s node_id=%s agent_id=%s", scope.workflow_run_id, scope.node_id, scope.agent_id, exc_info=True, ) return try: status_response = self._agent_backend_client.wait_run( response.run_id, timeout_seconds=self._cleanup_wait_timeout_seconds ) except AgentBackendError: logger.warning( "Agent backend session cleanup wait_run failed: " "workflow_run_id=%s node_id=%s agent_id=%s cleanup_run_id=%s", scope.workflow_run_id, scope.node_id, scope.agent_id, response.run_id, exc_info=True, ) return if status_response.status != "succeeded": logger.warning( "Agent backend session cleanup did not succeed: status=%s error=%s " "workflow_run_id=%s node_id=%s agent_id=%s cleanup_run_id=%s", status_response.status, status_response.error, scope.workflow_run_id, scope.node_id, scope.agent_id, response.run_id, ) return self._session_store.mark_cleaned(scope=scope, backend_run_id=response.run_id) def build_workflow_agent_session_cleanup_layer() -> WorkflowAgentSessionCleanupLayer: """Wire the cleanup layer with the standard production dependencies. The agent backend client is constructed only when ``AGENT_BACKEND_BASE_URL`` is configured (or the deterministic fake is explicitly enabled). When neither is set — for example unit tests that bring up the workflow runner without an Agent node — we pass ``None`` so the layer stays harmless. With ``_HTTP_CLEANUP_SUPPORTED = False`` the local-retire branch never touches the client anyway, but keeping it ``None`` avoids importing httpx and lets test harnesses skip backend configuration. """ agent_backend_client: AgentBackendRunClient | None if dify_config.AGENT_BACKEND_USE_FAKE or dify_config.AGENT_BACKEND_BASE_URL: agent_backend_client = create_agent_backend_run_client( base_url=dify_config.AGENT_BACKEND_BASE_URL, use_fake=dify_config.AGENT_BACKEND_USE_FAKE, fake_scenario=dify_config.AGENT_BACKEND_FAKE_SCENARIO, ) else: agent_backend_client = None return WorkflowAgentSessionCleanupLayer( session_store=WorkflowAgentRuntimeSessionStore(), request_builder=AgentBackendRunRequestBuilder(), agent_backend_client=agent_backend_client, )