fix(api): complete chatflow streams after direct replies (#45552)

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)

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(

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

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,