feat(api): abort active workflow runs during Celery warm shutdown

[linw1995]

Important

1. Make sure you have read our contribution guidelines
2. Ensure there is an associated issue and you have been assigned to it
3. Use the correct syntax to link this PR: Fixes #<issue number>.

Summary

Screenshots

Before	After
...	...

Checklist

• This change requires a documentation update, included: Dify Document
• I understand that this PR may be closed in case there was no previous discussion or issues. (This doesn't apply to typos!)
• I've added a test for each change that was introduced, and I tried as much as possible to make a single atomic change.
• I've updated the documentation accordingly.
• I ran make lint && make type-check (backend) and cd web && pnpm exec vp staged (frontend) to appease the lint gods

[github-actions]

Pyrefly Type Coverage

Metric	Base	PR	Delta
Type coverage	52.15%	52.18%	+0.03%
Strict coverage	51.67%	51.70%	+0.03%
Typed symbols	31,816	31,865	+49
Untyped symbols	29,464	29,474	+10
Modules	2953	2959	+6

[github-actions]

Pyrefly Diff

base → PR

--- /tmp/pyrefly_base.txt	2026-07-03 02:41:29.840793690 +0000
+++ /tmp/pyrefly_pr.txt	2026-07-03 02:41:21.388797318 +0000
@@ -3764,6 +3764,8 @@
 ERROR Object of class `StreamResponse` has no attribute `data` [missing-attribute]
    --> tests/unit_tests/core/app/apps/test_workflow_pause_events.py:284:12
 ERROR Object of class `FileInputConfig` has no attribute `option_source`
+ERROR Generator function should return `Generator` [bad-return]
+  --> tests/unit_tests/core/app/apps/workflow/test_active_workflow_tasks.py:11:29
 ERROR Argument `SimpleNamespace` is not assignable to parameter `app_model` with type `App` in function `core.app.apps.workflow.app_config_manager.WorkflowAppConfigManager.get_app_config` [bad-argument-type]
   --> tests/unit_tests/core/app/apps/workflow/test_app_config_manager.py:23:66
 ERROR Argument `SimpleNamespace` is not assignable to parameter `workflow` with type `Workflow` in function `core.app.apps.workflow.app_config_manager.WorkflowAppConfigManager.get_app_config` [bad-argument-type]
@@ -3874,6 +3876,8 @@
    --> tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py:505:43
 ERROR Argument `SimpleNamespace` is not assignable to parameter `workflow_node_execution_repository` with type `WorkflowNodeExecutionRepository` in function `core.app.apps.workflow.app_generator.WorkflowAppGenerator._generate_worker` [bad-argument-type]
    --> tests/unit_tests/core/app/apps/workflow/test_app_generator_extra.py:506:48
+ERROR Argument `tuple[SimpleNamespace, _CommandChannelStub]` is not assignable to parameter `command_channels` with type `Sequence[CommandChannel]` in function `core.app.apps.workflow.command_channels.CombinedCommandChannel.__init__` [bad-argument-type]
+  --> tests/unit_tests/core/app/apps/workflow/test_command_channels.py:64:39
 ERROR `in` is not supported between `Literal['blocking response mode']` and `None` [not-iterable]
  --> tests/unit_tests/core/app/apps/workflow/test_errors.py:9:16
 ERROR Cannot index into `str` [bad-index]
@@ -6416,6 +6420,8 @@
   --> tests/unit_tests/extensions/test_set_secretkey.py:49:32
 ERROR Argument `Flask` is not assignable to parameter `app` with type `DifyApp` in function `extensions.ext_set_secretkey.init_app` [bad-argument-type]
   --> tests/unit_tests/extensions/test_set_secretkey.py:69:32
+ERROR Generator function should return `Generator` [bad-return]
+  --> tests/unit_tests/extensions/test_workflow_warm_shutdown.py:13:36
 ERROR Object of class `NoneType` has no attribute `storage_key` [missing-attribute]
    --> tests/unit_tests/factories/test_build_from_mapping.py:132:12
 ERROR Object of class `NoneType` has no attribute `storage_key` [missing-attribute]

[codecov]

Codecov Report

❌ Patch coverage is 94.02985% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 85.37%. Comparing base (44e85c0) to head (4908c51).
⚠️ Report is 64 commits behind head on main.

Files with missing lines	Patch %	Lines
api/core/app/workflow/layers/observability.py	78.57%	1 Missing and 2 partials ⚠️
api/extensions/workflow_warm_shutdown.py	91.89%	2 Missing and 1 partial ⚠️
...pi/core/app/apps/workflow/active_workflow_tasks.py	91.30%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #38220      +/-   ##
==========================================
+ Coverage   85.36%   85.37%   +0.01%     
==========================================
  Files        4970     4981      +11     
  Lines      258774   261668    +2894     
  Branches    49076    49292     +216     
==========================================
+ Hits       220893   223411    +2518     
- Misses      33576    33866     +290     
- Partials     4305     4391      +86     

Flag	Coverage Δ
api	85.54% <94.02%> (+0.02%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[QuantumGhost]

The overall approach seems a bit weird, for the following reasons:

1. The cancellation logic (api/extensions/workflow_warm_shutdown.py) is separated from where the cancellation is registered.
2. The original CommandChanenl is mainly used for cross-machine cancellation scenario. However, in the case of this PR, the cancelled workers and the cancellation does live inside a same machine. It seems like a waste to send Redis command to trigger cancellation.
3. The CommandChanenl is designed as an interface intentionally to provide other command source.

I would propose the following solutions:

1. Introduces a CombinedCommandChannel, to combine multiple CommandChannel into one:

class CombinedCommandChannel(CommandChannel):
    def __init__(self, command_channels: Sequence[CommandChannel]): ...

2. Introduce CelerySignalCommandChannel, to convert Celery worker_shutting_down signals to Graphon commands
3. Combine classes introduced in 1 and 2, and rely on the command mechanism of Graphon.

[QuantumGhost]

I think the shutdown modeling can be simplified further.

Instead of storing a process-global AbortCommand, I would suggest modeling Celery warm shutdown as a process-local threading.Event owned by a Celery adapter module, for example:

_shutdown_started = threading.Event()

def _on_shutting_down(*args: object, **kwargs: object) -> None:
    _shutdown_started.set()

def shutdown_started() -> bool:
    return _shutdown_started.is_set()

Then CelerySignalCommandChannel can just depend on a shutdown-state getter and translate that state into a Graphon abort command, e.g.:

class CelerySignalCommandChannel(CommandChannel):
    def __init__(
        self,
        shutdown_state_getter: Callable[[], bool],
        abort_reason: str = WORKFLOW_WARM_SHUTDOWN_ABORT_REASON,
    ) -> None:
        ...

This gives a cleaner boundary:

• the Celery adapter owns the worker shutdown lifecycle state
• CelerySignalCommandChannel only translates that state into Graphon commands
• workflow runners / WorkflowEntry continue to depend only on the CommandChannel abstraction

I would also suggest keeping an instance-level flag inside CelerySignalCommandChannel (for example self._abort_emitted) so each channel emits the abort command at most once. That avoids the module-level _abort_command, avoids runtime reset, and still preserves the late-created channel behavior after shutdown has already started.

With this design, active_workflow_tasks.py also looks unnecessary for the cancellation path, since the shutdown signal can be injected into CelerySignalCommandChannel directly instead of being split across multiple modules.

If the active task registry is only there for cancellation, I think it can be removed. If you still want to log the number of in-flight workflow runs during worker shutdown, then keeping the current counting mechanism makes sense as an observability / diagnostic concern, but it should stay separate from the command propagation mechanism.