From 104f6e3bcff386177faac1fab52a7655d9e8c414 Mon Sep 17 00:00:00 2001 From: Yansong Zhang <916125788@qq.com> Date: Fri, 17 Jul 2026 14:08:25 +0800 Subject: [PATCH] fix(agent): make run cancellation terminal --- .../src/dify_agent/runtime/run_scheduler.py | 39 ++++++++++----- dify-agent/src/dify_agent/runtime/runner.py | 14 ++++++ .../dify_agent/runtime/test_run_scheduler.py | 50 +++++++++++++++++++ .../local/dify_agent/runtime/test_runner.py | 39 ++++++++++++++- 4 files changed, 128 insertions(+), 14 deletions(-) diff --git a/dify-agent/src/dify_agent/runtime/run_scheduler.py b/dify-agent/src/dify_agent/runtime/run_scheduler.py index dea0dfcc64c..c02799a66ec 100644 --- a/dify-agent/src/dify_agent/runtime/run_scheduler.py +++ b/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" diff --git a/dify-agent/src/dify_agent/runtime/runner.py b/dify-agent/src/dify_agent/runtime/runner.py index 97c1ce7ef1f..315f4484c4b 100644 --- a/dify-agent/src/dify_agent/runtime/runner.py +++ b/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, diff --git a/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py b/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py index 4ddc55a2f14..ac72d0502f1 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py +++ b/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() diff --git a/dify-agent/tests/local/dify_agent/runtime/test_runner.py b/dify-agent/tests/local/dify_agent/runtime/test_runner.py index 25814034220..07886a9dfa2 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_runner.py +++ b/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", *,