From a679ee56444cb15b7b0eccbe6ffbbb4b918dedd6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=9B=90=E7=B2=92=20Yanli?= Date: Mon, 17 Aug 2026 10:11:22 +0000 Subject: [PATCH] fix(dify-agent): preserve snapshots for failed and cancelled runs (#40876) --- api/clients/agent_backend/client.py | 32 ++ api/clients/agent_backend/event_adapter.py | 4 + api/clients/agent_backend/fake_client.py | 30 +- api/core/app/apps/agent_app/app_runner.py | 72 +++- .../workflow/nodes/agent_v2/agent_node.py | 68 +++- .../clients/agent_backend/test_client.py | 25 ++ .../app/apps/agent_app/test_app_runner.py | 130 ++++++- .../nodes/agent_v2/test_agent_node.py | 153 +++++++- .../concepts/run-lifecycle/index.md | 5 +- dify-agent/docs/dify-agent/guide/index.md | 51 ++- dify-agent/src/dify_agent/client/_client.py | 43 ++- dify-agent/src/dify_agent/protocol/schemas.py | 9 +- .../src/dify_agent/runtime/cancellation.py | 19 + .../src/dify_agent/runtime/event_sink.py | 41 +-- .../src/dify_agent/runtime/run_scheduler.py | 154 +++++--- dify-agent/src/dify_agent/runtime/runner.py | 15 +- .../src/dify_agent/storage/redis_keys.py | 7 +- .../src/dify_agent/storage/redis_run_store.py | 198 +++++++++- .../storage/test_terminal_finalization.py | 189 ++++++++-- .../local/dify_agent/client/test_client.py | 104 ++++++ .../protocol/test_protocol_schemas.py | 28 +- .../dify_agent/runtime/test_run_scheduler.py | 234 +++++++++--- .../local/dify_agent/runtime/test_runner.py | 63 +++- .../storage/test_redis_run_store.py | 341 +++++++----------- 24 files changed, 1575 insertions(+), 440 deletions(-) create mode 100644 dify-agent/src/dify_agent/runtime/cancellation.py diff --git a/api/clients/agent_backend/client.py b/api/clients/agent_backend/client.py index 6738b33bcf4..297bb523c0e 100644 --- a/api/clients/agent_backend/client.py +++ b/api/clients/agent_backend/client.py @@ -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, diff --git a/api/clients/agent_backend/event_adapter.py b/api/clients/agent_backend/event_adapter.py index e72d2fbfb51..14f52742d2c 100644 --- a/api/clients/agent_backend/event_adapter.py +++ b/api/clients/agent_backend/event_adapter.py @@ -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__}") diff --git a/api/clients/agent_backend/fake_client.py b/api/clients/agent_backend/fake_client.py index 2d0881c03ff..c016f2ce6bf 100644 --- a/api/clients/agent_backend/fake_client.py +++ b/api/clients/agent_backend/fake_client.py @@ -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: diff --git a/api/core/app/apps/agent_app/app_runner.py b/api/core/app/apps/agent_app/app_runner.py index da1008ddfe8..191c4223f70 100644 --- a/api/core/app/apps/agent_app/app_runner.py +++ b/api/core/app/apps/agent_app/app_runner.py @@ -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 diff --git a/api/core/workflow/nodes/agent_v2/agent_node.py b/api/core/workflow/nodes/agent_v2/agent_node.py index 9911ef57ad7..73288b4bcc6 100644 --- a/api/core/workflow/nodes/agent_v2/agent_node.py +++ b/api/core/workflow/nodes/agent_v2/agent_node.py @@ -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: diff --git a/api/tests/unit_tests/clients/agent_backend/test_client.py b/api/tests/unit_tests/clients/agent_backend/test_client.py index 5fb54117306..34b66b83211 100644 --- a/api/tests/unit_tests/clients/agent_backend/test_client.py +++ b/api/tests/unit_tests/clients/agent_backend/test_client.py @@ -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) diff --git a/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py b/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py index 495f8073f92..5de8b8bd9e7 100644 --- a/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py +++ b/api/tests/unit_tests/core/app/apps/agent_app/test_app_runner.py @@ -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 == [] diff --git a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py index 71f03fb92a7..a8c254ae604 100644 --- a/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py +++ b/api/tests/unit_tests/core/workflow/nodes/agent_v2/test_agent_node.py @@ -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 " ) 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(): diff --git a/dify-agent/docs/dify-agent/concepts/run-lifecycle/index.md b/dify-agent/docs/dify-agent/concepts/run-lifecycle/index.md index bbc9bd2fc31..e4a3593d986 100644 --- a/dify-agent/docs/dify-agent/concepts/run-lifecycle/index.md +++ b/dify-agent/docs/dify-agent/concepts/run-lifecycle/index.md @@ -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 diff --git a/dify-agent/docs/dify-agent/guide/index.md b/dify-agent/docs/dify-agent/guide/index.md index fd436822d77..f24fd846e76 100644 --- a/dify-agent/docs/dify-agent/guide/index.md +++ b/dify-agent/docs/dify-agent/guide/index.md @@ -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 diff --git a/dify-agent/src/dify_agent/client/_client.py b/dify-agent/src/dify_agent/client/_client.py index 0cd14a68711..80af37d12fe 100644 --- a/dify-agent/src/dify_agent/client/_client.py +++ b/dify-agent/src/dify_agent/client/_client.py @@ -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: diff --git a/dify-agent/src/dify_agent/protocol/schemas.py b/dify-agent/src/dify_agent/protocol/schemas.py index ca23f8bfd59..a691650bfba 100644 --- a/dify-agent/src/dify_agent/protocol/schemas.py +++ b/dify-agent/src/dify_agent/protocol/schemas.py @@ -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") diff --git a/dify-agent/src/dify_agent/runtime/cancellation.py b/dify-agent/src/dify_agent/runtime/cancellation.py new file mode 100644 index 00000000000..60ad2750734 --- /dev/null +++ b/dify-agent/src/dify_agent/runtime/cancellation.py @@ -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"] diff --git a/dify-agent/src/dify_agent/runtime/event_sink.py b/dify-agent/src/dify_agent/runtime/event_sink.py index 121d58cecde..5e5995045f0 100644 --- a/dify-agent/src/dify_agent/runtime/event_sink.py +++ b/dify-agent/src/dify_agent/runtime/event_sink.py @@ -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", diff --git a/dify-agent/src/dify_agent/runtime/run_scheduler.py b/dify-agent/src/dify_agent/runtime/run_scheduler.py index b22d48e5626..b197d7dcb3e 100644 --- a/dify-agent/src/dify_agent/runtime/run_scheduler.py +++ b/dify-agent/src/dify_agent/runtime/run_scheduler.py @@ -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"] diff --git a/dify-agent/src/dify_agent/runtime/runner.py b/dify-agent/src/dify_agent/runtime/runner.py index ef4228ac982..c6c8d829312 100644 --- a/dify-agent/src/dify_agent/runtime/runner.py +++ b/dify-agent/src/dify_agent/runtime/runner.py @@ -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.") diff --git a/dify-agent/src/dify_agent/storage/redis_keys.py b/dify-agent/src/dify_agent/storage/redis_keys.py index c93e69ba825..bc8d10ef739 100644 --- a/dify-agent/src/dify_agent/storage/redis_keys.py +++ b/dify-agent/src/dify_agent/storage/redis_keys.py @@ -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"] diff --git a/dify-agent/src/dify_agent/storage/redis_run_store.py b/dify-agent/src/dify_agent/storage/redis_run_store.py index f488749c0ff..70a6adc85d6 100644 --- a/dify-agent/src/dify_agent/storage/redis_run_store.py +++ b/dify-agent/src/dify_agent/storage/redis_run_store.py @@ -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) diff --git a/dify-agent/tests/integration/dify_agent/storage/test_terminal_finalization.py b/dify-agent/tests/integration/dify_agent/storage/test_terminal_finalization.py index 396e7736f68..af3a77c687d 100644 --- a/dify-agent/tests/integration/dify_agent/storage/test_terminal_finalization.py +++ b/dify-agent/tests/integration/dify_agent/storage/test_terminal_finalization.py @@ -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: diff --git a/dify-agent/tests/local/dify_agent/client/test_client.py b/dify-agent/tests/local/dify_agent/client/test_client.py index 7589eaa53e2..54c905625ac 100644 --- a/dify-agent/tests/local/dify_agent/client/test_client.py +++ b/dify-agent/tests/local/dify_agent/client/test_client.py @@ -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": diff --git a/dify-agent/tests/local/dify_agent/protocol/test_protocol_schemas.py b/dify-agent/tests/local/dify_agent/protocol/test_protocol_schemas.py index 8eac3a6023e..fbbf9077d73 100644 --- a/dify-agent/tests/local/dify_agent/protocol/test_protocol_schemas.py +++ b/dify-agent/tests/local/dify_agent/protocol/test_protocol_schemas.py @@ -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( { diff --git a/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py b/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py index f7fc7a1e94e..44643d6a312 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_run_scheduler.py @@ -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()) diff --git a/dify-agent/tests/local/dify_agent/runtime/test_runner.py b/dify-agent/tests/local/dify_agent/runtime/test_runner.py index 371a94e129c..c3199922922 100644 --- a/dify-agent/tests/local/dify_agent/runtime/test_runner.py +++ b/dify-agent/tests/local/dify_agent/runtime/test_runner.py @@ -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: diff --git a/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py b/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py index e33e1ec8188..13930624dd1 100644 --- a/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py +++ b/dify-agent/tests/local/dify_agent/storage/test_redis_run_store.py @@ -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)