fix(workflow): honor stop between non-streaming nodes (#41090)

Co-authored-by: Crazywoola <100913391+crazywoola@users.noreply.github.com>
Co-authored-by: crazywoola <427733928@qq.com>

api/core/app/apps/advanced_chat/app_runner.py

@@ -11,7 +11,9 @@ from core.app.apps.base_app_queue_manager import AppQueueManager
 from core.app.apps.workflow.command_channels import (
     CelerySignalCommandChannel,
     CombinedCommandChannel,
+    StopFlagCommandChannel,
 )
+from core.app.apps.workflow.stop_aware_ready_queue import attach_stop_aware_ready_queue
 from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner
 from core.app.entities.app_invoke_entities import (
     AdvancedChatAppGenerateEntity,
@@ -229,9 +231,11 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
             shutdown_state_getter=celery_warm_shutdown_started,
             abort_reason=WORKFLOW_WARM_SHUTDOWN_ABORT_REASON,
         )
+        attach_stop_aware_ready_queue(graph_runtime_state, task_id=task_id)
         command_channel = CombinedCommandChannel(
             (
                 RedisChannel(redis_client, channel_key),
+                StopFlagCommandChannel(task_id=task_id),
                 celery_signal_channel,
             )
         )

api/core/app/apps/execution_coordinator.py

@@ -26,11 +26,24 @@ def app_task_command_channel_key(task_id: str) -> str:
     return f"workflow:{task_id}:commands"
 
 
+def app_task_stop_flag_key(task_id: str) -> str:
+    """Redis key of the legacy generate-task stop flag."""
+    return f"generate_task_stopped:{task_id}"
+
+
 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)
+    redis_client.setex(app_task_stop_flag_key(task_id), 600, 1)
+
+
+def is_app_task_stop_flag_set(task_id: str) -> bool:
+    """Return whether the legacy Redis stop flag is currently armed."""
+    if not task_id:
+        return False
+
+    return redis_client.get(app_task_stop_flag_key(task_id)) is not None
 
 
 def clear_app_task_cancellation_signals(task_id: str) -> None:
@@ -47,7 +60,7 @@ def clear_app_task_cancellation_signals(task_id: str) -> None:
         return
 
     try:
-        redis_client.delete(f"generate_task_stopped:{task_id}")
+        redis_client.delete(app_task_stop_flag_key(task_id))
     except Exception:
         logger.exception("Failed to clear stop flag for app task %s", task_id)
 

api/core/app/apps/workflow/app_runner.py

@@ -9,7 +9,9 @@ from core.app.apps.workflow.app_config_manager import WorkflowAppConfig
 from core.app.apps.workflow.command_channels import (
     CelerySignalCommandChannel,
     CombinedCommandChannel,
+    StopFlagCommandChannel,
 )
+from core.app.apps.workflow.stop_aware_ready_queue import attach_stop_aware_ready_queue
 from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner
 from core.app.entities.app_invoke_entities import DifyRunContext, InvokeFrom, WorkflowAppGenerateEntity
 from core.app.workflow.layers.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer
@@ -159,9 +161,11 @@ class WorkflowAppRunner(WorkflowBasedAppRunner):
             shutdown_state_getter=celery_warm_shutdown_started,
             abort_reason=WORKFLOW_WARM_SHUTDOWN_ABORT_REASON,
         )
+        attach_stop_aware_ready_queue(graph_runtime_state, task_id=task_id)
         command_channel = CombinedCommandChannel(
             (
                 RedisChannel(redis_client, channel_key),
+                StopFlagCommandChannel(task_id=task_id),
                 celery_signal_channel,
             )
         )

api/core/app/apps/workflow/command_channels.py

@@ -4,6 +4,7 @@ import logging
 from collections.abc import Callable, Sequence
 from typing import final, override
 
+from core.app.apps.execution_coordinator import is_app_task_stop_flag_set
 from graphon.graph_engine.command_channels import CommandChannel
 from graphon.graph_engine.entities.commands import AbortCommand, GraphEngineCommand
 
@@ -65,3 +66,34 @@ class CelerySignalCommandChannel(CommandChannel):
     @override
     def send_command(self, command: GraphEngineCommand) -> None:
         _ = command
+
+
+@final
+class StopFlagCommandChannel(CommandChannel):
+    """Translate the legacy Redis stop flag into one GraphEngine abort command."""
+
+    _task_id: str
+    _abort_reason: str
+    _abort_emitted: bool
+
+    def __init__(
+        self,
+        *,
+        task_id: str,
+        abort_reason: str = "User requested stop",
+    ) -> None:
+        self._task_id = task_id
+        self._abort_reason = abort_reason
+        self._abort_emitted = False
+
+    @override
+    def fetch_commands(self) -> list[GraphEngineCommand]:
+        if self._abort_emitted or not is_app_task_stop_flag_set(self._task_id):
+            return []
+
+        self._abort_emitted = True
+        return [AbortCommand(reason=self._abort_reason)]
+
+    @override
+    def send_command(self, command: GraphEngineCommand) -> None:
+        _ = command

