From c649b84bf6ea35fd00ff95067ae1f1cfc6d99c51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=9B=90=E7=B2=92=20Yanli?= Date: Tue, 18 Aug 2026 03:29:42 +0000 Subject: [PATCH] fix(dify-agent): isolate cancellation intent observer (#40896) --- .../src/dify_agent/runtime/run_scheduler.py | 19 ++- .../src/dify_agent/storage/redis_run_store.py | 46 ++----- .../dify_agent/runtime/test_run_scheduler.py | 21 ++-- .../storage/test_redis_run_store.py | 113 ++++-------------- 4 files changed, 47 insertions(+), 152 deletions(-) diff --git a/dify-agent/src/dify_agent/runtime/run_scheduler.py b/dify-agent/src/dify_agent/runtime/run_scheduler.py index b197d7dcb3e..8bfd2316085 100644 --- a/dify-agent/src/dify_agent/runtime/run_scheduler.py +++ b/dify-agent/src/dify_agent/runtime/run_scheduler.py @@ -54,8 +54,8 @@ class RunStore(RunEventSink, Protocol): """Return the accepted cancellation intent, if one exists.""" ... - async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None: - """Wait for a cancellation intent or a different terminal state.""" + async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent: + """Wait for an accepted cancellation intent.""" ... async def finalize_cancellation( @@ -209,15 +209,12 @@ class RunScheduler: ) raise - if intent is not None: - await cancel_runner_and_wait() - _ = await self.store.finalize_cancellation( - record.run_id, - intent, - session_snapshot=runner.terminal_session_snapshot, - ) - else: - await runner_task + await cancel_runner_and_wait() + _ = await self.store.finalize_cancellation( + record.run_id, + intent, + session_snapshot=runner.terminal_session_snapshot, + ) else: runner_error: Exception | None = None try: diff --git a/dify-agent/src/dify_agent/storage/redis_run_store.py b/dify-agent/src/dify_agent/storage/redis_run_store.py index 70a6adc85d6..5422e93bd2d 100644 --- a/dify-agent/src/dify_agent/storage/redis_run_store.py +++ b/dify-agent/src/dify_agent/storage/redis_run_store.py @@ -272,44 +272,14 @@ class RedisRunStore(RunEventSink): return None return self._decode_cancellation_intent(entries[0][1]) - async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None: - """Wait until cancellation intent or another terminal state wins. - - The stream cursor is captured before reading the record so a terminal - transition cannot fall between the initial status check and blocking - stream read. - """ - events_key = run_events_key(self.prefix, run_id) - latest_events = await self.redis.xrevrange(events_key, count=1) - cursor = _decode_redis_text(latest_events[0][0]) if latest_events else "0-0" - record = await self.get_run(run_id) - if record.status != "running": - return None - - intent = await self.get_cancellation_intent(run_id) - if intent is not None: - return intent - - while True: - response = await self.redis.xread( - { - run_cancel_intent_key(self.prefix, run_id): "0-0", - events_key: cursor, - }, - block=0, - count=100, - ) - for stream_name, entries in response: - if _decode_redis_text(stream_name) == run_cancel_intent_key(self.prefix, run_id): - return self._decode_cancellation_intent(entries[0][1]) - for raw_id, fields in entries: - event = self._decode_event(run_id, raw_id, fields) - if event.id is not None: - cursor = event.id - if event.type == "run_cancelled": - return None - if event.type in {"run_succeeded", "run_failed"}: - return None + async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent: + """Wait until the first accepted private cancellation intent is available.""" + response = await self.redis.xread( + {run_cancel_intent_key(self.prefix, run_id): "0-0"}, + block=0, + count=1, + ) + return self._decode_cancellation_intent(response[0][1][0][1]) async def finalize_cancellation( self, 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 44643d6a312..493eeab37bb 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 @@ -99,7 +99,7 @@ class FakeStore: statuses: dict[str, RunStatus] errors: dict[str, str | None] error_types: dict[str, RunFailureType | None] - terminal_changes: dict[str, asyncio.Event] + cancellation_changes: dict[str, asyncio.Event] cancellation_intents: dict[str, RunCancellationIntent] def __init__(self) -> None: @@ -108,7 +108,7 @@ class FakeStore: self.statuses = {} self.errors = {} self.error_types = {} - self.terminal_changes = {} + self.cancellation_changes = {} self.cancellation_intents = {} async def create_run(self) -> RunRecord: @@ -116,7 +116,7 @@ class FakeStore: record = RunRecord(run_id=run_id, status="running") self.records[run_id] = record self.statuses[run_id] = "running" - self.terminal_changes[run_id] = asyncio.Event() + self.cancellation_changes[run_id] = asyncio.Event() return record async def append_event(self, event: NonTerminalRunEvent) -> str: @@ -146,7 +146,6 @@ class FakeStore: self.statuses[event.run_id] = status self.errors[event.run_id] = error self.error_types[event.run_id] = error_type - self.terminal_changes[event.run_id].set() return RunFinalizationResult(applied=True, status=status, event_id=event_id) async def request_cancellation(self, run_id: str, request: CancelRunRequest) -> RunStatus: @@ -159,16 +158,15 @@ class FakeStore: message=request.message, requested_at=utc_now(), ) - self.terminal_changes[run_id].set() + self.cancellation_changes[run_id].set() return "running" async def get_cancellation_intent(self, run_id: str) -> RunCancellationIntent | None: return self.cancellation_intents.get(run_id) - async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None: - while self.statuses[run_id] == "running" and run_id not in self.cancellation_intents: - await self.terminal_changes[run_id].wait() - return self.cancellation_intents.get(run_id) + async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent: + await self.cancellation_changes[run_id].wait() + return self.cancellation_intents[run_id] async def finalize_cancellation( self, @@ -196,7 +194,6 @@ class FakeStore: self.errors[run_id] = intent.message or intent.reason self.error_types[run_id] = None del self.cancellation_intents[run_id] - self.terminal_changes[run_id].set() return RunFinalizationResult(applied=True, status="cancelled", event_id=event_id) @@ -228,7 +225,7 @@ class TrackingStore(FakeStore): if not pause_observer: self.release_observer.set() - async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None: + async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent: self.observer_started.set() try: await self.release_observer.wait() @@ -246,7 +243,7 @@ class FailingObserverStore(FakeStore): self.fail_observer = fail_observer self.observer_finished = asyncio.Event() - async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None: + async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent: del run_id try: await self.fail_observer.wait() diff --git a/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py b/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py index 13930624dd1..83fefac8b55 100644 --- a/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py +++ b/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py @@ -62,18 +62,6 @@ class FakeRedis: self.stream_changed.set() return event_id - async def xrevrange( - self, - key: str, - max: str = "+", - min: str = "-", - *, - count: int | None = None, - ) -> list[tuple[str, dict[str, object]]]: - self.commands.append(("xrevrange", key, max, min, count)) - entries = list(reversed(self.streams.get(key, []))) - return entries[:count] if count is not None else entries - async def xread( self, streams: Mapping[str, str], @@ -310,73 +298,43 @@ def test_request_cancellation_raises_when_record_is_missing() -> None: asyncio.run(store.request_cancellation("missing", CancelRunRequest(reason="cancelled"))) -def test_wait_for_cancellation_observes_terminal_record_before_starting() -> None: +def test_wait_for_cancellation_reads_existing_intent_from_stream_start() -> None: redis = FakeRedis() store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] + intent = RunCancellationIntent( + reason="cancelled", + requested_at=utc_now(), + ) - async def scenario() -> object: - record = await store.create_run() - redis.values[f"test:runs:{record.run_id}:record"] = record.model_copy( - update={"status": "cancelled"} - ).model_dump_json() - redis.commands.clear() - return await store.wait_for_cancellation(record.run_id) - - assert asyncio.run(scenario()) is None - assert [command[0] for command in redis.commands] == ["xrevrange", "get"] - - -def test_wait_for_cancellation_covers_intent_transition_during_initialization() -> None: - class PausingRecordReadRedis(FakeRedis): - record_read_started: asyncio.Event - release_record_read: asyncio.Event - pause_next_record_read: bool - - def __init__(self) -> None: - super().__init__() - self.record_read_started = asyncio.Event() - self.release_record_read = asyncio.Event() - self.pause_next_record_read = True - - async def get(self, key: str) -> object | None: - if self.pause_next_record_read and key.endswith(":record"): - self.pause_next_record_read = False - self.record_read_started.set() - await self.release_record_read.wait() - return await super().get(key) - - redis = PausingRecordReadRedis() - observer_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - - async def scenario() -> object: - record = await observer_store.create_run() - observer = asyncio.create_task(observer_store.wait_for_cancellation(record.run_id)) - await asyncio.wait_for(redis.record_read_started.wait(), timeout=1) + async def scenario() -> RunCancellationIntent: _ = redis._append_stream_entry( - f"test:runs:{record.run_id}:cancel-intent", - { - "payload": RunCancellationIntent( - reason="cancelled", - requested_at=utc_now(), - ).model_dump_json() - }, + "test:runs:run-1:cancel-intent", + {"payload": intent.model_dump_json()}, ) - redis.release_record_read.set() - return await asyncio.wait_for(observer, timeout=1) + return await asyncio.wait_for(store.wait_for_cancellation("run-1"), timeout=1) - assert asyncio.run(scenario()) is not None + assert asyncio.run(scenario()) == intent + assert redis.commands == [ + ("xread", {"test:runs:run-1:cancel-intent": "0-0"}, 1, 0), + ] -def test_wait_for_cancellation_advances_past_non_terminal_events() -> None: +def test_wait_for_cancellation_ignores_public_events() -> None: redis = FakeRedis() store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - async def scenario() -> object: + async def scenario() -> RunCancellationIntent: record = await store.create_run() + redis.commands.clear() observer = asyncio.create_task(store.wait_for_cancellation(record.run_id)) await asyncio.sleep(0) _ = await store.append_event(RunStartedEvent(run_id=record.run_id)) await asyncio.sleep(0) + assert observer.done() is False + xread_commands = [command for command in redis.commands if command[0] == "xread"] + assert xread_commands == [ + ("xread", {f"test:runs:{record.run_id}:cancel-intent": "0-0"}, 1, 0), + ] _ = redis._append_stream_entry( f"test:runs:{record.run_id}:cancel-intent", { @@ -388,34 +346,7 @@ def test_wait_for_cancellation_advances_past_non_terminal_events() -> None: ) return await asyncio.wait_for(observer, timeout=1) - assert asyncio.run(scenario()) is not None - cursors = [command[1] for command in redis.commands if command[0] == "xread"] - assert any("0-0" in streams.values() for streams in cursors if isinstance(streams, dict)) - assert any("1-0" in streams.values() for streams in cursors if isinstance(streams, dict)) - - -def test_wait_for_cancellation_returns_none_when_success_wins() -> None: - redis = FakeRedis() - store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType] - - async def scenario() -> object: - record = await store.create_run() - observer = asyncio.create_task(store.wait_for_cancellation(record.run_id)) - await asyncio.sleep(0) - event = RunSucceededEvent( - run_id=record.run_id, - data=RunSucceededEventData( - output="done", - session_snapshot=CompositorSessionSnapshot(layers=[]), - ), - ) - _ = redis._append_stream_entry( - f"test:runs:{record.run_id}:events", - {"payload": RUN_EVENT_ADAPTER.dump_json(event, exclude={"id"}).decode()}, - ) - return await asyncio.wait_for(observer, timeout=1) - - assert asyncio.run(scenario()) is None + assert asyncio.run(scenario()).reason == "cancelled" def test_append_event_serializes_typed_event_without_id_and_expires_run_keys() -> None: