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

fix(dify-agent): preserve snapshots for failed and cancelled runs (#40876)

Browse Source
This commit is contained in:
盐粒 Yanli 2026-08-17 10:11:22 +00:00 committed by GitHub
parent 96809473fc
commit a679ee5644
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
24 changed files with 1575 additions and 440 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
api/clients/agent_backend/client.py
View File

@ -23,6 +23,7 @@ from dify_agent.protocol import (
CancelRunResponse,
CreateRunRequest,
CreateRunResponse,
RunCancelledEvent,
RunEvent,
RunStatusResponse,
)
@ -45,6 +46,15 @@ class AgentBackendRunClient(Protocol):
def cancel_run(self, run_id: str, request: CancelRunRequest | None = None) -> CancelRunResponse:
"""Request explicit cancellation for one Agent backend run."""
def cancel_run_and_wait(
self,
run_id: str,
request: CancelRunRequest | None = None,
*,
after: str | None = None,
) -> RunCancelledEvent:
"""Request cancellation and wait for runner cleanup to finish."""
def stream_events(
self,
run_id: str,
@ -67,6 +77,15 @@ class _DifyAgentSyncClient(Protocol):
def cancel_run_sync(self, run_id: str, request: CancelRunRequest | None = None) -> CancelRunResponse:
"""Cancel one run synchronously."""
def cancel_run_and_wait_sync(
self,
run_id: str,
request: CancelRunRequest | None = None,
*,
after: str | None = None,
) -> RunCancelledEvent:
"""Cancel one run and wait for its terminal event synchronously."""
def stream_events_sync(
self,
run_id: str,
@ -110,6 +129,19 @@ class DifyAgentBackendRunClient:
except Exception as exc:
raise _normalize_dify_agent_error(exc) from exc
def cancel_run_and_wait(
self,
run_id: str,
request: CancelRunRequest | None = None,
*,
after: str | None = None,
) -> RunCancelledEvent:
"""Cancel one run, then wait for the cleanup-complete terminal event."""
try:
return self.client.cancel_run_and_wait_sync(run_id, request=request, after=after)
except Exception as exc:
raise _normalize_dify_agent_error(exc) from exc
def stream_events(
self,
run_id: str,

4
api/clients/agent_backend/event_adapter.py
View File

@ -99,6 +99,7 @@ class AgentBackendRunFailedInternalEvent(AgentBackendInternalEventBase):
error: str
error_type: RunFailureType | None = None
reason: str | None = None
session_snapshot: CompositorSessionSnapshot | None = None
class AgentBackendRunCancelledInternalEvent(AgentBackendInternalEventBase):
@ -107,6 +108,7 @@ class AgentBackendRunCancelledInternalEvent(AgentBackendInternalEventBase):
type: Literal[AgentBackendInternalEventType.RUN_CANCELLED] = AgentBackendInternalEventType.RUN_CANCELLED
reason: str | None = None
message: str | None = None
session_snapshot: CompositorSessionSnapshot | None = None
type AgentBackendInternalEvent = Annotated[
@ -184,6 +186,7 @@ class AgentBackendRunEventAdapter:
error=event.data.error,
error_type=event.data.error_type,
reason=event.data.reason,
session_snapshot=event.data.session_snapshot,
)
]
case RunCancelledEvent():
@ -193,6 +196,7 @@ class AgentBackendRunEventAdapter:
source_event_id=event.id,
reason=event.data.reason,
message=event.data.message,
session_snapshot=event.data.session_snapshot,
)
]
raise TypeError(f"unsupported agent backend run event: {type(event).__name__}")

30
api/clients/agent_backend/fake_client.py
View File

@ -18,6 +18,8 @@ from dify_agent.protocol import (
CreateRunRequest,
CreateRunResponse,
DeferredToolCallPayload,
RunCancelledEvent,
RunCancelledEventData,
RunEvent,
RunFailedEvent,
RunFailedEventData,
@ -69,6 +71,28 @@ class FakeAgentBackendRunClient:
del request
return CancelRunResponse(run_id=run_id, status="cancelled")
def cancel_run_and_wait(
self,
run_id: str,
request: CancelRunRequest | None = None,
*,
after: str | None = None,
) -> RunCancelledEvent:
"""Return a deterministic cleanup-complete cancellation event."""
del after
request = request or CancelRunRequest()
_ = self.cancel_run(run_id, request)
return RunCancelledEvent(
id="cancel-0",
run_id=run_id,
created_at=_FIXED_TIME,
data=RunCancelledEventData(
reason=request.reason,
message=request.message,
session_snapshot=CompositorSessionSnapshot(layers=[]),
),
)
def stream_events(
self,
run_id: str,
@ -133,7 +157,11 @@ class FakeAgentBackendRunClient:
id="2-0",
run_id=run_id,
created_at=_FIXED_TIME,
data=RunFailedEventData(error="fake failure", reason="unit_test"),
data=RunFailedEventData(
error="fake failure",
reason="unit_test",
session_snapshot=CompositorSessionSnapshot(layers=[]),
),
),
)
case FakeAgentBackendScenario.PAUSED:

72
api/core/app/apps/agent_app/app_runner.py
View File

@ -24,6 +24,7 @@ from clients.agent_backend import (
AgentBackendDeferredToolCallInternalEvent,
AgentBackendError,
AgentBackendInternalEventType,
AgentBackendRunCancelledInternalEvent,
AgentBackendRunClient,
AgentBackendRunEventAdapter,
AgentBackendRunFailedError,
@ -681,6 +682,8 @@ class AgentAppRunner:
queue_manager=queue_manager,
model_name=model_name,
query=query,
session_scope=scope,
binding_id=runtime.binding_id,
)
if isinstance(terminal, AgentBackendDeferredToolCallInternalEvent):
@ -700,6 +703,15 @@ class AgentAppRunner:
)
return
if isinstance(terminal, AgentBackendRunFailedInternalEvent | AgentBackendRunCancelledInternalEvent):
# None means no post-exit snapshot was produced; leave the previously stored session snapshot untouched.
if terminal.session_snapshot is not None:
self._save_session(
scope=scope,
binding_id=runtime.binding_id,
snapshot=terminal.session_snapshot,
)
if not isinstance(terminal, AgentBackendRunSucceededInternalEvent):
if isinstance(terminal, AgentBackendRunFailedInternalEvent):
reason = terminal.reason
@ -895,6 +907,8 @@ class AgentAppRunner:
queue_manager: AppQueueManager,
model_name: str,
query: str | None,
session_scope: AgentAppSessionScope,
binding_id: str,
):
"""Consume backend events while preserving raw recorder granularity."""
terminal = None
@ -904,6 +918,7 @@ class AgentAppRunner:
queue_manager=queue_manager,
)
text_delta_debouncer = _TextDeltaDebouncer(debounce_seconds=self._text_delta_debounce_seconds)
last_event_id: str | None = None
def persist_answer_text(content_delta: str) -> None:
try:
@ -934,14 +949,26 @@ class AgentAppRunner:
should_stop=queue_manager.is_stopped,
)
for public_event in public_events:
if public_event.id is not None:
last_event_id = public_event.id
if queue_manager.is_stopped():
flush_pending_agent_message_text()
self._cancel_run(run_id)
self._cancel_run(
run_id,
after=last_event_id,
session_scope=session_scope,
binding_id=binding_id,
)
raise GenerateTaskStoppedError()
for internal_event in self._event_adapter.adapt(public_event):
if queue_manager.is_stopped():
flush_pending_agent_message_text()
self._cancel_run(run_id)
self._cancel_run(
run_id,
after=last_event_id,
session_scope=session_scope,
binding_id=binding_id,
)
raise GenerateTaskStoppedError()
if internal_event.type in (
AgentBackendInternalEventType.RUN_STARTED,
@ -978,21 +1005,52 @@ class AgentAppRunner:
raise
except Exception as error:
flush_pending_agent_message_text()
self._cancel_run(run_id)
self._cancel_run(
run_id,
after=last_event_id,
session_scope=session_scope,
binding_id=binding_id,
)
if queue_manager.is_stopped():
raise GenerateTaskStoppedError() from error
raise
flush_pending_agent_message_text()
if queue_manager.is_stopped():
self._cancel_run(run_id)
self._cancel_run(
run_id,
after=last_event_id,
session_scope=session_scope,
binding_id=binding_id,
)
raise GenerateTaskStoppedError()
return terminal, process_recorder
def _cancel_run(self, run_id: str) -> None:
def _cancel_run(
self,
run_id: str,
*,
after: str | None,
session_scope: AgentAppSessionScope,
binding_id: str,
) -> None:
try:
self._agent_backend_client.cancel_run(run_id)
public_event = self._agent_backend_client.cancel_run_and_wait(run_id, after=after)
for internal_event in self._event_adapter.adapt(public_event):
if (
isinstance(internal_event, AgentBackendRunCancelledInternalEvent)
and internal_event.session_snapshot is not None
):
self._save_session(
scope=session_scope,
binding_id=binding_id,
snapshot=internal_event.session_snapshot,
)
except Exception:
logger.warning("Failed to cancel stopped Agent App backend run: run_id=%s", run_id, exc_info=True)
logger.warning(
"Failed to finish cancelling stopped Agent App backend run: run_id=%s",
run_id,
exc_info=True,
)
def _publish_answer(
self, *, queue_manager: AppQueueManager, model_name: str, answer: str, query: str | None

68
api/core/workflow/nodes/agent_v2/agent_node.py
View File

@ -340,6 +340,17 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
process_data=process_data,
metadata=metadata,
)
# None means no post-exit snapshot was produced; leave the previously stored session snapshot untouched.
if (
isinstance(terminal_event, AgentBackendRunFailedInternalEvent | AgentBackendRunCancelledInternalEvent)
and terminal_event.session_snapshot is not None
):
self._save_session_snapshot(
session_scope=session_scope,
binding_id=stored_session.binding_id,
snapshot=terminal_event.session_snapshot,
metadata=metadata,
)
if exhausted is not None:
# Streaming error / unexpected end — surface immediately without
# retrying because the failure is transport-level.
@ -516,16 +527,20 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
- ``terminal_event``: the first non-stream/non-started internal event,
or ``None`` if the stream ended without one.
- ``transport_failure``: a populated ``StreamCompletedEvent`` when the
stream itself errored (backend/HTTP/protocol fault). Mutually
exclusive with ``terminal_event``.
stream itself errored (backend/HTTP/protocol fault). A cancellation
terminal may accompany it so the caller can persist the final session
snapshot while preserving the original transport failure.
"""
stream_event_count = 0
last_event_id: str | None = None
try:
for public_event in self._agent_backend_client.stream_events(
run_id,
should_stop=self._is_graph_aborted,
):
stream_event_count += 1
if public_event.id is not None:
last_event_id = public_event.id
for internal_event in self._event_adapter.adapt(public_event):
if internal_event.type == AgentBackendInternalEventType.RUN_STARTED:
continue
@ -552,8 +567,12 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
| AgentBackendDeferredToolCallInternalEvent,
):
return internal_event, None
self._cancel_backend_run(run_id, reason="unexpected_event")
return None, self._failure_event(
cancellation = self._cancel_backend_run(
run_id,
reason="unexpected_event",
after=last_event_id,
)
return cancellation, self._failure_event(
inputs=inputs,
process_data=process_data,
metadata=metadata,
@ -561,8 +580,12 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
error_type="agent_backend_stream_error",
)
except AgentBackendError as error:
self._cancel_backend_run(run_id, reason=self._stream_stop_reason())
return None, self._failure_event(
cancellation = self._cancel_backend_run(
run_id,
reason=self._stream_stop_reason(),
after=last_event_id,
)
return cancellation, self._failure_event(
inputs=inputs,
process_data=process_data,
metadata=metadata,
@ -570,8 +593,12 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
error_type=self._agent_backend_error_type(error),
)
except Exception as error:
self._cancel_backend_run(run_id, reason=self._stream_stop_reason())
return None, self._failure_event(
cancellation = self._cancel_backend_run(
run_id,
reason=self._stream_stop_reason(),
after=last_event_id,
)
return cancellation, self._failure_event(
inputs=inputs,
process_data=process_data,
metadata=metadata,
@ -579,8 +606,12 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
error_type="agent_backend_stream_error",
)
self._cancel_backend_run(run_id, reason="stream_ended_without_terminal_event")
return None, None
cancellation = self._cancel_backend_run(
run_id,
reason=self._stream_stop_reason() if self._is_graph_aborted() else "stream_ended_without_terminal_event",
after=last_event_id,
)
return cancellation, None
def _is_graph_aborted(self) -> bool:
"""Let Agent SSE consumption observe GraphEngine's cooperative abort state."""
@ -592,14 +623,25 @@ class DifyAgentNode(Node[DifyAgentNodeData]):
def _stream_stop_reason(self) -> str:
return "workflow_graph_aborted" if self._is_graph_aborted() else "event_stream_failed"
def _cancel_backend_run(self, run_id: str, *, reason: str) -> None:
def _cancel_backend_run(
self,
run_id: str,
*,
reason: str,
after: str | None,
) -> AgentBackendRunCancelledInternalEvent | None:
try:
self._agent_backend_client.cancel_run(
public_event = self._agent_backend_client.cancel_run_and_wait(
run_id,
CancelRunRequest(reason=reason, message="Workflow Agent event consumption stopped"),
after=after,
)
for internal_event in self._event_adapter.adapt(public_event):
if isinstance(internal_event, AgentBackendRunCancelledInternalEvent):
return internal_event
except Exception:
logger.warning("Failed to cancel Workflow Agent backend run: run_id=%s", run_id, exc_info=True)
logger.warning("Failed to finish cancelling Workflow Agent backend run: run_id=%s", run_id, exc_info=True)
return None
@staticmethod
def _record_type_check_metadata(metadata: dict[str, Any], outcome: OutputTypeCheckOutcome) -> None:

25
api/tests/unit_tests/clients/agent_backend/test_client.py
View File

@ -9,6 +9,8 @@ from dify_agent.protocol import (
CancelRunResponse,
CreateRunRequest,
CreateRunResponse,
RunCancelledEvent,
RunCancelledEventData,
RunEvent,
RunStartedEvent,
RunStatusResponse,
@ -51,6 +53,7 @@ def _request() -> CreateRunRequest:
class _SuccessfulClient:
stream_options: tuple[int | None, object, Callable[[], bool] | None] | None = None
cancel_after: str | None = None
def create_run_sync(self, request: CreateRunRequest) -> CreateRunResponse:
assert isinstance(request, CreateRunRequest)
@ -60,6 +63,21 @@ class _SuccessfulClient:
del request
return CancelRunResponse(run_id=run_id, status="cancelled")
def cancel_run_and_wait_sync(
self,
run_id: str,
request: CancelRunRequest | None = None,
*,
after: str | None = None,
) -> RunCancelledEvent:
self.cancel_after = after
request = request or CancelRunRequest()
return RunCancelledEvent(
id="2-0",
run_id=run_id,
data=RunCancelledEventData(reason=request.reason, message=request.message),
)
def stream_events_sync(
self,
run_id: str,
@ -94,13 +112,20 @@ def test_dify_agent_backend_run_client_delegates_sync_methods() -> None:
created = client.create_run(_request())
cancelled = client.cancel_run(created.run_id)
cancelled_event = client.cancel_run_and_wait(
created.run_id,
CancelRunRequest(reason="stopped"),
after="1-0",
)
events = list(client.stream_events(created.run_id, should_stop=should_stop))
status = client.wait_run(created.run_id)
assert created.run_id == "run-1"
assert cancelled.status == "cancelled"
assert cancelled_event.data.reason == "stopped"
assert events[0].type == "run_started"
assert status.status == "succeeded"
assert wrapped.cancel_after == "1-0"
assert wrapped.stream_options == (2, _STREAM_TIMEOUT_UNSET, should_stop)

130
api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py
View File

@ -18,6 +18,8 @@ from dify_agent.protocol import (
CancelRunRequest,
CancelRunResponse,
PydanticAIStreamRunEvent,
RunCancelledEvent,
RunCancelledEventData,
RunEvent,
RunFailedEvent,
RunFailedEventData,
@ -39,6 +41,7 @@ from sqlalchemy import event, select
from sqlalchemy.orm import Session
from clients.agent_backend import (
AgentBackendError,
AgentBackendRunEventAdapter,
AgentBackendRunFailedError,
AgentBackendRunFailedInternalEvent,
@ -112,12 +115,38 @@ class _RecordingFakeAgentBackendRunClient(FakeAgentBackendRunClient):
def __init__(self, **kwargs: Any) -> None:
super().__init__(**kwargs)
self.cancelled_run_ids: list[str] = []
self.cancel_after: list[str | None] = []
@override
def cancel_run(self, run_id: str, request: CancelRunRequest | None = None) -> CancelRunResponse:
self.cancelled_run_ids.append(run_id)
return super().cancel_run(run_id, request=request)
@override
def cancel_run_and_wait(
self,
run_id: str,
request: CancelRunRequest | None = None,
*,
after: str | None = None,
) -> RunCancelledEvent:
self.cancel_after.append(after)
return super().cancel_run_and_wait(run_id, request=request, after=after)
class _CancelAndWaitFailingClient(_RecordingFakeAgentBackendRunClient):
@override
def cancel_run_and_wait(
self,
run_id: str,
request: CancelRunRequest | None = None,
*,
after: str | None = None,
) -> RunCancelledEvent:
del request
self.cancel_after.append(after)
raise RuntimeError(f"failed to finish cancelling {run_id}")
class _RunLimitBindingLostFakeAgentBackendRunClient(FakeAgentBackendRunClient):
@override
@ -143,6 +172,38 @@ class _RunLimitBindingLostFakeAgentBackendRunClient(FakeAgentBackendRunClient):
)
class _TerminalWithoutSnapshotFakeAgentBackendRunClient(FakeAgentBackendRunClient):
def __init__(self, *, terminal_type: str) -> None:
super().__init__()
self.terminal_type = terminal_type
@override
def stream_events(
self,
run_id: str,
*,
after: str | None = None,
should_stop: Callable[[], bool] | None = None,
) -> Iterator[RunEvent]:
del after, should_stop
created_at = datetime(2026, 1, 1, tzinfo=UTC)
yield RunStartedEvent(id="1-0", run_id=run_id, created_at=created_at)
if self.terminal_type == "failed":
yield RunFailedEvent(
id="2-0",
run_id=run_id,
created_at=created_at,
data=RunFailedEventData(error="failed without snapshot"),
)
else:
yield RunCancelledEvent(
id="2-0",
run_id=run_id,
created_at=created_at,
data=RunCancelledEventData(reason="cancelled without snapshot"),
)
class _StreamingFakeAgentBackendRunClient(FakeAgentBackendRunClient):
@override
def stream_events(
@ -443,6 +504,26 @@ class _FakeSessionStore:
self.saved.append((scope, binding_id, snapshot, pending_form_id, pending_tool_call_id))
class _ExplodingSessionStore(_FakeSessionStore):
def __init__(self, loaded: CompositorSessionSnapshot | None = None) -> None:
super().__init__(loaded=loaded)
self.save_attempts: list[CompositorSessionSnapshot | None] = []
@override
def save_active_snapshot(
self,
*,
scope: AgentAppSessionScope,
binding_id: str,
snapshot: CompositorSessionSnapshot | None,
pending_form_id: str | None = None,
pending_tool_call_id: str | None = None,
) -> None:
del scope, binding_id, pending_form_id, pending_tool_call_id
self.save_attempts.append(snapshot)
raise RuntimeError("session save failed")
class _MonotonicClock:
def __init__(self, *values: float) -> None:
self._values = list(values)
@ -772,6 +853,7 @@ def test_streaming_turn_cancels_after_persisting_seen_agent_answer(
assert len(rows) == 1
assert rows[0].answer == "hello "
assert client.cancelled_run_ids == ["fake-run-1"]
assert client.cancel_after == ["3-0"]
def test_tool_result_without_identity_does_not_attach_to_previous_tool(
@ -1162,8 +1244,40 @@ def test_failed_run_raises_agent_backend_error() -> None:
with pytest.raises(AgentBackendRunFailedError, match="fake failure .*agent_run_id=fake-run-1"):
_run(_runner(client, store), qm)
# No message-end on failure; no snapshot saved.
# No message-end on failure; post-exit session state is still saved.
assert not [e for e in qm.events if isinstance(e, QueueMessageEndEvent)]
assert store.saved[0][2] == CompositorSessionSnapshot(layers=[])
@pytest.mark.parametrize("outcome", ["failed", "stopped"])
def test_snapshot_save_failure_preserves_original_app_outcome(outcome: str) -> None:
store = _ExplodingSessionStore()
queue_manager: _FakeQueueManager = _FakeQueueManager() if outcome == "failed" else _StoppedQueueManager()
client = FakeAgentBackendRunClient(
scenario=FakeAgentBackendScenario.FAILED if outcome == "failed" else FakeAgentBackendScenario.SUCCESS
)
expected_error = AgentBackendRunFailedError if outcome == "failed" else GenerateTaskStoppedError
with pytest.raises(expected_error, match="fake failure" if outcome == "failed" else None):
_run(_runner(client, store), queue_manager)
assert store.save_attempts == [CompositorSessionSnapshot(layers=[])]
@pytest.mark.parametrize(
("terminal_type", "expected_error"),
[("failed", AgentBackendRunFailedError), ("cancelled", AgentBackendError)],
)
def test_terminal_without_snapshot_preserves_prior_app_session_without_write(
terminal_type: str,
expected_error: type[Exception],
) -> None:
store = _FakeSessionStore()
client = _TerminalWithoutSnapshotFakeAgentBackendRunClient(terminal_type=terminal_type)
with pytest.raises(expected_error):
_run(_runner(client, store), _FakeQueueManager())
assert store.saved == []
@ -1233,7 +1347,7 @@ def test_agent_backend_failure_to_exception_prefers_run_failure_type_over_known_
}
def test_stopped_task_cancels_agent_backend_run_and_skips_session_save() -> None:
def test_stopped_task_waits_for_cancelled_snapshot_and_saves_session() -> None:
client = _RecordingFakeAgentBackendRunClient()
store = _FakeSessionStore()
qm = _StoppedQueueManager()
@ -1242,6 +1356,18 @@ def test_stopped_task_cancels_agent_backend_run_and_skips_session_save() -> None
_run(_runner(client, store), qm)
assert client.cancelled_run_ids == ["fake-run-1"]
assert len(store.saved) == 1
assert store.saved[0][2] == CompositorSessionSnapshot(layers=[])
def test_cancel_and_wait_failure_preserves_stopped_app_outcome() -> None:
client = _CancelAndWaitFailingClient()
store = _FakeSessionStore()
with pytest.raises(GenerateTaskStoppedError):
_run(_runner(client, store), _StoppedQueueManager())
assert client.cancel_after == [None]
assert store.saved == []

153
api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py
View File

@ -12,7 +12,11 @@ from dify_agent.protocol import (
CancelRunRequest,
CancelRunResponse,
PydanticAIStreamRunEvent,
RunCancelledEvent,
RunCancelledEventData,
RunEvent,
RunFailedEvent,
RunFailedEventData,
RunStartedEvent,
RunSucceededEvent,
RunSucceededEventData,
@ -21,6 +25,7 @@ from pydantic_ai.messages import PartDeltaEvent, TextPartDelta
from clients.agent_backend import (
AgentBackendInternalEventType,
AgentBackendRunCancelledInternalEvent,
AgentBackendRunEventAdapter,
AgentBackendStreamError,
AgentBackendStreamInternalEvent,
@ -205,6 +210,25 @@ class FakeSessionStore:
self.saved.append((scope, binding_id, snapshot, pending_form_id, pending_tool_call_id))
class ExplodingSessionStore(FakeSessionStore):
def __init__(self, snapshot: CompositorSessionSnapshot | None = None) -> None:
super().__init__(snapshot=snapshot)
self.save_attempts: list[CompositorSessionSnapshot | None] = []
def save_active_snapshot(
self,
*,
scope: WorkflowAgentSessionScope,
binding_id: str,
snapshot: CompositorSessionSnapshot | None,
pending_form_id: str | None = None,
pending_tool_call_id: str | None = None,
) -> None:
del scope, binding_id, pending_form_id, pending_tool_call_id
self.save_attempts.append(snapshot)
raise RuntimeError("simulated DB failure")
class FileOutputBackendClient(FakeAgentBackendRunClient):
output_payload: dict[str, object]
@ -259,6 +283,7 @@ class FailingStreamBackendClient(FakeAgentBackendRunClient):
def __init__(self) -> None:
super().__init__()
self.cancel_requests: list[CancelRunRequest | None] = []
self.cancel_after: list[str | None] = []
def stream_events(
self,
@ -275,6 +300,50 @@ class FailingStreamBackendClient(FakeAgentBackendRunClient):
self.cancel_requests.append(request)
return CancelRunResponse(run_id=run_id, status="cancelled")
def cancel_run_and_wait(
self,
run_id: str,
request: CancelRunRequest | None = None,
*,
after: str | None = None,
) -> RunCancelledEvent:
self.cancel_after.append(after)
return super().cancel_run_and_wait(run_id, request=request, after=after)
class FailingAfterStartedStreamBackendClient(FailingStreamBackendClient):
def stream_events(
self,
run_id: str,
*,
after: str | None = None,
should_stop: Callable[[], bool] | None = None,
) -> Iterator[RunEvent]:
del after, should_stop
yield RunStartedEvent(id="cursor-1", run_id=run_id)
raise AgentBackendStreamError("stream failed after started")
class TerminalWithoutSnapshotBackendClient(FakeAgentBackendRunClient):
def __init__(self, *, terminal_type: str) -> None:
super().__init__()
self.terminal_type = terminal_type
def _events(self, run_id: str):
if self.terminal_type == "failed":
terminal: RunEvent = RunFailedEvent(
id="2-0",
run_id=run_id,
data=RunFailedEventData(error="failed without snapshot"),
)
else:
terminal = RunCancelledEvent(
id="2-0",
run_id=run_id,
data=RunCancelledEventData(reason="cancelled without snapshot"),
)
return (RunStartedEvent(id="1-0", run_id=run_id), terminal)
class EmptyStreamBackendClient(FailingStreamBackendClient):
def stream_events(
@ -601,13 +670,15 @@ def test_agent_node_run_normalizes_declared_array_file_output_with_canonical_map
def test_agent_node_run_maps_failed_agent_backend_run_to_node_result():
events = list(_node(scenario=FakeAgentBackendScenario.FAILED)._run())
store = FakeSessionStore()
events = list(_node(scenario=FakeAgentBackendScenario.FAILED, session_store=store)._run())
assert len(events) == 1
result = cast(StreamCompletedEvent, events[0]).node_run_result
assert result.status == WorkflowNodeExecutionStatus.FAILED
assert result.error == "fake failure"
assert result.error_type == "unit_test"
assert store.saved[0][2] == CompositorSessionSnapshot(layers=[])
def test_agent_node_saves_success_snapshot_and_reuses_existing_snapshot():
@ -636,12 +707,7 @@ def test_agent_node_run_when_session_store_save_raises_records_persist_error_in_
``session_snapshot_persist_error`` in the agent_backend metadata so the
incident is observable from the workflow_node_executions record."""
class _ExplodingSessionStore(FakeSessionStore):
def save_active_snapshot(self, **kwargs): # type: ignore[override]
del kwargs
raise RuntimeError("simulated DB failure")
store = _ExplodingSessionStore()
store = ExplodingSessionStore()
events = list(_node(session_store=store)._run())
assert len(events) == 1
@ -652,6 +718,46 @@ def test_agent_node_run_when_session_store_save_raises_records_persist_error_in_
assert agent_backend["session_snapshot_persist_error"] == "workflow_agent_workspace_store_error"
@pytest.mark.parametrize("failure_kind", ["backend", "transport"])
def test_agent_node_snapshot_save_failure_preserves_original_failure(failure_kind: str) -> None:
store = ExplodingSessionStore()
client = (
FakeAgentBackendRunClient(scenario=FakeAgentBackendScenario.FAILED)
if failure_kind == "backend"
else FailingStreamBackendClient()
)
events = list(_node(agent_backend_client=client, session_store=store)._run())
result = cast(StreamCompletedEvent, events[0]).node_run_result
assert result.status == WorkflowNodeExecutionStatus.FAILED
if failure_kind == "backend":
assert (result.error, result.error_type) == ("fake failure", "unit_test")
else:
assert result.error == "stream reconnect attempts exhausted"
assert result.error_type == "agent_backend_stream_error"
agent_backend = result.metadata[WorkflowNodeExecutionMetadataKey.AGENT_LOG]["agent_backend"]
assert agent_backend["session_snapshot_persisted"] is False
assert agent_backend["session_snapshot_persist_error"] == "workflow_agent_workspace_store_error"
assert store.save_attempts == [CompositorSessionSnapshot(layers=[])]
@pytest.mark.parametrize("terminal_type", ["failed", "cancelled"])
def test_agent_node_terminal_without_snapshot_preserves_prior_session_without_write(terminal_type: str) -> None:
store = FakeSessionStore()
events = list(
_node(
agent_backend_client=TerminalWithoutSnapshotBackendClient(terminal_type=terminal_type),
session_store=store,
)._run()
)
result = cast(StreamCompletedEvent, events[0]).node_run_result
assert result.status == WorkflowNodeExecutionStatus.FAILED
assert store.saved == []
def test_agent_node_paused_run_requests_workflow_pause_and_persists_snapshot():
store = FakeSessionStore()
node = _node(scenario=FakeAgentBackendScenario.PAUSED, session_store=store)
@ -819,7 +925,7 @@ def test_agent_node_cancels_backend_run_when_stream_fails():
metadata={"agent_backend": {}},
)
assert terminal is None
assert isinstance(terminal, AgentBackendRunCancelledInternalEvent)
assert failure is not None
assert failure.node_run_result.process_data == {"workflow_agent_binding_id": "binding-1"}
assert len(client.cancel_requests) == 1
@ -827,6 +933,23 @@ def test_agent_node_cancels_backend_run_when_stream_fails():
assert client.cancel_requests[0].reason == "event_stream_failed"
def test_agent_node_forwards_last_stream_cursor_when_cancelling_after_failure() -> None:
client = FailingAfterStartedStreamBackendClient()
node = _node(agent_backend_client=client)
terminal, failure = node._consume_event_stream(
"run-1",
inputs={},
process_data={"workflow_agent_binding_id": "binding-1"},
metadata={"agent_backend": {}},
)
assert isinstance(terminal, AgentBackendRunCancelledInternalEvent)
assert failure is not None
assert failure.node_run_result.error == "stream failed after started"
assert client.cancel_after == ["cursor-1"]
def test_agent_node_cancels_backend_run_when_stream_ends_without_terminal_event():
client = EmptyStreamBackendClient()
node = _node(agent_backend_client=client)
@ -838,7 +961,7 @@ def test_agent_node_cancels_backend_run_when_stream_ends_without_terminal_event(
metadata={"agent_backend": {}},
)
assert terminal is None
assert isinstance(terminal, AgentBackendRunCancelledInternalEvent)
assert failure is None
assert client.cancel_requests[0] is not None
assert client.cancel_requests[0].reason == "stream_ended_without_terminal_event"
@ -855,7 +978,7 @@ def test_agent_node_cancels_backend_run_when_stream_raises_unexpected_error():
metadata={"agent_backend": {}},
)
assert terminal is None
assert isinstance(terminal, AgentBackendRunCancelledInternalEvent)
assert failure is not None
assert failure.node_run_result.error == "unexpected stream failure"
assert failure.node_run_result.process_data == {"workflow_agent_binding_id": "binding-1"}
@ -877,16 +1000,18 @@ def test_agent_node_uses_graph_abort_reason_when_cancel_request_fails(caplog):
assert terminal is None
assert failure is not None
assert failure.node_run_result.error == "stream reconnect attempts exhausted"
assert failure.node_run_result.error_type == "agent_backend_stream_error"
assert client.cancel_requests[0] is not None
assert client.cancel_requests[0].reason == "workflow_graph_aborted"
assert "Failed to cancel Workflow Agent backend run" in caplog.text
assert "Failed to finish cancelling Workflow Agent backend run" in caplog.text
def test_agent_node_cancels_backend_run_for_unexpected_internal_event():
client = FakeAgentBackendRunClient()
node = _node(agent_backend_client=client)
node._agent_backend_client.cancel_run = MagicMock( # type: ignore[method-assign]
return_value=CancelRunResponse(run_id="run-1", status="cancelled")
node._agent_backend_client.cancel_run_and_wait = MagicMock( # type: ignore[method-assign]
return_value=RunCancelledEvent(run_id="run-1")
)
node._event_adapter.adapt = MagicMock( # type: ignore[method-assign]
return_value=[SimpleNamespace(type=AgentBackendInternalEventType.RUN_FAILED)]
@ -905,7 +1030,7 @@ def test_agent_node_cancels_backend_run_for_unexpected_internal_event():
"Unexpected internal event type <AgentBackendInternalEventType.RUN_FAILED: 'run_failed'>"
)
assert failure.node_run_result.process_data == {"workflow_agent_binding_id": "binding-1"}
node._agent_backend_client.cancel_run.assert_called_once()
node._agent_backend_client.cancel_run_and_wait.assert_called_once()
def test_agent_node_records_stream_usage_metadata():

5
dify-agent/docs/dify-agent/concepts/run-lifecycle/index.md
View File

@ -76,8 +76,9 @@ current run. Callers control whether each layer is suspended or deleted through
`CreateRunRequest.on_exit`.
Exit signals control the **layer lifecycle state**, not the execution state of an
`agent run`. The default policy is `suspend`, so a successful `agent run` returns
a reusable `session_snapshot`.
`agent run`. The default policy is `suspend`, so any run that enters and exits its
compositor context can return a reusable `session_snapshot`, including failed or
cancelled runs. Failures or cancellations before entry have no new snapshot.
### Default: suspend layers

51
dify-agent/docs/dify-agent/guide/index.md
View File

@ -279,26 +279,29 @@ run as failed with `error_type: "agent_run_limit_exceeded"`.
During FastAPI shutdown the scheduler rejects new runs, waits up to
`DIFY_AGENT_SHUTDOWN_GRACE_SECONDS` for active tasks, then cancels remaining tasks
and attempts to finalize them as failed. Success, failure, cancellation, and this
shutdown path all use one atomic Redis transition: only the first transition from
`running` appends a terminal event and updates the run record. A later terminal
attempt leaves both the record and event stream unchanged. A hard process crash
can still leave active runs stuck as `running`; there is no in-service recovery
or worker handoff.
and attempts to finalize them as failed. Success and failure use an atomic Redis
transition. Cancellation first atomically records a private intent; after the
owner exits the runner, a second atomic transition appends `run_cancelled`,
updates the run record, and deletes the intent. The first accepted success,
failure, or cancellation intent wins. A hard process crash can still leave
active runs, including runs with accepted cancellation intent, stuck as
`running`; there is no in-service recovery or worker handoff.
Horizontal scaling is possible by running multiple API processes against the same
Redis prefix, but each process executes only the runs it accepted. Redis provides
shared status/event visibility, not load balancing or queued-job recovery. The
cancel endpoint can atomically accept a running run on any process. The process
that owns the runner observes the shared `run_cancelled` event, then cancels and
cleans up its local task. The HTTP response confirms that logical cancellation is
durable; local runner cleanup may still be in progress. Retrying a cancellation
after the run is already `cancelled` is idempotent.
that owns the runner observes the private cancellation-intent stream, cancels
and cleans up its local task, and only then emits `run_cancelled`. The HTTP
response confirms that cancellation intent is durable; `GET /runs/{run_id}` may
still report `running` until cleanup finishes. Retrying an accepted or completed
cancellation is idempotent.
Atomic terminal finalization currently assumes the configured Redis URL targets
one Redis deployment that can execute both run keys in a Lua script. The existing
record and event key names are unchanged and do not contain a shared Redis
Cluster hash tag, so Redis Cluster is not supported for this transition. During
one Redis deployment that can execute all run-coordination keys in a Lua script.
The record and event key names are unchanged, and cancellation adds a private
cancel-intent key. These keys do not contain a shared Redis Cluster hash tag, so
Redis Cluster is not supported for this transition. During
a rolling upgrade, older processes can still use the former split event/status
writes; treat the single-terminal invariant as active only after those processes
have exited. Deploy atomic terminal finalization everywhere first, then ensure
@ -336,8 +339,10 @@ effective prompts are rejected during create-run validation before the run is
persisted or scheduled.
There is no Pydantic AI history layer. To resume Agenton layer state, pass the
`session_snapshot` from a previous `run_succeeded.data` payload together with a
composition that has the same layer names and order.
`session_snapshot` from a previous terminal event together with a composition
that has the same layer names and order. Success always contains a snapshot.
Failure and cancellation contain one only when compositor entry succeeded and
layer exit completed; otherwise callers should retain their previous snapshot.
## Observing runs
@ -349,8 +354,11 @@ progress:
Failed records can also expose a stable machine-readable `error_type` alongside
the diagnostic `error` text.
- `POST /runs/{run_id}/cancel` atomically accepts cancellation on any API process
and emits `run_cancelled`; it returns `409` only when a success/failure terminal
already won. Runner cleanup continues asynchronously on the owner process.
and returns immediately. `CancelRunResponse.status == "cancelled"` acknowledges
a durably accepted cancellation intent, not completed runner cleanup. Callers
that require cleanup-complete state or its session snapshot must await the
public `run_cancelled` event or use `cancel_run_and_wait`. The endpoint returns
`409` only when a success/failure terminal already won.
- `GET /runs/{run_id}/events` polls the Redis Stream event log with `after` and
`next_cursor` cursors.
- `GET /runs/{run_id}/events/sse` replays and streams events over SSE. The SSE
@ -366,12 +374,15 @@ end with `run_cancelled`. Each run can append at most one of these terminal
events. Event envelopes retain `id`, `run_id`, `type`, `data`, and `created_at`;
`data` is typed per event type,
including Pydantic AI's `AgentStreamEvent` payload for `pydantic_ai_event` and a
terminal `run_succeeded.data` object containing a `CompositorSessionSnapshot` for
resumption. A successful run has exactly one active result branch: JSON-safe
terminal event may contain a `CompositorSessionSnapshot` for resumption.
`run_succeeded` always contains it; `run_failed` and `run_cancelled` contain it
only when compositor entry succeeded, layer exit completed, and a post-exit
snapshot was actually produced. A successful run has exactly one active result branch: JSON-safe
`output` for final answers, or `deferred_tool_call` when a layer such as
`dify.ask_human` ends the current agent run with an external deferred tool call.
Failed event payloads contain the diagnostic `error`, optional source-specific
`reason`, and optional stable `error_type`. Pydantic AI request/step budget
`reason`, optional stable `error_type`, and optional `session_snapshot`.
Cancelled payloads likewise may contain `session_snapshot`. Pydantic AI request/step budget
exhaustion enforced by Dify Agent is reported as
`error_type: "agent_run_limit_exceeded"`; consumers should branch on that value
rather than parsing the error text. The Dify Agent-owned wall-clock run deadline

43
dify-agent/src/dify_agent/client/_client.py
View File

@ -43,6 +43,7 @@ from dify_agent.protocol import (
DestroyExecutionBindingRequest,
HomeSnapshotResponse,
RUN_EVENT_ADAPTER,
RunCancelledEvent,
RunEvent,
RunEventsResponse,
RunStatusResponse,
@ -384,8 +385,8 @@ class Client:
async def cancel_run(self, run_id: str, request: CancelRunRequest | None = None) -> CancelRunResponse:
"""Request explicit cancellation for ``run_id``.
Acceptance atomically persists the cancelled state. The process executing
the run observes that state and performs runner cleanup asynchronously.
Acceptance atomically persists cancellation intent. The process executing
the run publishes ``run_cancelled`` after runner cleanup completes.
"""
request_model = request or CancelRunRequest()
try:
@ -417,6 +418,44 @@ class Client:
raise DifyAgentClientError(f"cancel_run_sync request failed: {exc}") from exc
return _parse_model_response(response, CancelRunResponse)
async def cancel_run_and_wait(
self,
run_id: str,
request: CancelRunRequest | None = None,
*,
after: str | None = None,
) -> RunCancelledEvent:
"""Request cancellation and wait for its public terminal event."""
_ = await self.cancel_run(run_id, request)
resume_after = after
if after is not None and (await self.get_run(run_id)).status == "cancelled":
resume_after = None
async for event in self.stream_events(run_id, after=resume_after):
if isinstance(event, RunCancelledEvent):
return event
if event.type in _TERMINAL_EVENT_TYPES:
raise DifyAgentClientError(f"run {run_id!r} finished with {event.type!r} before cancellation")
raise DifyAgentStreamError(f"run {run_id!r} stream ended before run_cancelled")
def cancel_run_and_wait_sync(
self,
run_id: str,
request: CancelRunRequest | None = None,
*,
after: str | None = None,
) -> RunCancelledEvent:
"""Synchronous variant of ``cancel_run_and_wait``."""
_ = self.cancel_run_sync(run_id, request)
resume_after = after
if after is not None and self.get_run_sync(run_id).status == "cancelled":
resume_after = None
for event in self.stream_events_sync(run_id, after=resume_after):
if isinstance(event, RunCancelledEvent):
return event
if event.type in _TERMINAL_EVENT_TYPES:
raise DifyAgentClientError(f"run {run_id!r} finished with {event.type!r} before cancellation")
raise DifyAgentStreamError(f"run {run_id!r} stream ended before run_cancelled")
async def get_run(self, run_id: str) -> RunStatusResponse:
"""Return the current status for ``run_id`` or raise a mapped client error."""
try:

9
dify-agent/src/dify_agent/protocol/schemas.py
View File

@ -21,9 +21,10 @@ by ``DIFY_AGENT_MODEL_LAYER_ID``, the optional history layer named by
``DIFY_AGENT_HISTORY_LAYER_ID``, and the optional structured output layer named
by ``DIFY_AGENT_OUTPUT_LAYER_ID``. Request-level ``on_exit`` signals decide
whether each active layer is suspended or deleted when the run exits, with
suspend as the default so successful terminal events can include resumable
snapshots. Successful runs always publish the resumable Agenton session snapshot
on the terminal ``run_succeeded`` event together with either the final JSON-safe
suspend as the default so terminal events can include resumable snapshots.
Successful runs always publish the resumable Agenton session snapshot on the
terminal ``run_succeeded`` event; failed and cancelled runs publish it when the
compositor context was entered and exited. Success includes either the final JSON-safe
``output`` or a deferred external ``deferred_tool_call`` payload. Session
snapshots carry only layer lifecycle/runtime state in
compositor order; they do not persist output-layer config. Resumed
@ -329,6 +330,7 @@ class RunFailedEventData(BaseModel):
error: str
error_type: RunFailureType | None = None
reason: str | None = None
session_snapshot: CompositorSessionSnapshot | None = None
model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")
@ -338,6 +340,7 @@ class RunCancelledEventData(BaseModel):
reason: str | None = None
message: str | None = None
session_snapshot: CompositorSessionSnapshot | None = None
model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")

19
dify-agent/src/dify_agent/runtime/cancellation.py Normal file
View File

@ -0,0 +1,19 @@
"""Private cancellation coordination types shared by schedulers and run stores."""
from datetime import datetime
from typing import ClassVar
from pydantic import BaseModel, ConfigDict
class RunCancellationIntent(BaseModel):
"""The first accepted request to cancel one running run."""
reason: str | None = None
message: str | None = None
requested_at: datetime
model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")
__all__ = ["RunCancellationIntent"]

41
dify-agent/src/dify_agent/runtime/event_sink.py
View File

@ -1,9 +1,10 @@
"""Event sink contracts used by the runner and storage adapters.
Non-terminal events remain append-only. Terminal events use ``finalize_run`` so
the event and matching run status are committed as one compare-and-set
transition. Tests can use ``InMemoryRunEventSink`` without Redis; production
storage implements the same contract with Redis streams in
Non-terminal events remain append-only. Successful and failed terminal events
use ``finalize_run`` so the event and matching run status are committed as one
compare-and-set transition. Cancellation has a dedicated intent-aware finalizer.
Tests can use ``InMemoryRunEventSink`` without Redis; production storage
implements the same contract with Redis streams in
``dify_agent.storage.redis_run_store``.
"""
@ -21,8 +22,6 @@ from dify_agent.protocol.schemas import (
EmptyRunEventData,
PydanticAIStreamRunEvent,
RunEvent,
RunCancelledEvent,
RunCancelledEventData,
RunFailedEvent,
RunFailedEventData,
RunFailureType,
@ -35,7 +34,7 @@ from dify_agent.protocol.schemas import (
_UNSET = object()
TerminalRunEvent: TypeAlias = RunSucceededEvent | RunFailedEvent | RunCancelledEvent
TerminalRunEvent: TypeAlias = RunSucceededEvent | RunFailedEvent
NonTerminalRunEvent: TypeAlias = RunStartedEvent | PydanticAIStreamRunEvent
@ -105,8 +104,6 @@ def terminal_event_status_fields(
return "succeeded", None, None
case RunFailedEvent():
return "failed", event.data.error, event.data.error_type
case RunCancelledEvent():
return "cancelled", event.data.message or event.data.reason, None
async def emit_run_event(
@ -188,29 +185,18 @@ async def emit_run_failed(
error: str,
error_type: RunFailureType | None = None,
reason: str | None = None,
session_snapshot: CompositorSessionSnapshot | None = None,
) -> RunFinalizationResult:
"""Finalize a run with a failed terminal event."""
return await sink.finalize_run(
RunFailedEvent(
run_id=run_id,
data=RunFailedEventData(error=error, error_type=error_type, reason=reason),
created_at=utc_now(),
),
)
async def emit_run_cancelled(
sink: RunEventSink,
*,
run_id: str,
reason: str | None = None,
message: str | None = None,
) -> RunFinalizationResult:
"""Finalize a run with a cancelled terminal event."""
return await sink.finalize_run(
RunCancelledEvent(
run_id=run_id,
data=RunCancelledEventData(reason=reason, message=message),
data=RunFailedEventData(
error=error,
error_type=error_type,
reason=reason,
session_snapshot=session_snapshot,
),
created_at=utc_now(),
),
)
@ -223,7 +209,6 @@ __all__ = [
"RunFinalizationResult",
"TerminalRunEvent",
"emit_pydantic_ai_event",
"emit_run_cancelled",
"emit_run_event",
"emit_run_failed",
"emit_run_started",

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

@ -20,10 +20,11 @@ from typing import Protocol
import httpx
from agenton.compositor import LayerProviderInput
from dify_agent.protocol.schemas import CancelRunRequest, CancelRunResponse, CreateRunRequest
from agenton.compositor import CompositorSessionSnapshot, LayerProviderInput
from dify_agent.protocol.schemas import CancelRunRequest, CancelRunResponse, CreateRunRequest, RunStatus
from dify_agent.runtime.cancellation import RunCancellationIntent
from dify_agent.runtime.compositor_factory import create_default_layer_providers
from dify_agent.runtime.event_sink import RunEventSink, emit_run_cancelled, emit_run_failed
from dify_agent.runtime.event_sink import RunEventSink, RunFinalizationResult, emit_run_failed
from dify_agent.runtime.runner import DEFAULT_AGENT_RUN_TIMEOUT_SECONDS, AgentRunRunner
from dify_agent.server.schemas import RunRecord
@ -45,14 +46,37 @@ class RunStore(RunEventSink, Protocol):
"""Persist a new run record and return it with status ``running``."""
...
async def wait_for_cancellation(self, run_id: str) -> bool:
"""Wait for a terminal state and report whether cancellation won."""
async def request_cancellation(self, run_id: str, request: CancelRunRequest) -> RunStatus:
"""Persist the first cancellation intent and return the current status."""
...
async def get_cancellation_intent(self, run_id: str) -> RunCancellationIntent | None:
"""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 finalize_cancellation(
self,
run_id: str,
intent: RunCancellationIntent,
*,
session_snapshot: CompositorSessionSnapshot | None = None,
) -> RunFinalizationResult:
"""Publish cancellation after the owner runner has exited."""
...
class RunnableRun(Protocol):
"""Executable unit for one scheduled run."""
@property
def terminal_session_snapshot(self) -> CompositorSessionSnapshot | None:
"""Return the post-exit snapshot for the current invocation, if available."""
...
async def run(self) -> None:
"""Run until terminal status/events have been written or cancellation occurs."""
...
@ -125,14 +149,9 @@ class RunScheduler:
async def cancel_run(self, run_id: str, request: CancelRunRequest) -> CancelRunResponse:
"""Persist an idempotent cancellation without relying on local task ownership."""
finalization = await emit_run_cancelled(
self.store,
run_id=run_id,
reason=request.reason,
message=request.message,
)
if finalization.status != "cancelled":
raise RunCancellationConflictError(f"run already finished with status {finalization.status!r}")
status = await self.store.request_cancellation(run_id, request)
if status in {"succeeded", "failed"}:
raise RunCancellationConflictError(f"run already finished with status {status!r}")
return CancelRunResponse(run_id=run_id, status="cancelled")
async def shutdown(self) -> None:
@ -141,18 +160,14 @@ class RunScheduler:
self.stopping = True
if not self.active_tasks:
return
tasks_by_run_id = dict(self.active_tasks)
done, pending = await asyncio.wait(tasks_by_run_id.values(), timeout=self.shutdown_grace_seconds)
del done
tasks = tuple(self.active_tasks.values())
_done, pending = await asyncio.wait(tasks, timeout=self.shutdown_grace_seconds)
if not pending:
return
pending_run_ids = [run_id for run_id, task in tasks_by_run_id.items() if task in pending]
for task in pending:
_ = task.cancel()
_ = await asyncio.gather(*pending, return_exceptions=True)
for run_id in pending_run_ids:
await self._mark_cancelled_run_failed(run_id)
async def _run_record(self, record: RunRecord, request: CreateRunRequest) -> None:
"""Supervise one local runner and its durable cancellation observer."""
@ -163,41 +178,92 @@ class RunScheduler:
self.store.wait_for_cancellation(record.run_id),
name=f"dify-agent-cancellation-observer-{record.run_id}",
)
async def cancel_runner_and_wait() -> None:
if not cancel_requested.is_set() and not runner_task.done():
cancel_requested.set()
_ = runner_task.cancel()
_ = await asyncio.shield(asyncio.gather(runner_task, return_exceptions=True))
try:
_ = await asyncio.wait((runner_task, observer_task), return_when=asyncio.FIRST_COMPLETED)
if observer_task.done():
try:
cancellation_won = observer_task.result()
intent = observer_task.result()
except Exception as exc:
cancel_requested.set()
await self._cancel_and_wait(runner_task, reinject=True)
_ = await emit_run_failed(
await cancel_runner_and_wait()
finalization = await emit_run_failed(
self.store,
run_id=record.run_id,
error=f"run cancellation observer failed: {exc}",
reason="cancellation_observer",
session_snapshot=runner.terminal_session_snapshot,
)
if not finalization.applied and finalization.status == "running":
intent = await self.store.get_cancellation_intent(record.run_id)
if intent is not None:
_ = await self.store.finalize_cancellation(
record.run_id,
intent,
session_snapshot=runner.terminal_session_snapshot,
)
raise
if cancellation_won:
cancel_requested.set()
await self._cancel_and_wait(runner_task, reinject=True)
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
else:
await runner_task
runner_error: Exception | None = None
try:
await runner_task
except Exception as exc:
runner_error = exc
intent = await self.store.get_cancellation_intent(record.run_id)
if intent is not None:
_ = await self.store.finalize_cancellation(
record.run_id,
intent,
session_snapshot=runner.terminal_session_snapshot,
)
if runner_error is not None:
raise runner_error
except asyncio.CancelledError:
cancel_requested.set()
await self._cancel_and_wait(observer_task)
await self._cancel_and_wait(runner_task, reinject=True)
await cancel_runner_and_wait()
intent = await self.store.get_cancellation_intent(record.run_id)
if intent is not None:
_ = await self.store.finalize_cancellation(
record.run_id,
intent,
session_snapshot=runner.terminal_session_snapshot,
)
else:
finalization = await self._mark_cancelled_run_failed(
record.run_id,
session_snapshot=runner.terminal_session_snapshot,
)
if finalization is not None and not finalization.applied and finalization.status == "running":
intent = await self.store.get_cancellation_intent(record.run_id)
if intent is not None:
_ = await self.store.finalize_cancellation(
record.run_id,
intent,
session_snapshot=runner.terminal_session_snapshot,
)
raise
except Exception:
logger.exception("scheduled run failed", extra={"run_id": record.run_id})
finally:
await self._cancel_and_wait(observer_task)
if not runner_task.done():
cancel_requested.set()
await self._cancel_and_wait(runner_task, reinject=True)
await cancel_runner_and_wait()
def _create_runner(
self,
@ -234,25 +300,31 @@ class RunScheduler:
_ = self.active_tasks.pop(run_id, None)
@staticmethod
async def _cancel_and_wait(task: asyncio.Task[object], *, reinject: bool = False) -> None:
"""Cancel and reap a child task, with bounded reinjection for runners."""
async def _cancel_and_wait(task: asyncio.Task[object]) -> None:
"""Cancel a child task once and await its complete exit."""
if not task.done():
_ = task.cancel()
if reinject:
for _attempt in range(2):
await asyncio.sleep(0)
if task.done():
break
_ = task.cancel()
_ = await asyncio.gather(task, return_exceptions=True)
async def _mark_cancelled_run_failed(self, run_id: str) -> None:
async def _mark_cancelled_run_failed(
self,
run_id: str,
*,
session_snapshot: CompositorSessionSnapshot | None = None,
) -> RunFinalizationResult | None:
"""Best-effort failure event/status for shutdown-cancelled runs."""
message = "run cancelled during server shutdown"
try:
_ = await emit_run_failed(self.store, run_id=run_id, error=message, reason="shutdown")
return await emit_run_failed(
self.store,
run_id=run_id,
error=message,
reason="shutdown",
session_snapshot=session_snapshot,
)
except Exception:
logger.exception("failed to mark cancelled run failed", extra={"run_id": run_id})
return None
__all__ = ["RunCancellationConflictError", "RunScheduler", "SchedulerStoppingError"]

15
dify-agent/src/dify_agent/runtime/runner.py
View File

@ -184,6 +184,7 @@ class AgentRunRunner:
dify_api_http_client: httpx.AsyncClient
is_cancelled: Callable[[], bool]
run_timeout_seconds: float
_terminal_session_snapshot: CompositorSessionSnapshot | None
def __init__(
self,
@ -205,9 +206,16 @@ class AgentRunRunner:
self.layer_providers = layer_providers if layer_providers is not None else create_default_layer_providers()
self.is_cancelled = is_cancelled or (lambda: False)
self.run_timeout_seconds = run_timeout_seconds
self._terminal_session_snapshot = None
@property
def terminal_session_snapshot(self) -> CompositorSessionSnapshot | None:
"""Return the snapshot captured after the current compositor context exited."""
return self._terminal_session_snapshot
async def run(self) -> None:
"""Execute the run and emit the documented event sequence."""
self._terminal_session_snapshot = None
if self.is_cancelled():
return
_ = await emit_run_started(self.sink, run_id=self.run_id)
@ -224,6 +232,7 @@ class AgentRunRunner:
error=message,
error_type=error_type,
reason=reason,
session_snapshot=self._terminal_session_snapshot,
)
if finalization.applied:
raise
@ -287,6 +296,7 @@ class AgentRunRunner:
deferred_tool_call: DeferredToolCallPayload | None = None
result_kind: Literal["output", "deferred_tool_call"] | None = None
usage: AgentRunUsage | None = None
run = None
try:
async with compositor.enter(configs=layer_configs, session_snapshot=self.request.session_snapshot) as run:
entered_run = True
@ -383,8 +393,11 @@ class AgentRunRunner:
if not entered_run:
raise AgentRunValidationError(str(exc)) from exc
raise
finally:
if entered_run and run is not None:
self._terminal_session_snapshot = run.session_snapshot
if run.session_snapshot is None:
if run is None or run.session_snapshot is None:
raise RuntimeError("Agenton run did not produce a session snapshot after exit.")
if result_kind is None:
raise RuntimeError("Agent run did not resolve either a final output or a deferred tool call.")

7
dify-agent/src/dify_agent/storage/redis_keys.py
View File

@ -11,4 +11,9 @@ def run_events_key(prefix: str, run_id: str) -> str:
return f"{prefix}:runs:{run_id}:events"
__all__ = ["run_events_key", "run_record_key"]
def run_cancel_intent_key(prefix: str, run_id: str) -> str:
"""Return the private Redis stream key holding one cancellation intent."""
return f"{prefix}:runs:{run_id}:cancel-intent"
__all__ = ["run_cancel_intent_key", "run_events_key", "run_record_key"]

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

@ -1,8 +1,8 @@
"""Redis-backed run records and per-run event streams.
"""Redis-backed run records, event streams, and private cancellation intents.
The store writes status-only run records as JSON strings and events as Redis
streams. HTTP event cursors are Redis stream ids; ``0-0`` means replay from the
beginning for polling and SSE. Records and streams share one retention window
beginning for polling and SSE. Records, event streams, and intents share one retention window
that is refreshed when status or event data is written. Execution is scheduled
in-process by ``dify_agent.runtime.run_scheduler``; Redis is not a job queue, and
create-run payloads are never persisted because layer config may include
@ -14,7 +14,18 @@ from typing import cast
from redis.asyncio import Redis
from dify_agent.protocol.schemas import RUN_EVENT_ADAPTER, RunEvent, RunEventsResponse, RunStatus
from agenton.compositor import CompositorSessionSnapshot
from dify_agent.protocol.schemas import (
RUN_EVENT_ADAPTER,
CancelRunRequest,
RunCancelledEvent,
RunCancelledEventData,
RunEvent,
RunEventsResponse,
RunStatus,
utc_now,
)
from dify_agent.runtime.cancellation import RunCancellationIntent
from dify_agent.runtime.event_sink import (
NonTerminalRunEvent,
RunEventSink,
@ -24,7 +35,7 @@ from dify_agent.runtime.event_sink import (
)
from dify_agent.server.schemas import RunRecord, new_run_id
from dify_agent.server.settings import DEFAULT_RUN_RETENTION_SECONDS
from dify_agent.storage.redis_keys import run_events_key, run_record_key
from dify_agent.storage.redis_keys import run_cancel_intent_key, run_events_key, run_record_key
_TERMINAL_RUN_EVENT_TYPES = {"run_succeeded", "run_failed", "run_cancelled"}
@ -44,6 +55,10 @@ if record.status ~= "running" then
return {0, tostring(record.status), ""}
end
if redis.call("EXISTS", KEYS[3]) == 1 then
return {-2, "running", ""}
end
record.status = ARGV[1]
record.updated_at = ARGV[2]
if ARGV[3] == "1" then
@ -66,6 +81,67 @@ return {1, ARGV[1], event_id}
"""
_REQUEST_CANCELLATION_SCRIPT = """
local record_json = redis.call("GET", KEYS[1])
if not record_json then
return {-1, ""}
end
local record = cjson.decode(record_json)
if record.status == "succeeded" or record.status == "failed" then
return {0, tostring(record.status)}
end
if record.status == "cancelled" then
return {1, "cancelled"}
end
if redis.call("EXISTS", KEYS[2]) == 1 then
return {1, "running"}
end
local ttl = tonumber(ARGV[2])
redis.call("XADD", KEYS[2], "*", "payload", ARGV[1])
redis.call("EXPIRE", KEYS[2], ttl)
redis.call("EXPIRE", KEYS[1], ttl)
redis.call("EXPIRE", KEYS[3], ttl)
return {1, "running"}
"""
_FINALIZE_CANCELLATION_SCRIPT = """
local record_json = redis.call("GET", KEYS[1])
if not record_json then
return {-1, "", ""}
end
local record = cjson.decode(record_json)
if record.status == "cancelled" then
return {0, "cancelled", ""}
end
if record.status ~= "running" then
return {0, tostring(record.status), ""}
end
if redis.call("EXISTS", KEYS[2]) == 0 then
return {-2, "running", ""}
end
record.status = "cancelled"
record.updated_at = ARGV[1]
if ARGV[2] == "1" then
record.error = ARGV[3]
else
record.error = cjson.null
end
record.error_type = cjson.null
local ttl = tonumber(ARGV[5])
local event_id = redis.call("XADD", KEYS[3], "*", "payload", ARGV[4])
redis.call("DEL", KEYS[2])
redis.call("EXPIRE", KEYS[3], ttl)
redis.call("SET", KEYS[1], cjson.encode(record), "EX", ttl)
return {1, "cancelled", event_id}
"""
class RedisRunStore(RunEventSink):
"""Async Redis implementation for run records and event logs.
@ -130,16 +206,17 @@ class RedisRunStore(RunEventSink):
return event_id.decode() if isinstance(event_id, bytes) else str(event_id)
async def finalize_run(self, event: TerminalRunEvent) -> RunFinalizationResult:
"""Atomically append the first terminal event and update its run record."""
"""Atomically append the first success/failure event and update its run record."""
status, error, error_type = terminal_event_status_fields(event)
payload = RUN_EVENT_ADAPTER.dump_json(event, exclude={"id"}).decode()
evaluation = cast(
Awaitable[object],
self.redis.eval(
_FINALIZE_RUN_SCRIPT,
2,
3,
run_record_key(self.prefix, event.run_id),
run_events_key(self.prefix, event.run_id),
run_cancel_intent_key(self.prefix, event.run_id),
status,
event.created_at.isoformat(),
"1" if error is not None else "0",
@ -164,8 +241,39 @@ class RedisRunStore(RunEventSink):
event_id=event_id,
)
async def wait_for_cancellation(self, run_id: str) -> bool:
"""Wait until cancellation or another terminal state wins for one run.
async def request_cancellation(self, run_id: str, request: CancelRunRequest) -> RunStatus:
"""Atomically persist the first cancellation intent for a running run."""
intent = RunCancellationIntent(
reason=request.reason,
message=request.message,
requested_at=utc_now(),
)
evaluation = cast(
Awaitable[object],
self.redis.eval(
_REQUEST_CANCELLATION_SCRIPT,
3,
run_record_key(self.prefix, run_id),
run_cancel_intent_key(self.prefix, run_id),
run_events_key(self.prefix, run_id),
intent.model_dump_json(),
str(self.run_retention_seconds),
),
)
result = cast(list[object], await evaluation)
if int(cast(int | bytes | str, result[0])) == -1:
raise RunNotFoundError(run_id)
return cast(RunStatus, _decode_redis_text(result[1]))
async def get_cancellation_intent(self, run_id: str) -> RunCancellationIntent | None:
"""Return the accepted private cancellation intent, if one exists."""
entries = await self.redis.xrange(run_cancel_intent_key(self.prefix, run_id), count=1)
if not entries:
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
@ -176,19 +284,76 @@ class RedisRunStore(RunEventSink):
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 record.status == "cancelled"
return None
intent = await self.get_cancellation_intent(run_id)
if intent is not None:
return intent
while True:
response = await self.redis.xread({events_key: cursor}, block=0, count=100)
for _stream_name, entries in response:
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 True
return None
if event.type in {"run_succeeded", "run_failed"}:
return False
return None
async def finalize_cancellation(
self,
run_id: str,
intent: RunCancellationIntent,
*,
session_snapshot: CompositorSessionSnapshot | None = None,
) -> RunFinalizationResult:
"""Atomically publish cancellation after the owner runner has exited."""
event = RunCancelledEvent(
run_id=run_id,
data=RunCancelledEventData(
reason=intent.reason,
message=intent.message,
session_snapshot=session_snapshot,
),
created_at=utc_now(),
)
payload = RUN_EVENT_ADAPTER.dump_json(event, exclude={"id"}).decode()
error = event.data.message or event.data.reason
evaluation = cast(
Awaitable[object],
self.redis.eval(
_FINALIZE_CANCELLATION_SCRIPT,
3,
run_record_key(self.prefix, run_id),
run_cancel_intent_key(self.prefix, run_id),
run_events_key(self.prefix, run_id),
event.created_at.isoformat(),
"1" if error is not None else "0",
error or "",
payload,
str(self.run_retention_seconds),
),
)
result = cast(list[object], await evaluation)
applied = int(cast(int | bytes | str, result[0]))
if applied == -1:
raise RunNotFoundError(run_id)
return RunFinalizationResult(
applied=applied == 1,
status=cast(RunStatus, _decode_redis_text(result[1])),
event_id=_decode_redis_text(result[2]) or None,
)
async def get_events(self, run_id: str, *, after: str = "0-0", limit: int = 100) -> RunEventsResponse:
"""Read a bounded page of events after ``after`` cursor."""
@ -235,6 +400,13 @@ class RedisRunStore(RunEventSink):
event = RUN_EVENT_ADAPTER.validate_json(cast(str, payload))
return event.model_copy(update={"id": event_id, "run_id": run_id})
@staticmethod
def _decode_cancellation_intent(fields: dict[object, object]) -> RunCancellationIntent:
payload = fields.get(b"payload") or fields.get("payload")
if isinstance(payload, bytes):
payload = payload.decode()
return RunCancellationIntent.model_validate_json(cast(str, payload))
def _decode_redis_text(value: object) -> str:
return value.decode() if isinstance(value, bytes) else str(value)

189
dify-agent/tests/integration/dify_agent/storage/test_terminal_finalization.py
View File

@ -16,24 +16,36 @@ from agenton.compositor import CompositorSessionSnapshot
from dify_agent.protocol.schemas import (
CancelRunRequest,
CreateRunRequest,
RunCancelledEvent,
RunCancelledEventData,
RunComposition,
RunFailedEvent,
RunFailedEventData,
RunFailureType,
RunStartedEvent,
RunSucceededEvent,
RunSucceededEventData,
utc_now,
)
from dify_agent.runtime.event_sink import TerminalRunEvent, terminal_event_status_fields
from dify_agent.runtime.cancellation import RunCancellationIntent
from dify_agent.runtime.run_scheduler import RunScheduler
from dify_agent.storage.redis_keys import run_events_key, run_record_key
from dify_agent.storage.redis_keys import run_cancel_intent_key, run_events_key, run_record_key
from dify_agent.storage.redis_run_store import RedisRunStore
pytestmark = pytest.mark.integration
def _success_or_failure_event(kind: str, run_id: str) -> RunSucceededEvent | RunFailedEvent:
if kind == "succeeded":
return RunSucceededEvent(
run_id=run_id,
data=RunSucceededEventData(
output="done",
session_snapshot=CompositorSessionSnapshot(layers=[]),
),
)
return RunFailedEvent(run_id=run_id, data=RunFailedEventData(error="model failed"))
@pytest.fixture
def redis_url() -> Iterator[str]:
"""Start an isolated Redis when the binary is available locally."""
@ -82,7 +94,7 @@ def redis_url() -> Iterator[str]:
process.wait(timeout=5)
def test_two_redis_clients_commit_exactly_one_matching_terminal(redis_url: str) -> None:
def test_success_and_cancel_intent_commit_exactly_one_matching_terminal(redis_url: str) -> None:
async def scenario() -> None:
first_client = Redis.from_url(redis_url)
second_client = Redis.from_url(redis_url)
@ -92,43 +104,42 @@ def test_two_redis_clients_commit_exactly_one_matching_terminal(redis_url: str)
second_store = RedisRunStore(second_client, prefix=prefix, run_retention_seconds=retention_seconds)
try:
record = await first_store.create_run()
terminal_events: tuple[TerminalRunEvent, TerminalRunEvent] = (
RunSucceededEvent(
run_id=record.run_id,
data=RunSucceededEventData(
output="done",
session_snapshot=CompositorSessionSnapshot(layers=[]),
),
),
RunCancelledEvent(
run_id=record.run_id,
data=RunCancelledEventData(
reason="concurrent_cancel",
message="cancel accepted",
),
success_event = RunSucceededEvent(
run_id=record.run_id,
data=RunSucceededEventData(
output="done",
session_snapshot=CompositorSessionSnapshot(layers=[]),
),
)
results = await asyncio.gather(
first_store.finalize_run(terminal_events[0]),
second_store.finalize_run(terminal_events[1]),
success_result, cancellation_status = await asyncio.gather(
first_store.finalize_run(success_event),
second_store.request_cancellation(
record.run_id,
CancelRunRequest(reason="concurrent_cancel", message="cancel accepted"),
),
)
assert sum(result.applied for result in results) == 1
winner_index = next(index for index, result in enumerate(results) if result.applied)
winner_event = terminal_events[winner_index]
winner_result = results[winner_index]
expected_status, expected_error, expected_error_type = terminal_event_status_fields(winner_event)
if success_result.applied:
assert cancellation_status == "succeeded"
winner_result = success_result
expected_status = "succeeded"
expected_event_type = "run_succeeded"
else:
assert success_result.status == "running"
assert cancellation_status == "running"
intent = await first_store.get_cancellation_intent(record.run_id)
assert intent is not None
winner_result = await first_store.finalize_cancellation(record.run_id, intent)
assert winner_result.applied is True
expected_status = "cancelled"
expected_event_type = "run_cancelled"
persisted = await first_store.get_run(record.run_id)
page = await second_store.get_events(record.run_id)
assert persisted.status == expected_status
assert persisted.error == expected_error
assert persisted.error_type == expected_error_type
assert persisted.updated_at == winner_event.created_at
assert len(page.events) == 1
assert page.events[0].type == winner_event.type
assert page.events[0].created_at == winner_event.created_at
assert page.events[0].type == expected_event_type
assert page.events[0].id == winner_result.event_id
record_ttl = await first_client.ttl(run_record_key(prefix, record.run_id))
@ -142,6 +153,116 @@ def test_two_redis_clients_commit_exactly_one_matching_terminal(redis_url: str)
asyncio.run(scenario())
@pytest.mark.parametrize("terminal_status", ["succeeded", "failed"])
def test_terminal_first_rejects_late_cancellation(redis_url: str, terminal_status: str) -> None:
async def scenario() -> None:
client = Redis.from_url(redis_url)
prefix = f"terminal-first-{terminal_status}-{uuid4().hex}"
store = RedisRunStore(client, prefix=prefix, run_retention_seconds=60)
try:
record = await store.create_run()
result = await store.finalize_run(_success_or_failure_event(terminal_status, record.run_id))
cancellation_status = await store.request_cancellation(
record.run_id,
CancelRunRequest(reason="late_cancel"),
)
assert result.applied is True
assert cancellation_status == terminal_status
assert await store.get_cancellation_intent(record.run_id) is None
events = await store.get_events(record.run_id)
assert [event.type for event in events.events] == [f"run_{terminal_status}"]
finally:
await client.aclose()
asyncio.run(scenario())
def test_cancellation_intent_lifecycle_and_terminal_exclusion(redis_url: str) -> None:
async def scenario() -> None:
client = Redis.from_url(redis_url)
prefix = f"cancel-intent-lifecycle-{uuid4().hex}"
retention_seconds = 60
store = RedisRunStore(client, prefix=prefix, run_retention_seconds=retention_seconds)
try:
record = await store.create_run()
_ = await store.append_event(RunStartedEvent(run_id=record.run_id))
record_key = run_record_key(prefix, record.run_id)
events_key = run_events_key(prefix, record.run_id)
intent_key = run_cancel_intent_key(prefix, record.run_id)
_ = await client.expire(record_key, 1)
_ = await client.expire(events_key, 1)
first_status = await store.request_cancellation(
record.run_id,
CancelRunRequest(reason="first", message="first message"),
)
duplicate_status = await store.request_cancellation(
record.run_id,
CancelRunRequest(reason="second", message="second message"),
)
intent = await store.get_cancellation_intent(record.run_id)
assert first_status == duplicate_status == "running"
assert intent is not None
assert (intent.reason, intent.message) == ("first", "first message")
for key in (record_key, events_key, intent_key):
assert 0 < await client.ttl(key) <= retention_seconds
success = await store.finalize_run(_success_or_failure_event("succeeded", record.run_id))
failure = await store.finalize_run(_success_or_failure_event("failed", record.run_id))
assert (success.applied, success.status) == (False, "running")
assert (failure.applied, failure.status) == (False, "running")
assert [event.type for event in (await store.get_events(record.run_id)).events] == ["run_started"]
first_finalization = await store.finalize_cancellation(
record.run_id,
intent,
session_snapshot=CompositorSessionSnapshot(layers=[]),
)
repeated_finalization = await store.finalize_cancellation(record.run_id, intent)
post_terminal_status = await store.request_cancellation(
record.run_id,
CancelRunRequest(reason="after_finalization"),
)
events = await store.get_events(record.run_id)
assert first_finalization.applied is True
assert repeated_finalization.applied is False
assert repeated_finalization.status == "cancelled"
assert post_terminal_status == "cancelled"
assert [event.type for event in events.events].count("run_cancelled") == 1
assert await client.exists(intent_key) == 0
for key in (record_key, events_key):
assert 0 < await client.ttl(key) <= retention_seconds
finally:
await client.aclose()
asyncio.run(scenario())
def test_cancellation_finalization_without_intent_is_unapplied(redis_url: str) -> None:
async def scenario() -> None:
client = Redis.from_url(redis_url)
store = RedisRunStore(client, prefix=f"cancel-without-intent-{uuid4().hex}", run_retention_seconds=60)
try:
record = await store.create_run()
result = await store.finalize_cancellation(
record.run_id,
RunCancellationIntent(reason="not-accepted", requested_at=utc_now()),
)
assert result.applied is False
assert result.status == "running"
assert (await store.get_events(record.run_id)).events == []
assert (await store.get_run(record.run_id)).status == "running"
finally:
await client.aclose()
asyncio.run(scenario())
def test_classified_failure_persists_matching_record_and_event_error_type(redis_url: str) -> None:
async def scenario() -> None:
client = Redis.from_url(redis_url)
@ -178,6 +299,10 @@ def test_non_owner_scheduler_cancellation_stops_owner_runner(redis_url: str) ->
self.started = started
self.stopped = stopped
@property
def terminal_session_snapshot(self) -> None:
return None
async def run(self) -> None:
self.started.set()
try:

104
dify-agent/tests/local/dify_agent/client/test_client.py
View File

@ -34,6 +34,7 @@ from dify_agent.protocol import (
DestroyExecutionBindingRequest,
RUN_EVENT_ADAPTER,
RunCancelledEvent,
RunCancelledEventData,
RunEvent,
RunEventsResponse,
RunFailedEvent,
@ -245,6 +246,109 @@ def test_async_methods_and_wait_run_parse_protocol_dtos() -> None:
asyncio.run(scenario())
def test_cancel_run_and_wait_sync_resumes_after_cursor_and_returns_cancelled_snapshot() -> None:
snapshot = CompositorSessionSnapshot(layers=[])
def handler(request: httpx.Request) -> httpx.Response:
if request.method == "POST":
return httpx.Response(202, json={"run_id": "run-1", "status": "cancelled"})
if request.url.path == "/runs/run-1":
return httpx.Response(
200,
json={
"run_id": "run-1",
"status": "running",
"created_at": "2026-08-17T00:00:00Z",
"updated_at": "2026-08-17T00:00:00Z",
},
)
assert request.url.params["after"] == "3-0"
event = RunCancelledEvent(
id="4-0",
run_id="run-1",
data=RunCancelledEventData(reason="stopped", session_snapshot=snapshot),
)
return httpx.Response(200, content=_event_frame(event))
client = Client(
base_url="http://testserver",
sync_http_client=httpx.Client(transport=httpx.MockTransport(handler)),
)
event = client.cancel_run_and_wait_sync(
"run-1",
CancelRunRequest(reason="stopped"),
after="3-0",
)
assert event.data.session_snapshot == snapshot
def test_cancel_run_and_wait_sync_replays_when_cursor_already_points_to_cancelled_event() -> None:
def handler(request: httpx.Request) -> httpx.Response:
if request.method == "POST":
return httpx.Response(202, json={"run_id": "run-1", "status": "cancelled"})
if request.url.path == "/runs/run-1":
return httpx.Response(
200,
json={
"run_id": "run-1",
"status": "cancelled",
"created_at": "2026-08-17T00:00:00Z",
"updated_at": "2026-08-17T00:00:01Z",
},
)
assert request.url.params["after"] == "0-0"
return httpx.Response(200, content=_event_frame(RunCancelledEvent(id="4-0", run_id="run-1")))
client = Client(
base_url="http://testserver",
sync_http_client=httpx.Client(transport=httpx.MockTransport(handler)),
)
event = client.cancel_run_and_wait_sync("run-1", after="4-0")
assert event.id == "4-0"
def test_cancel_run_and_wait_async_returns_cancelled_terminal() -> None:
def handler(request: httpx.Request) -> httpx.Response:
if request.method == "POST":
return httpx.Response(202, json={"run_id": "run-1", "status": "cancelled"})
if request.url.path == "/runs/run-1":
return httpx.Response(200, json=_run_status_json("running"))
assert request.url.params["after"] == "1-0"
return httpx.Response(200, content=_event_frame(RunCancelledEvent(id="2-0", run_id="run-1")))
async def scenario() -> None:
http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
client = Client(base_url="http://testserver", async_http_client=http_client)
event = await client.cancel_run_and_wait("run-1", after="1-0")
assert event.type == "run_cancelled"
await http_client.aclose()
asyncio.run(scenario())
def test_cancel_run_and_wait_async_replays_when_cursor_already_points_to_cancelled_event() -> None:
def handler(request: httpx.Request) -> httpx.Response:
if request.method == "POST":
return httpx.Response(202, json={"run_id": "run-1", "status": "cancelled"})
if request.url.path == "/runs/run-1":
return httpx.Response(200, json=_run_status_json("cancelled"))
assert request.url.params["after"] == "0-0"
return httpx.Response(200, content=_event_frame(RunCancelledEvent(id="4-0", run_id="run-1")))
async def scenario() -> None:
http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
client = Client(base_url="http://testserver", async_http_client=http_client)
event = await client.cancel_run_and_wait("run-1", after="4-0")
assert event.id == "4-0"
await http_client.aclose()
asyncio.run(scenario())
def test_sync_binding_file_methods_post_dtos_and_parse_responses() -> None:
def handler(request: httpx.Request) -> httpx.Response:
if request.url.path == "/execution-bindings/files/list":

28
dify-agent/tests/local/dify_agent/protocol/test_protocol_schemas.py
View File

@ -77,7 +77,12 @@ def test_run_event_adapter_round_trips_typed_variants() -> None:
reason="shutdown",
),
),
RunCancelledEvent(run_id="run-1", data=RunCancelledEventData(reason="user_cancelled")),
RunCancelledEvent(
run_id="run-1",
data=RunCancelledEventData(
reason="user_cancelled",
),
),
]
for event in events:
@ -113,6 +118,27 @@ def test_run_failed_event_error_type_is_optional_and_round_trips() -> None:
assert protocol_exports.RunFailureType is RunFailureType
@pytest.mark.parametrize("event_type", ["run_failed", "run_cancelled"])
def test_non_success_terminal_event_round_trips_optional_snapshot(event_type: str) -> None:
snapshot = CompositorSessionSnapshot(layers=[])
event: RunFailedEvent | RunCancelledEvent
if event_type == "run_failed":
event = RunFailedEvent(
run_id="run-1",
data=RunFailedEventData(error="boom", session_snapshot=snapshot),
)
else:
event = RunCancelledEvent(
run_id="run-1",
data=RunCancelledEventData(reason="stopped", session_snapshot=snapshot),
)
decoded = RUN_EVENT_ADAPTER.validate_json(RUN_EVENT_ADAPTER.dump_json(event))
assert isinstance(decoded, RunFailedEvent | RunCancelledEvent)
assert decoded.data.session_snapshot == snapshot
def test_pydantic_ai_event_data_uses_agent_stream_event_model() -> None:
event = RUN_EVENT_ADAPTER.validate_python(
{

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

@ -15,11 +15,16 @@ from dify_agent.protocol import DIFY_AGENT_MODEL_LAYER_ID, DIFY_AGENT_OUTPUT_LAY
from dify_agent.protocol.schemas import (
CancelRunRequest,
CreateRunRequest,
RunCancelledEvent,
RunCancelledEventData,
RunComposition,
RunEvent,
RunFailedEvent,
RunLayerSpec,
RunStatus,
utc_now,
)
from dify_agent.runtime.cancellation import RunCancellationIntent
from dify_agent.runtime.event_sink import (
NonTerminalRunEvent,
RunFinalizationResult,
@ -95,6 +100,7 @@ class FakeStore:
errors: dict[str, str | None]
error_types: dict[str, RunFailureType | None]
terminal_changes: dict[str, asyncio.Event]
cancellation_intents: dict[str, RunCancellationIntent]
def __init__(self) -> None:
self.records = {}
@ -103,6 +109,7 @@ class FakeStore:
self.errors = {}
self.error_types = {}
self.terminal_changes = {}
self.cancellation_intents = {}
async def create_run(self) -> RunRecord:
run_id = f"run-{len(self.records) + 1}"
@ -130,6 +137,8 @@ class FakeStore:
current_status = self.statuses[event.run_id]
if current_status != "running":
return RunFinalizationResult(applied=False, status=current_status)
if event.run_id in self.cancellation_intents:
return RunFinalizationResult(applied=False, status="running")
status, error, error_type = terminal_event_status_fields(event)
event_id = str(len(self.events[event.run_id]) + 1)
@ -140,10 +149,55 @@ class FakeStore:
self.terminal_changes[event.run_id].set()
return RunFinalizationResult(applied=True, status=status, event_id=event_id)
async def wait_for_cancellation(self, run_id: str) -> bool:
while self.statuses[run_id] == "running":
async def request_cancellation(self, run_id: str, request: CancelRunRequest) -> RunStatus:
status = self.statuses[run_id]
if status != "running":
return status
if run_id not in self.cancellation_intents:
self.cancellation_intents[run_id] = RunCancellationIntent(
reason=request.reason,
message=request.message,
requested_at=utc_now(),
)
self.terminal_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.statuses[run_id] == "cancelled"
return self.cancellation_intents.get(run_id)
async def finalize_cancellation(
self,
run_id: str,
intent: RunCancellationIntent,
*,
session_snapshot: CompositorSessionSnapshot | None = None,
) -> RunFinalizationResult:
current_status = self.statuses[run_id]
if current_status != "running":
return RunFinalizationResult(applied=False, status=current_status)
if run_id not in self.cancellation_intents:
return RunFinalizationResult(applied=False, status="running")
event = RunCancelledEvent(
run_id=run_id,
data=RunCancelledEventData(
reason=intent.reason,
message=intent.message,
session_snapshot=session_snapshot,
),
)
event_id = str(len(self.events[run_id]) + 1)
self.events[run_id].append(event.model_copy(update={"id": event_id}))
self.statuses[run_id] = "cancelled"
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)
class SlowCreateStore(FakeStore):
@ -174,7 +228,7 @@ class TrackingStore(FakeStore):
if not pause_observer:
self.release_observer.set()
async def wait_for_cancellation(self, run_id: str) -> bool:
async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None:
self.observer_started.set()
try:
await self.release_observer.wait()
@ -192,7 +246,7 @@ class FailingObserverStore(FakeStore):
self.fail_observer = fail_observer
self.observer_finished = asyncio.Event()
async def wait_for_cancellation(self, run_id: str) -> bool:
async def wait_for_cancellation(self, run_id: str) -> RunCancellationIntent | None:
del run_id
try:
await self.fail_observer.wait()
@ -201,10 +255,27 @@ class FailingObserverStore(FakeStore):
self.observer_finished.set()
class CancellationDuringShutdownFailureStore(FakeStore):
async def finalize_run(self, event: TerminalRunEvent) -> RunFinalizationResult:
if isinstance(event, RunFailedEvent) and event.data.reason == "shutdown":
_ = await self.request_cancellation(
event.run_id,
CancelRunRequest(reason="concurrent_shutdown_cancel"),
)
return await super().finalize_run(event)
class SnapshotlessRunner:
@property
def terminal_session_snapshot(self) -> CompositorSessionSnapshot | None:
return None
class ControlledRunner:
started: asyncio.Event
release: asyncio.Event
finished: asyncio.Event | None
_terminal_session_snapshot: CompositorSessionSnapshot
def __init__(
self,
@ -216,6 +287,11 @@ class ControlledRunner:
self.started = started
self.release = release
self.finished = finished
self._terminal_session_snapshot = CompositorSessionSnapshot(layers=[])
@property
def terminal_session_snapshot(self) -> CompositorSessionSnapshot:
return self._terminal_session_snapshot
async def run(self) -> None:
_ = self.started.set()
@ -226,24 +302,16 @@ class ControlledRunner:
self.finished.set()
class SwallowOneCancellationRunner:
started: asyncio.Event
first_cancellation: asyncio.Event
def __init__(self, *, started: asyncio.Event, first_cancellation: asyncio.Event) -> None:
class PreEnterBlockingRunner(SnapshotlessRunner):
def __init__(self, *, started: 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()
self.started.set()
await asyncio.Event().wait()
class SuccessThenWaitRunner:
class SuccessThenWaitRunner(SnapshotlessRunner):
def __init__(
self,
*,
@ -269,7 +337,7 @@ class SuccessThenWaitRunner:
await self.release.wait()
class IgnoreCancellationThenSucceedRunner:
class IgnoreCancellationThenSucceedRunner(SnapshotlessRunner):
def __init__(
self,
*,
@ -300,12 +368,12 @@ class IgnoreCancellationThenSucceedRunner:
session_snapshot=CompositorSessionSnapshot(layers=[]),
)
assert result.applied is False
assert result.status == "cancelled"
assert result.status == "running"
finally:
self.finished.set()
class ReleaseThenSucceedRunner:
class ReleaseThenSucceedRunner(SnapshotlessRunner):
def __init__(
self,
*,
@ -336,7 +404,7 @@ class ReleaseThenSucceedRunner:
self.finished.set()
class CompetingFailureRunner:
class CompetingFailureRunner(SnapshotlessRunner):
def __init__(
self,
*,
@ -362,7 +430,7 @@ class CompetingFailureRunner:
self.failure_attempted.set()
class FinalizeSuccessOnCancellationRunner:
class FinalizeSuccessOnCancellationRunner(SnapshotlessRunner):
def __init__(self, *, store: FakeStore, run_id: str, started: asyncio.Event) -> None:
self.store = store
self.run_id = run_id
@ -454,6 +522,37 @@ def test_shutdown_marks_unfinished_runs_failed_and_appends_event() -> None:
asyncio.run(scenario())
def test_shutdown_failure_finalization_yields_to_concurrent_cancellation_intent() -> None:
async def scenario() -> None:
store = CancellationDuringShutdownFailureStore()
started = asyncio.Event()
async with httpx.AsyncClient() as client:
scheduler = RunScheduler(
store=store,
plugin_daemon_http_client=client,
dify_api_http_client=client,
shutdown_grace_seconds=0,
runner_factory=lambda _record, _request: ControlledRunner(
started=started,
release=asyncio.Event(),
),
)
record = await scheduler.create_run(_request())
await asyncio.wait_for(started.wait(), timeout=1)
await scheduler.shutdown()
assert store.statuses[record.run_id] == "cancelled"
assert record.run_id not in store.cancellation_intents
assert [event.type for event in store.events[record.run_id]] == ["run_cancelled"]
terminal = store.events[record.run_id][0]
assert isinstance(terminal, RunCancelledEvent)
assert terminal.data.reason == "concurrent_shutdown_cancel"
assert terminal.data.session_snapshot == CompositorSessionSnapshot(layers=[])
asyncio.run(scenario())
def test_cancellation_observer_failure_stops_runner_and_finalizes_failed() -> None:
async def scenario() -> None:
fail_observer = asyncio.Event()
@ -489,6 +588,46 @@ def test_cancellation_observer_failure_stops_runner_and_finalizes_failed() -> No
asyncio.run(scenario())
def test_cancellation_observer_failure_finalizes_concurrent_intent_after_runner_exit() -> None:
async def scenario() -> None:
fail_observer = asyncio.Event()
store = FailingObserverStore(fail_observer=fail_observer)
runner_started = asyncio.Event()
runner_finished = 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: ControlledRunner(
started=runner_started,
release=asyncio.Event(),
finished=runner_finished,
),
)
record = await scheduler.create_run(_request())
supervisor_task = scheduler.active_tasks[record.run_id]
await asyncio.wait_for(runner_started.wait(), timeout=1)
response = await scheduler.cancel_run(
record.run_id,
CancelRunRequest(reason="workflow_aborted", message="outer workflow stopped"),
)
fail_observer.set()
await asyncio.wait_for(supervisor_task, timeout=1)
assert response.status == "cancelled"
assert runner_finished.is_set()
assert store.statuses[record.run_id] == "cancelled"
assert record.run_id not in store.cancellation_intents
assert [event.type for event in store.events[record.run_id]] == ["run_cancelled"]
terminal = store.events[record.run_id][0]
assert isinstance(terminal, RunCancelledEvent)
assert terminal.data.session_snapshot == CompositorSessionSnapshot(layers=[])
asyncio.run(scenario())
def test_non_owner_cancel_run_stops_owner_task_and_persists_cancelled_terminal() -> None:
async def scenario() -> None:
store = TrackingStore()
@ -522,10 +661,13 @@ def test_non_owner_cancel_run_stops_owner_task_and_persists_cancelled_terminal()
assert response.status == "cancelled"
assert remote_scheduler.active_tasks == {}
await asyncio.wait_for(owner_task, timeout=1)
assert store.statuses[record.run_id] == "cancelled"
assert store.errors[record.run_id] == "outer workflow stopped"
assert [event.type for event in store.events[record.run_id]] == ["run_cancelled"]
await asyncio.wait_for(owner_task, timeout=1)
terminal = store.events[record.run_id][0]
assert isinstance(terminal, RunCancelledEvent)
assert terminal.data.session_snapshot == CompositorSessionSnapshot(layers=[])
assert runner_finished.is_set()
assert store.observer_finished.is_set()
await asyncio.sleep(0)
@ -538,37 +680,38 @@ def test_non_owner_cancel_run_stops_owner_task_and_persists_cancelled_terminal()
asyncio.run(scenario())
def test_owner_observer_reinjects_cancellation_consumed_by_runner() -> None:
def test_pre_enter_cancellation_does_not_copy_input_session_snapshot() -> None:
async def scenario() -> None:
store = FakeStore()
started = asyncio.Event()
first_cancellation = asyncio.Event()
request = _request()
request.session_snapshot = CompositorSessionSnapshot(
layers=[
LayerSessionSnapshot(
name="prior",
lifecycle_state=LifecycleState.SUSPENDED,
runtime_state={"value": "prior"},
)
]
)
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,
),
runner_factory=lambda _record, _request: PreEnterBlockingRunner(started=started),
)
record = await scheduler.create_run(_request())
supervisor_task = scheduler.active_tasks[record.run_id]
record = await scheduler.create_run(request)
supervisor = scheduler.active_tasks[record.run_id]
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,
)
_ = await scheduler.cancel_run(record.run_id, CancelRunRequest(reason="pre_enter_cancel"))
await asyncio.wait_for(supervisor, timeout=1)
assert response.status == "cancelled"
assert store.statuses[record.run_id] == "cancelled"
assert [event.type for event in store.events[record.run_id]] == ["run_cancelled"]
await asyncio.wait_for(first_cancellation.wait(), timeout=1)
await asyncio.wait_for(supervisor_task, timeout=1)
await asyncio.sleep(0)
assert scheduler.active_tasks == {}
terminal = store.events[record.run_id][0]
assert isinstance(terminal, RunCancelledEvent)
assert request.session_snapshot is not None
assert terminal.data.session_snapshot is None
asyncio.run(scenario())
@ -634,8 +777,9 @@ def test_cancelled_terminal_survives_shutdown_while_runner_cleanup_is_pending()
response = await scheduler.cancel_run(record.run_id, CancelRunRequest(reason="workflow_aborted"))
assert response.status == "cancelled"
assert store.statuses[record.run_id] == "cancelled"
assert [event.type for event in store.events[record.run_id]] == ["run_cancelled"]
assert store.statuses[record.run_id] == "running"
assert store.events[record.run_id] == []
assert record.run_id in store.cancellation_intents
await asyncio.wait_for(store.observer_finished.wait(), timeout=1)
assert supervisor_task.done() is False
shutdown_task = asyncio.create_task(scheduler.shutdown())

63
dify-agent/tests/local/dify_agent/runtime/test_runner.py
View File

@ -67,7 +67,7 @@ from dify_agent.protocol.schemas import (
RunLayerSpec,
RunSucceededEvent,
)
from dify_agent.runtime.event_sink import InMemoryRunEventSink, emit_run_cancelled
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,
@ -224,7 +224,7 @@ def test_run_failed_error_payload_classifies_usage_limit() -> None:
assert reason is None
def test_cancelled_runner_does_not_overwrite_cancelled_status_with_late_failure(
def test_cancelled_runner_does_not_emit_late_failure(
monkeypatch: pytest.MonkeyPatch,
) -> None:
async def scenario() -> None:
@ -243,20 +243,13 @@ def test_cancelled_runner_does_not_overwrite_cancelled_status_with_late_failure(
async def fail_after_cancel() -> RunSuccessOutcome:
nonlocal cancelled
cancelled = True
_ = await emit_run_cancelled(
sink,
run_id="run-cancelled",
reason="workflow_aborted",
message="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", "run_cancelled"]
assert "run-cancelled" not in sink.statuses
assert [event.type for event in sink.events["run-cancelled"]] == ["run_started"]
asyncio.run(scenario())
@ -924,6 +917,44 @@ def test_runner_does_not_classify_nested_timeout_as_agent_limit(monkeypatch: pyt
assert sink.statuses["run-provider-timeout"] == "failed"
def test_runner_captures_post_exit_snapshot_when_task_is_cancelled(monkeypatch: pytest.MonkeyPatch) -> None:
started = asyncio.Event()
def fake_get_model(_self: DifyPluginLLMLayer, *, http_client: httpx.AsyncClient, agent_run_id: str):
return TestModel(custom_output_text="unused") # pyright: ignore[reportReturnType]
class FakeAgent:
async def run(self, *_args: object, **_kwargs: object) -> None:
started.set()
await asyncio.Event().wait()
monkeypatch.setattr(DifyPluginLLMLayer, "get_model", fake_get_model)
monkeypatch.setattr("dify_agent.runtime.runner.create_agent", lambda *_args, **_kwargs: FakeAgent())
sink = InMemoryRunEventSink()
async def scenario() -> AgentRunRunner:
async with httpx.AsyncClient() as client:
runner = AgentRunRunner(
sink=sink,
request=_request(),
run_id="run-cancel-snapshot",
plugin_daemon_http_client=client,
dify_api_http_client=client,
)
task = asyncio.create_task(runner.run())
await asyncio.wait_for(started.wait(), timeout=1)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
return runner
runner = asyncio.run(scenario())
assert runner.terminal_session_snapshot is not None
assert all(layer.lifecycle_state is LifecycleState.SUSPENDED for layer in runner.terminal_session_snapshot.layers)
assert [event.type for event in sink.events["run-cancel-snapshot"]] == ["run_started"]
def test_runner_emits_deferred_tool_call_and_persists_pending_history(monkeypatch: pytest.MonkeyPatch) -> None:
captured_output_types: list[object] = []
captured_user_prompts: list[object] = []
@ -2206,7 +2237,7 @@ def test_runner_with_empty_history_layer_uses_instructions_and_saves_full_histor
assert all(not isinstance(message, ModelRequest) or message.instructions is None for message in saved_history)
def test_runner_failure_with_history_layer_emits_failed_terminal_event_without_success_snapshot(
def test_runner_failure_with_history_layer_emits_post_exit_snapshot_without_new_history(
monkeypatch: pytest.MonkeyPatch,
) -> None:
model = RecordingTestModel(failure=RuntimeError("boom"))
@ -2239,6 +2270,10 @@ def test_runner_failure_with_history_layer_emits_failed_terminal_event_without_s
assert [event.type for event in sink.events["run-history-failure"]] == ["run_started", "run_failed"]
assert sink.statuses["run-history-failure"] == "failed"
terminal = sink.events["run-history-failure"][-1]
assert isinstance(terminal, RunFailedEvent)
assert terminal.data.session_snapshot is not None
assert _history_messages_from_snapshot(terminal.data.session_snapshot) == stored_history
assert request.session_snapshot is not None
assert _history_messages_from_snapshot(request.session_snapshot) == stored_history
@ -3015,6 +3050,10 @@ def test_runner_rejects_closed_session_snapshot_as_validation_error() -> None:
assert [event.type for event in sink.events["run-closed-snapshot"]] == ["run_started", "run_failed"]
assert sink.statuses["run-closed-snapshot"] == "failed"
terminal = sink.events["run-closed-snapshot"][-1]
assert isinstance(terminal, RunFailedEvent)
assert request.session_snapshot is not None
assert terminal.data.session_snapshot is None
def test_runner_treats_missing_runtime_dependency_as_validation_error() -> None:

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

@ -9,16 +9,19 @@ from pydantic import JsonValue
from agenton.compositor import CompositorSessionSnapshot, LayerSessionSnapshot
from agenton.layers import LifecycleState
from dify_agent.protocol.schemas import (
RUN_EVENT_ADAPTER,
CancelRunRequest,
RunCancelledEvent,
RunCancelledEventData,
RunFailedEvent,
RunFailedEventData,
RunFailureType,
RunStartedEvent,
RunStatus,
RunSucceededEvent,
RunSucceededEventData,
utc_now,
)
from dify_agent.runtime.cancellation import RunCancellationIntent
from dify_agent.runtime.event_sink import RunFinalizationResult
from dify_agent.storage.redis_run_store import DEFAULT_RUN_RETENTION_SECONDS, RedisRunStore, RunNotFoundError
@ -27,12 +30,14 @@ class FakeRedis:
commands: list[tuple[object, ...]]
values: dict[str, object]
streams: dict[str, list[tuple[str, dict[str, object]]]]
eval_result: list[object] | None
def __init__(self) -> None:
self.commands = []
self.values = {}
self.streams = {}
self.stream_changed = asyncio.Event()
self.eval_result = None
async def set(self, key: str, value: object, *, ex: int | None = None) -> None:
self.commands.append(("set", key, value, ex))
@ -109,29 +114,9 @@ class FakeRedis:
async def eval(self, script: str, numkeys: int, *keys_and_args: object) -> list[object]:
self.commands.append(("eval", script, numkeys, *keys_and_args))
assert numkeys == 2
record_key = str(keys_and_args[0])
events_key = str(keys_and_args[1])
status = str(keys_and_args[2])
updated_at = str(keys_and_args[3])
has_error = str(keys_and_args[4]) == "1"
error = str(keys_and_args[5]) if has_error else None
has_error_type = str(keys_and_args[6]) == "1"
error_type = str(keys_and_args[7]) if has_error_type else None
payload = str(keys_and_args[8])
record_json = self.values.get(record_key)
if record_json is None:
return [-1, "", ""]
if isinstance(record_json, bytes):
record_json = record_json.decode()
record = json.loads(cast(str, record_json))
if record["status"] != "running":
return [0, record["status"], ""]
record.update({"status": status, "updated_at": updated_at, "error": error, "error_type": error_type})
event_id = self._append_stream_entry(events_key, {"payload": payload})
self.values[record_key] = json.dumps(record, separators=(",", ":"))
return [1, status, event_id]
if self.eval_result is None:
raise AssertionError("test must configure FakeRedis.eval_result")
return list(self.eval_result)
@staticmethod
def _is_after_min(event_id: str, min_id: str) -> bool:
@ -224,75 +209,76 @@ def test_get_run_accepts_legacy_record_without_error_type() -> None:
assert loaded.error_type is None
def test_finalize_run_atomically_writes_terminal_event_and_status() -> None:
def test_request_cancellation_maps_eval_result_and_arguments() -> None:
redis = FakeRedis()
store = RedisRunStore(redis, prefix="test", run_retention_seconds=60) # pyright: ignore[reportArgumentType]
record = asyncio.run(store.create_run())
redis.commands.clear()
event = RunCancelledEvent(
run_id=record.run_id,
data=RunCancelledEventData(reason="workflow_aborted", message="workflow stopped"),
redis.eval_result = [1, "running"]
status = asyncio.run(
store.request_cancellation(
"run-1",
CancelRunRequest(reason="workflow_aborted", message="workflow stopped"),
)
)
result = asyncio.run(store.finalize_run(event))
updated = asyncio.run(store.get_run(record.run_id))
assert status == "running"
eval_command = redis.commands[-1]
assert eval_command[0] == "eval"
assert eval_command[2] == 3
assert eval_command[3:6] == (
"test:runs:run-1:record",
"test:runs:run-1:cancel-intent",
"test:runs:run-1:events",
)
intent_payload = json.loads(cast(str, eval_command[6]))
assert intent_payload["reason"] == "workflow_aborted"
assert intent_payload["message"] == "workflow stopped"
assert eval_command[7] == "60"
assert result.applied is True
assert result.status == "cancelled"
assert result.event_id == "1-0"
assert updated.status == "cancelled"
assert updated.error == "workflow stopped"
assert updated.error_type is None
assert updated.updated_at == event.created_at
stream_entry_id, stream_fields = redis.streams[f"test:runs:{record.run_id}:events"][0]
assert stream_entry_id == result.event_id
payload = json.loads(cast(str, stream_fields["payload"]))
def test_finalize_cancellation_maps_eval_result_and_arguments() -> None:
redis = FakeRedis()
store = RedisRunStore(redis, prefix="test", run_retention_seconds=60) # pyright: ignore[reportArgumentType]
redis.eval_result = [1, "cancelled", "7-0"]
intent = RunCancellationIntent(
reason="workflow_aborted",
message="workflow stopped",
requested_at=utc_now(),
)
result = asyncio.run(
store.finalize_cancellation(
"run-1",
intent,
session_snapshot=CompositorSessionSnapshot(layers=[]),
)
)
assert result == RunFinalizationResult(applied=True, status="cancelled", event_id="7-0")
eval_command = redis.commands[-1]
assert eval_command[2] == 3
assert eval_command[3:6] == (
"test:runs:run-1:record",
"test:runs:run-1:cancel-intent",
"test:runs:run-1:events",
)
payload = json.loads(cast(str, eval_command[9]))
assert "id" not in payload
assert payload["type"] == "run_cancelled"
assert payload["data"] == {"reason": "workflow_aborted", "message": "workflow stopped"}
assert payload["created_at"] == event.created_at.isoformat().replace("+00:00", "Z")
eval_command = redis.commands[0]
assert eval_command[0] == "eval"
assert eval_command[2] == 2
assert eval_command[-1] == "60"
assert payload["data"] == {
"reason": "workflow_aborted",
"message": "workflow stopped",
"session_snapshot": {"schema_version": 1, "layers": []},
}
assert eval_command[10] == "60"
def test_finalize_run_rejects_a_second_terminal_without_appending_event() -> None:
redis = FakeRedis()
store = RedisRunStore(redis, prefix="test", run_retention_seconds=60) # pyright: ignore[reportArgumentType]
record = asyncio.run(store.create_run())
snapshot = CompositorSessionSnapshot(layers=[])
first = asyncio.run(
store.finalize_run(
RunSucceededEvent(
run_id=record.run_id,
data=RunSucceededEventData(output="done", session_snapshot=snapshot),
)
)
)
second = asyncio.run(
store.finalize_run(
RunCancelledEvent(
run_id=record.run_id,
data=RunCancelledEventData(reason="late_cancel"),
)
)
)
assert first.applied is True
assert second.applied is False
assert second.status == "succeeded"
assert second.event_id is None
assert len(redis.streams[f"test:runs:{record.run_id}:events"]) == 1
def test_finalize_failed_run_derives_error_and_timestamp_from_event() -> None:
def test_finalize_failed_run_maps_eval_result_and_arguments() -> None:
redis = FakeRedis()
store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType]
record = asyncio.run(store.create_run())
redis.eval_result = [1, "failed", "8-0"]
event = RunFailedEvent(
run_id=record.run_id,
run_id="run-1",
data=RunFailedEventData(
error="model failed",
error_type=RunFailureType.AGENT_RUN_LIMIT_EXCEEDED,
@ -301,117 +287,46 @@ def test_finalize_failed_run_derives_error_and_timestamp_from_event() -> None:
)
result = asyncio.run(store.finalize_run(event))
updated = asyncio.run(store.get_run(record.run_id))
assert result.applied is True
assert result.status == "failed"
assert updated.status == "failed"
assert updated.error == "model failed"
assert updated.error_type is RunFailureType.AGENT_RUN_LIMIT_EXCEEDED
assert updated.updated_at == event.created_at
stream_entry = redis.streams[f"test:runs:{record.run_id}:events"][0]
payload = json.loads(cast(str, stream_entry[1]["payload"]))
assert result == RunFinalizationResult(applied=True, status="failed", event_id="8-0")
eval_command = redis.commands[-1]
assert eval_command[3:6] == (
"test:runs:run-1:record",
"test:runs:run-1:events",
"test:runs:run-1:cancel-intent",
)
assert eval_command[6] == "failed"
assert eval_command[8:12] == ("1", "model failed", "1", "agent_run_limit_exceeded")
payload = json.loads(cast(str, eval_command[12]))
assert payload["data"]["error_type"] == "agent_run_limit_exceeded"
def test_two_store_instances_choose_exactly_one_terminal_winner() -> None:
redis = FakeRedis()
first_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType]
second_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType]
async def scenario() -> tuple[list[RunFinalizationResult], RunStatus, list[str]]:
record = await first_store.create_run()
snapshot = CompositorSessionSnapshot(layers=[])
results = await asyncio.gather(
first_store.finalize_run(
RunSucceededEvent(
run_id=record.run_id,
data=RunSucceededEventData(output="done", session_snapshot=snapshot),
)
),
second_store.finalize_run(
RunCancelledEvent(
run_id=record.run_id,
data=RunCancelledEventData(reason="concurrent_cancel"),
)
),
)
persisted = await first_store.get_run(record.run_id)
page = await second_store.get_events(record.run_id)
return list(results), persisted.status, [event.type for event in page.events]
results, status, event_types = asyncio.run(scenario())
assert sum(result.applied for result in results) == 1
assert len(event_types) == 1
assert (status, event_types[0]) in {
("succeeded", "run_succeeded"),
("cancelled", "run_cancelled"),
}
def test_failure_and_cancellation_compete_for_one_terminal() -> None:
redis = FakeRedis()
failure_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType]
cancellation_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType]
async def scenario() -> tuple[list[RunFinalizationResult], RunStatus, list[str]]:
record = await failure_store.create_run()
results = await asyncio.gather(
failure_store.finalize_run(
RunFailedEvent(
run_id=record.run_id,
data=RunFailedEventData(error="model failed", reason="model_error"),
)
),
cancellation_store.finalize_run(
RunCancelledEvent(
run_id=record.run_id,
data=RunCancelledEventData(reason="concurrent_cancel"),
)
),
)
persisted = await failure_store.get_run(record.run_id)
page = await cancellation_store.get_events(record.run_id)
return list(results), persisted.status, [event.type for event in page.events]
results, status, event_types = asyncio.run(scenario())
assert sum(result.applied for result in results) == 1
assert len(event_types) == 1
assert (status, event_types[0]) in {
("failed", "run_failed"),
("cancelled", "run_cancelled"),
}
def test_finalize_run_raises_when_record_is_missing() -> None:
def test_request_cancellation_raises_when_record_is_missing() -> None:
redis = FakeRedis()
store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType]
redis.eval_result = [-1, ""]
with pytest.raises(RunNotFoundError):
asyncio.run(
store.finalize_run(RunCancelledEvent(run_id="missing", data=RunCancelledEventData(reason="cancelled")))
)
asyncio.run(store.request_cancellation("missing", CancelRunRequest(reason="cancelled")))
def test_wait_for_cancellation_observes_terminal_record_before_starting() -> None:
redis = FakeRedis()
store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType]
async def scenario() -> bool:
async def scenario() -> object:
record = await store.create_run()
_ = await store.finalize_run(
RunCancelledEvent(run_id=record.run_id, data=RunCancelledEventData(reason="cancelled"))
)
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 True
assert asyncio.run(scenario()) is None
assert [command[0] for command in redis.commands] == ["xrevrange", "get"]
def test_wait_for_cancellation_covers_terminal_transition_during_initialization() -> None:
def test_wait_for_cancellation_covers_intent_transition_during_initialization() -> None:
class PausingRecordReadRedis(FakeRedis):
record_read_started: asyncio.Event
release_record_read: asyncio.Event
@ -432,62 +347,75 @@ def test_wait_for_cancellation_covers_terminal_transition_during_initialization(
redis = PausingRecordReadRedis()
observer_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType]
cancelling_store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType]
async def scenario() -> bool:
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)
_ = await cancelling_store.finalize_run(
RunCancelledEvent(run_id=record.run_id, data=RunCancelledEventData(reason="cancelled"))
_ = redis._append_stream_entry(
f"test:runs:{record.run_id}:cancel-intent",
{
"payload": RunCancellationIntent(
reason="cancelled",
requested_at=utc_now(),
).model_dump_json()
},
)
redis.release_record_read.set()
return await asyncio.wait_for(observer, timeout=1)
assert asyncio.run(scenario()) is True
assert asyncio.run(scenario()) is not None
def test_wait_for_cancellation_advances_past_non_terminal_events() -> None:
redis = FakeRedis()
store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType]
async def scenario() -> bool:
async def scenario() -> object:
record = await store.create_run()
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)
_ = await store.finalize_run(
RunCancelledEvent(run_id=record.run_id, data=RunCancelledEventData(reason="cancelled"))
_ = redis._append_stream_entry(
f"test:runs:{record.run_id}:cancel-intent",
{
"payload": RunCancellationIntent(
reason="cancelled",
requested_at=utc_now(),
).model_dump_json()
},
)
return await asyncio.wait_for(observer, timeout=1)
assert asyncio.run(scenario()) is True
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_false_when_success_wins() -> None:
def test_wait_for_cancellation_returns_none_when_success_wins() -> None:
redis = FakeRedis()
store = RedisRunStore(redis, prefix="test") # pyright: ignore[reportArgumentType]
async def scenario() -> bool:
async def scenario() -> object:
record = await store.create_run()
observer = asyncio.create_task(store.wait_for_cancellation(record.run_id))
await asyncio.sleep(0)
_ = await store.finalize_run(
RunSucceededEvent(
run_id=record.run_id,
data=RunSucceededEventData(
output="done",
session_snapshot=CompositorSessionSnapshot(layers=[]),
),
)
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 False
assert asyncio.run(scenario()) is None
def test_append_event_serializes_typed_event_without_id_and_expires_run_keys() -> None:
@ -530,19 +458,20 @@ def test_get_events_round_trips_run_succeeded_output_and_session_snapshot() -> N
async def scenario() -> tuple[str, RunSucceededEvent]:
record = await store.create_run()
result = await store.finalize_run(
RunSucceededEvent(
id="local-only",
run_id=record.run_id,
data=RunSucceededEventData(output=output, session_snapshot=session_snapshot),
)
event = RunSucceededEvent(
id="local-only",
run_id=record.run_id,
data=RunSucceededEventData(output=output, session_snapshot=session_snapshot),
)
event_id = redis._append_stream_entry(
f"test:runs:{record.run_id}:events",
{"payload": RUN_EVENT_ADAPTER.dump_json(event, exclude={"id"}).decode()},
)
assert result.event_id is not None
page = await store.get_events(record.run_id, after="0-0", limit=10)
decoded = page.events[0]
assert isinstance(decoded, RunSucceededEvent)
assert page.next_cursor == result.event_id
return result.event_id, decoded
assert page.next_cursor == event_id
return event_id, decoded
event_id, decoded = asyncio.run(scenario())
@ -559,7 +488,11 @@ def test_iter_events_ends_after_replaying_terminal_event(terminal_type: str) ->
async def scenario() -> list[str]:
record = await store.create_run()
_ = await store.append_event(RunStartedEvent(run_id=record.run_id))
_ = await store.finalize_run(_terminal_event(terminal_type, record.run_id))
terminal = _terminal_event(terminal_type, record.run_id)
_ = redis._append_stream_entry(
f"test:runs:{record.run_id}:events",
{"payload": RUN_EVENT_ADAPTER.dump_json(terminal, exclude={"id"}).decode()},
)
redis.commands.clear()
async def collect_events() -> list[str]:
@ -586,7 +519,11 @@ def test_iter_events_ends_after_live_terminal_event(terminal_type: str) -> None:
assert not next_event.done()
assert "xread" in [command[0] for command in redis.commands]
_ = await store.finalize_run(_terminal_event(terminal_type, record.run_id))
terminal = _terminal_event(terminal_type, record.run_id)
_ = redis._append_stream_entry(
f"test:runs:{record.run_id}:events",
{"payload": RUN_EVENT_ADAPTER.dump_json(terminal, exclude={"id"}).decode()},
)
event = await asyncio.wait_for(next_event, timeout=1)
with pytest.raises(StopAsyncIteration):
_ = await anext(events)