diff --git a/api/.env.example b/api/.env.example index 11be68e43d7..c7e7abfe3ea 100644 --- a/api/.env.example +++ b/api/.env.example @@ -570,8 +570,8 @@ WORKFLOW_GENERATOR_NODE_BUILDER_MAX_WORKERS=6 GRAPH_ENGINE_MIN_WORKERS=3 # Maximum number of workers per GraphEngine instance (default: 10) GRAPH_ENGINE_MAX_WORKERS=10 -# Queue depth threshold that triggers worker scale up (default: 3) -GRAPH_ENGINE_SCALE_UP_THRESHOLD=3 +# Pending task threshold that triggers worker scale up (default: 0) +GRAPH_ENGINE_SCALE_UP_THRESHOLD=0 # Seconds of idle time before scaling down workers (default: 5.0) GRAPH_ENGINE_SCALE_DOWN_IDLE_TIME=5.0 diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 5dc8aed12c8..6cf125543ac 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -924,9 +924,9 @@ class WorkflowConfig(BaseSettings): default=10, ) - GRAPH_ENGINE_SCALE_UP_THRESHOLD: PositiveInt = Field( - description="Queue depth threshold that triggers worker scale up", - default=3, + GRAPH_ENGINE_SCALE_UP_THRESHOLD: NonNegativeInt = Field( + description="Pending task threshold that triggers worker scale up", + default=0, ) GRAPH_ENGINE_SCALE_DOWN_IDLE_TIME: float = Field( diff --git a/api/controllers/console/human_input_form.py b/api/controllers/console/human_input_form.py index 195254e1145..4cbf29a89f1 100644 --- a/api/controllers/console/human_input_form.py +++ b/api/controllers/console/human_input_form.py @@ -215,6 +215,7 @@ class ConsoleWorkflowEventsApi(Resource): raise InvalidArgumentError(f"cannot subscribe to workflow run, workflow_run_id={workflow_run.id}") include_state_snapshot = request.args.get("include_state_snapshot", "false").lower() == "true" + continue_on_pause = request.args.get("continue_on_pause", "false").lower() == "true" def _generate_stream_events(): if include_state_snapshot: @@ -225,6 +226,8 @@ class ConsoleWorkflowEventsApi(Resource): tenant_id=workflow_run.tenant_id, app_id=workflow_run.app_id, session_maker=session_maker, + human_input_surface=HumanInputSurface.CONSOLE, + close_on_pause=not continue_on_pause, ) ) return generator.convert_to_event_stream( diff --git a/api/controllers/web/workflow_events.py b/api/controllers/web/workflow_events.py index 48eba33f04c..2954b9a81de 100644 --- a/api/controllers/web/workflow_events.py +++ b/api/controllers/web/workflow_events.py @@ -86,6 +86,7 @@ class WorkflowEventsApi(WebApiResource): raise InvalidArgumentError(f"cannot subscribe to workflow run, workflow_run_id={workflow_run.id}") include_state_snapshot = request.args.get("include_state_snapshot", "false").lower() == "true" + continue_on_pause = request.args.get("continue_on_pause", "false").lower() == "true" def _generate_stream_events(): if include_state_snapshot: @@ -96,6 +97,7 @@ class WorkflowEventsApi(WebApiResource): tenant_id=app_model.tenant_id, app_id=app_model.id, session_maker=session_maker, + close_on_pause=not continue_on_pause, ) ) return generator.convert_to_event_stream( diff --git a/api/core/app/apps/advanced_chat/app_generator.py b/api/core/app/apps/advanced_chat/app_generator.py index 912641b5e8d..aa963683734 100644 --- a/api/core/app/apps/advanced_chat/app_generator.py +++ b/api/core/app/apps/advanced_chat/app_generator.py @@ -685,6 +685,12 @@ class AdvancedChatAppGenerator(MessageBasedAppGenerator): ) if workflow is None: raise ValueError("Workflow not found") + if graph_runtime_state is not None: + self._restore_workflow_run_graph( + session=session, + workflow=workflow, + workflow_run_id=application_generate_entity.workflow_run_id, + ) # Determine system_user_id based on invocation source is_external_api_call = application_generate_entity.invoke_from in { diff --git a/api/core/app/apps/base_app_generator.py b/api/core/app/apps/base_app_generator.py index 71831bbb6a7..aa49d90d1d9 100644 --- a/api/core/app/apps/base_app_generator.py +++ b/api/core/app/apps/base_app_generator.py @@ -5,6 +5,7 @@ from contextlib import AbstractContextManager, nullcontext from typing import TYPE_CHECKING, Any, Union, final from sqlalchemy.orm import Session +from sqlalchemy.orm.attributes import set_committed_value from core.app.apps.draft_variable_saver import ( DraftVariableSaver, @@ -19,7 +20,7 @@ from graphon.enums import NodeType from graphon.file import File, FileUploadConfig from graphon.variables.input_entities import VariableEntityType from libs.orjson import orjson_dumps -from models import Account, EndUser +from models import Account, EndUser, Workflow, WorkflowRun from services.workflow_draft_variable_service import DraftVariableSaver as DraftVariableSaverImpl if TYPE_CHECKING: @@ -70,6 +71,15 @@ class _DebuggerDraftVariableSaver: class BaseAppGenerator: _file_access_controller: DatabaseFileAccessController = DatabaseFileAccessController() + @staticmethod + def _restore_workflow_run_graph(*, session: Session, workflow: Workflow, workflow_run_id: str | None) -> None: + if workflow_run_id is None: + raise ValueError("Workflow run id is required when resuming") + workflow_run = session.get(WorkflowRun, workflow_run_id) + if workflow_run is None or workflow_run.graph is None: + raise ValueError(f"Workflow run graph not found: {workflow_run_id}") + set_committed_value(workflow, "graph", workflow_run.graph) + @staticmethod def _join_worker_thread(worker_thread: threading.Thread) -> None: # Bound the wait so a leaked app worker cannot occupy an execution slot indefinitely. diff --git a/api/core/app/apps/message_based_app_queue_manager.py b/api/core/app/apps/message_based_app_queue_manager.py index b253d93ee52..1be1e532694 100644 --- a/api/core/app/apps/message_based_app_queue_manager.py +++ b/api/core/app/apps/message_based_app_queue_manager.py @@ -10,6 +10,7 @@ from core.app.entities.queue_entities import ( QueueErrorEvent, QueueMessageEndEvent, QueueStopEvent, + QueueWorkflowPausedEvent, ) from models.model import AppMode @@ -43,7 +44,12 @@ class MessageBasedAppQueueManager(AppQueueManager): self._q.put(message) if isinstance( - event, QueueStopEvent | QueueErrorEvent | QueueMessageEndEvent | QueueAdvancedChatMessageEndEvent + event, + QueueStopEvent + | QueueErrorEvent + | QueueMessageEndEvent + | QueueAdvancedChatMessageEndEvent + | QueueWorkflowPausedEvent, ): self.stop_listen(execution_terminal=True) diff --git a/api/core/app/apps/workflow/app_generator.py b/api/core/app/apps/workflow/app_generator.py index db1c49cd6d0..2d6c1512c48 100644 --- a/api/core/app/apps/workflow/app_generator.py +++ b/api/core/app/apps/workflow/app_generator.py @@ -648,6 +648,12 @@ class WorkflowAppGenerator(BaseAppGenerator): raise ValueError("Workflow not found") workflow = self._ensure_snippet_start_node_in_worker(session=session, workflow=workflow) + if graph_runtime_state is not None: + self._restore_workflow_run_graph( + session=session, + workflow=workflow, + workflow_run_id=application_generate_entity.workflow_execution_id, + ) # Determine system_user_id based on invocation source is_external_api_call = application_generate_entity.invoke_from in { diff --git a/api/core/app/apps/workflow_app_runner.py b/api/core/app/apps/workflow_app_runner.py index 3d2857f130a..84d740774b1 100644 --- a/api/core/app/apps/workflow_app_runner.py +++ b/api/core/app/apps/workflow_app_runner.py @@ -55,6 +55,7 @@ from core.workflow.variable_pool_initializer import add_variables_to_pool from core.workflow.workflow_entry import WorkflowEntry from core.workflow.workflow_run_outputs import project_node_outputs_for_workflow_run from graphon.entities.graph_config import NodeConfigDictAdapter +from graphon.entities.pause_reason import HitlRequired from graphon.graph import Graph from graphon.graph_engine.layers import GraphEngineLayer from graphon.graph_events import ( @@ -433,7 +434,9 @@ class WorkflowBasedAppRunner: ) case GraphRunPausedEvent(): runtime_state = workflow_entry.graph_engine.graph_runtime_state - paused_nodes = runtime_state.get_paused_nodes() + paused_nodes = list( + dict.fromkeys(reason.node_id for reason in event.reasons if isinstance(reason, HitlRequired)) + ) enriched_reasons = enrich_graph_pause_reasons( reasons=event.reasons, form_repository=HumanInputFormSubmissionRepository(), diff --git a/api/core/app/workflow/layers/persistence.py b/api/core/app/workflow/layers/persistence.py index 415c8a1826d..8b04661deb3 100644 --- a/api/core/app/workflow/layers/persistence.py +++ b/api/core/app/workflow/layers/persistence.py @@ -24,7 +24,7 @@ from core.workflow.node_execution_process_data import preserve_workflow_agent_bi from core.workflow.system_variables import SystemVariableKey from core.workflow.variable_prefixes import SYSTEM_VARIABLE_NODE_ID from core.workflow.workflow_run_outputs import project_node_outputs_for_workflow_run -from graphon.entities import WorkflowExecution, WorkflowNodeExecution +from graphon.entities import WorkflowExecution, WorkflowNodeExecution, WorkflowStartReason from graphon.enums import ( BuiltinNodeTypes, WorkflowExecutionStatus, @@ -118,7 +118,7 @@ class WorkflowPersistenceLayer(GraphEngineLayer): def on_event(self, event: GraphEngineEvent) -> None: match event: case GraphRunStartedEvent(): - self._handle_graph_run_started() + self._handle_graph_run_started(event) case GraphRunSucceededEvent(): self._handle_graph_run_succeeded(event) case GraphRunPartialSucceededEvent(): @@ -149,7 +149,7 @@ class WorkflowPersistenceLayer(GraphEngineLayer): # ------------------------------------------------------------------ # Graph-level handlers # ------------------------------------------------------------------ - def _handle_graph_run_started(self) -> None: + def _handle_graph_run_started(self, event: GraphRunStartedEvent | None = None) -> None: execution_id = self._get_execution_id() workflow_execution = WorkflowExecution.new( id_=execution_id, @@ -163,6 +163,10 @@ class WorkflowPersistenceLayer(GraphEngineLayer): self._workflow_execution_repository.save(workflow_execution) self._workflow_execution = workflow_execution + if event is not None and event.reason == WorkflowStartReason.RESUMPTION: + node_executions = self._workflow_node_execution_repository.get_by_workflow_execution(execution_id) + self._node_execution_cache = {execution.id: execution for execution in node_executions} + self._node_sequence = max((execution.index for execution in node_executions), default=0) def _handle_graph_run_succeeded(self, event: GraphRunSucceededEvent) -> None: execution = self._get_workflow_execution() diff --git a/api/core/plugin/impl/model.py b/api/core/plugin/impl/model.py index c69be8a3933..ee1bd11901e 100644 --- a/api/core/plugin/impl/model.py +++ b/api/core/plugin/impl/model.py @@ -19,6 +19,7 @@ from graphon.model_runtime.entities.message_entities import PromptMessage, Promp from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelType from graphon.model_runtime.entities.rerank_entities import MultimodalRerankInput, RerankResult from graphon.model_runtime.entities.text_embedding_entities import EmbeddingResult +from graphon.model_runtime.protocols.tts_runtime import TTSModelVoice from graphon.model_runtime.utils.encoders import jsonable_encoder _POLLING_UNSUPPORTED_INVOKE_ERROR_TYPES = frozenset((NotImplementedError.__name__,)) @@ -614,7 +615,7 @@ class PluginModelClient(BasePluginClient): model: str, credentials: dict[str, Any], language: str | None = None, - ): + ) -> list[TTSModelVoice]: """ Get tts model voices """ @@ -641,7 +642,7 @@ class PluginModelClient(BasePluginClient): ) for resp in response: - voices = [] + voices: list[TTSModelVoice] = [] for voice in resp.voices: voices.append({"name": voice.name, "value": voice.value}) diff --git a/api/core/plugin/impl/model_runtime.py b/api/core/plugin/impl/model_runtime.py index 454bd38958a..c0fe1fe3d50 100644 --- a/api/core/plugin/impl/model_runtime.py +++ b/api/core/plugin/impl/model_runtime.py @@ -31,6 +31,7 @@ from graphon.model_runtime.entities.rerank_entities import MultimodalRerankInput from graphon.model_runtime.entities.text_embedding_entities import EmbeddingInputType, EmbeddingResult from graphon.model_runtime.model_providers.base.large_language_model import normalize_non_stream_runtime_result from graphon.model_runtime.protocols.runtime import ModelRuntime +from graphon.model_runtime.protocols.tts_runtime import TTSModelVoice from models.provider_ids import ModelProviderID logger = logging.getLogger(__name__) @@ -658,7 +659,7 @@ class PluginModelRuntime(ModelRuntime): model: str, credentials: dict[str, Any], language: str | None, - ) -> Any: + ) -> list[TTSModelVoice]: plugin_id, provider_name = self._split_provider(provider) return self.client.get_tts_model_voices( tenant_id=self.tenant_id, diff --git a/api/core/prompt/advanced_prompt_transform.py b/api/core/prompt/advanced_prompt_transform.py index e7c88811fe5..5f9f2384568 100644 --- a/api/core/prompt/advanced_prompt_transform.py +++ b/api/core/prompt/advanced_prompt_transform.py @@ -19,6 +19,7 @@ from graphon.model_runtime.entities import ( ) from graphon.model_runtime.entities.message_entities import ImagePromptMessageContent, PromptMessageContentUnionTypes from graphon.runtime import VariablePool +from graphon.variables.template_resolution import convert_template class AdvancedPromptTransform(PromptTransform): @@ -171,7 +172,7 @@ class AdvancedPromptTransform(PromptTransform): if k.startswith("#"): vp.add(k[1:-1].split("."), v) raw_prompt = raw_prompt.replace("{{#context#}}", context or "") - prompt = vp.convert_template(raw_prompt).text + prompt = convert_template(vp, raw_prompt).text else: parser = PromptTemplateParser(template=raw_prompt, with_variable_tmpl=self.with_variable_tmpl) prompt_inputs: Mapping[str, str] = {k: inputs[k] for k in parser.variable_keys if k in inputs} diff --git a/api/core/repositories/celery_workflow_node_execution_repository.py b/api/core/repositories/celery_workflow_node_execution_repository.py index 697429eccdc..c393f6afd51 100644 --- a/api/core/repositories/celery_workflow_node_execution_repository.py +++ b/api/core/repositories/celery_workflow_node_execution_repository.py @@ -39,7 +39,7 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): Key features: - Asynchronous save operations using Celery tasks - - In-memory cache for immediate reads + - In-memory cache for immediate reads with database backfill across Celery tasks - Support for multi-tenancy through tenant/app filtering - Automatic retry and error handling through Celery """ @@ -52,6 +52,7 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): _creator_user_role: CreatorUserRole _execution_cache: dict[str, WorkflowNodeExecution] _workflow_execution_mapping: dict[str, list[str]] + _database_loaded_workflow_executions: set[str] _sql_repository: SQLAlchemyWorkflowNodeExecutionRepository def __init__( @@ -102,8 +103,9 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): # Cache for mapping workflow_execution_ids to execution IDs for efficient retrieval self._workflow_execution_mapping = {} + self._database_loaded_workflow_executions = set() self._sql_repository = SQLAlchemyWorkflowNodeExecutionRepository( - session_factory=session_factory, + session_factory=self._session_factory, tenant_id=tenant_id, user=user, app_id=app_id, @@ -178,7 +180,7 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): order_config: OrderConfig | None = None, ) -> Sequence[WorkflowNodeExecution]: """ - Retrieve all workflow node executions for a workflow execution from cache. + Retrieve workflow node executions from cache after loading persisted history once. Args: workflow_execution_id: The workflow execution identifier @@ -188,6 +190,25 @@ class CeleryWorkflowNodeExecutionRepository(WorkflowNodeExecutionRepository): A sequence of WorkflowNodeExecution instances """ try: + if workflow_execution_id not in self._database_loaded_workflow_executions: + try: + persisted_executions = self._sql_repository.get_by_workflow_execution( + workflow_execution_id, + order_config, + ) + except Exception: + logger.exception( + "Failed to load persisted workflow node executions for execution %s", + workflow_execution_id, + ) + else: + execution_ids = self._workflow_execution_mapping.setdefault(workflow_execution_id, []) + for execution in persisted_executions: + self._execution_cache.setdefault(execution.id, execution) + if execution.id not in execution_ids: + execution_ids.append(execution.id) + self._database_loaded_workflow_executions.add(workflow_execution_id) + # Get execution IDs for this workflow execution from cache execution_ids = self._workflow_execution_mapping.get(workflow_execution_id, []) diff --git a/api/core/repositories/human_input_repository.py b/api/core/repositories/human_input_repository.py index ac3aa46fc21..b144d1f88d4 100644 --- a/api/core/repositories/human_input_repository.py +++ b/api/core/repositories/human_input_repository.py @@ -67,6 +67,7 @@ class FormCreateParams: # workflow_execution_id for chatflow runs; set alone (workflow_execution_id None) # for Agent v2 chat ask_human forms, which have no workflow run. conversation_id: str | None = None + form_id: str | None = None class HumanInputFormRecipientEntity(Protocol): @@ -110,7 +111,7 @@ class HumanInputFormEntity(Protocol): class HumanInputFormRepository(Protocol): - def get_form(self, node_id: str) -> HumanInputFormEntity | None: ... + def get_form(self, node_id: str, *, form_id: str | None = None) -> HumanInputFormEntity | None: ... def create_form(self, params: FormCreateParams) -> HumanInputFormEntity: ... @@ -460,8 +461,7 @@ class HumanInputFormRepositoryImpl: raise ValueError("a runtime human input form requires a workflow_execution_id or conversation_id") with session_factory.create_session() as session, session.begin(): - # Generate unique form ID - form_id = str(uuidv7()) + form_id = params.form_id or str(uuidv7()) start_time = naive_utc_now() node_expiration = form_config.expiration_time(start_time) form_definition = FormDefinition( @@ -546,7 +546,7 @@ class HumanInputFormRepositoryImpl: return _HumanInputFormEntityImpl(form_model=form_model, recipient_models=recipient_models) - def get_form(self, node_id: str) -> HumanInputFormEntity | None: + def get_form(self, node_id: str, *, form_id: str | None = None) -> HumanInputFormEntity | None: if self._workflow_execution_id is None: raise ValueError("workflow_execution_id is required to load runtime human input forms") @@ -555,6 +555,8 @@ class HumanInputFormRepositoryImpl: HumanInputForm.node_id == node_id, HumanInputForm.tenant_id == self._tenant_id, ) + if form_id is not None: + form_query = form_query.where(HumanInputForm.id == form_id) with session_factory.create_session() as session: form_model: HumanInputForm | None = session.scalars(form_query).first() if form_model is None: diff --git a/api/core/tools/tool_manager.py b/api/core/tools/tool_manager.py index fc85f20bdd4..0f33b104914 100644 --- a/api/core/tools/tool_manager.py +++ b/api/core/tools/tool_manager.py @@ -55,6 +55,7 @@ from core.tools.workflow_as_tool.provider import WorkflowToolProviderController from core.tools.workflow_as_tool.tool import WorkflowTool from extensions.ext_database import db from graphon.runtime import VariablePool +from graphon.variables.template_resolution import convert_template from models.provider_ids import ToolProviderID from models.tools import ApiToolProvider, BuiltinToolProvider, WorkflowToolProvider from services.tools.mcp_tools_manage_service import MCPToolManageService @@ -1113,7 +1114,7 @@ class ToolManager: elif tool_input.type == "constant": parameter_value = tool_input.value elif tool_input.type == "mixed": - segment_group = variable_pool.convert_template(str(tool_input.value)) + segment_group = convert_template(variable_pool, str(tool_input.value)) parameter_value = segment_group.text else: raise ToolParameterError(f"Unknown tool input type '{tool_input.type}'") diff --git a/api/core/workflow/human_input_adapter.py b/api/core/workflow/human_input_adapter.py index 0865365ea68..52021f9f49b 100644 --- a/api/core/workflow/human_input_adapter.py +++ b/api/core/workflow/human_input_adapter.py @@ -21,6 +21,7 @@ from graphon.enums import BuiltinNodeTypes from graphon.nodes.base.variable_template_parser import VariableTemplateParser from graphon.runtime import VariablePool from graphon.variables.consts import SELECTORS_LENGTH +from graphon.variables.template_resolution import convert_template class DeliveryMethodType(enum.StrEnum): @@ -116,7 +117,7 @@ class EmailDeliveryConfig(BaseModel): templated_body = cls.replace_url_placeholder(body, url) if variable_pool is None: return templated_body - return variable_pool.convert_template(templated_body).text + return convert_template(variable_pool, templated_body).text @classmethod def render_markdown_body(cls, body: str) -> str: diff --git a/api/core/workflow/node_factory.py b/api/core/workflow/node_factory.py index 5b55af4b36b..4ad5709b776 100644 --- a/api/core/workflow/node_factory.py +++ b/api/core/workflow/node_factory.py @@ -361,6 +361,12 @@ class DifyNodeFactory(NodeFactory): self._agent_runtime_support = AgentRuntimeSupport() self._agent_message_transformer = AgentMessageTransformer() + def with_runtime_state(self, graph_runtime_state: "GraphRuntimeState") -> "DifyNodeFactory": + return DifyNodeFactory( + graph_init_params=self.graph_init_params, + graph_runtime_state=graph_runtime_state, + ) + @staticmethod def _resolve_dify_context(run_context: Mapping[str, Any]) -> DifyRunContext: raw_ctx = run_context.get(DIFY_RUN_CONTEXT_KEY) @@ -394,6 +400,7 @@ class DifyNodeFactory(NodeFactory): # stay explicit and constructors receive the concrete typed payload. resolved_node_data = self._validate_resolved_node_data(node_class, node_data) node_type = node_data.type + node: Node | None = None node_init_kwargs_factories: Mapping[NodeType, Callable[[], dict[str, object]]] = { BuiltinNodeTypes.CODE: lambda: { "code_executor": self._code_executor, @@ -412,7 +419,8 @@ class DifyNodeFactory(NodeFactory): }, BuiltinNodeTypes.HUMAN_INPUT: lambda: { "hitl_callback": self._build_human_input_callback( - node_data=DifyHumanInputNodeData.model_validate(adapted_node_config["data"]) + node_data=DifyHumanInputNodeData.model_validate(adapted_node_config["data"]), + execution_id_getter=lambda: node.execution_id if node is not None else None, ), }, BuiltinNodeTypes.LLM: lambda: self._build_llm_compatible_node_init_kwargs( @@ -457,13 +465,14 @@ class DifyNodeFactory(NodeFactory): } node_init_kwargs = node_init_kwargs_factories.get(node_type, lambda: {})() constructor_node_data = resolved_node_data.model_dump(mode="python", by_alias=True) - return node_class( + node = node_class( node_id=node_id, data=constructor_node_data, graph_init_params=self.graph_init_params, graph_runtime_state=self.graph_runtime_state, **node_init_kwargs, ) + return node @staticmethod def _validate_resolved_node_data(node_class: type[Node], node_data: BaseNodeData) -> BaseNodeData: @@ -525,6 +534,7 @@ class DifyNodeFactory(NodeFactory): self, *, node_data: DifyHumanInputNodeData, + execution_id_getter: Callable[[], str | None], ) -> DifyHITLCallback: return DifyHITLCallback( form_repository=self._human_input_runtime.build_form_repository(), @@ -533,6 +543,7 @@ class DifyNodeFactory(NodeFactory): delivery_methods=self._human_input_runtime._resolve_delivery_methods(node_data=node_data), display_in_ui=self._human_input_runtime._display_in_ui(node_data=node_data), file_reference_factory=self._file_reference_factory, + execution_id_getter=execution_id_getter, ) def _build_llm_compatible_node_init_kwargs( diff --git a/api/core/workflow/nodes/agent/runtime_support.py b/api/core/workflow/nodes/agent/runtime_support.py index 9a36e87e015..6ffccd70777 100644 --- a/api/core/workflow/nodes/agent/runtime_support.py +++ b/api/core/workflow/nodes/agent/runtime_support.py @@ -21,6 +21,7 @@ from core.workflow.system_variables import SystemVariableKey, get_system_text from extensions.ext_database import db from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelType from graphon.runtime import VariablePool +from graphon.variables.template_resolution import convert_template from models.model import Conversation from .entities import AgentNodeData, AgentOldVersionModelFeatures, ParamsAutoGenerated @@ -67,7 +68,7 @@ class AgentRuntimeSupport: except TypeError: parameter_value = str(agent_input.value) - segment_group = variable_pool.convert_template(parameter_value) + segment_group = convert_template(variable_pool, parameter_value) parameter_value = segment_group.log if for_log else segment_group.text try: if not isinstance(agent_input.value, str): diff --git a/api/core/workflow/nodes/human_input/callback.py b/api/core/workflow/nodes/human_input/callback.py index 46c6f3ceeee..8095078b455 100644 --- a/api/core/workflow/nodes/human_input/callback.py +++ b/api/core/workflow/nodes/human_input/callback.py @@ -2,7 +2,7 @@ from __future__ import annotations import json import logging -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from datetime import datetime, timedelta from typing import Any @@ -11,10 +11,10 @@ from core.repositories.human_input_repository import FormCreateParams, HumanInpu from core.workflow.human_input_adapter import DeliveryChannelConfig from core.workflow.node_runtime import DifyFileReferenceFactory from graphon.nodes.human_input.entities import Completed, Expired, HITLContext, HITLDecision, PauseRequested -from graphon.runtime import VariablePool from graphon.runtime.graph_runtime_state_protocol import ReadOnlyVariablePool from graphon.variables.factory import build_segment from graphon.variables.segments import Segment +from graphon.variables.template_resolution import convert_template from libs.datetime_utils import ensure_naive_utc, naive_utc_now from .entities import ( @@ -31,24 +31,13 @@ from .session_binding import default_session_binding logger = logging.getLogger(__name__) -def _require_template_variable_pool(pool: ReadOnlyVariablePool) -> VariablePool: - """Return the concrete graphon pool required for template expansion.""" - if isinstance(pool, VariablePool): - return pool - - msg = "human input rendering requires graphon.runtime.VariablePool for template expansion" - raise TypeError(msg) - - def render_form_content_before_submission( node_data: HumanInputNodeData, *, variable_pool: ReadOnlyVariablePool, ) -> str: """Process form content by substituting runtime variables before pause.""" - # NOTE(QuantumGhost): This is not ideal, we should expose - # VariablePool method in Graphon. - rendered_form_content = _require_template_variable_pool(variable_pool).convert_template(node_data.form_content) + rendered_form_content = convert_template(variable_pool, node_data.form_content) return rendered_form_content.markdown @@ -91,6 +80,7 @@ class DifyHITLCallback: delivery_methods: Sequence[DeliveryChannelConfig] = (), display_in_ui: bool = False, file_reference_factory: DifyFileReferenceFactory | None = None, + execution_id_getter: Callable[[], str | None] | None = None, ) -> None: self._form_repository = form_repository self._session_binding = default_session_binding @@ -100,11 +90,17 @@ class DifyHITLCallback: self._delivery_methods = tuple(delivery_methods) self._display_in_ui = display_in_ui self._file_reference_factory = file_reference_factory + self._execution_id_getter = execution_id_getter def __call__(self, ctx: HITLContext) -> HITLDecision: - form = self._form_repository.get_form(ctx.node_id) + form_id = self._execution_id_getter() if self._execution_id_getter is not None else None + form = ( + self._form_repository.get_form(ctx.node_id, form_id=form_id) + if form_id is not None + else self._form_repository.get_form(ctx.node_id) + ) if form is None: - created = self._create_form(ctx) + created = self._create_form(ctx, form_id=form_id) return PauseRequested(session_id=self._session_binding.issue_session_id_for_form(form_id=created.id)) status = self._normalize_status(form.status) @@ -163,7 +159,7 @@ class DifyHITLCallback: outputs=outputs, ) - def _create_form(self, ctx: HITLContext) -> HumanInputFormEntity: + def _create_form(self, ctx: HITLContext, *, form_id: str | None = None) -> HumanInputFormEntity: params = FormCreateParams( workflow_execution_id=self._workflow_execution_id or ctx.workflow_execution_id, conversation_id=self._conversation_id, @@ -181,6 +177,7 @@ class DifyHITLCallback: variable_pool=ctx.variable_pool, ) ), + form_id=form_id, ) return self._form_repository.create_form(params) diff --git a/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py b/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py index 11082d53fa2..b4e975dcfa6 100644 --- a/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py +++ b/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py @@ -25,7 +25,6 @@ from graphon.enums import ( from graphon.model_runtime.entities.llm_entities import LLMUsage from graphon.model_runtime.utils.encoders import jsonable_encoder from graphon.node_events import NodeRunResult -from graphon.nodes.base import LLMUsageTrackingMixin from graphon.nodes.base.node import Node from graphon.variables import ( ArrayFileSegment, @@ -33,6 +32,7 @@ from graphon.variables import ( StringSegment, ) from graphon.variables.segments import ArrayObjectSegment +from graphon.variables.template_resolution import convert_template from .entities import ( Condition, @@ -64,7 +64,7 @@ def _normalize_metadata_filter_sequence_item(value: object) -> str: return value if isinstance(value, str) else str(value) -class KnowledgeRetrievalNode(LLMUsageTrackingMixin, Node[KnowledgeRetrievalNodeData]): +class KnowledgeRetrievalNode(Node[KnowledgeRetrievalNodeData]): node_type = BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL # Instance attributes specific to LLMNode. @@ -309,7 +309,7 @@ class KnowledgeRetrievalNode(LLMUsageTrackingMixin, Node[KnowledgeRetrievalNodeD resolved_value: str | Sequence[str] | int | float | None match value: case str(): - segment_group = variable_pool.convert_template(value) + segment_group = convert_template(variable_pool, value) if len(segment_group.value) == 1: resolved_value = _normalize_metadata_filter_scalar(segment_group.value[0].to_object()) else: @@ -317,7 +317,7 @@ class KnowledgeRetrievalNode(LLMUsageTrackingMixin, Node[KnowledgeRetrievalNodeD case _ if isinstance(value, Sequence) and all(isinstance(v, str) for v in value): resolved_values: list[str] = [] for v in value: - segment_group = variable_pool.convert_template(v) + segment_group = convert_template(variable_pool, v) if len(segment_group.value) == 1: resolved_values.append( _normalize_metadata_filter_sequence_item(segment_group.value[0].to_object()) diff --git a/api/core/workflow/workflow_entry.py b/api/core/workflow/workflow_entry.py index fb12922ed7f..0377c34eb96 100644 --- a/api/core/workflow/workflow_entry.py +++ b/api/core/workflow/workflow_entry.py @@ -2,6 +2,7 @@ import logging import time from collections.abc import Generator, Mapping, Sequence from typing import Any, TypedDict +from uuid import uuid4 from configs import dify_config from context import capture_current_context @@ -26,7 +27,6 @@ from core.workflow.variable_pool_initializer import add_node_inputs_to_pool, add from core.workflow.variable_prefixes import ENVIRONMENT_VARIABLE_NODE_ID from extensions.otel.runtime import is_instrument_flag_enabled from factories import file_factory -from graphon.entities import GraphInitParams from graphon.entities.graph_config import NodeConfigDictAdapter from graphon.errors import WorkflowNodeRunFailedError from graphon.file import File @@ -38,7 +38,8 @@ from graphon.graph_engine.layers import DebugLoggingLayer, ExecutionLimitsLayer from graphon.graph_events import GraphEngineEvent, GraphNodeEventBase, GraphRunFailedEvent from graphon.nodes import BuiltinNodeTypes from graphon.nodes.base.node import Node -from graphon.runtime import ChildGraphNotFoundError, GraphRuntimeState, VariablePool +from graphon.nodes.container_effects import ContainerAwaitRequest +from graphon.runtime import GraphRuntimeState, VariablePool from graphon.variable_loader import DUMMY_VARIABLE_LOADER, VariableLoader, load_into_variable_pool from models.workflow import Workflow @@ -69,77 +70,6 @@ def iter_dify_graph_engine_events( ) -class _WorkflowChildEngineBuilder: - tenant_id: str - - def __init__(self, *, tenant_id: str) -> None: - self.tenant_id = tenant_id - - @staticmethod - def _has_node_id(graph_config: Mapping[str, Any], node_id: str) -> bool | None: - """ - Return whether `graph_config["nodes"]` contains the given node id. - - Returns `None` when the nodes payload shape is unexpected, so graph-level - validation can surface the original configuration error. - """ - nodes = graph_config.get("nodes") - if not isinstance(nodes, list): - return None - - for node in nodes: - if not isinstance(node, Mapping): - return None - current_id = node.get("id") - if isinstance(current_id, str) and current_id == node_id: - return True - return False - - def build_child_engine( - self, - *, - workflow_id: str, - graph_init_params: GraphInitParams, - parent_graph_runtime_state: GraphRuntimeState, - root_node_id: str, - variable_pool: VariablePool | None = None, - ) -> GraphEngine: - """Build a child engine with a fresh runtime state and only child-safe layers.""" - child_graph_runtime_state = GraphRuntimeState( - variable_pool=variable_pool if variable_pool is not None else parent_graph_runtime_state.variable_pool, - start_at=time.perf_counter(), - execution_context=parent_graph_runtime_state.execution_context, - ) - node_factory = DifyNodeFactory( - graph_init_params=graph_init_params, - graph_runtime_state=child_graph_runtime_state, - ) - - graph_config = graph_init_params.graph_config - has_root_node = self._has_node_id(graph_config=graph_config, node_id=root_node_id) - if has_root_node is False: - raise ChildGraphNotFoundError(f"child graph root node '{root_node_id}' not found") - - child_graph = Graph.init( - graph_config=graph_config, - node_factory=node_factory, - root_node_id=root_node_id, - ) - - command_channel = InMemoryChannel() - config = GraphEngineConfig() - child_engine = GraphEngine( - workflow_id=workflow_id, - graph=child_graph, - graph_runtime_state=child_graph_runtime_state, - command_channel=command_channel, - config=config, - child_engine_builder=self, - ) - child_engine.layer(LLMQuotaLayer(tenant_id=self.tenant_id)) - return child_engine - - class _NodeConfigDict(TypedDict): id: str width: int @@ -208,8 +138,8 @@ class WorkflowEntry: self.command_channel = command_channel self._response_stream_filter = response_stream_filter or ResponseStreamFilter() execution_context = capture_current_context() - graph_runtime_state.execution_context = execution_context - self._child_engine_builder = _WorkflowChildEngineBuilder(tenant_id=tenant_id) + # ponytail: Graphon snapshots omit process-local context; use a public rebind API when Graphon exposes one. + graph_runtime_state._execution_context = execution_context self.graph_engine = GraphEngine( workflow_id=workflow_id, graph=graph, @@ -221,7 +151,6 @@ class WorkflowEntry: scale_up_threshold=dify_config.GRAPH_ENGINE_SCALE_UP_THRESHOLD, scale_down_idle_time=dify_config.GRAPH_ENGINE_SCALE_DOWN_IDLE_TIME, ), - child_engine_builder=self._child_engine_builder, ) # Add debug logging layer when in debug mode @@ -271,7 +200,7 @@ class WorkflowEntry: user_inputs: Mapping[str, Any], variable_pool: VariablePool, variable_loader: VariableLoader = DUMMY_VARIABLE_LOADER, - ) -> tuple[Node, Generator[GraphNodeEventBase, None, None]]: + ) -> tuple[Node, Generator[GraphNodeEventBase | ContainerAwaitRequest, None, None]]: """ Single step run workflow node :param workflow: Workflow instance @@ -285,6 +214,8 @@ class WorkflowEntry: # Get node type node_type = node_config_data.type + if node_type in {BuiltinNodeTypes.LOOP, BuiltinNodeTypes.ITERATION}: + raise ValueError("Loop and Iteration nodes must use their engine-backed debug endpoints") node_version = str(node_config_data.version) node_cls = resolve_workflow_node_class(node_type=node_type, node_version=node_version) @@ -419,7 +350,7 @@ class WorkflowEntry: @classmethod def run_free_node( cls, node_data: dict[str, Any], node_id: str, tenant_id: str, user_id: str, user_inputs: dict[str, Any] - ) -> tuple[Node, Generator[GraphNodeEventBase, None, None]]: + ) -> tuple[Node, Generator[GraphNodeEventBase | ContainerAwaitRequest, None, None]]: """ Run free node @@ -613,14 +544,14 @@ class WorkflowEntry: variable_pool.add([variable_node_id] + variable_key_list, input_value) @staticmethod - def _traced_node_run(node: Node) -> Generator[GraphNodeEventBase, None, None]: + def _traced_node_run(node: Node) -> Generator[GraphNodeEventBase | ContainerAwaitRequest, None, None]: """ Wraps a node's run method with OpenTelemetry tracing and returns a generator. """ # Wrap node.run() with ObservabilityLayer hooks to produce node-level spans layer = ObservabilityLayer() layer.on_graph_start() - node.ensure_execution_id() + node.bind_execution_id(str(uuid4())) def _gen(): error: Exception | None = None diff --git a/api/libs/broadcast_channel/redis/_subscription.py b/api/libs/broadcast_channel/redis/_subscription.py index 01a9e668bcc..c155b9e1973 100644 --- a/api/libs/broadcast_channel/redis/_subscription.py +++ b/api/libs/broadcast_channel/redis/_subscription.py @@ -117,9 +117,12 @@ class RedisSubscriptionBase(Subscription): ) continue - self._enqueue_message(payload_bytes) if payload_bytes == SIG_CLOSE: - break + # Close signals are broadcast to every subscriber on the topic. + # The closing subscription is already handled by the _closed check above. + continue + + self._enqueue_message(payload_bytes) _logger.debug("%s listener thread stopped for channel %s", self._get_subscription_type().title(), self._topic) try: diff --git a/api/libs/broadcast_channel/redis/streams_channel.py b/api/libs/broadcast_channel/redis/streams_channel.py index b3385b05388..c86d77ada6b 100644 --- a/api/libs/broadcast_channel/redis/streams_channel.py +++ b/api/libs/broadcast_channel/redis/streams_channel.py @@ -128,11 +128,17 @@ class _StreamsSubscription(Subscription): data_bytes = data.encode() case bytes() | bytearray(): data_bytes = bytes(data) - if data_bytes is not None: - if data_bytes == SIG_CLOSE: - break - self._queue.put_nowait(data_bytes) last_id = entry_id + if data_bytes is None: + continue + if data_bytes == SIG_CLOSE: + # Close signals share the stream with normal events. Ignore signals + # emitted by another subscription while this one is still open. + with self._lock: + if self._closed: + break + continue + self._queue.put_nowait(data_bytes) finally: self._queue.put_nowait(self._SENTINEL) with self._lock: diff --git a/api/pyproject.toml b/api/pyproject.toml index 27070217e63..275739ed64e 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -45,7 +45,7 @@ dependencies = [ "zstandard==0.25.0", # Emerging: newer and fast-moving, use compatible pins "fastopenapi[flask]==0.7.0", - "graphon==0.6.0", + "graphon==0.7.0", "httpx-sse==0.4.3", "json-repair==0.60.1", ] diff --git a/api/services/legacy_model_type_migration.py b/api/services/legacy_model_type_migration.py index 1465fc0912f..2034246991f 100644 --- a/api/services/legacy_model_type_migration.py +++ b/api/services/legacy_model_type_migration.py @@ -38,7 +38,7 @@ from typing import Protocol, cast, override import sqlalchemy as sa from sqlalchemy.exc import OperationalError -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, defer from sqlalchemy.sql import select from core.helper.model_provider_cache import ProviderCredentialsCache, ProviderCredentialsCacheType @@ -747,6 +747,7 @@ class Migration: with _session_factory(self._engine) as session: stmt = ( select(ProviderModel, raw_model_type) + .options(defer(ProviderModel.model_type)) .where( ProviderModel.tenant_id == self._tenant_id, sa.type_coerce(ProviderModel.model_type, sa.String()).in_(self._selected_legacy_values()), @@ -787,6 +788,7 @@ class Migration: raw_model_type = sa.type_coerce(ProviderModel.model_type, sa.String()).label(_RAW_MODEL_TYPE_COLUMN) stmt = ( select(ProviderModel, raw_model_type) + .options(defer(ProviderModel.model_type)) .where( ProviderModel.tenant_id == candidate.row.tenant_id, ProviderModel.provider_name == candidate.row.provider_name, @@ -983,6 +985,7 @@ class Migration: with _session_factory(self._engine) as session: stmt = ( select(TenantDefaultModel, raw_model_type) + .options(defer(TenantDefaultModel.model_type)) .where( TenantDefaultModel.tenant_id == self._tenant_id, sa.type_coerce(TenantDefaultModel.model_type, sa.String()).in_(self._selected_legacy_values()), @@ -1023,6 +1026,7 @@ class Migration: raw_model_type = sa.type_coerce(TenantDefaultModel.model_type, sa.String()).label(_RAW_MODEL_TYPE_COLUMN) stmt = ( select(TenantDefaultModel, raw_model_type) + .options(defer(TenantDefaultModel.model_type)) .where( TenantDefaultModel.tenant_id == candidate.row.tenant_id, sa.type_coerce(TenantDefaultModel.model_type, sa.String()).in_( @@ -1193,6 +1197,7 @@ class Migration: with _session_factory(self._engine) as session: stmt = ( select(ProviderModelSetting, raw_model_type) + .options(defer(ProviderModelSetting.model_type)) .where( ProviderModelSetting.tenant_id == self._tenant_id, sa.type_coerce(ProviderModelSetting.model_type, sa.String()).in_(self._selected_legacy_values()), @@ -1233,6 +1238,7 @@ class Migration: raw_model_type = sa.type_coerce(ProviderModelSetting.model_type, sa.String()).label(_RAW_MODEL_TYPE_COLUMN) stmt = ( select(ProviderModelSetting, raw_model_type) + .options(defer(ProviderModelSetting.model_type)) .where( ProviderModelSetting.tenant_id == candidate.row.tenant_id, ProviderModelSetting.provider_name == candidate.row.provider_name, @@ -1437,6 +1443,7 @@ class Migration: with _session_factory(self._engine) as session: stmt = ( select(LoadBalancingModelConfig, raw_model_type) + .options(defer(LoadBalancingModelConfig.model_type)) .where( LoadBalancingModelConfig.tenant_id == self._tenant_id, LoadBalancingModelConfig.name == "__inherit__", @@ -1485,6 +1492,7 @@ class Migration: raw_model_type = sa.type_coerce(LoadBalancingModelConfig.model_type, sa.String()).label(_RAW_MODEL_TYPE_COLUMN) stmt = ( select(LoadBalancingModelConfig, raw_model_type) + .options(defer(LoadBalancingModelConfig.model_type)) .where( LoadBalancingModelConfig.tenant_id == candidate.row.tenant_id, LoadBalancingModelConfig.provider_name == candidate.row.provider_name, @@ -1617,6 +1625,7 @@ class Migration: with _session_factory(self._engine) as session: stmt = ( select(LoadBalancingModelConfig, raw_model_type) + .options(defer(LoadBalancingModelConfig.model_type)) .where( LoadBalancingModelConfig.tenant_id == self._tenant_id, sa.type_coerce(LoadBalancingModelConfig.model_type, sa.String()).in_( @@ -1660,9 +1669,13 @@ class Migration: lock_rows: bool, ) -> _RowWithRawModelType[LoadBalancingModelConfig] | None: raw_model_type = sa.type_coerce(LoadBalancingModelConfig.model_type, sa.String()).label(_RAW_MODEL_TYPE_COLUMN) - stmt = select(LoadBalancingModelConfig, raw_model_type).where( - LoadBalancingModelConfig.id == candidate.row.id, - LoadBalancingModelConfig.tenant_id == self._tenant_id, + stmt = ( + select(LoadBalancingModelConfig, raw_model_type) + .options(defer(LoadBalancingModelConfig.model_type)) + .where( + LoadBalancingModelConfig.id == candidate.row.id, + LoadBalancingModelConfig.tenant_id == self._tenant_id, + ) ) if lock_rows: stmt = stmt.with_for_update() @@ -1818,6 +1831,7 @@ class Migration: with _session_factory(self._engine) as session: stmt = ( select(ProviderModelCredential, raw_model_type) + .options(defer(ProviderModelCredential.model_type)) .where( ProviderModelCredential.tenant_id == self._tenant_id, sa.type_coerce(ProviderModelCredential.model_type, sa.String()).in_(self._selected_legacy_values()), @@ -1858,6 +1872,7 @@ class Migration: raw_model_type = sa.type_coerce(ProviderModelCredential.model_type, sa.String()).label(_RAW_MODEL_TYPE_COLUMN) stmt = ( select(ProviderModelCredential, raw_model_type) + .options(defer(ProviderModelCredential.model_type)) .where( ProviderModelCredential.tenant_id == candidate.row.tenant_id, ProviderModelCredential.provider_name == candidate.row.provider_name, @@ -2147,6 +2162,7 @@ class Migration: stmt = ( select(ProviderModel) + .options(defer(ProviderModel.model_type)) .where( ProviderModel.tenant_id == self._tenant_id, ProviderModel.credential_id.in_(loser_ids), @@ -2182,6 +2198,7 @@ class Migration: stmt = ( select(LoadBalancingModelConfig) + .options(defer(LoadBalancingModelConfig.model_type)) .where( LoadBalancingModelConfig.tenant_id == self._tenant_id, LoadBalancingModelConfig.credential_id.in_(loser_ids), @@ -2295,9 +2312,14 @@ class Migration: def _row_to_dict(self, row: TypeBase, *, raw_model_type: str | None = None) -> dict[str, object]: mapper = sa.inspect(row).mapper - row_dict = {column.key: row.__dict__[column.key] for column in mapper.column_attrs} - if raw_model_type is not None and "model_type" in row_dict: - row_dict["model_type"] = raw_model_type + row_dict = { + column.key: ( + raw_model_type + if column.key == "model_type" and raw_model_type is not None + else row.__dict__[column.key] + ) + for column in mapper.column_attrs + } return _normalize_log_mapping(row_dict) def _log_row_deleted[T: TypeBase]( diff --git a/api/services/rag_pipeline/rag_pipeline.py b/api/services/rag_pipeline/rag_pipeline.py index 993f6d79492..54f0c16b594 100644 --- a/api/services/rag_pipeline/rag_pipeline.py +++ b/api/services/rag_pipeline/rag_pipeline.py @@ -49,6 +49,7 @@ from graphon.errors import WorkflowNodeRunFailedError from graphon.graph_events import GraphNodeEventBase, NodeRunFailedEvent, NodeRunSucceededEvent from graphon.node_events import NodeRunResult from graphon.nodes.base.node import Node +from graphon.nodes.container_effects import ContainerAwaitRequest from graphon.nodes.http_request import HTTP_REQUEST_CONFIG_FILTER_KEY, build_http_request_config from graphon.runtime import VariablePool from graphon.variables.variables import Variable, VariableBase @@ -910,7 +911,10 @@ class RagPipelineService: def _handle_node_run_result( self, - getter: Callable[[], tuple[Node, Generator[GraphNodeEventBase, None, None]]], + getter: Callable[ + [], + tuple[Node, Generator[GraphNodeEventBase | ContainerAwaitRequest, None, None]], + ], start_at: float, tenant_id: str, node_id: str, diff --git a/api/services/workflow_event_snapshot_service.py b/api/services/workflow_event_snapshot_service.py index 58f42badc66..5fb86206ec9 100644 --- a/api/services/workflow_event_snapshot_service.py +++ b/api/services/workflow_event_snapshot_service.py @@ -565,7 +565,6 @@ def _build_pause_event( variable_pool: ReadOnlyVariablePool | None = None if resumption_context is not None: state = GraphRuntimeState.from_snapshot(resumption_context.serialized_graph_runtime_state) - paused_nodes = state.get_paused_nodes() outputs = dict(WorkflowRuntimeTypeConverter().to_json_encodable(state.outputs or {})) variable_pool = state.variable_pool @@ -573,6 +572,9 @@ def _build_pause_event( pause_entity.get_pause_reasons(), variable_pool=variable_pool, ) + paused_nodes = list( + dict.fromkeys(reason.node_id for reason in resolved_pause_reasons if isinstance(reason, HumanInputRequired)) + ) reasons = [reason.model_dump(mode="json") for reason in resolved_pause_reasons] human_input_form_ids = [ form_id diff --git a/api/services/workflow_service.py b/api/services/workflow_service.py index b1f29c2ef7e..c2700fe8c8c 100644 --- a/api/services/workflow_service.py +++ b/api/services/workflow_service.py @@ -62,6 +62,7 @@ from graphon.graph_events import GraphNodeEventBase, NodeRunFailedEvent, NodeRun from graphon.node_events import NodeRunResult from graphon.nodes import BuiltinNodeTypes from graphon.nodes.base.node import Node +from graphon.nodes.container_effects import ContainerAwaitRequest from graphon.nodes.http_request import HTTP_REQUEST_CONFIG_FILTER_KEY, build_http_request_config from graphon.nodes.start.entities import StartNodeData from graphon.runtime import VariablePool @@ -1472,7 +1473,10 @@ class WorkflowService: def _handle_single_step_result( self, - invoke_node_fn: Callable[[], tuple[Node, Generator[GraphNodeEventBase, None, None]]], + invoke_node_fn: Callable[ + [], + tuple[Node, Generator[GraphNodeEventBase | ContainerAwaitRequest, None, None]], + ], start_at: float, node_id: str, ) -> WorkflowNodeExecution: @@ -1508,7 +1512,11 @@ class WorkflowService: return node_execution def _execute_node_safely( - self, invoke_node_fn: Callable[[], tuple[Node, Generator[GraphNodeEventBase, None, None]]] + self, + invoke_node_fn: Callable[ + [], + tuple[Node, Generator[GraphNodeEventBase | ContainerAwaitRequest, None, None]], + ], ) -> tuple[Node, NodeRunResult | None, bool, str | None]: """ Execute node safely and handle errors according to error strategy. diff --git a/api/tests/integration_tests/workflow/nodes/test_tool.py b/api/tests/integration_tests/workflow/nodes/test_tool.py index c109be9fae4..d3e248cc28c 100644 --- a/api/tests/integration_tests/workflow/nodes/test_tool.py +++ b/api/tests/integration_tests/workflow/nodes/test_tool.py @@ -70,6 +70,7 @@ def init_tool_node(config: dict): tool_file_manager=tool_file_manager, runtime=DifyToolNodeRuntime(init_params.run_context), ) + node.bind_execution_id(str(uuid.uuid4())) return node diff --git a/api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py b/api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py index 84f01ea52ee..2d5e3ebcca9 100644 --- a/api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py +++ b/api/tests/test_containers_integration_tests/core/app/layers/test_pause_state_persist_layer.py @@ -17,7 +17,6 @@ These tests use TestContainers to spin up real services for integration testing, providing more reliable and realistic test scenarios than mocks. """ -import json import uuid from time import time from unittest.mock import Mock @@ -237,12 +236,12 @@ class TestPauseStatePersistenceLayerTestContainers: # Create LLM usage llm_usage = LLMUsage.empty_usage() + llm_usage.total_tokens = total_tokens # Create graph runtime state graph_runtime_state = GraphRuntimeState( variable_pool=variable_pool, start_at=start_at, - total_tokens=total_tokens, llm_usage=llm_usage, outputs=outputs or {}, node_run_steps=node_run_steps, @@ -366,9 +365,6 @@ class TestPauseStatePersistenceLayerTestContainers: resumption_context = WorkflowResumptionContext.loads(storage_content) assert resumption_context.version == "1" assert resumption_context.serialized_graph_runtime_state == graph_runtime_state.dumps() - expected_state = json.loads(graph_runtime_state.dumps()) - actual_state = json.loads(resumption_context.serialized_graph_runtime_state) - assert actual_state == expected_state persisted_entity = resumption_context.get_generate_entity() assert isinstance(persisted_entity, WorkflowAppGenerateEntity) assert persisted_entity.workflow_execution_id == self.test_workflow_run_id @@ -414,13 +410,11 @@ class TestPauseStatePersistenceLayerTestContainers: state_bytes = pause_entity.get_state() resumption_context = WorkflowResumptionContext.loads(state_bytes.decode()) - retrieved_state = json.loads(resumption_context.serialized_graph_runtime_state) - expected_state = json.loads(graph_runtime_state.dumps()) + retrieved_state = GraphRuntimeState.from_snapshot(resumption_context.serialized_graph_runtime_state) - assert retrieved_state == expected_state - assert retrieved_state["outputs"] == complex_outputs - assert retrieved_state["total_tokens"] == 250 - assert retrieved_state["node_run_steps"] == 10 + assert retrieved_state.outputs == complex_outputs + assert retrieved_state.total_tokens == 250 + assert retrieved_state.node_run_steps == 10 assert resumption_context.get_generate_entity().workflow_execution_id == self.test_workflow_run_id def test_database_transaction_handling(self, db_session_with_containers: Session): diff --git a/api/tests/test_containers_integration_tests/models/test_types_enum_text.py b/api/tests/test_containers_integration_tests/models/test_types_enum_text.py index b325c97f7d1..cb114537e63 100644 --- a/api/tests/test_containers_integration_tests/models/test_types_enum_text.py +++ b/api/tests/test_containers_integration_tests/models/test_types_enum_text.py @@ -210,7 +210,7 @@ class TestEnumText: assert str(exc.value) == "'invalid' is not a valid _UserType" - def test_select_legacy_model_type_values(self, engine_with_containers: Engine): + def test_select_rejects_legacy_model_type_values(self, engine_with_containers: Engine): insertion_sql = """ INSERT INTO enum_text_legacy_model_type_test (id, model_type) VALUES (1, 'text-generation'), @@ -221,11 +221,9 @@ class TestEnumText: session.execute(sa.text(insertion_sql)) session.commit() - with Session(engine_with_containers) as session: - records = session.scalars(select(_LegacyModelTypeRecord).order_by(_LegacyModelTypeRecord.id)).all() + for record_id, legacy_value in enumerate(("text-generation", "embeddings", "reranking"), 1): + with pytest.raises(ValueError) as exc: + with Session(engine_with_containers) as session: + session.scalar(select(_LegacyModelTypeRecord).where(_LegacyModelTypeRecord.id == record_id)) - assert [record.model_type for record in records] == [ - ModelType.LLM, - ModelType.TEXT_EMBEDDING, - ModelType.RERANK, - ] + assert str(exc.value) == f"'{legacy_value}' is not a valid ModelType" diff --git a/api/tests/unit_tests/configs/test_dify_config.py b/api/tests/unit_tests/configs/test_dify_config.py index 3588f6beef1..f8e6e890eef 100644 --- a/api/tests/unit_tests/configs/test_dify_config.py +++ b/api/tests/unit_tests/configs/test_dify_config.py @@ -83,6 +83,7 @@ def test_dify_config(monkeypatch: pytest.MonkeyPatch): assert config.AGENT_SHELL_ENABLED is True assert config.SENTRY_TRACES_SAMPLE_RATE == 1.0 assert config.TEMPLATE_TRANSFORM_MAX_LENGTH == 400_000 + assert config.GRAPH_ENGINE_SCALE_UP_THRESHOLD == 0 # annotated field with custom configured value assert config.HTTP_REQUEST_MAX_READ_TIMEOUT == 300 diff --git a/api/tests/unit_tests/controllers/console/test_human_input_form.py b/api/tests/unit_tests/controllers/console/test_human_input_form.py index a9e847d7d7c..e40d85c331a 100644 --- a/api/tests/unit_tests/controllers/console/test_human_input_form.py +++ b/api/tests/unit_tests/controllers/console/test_human_input_form.py @@ -4,7 +4,7 @@ import json from datetime import UTC, datetime from inspect import unwrap from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import ANY, Mock import pytest from flask import Flask, Response @@ -18,6 +18,7 @@ from controllers.console.human_input_form import ( WorkflowResponseConverter, _jsonify_form_definition, ) +from core.workflow.human_input_policy import HumanInputSurface from models.account import AccountStatus from models.enums import CreatorUserRole from models.human_input import RecipientType @@ -344,3 +345,62 @@ def test_workflow_events_finished(app: Flask, monkeypatch: pytest.MonkeyPatch) - assert response.mimetype == "text/event-stream" assert "data" in response.get_data(as_text=True) + + +def test_workflow_events_snapshot_can_continue_across_pauses(app: Flask, monkeypatch: pytest.MonkeyPatch) -> None: + workflow_run = SimpleNamespace( + id="run-1", + created_by_role=CreatorUserRole.ACCOUNT, + created_by="user-1", + tenant_id="t1", + app_id="app-1", + finished_at=None, + ) + app_model = SimpleNamespace(mode=AppMode.WORKFLOW) + + class _RepoStub: + def get_workflow_run_by_id_and_tenant_id(self, **_kwargs): + return workflow_run + + workflow_generator = Mock() + workflow_generator.convert_to_event_stream.return_value = iter(["data: snapshot\n\n"]) + snapshot_builder = Mock(return_value=["snapshot-events"]) + + monkeypatch.setattr( + DifyAPIRepositoryFactory, + "create_api_workflow_run_repository", + lambda *_args, **_kwargs: _RepoStub(), + ) + monkeypatch.setattr( + "controllers.console.human_input_form._retrieve_app_for_workflow_run", + lambda *_args, **_kwargs: app_model, + ) + monkeypatch.setattr( + "controllers.console.human_input_form.WorkflowAppGenerator", + lambda: workflow_generator, + ) + monkeypatch.setattr( + "controllers.console.human_input_form.build_workflow_event_stream", + snapshot_builder, + ) + monkeypatch.setattr("controllers.console.human_input_form.db", SimpleNamespace(engine=object())) + + api = ConsoleWorkflowEventsApi() + handler = unwrap(api.get) + + with app.test_request_context( + "/console/api/workflow/run-1/events?include_state_snapshot=true&continue_on_pause=true", + method="GET", + ): + response = handler(api, "t1", SimpleNamespace(id="user-1"), workflow_run_id="run-1") + + assert response.get_data(as_text=True) == "data: snapshot\n\n" + snapshot_builder.assert_called_once_with( + app_mode=AppMode.WORKFLOW, + workflow_run=workflow_run, + tenant_id="t1", + app_id="app-1", + session_maker=ANY, + human_input_surface=HumanInputSurface.CONSOLE, + close_on_pause=False, + ) diff --git a/api/tests/unit_tests/controllers/service_api/app/test_hitl_service_api.py b/api/tests/unit_tests/controllers/service_api/app/test_hitl_service_api.py index c5072ec2e70..005d1ea7502 100644 --- a/api/tests/unit_tests/controllers/service_api/app/test_hitl_service_api.py +++ b/api/tests/unit_tests/controllers/service_api/app/test_hitl_service_api.py @@ -271,8 +271,7 @@ def _build_resumption_context(task_id: str) -> WorkflowResumptionContext: workflow_execution_id="run-1", ) runtime_state = GraphRuntimeState(variable_pool=VariablePool(), start_at=0.0) - runtime_state.register_paused_node("node-1") - runtime_state.outputs = {"result": "value"} + runtime_state.set_output("result", "value") wrapper = _WorkflowGenerateEntityWrapper(entity=generate_entity) return WorkflowResumptionContext( generate_entity=wrapper, diff --git a/api/tests/unit_tests/controllers/web/test_workflow_events.py b/api/tests/unit_tests/controllers/web/test_workflow_events.py index ab2ca7dde67..d9eb5fc7e73 100644 --- a/api/tests/unit_tests/controllers/web/test_workflow_events.py +++ b/api/tests/unit_tests/controllers/web/test_workflow_events.py @@ -3,7 +3,7 @@ from __future__ import annotations from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, Mock, patch import pytest from flask import Flask @@ -11,6 +11,7 @@ from flask import Flask from controllers.common.errors import NotFoundError from controllers.web.workflow_events import WorkflowEventsApi from models.enums import CreatorUserRole +from models.model import AppMode def _workflow_app() -> SimpleNamespace: @@ -125,3 +126,39 @@ class TestWorkflowEventsApi: response = WorkflowEventsApi().get(_workflow_app(), _end_user(), "run-1") assert response.mimetype == "text/event-stream" + + @patch("controllers.web.workflow_events.DifyAPIRepositoryFactory") + @patch("controllers.web.workflow_events.db") + def test_snapshot_stream_can_continue_across_pauses( + self, mock_db: MagicMock, mock_factory: MagicMock, app: Flask, monkeypatch: pytest.MonkeyPatch + ) -> None: + mock_db.engine = "engine" + run = SimpleNamespace( + id="run-1", + app_id="app-1", + created_by_role=CreatorUserRole.END_USER, + created_by="eu-1", + finished_at=None, + ) + mock_repo = MagicMock() + mock_repo.get_workflow_run_by_id_and_tenant_id.return_value = run + mock_factory.create_api_workflow_run_repository.return_value = mock_repo + + workflow_generator = Mock() + workflow_generator.convert_to_event_stream.return_value = iter(["data: snapshot\n\n"]) + snapshot_builder = Mock(return_value=["snapshot-events"]) + monkeypatch.setattr("controllers.web.workflow_events.WorkflowAppGenerator", lambda: workflow_generator) + monkeypatch.setattr("controllers.web.workflow_events.build_workflow_event_stream", snapshot_builder) + + with app.test_request_context("/workflow/run-1/events?include_state_snapshot=true&continue_on_pause=true"): + response = WorkflowEventsApi().get(_workflow_app(), _end_user(), "run-1") + + assert response.get_data(as_text=True) == "data: snapshot\n\n" + snapshot_builder.assert_called_once_with( + app_mode=AppMode.WORKFLOW, + workflow_run=run, + tenant_id="tenant-1", + app_id="app-1", + session_maker=ANY, + close_on_pause=False, + ) diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py index 004906a5de3..915bbdc556e 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_app_generator.py @@ -748,11 +748,13 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.preserve_flask_contexts", _fake_context) + workflow = SimpleNamespace(id="workflow-id", tenant_id="tenant", app_id="app") + class _Session: def __init__(self, *args, **kwargs): self.scalar = MagicMock( side_effect=[ - SimpleNamespace(id="workflow-id", tenant_id="tenant", app_id="app"), + workflow, SimpleNamespace(id="app"), ] ) @@ -772,6 +774,8 @@ class TestAdvancedChatAppGeneratorInternals: monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.Session", _Session) monkeypatch.setattr("core.app.apps.advanced_chat.app_generator.AdvancedChatAppRunner", _Runner) + restore_workflow_run_graph = MagicMock() + monkeypatch.setattr(generator, "_restore_workflow_run_graph", restore_workflow_run_graph) monkeypatch.setattr( "core.app.apps.advanced_chat.app_generator.db", SimpleNamespace(engine=object(), session=SimpleNamespace(close=lambda: None)), @@ -788,10 +792,12 @@ class TestAdvancedChatAppGeneratorInternals: workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), graph_engine_layers=(), - graph_runtime_state=None, + graph_runtime_state=SimpleNamespace(), ) queue_manager.publish_error.assert_not_called() + assert restore_workflow_run_graph.call_args.kwargs["workflow"] is workflow + assert restore_workflow_run_graph.call_args.kwargs["workflow_run_id"] == "run-id" def test_generate_worker_handles_validation_error(self, monkeypatch: pytest.MonkeyPatch): generator = AdvancedChatAppGenerator() diff --git a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py b/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py index 3633b75f217..f9995597279 100644 --- a/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py +++ b/api/tests/unit_tests/core/app/apps/advanced_chat/test_generate_task_pipeline_core.py @@ -56,6 +56,7 @@ from core.workflow.nodes.human_input.pause_reason import DifyHITLEventType from core.workflow.system_variables import build_system_variables from graphon.enums import BuiltinNodeTypes from graphon.file import FileTransferMethod, FileType +from graphon.model_runtime.entities.llm_entities import LLMUsage from graphon.runtime import GraphRuntimeState, VariablePool from libs.datetime_utils import naive_utc_now from models.enums import MessageStatus @@ -174,7 +175,7 @@ class TestAdvancedChatGenerateTaskPipeline: variables=build_system_variables(workflow_execution_id="run-id"), ), start_at=0.0, - total_tokens=7, + llm_usage=LLMUsage.empty_usage().model_copy(update={"total_tokens": 7}), node_run_steps=3, ) diff --git a/api/tests/unit_tests/core/app/apps/test_base_app_generator.py b/api/tests/unit_tests/core/app/apps/test_base_app_generator.py index fe07e420198..0b2bd95f436 100644 --- a/api/tests/unit_tests/core/app/apps/test_base_app_generator.py +++ b/api/tests/unit_tests/core/app/apps/test_base_app_generator.py @@ -1,10 +1,39 @@ import logging +from types import SimpleNamespace from unittest.mock import Mock import pytest +from sqlalchemy import inspect from core.app.apps.base_app_generator import BaseAppGenerator from graphon.variables.input_entities import VariableEntity, VariableEntityType +from models import Workflow, WorkflowRun + + +def test_restore_workflow_run_graph(): + workflow = Workflow(graph='{"nodes": [{"id": "edited"}]}') + session = SimpleNamespace(get=Mock(return_value=SimpleNamespace(graph='{"nodes": [{"id": "paused"}]}'))) + + BaseAppGenerator._restore_workflow_run_graph(session=session, workflow=workflow, workflow_run_id="run-id") + + session.get.assert_called_once_with(WorkflowRun, "run-id") + assert workflow.graph == '{"nodes": [{"id": "paused"}]}' + assert not inspect(workflow).attrs.graph.history.has_changes() + + +@pytest.mark.parametrize( + ("workflow_run_id", "workflow_run"), + [(None, None), ("run-id", None), ("run-id", SimpleNamespace(graph=None))], +) +def test_restore_workflow_run_graph_requires_persisted_snapshot(workflow_run_id, workflow_run): + session = SimpleNamespace(get=Mock(return_value=workflow_run)) + + with pytest.raises(ValueError): + BaseAppGenerator._restore_workflow_run_graph( + session=session, + workflow=Workflow(graph="{}"), + workflow_run_id=workflow_run_id, + ) def test_validate_inputs_with_zero(): diff --git a/api/tests/unit_tests/core/app/apps/test_message_based_app_queue_manager.py b/api/tests/unit_tests/core/app/apps/test_message_based_app_queue_manager.py index 847ad0ce9bc..0bc4753752d 100644 --- a/api/tests/unit_tests/core/app/apps/test_message_based_app_queue_manager.py +++ b/api/tests/unit_tests/core/app/apps/test_message_based_app_queue_manager.py @@ -6,7 +6,12 @@ from core.app.apps.base_app_queue_manager import PublishFrom from core.app.apps.exc import GenerateTaskStoppedError from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager from core.app.entities.app_invoke_entities import InvokeFrom -from core.app.entities.queue_entities import QueueErrorEvent, QueueMessageEndEvent, QueueStopEvent +from core.app.entities.queue_entities import ( + QueueErrorEvent, + QueueMessageEndEvent, + QueueStopEvent, + QueueWorkflowPausedEvent, +) class TestMessageBasedAppQueueManager: @@ -63,3 +68,21 @@ class TestMessageBasedAppQueueManager: manager._publish(QueueMessageEndEvent(), PublishFrom.TASK_PIPELINE) assert manager._q.qsize() == 1 + + def test_publish_pause_event_stops_listener_without_aborting_execution(self): + with patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis: + mock_redis.setex.return_value = True + manager = MessageBasedAppQueueManager( + task_id="t1", + user_id="u1", + invoke_from=InvokeFrom.DEBUGGER, + conversation_id="c1", + app_mode="advanced-chat", + message_id="m1", + ) + manager.stop_listen = Mock() + manager._is_stopped = Mock(return_value=False) + + manager._publish(QueueWorkflowPausedEvent(), PublishFrom.APPLICATION_MANAGER) + + manager.stop_listen.assert_called_once_with(execution_terminal=True) diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_core.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_core.py index fd643893f69..f32546d1b23 100644 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_core.py +++ b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_core.py @@ -334,7 +334,6 @@ class TestWorkflowBasedAppRunner: variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables()), start_at=0.0, ) - graph_runtime_state.register_paused_node("node-1") workflow_entry = SimpleNamespace(graph_engine=SimpleNamespace(graph_runtime_state=graph_runtime_state)) emails: list[dict] = [] diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_notifications.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_notifications.py index 778c6482635..f7f9bdb4ba2 100644 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_notifications.py +++ b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_notifications.py @@ -20,9 +20,6 @@ class _DummyQueueManager: class _DummyRuntimeState: variable_pool = object() - def get_paused_nodes(self): - return ["node-1"] - class _DummyGraphEngine: def __init__(self): diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_single_node.py b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_single_node.py index f02994fd61c..dd40e55e568 100644 --- a/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_single_node.py +++ b/api/tests/unit_tests/core/app/apps/test_workflow_app_runner_single_node.py @@ -130,6 +130,7 @@ def test_single_node_run_validates_target_node_config(monkeypatch: pytest.Monkey "type": "loop", "title": "Loop", "loop_count": 1, + "start_node_id": "loop-start", "break_conditions": [], "logical_operator": "and", }, diff --git a/api/tests/unit_tests/core/app/apps/test_workflow_pause_events.py b/api/tests/unit_tests/core/app/apps/test_workflow_pause_events.py index 41064ea8830..ca0c69e053a 100644 --- a/api/tests/unit_tests/core/app/apps/test_workflow_pause_events.py +++ b/api/tests/unit_tests/core/app/apps/test_workflow_pause_events.py @@ -40,9 +40,6 @@ class _RecordingWorkflowAppRunner(WorkflowAppRunner): class _FakeRuntimeState: variable_pool = object() - def get_paused_nodes(self): - return ["node-pause-1"] - @pytest.fixture def sqlite_pause_session(sqlite_session: Session, monkeypatch: pytest.MonkeyPatch) -> Session: @@ -140,6 +137,7 @@ def test_graph_run_paused_event_emits_queue_pause_event(monkeypatch: pytest.Monk "core.app.apps.workflow_app_runner.enrich_graph_pause_reasons", lambda **_: [enriched_reason], ) + monkeypatch.setattr("core.app.apps.workflow_app_runner.dispatch_human_input_email_task", MagicMock()) runner._handle_event(workflow_entry, event) @@ -148,7 +146,7 @@ def test_graph_run_paused_event_emits_queue_pause_event(monkeypatch: pytest.Monk assert isinstance(queue_event, QueueWorkflowPausedEvent) assert queue_event.reasons == [enriched_reason] assert queue_event.outputs == {"foo": "bar"} - assert queue_event.paused_nodes == ["node-pause-1"] + assert queue_event.paused_nodes == ["node-human"] def _build_converter(*, invoke_from: InvokeFrom = InvokeFrom.SERVICE_API): diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py b/api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py index 279b886cdf0..2c7573f82b9 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py @@ -543,6 +543,8 @@ class TestWorkflowAppGeneratorWorker: lambda self, *, session, workflow: workflow, ) monkeypatch.setattr("core.app.apps.workflow.app_generator.WorkflowAppRunner", _Runner) + restore_workflow_run_graph = Mock() + monkeypatch.setattr(generator, "_restore_workflow_run_graph", restore_workflow_run_graph) app_config = WorkflowUIBasedAppConfig( tenant_id="tenant", @@ -574,6 +576,12 @@ class TestWorkflowAppGeneratorWorker: variable_loader=SimpleNamespace(), workflow_execution_repository=SimpleNamespace(), workflow_node_execution_repository=SimpleNamespace(), + graph_runtime_state=SimpleNamespace(), ) assert runner_kwargs["system_user_id"] == "session-id" + restore_workflow_run_graph.assert_called_once_with( + session=session, + workflow=workflow, + workflow_run_id="run-id", + ) diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_app_queue_manager.py b/api/tests/unit_tests/core/app/apps/workflow/test_app_queue_manager.py index 0867a92f9d4..c33860b3c14 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_app_queue_manager.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_app_queue_manager.py @@ -1,6 +1,6 @@ from __future__ import annotations -from unittest.mock import patch +from unittest.mock import Mock, patch from core.app.apps.base_app_queue_manager import PublishFrom from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager @@ -41,6 +41,19 @@ class TestWorkflowAppQueueManager: manager._publish(QueuePingEvent(), PublishFrom.TASK_PIPELINE) + def test_publish_pause_event_stops_listener_without_aborting_execution(self): + manager = WorkflowAppQueueManager( + task_id="task", + user_id="user", + invoke_from=InvokeFrom.DEBUGGER, + app_mode="workflow", + ) + manager.stop_listen = Mock() + + manager._publish(QueueWorkflowPausedEvent(), PublishFrom.APPLICATION_MANAGER) + + manager.stop_listen.assert_called_once_with(execution_terminal=True) + def test_listener_close_aborts_unfinished_execution(self): with ( patch("core.app.apps.base_app_queue_manager.redis_client") as redis_client, diff --git a/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline_core.py b/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline_core.py index 9f83e657cfb..884b7a7a8aa 100644 --- a/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline_core.py +++ b/api/tests/unit_tests/core/app/apps/workflow/test_generate_task_pipeline_core.py @@ -50,6 +50,7 @@ from core.app.entities.task_entities import ( from core.base.tts.app_generator_tts_publisher import AudioTrunk from core.workflow.system_variables import build_system_variables, system_variables_to_mapping from graphon.enums import BuiltinNodeTypes, WorkflowExecutionStatus +from graphon.model_runtime.entities.llm_entities import LLMUsage from graphon.runtime import GraphRuntimeState, VariablePool from libs.datetime_utils import naive_utc_now from models.enums import CreatorUserRole @@ -103,7 +104,7 @@ class TestWorkflowGenerateTaskPipeline: variables=build_system_variables(workflow_execution_id="run-id"), ), start_at=0.0, - total_tokens=5, + llm_usage=LLMUsage.empty_usage().model_copy(update={"total_tokens": 5}), node_run_steps=2, ) diff --git a/api/tests/unit_tests/core/app/workflow/test_persistence_layer.py b/api/tests/unit_tests/core/app/workflow/test_persistence_layer.py index 52d0db414c2..bc67bfc12d5 100644 --- a/api/tests/unit_tests/core/app/workflow/test_persistence_layer.py +++ b/api/tests/unit_tests/core/app/workflow/test_persistence_layer.py @@ -9,7 +9,7 @@ from core.app.entities.app_invoke_entities import WorkflowAppGenerateEntity from core.app.workflow.layers.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer from core.ops.ops_trace_manager import TraceTask, TraceTaskName from core.workflow.system_variables import SystemVariableKey, build_system_variables -from graphon.entities import WorkflowNodeExecution +from graphon.entities import WorkflowNodeExecution, WorkflowStartReason from graphon.entities.pause_reason import SchedulingPause from graphon.enums import ( BuiltinNodeTypes, @@ -32,6 +32,7 @@ from graphon.graph_events import ( NodeRunStartedEvent, NodeRunSucceededEvent, ) +from graphon.model_runtime.entities.llm_entities import LLMUsage from graphon.node_events import NodeRunResult from graphon.runtime import GraphRuntimeState, ReadOnlyGraphRuntimeStateWrapper, VariablePool @@ -41,6 +42,7 @@ class _RepoRecorder: self.saved: list[object] = [] self.synchronously_saved: list[object] = [] self.saved_exec_data: list[object] = [] + self.loaded: list[object] = [] def save(self, entity): self.saved.append(entity) @@ -51,6 +53,9 @@ class _RepoRecorder: def save_execution_data(self, entity): self.saved_exec_data.append(entity) + def get_by_workflow_execution(self, _workflow_execution_id): + return self.loaded + def _naive_utc_now() -> datetime: return datetime.now(UTC).replace(tzinfo=None) @@ -169,12 +174,45 @@ class TestWorkflowPersistenceLayer: assert exec_repo.saved + def test_resumption_restores_container_execution_before_terminal_event(self): + layer, _, node_repo, _ = _make_layer() + started_at = _naive_utc_now() + execution = WorkflowNodeExecution( + id="loop-exec", + workflow_id="workflow-id", + workflow_execution_id="run-id", + index=4, + node_id="loop", + node_type=BuiltinNodeTypes.LOOP, + title="Loop", + status=WorkflowNodeExecutionStatus.RUNNING, + created_at=started_at, + ) + node_repo.loaded = [execution] + + layer.on_event(GraphRunStartedEvent(reason=WorkflowStartReason.RESUMPTION)) + layer.on_event( + NodeRunSucceededEvent( + id=execution.id, + node_id=execution.node_id, + node_type=execution.node_type, + start_at=started_at, + node_run_result=NodeRunResult(status=WorkflowNodeExecutionStatus.SUCCEEDED), + ) + ) + + assert execution.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert layer._next_node_sequence() == 5 + def test_handle_graph_run_succeeded_updates_execution(self): layer, exec_repo, _, runtime_state = _make_layer() layer._handle_graph_run_started() - runtime_state.total_tokens = 3 - runtime_state.node_run_steps = 2 - runtime_state.outputs = {"out": "v"} + usage = LLMUsage.empty_usage() + usage.total_tokens = 3 + runtime_state.add_llm_usage(usage) + for _ in range(2): + runtime_state.increment_node_run_steps() + runtime_state.set_output("out", "v") layer._handle_graph_run_succeeded(GraphRunSucceededEvent(outputs={"ok": True})) @@ -186,8 +224,11 @@ class TestWorkflowPersistenceLayer: def test_handle_graph_run_partial_succeeded_updates_execution(self): layer, exec_repo, _, runtime_state = _make_layer() layer._handle_graph_run_started() - runtime_state.total_tokens = 5 - runtime_state.node_run_steps = 4 + usage = LLMUsage.empty_usage() + usage.total_tokens = 5 + runtime_state.add_llm_usage(usage) + for _ in range(4): + runtime_state.increment_node_run_steps() runtime_state._graph_execution = SimpleNamespace(exceptions_count=2) layer._handle_graph_run_partial_succeeded( @@ -293,8 +334,11 @@ class TestWorkflowPersistenceLayer: def test_handle_graph_run_paused_updates_outputs(self): layer, exec_repo, _, runtime_state = _make_layer() layer._handle_graph_run_started() - runtime_state.total_tokens = 7 - runtime_state.node_run_steps = 5 + usage = LLMUsage.empty_usage() + usage.total_tokens = 7 + runtime_state.add_llm_usage(usage) + for _ in range(5): + runtime_state.increment_node_run_steps() layer._handle_graph_run_paused(GraphRunPausedEvent(outputs={"pause": True})) diff --git a/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py b/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py index b502144b0cf..8e798364756 100644 --- a/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py +++ b/api/tests/unit_tests/core/repositories/test_celery_workflow_node_execution_repository.py @@ -262,6 +262,60 @@ class TestCeleryWorkflowNodeExecutionRepository: # Should return empty list since nothing in cache assert len(result) == 0 + def test_get_by_workflow_execution_loads_persisted_executions_on_cache_miss( + self, mock_session_factory, mock_account, sample_workflow_node_execution + ): + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=mock_account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + repo._sql_repository = Mock() + repo._sql_repository.get_by_workflow_execution.return_value = [sample_workflow_node_execution] + + result = repo.get_by_workflow_execution(sample_workflow_node_execution.workflow_execution_id) + + assert result == [sample_workflow_node_execution] + assert repo._execution_cache[sample_workflow_node_execution.id] is sample_workflow_node_execution + assert repo._workflow_execution_mapping[sample_workflow_node_execution.workflow_execution_id] == [ + sample_workflow_node_execution.id + ] + + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") + def test_get_by_workflow_execution_merges_database_and_newer_cache( + self, mock_task, mock_session_factory, mock_account, sample_workflow_node_execution + ): + repo = CeleryWorkflowNodeExecutionRepository( + session_factory=mock_session_factory, + tenant_id=RESOURCE_TENANT_ID, + user=mock_account, + app_id="test-app", + triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN, + ) + persisted_current = sample_workflow_node_execution.model_copy(deep=True) + historical = sample_workflow_node_execution.model_copy( + update={ + "id": str(uuid4()), + "node_execution_id": str(uuid4()), + "index": 0, + "node_id": "start", + } + ) + sample_workflow_node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED + repo.save(sample_workflow_node_execution) + repo._sql_repository = Mock() + repo._sql_repository.get_by_workflow_execution.return_value = [persisted_current, historical] + + result = repo.get_by_workflow_execution( + sample_workflow_node_execution.workflow_execution_id, + OrderConfig(order_by=["index"], order_direction="asc"), + ) + + assert [execution.id for execution in result] == [historical.id, sample_workflow_node_execution.id] + assert result[1] is sample_workflow_node_execution + @patch("core.repositories.celery_workflow_node_execution_repository.save_workflow_node_execution_task") def test_cache_operations(self, mock_task, mock_session_factory, mock_account, sample_workflow_node_execution): """Test cache operations work correctly.""" diff --git a/api/tests/unit_tests/core/tools/test_tool_manager.py b/api/tests/unit_tests/core/tools/test_tool_manager.py index 1f3109aeac0..a2947dbd94a 100644 --- a/api/tests/unit_tests/core/tools/test_tool_manager.py +++ b/api/tests/unit_tests/core/tools/test_tool_manager.py @@ -1083,14 +1083,14 @@ def test_convert_tool_parameters_type_agent_and_workflow_branches(): variable_pool = Mock() variable_pool.get.return_value = SimpleNamespace(value="from-variable") - variable_pool.convert_template.return_value = SimpleNamespace(text="from-template") - mixed = ToolManager._convert_tool_parameters_type( - parameters=[text_param], - variable_pool=variable_pool, - tool_configurations={"text": {"type": "mixed", "value": "Hello {{name}}"}}, - typ="workflow", - ) + with patch("core.tools.tool_manager.convert_template", return_value=SimpleNamespace(text="from-template")): + mixed = ToolManager._convert_tool_parameters_type( + parameters=[text_param], + variable_pool=variable_pool, + tool_configurations={"text": {"type": "mixed", "value": "Hello {{name}}"}}, + typ="workflow", + ) assert mixed == {"text": "from-template"} variable = ToolManager._convert_tool_parameters_type( diff --git a/api/tests/unit_tests/core/variables/test_segment.py b/api/tests/unit_tests/core/variables/test_segment.py index 9e07ea1b6db..f65e5bbde75 100644 --- a/api/tests/unit_tests/core/variables/test_segment.py +++ b/api/tests/unit_tests/core/variables/test_segment.py @@ -28,6 +28,7 @@ from graphon.variables.segments import ( StringSegment, get_segment_discriminator, ) +from graphon.variables.template_resolution import convert_template from graphon.variables.types import SegmentType from graphon.variables.utils import ( dumps_with_segments, @@ -98,7 +99,7 @@ def test_segment_group_to_text(): template = ( "Hello, {{#sys.user_id#}}! Your query is {{#node_id.custom_query#}}. And your key is {{#env.secret_key#}}." ) - segments_group = variable_pool.convert_template(template) + segments_group = convert_template(variable_pool, template) assert segments_group.text == "Hello, fake-user-id! Your query is fake-user-query. And your key is fake-secret-key." assert segments_group.log == ( @@ -112,7 +113,7 @@ def test_convert_constant_to_segment_group(): system_variables=build_system_variables(user_id="1", app_id="1", workflow_id="1"), ) template = "Hello, world!" - segments_group = variable_pool.convert_template(template) + segments_group = convert_template(variable_pool, template) assert segments_group.text == "Hello, world!" assert segments_group.log == "Hello, world!" @@ -120,7 +121,7 @@ def test_convert_constant_to_segment_group(): def test_convert_variable_to_segment_group(): variable_pool = _build_variable_pool(system_variables=build_system_variables(user_id="fake-user-id")) template = "{{#sys.user_id#}}" - segments_group = variable_pool.convert_template(template) + segments_group = convert_template(variable_pool, template) assert segments_group.text == "fake-user-id" assert segments_group.log == "fake-user-id" assert isinstance(segments_group.value[0], StringVariable) diff --git a/api/tests/unit_tests/core/workflow/graph_engine/test_mock_factory.py b/api/tests/unit_tests/core/workflow/graph_engine/test_mock_factory.py index c721c7b0ebd..9c507bb1619 100644 --- a/api/tests/unit_tests/core/workflow/graph_engine/test_mock_factory.py +++ b/api/tests/unit_tests/core/workflow/graph_engine/test_mock_factory.py @@ -5,7 +5,7 @@ The factory follows the same config adaptation path as production implementations before instantiation. """ -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, override from core.workflow.human_input_adapter import adapt_node_config_for_graph from core.workflow.node_factory import DifyNodeFactory @@ -76,6 +76,14 @@ class MockNodeFactory(DifyNodeFactory): BuiltinNodeTypes.CODE: MockCodeNode, } + @override + def with_runtime_state(self, graph_runtime_state: "GraphRuntimeState") -> "MockNodeFactory": + return MockNodeFactory( + graph_init_params=self.graph_init_params, + graph_runtime_state=graph_runtime_state, + mock_config=self.mock_config, + ) + def create_node(self, node_config: dict[str, Any] | NodeConfigDict) -> Node: """ Create a node instance, using mock implementations for third-party service nodes. diff --git a/api/tests/unit_tests/core/workflow/graph_engine/test_mock_nodes.py b/api/tests/unit_tests/core/workflow/graph_engine/test_mock_nodes.py index 55dcbdb7a11..67b94e46399 100644 --- a/api/tests/unit_tests/core/workflow/graph_engine/test_mock_nodes.py +++ b/api/tests/unit_tests/core/workflow/graph_engine/test_mock_nodes.py @@ -615,69 +615,6 @@ class MockIterationNode(MockNodeMixin, IterationNode): """Return the version of this mock node.""" return "1" - def _create_graph_engine(self, index: int, item: Any): - """Create a graph engine with MockNodeFactory instead of DifyNodeFactory.""" - # Import dependencies - from graphon.entities import GraphInitParams - from graphon.graph import Graph - from graphon.graph_engine import GraphEngine, GraphEngineConfig - from graphon.graph_engine.command_channels import InMemoryChannel - from graphon.runtime import GraphRuntimeState - - # Import our MockNodeFactory instead of DifyNodeFactory - from .test_mock_factory import MockNodeFactory - - # Create GraphInitParams from node attributes - graph_init_params = GraphInitParams( - workflow_id=self.workflow_id, - graph_config=self.graph_config, - run_context=self.run_context, - call_depth=self.workflow_call_depth, - ) - - # Create a deep copy of the variable pool for each iteration - variable_pool_copy = self.graph_runtime_state.variable_pool.model_copy(deep=True) - - # append iteration variable (item, index) to variable pool - variable_pool_copy.add([self._node_id, "index"], index) - variable_pool_copy.add([self._node_id, "item"], item) - - # Create a new GraphRuntimeState for this iteration - graph_runtime_state_copy = GraphRuntimeState( - variable_pool=variable_pool_copy, - start_at=self.graph_runtime_state.start_at, - total_tokens=0, - node_run_steps=0, - ) - - # Create a MockNodeFactory with the same mock_config - node_factory = MockNodeFactory( - graph_init_params=graph_init_params, - graph_runtime_state=graph_runtime_state_copy, - mock_config=self.mock_config, # Pass the mock configuration - ) - - # Initialize the iteration graph with the mock node factory - iteration_graph = Graph.init( - graph_config=self.graph_config, node_factory=node_factory, root_node_id=self._node_data.start_node_id - ) - - if not iteration_graph: - from graphon.nodes.iteration.exc import IterationGraphNotFoundError - - raise IterationGraphNotFoundError("iteration graph not found") - - # Create a new GraphEngine for this iteration - graph_engine = GraphEngine( - workflow_id=self.workflow_id, - graph=iteration_graph, - graph_runtime_state=graph_runtime_state_copy, - command_channel=InMemoryChannel(), # Use InMemoryChannel for sub-graphs - config=GraphEngineConfig(), - ) - - return graph_engine - class MockLoopNode(MockNodeMixin, LoopNode): """Mock implementation of LoopNode that preserves mock configuration.""" @@ -687,56 +624,6 @@ class MockLoopNode(MockNodeMixin, LoopNode): """Return the version of this mock node.""" return "1" - def _create_graph_engine(self, start_at, root_node_id: str): - """Create a graph engine with MockNodeFactory instead of DifyNodeFactory.""" - # Import dependencies - from graphon.entities import GraphInitParams - from graphon.graph import Graph - from graphon.graph_engine import GraphEngine, GraphEngineConfig - from graphon.graph_engine.command_channels import InMemoryChannel - from graphon.runtime import GraphRuntimeState - - # Import our MockNodeFactory instead of DifyNodeFactory - from .test_mock_factory import MockNodeFactory - - # Create GraphInitParams from node attributes - graph_init_params = GraphInitParams( - workflow_id=self.workflow_id, - graph_config=self.graph_config, - run_context=self.run_context, - call_depth=self.workflow_call_depth, - ) - - # Create a new GraphRuntimeState for this iteration - graph_runtime_state_copy = GraphRuntimeState( - variable_pool=self.graph_runtime_state.variable_pool, - start_at=start_at.timestamp(), - ) - - # Create a MockNodeFactory with the same mock_config - node_factory = MockNodeFactory( - graph_init_params=graph_init_params, - graph_runtime_state=graph_runtime_state_copy, - mock_config=self.mock_config, # Pass the mock configuration - ) - - # Initialize the loop graph with the mock node factory - loop_graph = Graph.init(graph_config=self.graph_config, node_factory=node_factory, root_node_id=root_node_id) - - if not loop_graph: - raise ValueError("loop graph not found") - - # Create a new GraphEngine for this iteration - graph_engine = GraphEngine( - workflow_id=self.workflow_id, - graph=loop_graph, - graph_runtime_state=graph_runtime_state_copy, - command_channel=InMemoryChannel(), # Use InMemoryChannel for sub-graphs - config=GraphEngineConfig(), - ) - - return graph_engine - class MockTemplateTransformNode(MockNodeMixin, TemplateTransformNode): """Mock implementation of TemplateTransformNode for testing.""" diff --git a/api/tests/unit_tests/core/workflow/graph_engine/test_table_runner.py b/api/tests/unit_tests/core/workflow/graph_engine/test_table_runner.py index e9c9e04e17b..175b5da8a1b 100644 --- a/api/tests/unit_tests/core/workflow/graph_engine/test_table_runner.py +++ b/api/tests/unit_tests/core/workflow/graph_engine/test_table_runner.py @@ -51,53 +51,6 @@ from .test_mock_factory import MockNodeFactory logger = logging.getLogger(__name__) -class _TableTestChildEngineBuilder: - def __init__(self, *, use_mock_factory: bool, mock_config: MockConfig | None) -> None: - self._use_mock_factory = use_mock_factory - self._mock_config = mock_config - - def build_child_engine( - self, - *, - workflow_id: str, - graph_init_params: GraphInitParams, - parent_graph_runtime_state: GraphRuntimeState, - root_node_id: str, - variable_pool: VariablePool | None = None, - ) -> GraphEngine: - child_graph_runtime_state = GraphRuntimeState( - variable_pool=variable_pool if variable_pool is not None else parent_graph_runtime_state.variable_pool, - start_at=time.perf_counter(), - execution_context=parent_graph_runtime_state.execution_context, - ) - if self._use_mock_factory: - node_factory = MockNodeFactory( - graph_init_params=graph_init_params, - graph_runtime_state=child_graph_runtime_state, - mock_config=self._mock_config, - ) - else: - node_factory = DifyNodeFactory( - graph_init_params=graph_init_params, - graph_runtime_state=child_graph_runtime_state, - ) - - graph_config = graph_init_params.graph_config - child_graph = Graph.init(graph_config=graph_config, node_factory=node_factory, root_node_id=root_node_id) - if not child_graph: - raise ValueError("child graph not found") - - child_engine = GraphEngine( - workflow_id=workflow_id, - graph=child_graph, - graph_runtime_state=child_graph_runtime_state, - command_channel=InMemoryChannel(), - config=GraphEngineConfig(), - child_engine_builder=self, - ) - return child_engine - - @dataclass class WorkflowTestCase: """Represents a single test case for table-driven testing.""" @@ -379,10 +332,6 @@ class TableTestRunner: scale_up_threshold=self.graph_engine_scale_up_threshold, scale_down_idle_time=self.graph_engine_scale_down_idle_time, ), - child_engine_builder=_TableTestChildEngineBuilder( - use_mock_factory=test_case.use_auto_mock, - mock_config=test_case.mock_config, - ), ) # Execute and collect events diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py index 4dd7dc9fd80..d1df3ac95c9 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py @@ -3,7 +3,7 @@ from datetime import UTC, datetime from types import SimpleNamespace from typing import cast from unittest.mock import MagicMock, patch -from uuid import UUID +from uuid import UUID, uuid4 from agenton.compositor import CompositorSessionSnapshot from dify_agent.layers.ask_human import AskHumanToolResult @@ -347,7 +347,7 @@ def _node( } ) - return DifyAgentNode( + node = DifyAgentNode( node_id="agent-node", data=DifyAgentNodeData.model_validate({"type": BuiltinNodeTypes.AGENT, "version": "2"}), graph_init_params=graph_init_params, @@ -355,7 +355,7 @@ def _node( GraphRuntimeState, SimpleNamespace( variable_pool=FakeVariablePool(), - graph_execution=SimpleNamespace(node_executions={}), + graph_execution=SimpleNamespace(aborted=False), ), ), binding_resolver=binding_resolver, @@ -368,6 +368,8 @@ def _node( failure_orchestrator=OutputFailureOrchestrator(), session_store=cast(WorkflowAgentWorkspaceStore, session_store or FakeSessionStore()), ) + node.bind_execution_id(str(uuid4())) + return node def test_extract_variable_selector_to_variable_mapping_uses_frontend_agent_task_markers(): @@ -465,7 +467,7 @@ def test_agent_node_passes_execution_id_to_session_store_and_runtime_request_bui store = FakeSessionStore() request_builder = WorkflowAgentRuntimeRequestBuilder(credentials_provider=FakeCredentialsProvider()) node = _node(session_store=store, runtime_request_builder=request_builder) - execution_id = node.ensure_execution_id() + execution_id = node.execution_id with patch.object(request_builder, "build", wraps=request_builder.build) as build: list(node._run()) diff --git a/api/tests/unit_tests/core/workflow/nodes/human_input/test_human_input_form_filled_event.py b/api/tests/unit_tests/core/workflow/nodes/human_input/test_human_input_form_filled_event.py index a1cb0af85fe..33d1a41fad8 100644 --- a/api/tests/unit_tests/core/workflow/nodes/human_input/test_human_input_form_filled_event.py +++ b/api/tests/unit_tests/core/workflow/nodes/human_input/test_human_input_form_filled_event.py @@ -73,13 +73,15 @@ def _create_human_input_node( node_data=node_data, file_reference_factory=_TestFileReferenceFactory(), ) - return HumanInputNode( + node = HumanInputNode( node_id=config["id"], data=node_data, graph_init_params=graph_init_params, graph_runtime_state=graph_runtime_state, hitl_callback=callback, ) + node.bind_execution_id("00000000-0000-4000-8000-000000000001") + return node def _build_node( diff --git a/api/tests/unit_tests/core/workflow/nodes/iteration/test_iteration_child_engine_errors.py b/api/tests/unit_tests/core/workflow/nodes/iteration/test_iteration_child_engine_errors.py deleted file mode 100644 index 18ed7a0b1d6..00000000000 --- a/api/tests/unit_tests/core/workflow/nodes/iteration/test_iteration_child_engine_errors.py +++ /dev/null @@ -1,95 +0,0 @@ -from collections.abc import Mapping -from typing import Any - -import pytest - -from core.workflow.system_variables import default_system_variables -from graphon.entities import GraphInitParams -from graphon.nodes.iteration.entities import IterationNodeData -from graphon.nodes.iteration.exc import IterationGraphNotFoundError -from graphon.nodes.iteration.iteration_node import IterationNode -from graphon.runtime import ( - ChildEngineBuilderNotConfiguredError, - ChildGraphNotFoundError, - GraphRuntimeState, - VariablePool, -) -from tests.workflow_test_utils import build_test_graph_init_params - - -class _MissingGraphBuilder: - def build_child_engine( - self, - *, - workflow_id: str, - graph_init_params: GraphInitParams, - parent_graph_runtime_state: GraphRuntimeState, - root_node_id: str, - variable_pool: VariablePool | None = None, - ) -> object: - raise ChildGraphNotFoundError(f"child graph root node '{root_node_id}' not found") - - -def _build_runtime_state() -> GraphRuntimeState: - return GraphRuntimeState( - variable_pool=VariablePool.from_bootstrap(system_variables=default_system_variables(), user_inputs={}), - start_at=0.0, - ) - - -def _build_iteration_node( - *, - graph_config: Mapping[str, Any], - runtime_state: GraphRuntimeState, - start_node_id: str, -) -> IterationNode: - init_params = build_test_graph_init_params(graph_config=graph_config) - return IterationNode( - node_id="iteration-node", - data=IterationNodeData( - type="iteration", - title="Iteration", - iterator_selector=["start", "items"], - output_selector=["iteration-node", "output"], - start_node_id=start_node_id, - ), - graph_init_params=init_params, - graph_runtime_state=runtime_state, - ) - - -def test_graph_runtime_state_raises_specific_error_when_child_builder_is_missing(): - runtime_state = _build_runtime_state() - graph_init_params = build_test_graph_init_params() - - with pytest.raises(ChildEngineBuilderNotConfiguredError): - runtime_state.create_child_engine( - workflow_id="workflow", - graph_init_params=graph_init_params, - root_node_id="root", - ) - - -def test_iteration_node_only_translates_child_graph_not_found_error(): - runtime_state = _build_runtime_state() - runtime_state.bind_child_engine_builder(_MissingGraphBuilder()) - node = _build_iteration_node( - graph_config={"nodes": [{"id": "present-node"}], "edges": []}, - runtime_state=runtime_state, - start_node_id="missing-node", - ) - - with pytest.raises(IterationGraphNotFoundError): - node._create_graph_engine(index=0, item="item") - - -def test_iteration_node_propagates_non_graph_not_found_errors(): - runtime_state = _build_runtime_state() - node = _build_iteration_node( - graph_config={"nodes": [{"id": "start-node"}], "edges": []}, - runtime_state=runtime_state, - start_node_id="start-node", - ) - - with pytest.raises(ChildEngineBuilderNotConfiguredError): - node._create_graph_engine(index=0, item="item") diff --git a/api/tests/unit_tests/core/workflow/nodes/list_operator/node_spec.py b/api/tests/unit_tests/core/workflow/nodes/list_operator/node_spec.py index 20b94d5d509..2c47ba93926 100644 --- a/api/tests/unit_tests/core/workflow/nodes/list_operator/node_spec.py +++ b/api/tests/unit_tests/core/workflow/nodes/list_operator/node_spec.py @@ -1,4 +1,3 @@ -from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -36,7 +35,6 @@ class TestListOperatorNode: """Create mock GraphRuntimeState.""" mock_state = MagicMock(spec=GraphRuntimeState) mock_variable_pool = MagicMock() - mock_variable_pool.convert_template.side_effect = lambda value: SimpleNamespace(text=value) mock_state.variable_pool = mock_variable_pool return mock_state diff --git a/api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node.py b/api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node.py index 0ee70256d7f..f898a8a8f4b 100644 --- a/api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/tool/test_tool_node.py @@ -239,7 +239,6 @@ def test_image_link_messages_use_tool_file_id_metadata(tool_node: ToolNode): def test_tool_node_passes_node_execution_id_when_runtime_accepts_it(tool_node: ToolNode): runtime_handle = ToolRuntimeHandle(raw=object()) tool_node._runtime.get_runtime = MagicMock(return_value=runtime_handle) - tool_node.ensure_execution_id = MagicMock(return_value="node-execution-id") result = tool_node._get_tool_runtime( variable_pool=tool_node.graph_runtime_state.variable_pool, diff --git a/api/tests/unit_tests/core/workflow/test_human_input_adapter.py b/api/tests/unit_tests/core/workflow/test_human_input_adapter.py index 7a6328ffb4b..30352738d0b 100644 --- a/api/tests/unit_tests/core/workflow/test_human_input_adapter.py +++ b/api/tests/unit_tests/core/workflow/test_human_input_adapter.py @@ -18,12 +18,12 @@ from core.workflow.human_input_adapter import ( ) from graphon.enums import BuiltinNodeTypes from graphon.nodes.base.variable_template_parser import VariableTemplateParser +from graphon.runtime import VariablePool def test_email_delivery_config_helpers_render_and_sanitize_text() -> None: - variable_pool = SimpleNamespace( - convert_template=lambda body: SimpleNamespace(text=body.replace("{{#node.value#}}", "42")) - ) + variable_pool = VariablePool() + variable_pool.add(["node", "value"], "42") rendered = EmailDeliveryConfig.render_body_template( body="Open {{#url#}} and use {{#node.value#}}", diff --git a/api/tests/unit_tests/core/workflow/test_human_input_callback.py b/api/tests/unit_tests/core/workflow/test_human_input_callback.py index 4e4479b1594..88477049d84 100644 --- a/api/tests/unit_tests/core/workflow/test_human_input_callback.py +++ b/api/tests/unit_tests/core/workflow/test_human_input_callback.py @@ -59,6 +59,27 @@ def test_dify_hitl_callback_creates_pause_requested_for_new_form() -> None: assert params.node_id == "node-1" +def test_dify_hitl_callback_scopes_form_to_node_execution() -> None: + repository = MagicMock(spec=HumanInputFormRepository) + repository.get_form.return_value = None + repository.create_form.return_value = SimpleNamespace(id="execution-1") + callback = DifyHITLCallback( + form_repository=repository, + node_data=HumanInputNodeData( + title="Approval", + form_content="Please approve", + user_actions=[UserActionConfig(id="approve", title="Approve")], + ), + execution_id_getter=lambda: "execution-1", + ) + + callback(_ctx("run-1", "node-1")) + + repository.get_form.assert_called_once_with("node-1", form_id="execution-1") + params: FormCreateParams = repository.create_form.call_args.args[0] + assert params.form_id == "execution-1" + + def test_dify_hitl_callback_returns_completed_for_submitted_form() -> None: repository = MagicMock(spec=HumanInputFormRepository) repository.get_form.return_value = SimpleNamespace( diff --git a/api/tests/unit_tests/core/workflow/test_node_factory.py b/api/tests/unit_tests/core/workflow/test_node_factory.py index 3d305d3a8f4..3bff332b8ae 100644 --- a/api/tests/unit_tests/core/workflow/test_node_factory.py +++ b/api/tests/unit_tests/core/workflow/test_node_factory.py @@ -324,6 +324,19 @@ class TestDifyNodeFactoryInit: graph_runtime_state=sentinel.graph_runtime_state, ) + def test_with_runtime_state_rebinds_factory(self): + factory = object.__new__(node_factory.DifyNodeFactory) + factory.graph_init_params = sentinel.graph_init_params + + with patch.object(node_factory, "DifyNodeFactory", return_value=sentinel.factory) as factory_cls: + rebound = factory.with_runtime_state(sentinel.graph_runtime_state) + + assert rebound is sentinel.factory + factory_cls.assert_called_once_with( + graph_init_params=sentinel.graph_init_params, + graph_runtime_state=sentinel.graph_runtime_state, + ) + def test_init_builds_default_dependencies(self): graph_init_params = SimpleNamespace(run_context={"context": "value"}) graph_runtime_state = sentinel.graph_runtime_state diff --git a/api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py b/api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py index 41037233b8c..3a06ea3f456 100644 --- a/api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py +++ b/api/tests/unit_tests/core/workflow/test_workflow_entry_helpers.py @@ -1,5 +1,4 @@ from collections import UserString -from contextlib import nullcontext from types import SimpleNamespace from unittest.mock import MagicMock, patch, sentinel @@ -10,238 +9,20 @@ from core.app.entities.app_invoke_entities import InvokeFrom, UserFrom from core.workflow import workflow_entry from core.workflow.system_variables import default_system_variables from graphon.entities.base_node_data import BaseNodeData -from graphon.enums import NodeType, WorkflowNodeExecutionStatus +from graphon.enums import NodeType from graphon.errors import WorkflowNodeRunFailedError from graphon.file import File, FileTransferMethod, FileType from graphon.filters import ResponseStreamFilter -from graphon.graph import Graph from graphon.graph_events import GraphRunFailedEvent -from graphon.model_runtime.entities.llm_entities import LLMMode, LLMUsage -from graphon.node_events import NodeRunResult from graphon.nodes import BuiltinNodeTypes -from graphon.nodes.base.node import Node -from graphon.nodes.llm.entities import ContextConfig, LLMNodeData, ModelConfig -from graphon.nodes.question_classifier.entities import QuestionClassifierNodeData -from graphon.runtime import ChildGraphNotFoundError, VariablePool +from graphon.runtime import VariablePool from graphon.variables.variables import StringVariable -from tests.workflow_test_utils import build_test_graph_init_params, build_test_variable_pool def _build_typed_node_config(node_type: NodeType): return {"id": "node-id", "data": BaseNodeData(type=node_type)} -def _build_model_config(*, provider: str = "openai", model_name: str = "gpt-4o") -> ModelConfig: - return ModelConfig(provider=provider, name=model_name, mode=LLMMode.CHAT) - - -def _build_llm_node_data(*, provider: str = "openai", model_name: str = "gpt-4o") -> LLMNodeData: - return LLMNodeData( - type=BuiltinNodeTypes.LLM, - title="Child Model", - model=_build_model_config(provider=provider, model_name=model_name), - prompt_template=[], - context=ContextConfig(enabled=False), - ) - - -def _build_question_classifier_node_data( - *, provider: str = "openai", model_name: str = "gpt-4o" -) -> QuestionClassifierNodeData: - return QuestionClassifierNodeData( - type=BuiltinNodeTypes.QUESTION_CLASSIFIER, - title="Child Model", - query_variable_selector=["sys", "query"], - model=_build_model_config(provider=provider, model_name=model_name), - classes=[], - ) - - -class _FakeModelNodeMixin: - @classmethod - def version(cls) -> str: - return "1" - - def post_init(self) -> None: - self.model_instance = SimpleNamespace(provider="stale-provider", model_name="stale-model") - self.usage_snapshot = LLMUsage.empty_usage() - self.usage_snapshot.total_tokens = 1 - - def _run(self) -> NodeRunResult: - return NodeRunResult( - status=WorkflowNodeExecutionStatus.SUCCEEDED, - inputs={ - "model_provider": self.node_data.model.provider, - "model_name": self.node_data.model.name, - }, - llm_usage=self.usage_snapshot, - ) - - -class _FakeLLMNode(_FakeModelNodeMixin, Node[LLMNodeData]): - node_type = BuiltinNodeTypes.LLM - - -class _FakeQuestionClassifierNode(_FakeModelNodeMixin, Node[QuestionClassifierNodeData]): - node_type = BuiltinNodeTypes.QUESTION_CLASSIFIER - - -class TestWorkflowChildEngineBuilder: - @pytest.mark.parametrize( - ("graph_config", "node_id", "expected"), - [ - ({"nodes": [{"id": "root"}]}, "root", True), - ({"nodes": [{"id": "root"}]}, "other", False), - ({"nodes": "invalid"}, "root", None), - ({"nodes": ["invalid"]}, "root", None), - ], - ) - def test_has_node_id(self, graph_config, node_id, expected): - result = workflow_entry._WorkflowChildEngineBuilder._has_node_id(graph_config, node_id) - - assert result is expected - - def test_build_child_engine_raises_when_root_node_is_missing(self): - builder = workflow_entry._WorkflowChildEngineBuilder(tenant_id="tenant-id") - graph_init_params = SimpleNamespace(graph_config={"nodes": []}) - parent_graph_runtime_state = SimpleNamespace( - execution_context=sentinel.execution_context, - variable_pool=sentinel.variable_pool, - ) - - with patch.object(workflow_entry, "DifyNodeFactory", return_value=sentinel.factory): - with pytest.raises(ChildGraphNotFoundError, match="child graph root node 'missing' not found"): - builder.build_child_engine( - workflow_id="workflow-id", - graph_init_params=graph_init_params, - parent_graph_runtime_state=parent_graph_runtime_state, - root_node_id="missing", - ) - - def test_build_child_engine_constructs_graph_engine_with_quota_layer_only(self): - builder = workflow_entry._WorkflowChildEngineBuilder(tenant_id="tenant-id") - graph_init_params = SimpleNamespace(graph_config={"nodes": [{"id": "root"}]}) - parent_graph_runtime_state = SimpleNamespace( - execution_context=sentinel.execution_context, - variable_pool=sentinel.parent_variable_pool, - ) - child_graph = sentinel.child_graph - child_graph_runtime_state = sentinel.child_graph_runtime_state - child_engine = MagicMock() - - with ( - patch.object(workflow_entry.time, "perf_counter", return_value=123.0), - patch.object( - workflow_entry, - "GraphRuntimeState", - return_value=child_graph_runtime_state, - ) as graph_runtime_state_cls, - patch.object(workflow_entry, "DifyNodeFactory", return_value=sentinel.factory) as dify_node_factory, - patch.object(workflow_entry.Graph, "init", return_value=child_graph) as graph_init, - patch.object(workflow_entry, "GraphEngine", return_value=child_engine) as graph_engine_cls, - patch.object(workflow_entry, "GraphEngineConfig", return_value=sentinel.graph_engine_config), - patch.object(workflow_entry, "InMemoryChannel", return_value=sentinel.command_channel), - patch.object(workflow_entry, "LLMQuotaLayer", return_value=sentinel.llm_quota_layer) as llm_quota_layer_cls, - ): - result = builder.build_child_engine( - workflow_id="workflow-id", - graph_init_params=graph_init_params, - parent_graph_runtime_state=parent_graph_runtime_state, - root_node_id="root", - variable_pool=sentinel.child_variable_pool, - ) - - assert result is child_engine - graph_runtime_state_cls.assert_called_once_with( - variable_pool=sentinel.child_variable_pool, - start_at=123.0, - execution_context=sentinel.execution_context, - ) - dify_node_factory.assert_called_once_with( - graph_init_params=graph_init_params, - graph_runtime_state=child_graph_runtime_state, - ) - graph_init.assert_called_once_with( - graph_config={"nodes": [{"id": "root"}]}, - node_factory=sentinel.factory, - root_node_id="root", - ) - graph_engine_cls.assert_called_once_with( - workflow_id="workflow-id", - graph=child_graph, - graph_runtime_state=child_graph_runtime_state, - command_channel=sentinel.command_channel, - config=sentinel.graph_engine_config, - child_engine_builder=builder, - ) - llm_quota_layer_cls.assert_called_once_with(tenant_id="tenant-id") - assert child_engine.layer.call_args_list == [((sentinel.llm_quota_layer,), {})] - - @pytest.mark.parametrize("node_cls", [_FakeLLMNode, _FakeQuestionClassifierNode]) - def test_build_child_engine_runs_llm_quota_layer_for_child_model_nodes(self, node_cls): - builder = workflow_entry._WorkflowChildEngineBuilder(tenant_id="tenant-id") - graph_init_params = build_test_graph_init_params( - graph_config={"nodes": [{"id": "root"}], "edges": []}, - ) - parent_graph_runtime_state = SimpleNamespace( - execution_context=nullcontext(None), - variable_pool=build_test_variable_pool(), - ) - created_node: dict[str, _FakeLLMNode | _FakeQuestionClassifierNode] = {} - - def build_graph(*, graph_config, node_factory, root_node_id): - _ = graph_config - node_data = _build_llm_node_data() if node_cls is _FakeLLMNode else _build_question_classifier_node_data() - node = node_cls( - node_id=root_node_id, - data=node_data, - graph_init_params=node_factory.graph_init_params, - graph_runtime_state=node_factory.graph_runtime_state, - ) - created_node["node"] = node - return Graph( - nodes={root_node_id: node}, - edges={}, - in_edges={}, - out_edges={}, - root_node=node, - ) - - with ( - patch.object( - workflow_entry, - "DifyNodeFactory", - side_effect=lambda graph_init_params, graph_runtime_state: SimpleNamespace( - graph_init_params=graph_init_params, - graph_runtime_state=graph_runtime_state, - ), - ), - patch.object(workflow_entry.Graph, "init", side_effect=build_graph), - patch("core.app.workflow.layers.llm_quota.ensure_llm_quota_available_for_model") as ensure_quota, - patch("core.app.workflow.layers.llm_quota.deduct_llm_quota_for_model") as deduct_quota, - ): - child_engine = builder.build_child_engine( - workflow_id="workflow-id", - graph_init_params=graph_init_params, - parent_graph_runtime_state=parent_graph_runtime_state, - root_node_id="root", - ) - list(child_engine.run()) - - node = created_node["node"] - ensure_quota.assert_called_once_with( - tenant_id="tenant-id", - provider=node.node_data.model.provider, - model=node.node_data.model.name, - ) - deduct_quota.assert_called_once_with( - tenant_id="tenant-id", - provider=node.node_data.model.provider, - model=node.node_data.model.name, - usage=node.usage_snapshot, - ) - - def _build_minimal_workflow_entry( monkeypatch: pytest.MonkeyPatch, *, @@ -249,7 +30,7 @@ def _build_minimal_workflow_entry( ) -> workflow_entry.WorkflowEntry: """Construct a minimal WorkflowEntry with GraphEngine construction mocked out.""" graph_engine = MagicMock() - graph_runtime_state = SimpleNamespace(execution_context=None) + graph_runtime_state = SimpleNamespace(_execution_context=None) monkeypatch.setattr(workflow_entry, "capture_current_context", lambda: sentinel.execution_context) monkeypatch.setattr(workflow_entry, "GraphEngine", MagicMock(return_value=graph_engine)) @@ -294,7 +75,7 @@ class TestWorkflowEntryInit: def test_applies_debug_and_observability_layers(self): graph_engine = MagicMock() - graph_runtime_state = SimpleNamespace(execution_context=None) + graph_runtime_state = SimpleNamespace(_execution_context=None) debug_layer = sentinel.debug_layer execution_limits_layer = sentinel.execution_limits_layer llm_quota_layer = sentinel.llm_quota_layer @@ -339,9 +120,8 @@ class TestWorkflowEntryInit: graph_runtime_state=graph_runtime_state, command_channel=sentinel.command_channel, config=sentinel.graph_engine_config, - child_engine_builder=entry._child_engine_builder, ) - assert graph_runtime_state.execution_context is sentinel.execution_context + assert graph_runtime_state._execution_context is sentinel.execution_context debug_logging_layer.assert_called_once_with( level="DEBUG", include_inputs=True, @@ -443,6 +223,21 @@ class TestWorkflowEntryRun: class TestWorkflowEntrySingleStepRun: + @pytest.mark.parametrize("node_type", [BuiltinNodeTypes.LOOP, BuiltinNodeTypes.ITERATION]) + def test_rejects_container_nodes(self, node_type): + workflow = SimpleNamespace( + get_node_config_by_id=lambda _node_id: _build_typed_node_config(node_type), + ) + + with pytest.raises(ValueError, match="engine-backed debug endpoints"): + workflow_entry.WorkflowEntry.single_step_run( + workflow=workflow, + node_id="node-id", + user_id="user-id", + user_inputs={}, + variable_pool=sentinel.variable_pool, + ) + def test_preloads_constructor_variables_before_creating_memory_node(self): class FakeLLMNode: id = "node-id" @@ -958,7 +753,7 @@ class TestWorkflowEntryTracing: layer = MagicMock() class FakeNode: - def ensure_execution_id(self): + def bind_execution_id(self, _execution_id): return None def run(self): @@ -979,7 +774,7 @@ class TestWorkflowEntryTracing: layer = MagicMock() class FakeNode: - def ensure_execution_id(self): + def bind_execution_id(self, _execution_id): return None def run(self): diff --git a/api/tests/unit_tests/libs/broadcast_channel/redis/test_channel_unit_tests.py b/api/tests/unit_tests/libs/broadcast_channel/redis/test_channel_unit_tests.py index b74d494134b..b323e4c950f 100644 --- a/api/tests/unit_tests/libs/broadcast_channel/redis/test_channel_unit_tests.py +++ b/api/tests/unit_tests/libs/broadcast_channel/redis/test_channel_unit_tests.py @@ -27,6 +27,7 @@ from libs.broadcast_channel.redis.sharded_channel import ( ShardedTopic, _RedisShardedSubscription, ) +from libs.broadcast_channel.signals import SIG_CLOSE class TestBroadcastChannel: @@ -1239,6 +1240,30 @@ class TestRedisSubscriptionCommon: subscription_type, _ = subscription_params assert subscription._get_subscription_type() == subscription_type + def test_listener_ignores_close_signal_from_another_subscription(self, subscription, subscription_params): + subscription_type, _ = subscription_params + topic = f"test-{subscription_type}-topic" + message_type = "message" if subscription_type == "regular" else "smessage" + messages = iter( + [ + {"type": message_type, "channel": topic, "data": SIG_CLOSE}, + {"type": message_type, "channel": topic, "data": b"next-event"}, + ] + ) + + def get_message(): + try: + return next(messages) + except StopIteration: + subscription._closed.set() + return None + + subscription._get_message = get_message + subscription._listen() + + assert subscription._queue.get_nowait() == b"next-event" + assert subscription._queue.empty() + # ==================== Lifecycle Tests ==================== def test_start_if_needed_first_call(self, subscription, subscription_params, mock_pubsub: MagicMock): diff --git a/api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel_unit_tests.py b/api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel_unit_tests.py index 9a8cb861abf..5022f572153 100644 --- a/api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel_unit_tests.py +++ b/api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel_unit_tests.py @@ -12,6 +12,7 @@ from libs.broadcast_channel.redis.streams_channel import ( StreamsTopic, _StreamsSubscription, ) +from libs.broadcast_channel.signals import SIG_CLOSE class FakeStreamsRedis: @@ -282,6 +283,34 @@ class TestStreamsSubscription: assert received == case.expected_messages + def test_listener_ignores_close_signal_from_another_subscription(self): + class OneShotRedis: + def __init__(self) -> None: + self._calls = 0 + + def xread(self, streams: dict[str, Any], block: int | None = None, count: int | None = None): + self._calls += 1 + if self._calls == 1: + key = next(iter(streams)) + return [ + ( + key, + [ + ("1-0", {b"data": SIG_CLOSE}), + ("2-0", {b"data": b"next-event"}), + ], + ) + ] + subscription._closed = True + return [] + + subscription = _StreamsSubscription(OneShotRedis(), "stream:close-signal") + subscription._listen() + + assert subscription._queue.get_nowait() == b"next-event" + assert subscription._queue.get_nowait() is subscription._SENTINEL + assert subscription._queue.empty() + def test_iterator_yields_messages_until_subscription_is_closed(self, streams_channel: StreamsBroadcastChannel): topic = streams_channel.topic("iter") subscription = topic.subscribe() diff --git a/api/tests/unit_tests/services/test_model_load_balancing_service.py b/api/tests/unit_tests/services/test_model_load_balancing_service.py index 743e6e797a3..e9a05dce6b1 100644 --- a/api/tests/unit_tests/services/test_model_load_balancing_service.py +++ b/api/tests/unit_tests/services/test_model_load_balancing_service.py @@ -123,12 +123,11 @@ def test_enable_disable_model_load_balancing_uses_model_type_constructor_directl method_name: str, expected_provider_method: str, service: ModelLoadBalancingService, - monkeypatch: pytest.MonkeyPatch, ) -> None: provider_configuration = _build_provider_configuration(provider_schema=_build_provider_credential_schema()) service.provider_manager.get_configurations.return_value = {"openai": provider_configuration} - getattr(service, method_name)("tenant-1", "openai", "gpt-4o-mini", "text-generation") + getattr(service, method_name)("tenant-1", "openai", "gpt-4o-mini", "llm") getattr(provider_configuration, expected_provider_method).assert_called_once_with( model="gpt-4o-mini", model_type=ModelType.LLM diff --git a/api/tests/unit_tests/services/test_model_provider_service.py b/api/tests/unit_tests/services/test_model_provider_service.py index a8a976f4b07..12d9404ff79 100644 --- a/api/tests/unit_tests/services/test_model_provider_service.py +++ b/api/tests/unit_tests/services/test_model_provider_service.py @@ -377,7 +377,7 @@ class TestModelProviderServiceDelegation: { "tenant_id": "tenant-1", "provider": "openai", - "model_type": "text-generation", + "model_type": "llm", "model": "gpt-4o", "credential_id": "cred-1", }, @@ -389,7 +389,7 @@ class TestModelProviderServiceDelegation: { "tenant_id": "tenant-1", "provider": "openai", - "model_type": "text-generation", + "model_type": "llm", "model": "gpt-4o", "credentials": {"api_key": "x"}, "credential_name": "cred-a", @@ -407,7 +407,7 @@ class TestModelProviderServiceDelegation: { "tenant_id": "tenant-1", "provider": "openai", - "model_type": "text-generation", + "model_type": "llm", "model": "gpt-4o", }, "delete_custom_model", diff --git a/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service.py b/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service.py index 80bc90b72b9..63aa53c0bf6 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service.py @@ -144,8 +144,7 @@ def _build_resumption_context(task_id: str, *, select_options: list[str] | None runtime_state = GraphRuntimeState(variable_pool=VariablePool(), start_at=0.0) if select_options is not None: runtime_state.variable_pool.add(("start", "options"), select_options) - runtime_state.register_paused_node("node-1") - runtime_state.outputs = {"result": "value"} + runtime_state.set_output("result", "value") wrapper = _WorkflowGenerateEntityWrapper(entity=generate_entity) return WorkflowResumptionContext( generate_entity=wrapper, @@ -250,7 +249,7 @@ def _build_resumption_context_additional(task_id: str) -> WorkflowResumptionCont workflow_execution_id="run-1", ) runtime_state = GraphRuntimeState(variable_pool=VariablePool(), start_at=0.0) - runtime_state.outputs = {"answer": "ok"} + runtime_state.set_output("answer", "ok") wrapper = _WorkflowGenerateEntityWrapper(entity=generate_entity) return WorkflowResumptionContext( generate_entity=wrapper, diff --git a/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py b/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py index 8efd7370a73..924d272f686 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service_additional.py @@ -74,7 +74,7 @@ def _build_resumption_context(task_id: str) -> WorkflowResumptionContext: workflow_execution_id="run-1", ) runtime_state = GraphRuntimeState(variable_pool=VariablePool(), start_at=0.0) - runtime_state.outputs = {"answer": "ok"} + runtime_state.set_output("answer", "ok") wrapper = _WorkflowGenerateEntityWrapper(entity=generate_entity) return WorkflowResumptionContext( generate_entity=wrapper, diff --git a/api/uv.lock b/api/uv.lock index c5214c70e40..cc0061bda8c 100644 --- a/api/uv.lock +++ b/api/uv.lock @@ -1639,7 +1639,7 @@ requires-dist = [ { name = "gmpy2", specifier = ">=2.3.0,<3.0.0" }, { name = "google-api-python-client", specifier = ">=2.198.0,<3.0.0" }, { name = "google-cloud-aiplatform", specifier = ">=1.160.0,<2.0.0" }, - { name = "graphon", specifier = "==0.6.0" }, + { name = "graphon", specifier = "==0.7.0" }, { name = "gunicorn", specifier = ">=26.0.0,<27.0.0" }, { name = "httpx", extras = ["socks"], specifier = "==0.28.1" }, { name = "httpx-sse", specifier = "==0.4.3" }, @@ -2992,7 +2992,7 @@ httpx = [ [[package]] name = "graphon" -version = "0.6.0" +version = "0.7.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "charset-normalizer" }, @@ -3014,9 +3014,9 @@ dependencies = [ { name = "unstructured", extra = ["docx", "epub", "md", "ppt", "pptx"] }, { name = "webvtt-py" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ee/6c/9ea051ed30dc3306e9e77c4486b5a2e5462af45e35daf230d9ec886eb07e/graphon-0.6.0.tar.gz", hash = "sha256:2d3a386899dc7ab8e9767ab96c694ff7e6eb454c045a1e801505cba9c615160d", size = 264404, upload-time = "2026-06-29T15:26:27.437Z" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/c6/6f16398e28bceb11304f8dd743d5f4f2c1e1aedf20eef6531764edb40a7e/graphon-0.7.0.tar.gz", hash = "sha256:e3f19284432d6e4a947fb285ef36bc79eb48c41af36867b03c0d70f062fe8563", size = 266715, upload-time = "2026-07-29T10:00:51.303Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/32/36/5b0ece2d61fa091f7d74aea5a48f9cf927202099e8d22b856e0319218e9d/graphon-0.6.0-py3-none-any.whl", hash = "sha256:f1445ccef40c0d0eb50a60af85c1028b26d2d782f357315047acee816acbcb33", size = 376038, upload-time = "2026-06-29T15:26:26.186Z" }, + { url = "https://files.pythonhosted.org/packages/9c/9b/3b02d944a20b3693bb5bb77cbb16b3afed44a965096970c47625560c0235/graphon-0.7.0-py3-none-any.whl", hash = "sha256:6f7608aaaa65607c4935735137ed4b7029b1e6ca80202cbe829af3370cf4404e", size = 380138, upload-time = "2026-07-29T10:00:49.654Z" }, ] [[package]] diff --git a/docker/envs/core-services/shared.env.example b/docker/envs/core-services/shared.env.example index e68545d1a71..c0fed86f899 100644 --- a/docker/envs/core-services/shared.env.example +++ b/docker/envs/core-services/shared.env.example @@ -206,7 +206,7 @@ WORKFLOW_GENERATION_TIMEOUT_MS=180000 WORKFLOW_FILE_UPLOAD_LIMIT=10 GRAPH_ENGINE_MIN_WORKERS=3 GRAPH_ENGINE_MAX_WORKERS=10 -GRAPH_ENGINE_SCALE_UP_THRESHOLD=3 +GRAPH_ENGINE_SCALE_UP_THRESHOLD=0 GRAPH_ENGINE_SCALE_DOWN_IDLE_TIME=5.0 ALIYUN_SLS_ACCESS_KEY_ID= ALIYUN_SLS_ACCESS_KEY_SECRET= diff --git a/oxlint-suppressions.json b/oxlint-suppressions.json index f42c0310ed5..1dc6879617d 100644 --- a/oxlint-suppressions.json +++ b/oxlint-suppressions.json @@ -5385,11 +5385,6 @@ "count": 1 } }, - "web/app/components/workflow/panel/human-input-form-list.tsx": { - "typescript/no-explicit-any": { - "count": 1 - } - }, "web/app/components/workflow/panel/inputs-panel.tsx": { "jsx_a11y/no-autofocus": { "count": 1 diff --git a/web/app/components/base/chat/chat-with-history/__tests__/chat-wrapper.spec.tsx b/web/app/components/base/chat/chat-with-history/__tests__/chat-wrapper.spec.tsx index 680915c7138..ce5bce5bb19 100644 --- a/web/app/components/base/chat/chat-with-history/__tests__/chat-wrapper.spec.tsx +++ b/web/app/components/base/chat/chat-with-history/__tests__/chat-wrapper.spec.tsx @@ -3,7 +3,7 @@ import type { ChatWithHistoryContextValue } from '../context' import type { FileEntity } from '@/app/components/base/file-uploader/types' import type { AppData, AppMeta, ConversationItem } from '@/models/share' import type { HumanInputFormData } from '@/types/workflow' -import { fireEvent, render, screen, waitFor } from '@testing-library/react' +import { act, fireEvent, render, screen, waitFor } from '@testing-library/react' import { InputVarType } from '@/app/components/workflow/types' import { fetchChatList, @@ -144,6 +144,7 @@ const defaultChatHookReturn: Partial = { handleSend: vi.fn(), handleStop: vi.fn(), handleSwitchSibling: vi.fn(), + prepareHumanInputSubmission: vi.fn().mockResolvedValue(true), isResponding: false, suggestedQuestions: [], } @@ -910,6 +911,13 @@ describe('ChatWrapper', () => { it('should handle human input form submission for installed app', async () => { const { submitHumanInputForm: submitWorkflowForm } = await import('@/service/workflow') vi.mocked(submitWorkflowForm).mockResolvedValue({} as unknown as void) + let resolveWorkflowEventsReady: (isReady: boolean) => void = () => {} + const prepareHumanInputSubmission = vi.fn( + () => + new Promise((resolve) => { + resolveWorkflowEventsReady = resolve + }), + ) vi.mocked(useChatWithHistoryContext).mockReturnValue({ ...defaultContextValue, @@ -918,6 +926,7 @@ describe('ChatWrapper', () => { vi.mocked(useChat).mockReturnValue({ ...defaultChatHookReturn, + prepareHumanInputSubmission, chatList: [ { id: 'q1', content: 'Question' }, { @@ -961,6 +970,12 @@ describe('ChatWrapper', () => { const runButton = screen.getByText('Run') fireEvent.click(runButton) + expect(prepareHumanInputSubmission).toHaveBeenCalledOnce() + expect(submitWorkflowForm).not.toHaveBeenCalled() + + await act(async () => { + resolveWorkflowEventsReady(true) + }) await waitFor(() => { expect(submitWorkflowForm).toHaveBeenCalled() }) diff --git a/web/app/components/base/chat/chat-with-history/chat-wrapper.tsx b/web/app/components/base/chat/chat-with-history/chat-wrapper.tsx index 3fc3eb69aa9..7174cf9041a 100644 --- a/web/app/components/base/chat/chat-with-history/chat-wrapper.tsx +++ b/web/app/components/base/chat/chat-with-history/chat-wrapper.tsx @@ -86,6 +86,7 @@ const ChatWrapper = () => { handleSend, handleStop, handleSwitchSibling, + prepareHumanInputSubmission, isResponding: respondingState, suggestedQuestions, } = useChat( @@ -284,10 +285,12 @@ const ChatWrapper = () => { const handleSubmitHumanInputForm = useCallback( async (formToken: string, formData: any) => { + if (!(await prepareHumanInputSubmission())) return + if (isInstalledApp) await submitHumanInputFormService(formToken, formData) else await submitHumanInputForm(formToken, formData) }, - [isInstalledApp], + [isInstalledApp, prepareHumanInputSubmission], ) const [collapsed, setCollapsed] = useState(!!currentConversationId) diff --git a/web/app/components/base/chat/chat/__tests__/hooks.spec.tsx b/web/app/components/base/chat/chat/__tests__/hooks.spec.tsx index 354978bde8e..0a1f90dc7f9 100644 --- a/web/app/components/base/chat/chat/__tests__/hooks.spec.tsx +++ b/web/app/components/base/chat/chat/__tests__/hooks.spec.tsx @@ -723,6 +723,151 @@ describe('useChat', () => { expect(result.current.isResponding).toBe(true) }) + it('should only allow submission after the continuation stream observes the pause', async () => { + let postCallbacks: HookCallbacks + let continuationCallbacks: HookCallbacks + vi.mocked(ssePost).mockImplementation(async (_url, _params, options) => { + postCallbacks = options as HookCallbacks + }) + vi.mocked(sseGet).mockImplementation(async (_url, _params, options) => { + continuationCallbacks = options as HookCallbacks + }) + + const { result } = renderHook(() => useChat()) + act(() => { + result.current.handleSend('test-url', { query: 'human input test' }, {}) + postCallbacks.onWorkflowStarted({ workflow_run_id: 'wr-1', task_id: 't-1' }) + postCallbacks.onHumanInputRequired({ + workflow_run_id: 'wr-1', + data: { node_id: 'human-1' }, + }) + }) + expect(sseGet).not.toHaveBeenCalled() + + let isReady: boolean | undefined + const readyPromise = result.current + .prepareHumanInputSubmission() + .then((ready) => (isReady = ready)) + await act(async () => Promise.resolve()) + expect(isReady).toBeUndefined() + + act(() => { + postCallbacks.onWorkflowPaused({ data: { workflow_run_id: 'wr-1' } }) + }) + expect(sseGet).toHaveBeenCalledWith( + '/workflow/wr-1/events?include_state_snapshot=true&continue_on_pause=true', + expect.any(Object), + expect.any(Object), + ) + await act(async () => Promise.resolve()) + expect(isReady).toBeUndefined() + + act(() => { + continuationCallbacks.onWorkflowPaused({ data: { workflow_run_id: 'wr-1' } }) + }) + await act(async () => readyPromise) + expect(isReady).toBe(true) + }) + + it('should register a new paused conversation without running final completion twice', async () => { + let postCallbacks: HookCallbacks + let continuationCallbacks: HookCallbacks + vi.mocked(ssePost).mockImplementation(async (_url, _params, options) => { + postCallbacks = options as HookCallbacks + }) + vi.mocked(sseGet).mockImplementation(async (_url, _params, options) => { + continuationCallbacks = options as HookCallbacks + }) + const onConversationComplete = vi.fn() + const onGetConversationMessages = vi.fn().mockResolvedValue({ + data: [ + { + id: 'm-1', + answer: 'completed answer', + message: [], + workflow_run_id: 'wr-1', + inputs: {}, + query: 'human input test', + }, + ], + }) + const onGetSuggestedQuestions = vi.fn().mockResolvedValue({ data: ['Next question'] }) + const config = { suggested_questions_after_answer: { enabled: true } } + + const { result } = renderHook(() => useChat(config as ChatConfig)) + act(() => { + result.current.handleSend( + 'test-url', + { query: 'human input test' }, + { + onConversationComplete, + onGetConversationMessages, + onGetSuggestedQuestions, + }, + ) + postCallbacks.onWorkflowStarted({ + workflow_run_id: 'wr-1', + task_id: 't-1', + conversation_id: 'c-1', + message_id: 'm-1', + }) + postCallbacks.onHumanInputRequired({ + workflow_run_id: 'wr-1', + data: { node_id: 'human-1' }, + }) + postCallbacks.onWorkflowPaused({ data: { workflow_run_id: 'wr-1' } }) + }) + + await act(async () => { + await postCallbacks.onCompleted() + }) + + expect(onConversationComplete).toHaveBeenCalledOnce() + expect(onConversationComplete).toHaveBeenCalledWith('c-1', 'wr-1') + expect(onGetConversationMessages).not.toHaveBeenCalled() + expect(onGetSuggestedQuestions).not.toHaveBeenCalled() + + await act(async () => { + continuationCallbacks.onWorkflowPaused({ data: { workflow_run_id: 'wr-1' } }) + continuationCallbacks.onWorkflowFinished({ data: { status: 'succeeded' } }) + await continuationCallbacks.onCompleted() + }) + + expect(onConversationComplete).toHaveBeenCalledOnce() + expect(onGetConversationMessages).toHaveBeenCalledOnce() + expect(onGetSuggestedQuestions).toHaveBeenCalledOnce() + }) + + it('should reject a pending form submission if the initial stream fails before pausing', async () => { + let postCallbacks: HookCallbacks + vi.mocked(ssePost).mockImplementation(async (_url, _params, options) => { + postCallbacks = options as HookCallbacks + }) + + const { result } = renderHook(() => useChat()) + act(() => { + result.current.handleSend('test-url', { query: 'human input test' }, {}) + postCallbacks.onWorkflowStarted({ workflow_run_id: 'wr-1', task_id: 't-1' }) + postCallbacks.onNodeStarted({ data: { node_id: 'human-1', id: 'human-1' } }) + postCallbacks.onHumanInputRequired({ + workflow_run_id: 'wr-1', + data: { node_id: 'human-1' }, + }) + }) + + const readyPromise = result.current.prepareHumanInputSubmission() + await act(async () => Promise.resolve()) + expect(result.current.chatList[1]!.humanInputFormDataList).toHaveLength(1) + + act(() => { + postCallbacks.onError('stream failed') + }) + + await expect(readyPromise).resolves.toBe(false) + expect(result.current.chatList[1]!.humanInputFormDataList).toHaveLength(0) + expect(result.current.isResponding).toBe(false) + }) + it('should handle file uploads in onFile', () => { let callbacks: HookCallbacks @@ -1344,7 +1489,7 @@ describe('useChat', () => { }) expect(sseGet).toHaveBeenCalledWith( - '/workflow/wr-1/events?include_state_snapshot=true', + '/workflow/wr-1/events?include_state_snapshot=true&continue_on_pause=true', expect.any(Object), expect.any(Object), ) @@ -1397,6 +1542,7 @@ describe('useChat', () => { }) callbacks.onMessageReplace({ answer: 'replaced resume' }) + callbacks.onWorkflowPaused({ data: { workflow_run_id: 'wr-1' } }) callbacks.onWorkflowPaused({ data: { workflow_run_id: 'wr-1' } }) callbacks.onError() @@ -1414,6 +1560,164 @@ describe('useChat', () => { expect(lastResponse!.humanInputFilledFormDataList).toHaveLength(1) expect(lastResponse!.humanInputFormDataList).toHaveLength(0) expect(lastResponse!.content).toBe('replaced resume') + expect(sseGet).toHaveBeenCalledTimes(1) + }) + + it('should wait for the resumed event stream before allowing a restored form submission', async () => { + let callbacks: HookCallbacks + vi.mocked(sseGet).mockImplementation(async (_url, _params, options) => { + callbacks = options as HookCallbacks + }) + + const prevChatTree = [ + { + id: 'q-1', + content: 'query', + isAnswer: false, + children: [ + { + id: 'm-1', + content: '', + isAnswer: true, + workflow_run_id: 'wr-1', + humanInputFormDataList: [{ node_id: 'human-1' }], + workflowProcess: { status: WorkflowRunningStatus.Paused, tracing: [] }, + siblingIndex: 0, + }, + ], + }, + ] + const { result } = renderHook(() => + useChat(undefined, undefined, prevChatTree as unknown as ChatItemInTree[]), + ) + + let isReady: boolean | undefined + const readyPromise = result.current + .prepareHumanInputSubmission() + .then((ready) => (isReady = ready)) + + await act(async () => Promise.resolve()) + expect(isReady).toBeUndefined() + + act(() => { + result.current.handleSwitchSibling('m-1', { isPublicAPI: true }) + }) + expect(sseGet).toHaveBeenCalledWith( + '/workflow/wr-1/events?include_state_snapshot=true&continue_on_pause=true', + expect.any(Object), + expect.any(Object), + ) + + act(() => { + callbacks.onWorkflowPaused({ data: { workflow_run_id: 'wr-1' } }) + }) + await act(async () => readyPromise) + + expect(isReady).toBe(true) + expect(sseGet).toHaveBeenCalledTimes(1) + }) + + it('should reconnect an idle paused event stream before the next submission', async () => { + const callbacksList: HookCallbacks[] = [] + const onConversationComplete = vi.fn() + vi.mocked(sseGet).mockImplementation(async (_url, _params, options) => { + callbacksList.push(options as HookCallbacks) + }) + + const prevChatTree = [ + { + id: 'q-1', + content: 'query', + isAnswer: false, + children: [ + { + id: 'm-1', + content: '', + isAnswer: true, + workflow_run_id: 'wr-1', + humanInputFormDataList: [{ node_id: 'human-1' }], + workflowProcess: { status: WorkflowRunningStatus.Paused, tracing: [] }, + siblingIndex: 0, + }, + ], + }, + ] + const { result } = renderHook(() => + useChat(undefined, undefined, prevChatTree as unknown as ChatItemInTree[]), + ) + + act(() => { + result.current.handleResume('m-1', 'wr-1', { + isPublicAPI: true, + onConversationComplete, + }) + callbacksList[0]!.onWorkflowPaused({ data: { workflow_run_id: 'wr-1' } }) + }) + await act(async () => { + await callbacksList[0]!.onCompleted() + }) + expect(onConversationComplete).not.toHaveBeenCalled() + + let isReady: boolean | undefined + const readyPromise = result.current + .prepareHumanInputSubmission() + .then((ready) => (isReady = ready)) + expect(sseGet).toHaveBeenCalledTimes(2) + await act(async () => Promise.resolve()) + expect(isReady).toBeUndefined() + + act(() => { + callbacksList[1]!.onWorkflowPaused({ data: { workflow_run_id: 'wr-1' } }) + }) + await act(async () => readyPromise) + + expect(isReady).toBe(true) + expect(onConversationComplete).not.toHaveBeenCalled() + }) + + it('should reconnect immediately if the event stream idles after submission resumes the run', async () => { + const callbacksList: HookCallbacks[] = [] + vi.mocked(sseGet).mockImplementation(async (_url, _params, options) => { + callbacksList.push(options as HookCallbacks) + }) + + const prevChatTree = [ + { + id: 'q-1', + content: 'query', + isAnswer: false, + children: [ + { + id: 'm-1', + content: '', + isAnswer: true, + workflow_run_id: 'wr-1', + humanInputFormDataList: [{ node_id: 'human-1' }], + workflowProcess: { status: WorkflowRunningStatus.Paused, tracing: [] }, + siblingIndex: 0, + }, + ], + }, + ] + const { result } = renderHook(() => + useChat(undefined, undefined, prevChatTree as unknown as ChatItemInTree[]), + ) + + act(() => { + result.current.handleResume('m-1', 'wr-1', { isPublicAPI: true }) + callbacksList[0]!.onWorkflowPaused({ data: { workflow_run_id: 'wr-1' } }) + }) + await act(async () => { + await result.current.prepareHumanInputSubmission() + await callbacksList[0]!.onCompleted() + }) + + expect(sseGet).toHaveBeenCalledTimes(2) + expect(sseGet).toHaveBeenLastCalledWith( + '/workflow/wr-1/events?include_state_snapshot=true&continue_on_pause=true', + expect.any(Object), + expect.any(Object), + ) }) it('should handle non-agent mode resume', async () => { @@ -1674,6 +1978,7 @@ describe('useChat', () => { conversationId: 'c-resume', taskId: 't-resume', }) + callbacks.onWorkflowFinished({ data: { status: 'succeeded' } }) await callbacks.onCompleted() }) @@ -1874,7 +2179,7 @@ describe('useChat', () => { }) expect(sseGet).toHaveBeenCalledWith( - '/workflow/wr-tts-app/events?include_state_snapshot=true', + '/workflow/wr-tts-app/events?include_state_snapshot=true&continue_on_pause=true', expect.any(Object), expect.any(Object), ) @@ -2004,6 +2309,38 @@ describe('useChat', () => { expect(suggestedAbort.abort).toHaveBeenCalledTimes(1) }) + it('should mark an unmounted continuation stream as an intentional abort', () => { + let callbacks: HookCallbacks + vi.mocked(sseGet).mockImplementation(async (_url, _params, options) => { + callbacks = options as HookCallbacks + }) + const workflowAbort = createAbortControllerMock() + const prevChatTree = [ + { + id: 'q-1', + content: 'query', + isAnswer: false, + children: [{ id: 'm-1', content: '', isAnswer: true, siblingIndex: 0 }], + }, + ] + const { result, unmount } = renderHook(() => + useChat(undefined, undefined, prevChatTree as ChatItemInTree[]), + ) + + act(() => { + result.current.handleResume('m-1', 'wr-1', { isPublicAPI: true }) + callbacks.getAbortController(workflowAbort) + }) + unmount() + + expect(workflowAbort.abort).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'AbortError', + message: 'The user aborted a request.', + }), + ) + }) + it('should clear chat list when clearChatList flag is true and reset flag via callback', () => { const clearChatListCallback = vi.fn() @@ -2108,7 +2445,7 @@ describe('useChat', () => { // Should automatically call handleResume -> sseGet for human input expect(sseGet).toHaveBeenCalledWith( - '/workflow/wr-1/events?include_state_snapshot=true', + '/workflow/wr-1/events?include_state_snapshot=true&continue_on_pause=true', expect.any(Object), expect.any(Object), ) @@ -3141,6 +3478,7 @@ describe('useChat', () => { }) await act(async () => { + resumeCallbacks.onWorkflowFinished({ data: { status: 'succeeded' } }) await resumeCallbacks.onCompleted() }) expect(result.current.suggestedQuestions).toEqual(['Suggested 1', 'Suggested 2']) diff --git a/web/app/components/base/chat/chat/hooks.ts b/web/app/components/base/chat/chat/hooks.ts index 13a7c70613c..408cb44ede1 100644 --- a/web/app/components/base/chat/chat/hooks.ts +++ b/web/app/components/base/chat/chat/hooks.ts @@ -180,6 +180,10 @@ function getConversationMessagesData(response: unknown): ConversationMessagesRes return Array.isArray(data) ? data.filter(isHistoryConversationMessage) : [] } +function abortWorkflowEventsRequest(abortController: AbortController | null) { + abortController?.abort(new DOMException('The user aborted a request.', 'AbortError')) +} + export const useChat = ( config?: ChatConfig, formSettings?: { @@ -206,6 +210,26 @@ export const useChat = ( const conversationMessagesAbortControllerRef = useRef(null) const suggestedQuestionsAbortControllerRef = useRef(null) const workflowEventsAbortControllerRef = useRef(null) + const pausedWorkflowEventsAbortControllerRef = useRef(null) + const pausedWorkflowEventsRef = useRef<{ + workflowRunId: string + options: IOtherOptions + } | null>(null) + const workflowEventsSubscriptionActiveRef = useRef(false) + const workflowEventsSubscriptionRunIdRef = useRef(null) + const workflowEventsSubscriptionGenerationRef = useRef(0) + const workflowRequestGenerationRef = useRef(0) + const workflowEventsReadyRef = useRef(false) + const workflowPauseConfirmedRef = useRef(false) + const workflowEventsReadyWaitersRef = useRef< + Array<{ + workflowRunId: string | null + resolve: (isReady: boolean) => void + }> + >([]) + const startWorkflowEventsSubscriptionRef = useRef< + ((workflowRunId: string, options: IOtherOptions) => void) | null + >(null) const params = useParams() const pathname = usePathname() @@ -331,6 +355,172 @@ export const useChat = ( isRespondingRef.current = isResponding }, []) + const resolveWorkflowEventsReadyWaiters = useCallback((isReady: boolean) => { + const waiters = workflowEventsReadyWaitersRef.current.splice(0) + waiters.forEach(({ resolve }) => resolve(isReady)) + }, []) + + const bindWorkflowEventsReadyWaiters = useCallback((workflowRunId: string) => { + workflowEventsReadyWaitersRef.current = workflowEventsReadyWaitersRef.current.filter( + (waiter) => { + if (waiter.workflowRunId && waiter.workflowRunId !== workflowRunId) { + waiter.resolve(false) + return false + } + + waiter.workflowRunId = workflowRunId + return true + }, + ) + }, []) + + const markWorkflowEventsPending = useCallback(() => { + workflowEventsReadyRef.current = false + }, []) + + const startWorkflowEventsSubscription = useCallback( + (workflowRunId: string, options: IOtherOptions) => { + const generation = ++workflowEventsSubscriptionGenerationRef.current + abortWorkflowEventsRequest(pausedWorkflowEventsAbortControllerRef.current) + pausedWorkflowEventsAbortControllerRef.current = null + pausedWorkflowEventsRef.current = { workflowRunId, options } + bindWorkflowEventsReadyWaiters(workflowRunId) + workflowEventsSubscriptionActiveRef.current = true + workflowEventsSubscriptionRunIdRef.current = workflowRunId + markWorkflowEventsPending() + + let hasWorkflowFinished = false + const releaseSubscription = () => { + if (generation !== workflowEventsSubscriptionGenerationRef.current) return false + + workflowEventsSubscriptionActiveRef.current = false + workflowEventsSubscriptionRunIdRef.current = null + pausedWorkflowEventsAbortControllerRef.current = null + return true + } + const subscriptionOptions: IOtherOptions = { + ...options, + getAbortController: (abortController) => { + if (generation !== workflowEventsSubscriptionGenerationRef.current) { + abortWorkflowEventsRequest(abortController) + return + } + pausedWorkflowEventsAbortControllerRef.current = abortController + }, + onHumanInputRequired: (event) => { + if (generation !== workflowEventsSubscriptionGenerationRef.current) return + options.onHumanInputRequired?.(event) + }, + onWorkflowFinished: (event) => { + if (generation !== workflowEventsSubscriptionGenerationRef.current) return + hasWorkflowFinished = true + options.onWorkflowFinished?.(event) + }, + onWorkflowPaused: (event) => { + if (generation !== workflowEventsSubscriptionGenerationRef.current) return + + options.onWorkflowPaused?.(event) + workflowEventsReadyRef.current = true + resolveWorkflowEventsReadyWaiters(true) + }, + onError: (...args) => { + if (!releaseSubscription()) return + + markWorkflowEventsPending() + resolveWorkflowEventsReadyWaiters(false) + options.onError?.(...args) + }, + async onCompleted(hasError?: boolean, errorMessage?: string) { + if (!releaseSubscription()) return + + markWorkflowEventsPending() + if (!hasWorkflowFinished) { + if (hasError) { + resolveWorkflowEventsReadyWaiters(false) + await options.onCompleted?.(hasError, errorMessage) + } else { + resolveWorkflowEventsReadyWaiters(false) + if (!workflowPauseConfirmedRef.current) + startWorkflowEventsSubscriptionRef.current?.(workflowRunId, options) + } + return + } + + workflowPauseConfirmedRef.current = false + pausedWorkflowEventsRef.current = null + resolveWorkflowEventsReadyWaiters(false) + await options.onCompleted?.(hasError, errorMessage) + }, + } + + void sseGet( + `/workflow/${workflowRunId}/events?include_state_snapshot=true&continue_on_pause=true`, + {}, + subscriptionOptions, + ) + }, + [bindWorkflowEventsReadyWaiters, markWorkflowEventsPending, resolveWorkflowEventsReadyWaiters], + ) + startWorkflowEventsSubscriptionRef.current = startWorkflowEventsSubscription + + const ensureWorkflowEventsSubscription = useCallback( + (workflowRunId: string, options: IOtherOptions) => { + pausedWorkflowEventsRef.current = { workflowRunId, options } + if ( + workflowEventsSubscriptionActiveRef.current && + workflowEventsSubscriptionRunIdRef.current === workflowRunId + ) + return + + startWorkflowEventsSubscription(workflowRunId, options) + }, + [startWorkflowEventsSubscription], + ) + + const prepareHumanInputSubmission = useCallback(async () => { + if (workflowEventsReadyRef.current) { + workflowPauseConfirmedRef.current = false + return true + } + + const isReady = await new Promise((resolve) => { + const pausedWorkflowEvents = pausedWorkflowEventsRef.current + workflowEventsReadyWaitersRef.current.push({ + workflowRunId: pausedWorkflowEvents?.workflowRunId ?? null, + resolve, + }) + if ( + pausedWorkflowEvents && + workflowPauseConfirmedRef.current && + !workflowEventsSubscriptionActiveRef.current + ) { + startWorkflowEventsSubscription( + pausedWorkflowEvents.workflowRunId, + pausedWorkflowEvents.options, + ) + } + }) + if (isReady) { + workflowPauseConfirmedRef.current = false + } + return isReady + }, [startWorkflowEventsSubscription]) + + const resetWorkflowEventsSubscription = useCallback(() => { + workflowRequestGenerationRef.current += 1 + workflowEventsSubscriptionGenerationRef.current += 1 + workflowEventsSubscriptionActiveRef.current = false + workflowEventsSubscriptionRunIdRef.current = null + workflowEventsReadyRef.current = false + workflowPauseConfirmedRef.current = false + pausedWorkflowEventsRef.current = null + abortWorkflowEventsRequest(pausedWorkflowEventsAbortControllerRef.current) + pausedWorkflowEventsAbortControllerRef.current = null + resolveWorkflowEventsReadyWaiters(false) + }, [resolveWorkflowEventsReadyWaiters]) + + useEffect(() => resetWorkflowEventsSubscription, [resetWorkflowEventsSubscription]) + const handleStop = useCallback(() => { hasStopRespondedRef.current = true handleResponding(false) @@ -340,7 +530,8 @@ export const useChat = ( if (suggestedQuestionsAbortControllerRef.current) suggestedQuestionsAbortControllerRef.current.abort() if (workflowEventsAbortControllerRef.current) workflowEventsAbortControllerRef.current.abort() - }, [stopChat, handleResponding]) + resetWorkflowEventsSubscription() + }, [stopChat, handleResponding, resetWorkflowEventsSubscription]) const handleRestart = useCallback( (cb?: any) => { @@ -389,6 +580,16 @@ export const useChat = ( workflowRunId: string, { onGetSuggestedQuestions, onConversationComplete, onSendSettled, isPublicAPI }: SendCallback, ) => { + const hasActiveSubscription = + workflowEventsSubscriptionActiveRef.current && + workflowEventsSubscriptionRunIdRef.current === workflowRunId + const requestGeneration = hasActiveSubscription + ? workflowRequestGenerationRef.current + : ++workflowRequestGenerationRef.current + if (!hasActiveSubscription) { + workflowEventsAbortControllerRef.current?.abort() + workflowEventsAbortControllerRef.current = null + } const getOrCreatePlayer = createAudioPlayerManager() let hasSettled = false const settleSend = (hasError?: boolean) => { @@ -397,9 +598,6 @@ export const useChat = ( hasSettled = true onSendSettled?.(hasError) } - // Re-subscribe to workflow events for the specific message - const url = `/workflow/${workflowRunId}/events?include_state_snapshot=true` - const otherOptions: IOtherOptions = { isPublicAPI, getAbortController: (abortController) => { @@ -440,6 +638,8 @@ export const useChat = ( }) }, async onCompleted(hasError?: boolean) { + if (requestGeneration !== workflowRequestGenerationRef.current) return + handleResponding(false) try { @@ -575,10 +775,14 @@ export const useChat = ( }) }, onError() { + if (requestGeneration !== workflowRequestGenerationRef.current) return + handleResponding(false) settleSend(true) }, onWorkflowStarted: ({ workflow_run_id, task_id }) => { + if (requestGeneration !== workflowRequestGenerationRef.current) return + handleResponding(true) hasStopRespondedRef.current = false updateChatTreeNode(messageId, (responseItem) => { @@ -599,6 +803,9 @@ export const useChat = ( }) }, onWorkflowFinished: ({ data: workflowFinishedData }) => { + if (requestGeneration !== workflowRequestGenerationRef.current) return + + pausedStateRef.current = false updateChatTreeNode(messageId, (responseItem) => { if (responseItem.workflowProcess) { responseItem.workflowProcess = { @@ -724,7 +931,18 @@ export const useChat = ( } }) }, - onHumanInputRequired: ({ data: humanInputRequiredData }) => { + onHumanInputRequired: ({ + workflow_run_id: pausedWorkflowRunId, + data: humanInputRequiredData, + }) => { + if (requestGeneration !== workflowRequestGenerationRef.current) return + + markWorkflowEventsPending() + workflowPauseConfirmedRef.current = false + pausedWorkflowEventsRef.current = { + workflowRunId: pausedWorkflowRunId || workflowRunId, + options: otherOptions, + } updateChatTreeNode(messageId, (responseItem) => { if (!responseItem.humanInputFormDataList) { responseItem.humanInputFormDataList = [humanInputRequiredData] @@ -749,6 +967,8 @@ export const useChat = ( }) }, onHumanInputFormFilled: ({ data: humanInputFilledFormData }) => { + workflowPauseConfirmedRef.current = false + handleResponding(true) updateChatTreeNode(messageId, (responseItem) => { let requiredFormData: | NonNullable[number] @@ -785,18 +1005,20 @@ export const useChat = ( }) }, onWorkflowPaused: ({ data: workflowPausedData }) => { - const resumeUrl = `/workflow/${workflowPausedData.workflow_run_id}/events` + if (requestGeneration !== workflowRequestGenerationRef.current) return + pausedStateRef.current = true - sseGet(resumeUrl, {}, otherOptions) + workflowPauseConfirmedRef.current = true + handleResponding(false) + ensureWorkflowEventsSubscription(workflowPausedData.workflow_run_id, otherOptions) updateChatTreeNode(messageId, (responseItem) => { responseItem.workflowProcess!.status = WorkflowRunningStatus.Paused }) }, } - if (workflowEventsAbortControllerRef.current) workflowEventsAbortControllerRef.current.abort() - - sseGet(url, {}, otherOptions) + workflowPauseConfirmedRef.current = true + ensureWorkflowEventsSubscription(workflowRunId, otherOptions) }, [ updateChatTreeNode, @@ -804,6 +1026,8 @@ export const useChat = ( createAudioPlayerManager, config?.suggested_questions_after_answer, options.isNewAgent, + ensureWorkflowEventsSubscription, + markWorkflowEventsPending, ], ) @@ -871,6 +1095,10 @@ export const useChat = ( return false } + pausedStateRef.current = false + resetWorkflowEventsSubscription() + const requestGeneration = ++workflowRequestGenerationRef.current + const parentMessage = threadMessages.find((item) => item.id === data.parent_message_id) const placeholderQuestionId = `question-${Date.now()}` @@ -938,12 +1166,21 @@ export const useChat = ( let isAgentMode = false let hasSetResponseId = false let hasSettled = false + let hasPaused = false + let hasNotifiedConversationComplete = false + let currentWorkflowRunId = '' const settleSend = (hasError?: boolean) => { if (hasSettled) return hasSettled = true onSendSettled?.(hasError) } + const notifyConversationComplete = (workflowRunId?: string) => { + if (hasNotifiedConversationComplete) return + + hasNotifiedConversationComplete = true + onConversationComplete?.(conversationIdRef.current, workflowRunId) + } const getOrCreatePlayer = createAudioPlayerManager() @@ -1005,6 +1242,8 @@ export const useChat = ( }) }, async onCompleted(hasError?: boolean) { + if (requestGeneration !== workflowRequestGenerationRef.current) return + handleResponding(false) try { @@ -1025,8 +1264,7 @@ export const useChat = ( const data = getConversationMessagesData(conversationMessagesResponse) const newResponseItem = data.find((item) => item.id === responseItem.id) completedWorkflowRunId = newResponseItem?.workflow_run_id ?? completedWorkflowRunId - if (!newResponseItem) - return onConversationComplete?.(conversationIdRef.current, completedWorkflowRunId) + if (!newResponseItem) return notifyConversationComplete(completedWorkflowRunId) const historyAgentThoughts = getHistoryAgentThoughts(newResponseItem) const lastHistoryAgentThought = historyAgentThoughts.at(-1) @@ -1082,7 +1320,7 @@ export const useChat = ( }) } - onConversationComplete?.(conversationIdRef.current, completedWorkflowRunId) + notifyConversationComplete(completedWorkflowRunId) if ( config?.suggested_questions_after_answer?.enabled && @@ -1236,6 +1474,8 @@ export const useChat = ( responseItem.content = messageReplace.answer }, onError() { + if (requestGeneration !== workflowRequestGenerationRef.current) return + handleResponding(false) settleSend(true) updateCurrentQAOnTree({ @@ -1246,6 +1486,9 @@ export const useChat = ( }) }, onWorkflowStarted: ({ workflow_run_id, task_id, conversation_id, message_id }) => { + if (requestGeneration !== workflowRequestGenerationRef.current) return + + currentWorkflowRunId = workflow_run_id handleResponding(true) // If there are no streaming messages, we still need to set the conversation_id to avoid create a new conversation when regeneration in chat-flow. if (conversation_id) { @@ -1280,6 +1523,8 @@ export const useChat = ( }) }, onWorkflowFinished: ({ data: workflowFinishedData }) => { + if (requestGeneration !== workflowRequestGenerationRef.current) return + if (pausedStateRef.current) pausedStateRef.current = false responseItem.workflowProcess = { ...responseItem.workflowProcess!, @@ -1425,7 +1670,18 @@ export const useChat = ( parentId: data.parent_message_id, }) }, - onHumanInputRequired: ({ data: humanInputRequiredData }) => { + onHumanInputRequired: ({ + workflow_run_id: pausedWorkflowRunId, + data: humanInputRequiredData, + }) => { + if (requestGeneration !== workflowRequestGenerationRef.current) return + + markWorkflowEventsPending() + workflowPauseConfirmedRef.current = false + pausedWorkflowEventsRef.current = { + workflowRunId: pausedWorkflowRunId || currentWorkflowRunId, + options: otherOptions, + } if (!responseItem.humanInputFormDataList) { responseItem.humanInputFormDataList = [humanInputRequiredData] } else { @@ -1453,6 +1709,8 @@ export const useChat = ( } }, onHumanInputFormFilled: ({ data: humanInputFilledFormData }) => { + workflowPauseConfirmedRef.current = false + handleResponding(true) let requiredFormData: NonNullable[number] | undefined if (responseItem.humanInputFormDataList?.length) { const currentFormIndex = responseItem.humanInputFormDataList!.findIndex( @@ -1495,9 +1753,13 @@ export const useChat = ( }) }, onWorkflowPaused: ({ data: workflowPausedData }) => { - const url = `/workflow/${workflowPausedData.workflow_run_id}/events` + if (requestGeneration !== workflowRequestGenerationRef.current) return + + hasPaused = true pausedStateRef.current = true - sseGet(url, {}, otherOptions) + workflowPauseConfirmedRef.current = true + handleResponding(false) + ensureWorkflowEventsSubscription(workflowPausedData.workflow_run_id, otherOptions) responseItem.workflowProcess!.status = WorkflowRunningStatus.Paused updateCurrentQAOnTree({ placeholderQuestionId, @@ -1511,12 +1773,34 @@ export const useChat = ( // Abort the previous workflow events SSE request if (workflowEventsAbortControllerRef.current) workflowEventsAbortControllerRef.current.abort() + const postOptions: IOtherOptions = { + ...otherOptions, + onError: (...args) => { + if (requestGeneration !== workflowRequestGenerationRef.current) return + + if (!hasPaused) { + markWorkflowEventsPending() + workflowPauseConfirmedRef.current = false + pausedWorkflowEventsRef.current = null + resolveWorkflowEventsReadyWaiters(false) + responseItem.humanInputFormDataList = [] + } + otherOptions.onError?.(...args) + }, + onCompleted: (hasError?: boolean, errorMessage?: string) => { + if (hasPaused && !hasError) { + notifyConversationComplete(currentWorkflowRunId) + return + } + return otherOptions.onCompleted?.(hasError, errorMessage) + }, + } ssePost( url, { body: bodyParams, }, - otherOptions, + postOptions, ) return true }, @@ -1532,6 +1816,10 @@ export const useChat = ( createAudioPlayerManager, formSettings, options.isNewAgent, + ensureWorkflowEventsSubscription, + markWorkflowEventsPending, + resetWorkflowEventsSubscription, + resolveWorkflowEventsReadyWaiters, ], ) @@ -1640,6 +1928,7 @@ export const useChat = ( handleSend, handleResume, handleSwitchSibling, + prepareHumanInputSubmission, suggestedQuestions, handleRestart, handleStop, diff --git a/web/app/components/base/chat/embedded-chatbot/__tests__/chat-wrapper.spec.tsx b/web/app/components/base/chat/embedded-chatbot/__tests__/chat-wrapper.spec.tsx index bdff3289b74..b403168104c 100644 --- a/web/app/components/base/chat/embedded-chatbot/__tests__/chat-wrapper.spec.tsx +++ b/web/app/components/base/chat/embedded-chatbot/__tests__/chat-wrapper.spec.tsx @@ -3,7 +3,7 @@ import type { HumanInputFieldValue } from '../../chat/answer/human-input-content import type { ChatConfig, ChatItem, ChatItemInTree } from '../../types' import type { EmbeddedChatbotContextValue } from '../context' import type { ConversationItem } from '@/models/share' -import { cleanup, fireEvent, render, screen, waitFor } from '@testing-library/react' +import { act, cleanup, fireEvent, render, screen, waitFor } from '@testing-library/react' import { InputVarType } from '@/app/components/workflow/types' import { AppSourceType, fetchSuggestedQuestions, submitHumanInputForm } from '@/service/share' import { submitHumanInputForm as submitHumanInputFormService } from '@/service/workflow' @@ -210,6 +210,7 @@ const createUseChatReturn = (overrides: Partial = {}): UseChatRet setIsResponding: vi.fn() as UseChatReturn['setIsResponding'], handleStop: vi.fn(), handleSwitchSibling: vi.fn(), + prepareHumanInputSubmission: vi.fn().mockResolvedValue(true), isResponding: false, suggestedQuestions: [], handleRestart: vi.fn(), @@ -457,6 +458,18 @@ describe('EmbeddedChatbot chat-wrapper', () => { describe('Human input submit behavior', () => { it('should submit via installed app service when the app is installed', async () => { + let resolveWorkflowEventsReady: (isReady: boolean) => void = () => {} + const prepareHumanInputSubmission = vi.fn( + () => + new Promise((resolve) => { + resolveWorkflowEventsReady = resolve + }), + ) + vi.mocked(useChat).mockReturnValue( + createUseChatReturn({ + prepareHumanInputSubmission, + }), + ) vi.mocked(useEmbeddedChatbotContext).mockReturnValue( createContextValue({ isInstalledApp: true, @@ -466,6 +479,12 @@ describe('EmbeddedChatbot chat-wrapper', () => { render() fireEvent.click(screen.getByRole('button', { name: 'submit human input' })) + expect(prepareHumanInputSubmission).toHaveBeenCalledOnce() + expect(submitHumanInputFormService).not.toHaveBeenCalled() + + await act(async () => { + resolveWorkflowEventsReady(true) + }) await waitFor(() => { expect(submitHumanInputFormService).toHaveBeenCalledWith('form-token', { inputs: { answer: 'ok' }, diff --git a/web/app/components/base/chat/embedded-chatbot/chat-wrapper.tsx b/web/app/components/base/chat/embedded-chatbot/chat-wrapper.tsx index b4be34faf3c..f2c364012c2 100644 --- a/web/app/components/base/chat/embedded-chatbot/chat-wrapper.tsx +++ b/web/app/components/base/chat/embedded-chatbot/chat-wrapper.tsx @@ -88,6 +88,7 @@ const ChatWrapper = () => { handleSend, handleStop, handleSwitchSibling, + prepareHumanInputSubmission, isResponding: respondingState, suggestedQuestions, } = useChat( @@ -344,10 +345,12 @@ const ChatWrapper = () => { const handleSubmitHumanInputForm = useCallback( async (formToken: string, formData: HumanInputFormSubmitData) => { + if (!(await prepareHumanInputSubmission())) return + if (isInstalledApp) await submitHumanInputFormService(formToken, formData) else await submitHumanInputForm(formToken, formData) }, - [isInstalledApp], + [isInstalledApp, prepareHumanInputSubmission], ) const welcome = useMemo(() => { diff --git a/web/app/components/share/text-generation/result/__tests__/workflow-stream-handlers.spec.ts b/web/app/components/share/text-generation/result/__tests__/workflow-stream-handlers.spec.ts index f8ad4bed78c..335ce562c2e 100644 --- a/web/app/components/share/text-generation/result/__tests__/workflow-stream-handlers.spec.ts +++ b/web/app/components/share/text-generation/result/__tests__/workflow-stream-handlers.spec.ts @@ -594,6 +594,17 @@ describe('createWorkflowStreamHandlers', () => { workflow_run_id: 'run-1', }, }) + handlers.onWorkflowPaused({ + task_id: 'task-1', + workflow_run_id: 'run-1', + event: 'workflow_paused', + data: { + outputs: {}, + paused_nodes: [], + reasons: [], + workflow_run_id: 'run-1', + }, + }) handlers.onWorkflowFinished({ task_id: 'task-1', workflow_run_id: 'run-1', @@ -627,16 +638,37 @@ describe('createWorkflowStreamHandlers', () => { }), ) expect(sseGetMock).toHaveBeenCalledWith( - '/workflow/run-1/events', + '/workflow/run-1/events?include_state_snapshot=true&continue_on_pause=true', {}, expect.objectContaining({ isPublicAPI: true }), ) + expect(sseGetMock).toHaveBeenCalledTimes(1) expect(setup.messageId()).toBe('run-1') expect(setup.onCompleted).toHaveBeenCalledWith('{"answer":"Hello"}', 3, true) expect(setup.setRespondingFalse).toHaveBeenCalled() expect(setup.resetRunState).toHaveBeenCalled() }) + it('should keep one resumable stream for installed apps', () => { + const { handlers } = setupHandlers({ isPublicAPI: false }) + const onWorkflowPaused = handlers.onWorkflowPaused! + const pausedEvent = { + data: { + workflow_run_id: 'run-installed', + }, + } as never + + onWorkflowPaused(pausedEvent) + onWorkflowPaused(pausedEvent) + + expect(sseGetMock).toHaveBeenCalledWith( + '/workflow/run-installed/events?include_state_snapshot=true&continue_on_pause=true', + {}, + expect.objectContaining({ isPublicAPI: false }), + ) + expect(sseGetMock).toHaveBeenCalledTimes(1) + }) + it('should finish timed-out workflow state and warn without applying late outputs', () => { const timeoutSetup = setupHandlers({ isTimedOut: () => true, diff --git a/web/app/components/share/text-generation/result/workflow-stream-handlers.ts b/web/app/components/share/text-generation/result/workflow-stream-handlers.ts index 64ff5fc7dc8..7ab402aba93 100644 --- a/web/app/components/share/text-generation/result/workflow-stream-handlers.ts +++ b/web/app/components/share/text-generation/result/workflow-stream-handlers.ts @@ -280,6 +280,7 @@ export const createWorkflowStreamHandlers = ({ taskId, }: CreateWorkflowStreamHandlersParams): IOtherOptions => { let tempMessageId = '' + let hasStartedResumeStream = false const finishWithFailure = () => { setRespondingFalse() @@ -420,8 +421,14 @@ export const createWorkflowStreamHandlers = ({ }, onWorkflowPaused: ({ data }) => { tempMessageId = data.workflow_run_id - // WebApp workflows must keep using the public API namespace after pause/resume. - void sseGet(`/workflow/${data.workflow_run_id}/events`, {}, otherOptions) + if (!hasStartedResumeStream) { + hasStartedResumeStream = true + void sseGet( + `/workflow/${data.workflow_run_id}/events?include_state_snapshot=true&continue_on_pause=true`, + {}, + otherOptions, + ) + } setWorkflowProcessData(applyWorkflowPaused(getWorkflowProcessData())) }, } diff --git a/web/app/components/workflow-app/hooks/__tests__/use-workflow-run-callbacks.spec.ts b/web/app/components/workflow-app/hooks/__tests__/use-workflow-run-callbacks.spec.ts index eff299a56c2..c87c9a403d7 100644 --- a/web/app/components/workflow-app/hooks/__tests__/use-workflow-run-callbacks.spec.ts +++ b/web/app/components/workflow-app/hooks/__tests__/use-workflow-run-callbacks.spec.ts @@ -141,10 +141,16 @@ describe('useWorkflowRun callbacks helpers', () => { expect(player.playAudioWithAudio).toHaveBeenCalledWith('audio-chunk', true) expect(mockResetMsgId).toHaveBeenCalledWith('message-1') + callbacks.onWorkflowPaused?.({ workflow_run_id: 'run-2' } as never) callbacks.onWorkflowPaused?.({ workflow_run_id: 'run-2' } as never) expect(handlers.handleWorkflowPaused).toHaveBeenCalled() expect(userOnWorkflowPaused).toHaveBeenCalled() - expect(mockSseGet).toHaveBeenCalledWith('/workflow/run-2/events', {}, callbacks) + expect(mockSseGet).toHaveBeenCalledWith( + '/workflow/run-2/events?include_state_snapshot=true&continue_on_pause=true', + {}, + callbacks, + ) + expect(mockSseGet).toHaveBeenCalledTimes(1) }) it('should create final callbacks that preserve rest callback override order and eager abort-controller wiring', () => { @@ -267,6 +273,7 @@ describe('useWorkflowRun callbacks helpers', () => { callbacks.onTTSChunk?.('message-1', 'audio-chunk') callbacks.onTTSEnd?.('message-1', 'audio-finished') callbacks.onWorkflowPaused?.({ workflow_run_id: 'run-2' } as never) + callbacks.onWorkflowPaused?.({ workflow_run_id: 'run-2' } as never) callbacks.onError?.({ error: 'failed', node_type: 'llm' } as never, '500') expect(handlers.handleWorkflowStarted).toHaveBeenCalled() @@ -320,7 +327,12 @@ describe('useWorkflowRun callbacks helpers', () => { expect(mockResetMsgId).toHaveBeenCalledWith('message-1') expect(handlers.handleWorkflowPaused).toHaveBeenCalled() expect(userCallbacks.onWorkflowPaused).toHaveBeenCalled() - expect(mockSseGet).toHaveBeenCalledWith('/workflow/run-2/events', {}, callbacks) + expect(mockSseGet).toHaveBeenCalledWith( + '/workflow/run-2/events?include_state_snapshot=true&continue_on_pause=true', + {}, + callbacks, + ) + expect(mockSseGet).toHaveBeenCalledTimes(1) expect(clearAbortController).toHaveBeenCalled() expect(handlers.handleWorkflowFailed).toHaveBeenCalled() expect(userCallbacks.onError).toHaveBeenCalledWith({ error: 'failed', node_type: 'llm' }, '500') @@ -439,6 +451,7 @@ describe('useWorkflowRun callbacks helpers', () => { finalCallbacks.onHumanInputFormFilled?.({ node_id: 'node-1' } as never) finalCallbacks.onHumanInputFormTimeout?.({ node_id: 'node-1' } as never) finalCallbacks.onWorkflowPaused?.({ workflow_run_id: 'run-2' } as never) + finalCallbacks.onWorkflowPaused?.({ workflow_run_id: 'run-2' } as never) finalCallbacks.onTTSChunk?.('message-2', 'audio-chunk') finalCallbacks.onTTSEnd?.('message-2', 'audio-finished') await finalCallbacks.onCompleted?.(true, 'done') @@ -482,7 +495,12 @@ describe('useWorkflowRun callbacks helpers', () => { expect(userCallbacks.onHumanInputFormTimeout).toHaveBeenCalled() expect(handlers.handleWorkflowPaused).toHaveBeenCalled() expect(userCallbacks.onWorkflowPaused).toHaveBeenCalled() - expect(mockSseGet).toHaveBeenCalledWith('/workflow/run-2/events', {}, finalCallbacks) + expect(mockSseGet).toHaveBeenCalledWith( + '/workflow/run-2/events?include_state_snapshot=true&continue_on_pause=true', + {}, + finalCallbacks, + ) + expect(mockSseGet).toHaveBeenCalledTimes(1) expect(player.playAudioWithAudio).toHaveBeenCalledWith('audio-chunk', true) expect(player.playAudioWithAudio).toHaveBeenCalledWith('audio-finished', false) expect(clearAbortController).toHaveBeenCalled() diff --git a/web/app/components/workflow-app/hooks/use-workflow-run-callbacks.ts b/web/app/components/workflow-app/hooks/use-workflow-run-callbacks.ts index c0f7524a8c9..01259976324 100644 --- a/web/app/components/workflow-app/hooks/use-workflow-run-callbacks.ts +++ b/web/app/components/workflow-app/hooks/use-workflow-run-callbacks.ts @@ -147,6 +147,7 @@ export const createBaseWorkflowRunCallbacks = ({ onHumanInputFormTimeout, onCompleted, } = callbacks + let hasStartedResumeStream = false const wrappedOnError: IOtherOptions['onError'] = (params, code) => { clearAbortController() @@ -260,8 +261,11 @@ export const createBaseWorkflowRunCallbacks = ({ handleWorkflowPaused() invalidateRunHistory(runHistoryUrl) if (onWorkflowPaused) onWorkflowPaused(params) - const url = `/workflow/${params.workflow_run_id}/events` - sseGet(url, {}, baseSseOptions) + if (!hasStartedResumeStream) { + hasStartedResumeStream = true + const url = `/workflow/${params.workflow_run_id}/events?include_state_snapshot=true&continue_on_pause=true` + sseGet(url, {}, baseSseOptions) + } }, onHumanInputRequired: (params) => { handleWorkflowNodeHumanInputRequired(params) @@ -340,6 +344,7 @@ export const createFinalWorkflowRunCallbacks = ({ onHumanInputFormFilled, onHumanInputFormTimeout, } = callbacks + let hasStartedResumeStream = false const finalCallbacks: IOtherOptions = { ...baseSseOptions, @@ -437,8 +442,11 @@ export const createFinalWorkflowRunCallbacks = ({ handleWorkflowPaused() invalidateRunHistory(runHistoryUrl) if (onWorkflowPaused) onWorkflowPaused(params) - const url = `/workflow/${params.workflow_run_id}/events` - sseGet(url, {}, finalCallbacks) + if (!hasStartedResumeStream) { + hasStartedResumeStream = true + const url = `/workflow/${params.workflow_run_id}/events?include_state_snapshot=true&continue_on_pause=true` + sseGet(url, {}, finalCallbacks) + } }, onHumanInputRequired: (params) => { handleWorkflowNodeHumanInputRequired(params) diff --git a/web/app/components/workflow/panel/__tests__/human-input-form-list.spec.tsx b/web/app/components/workflow/panel/__tests__/human-input-form-list.spec.tsx index b22057985aa..91812d40a78 100644 --- a/web/app/components/workflow/panel/__tests__/human-input-form-list.spec.tsx +++ b/web/app/components/workflow/panel/__tests__/human-input-form-list.spec.tsx @@ -143,6 +143,23 @@ describe('HumanInputFormList', () => { expect(screen.queryByTestId('tips')).not.toBeInTheDocument() }) + it('should reset inputs when the same node produces a new form', async () => { + const user = userEvent.setup() + const { rerender } = render() + + const input = screen.getByTestId('content-item-textarea') + await user.clear(input) + await user.type(input, 'previous response') + + rerender( + , + ) + + expect(screen.getByTestId('content-item-textarea')).toHaveValue('prefill') + }) + it('should render an empty container when there are no visible forms', () => { render() diff --git a/web/app/components/workflow/panel/debug-and-preview/__tests__/hooks/handle-resume.spec.ts b/web/app/components/workflow/panel/debug-and-preview/__tests__/hooks/handle-resume.spec.ts index 1809a7f79a1..56a8ab97f6d 100644 --- a/web/app/components/workflow/panel/debug-and-preview/__tests__/hooks/handle-resume.spec.ts +++ b/web/app/components/workflow/panel/debug-and-preview/__tests__/hooks/handle-resume.spec.ts @@ -125,7 +125,7 @@ describe('useChat – handleResume', () => { }) expect(mockSseGet).toHaveBeenCalledWith( - '/workflow/wfr-1/events?include_state_snapshot=true', + '/workflow/wfr-1/events?include_state_snapshot=true&continue_on_pause=true', {}, expect.any(Object), ) @@ -889,7 +889,7 @@ describe('useChat – handleResume', () => { }) describe('onWorkflowPaused', () => { - it('should re-subscribe via sseGet and set status to Paused', async () => { + it('should keep the resumable stream and set status to Paused', async () => { const { result } = await setupResumeWithTree() const sseGetCallsBefore = mockSseGet.mock.calls.length @@ -899,7 +899,7 @@ describe('useChat – handleResume', () => { }) }) - expect(mockSseGet.mock.calls.length).toBeGreaterThan(sseGetCallsBefore) + expect(mockSseGet.mock.calls.length).toBe(sseGetCallsBefore) const answer = result.current.chatList.find((item) => item.id === 'msg-resume') expect(answer!.workflowProcess!.status).toBe('paused') }) diff --git a/web/app/components/workflow/panel/debug-and-preview/hooks.ts b/web/app/components/workflow/panel/debug-and-preview/hooks.ts index 69c91d2c397..da38afa91d0 100644 --- a/web/app/components/workflow/panel/debug-and-preview/hooks.ts +++ b/web/app/components/workflow/panel/debug-and-preview/hooks.ts @@ -732,7 +732,7 @@ export const useChat = ( const handleResume = useCallback( (messageId: string, workflowRunId: string, { onGetSuggestedQuestions }: SendCallback) => { // Re-subscribe to workflow events for the specific message - const url = `/workflow/${workflowRunId}/events?include_state_snapshot=true` + const url = `/workflow/${workflowRunId}/events?include_state_snapshot=true&continue_on_pause=true` const otherOptions: IOtherOptions = { getAbortController: (abortController) => { @@ -1002,9 +1002,7 @@ export const useChat = ( } }) }, - onWorkflowPaused: ({ data: workflowPausedData }) => { - const resumeUrl = `/workflow/${workflowPausedData.workflow_run_id}/events` - sseGet(resumeUrl, {}, otherOptions) + onWorkflowPaused: () => { updateChatTreeNode(messageId, (responseItem) => { responseItem.workflowProcess!.status = WorkflowRunningStatus.Paused }) diff --git a/web/app/components/workflow/panel/human-input-form-list.tsx b/web/app/components/workflow/panel/human-input-form-list.tsx index 9d9ed38cf92..554cd11f803 100644 --- a/web/app/components/workflow/panel/human-input-form-list.tsx +++ b/web/app/components/workflow/panel/human-input-form-list.tsx @@ -1,3 +1,4 @@ +import type { HumanInputFormSubmitData } from '@/app/components/base/chat/chat/answer/human-input-content/type' import type { DeliveryMethod } from '@/app/components/workflow/nodes/human-input/types' import type { HumanInputFormData } from '@/types/workflow' import { useCallback, useMemo } from 'react' @@ -9,7 +10,7 @@ import { DeliveryMethodType } from '@/app/components/workflow/nodes/human-input/ type HumanInputFormListProps = { humanInputFormDataList: HumanInputFormData[] - onHumanInputFormSubmit?: (formToken: string, formData: any) => Promise + onHumanInputFormSubmit?: (formToken: string, formData: HumanInputFormSubmitData) => Promise } const HumanInputFormList = ({ @@ -77,12 +78,12 @@ const HumanInputFormList = ({
{filteredHumanInputFormDataList.map((formData) => (