opensource/dify
2
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-09-05 08:48:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(dify-agent): isolate cancellation intent observer (#40896)

Browse Source
This commit is contained in:
盐粒 Yanli 2026-08-18 03:29:42 +00:00 committed by GitHub
parent 09f8f43d20
commit c649b84bf6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 47 additions and 152 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
dify-agent/src/dify_agent/runtime/run_scheduler.py
View File

@ -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:

46
dify-agent/src/dify_agent/storage/redis_run_store.py
View File

@ -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,

21
dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py
View File

@ -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()

113
dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py
View File

@ -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: