fix(agent): make run cancellation terminal

dify-agent/src/dify_agent/runtime/run_scheduler.py

@@ -73,6 +73,7 @@ class RunScheduler:
     store: RunStore
     shutdown_grace_seconds: float
     active_tasks: dict[str, asyncio.Task[None]]
+    cancelled_run_ids: set[str]
     stopping: bool
     runner_factory: RunRunnerFactory
     layer_providers: tuple[LayerProviderInput, ...]
@@ -93,6 +94,7 @@ class RunScheduler:
         self.store = store
         self.shutdown_grace_seconds = shutdown_grace_seconds
         self.active_tasks = {}
+        self.cancelled_run_ids = set()
         self.stopping = False
         self.plugin_daemon_http_client = plugin_daemon_http_client
         self.dify_api_http_client = dify_api_http_client
@@ -114,7 +116,7 @@ class RunScheduler:
             record = await self.store.create_run()
             task = asyncio.create_task(self._run_record(record, request), name=f"dify-agent-run-{record.run_id}")
             self.active_tasks[record.run_id] = task
-            task.add_done_callback(lambda _task, run_id=record.run_id: self.active_tasks.pop(run_id, None))
+            task.add_done_callback(lambda _task, run_id=record.run_id: self._discard_active_run(run_id))
             return record
 
     async def cancel_run(self, run_id: str, request: CancelRunRequest) -> CancelRunResponse:
@@ -129,17 +131,8 @@ class RunScheduler:
             task = self.active_tasks.get(run_id)
             if task is None:
                 raise RunCancellationConflictError("run is not active in this scheduler process")
+            self.cancelled_run_ids.add(run_id)
             _ = task.cancel(request.message or request.reason)
-
-        _ = await asyncio.gather(task, return_exceptions=True)
-
-        async with self._lifecycle_lock:
-            latest = await self.store.get_run(run_id)
-            if latest.status == "cancelled":
-                return CancelRunResponse(run_id=run_id, status="cancelled")
-            if latest.status != "running":
-                raise RunCancellationConflictError(f"run already finished with status {latest.status!r}")
-
             _ = await emit_run_cancelled(
                 self.store,
                 run_id=run_id,
@@ -147,7 +140,18 @@ class RunScheduler:
                 message=request.message,
             )
             await self.store.update_status(run_id, "cancelled", request.message or request.reason)
-            return CancelRunResponse(run_id=run_id, status="cancelled")
+
+        # Some model/tool stacks can consume one CancelledError. Re-inject it
+        # after the terminal state is durable without making the HTTP request
+        # wait for arbitrary third-party cleanup.
+        for _attempt in range(2):
+            if task.done():
+                break
+            _ = task.cancel(request.message or request.reason)
+            await asyncio.sleep(0)
+        if task.done():
+            self._discard_active_run(run_id)
+        return CancelRunResponse(run_id=run_id, status="cancelled")
 
     async def shutdown(self) -> None:
         """Stop accepting runs, wait briefly, then cancel and fail unfinished runs."""
@@ -161,7 +165,11 @@ class RunScheduler:
         if not pending:
             return
 
-        pending_run_ids = [run_id for run_id, task in tasks_by_run_id.items() if task in pending]
+        pending_run_ids = [
+            run_id
+            for run_id, task in tasks_by_run_id.items()
+            if task in pending and run_id not in self.cancelled_run_ids
+        ]
         for task in pending:
             _ = task.cancel()
         _ = await asyncio.gather(*pending, return_exceptions=True)
@@ -186,8 +194,13 @@ class RunScheduler:
             plugin_daemon_http_client=self.plugin_daemon_http_client,
             dify_api_http_client=self.dify_api_http_client,
             layer_providers=self.layer_providers,
+            is_cancelled=lambda: record.run_id in self.cancelled_run_ids,
         )
 
+    def _discard_active_run(self, run_id: str) -> None:
+        _ = self.active_tasks.pop(run_id, None)
+        self.cancelled_run_ids.discard(run_id)
+
     async def _mark_cancelled_run_failed(self, run_id: str) -> None:
         """Best-effort failure event/status for shutdown-cancelled runs."""
         message = "run cancelled during server shutdown"

dify-agent/src/dify_agent/runtime/runner.py

@@ -31,6 +31,7 @@ both the JSON-safe final output or deferred tool call and the session snapshot;
 there are no separate output or snapshot events to correlate.
 """
 
+import asyncio
 from collections.abc import AsyncIterable, Callable, Mapping
 from collections import Counter
 from dataclasses import dataclass
@@ -170,6 +171,7 @@ class AgentRunRunner:
     layer_providers: tuple[LayerProviderInput, ...]
     plugin_daemon_http_client: httpx.AsyncClient
     dify_api_http_client: httpx.AsyncClient
+    is_cancelled: Callable[[], bool]
 
     def __init__(
         self,
@@ -180,6 +182,7 @@ class AgentRunRunner:
         plugin_daemon_http_client: httpx.AsyncClient,
         dify_api_http_client: httpx.AsyncClient,
         layer_providers: tuple[LayerProviderInput, ...] | None = None,
+        is_cancelled: Callable[[], bool] | None = None,
     ) -> None:
         self.sink = sink
         self.request = request
@@ -187,20 +190,29 @@ class AgentRunRunner:
         self.plugin_daemon_http_client = plugin_daemon_http_client
         self.dify_api_http_client = dify_api_http_client
         self.layer_providers = layer_providers if layer_providers is not None else create_default_layer_providers()
+        self.is_cancelled = is_cancelled or (lambda: False)
 
     async def run(self) -> None:
         """Execute the run and emit the documented event sequence."""
+        if self.is_cancelled():
+            return
         await self.sink.update_status(self.run_id, "running")
+        if self.is_cancelled():
+            return
         _ = await emit_run_started(self.sink, run_id=self.run_id)
 
         try:
             outcome = await self._run_agent()
         except Exception as exc:
+            if self.is_cancelled():
+                return
             message, reason = _run_failed_error_payload(exc)
             _ = await emit_run_failed(self.sink, run_id=self.run_id, error=message, reason=reason)
             await self.sink.update_status(self.run_id, "failed", message)
             raise
 
+        if self.is_cancelled():
+            return
         _ = await emit_run_succeeded(
             self.sink,
             run_id=self.run_id,
@@ -309,6 +321,8 @@ class AgentRunRunner:
 
                 async def handle_events(_ctx: object, events: AsyncIterable[AgentStreamEvent]) -> None:
                     async for event in events:
+                        if self.is_cancelled():
+                            raise asyncio.CancelledError
                         text_delta = _extract_agent_message_delta(event)
                         _ = await emit_pydantic_ai_event(
                             self.sink,

dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py

@@ -117,6 +117,23 @@ class ControlledRunner:
         await self.release.wait()
 
 
+class SwallowOneCancellationRunner:
+    started: asyncio.Event
+    first_cancellation: asyncio.Event
+
+    def __init__(self, *, started: asyncio.Event, first_cancellation: asyncio.Event) -> None:
+        self.started = started
+        self.first_cancellation = first_cancellation
+
+    async def run(self) -> None:
+        _ = self.started.set()
+        try:
+            await asyncio.Event().wait()
+        except asyncio.CancelledError:
+            _ = self.first_cancellation.set()
+            await asyncio.Event().wait()
+
+
 def test_create_run_starts_background_task_and_returns_running() -> None:
     async def scenario() -> None:
         store = FakeStore()
@@ -201,6 +218,39 @@ def test_cancel_run_stops_task_and_persists_cancelled_terminal() -> None:
     asyncio.run(scenario())
 
 
+def test_cancel_run_reinjects_cancellation_without_waiting_for_runner_cleanup() -> None:
+    async def scenario() -> None:
+        store = FakeStore()
+        started = asyncio.Event()
+        first_cancellation = asyncio.Event()
+        async with httpx.AsyncClient() as client:
+            scheduler = RunScheduler(
+                store=store,
+                plugin_daemon_http_client=client,
+                dify_api_http_client=client,
+                runner_factory=lambda _record, _request: SwallowOneCancellationRunner(
+                    started=started,
+                    first_cancellation=first_cancellation,
+                ),
+            )
+            record = await scheduler.create_run(_request())
+            await asyncio.wait_for(started.wait(), timeout=1)
+
+            response = await asyncio.wait_for(
+                scheduler.cancel_run(record.run_id, CancelRunRequest(reason="workflow_aborted")),
+                timeout=1,
+            )
+
+            assert response.status == "cancelled"
+            assert first_cancellation.is_set()
+            assert store.statuses[record.run_id] == "cancelled"
+            assert [event.type for event in store.events[record.run_id]] == ["run_cancelled"]
+            await asyncio.sleep(0)
+            assert scheduler.active_tasks == {}
+
+    asyncio.run(scenario())
+
+
 def test_cancel_run_rejects_finished_run() -> None:
     async def scenario() -> None:
         store = FakeStore()

dify-agent/tests/local/dify_agent/runtime/test_runner.py

@@ -64,7 +64,12 @@ from dify_agent.protocol.schemas import (
 )
 from dify_agent.runtime.event_sink import InMemoryRunEventSink
 from dify_agent.runtime.compositor_factory import create_default_layer_providers
-from dify_agent.runtime.runner import AgentRunRunner, AgentRunValidationError, _run_failed_error_payload
+from dify_agent.runtime.runner import (
+    AgentRunRunner,
+    AgentRunValidationError,
+    RunSuccessOutcome,
+    _run_failed_error_payload,
+)
 from shellctl.shared import DeleteJobResponse, JobResult, JobStatusName, JobStatusView
 
 
@@ -164,6 +169,38 @@ def test_run_failed_error_payload_preserves_knowledge_error_code() -> None:
     assert reason == "dataset_not_found"
 
 
+def test_cancelled_runner_does_not_overwrite_cancelled_status_with_late_failure(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    async def scenario() -> None:
+        sink = InMemoryRunEventSink()
+        cancelled = False
+        async with httpx.AsyncClient() as client:
+            runner = AgentRunRunner(
+                sink=sink,
+                request=_request(),
+                run_id="run-cancelled",
+                plugin_daemon_http_client=client,
+                dify_api_http_client=client,
+                is_cancelled=lambda: cancelled,
+            )
+
+            async def fail_after_cancel() -> RunSuccessOutcome:
+                nonlocal cancelled
+                cancelled = True
+                await sink.update_status("run-cancelled", "cancelled", "workflow stopped")
+                raise RuntimeError("late model failure")
+
+            monkeypatch.setattr(runner, "_run_agent", fail_after_cancel)
+            await runner.run()
+
+        assert sink.statuses["run-cancelled"] == "cancelled"
+        assert sink.errors["run-cancelled"] == "workflow stopped"
+        assert [event.type for event in sink.events["run-cancelled"]] == ["run_started"]
+
+    asyncio.run(scenario())
+
+
 def _request(
     user: str | list[str] = "hello",
     *,