mirror of
https://github.com/langgenius/dify.git
synced 2026-08-15 04:59:46 +08:00
fix(api): decouple response listeners from execution lifecycle (#39813)
This commit is contained in:
zyssyz123
GitHub
parent
211344d0b0
commit
6e80f42f38
@ -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:
|
||||
|
||||
173
api/core/app/apps/execution_coordinator.py
Normal file
173
api/core/app/apps/execution_coordinator.py
Normal file
@ -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
|
||||
@ -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:
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
@ -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()
|
||||
|
||||
@ -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()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user