From 6e80f42f3864da36c2faa7d78b5c496a87c80e97 Mon Sep 17 00:00:00 2001 From: zyssyz123 <916125788@qq.com> Date: Thu, 13 Aug 2026 05:50:43 +0000 Subject: [PATCH] fix(api): decouple response listeners from execution lifecycle (#39813) --- api/core/app/apps/base_app_queue_manager.py | 70 ++++--- api/core/app/apps/execution_coordinator.py | 173 ++++++++++++++++++ .../apps/message_based_app_queue_manager.py | 19 +- .../apps/pipeline/pipeline_queue_manager.py | 3 +- .../app/apps/workflow/app_queue_manager.py | 13 +- .../app/apps/test_base_app_queue_manager.py | 31 ++-- .../app/apps/test_execution_coordinator.py | 92 ++++++++++ .../test_message_based_app_queue_manager.py | 34 +++- .../apps/workflow/test_app_queue_manager.py | 75 ++++++-- 9 files changed, 418 insertions(+), 92 deletions(-) create mode 100644 api/core/app/apps/execution_coordinator.py create mode 100644 api/tests/unit_tests/core/app/apps/test_execution_coordinator.py diff --git a/api/core/app/apps/base_app_queue_manager.py b/api/core/app/apps/base_app_queue_manager.py index 365d66ad0e3..d8053f2797b 100644 --- a/api/core/app/apps/base_app_queue_manager.py +++ b/api/core/app/apps/base_app_queue_manager.py @@ -10,7 +10,11 @@ from cachetools import TTLCache, cachedmethod from redis.exceptions import RedisError from sqlalchemy.orm import DeclarativeMeta -from configs import dify_config +from core.app.apps.execution_coordinator import ( + AppExecutionCoordinator, + AppExecutionState, + set_app_task_stop_flag, +) from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import ( AppQueueEvent, @@ -21,7 +25,6 @@ from core.app.entities.queue_entities import ( WorkflowQueueMessage, ) from extensions.ext_redis import redis_client -from graphon.graph_engine.manager import GraphEngineManager from graphon.runtime import GraphRuntimeState logger = logging.getLogger(__name__) @@ -52,17 +55,18 @@ class AppQueueManager(ABC): self._graph_runtime_state: GraphRuntimeState | None = None self._stopped_cache: TTLCache[tuple, bool] = TTLCache(maxsize=1, ttl=1) self._cache_lock = threading.Lock() - self._execution_terminal = threading.Event() - self._abort_sent = threading.Event() - self._lifecycle_lock = threading.Lock() + self._listener_segment_completed = threading.Event() + self._execution_coordinator = AppExecutionCoordinator( + task_id=self._task_id, + on_timeout=self._publish_timeout_stop, + ) def listen(self): """ Listen to queue :return: """ - # wait for APP_MAX_EXECUTION_TIME seconds to stop listen - listen_timeout = dify_config.APP_MAX_EXECUTION_TIME + self._execution_coordinator.start_watchdog() start_time = time.monotonic() last_ping_time: int | float = 0 try: @@ -77,13 +81,8 @@ class AppQueueManager(ABC): continue finally: elapsed_time = time.monotonic() - start_time - timed_out = elapsed_time >= listen_timeout manually_stopped = self._is_stopped() - if not self._execution_terminal.is_set() and (timed_out or manually_stopped): - reason = ( - f"App execution exceeded {listen_timeout} seconds" if timed_out else "App task was stopped" - ) - self._abort_execution(reason) + if manually_stopped and self._execution_coordinator.request_abort("App task was stopped"): # publish two messages to make sure the client can receive the stop signal # and stop listening after the stop signal processed self.publish( @@ -94,32 +93,31 @@ class AppQueueManager(ABC): self.publish(QueuePingEvent(), PublishFrom.TASK_PIPELINE) last_ping_time = elapsed_time // 10 finally: - if not self._execution_terminal.is_set(): - self._abort_execution("Client response stream closed before app execution completed") + self._execution_coordinator.listener_closed(segment_completed=self._listener_segment_completed.is_set()) self._graph_runtime_state = None # Release reference once consumers finish or close the generator. - def stop_listen(self, *, execution_terminal: bool = False): - """ - Stop listen to queue - :return: - """ - if execution_terminal: - self._execution_terminal.set() + def stop_listen(self, *, execution_state: AppExecutionState) -> None: + """Complete the current listener segment with an explicit execution state.""" + if execution_state is AppExecutionState.PAUSED: + self._execution_coordinator.mark_paused() + elif execution_state is AppExecutionState.TERMINAL: + self._execution_coordinator.mark_terminal() + else: + raise ValueError(f"Unsupported listener completion state: {execution_state}") + + self._listener_segment_completed.set() self._clear_task_belong_cache() self._q.put(None) - def _abort_execution(self, reason: str) -> None: - """Propagate response timeout/disconnect to legacy and GraphEngine runners.""" - with self._lifecycle_lock: - if self._execution_terminal.is_set() or self._abort_sent.is_set(): - return - self._abort_sent.set() + @property + def execution_state(self) -> AppExecutionState: + return self._execution_coordinator.state - try: - self.set_stop_flag_no_user_check(self._task_id) - GraphEngineManager(redis_client).send_stop_command(self._task_id, reason=reason) - except Exception: - logger.exception("Failed to abort app execution for task %s", self._task_id) + def _publish_timeout_stop(self, reason: str) -> None: + self.publish( + QueueStopEvent(stopped_by=QueueStopEvent.StopBy.USER_MANUAL, reason=reason), + PublishFrom.TASK_PIPELINE, + ) def _clear_task_belong_cache(self) -> None: """ @@ -201,11 +199,7 @@ class AppQueueManager(ABC): :param task_id: The task ID to stop :return: """ - if not task_id: - return - - stopped_cache_key = cls._generate_stopped_cache_key(task_id) - redis_client.setex(stopped_cache_key, 600, 1) + set_app_task_stop_flag(task_id) @cachedmethod(lambda self: self._stopped_cache, lock=lambda self: self._cache_lock) def _is_stopped(self) -> bool: diff --git a/api/core/app/apps/execution_coordinator.py b/api/core/app/apps/execution_coordinator.py new file mode 100644 index 00000000000..ce8e8986307 --- /dev/null +++ b/api/core/app/apps/execution_coordinator.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import logging +import threading +import uuid +from collections.abc import Callable +from enum import Enum, auto + +from configs import dify_config +from extensions.ext_redis import redis_client +from graphon.graph_engine.manager import GraphEngineManager + +logger = logging.getLogger(__name__) + + +class AppExecutionState(Enum): + RUNNING = auto() + PAUSED = auto() + ABORTING = auto() + TERMINAL = auto() + + +def set_app_task_stop_flag(task_id: str) -> None: + if not task_id: + return + + redis_client.setex(f"generate_task_stopped:{task_id}", 600, 1) + + +class AppExecutionCoordinator: + """Own cancellation policy for one app execution attempt. + + A resumed workflow creates a new coordinator even when it reuses the stable + task ID. Listener segments only report lifecycle observations here; they do + not send cancellation commands themselves. Response detachment is not an + execution cancellation signal: streaming workflow execution may continue in + another process and publish durable events for a later subscriber. + """ + + def __init__( + self, + *, + task_id: str, + on_timeout: Callable[[str], None], + timeout_seconds: int | float | None = None, + ) -> None: + self._task_id = task_id + self._attempt_id = str(uuid.uuid4()) + self._on_timeout = on_timeout + self._timeout_seconds = timeout_seconds if timeout_seconds is not None else dify_config.APP_MAX_EXECUTION_TIME + self._state = AppExecutionState.RUNNING + self._abort_sent = False + self._watchdog_started = False + self._watchdog: threading.Timer | None = None + self._lock = threading.Lock() + + @property + def attempt_id(self) -> str: + return self._attempt_id + + @property + def state(self) -> AppExecutionState: + with self._lock: + return self._state + + def start_watchdog(self) -> None: + watchdog: threading.Timer | None = None + run_immediately = False + with self._lock: + if self._watchdog_started or self._state is not AppExecutionState.RUNNING: + return + + self._watchdog_started = True + if self._timeout_seconds <= 0: + run_immediately = True + else: + watchdog = threading.Timer(self._timeout_seconds, self._handle_timeout) + watchdog.daemon = True + self._watchdog = watchdog + + if run_immediately: + self._handle_timeout() + elif watchdog is not None: + watchdog.start() + + def mark_paused(self) -> None: + watchdog: threading.Timer | None = None + with self._lock: + if self._state is not AppExecutionState.RUNNING: + return + self._state = AppExecutionState.PAUSED + watchdog = self._detach_watchdog_locked() + + if watchdog is not None: + watchdog.cancel() + + def mark_terminal(self) -> None: + watchdog: threading.Timer | None = None + with self._lock: + self._state = AppExecutionState.TERMINAL + watchdog = self._detach_watchdog_locked() + + if watchdog is not None: + watchdog.cancel() + + def listener_closed(self, *, segment_completed: bool) -> None: + if segment_completed: + return + + logger.info( + "App response listener detached while execution continues task=%s attempt=%s", + self._task_id, + self._attempt_id, + ) + + def request_abort(self, reason: str) -> bool: + watchdog: threading.Timer | None = None + with self._lock: + if self._state is not AppExecutionState.RUNNING or self._abort_sent: + return False + self._abort_sent = True + self._state = AppExecutionState.ABORTING + watchdog = self._detach_watchdog_locked() + + if watchdog is not None: + watchdog.cancel() + + logger.info( + "Aborting app execution task=%s attempt=%s reason=%s", + self._task_id, + self._attempt_id, + reason, + ) + self._abort_execution(reason) + return True + + def _handle_timeout(self) -> None: + reason = f"App execution exceeded {self._timeout_seconds} seconds" + if not self.request_abort(reason): + return + + try: + self._on_timeout(reason) + except Exception: + logger.exception( + "Failed to publish timeout for app execution task=%s attempt=%s", + self._task_id, + self._attempt_id, + ) + + def _abort_execution(self, reason: str) -> None: + try: + set_app_task_stop_flag(self._task_id) + except Exception: + logger.exception( + "Failed to set stop flag for app execution task=%s attempt=%s", + self._task_id, + self._attempt_id, + ) + + try: + GraphEngineManager(redis_client).send_stop_command(self._task_id, reason=reason) + except Exception: + logger.exception( + "Failed to send stop command for app execution task=%s attempt=%s", + self._task_id, + self._attempt_id, + ) + + def _detach_watchdog_locked(self) -> threading.Timer | None: + watchdog = self._watchdog + self._watchdog = None + return watchdog 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 1be1e532694..548b00eab42 100644 --- a/api/core/app/apps/message_based_app_queue_manager.py +++ b/api/core/app/apps/message_based_app_queue_manager.py @@ -2,6 +2,7 @@ from typing import override from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom from core.app.apps.exc import GenerateTaskStoppedError +from core.app.apps.execution_coordinator import AppExecutionState from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import ( AppQueueEvent, @@ -43,15 +44,17 @@ class MessageBasedAppQueueManager(AppQueueManager): self._q.put(message) - if isinstance( - event, - QueueStopEvent - | QueueErrorEvent - | QueueMessageEndEvent - | QueueAdvancedChatMessageEndEvent - | QueueWorkflowPausedEvent, + if isinstance(event, QueueWorkflowPausedEvent): + self.stop_listen(execution_state=AppExecutionState.PAUSED) + elif isinstance( + event, QueueStopEvent | QueueErrorEvent | QueueMessageEndEvent | QueueAdvancedChatMessageEndEvent ): - self.stop_listen(execution_terminal=True) + execution_state = ( + AppExecutionState.PAUSED + if self.execution_state is AppExecutionState.PAUSED + else AppExecutionState.TERMINAL + ) + self.stop_listen(execution_state=execution_state) if pub_from == PublishFrom.APPLICATION_MANAGER and self._is_stopped(): if self._app_mode == AppMode.ADVANCED_CHAT.value: diff --git a/api/core/app/apps/pipeline/pipeline_queue_manager.py b/api/core/app/apps/pipeline/pipeline_queue_manager.py index d424266454e..13459fb94d1 100644 --- a/api/core/app/apps/pipeline/pipeline_queue_manager.py +++ b/api/core/app/apps/pipeline/pipeline_queue_manager.py @@ -2,6 +2,7 @@ from typing import override from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom from core.app.apps.exc import GenerateTaskStoppedError +from core.app.apps.execution_coordinator import AppExecutionState from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import ( AppQueueEvent, @@ -42,7 +43,7 @@ class PipelineQueueManager(AppQueueManager): | QueueWorkflowFailedEvent | QueueWorkflowPartialSuccessEvent, ): - self.stop_listen(execution_terminal=True) + self.stop_listen(execution_state=AppExecutionState.TERMINAL) if pub_from == PublishFrom.APPLICATION_MANAGER and self._is_stopped(): raise GenerateTaskStoppedError() diff --git a/api/core/app/apps/workflow/app_queue_manager.py b/api/core/app/apps/workflow/app_queue_manager.py index da99d034477..f837805f7b5 100644 --- a/api/core/app/apps/workflow/app_queue_manager.py +++ b/api/core/app/apps/workflow/app_queue_manager.py @@ -1,6 +1,7 @@ from typing import override from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom +from core.app.apps.execution_coordinator import AppExecutionState from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import ( AppQueueEvent, @@ -33,19 +34,15 @@ class WorkflowAppQueueManager(AppQueueManager): self._q.put(message) - # A pause ends only the current listener segment; the workflow stays PAUSED and - # resumes with the same task ID. Without this marker, listen() cleanup calls - # _abort_execution(), whose stop flag and abort command can stop the resumed run. - # This is a compatibility workaround: cancellation policy belongs to the execution - # owner, not the response-stream listener. - if isinstance( + if isinstance(event, QueueWorkflowPausedEvent): + self.stop_listen(execution_state=AppExecutionState.PAUSED) + elif isinstance( event, QueueStopEvent | QueueErrorEvent | QueueMessageEndEvent | QueueWorkflowSucceededEvent | QueueWorkflowFailedEvent - | QueueWorkflowPausedEvent | QueueWorkflowPartialSuccessEvent, ): - self.stop_listen(execution_terminal=True) + self.stop_listen(execution_state=AppExecutionState.TERMINAL) diff --git a/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py b/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py index a92953313f2..2431b9951ea 100644 --- a/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py +++ b/api/tests/unit_tests/core/app/apps/test_base_app_queue_manager.py @@ -4,6 +4,7 @@ from unittest.mock import patch import pytest from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom +from core.app.apps.execution_coordinator import AppExecutionState from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import QueueErrorEvent from models import Tenant @@ -39,7 +40,7 @@ class TestBaseAppQueueManager: mock_redis.setex.assert_called_once() def test_set_stop_flag_no_user_check(self): - with patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis: + with patch("core.app.apps.execution_coordinator.redis_client") as mock_redis: AppQueueManager.set_stop_flag_no_user_check(task_id="t1") mock_redis.setex.assert_called_once() @@ -61,10 +62,10 @@ class TestBaseAppQueueManager: with pytest.raises(TypeError): manager._check_for_sqlalchemy_models(bad) - def test_stop_listen_defers_graph_runtime_state_cleanup_until_listener_exits(self): + def test_completed_listener_defers_graph_runtime_state_cleanup_until_listener_exits(self): with ( patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis, - patch("core.app.apps.base_app_queue_manager.GraphEngineManager") as graph_engine_manager, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, ): mock_redis.setex.return_value = True mock_redis.get.return_value = None @@ -72,27 +73,27 @@ class TestBaseAppQueueManager: runtime_state = SimpleNamespace(name="runtime-state") manager.graph_runtime_state = runtime_state - manager.stop_listen() + manager.stop_listen(execution_state=AppExecutionState.TERMINAL) assert manager.graph_runtime_state is runtime_state assert list(manager.listen()) == [] assert manager.graph_runtime_state is None - graph_engine_manager.return_value.send_stop_command.assert_called_once_with( - "t1", - reason="Client response stream closed before app execution completed", - ) + graph_engine_manager.return_value.send_stop_command.assert_not_called() - def test_abort_execution_is_idempotent_when_graph_stop_command_fails(self, caplog): + def test_execution_coordinator_abort_is_idempotent_when_graph_stop_command_fails(self, caplog): with ( - patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis, - patch("core.app.apps.base_app_queue_manager.GraphEngineManager") as graph_engine_manager, + patch("core.app.apps.base_app_queue_manager.redis_client") as queue_redis, + patch("core.app.apps.execution_coordinator.redis_client") as execution_redis, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, ): - mock_redis.setex.return_value = True + queue_redis.setex.return_value = True + execution_redis.setex.return_value = True graph_engine_manager.return_value.send_stop_command.side_effect = RuntimeError("redis unavailable") manager = DummyQueueManager(task_id="t1", user_id="u1", invoke_from=InvokeFrom.SERVICE_API) - manager._abort_execution("stream closed") - manager._abort_execution("duplicate") + assert manager._execution_coordinator.request_abort("stream closed") is True + assert manager._execution_coordinator.request_abort("duplicate") is False + execution_redis.setex.assert_called_once_with("generate_task_stopped:t1", 600, 1) graph_engine_manager.return_value.send_stop_command.assert_called_once_with("t1", reason="stream closed") - assert "Failed to abort app execution for task t1" in caplog.text + assert "Failed to send stop command for app execution task=t1" in caplog.text diff --git a/api/tests/unit_tests/core/app/apps/test_execution_coordinator.py b/api/tests/unit_tests/core/app/apps/test_execution_coordinator.py new file mode 100644 index 00000000000..4e25532650d --- /dev/null +++ b/api/tests/unit_tests/core/app/apps/test_execution_coordinator.py @@ -0,0 +1,92 @@ +from unittest.mock import Mock, patch + +import pytest + +from core.app.apps.execution_coordinator import AppExecutionCoordinator, AppExecutionState + + +def test_listener_close_does_not_abort_running_attempt() -> None: + on_timeout = Mock() + with ( + patch("core.app.apps.execution_coordinator.redis_client") as redis_client, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, + ): + coordinator = AppExecutionCoordinator(task_id="task", on_timeout=on_timeout, timeout_seconds=1200) + + coordinator.listener_closed(segment_completed=False) + coordinator.listener_closed(segment_completed=False) + + assert coordinator.state is AppExecutionState.RUNNING + redis_client.setex.assert_not_called() + graph_engine_manager.return_value.send_stop_command.assert_not_called() + on_timeout.assert_not_called() + + +def test_paused_attempt_ignores_listener_close_and_timeout() -> None: + on_timeout = Mock() + with ( + patch("core.app.apps.execution_coordinator.redis_client") as redis_client, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, + ): + coordinator = AppExecutionCoordinator(task_id="task", on_timeout=on_timeout, timeout_seconds=0) + coordinator.mark_paused() + + coordinator.start_watchdog() + coordinator.listener_closed(segment_completed=False) + + assert coordinator.state is AppExecutionState.PAUSED + redis_client.setex.assert_not_called() + graph_engine_manager.return_value.send_stop_command.assert_not_called() + on_timeout.assert_not_called() + + +def test_watchdog_aborts_and_notifies_response_pipeline() -> None: + on_timeout = Mock() + with ( + patch("core.app.apps.execution_coordinator.redis_client") as redis_client, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, + ): + coordinator = AppExecutionCoordinator(task_id="task", on_timeout=on_timeout, timeout_seconds=0) + + coordinator.listener_closed(segment_completed=False) + coordinator.start_watchdog() + + assert coordinator.state is AppExecutionState.ABORTING + redis_client.setex.assert_called_once_with("generate_task_stopped:task", 600, 1) + graph_engine_manager.return_value.send_stop_command.assert_called_once_with( + "task", + reason="App execution exceeded 0 seconds", + ) + on_timeout.assert_called_once_with("App execution exceeded 0 seconds") + + +def test_pausing_started_attempt_cancels_watchdog() -> None: + on_timeout = Mock() + watchdog = Mock() + with patch("core.app.apps.execution_coordinator.threading.Timer", return_value=watchdog): + coordinator = AppExecutionCoordinator(task_id="task", on_timeout=on_timeout, timeout_seconds=1200) + + coordinator.start_watchdog() + coordinator.mark_paused() + + watchdog.start.assert_called_once() + watchdog.cancel.assert_called_once() + assert coordinator.state is AppExecutionState.PAUSED + + +def test_stop_flag_failure_does_not_block_graph_stop(caplog: pytest.LogCaptureFixture) -> None: + on_timeout = Mock() + with ( + patch("core.app.apps.execution_coordinator.redis_client") as redis_client, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, + ): + redis_client.setex.side_effect = RuntimeError("redis write failed") + coordinator = AppExecutionCoordinator(task_id="task", on_timeout=on_timeout, timeout_seconds=1200) + + coordinator.request_abort("test abort") + + graph_engine_manager.return_value.send_stop_command.assert_called_once_with( + "task", + reason="test abort", + ) + assert "Failed to set stop flag for app execution task=task" in caplog.text 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 0bc4753752d..9d480366334 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 @@ -4,9 +4,11 @@ import pytest from core.app.apps.base_app_queue_manager import PublishFrom from core.app.apps.exc import GenerateTaskStoppedError +from core.app.apps.execution_coordinator import AppExecutionState 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 ( + QueueAdvancedChatMessageEndEvent, QueueErrorEvent, QueueMessageEndEvent, QueueStopEvent, @@ -69,7 +71,7 @@ class TestMessageBasedAppQueueManager: assert manager._q.qsize() == 1 - def test_publish_pause_event_stops_listener_without_aborting_execution(self): + def test_publish_pause_event_marks_listener_as_paused(self): with patch("core.app.apps.base_app_queue_manager.redis_client") as mock_redis: mock_redis.setex.return_value = True manager = MessageBasedAppQueueManager( @@ -85,4 +87,32 @@ class TestMessageBasedAppQueueManager: manager._publish(QueueWorkflowPausedEvent(), PublishFrom.APPLICATION_MANAGER) - manager.stop_listen.assert_called_once_with(execution_terminal=True) + manager.stop_listen.assert_called_once_with(execution_state=AppExecutionState.PAUSED) + + def test_pause_state_is_preserved_when_advanced_chat_message_ends(self): + with ( + patch("core.app.apps.base_app_queue_manager.redis_client") as queue_redis, + patch("core.app.apps.execution_coordinator.redis_client") as execution_redis, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, + ): + queue_redis.get.return_value = None + manager = MessageBasedAppQueueManager( + task_id="t1", + user_id="u1", + invoke_from=InvokeFrom.DEBUGGER, + conversation_id="c1", + app_mode="advanced-chat", + message_id="m1", + ) + + manager.publish( + QueueWorkflowPausedEvent(reasons=[], outputs={}, paused_nodes=["human-input"]), + PublishFrom.APPLICATION_MANAGER, + ) + manager.publish(QueueAdvancedChatMessageEndEvent(), PublishFrom.TASK_PIPELINE) + messages = list(manager.listen()) + + assert isinstance(messages[0].event, QueueWorkflowPausedEvent) + assert manager.execution_state is AppExecutionState.PAUSED + execution_redis.setex.assert_not_called() + graph_engine_manager.return_value.send_stop_command.assert_not_called() 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 c33860b3c14..eeabb51c27a 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 @@ -3,6 +3,7 @@ from __future__ import annotations from unittest.mock import Mock, patch from core.app.apps.base_app_queue_manager import PublishFrom +from core.app.apps.execution_coordinator import AppExecutionState from core.app.apps.workflow.app_queue_manager import WorkflowAppQueueManager from core.app.entities.app_invoke_entities import InvokeFrom from core.app.entities.queue_entities import ( @@ -41,7 +42,7 @@ class TestWorkflowAppQueueManager: manager._publish(QueuePingEvent(), PublishFrom.TASK_PIPELINE) - def test_publish_pause_event_stops_listener_without_aborting_execution(self): + def test_publish_pause_event_marks_listener_as_paused(self): manager = WorkflowAppQueueManager( task_id="task", user_id="user", @@ -52,14 +53,15 @@ class TestWorkflowAppQueueManager: manager._publish(QueueWorkflowPausedEvent(), PublishFrom.APPLICATION_MANAGER) - manager.stop_listen.assert_called_once_with(execution_terminal=True) + manager.stop_listen.assert_called_once_with(execution_state=AppExecutionState.PAUSED) - def test_listener_close_aborts_unfinished_execution(self): + def test_listener_close_does_not_abort_unfinished_execution(self): with ( - patch("core.app.apps.base_app_queue_manager.redis_client") as redis_client, - patch("core.app.apps.base_app_queue_manager.GraphEngineManager") as graph_engine_manager, + patch("core.app.apps.base_app_queue_manager.redis_client") as queue_redis, + patch("core.app.apps.execution_coordinator.redis_client") as execution_redis, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, ): - redis_client.get.return_value = None + queue_redis.get.return_value = None manager = WorkflowAppQueueManager( task_id="task", user_id="user", @@ -72,18 +74,19 @@ class TestWorkflowAppQueueManager: assert isinstance(next(listener).event, QueuePingEvent) listener.close() - graph_engine_manager.return_value.send_stop_command.assert_called_once_with( - "task", - reason="Client response stream closed before app execution completed", - ) + assert manager.execution_state is AppExecutionState.RUNNING + execution_redis.setex.assert_not_called() + graph_engine_manager.return_value.send_stop_command.assert_not_called() + manager._execution_coordinator.mark_terminal() def test_execution_timeout_aborts_graph_before_stop_event(self): with ( - patch("core.app.apps.base_app_queue_manager.redis_client") as redis_client, - patch("core.app.apps.base_app_queue_manager.GraphEngineManager") as graph_engine_manager, - patch("core.app.apps.base_app_queue_manager.dify_config.APP_MAX_EXECUTION_TIME", 0), + patch("core.app.apps.base_app_queue_manager.redis_client") as queue_redis, + patch("core.app.apps.execution_coordinator.redis_client") as execution_redis, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, + patch("core.app.apps.execution_coordinator.dify_config.APP_MAX_EXECUTION_TIME", 0), ): - redis_client.get.return_value = None + queue_redis.get.return_value = None manager = WorkflowAppQueueManager( task_id="task", user_id="user", @@ -95,6 +98,7 @@ class TestWorkflowAppQueueManager: messages = list(manager.listen()) assert any(isinstance(message.event, QueueStopEvent) for message in messages) + execution_redis.setex.assert_called_once_with("generate_task_stopped:task", 600, 1) graph_engine_manager.return_value.send_stop_command.assert_called_once_with( "task", reason="App execution exceeded 0 seconds", @@ -102,10 +106,11 @@ class TestWorkflowAppQueueManager: def test_terminal_event_does_not_abort_completed_execution(self): with ( - patch("core.app.apps.base_app_queue_manager.redis_client") as redis_client, - patch("core.app.apps.base_app_queue_manager.GraphEngineManager") as graph_engine_manager, + patch("core.app.apps.base_app_queue_manager.redis_client") as queue_redis, + patch("core.app.apps.execution_coordinator.redis_client") as execution_redis, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, ): - redis_client.get.return_value = None + queue_redis.get.return_value = None manager = WorkflowAppQueueManager( task_id="task", user_id="user", @@ -116,14 +121,42 @@ class TestWorkflowAppQueueManager: _ = list(manager.listen()) + execution_redis.setex.assert_not_called() + graph_engine_manager.return_value.send_stop_command.assert_not_called() + + def test_pause_completes_listener_without_aborting_resumable_execution(self): + with ( + patch("core.app.apps.base_app_queue_manager.redis_client") as queue_redis, + patch("core.app.apps.execution_coordinator.redis_client") as execution_redis, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, + ): + queue_redis.get.return_value = None + manager = WorkflowAppQueueManager( + task_id="task", + user_id="user", + invoke_from=InvokeFrom.DEBUGGER, + app_mode="workflow", + ) + manager.publish( + QueueWorkflowPausedEvent(reasons=[], outputs={}, paused_nodes=["human-input"]), + PublishFrom.APPLICATION_MANAGER, + ) + + messages = list(manager.listen()) + + assert len(messages) == 1 + assert isinstance(messages[0].event, QueueWorkflowPausedEvent) + assert manager.execution_state is AppExecutionState.PAUSED + execution_redis.setex.assert_not_called() graph_engine_manager.return_value.send_stop_command.assert_not_called() def test_workflow_pause_does_not_abort_execution(self): with ( - patch("core.app.apps.base_app_queue_manager.redis_client") as redis_client, - patch("core.app.apps.base_app_queue_manager.GraphEngineManager") as graph_engine_manager, + patch("core.app.apps.base_app_queue_manager.redis_client") as queue_redis, + patch("core.app.apps.execution_coordinator.redis_client") as execution_redis, + patch("core.app.apps.execution_coordinator.GraphEngineManager") as graph_engine_manager, ): - redis_client.get.return_value = None + queue_redis.get.return_value = None manager = WorkflowAppQueueManager( task_id="task", user_id="user", @@ -136,4 +169,6 @@ class TestWorkflowAppQueueManager: assert isinstance(next(listener).event, QueueWorkflowPausedEvent) listener.close() + assert manager.execution_state is AppExecutionState.PAUSED + execution_redis.setex.assert_not_called() graph_engine_manager.return_value.send_stop_command.assert_not_called()