From 1ac9aa0d3c93733667d2095f3de2fa4cae0a3e1d Mon Sep 17 00:00:00 2001 From: "Byron.wang" Date: Wed, 4 Nov 2026 09:34:19 +0000 Subject: [PATCH] fix(api): complete chatflow streams after direct replies (#45552) --- api/core/app/apps/streaming_utils.py | 13 + .../app_generate/workflow_execute_task.py | 28 ++- .../core/app/apps/test_streaming_utils.py | 234 +++++++++++++++++- .../test_workflow_execute_task.py | 130 ++++++++++ 4 files changed, 397 insertions(+), 8 deletions(-) diff --git a/api/core/app/apps/streaming_utils.py b/api/core/app/apps/streaming_utils.py index a644f5ded27..db1980d0b34 100644 --- a/api/core/app/apps/streaming_utils.py +++ b/api/core/app/apps/streaming_utils.py @@ -9,6 +9,9 @@ from core.app.entities.task_entities import StreamEvent from libs.broadcast_channel.channel import Subscription from libs.broadcast_channel.exc import SubscriptionClosedError +# Internal transport signal, consumed before decoding public stream events. +STREAM_COMPLETE_SIGNAL = b"__dify_app_stream_complete__" + def stream_topic_events( *, @@ -18,6 +21,13 @@ def stream_topic_events( on_subscribe: Callable[[], None] | None = None, terminal_events: Iterable[str | StreamEvent] | None = None, ) -> Generator[Mapping[str, Any] | str]: + """Yield JSON topic events until a terminal event, transport completion, closure, or idle timeout. + + The subscription must deliver JSON payloads or ``STREAM_COMPLETE_SIGNAL`` in publish order. + The completion signal closes the subscription without becoming a public event, regardless + of the selected terminal events. Invalid JSON and subscription errors other than closure + propagate to the caller; an entered subscription is closed on exit. + """ # send a PING event immediately to prevent the connection staying in pending state for a long time. # # This simplify the debugging process as the DevTools in Chrome does not @@ -46,6 +56,9 @@ def stream_topic_events( last_ping_time = current_time continue + if msg == STREAM_COMPLETE_SIGNAL: + return + last_msg_time = time.time() last_ping_time = last_msg_time event = json.loads(msg) diff --git a/api/tasks/app_generate/workflow_execute_task.py b/api/tasks/app_generate/workflow_execute_task.py index 12eec51ec54..6b47a01beab 100644 --- a/api/tasks/app_generate/workflow_execute_task.py +++ b/api/tasks/app_generate/workflow_execute_task.py @@ -15,6 +15,7 @@ from sqlalchemy.orm import Session, sessionmaker from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator from core.app.apps.execution_coordinator import clear_app_task_cancellation_signals from core.app.apps.message_based_app_generator import MessageBasedAppGenerator +from core.app.apps.streaming_utils import STREAM_COMPLETE_SIGNAL from core.app.apps.workflow.app_generator import WorkflowAppGenerator from core.app.entities.app_invoke_entities import ( AdvancedChatAppGenerateEntity, @@ -365,12 +366,13 @@ def _publish_streaming_response( inputs: Mapping[str, Any], started_reason: WorkflowStartReason, ) -> None: - """Publish workflow stream events and close broken streams with a failed terminal event. + """Publish app stream events in order and terminate successful or broken streams. - `_AppRunner.run()` only handles failures before the generator is returned. - Once we start iterating the runtime stream, this helper becomes the last - place that can guarantee SSE consumers eventually see a terminal workflow - lifecycle event. + Accept runtime events for the given app mode and run. A normally exhausted Advanced Chat + stream with a published message end, no workflow start, and no errors completes through an + internal transport signal after all output is published. It must not invent a workflow run. + Other streams require a workflow terminal event; missing terminals produce a failed fallback. + Iteration and publish errors also produce a fallback when possible, then propagate. """ normalized_workflow_run_id = str(workflow_run_id) @@ -423,6 +425,8 @@ def _publish_streaming_response( topic = MessageBasedAppGenerator.get_response_topic(app_mode, normalized_workflow_run_id) started_published = False terminal_published = False + message_end_published = False + stream_failed = False last_task_id = normalized_workflow_run_id stream_error_message: str | None = None @@ -439,6 +443,7 @@ def _publish_streaming_response( else: payload = json.dumps(event, ensure_ascii=False, default=str) except (TypeError, ValueError): + stream_failed = True logger.exception("error while encoding event") continue @@ -448,8 +453,21 @@ def _publish_streaming_response( started_published = True elif event_name in terminal_events: terminal_published = True + elif event_name == "message_end": + message_end_published = True elif event_name == "error": + stream_failed = True stream_error_message = _get_error_message(event) or stream_error_message + + if ( + app_mode == AppMode.ADVANCED_CHAT + and message_end_published + and not started_published + and not terminal_published + and not stream_failed + ): + topic.publish(STREAM_COMPLETE_SIGNAL) + return except Exception as exc: if not terminal_published: logger.exception( diff --git a/api/tests/unit_tests/core/app/apps/test_streaming_utils.py b/api/tests/unit_tests/core/app/apps/test_streaming_utils.py index 66afebb3e60..95fec372f7c 100644 --- a/api/tests/unit_tests/core/app/apps/test_streaming_utils.py +++ b/api/tests/unit_tests/core/app/apps/test_streaming_utils.py @@ -2,14 +2,37 @@ from __future__ import annotations import json import queue +from collections.abc import Generator import pytest +from sqlalchemy.orm import Session +from core.app.app_config.entities import WorkflowUIBasedAppConfig +from core.app.apps.advanced_chat.generate_response_converter import AdvancedChatAppGenerateResponseConverter +from core.app.apps.advanced_chat.generate_task_pipeline import ( + AdvancedChatAppGenerateTaskPipeline, + ConversationSnapshot, + MessageSnapshot, + WorkflowSnapshot, +) +from core.app.apps.base_app_generator import BaseAppGenerator +from core.app.apps.base_app_queue_manager import PublishFrom from core.app.apps.message_based_app_generator import MessageBasedAppGenerator -from core.app.apps.streaming_utils import _normalize_terminal_events, stream_topic_events -from core.app.entities.task_entities import StreamEvent +from core.app.apps.message_based_app_queue_manager import MessageBasedAppQueueManager +from core.app.apps.streaming_utils import STREAM_COMPLETE_SIGNAL, _normalize_terminal_events, stream_topic_events +from core.app.entities.app_invoke_entities import AdvancedChatAppGenerateEntity, InvokeFrom +from core.app.entities.queue_entities import QueueStopEvent, QueueTextChunkEvent +from core.app.entities.task_entities import ( + ChatbotAppStreamResponse, + MessageEndStreamResponse, + MessageStreamResponse, + StreamEvent, +) +from graphon.entities import WorkflowStartReason from libs.broadcast_channel.channel import SupportsPreparedSubscription -from models.model import AppMode +from models.enums import ConversationFromSource, MessageStatus +from models.model import AppMode, EndUser, Message +from tasks.app_generate.workflow_execute_task import _publish_streaming_response class FakeSubscription: @@ -28,6 +51,10 @@ class FakeSubscription: def close(self) -> None: self._closed = True + @property + def closed(self) -> bool: + return self._closed + def receive(self, timeout: float | None = 0.1) -> bytes | None: if self._closed: return None @@ -220,3 +247,204 @@ def test_stream_topic_events_can_continue_past_pause(): assert next(generator)["event"] == StreamEvent.WORKFLOW_FINISHED.value with pytest.raises(StopIteration): next(generator) + + +def test_stream_topic_events_hides_completion_signal_and_delivers_prior_public_events(): + topic = FakeTopic() + message_end = {"event": StreamEvent.MESSAGE_END.value, "message_id": "message-id"} + tts_end = {"event": StreamEvent.TTS_MESSAGE_END.value, "message_id": "message-id"} + topic.publish(json.dumps(message_end).encode()) + topic.publish(json.dumps(tts_end).encode()) + topic.publish(STREAM_COMPLETE_SIGNAL) + topic.publish(json.dumps({"event": StreamEvent.MESSAGE.value, "answer": "too late"}).encode()) + subscription = topic.subscribe() + + events = list(stream_topic_events(subscription=subscription, idle_timeout=1.0)) + + assert events == [StreamEvent.PING.value, message_end, tts_end] + assert subscription.closed is True + + +def test_stream_completion_signal_ends_stream_when_public_terminal_events_are_disabled(): + topic = FakeTopic() + workflow_finished = {"event": StreamEvent.WORKFLOW_FINISHED.value} + topic.publish(json.dumps(workflow_finished).encode()) + topic.publish(STREAM_COMPLETE_SIGNAL) + topic.publish(json.dumps({"event": StreamEvent.MESSAGE.value, "answer": "too late"}).encode()) + subscription = topic.subscribe() + + events = list( + stream_topic_events( + subscription=subscription, + idle_timeout=1.0, + terminal_events=[], + ) + ) + + assert events == [StreamEvent.PING.value, workflow_finished] + assert subscription.closed is True + + +def test_annotation_shaped_response_stream_completes_without_workflow_events(monkeypatch: pytest.MonkeyPatch): + topic = FakeTopic() + + def fake_get_response_topic(cls, app_mode, workflow_run_id): + return topic + + monkeypatch.setattr(MessageBasedAppGenerator, "get_response_topic", classmethod(fake_get_response_topic)) + subscription = topic.subscribe() + + def runtime_responses() -> Generator[ChatbotAppStreamResponse]: + yield ChatbotAppStreamResponse( + conversation_id="conversation-id", + message_id="message-id", + created_at=1_700_000_000, + stream_response=MessageStreamResponse( + task_id="task-id", + id="message-id", + answer="annotation answer", + ), + ) + yield ChatbotAppStreamResponse( + conversation_id="conversation-id", + message_id="message-id", + created_at=1_700_000_000, + stream_response=MessageEndStreamResponse( + task_id="task-id", + id="message-id", + metadata={}, + ), + ) + + converted = AdvancedChatAppGenerateResponseConverter.convert(runtime_responses(), InvokeFrom.WEB_APP) + assert isinstance(converted, Generator) + _publish_streaming_response( + converted, + "workflow-run-id", + app_mode=AppMode.ADVANCED_CHAT, + workflow_id="workflow-id", + inputs={}, + started_reason=WorkflowStartReason.INITIAL, + ) + + # Two public events plus transport completion must be queued before consumption. + # Otherwise an idle timeout could masquerade as a successful EOF in this test. + assert topic._queue.qsize() == 3 + response_stream = stream_topic_events(subscription=subscription, idle_timeout=1.0) + assert next(response_stream) == StreamEvent.PING.value + message = next(response_stream) + assert isinstance(message, dict) + assert message["event"] == StreamEvent.MESSAGE.value + assert message["answer"] == "annotation answer" + message_end = next(response_stream) + assert isinstance(message_end, dict) + assert message_end["event"] == StreamEvent.MESSAGE_END.value + assert message_end["metadata"] == {} + with pytest.raises(StopIteration): + next(response_stream) + + assert subscription.closed is True + assert topic._queue.empty() + + +def test_input_moderation_direct_reply_completes_without_workflow_events( + monkeypatch: pytest.MonkeyPatch, + sqlite_session: Session, +): + """A moderation text/stop queue saves the reply and reaches EOF without a workflow or annotation.""" + moderation_reply = "This request was blocked by input moderation." + user = EndUser(tenant_id="tenant-id", type="session", session_id="session-id") + message = Message( + id="message-id", + app_id="app-id", + conversation_id="conversation-id", + query="blocked query", + message={}, + answer="", + from_source=ConversationFromSource.API, + from_end_user_id=user.id, + invoke_from=InvokeFrom.WEB_APP, + app_mode=AppMode.ADVANCED_CHAT, + ) + sqlite_session.add(message) + sqlite_session.commit() + app_config = WorkflowUIBasedAppConfig( + app_id=message.app_id, + tenant_id=user.tenant_id, + workflow_id="workflow-id", + app_mode=AppMode.ADVANCED_CHAT, + ) + generate_entity = AdvancedChatAppGenerateEntity( + task_id="task-id", + app_config=app_config, + inputs={}, + files=[], + user_id=user.id, + stream=True, + invoke_from=InvokeFrom.WEB_APP, + query=message.query, + workflow_run_id="workflow-run-id", + extras={"auto_generate_conversation_name": False}, + ) + queue_manager = MessageBasedAppQueueManager( + task_id=generate_entity.task_id, + user_id=user.id, + invoke_from=generate_entity.invoke_from, + conversation_id=message.conversation_id, + app_mode=AppMode.ADVANCED_CHAT, + message_id=message.id, + ) + pipeline = AdvancedChatAppGenerateTaskPipeline( + application_generate_entity=generate_entity, + workflow=WorkflowSnapshot(id=app_config.workflow_id, tenant_id=user.tenant_id, features_dict={}), + queue_manager=queue_manager, + conversation=ConversationSnapshot(id=message.conversation_id, mode=AppMode.ADVANCED_CHAT), + message=MessageSnapshot.from_message(message), + user=user, + stream=True, + dialogue_count=1, + draft_var_saver_factory=BaseAppGenerator._get_draft_var_saver_factory( + generate_entity.invoke_from, user, tenant_id=user.tenant_id + ), + ) + # These are the events emitted by the runner when input moderation returns a direct reply. + queue_manager.publish(QueueTextChunkEvent(text=moderation_reply), PublishFrom.APPLICATION_MANAGER) + queue_manager.publish( + QueueStopEvent(stopped_by=QueueStopEvent.StopBy.INPUT_MODERATION), PublishFrom.APPLICATION_MANAGER + ) + runtime_responses = pipeline.process() + assert isinstance(runtime_responses, Generator) + converted = AdvancedChatAppGenerateResponseConverter.convert(runtime_responses, InvokeFrom.WEB_APP) + assert isinstance(converted, Generator) + topic = FakeTopic() + monkeypatch.setattr(MessageBasedAppGenerator, "get_response_topic", lambda *_args, **_kwargs: topic) + subscription = topic.subscribe() + + _publish_streaming_response( + converted, + generate_entity.workflow_run_id, + app_mode=AppMode.ADVANCED_CHAT, + workflow_id=app_config.workflow_id, + inputs={}, + started_reason=WorkflowStartReason.INITIAL, + ) + + # A queued completion signal makes EOF independent of the idle timeout. + assert topic._queue.qsize() == 3 + events = list(stream_topic_events(subscription=subscription, idle_timeout=1.0)) + assert events[0] == StreamEvent.PING.value + message_event, message_end = events[1:] + assert isinstance(message_event, dict) + assert message_event["event"] == StreamEvent.MESSAGE.value + assert message_event["answer"] == moderation_reply + assert isinstance(message_end, dict) + assert message_end["event"] == StreamEvent.MESSAGE_END.value + assert "annotation_reply" not in message_end["metadata"] + assert subscription.closed is True + assert topic._queue.empty() + assert queue_manager.graph_runtime_state is None + + sqlite_session.refresh(message) + assert message.answer == moderation_reply + assert message.status == MessageStatus.NORMAL + assert message.workflow_run_id is None diff --git a/api/tests/unit_tests/tasks/app_generate/test_workflow_execute_task.py b/api/tests/unit_tests/tasks/app_generate/test_workflow_execute_task.py index 8532cb84f5c..c0b9f37278a 100644 --- a/api/tests/unit_tests/tasks/app_generate/test_workflow_execute_task.py +++ b/api/tests/unit_tests/tasks/app_generate/test_workflow_execute_task.py @@ -2,6 +2,7 @@ from __future__ import annotations import json import logging +import queue import uuid from collections.abc import Callable, Generator, Mapping from contextlib import nullcontext @@ -59,6 +60,25 @@ class _StreamEventModel(BaseModel): message: object | None = None +class _ResponsePublisher(queue.Queue[bytes]): + """Preserve actual published payloads in a FIFO queue without Redis I/O.""" + + def publish(self, payload: bytes) -> None: + """Enqueue the payload unchanged for inspection by the consumer.""" + self.put(payload) + + +@pytest.fixture +def response_publisher(monkeypatch: pytest.MonkeyPatch) -> _ResponsePublisher: + publisher = _ResponsePublisher() + monkeypatch.setattr( + workflow_execute_task_module.MessageBasedAppGenerator, + "get_response_topic", + lambda *_args, **_kwargs: publisher, + ) + return publisher + + def _build_advanced_chat_generate_entity(conversation_id: str | None) -> AdvancedChatAppGenerateEntity: return AdvancedChatAppGenerateEntity( task_id="task-id", @@ -679,6 +699,116 @@ def test_publish_streaming_response_does_not_publish_synthetic_failure_after_ter assert [payload["event"] for payload in payloads] == ["workflow_started", "workflow_finished"] +@pytest.mark.parametrize( + ("events", "app_mode"), + [ + pytest.param([], AppMode.ADVANCED_CHAT, id="empty-stream"), + pytest.param([{"event": "message", "answer": "answer"}], AppMode.ADVANCED_CHAT, id="missing-message-end"), + pytest.param( + [{"event": "workflow_started"}, {"event": "message_end"}], + AppMode.ADVANCED_CHAT, + id="incomplete-workflow", + ), + pytest.param( + [{"event": "message_end"}, {"event": "error", "message": "provider failed"}], + AppMode.ADVANCED_CHAT, + id="error-after-message-end", + ), + pytest.param( + [{"event": "message_end"}, {"event": "error", "message": ""}], + AppMode.ADVANCED_CHAT, + id="empty-error-message", + ), + pytest.param( + [{"event": "message_end"}, {"event": "error", "message": 123}], + AppMode.ADVANCED_CHAT, + id="invalid-error-message", + ), + pytest.param( + [{"event": "message_end"}, {"event": "error"}], + AppMode.ADVANCED_CHAT, + id="missing-error-message", + ), + pytest.param([{"event": "message_end"}], AppMode.WORKFLOW, id="workflow-app"), + ], +) +def test_publish_streaming_response_preserves_failure_for_incomplete_or_failed_chat( + response_publisher: _ResponsePublisher, + events: list[Mapping[str, object]], + app_mode: AppMode, +): + def response_stream() -> Generator[Mapping[str, object]]: + yield from events + + _publish_streaming_response( + response_stream(), + "workflow-run-id", + app_mode=app_mode, + workflow_id="workflow-id", + inputs={}, + started_reason=WorkflowStartReason.INITIAL, + ) + + payloads = [json.loads(response_publisher.get_nowait()) for _ in range(response_publisher.qsize())] + assert payloads[: len(events)] == events + assert [payload["event"] for payload in payloads].count("workflow_started") == 1 + assert payloads[-1]["event"] == "workflow_finished" + assert payloads[-1]["data"]["status"] == WorkflowExecutionStatus.FAILED + expected_error = ( + "provider failed" + if events and events[-1].get("message") == "provider failed" + else "Workflow stream ended without a terminal event" + ) + assert payloads[-1]["data"]["error"] == expected_error + + +def test_publish_streaming_response_does_not_complete_chat_after_event_encoding_failure( + response_publisher: _ResponsePublisher, +): + unencodable_event: dict[str, object] = {"event": "workflow_started"} + unencodable_event["data"] = unencodable_event + + def response_stream() -> Generator[Mapping[str, object]]: + yield unencodable_event + yield {"event": "message_end"} + + _publish_streaming_response( + response_stream(), + "workflow-run-id", + app_mode=AppMode.ADVANCED_CHAT, + workflow_id="workflow-id", + inputs={}, + started_reason=WorkflowStartReason.INITIAL, + ) + + payloads = [json.loads(response_publisher.get_nowait()) for _ in range(response_publisher.qsize())] + assert [payload["event"] for payload in payloads] == ["message_end", "workflow_started", "workflow_finished"] + assert payloads[-1]["data"]["status"] == WorkflowExecutionStatus.FAILED + + +def test_publish_streaming_response_preserves_failure_when_iteration_raises_after_message_end( + response_publisher: _ResponsePublisher, +): + def response_stream() -> Generator[Mapping[str, object]]: + yield {"event": "message_end"} + raise RuntimeError("stream exploded") + + with pytest.raises(RuntimeError, match="stream exploded"): + _publish_streaming_response( + response_stream(), + "workflow-run-id", + app_mode=AppMode.ADVANCED_CHAT, + workflow_id="workflow-id", + inputs={}, + started_reason=WorkflowStartReason.INITIAL, + ) + + payloads = [json.loads(response_publisher.get_nowait()) for _ in range(response_publisher.qsize())] + assert [payload["event"] for payload in payloads] == ["message_end", "workflow_started", "workflow_finished"] + assert payloads[-1]["data"]["status"] == WorkflowExecutionStatus.FAILED + assert payloads[-1]["data"]["error"] == "stream exploded" + + def test_app_runner_streaming_failure_publishes_started_then_failed_workflow_finished( mock_topic: MagicMock, monkeypatch: pytest.MonkeyPatch,