api/core/app/apps/workflow/stop_aware_ready_queue.py

@@ -0,0 +1,67 @@
+"""Ready-queue wrapper that honors stop before the next node is scheduled."""
+
+from __future__ import annotations
+
+from typing import final
+
+from core.app.apps.execution_coordinator import is_app_task_stop_flag_set
+from graphon.graph_engine.ready_queue import ReadyQueue, ReadyTask
+from graphon.runtime.graph_runtime_state import GraphExecutionProtocol, GraphRuntimeState
+
+
+@final
+class StopAwareReadyQueue:
+    """Reject newly ready nodes once the run has been stopped.
+
+    GraphEngine drain still enqueues successors after abort. Drop those puts so
+    later nodes do not start; the in-flight node can finish.
+    """
+
+    def __init__(
+        self,
+        inner: ReadyQueue,
+        *,
+        task_id: str,
+        graph_execution: GraphExecutionProtocol,
+    ) -> None:
+        self._inner = inner
+        self._task_id = task_id
+        self._graph_execution = graph_execution
+
+    def _should_reject(self) -> bool:
+        return self._graph_execution.aborted or is_app_task_stop_flag_set(self._task_id)
+
+    def put(self, item: ReadyTask) -> None:
+        if self._should_reject():
+            return
+        self._inner.put(item)
+
+    def get(self, timeout: float | None = None) -> ReadyTask:
+        return self._inner.get(timeout=timeout)
+
+    def task_done(self) -> None:
+        self._inner.task_done()
+
+    def qsize(self) -> int:
+        return self._inner.qsize()
+
+    def drain(self) -> list[ReadyTask]:
+        return self._inner.drain()
+
+    def dumps(self) -> str:
+        return self._inner.dumps()
+
+    def loads(self, data: str) -> None:
+        self._inner.loads(data)
+
+
+def attach_stop_aware_ready_queue(graph_runtime_state: GraphRuntimeState, *, task_id: str) -> None:
+    """Install stop-aware enqueue policy on an existing runtime state."""
+    current = graph_runtime_state.ready_queue
+    if isinstance(current, StopAwareReadyQueue):
+        return
+    graph_runtime_state._ready_queue = StopAwareReadyQueue(
+        current,
+        task_id=task_id,
+        graph_execution=graph_runtime_state.graph_execution,
+    )

api/tests/unit_tests/core/app/apps/test_execution_coordinator.py

@@ -6,7 +6,9 @@ from core.app.apps.execution_coordinator import (
     AppExecutionCoordinator,
     AppExecutionState,
     app_task_command_channel_key,
+    app_task_stop_flag_key,
     clear_app_task_cancellation_signals,
+    is_app_task_stop_flag_set,
 )
 
 
@@ -143,3 +145,22 @@ def test_stop_flag_failure_does_not_block_graph_stop(caplog: pytest.LogCaptureFi
         reason="test abort",
     )
     assert "Failed to set stop flag for app execution task=task" in caplog.text
+
+
+def test_stop_flag_key_matches_legacy_redis_flag() -> None:
+    assert app_task_stop_flag_key("task") == "generate_task_stopped:task"
+
+
+def test_is_app_task_stop_flag_set_reads_redis() -> None:
+    with patch("core.app.apps.execution_coordinator.redis_client") as redis_client:
+        redis_client.get.return_value = b"1"
+        assert is_app_task_stop_flag_set("task") is True
+        redis_client.get.assert_called_once_with("generate_task_stopped:task")
+
+
+def test_is_app_task_stop_flag_set_is_false_when_missing_or_empty() -> None:
+    with patch("core.app.apps.execution_coordinator.redis_client") as redis_client:
+        redis_client.get.return_value = None
+        assert is_app_task_stop_flag_set("task") is False
+        assert is_app_task_stop_flag_set("") is False
+        redis_client.get.assert_called_once_with("generate_task_stopped:task")

api/tests/unit_tests/core/app/apps/workflow/test_command_channels.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
 
 from core.app.apps.workflow.command_channels import (
     CelerySignalCommandChannel,
     CombinedCommandChannel,
+    StopFlagCommandChannel,
 )
 from graphon.graph_engine.entities.commands import AbortCommand, PauseCommand
 
@@ -105,3 +107,35 @@ def test_celery_signal_command_channel_send_command_is_noop() -> None:
     channel.send_command(command)
 
     assert channel.fetch_commands() == []
+
+
+def test_stop_flag_command_channel_emits_abort_when_flag_is_set() -> None:
+    channel = StopFlagCommandChannel(task_id="task-1", abort_reason="User requested stop")
+
+    with patch("core.app.apps.workflow.command_channels.is_app_task_stop_flag_set", return_value=False):
+        assert channel.fetch_commands() == []
+
+    with patch("core.app.apps.workflow.command_channels.is_app_task_stop_flag_set", return_value=True):
+        commands = channel.fetch_commands()
+
+    assert len(commands) == 1
+    assert isinstance(commands[0], AbortCommand)
+    assert commands[0].reason == "User requested stop"
+
+
+def test_stop_flag_command_channel_emits_abort_once_per_instance() -> None:
+    channel = StopFlagCommandChannel(task_id="task-1")
+
+    with patch("core.app.apps.workflow.command_channels.is_app_task_stop_flag_set", return_value=True):
+        assert len(channel.fetch_commands()) == 1
+        assert channel.fetch_commands() == []
+
+
+def test_stop_flag_command_channel_send_command_is_noop() -> None:
+    channel = StopFlagCommandChannel(task_id="task-1")
+    command = PauseCommand(reason="pause")
+
+    channel.send_command(command)
+
+    with patch("core.app.apps.workflow.command_channels.is_app_task_stop_flag_set", return_value=False):
+        assert channel.fetch_commands() == []

api/tests/unit_tests/core/app/apps/workflow/test_stop_aware_ready_queue.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from unittest.mock import Mock, patch
+
+from core.app.apps.workflow.stop_aware_ready_queue import (
+    StopAwareReadyQueue,
+    attach_stop_aware_ready_queue,
+)
+from graphon.graph_engine.domain.graph_execution import GraphExecution
+from graphon.graph_engine.ready_queue import StartTask
+from graphon.runtime.graph_runtime_state import GraphRuntimeState
+from graphon.runtime.variable_pool import VariablePool
+
+
+def _start_task(node_id: str = "next-node") -> StartTask:
+    return StartTask(frame_id="root", node_id=node_id)
+
+
+def _graph_execution(*, aborted: bool = False) -> GraphExecution:
+    return GraphExecution(workflow_id="workflow-1", aborted=aborted)
+
+
+def test_ready_queue_accepts_work_while_run_is_active() -> None:
+    inner = Mock()
+    queue = StopAwareReadyQueue(inner, task_id="task-1", graph_execution=_graph_execution())
+    task = _start_task()
+
+    with patch(
+        "core.app.apps.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set",
+        return_value=False,
+    ):
+        queue.put(task)
+
+    inner.put.assert_called_once_with(task)
+
+
+def test_ready_queue_rejects_next_node_after_graph_abort() -> None:
+    inner = Mock()
+    queue = StopAwareReadyQueue(inner, task_id="task-1", graph_execution=_graph_execution(aborted=True))
+
+    with patch(
+        "core.app.apps.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set",
+        return_value=False,
+    ):
+        queue.put(_start_task())
+
+    inner.put.assert_not_called()
+
+
+def test_ready_queue_rejects_next_node_when_stop_flag_is_set() -> None:
+    inner = Mock()
+    queue = StopAwareReadyQueue(inner, task_id="task-1", graph_execution=_graph_execution())
+
+    with patch(
+        "core.app.apps.workflow.stop_aware_ready_queue.is_app_task_stop_flag_set",
+        return_value=True,
+    ) as stop_flag:
+        queue.put(_start_task("code-node"))
+
+    stop_flag.assert_called_once_with("task-1")
+    inner.put.assert_not_called()
+
+
+def test_attach_stop_aware_ready_queue_wraps_once() -> None:
+    inner = Mock()
+    runtime_state = GraphRuntimeState(
+        variable_pool=VariablePool(),
+        start_at=0,
+        ready_queue=inner,
+        graph_execution=_graph_execution(),
+    )
+
+    attach_stop_aware_ready_queue(runtime_state, task_id="task-1")
+    first = runtime_state.ready_queue
+    attach_stop_aware_ready_queue(runtime_state, task_id="task-1")
+
+    assert isinstance(first, StopAwareReadyQueue)
+    assert runtime_state.ready_queue is first
+
+
+def test_stop_aware_queue_delegates_reads() -> None:
+    inner = Mock()
+    inner.get.return_value = _start_task("queued")
+    inner.qsize.return_value = 1
+    inner.drain.return_value = [_start_task("queued")]
+    inner.dumps.return_value = "{}"
+    queue = StopAwareReadyQueue(inner, task_id="task-1", graph_execution=_graph_execution())
+
+    assert queue.get(timeout=0.1) == _start_task("queued")
+    queue.task_done()
+    assert queue.qsize() == 1
+    assert queue.drain() == [_start_task("queued")]
+    assert queue.dumps() == "{}"
+    queue.loads("{}")
+
+    inner.get.assert_called_once_with(timeout=0.1)
+    inner.task_done.assert_called_once()
+    inner.loads.assert_called_once_with("{}")