fix(api): prevent dropped workflow_started events in Redis Streams (#40964)

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: QuantumGhost <obelisk.reg+git@gmail.com>

api/core/app/apps/message_based_app_generator.py

@@ -31,7 +31,7 @@ from core.prompt.utils.prompt_template_parser import PromptTemplateParser
 from core.workflow.file_reference import resolve_file_record_id
 from extensions.ext_database import db
 from extensions.ext_redis import get_pubsub_broadcast_channel
-from libs.broadcast_channel.channel import Topic
+from libs.broadcast_channel.channel import SupportsPreparedSubscription, Topic
 from libs.datetime_utils import naive_utc_now
 from models import Account
 from models.enums import ConversationFromSource, CreatorUserRole, MessageFileBelongsTo
@@ -323,8 +323,14 @@ class MessageBasedAppGenerator(BaseAppGenerator):
         on_subscribe: Callable[[], None] | None = None,
     ) -> Generator[Mapping | str, None, None]:
         topic = cls.get_response_topic(app_mode, workflow_run_id)
+        subscriber = topic.as_subscriber()
+        subscription = (
+            subscriber.prepare_subscription()
+            if isinstance(subscriber, SupportsPreparedSubscription)
+            else subscriber.subscribe()
+        )
         return stream_topic_events(
-            topic=topic,
+            subscription=subscription,
             idle_timeout=idle_timeout,
             on_subscribe=on_subscribe,
         )

api/core/app/apps/message_generator.py

@@ -3,7 +3,7 @@ from collections.abc import Callable, Generator, Iterable, Mapping
 from core.app.apps.streaming_utils import stream_topic_events
 from core.app.entities.task_entities import StreamEvent
 from extensions.ext_redis import get_pubsub_broadcast_channel
-from libs.broadcast_channel.channel import Topic
+from libs.broadcast_channel.channel import SupportsPreparedSubscription, Topic
 from models.model import AppMode
 
 
@@ -30,8 +30,14 @@ class MessageGenerator:
         terminal_events: Iterable[str | StreamEvent] | None = None,
     ) -> Generator[Mapping | str, None, None]:
         topic = cls.get_response_topic(app_mode, workflow_run_id)
+        subscriber = topic.as_subscriber()
+        subscription = (
+            subscriber.prepare_subscription()
+            if isinstance(subscriber, SupportsPreparedSubscription)
+            else subscriber.subscribe()
+        )
         return stream_topic_events(
-            topic=topic,
+            subscription=subscription,
             idle_timeout=idle_timeout,
             ping_interval=ping_interval,
             on_subscribe=on_subscribe,

api/core/app/apps/streaming_utils.py

@@ -6,13 +6,13 @@ from collections.abc import Callable, Generator, Iterable, Mapping
 from typing import Any
 
 from core.app.entities.task_entities import StreamEvent
-from libs.broadcast_channel.channel import Topic
+from libs.broadcast_channel.channel import Subscription
 from libs.broadcast_channel.exc import SubscriptionClosedError
 
 
 def stream_topic_events(
     *,
-    topic: Topic,
+    subscription: Subscription,
     idle_timeout: float,
     ping_interval: float | None = None,
     on_subscribe: Callable[[], None] | None = None,
@@ -27,7 +27,7 @@ def stream_topic_events(
     terminal_values = _normalize_terminal_events(terminal_events)
     last_msg_time = time.time()
     last_ping_time = last_msg_time
-    with topic.subscribe() as sub:
+    with subscription as sub:
         # on_subscribe fires only after the Redis subscription is active.
         # This is used to gate task start and reduce pub/sub race for the first event.
         if on_subscribe is not None:

api/libs/broadcast_channel/channel.py

@@ -4,6 +4,7 @@ Broadcast channel for Pub/Sub messaging.
 
 from __future__ import annotations
 
+import abc
 import types
 from abc import abstractmethod
 from collections.abc import Iterator
@@ -98,6 +99,15 @@ class Subscriber(Protocol):
         pass
 
 
+class SupportsPreparedSubscription(Subscriber, abc.ABC):
+    """Optional capability for fixing a subscription's delivery boundary before activation."""
+
+    @abstractmethod
+    def prepare_subscription(self) -> Subscription:
+        """Create an inactive subscription whose logical delivery boundary is already fixed."""
+        ...
+
+
 class Topic(Producer, Subscriber, Protocol):
     """A named channel for publishing and subscribing to messages.
 

api/libs/broadcast_channel/redis/streams_channel.py

@@ -7,7 +7,7 @@ from collections.abc import Iterator
 from typing import Self, override
 
 from extensions.redis_names import serialize_redis_name
-from libs.broadcast_channel.channel import Producer, Subscriber, Subscription
+from libs.broadcast_channel.channel import Producer, Subscriber, Subscription, SupportsPreparedSubscription
 from libs.broadcast_channel.exc import SubscriptionClosedError
 from libs.broadcast_channel.signals import SIG_CLOSE
 from redis import Redis, RedisCluster
@@ -68,18 +68,41 @@ class StreamsTopic:
                 logger.warning("Failed to set expire for stream key %s: %s", self._key, e, exc_info=True)
 
     def as_subscriber(self) -> Subscriber:
-        return self
+        return _StreamsSubscriber(self._client, self._key)
 
+    def subscribe(self) -> Subscription:
+        return self.as_subscriber().subscribe()
+
+
+class _StreamsSubscriber(SupportsPreparedSubscription):
+    def __init__(self, client: Redis | RedisCluster, key: str):
+        self._client = client
+        self._key = key
+
+    @override
     def subscribe(self) -> Subscription:
         return _StreamsSubscription(self._client, self._key)
 
+    @override
+    def prepare_subscription(self) -> Subscription:
+        entries = self._client.xrevrange(self._key, count=1)
+        start_id = entries[0][0] if entries else "0-0"
+        return _StreamsSubscription(self._client, self._key, start_id=start_id)
+
 
 class _StreamsSubscription(Subscription):
     _SENTINEL = object()
 
-    def __init__(self, client: Redis | RedisCluster, key: str):
+    def __init__(
+        self,
+        client: Redis | RedisCluster,
+        key: str,
+        *,
+        start_id: bytes | str = "$",
+    ):
         self._client = client
         self._key = key
+        self._start_id = start_id
 
         self._queue: queue.Queue[object] = queue.Queue()
 
@@ -104,10 +127,7 @@ class _StreamsSubscription(Subscription):
         # since this method runs in a dedicated thread, acquiring `_lock` inside this method won't cause
         # deadlock.
 
-        # Setting initial last id to `$` to signal redis that we only want new messages.
-        #
-        # ref: https://redis.io/docs/latest/commands/xread/#the-special--id
-        last_id = "$"
+        last_id = self._start_id
         try:
             while True:
                 with self._lock:
@@ -147,11 +167,9 @@ class _StreamsSubscription(Subscription):
 
     def _start_if_needed(self) -> None:
         """This method must be called with `_lock` held."""
-        if self._listener is not None:
-            return
-        # Ensure only one listener thread is created under concurrent calls
         if self._listener is not None or self._closed:
             return
+
         self._listener = threading.Thread(
             target=self._listen,
             name=f"redis-streams-sub-{self._key}",

api/services/app_generate_service.py

@@ -41,13 +41,14 @@ if TYPE_CHECKING:
 
 class AppGenerateService:
     @staticmethod
-    def _build_streaming_task_on_subscribe(start_task: Callable[[], None]) -> Callable[[], None]:
+    def _build_streaming_task_on_subscribe(
+        start_task: Callable[[], None],
+    ) -> Callable[[], None]:
         """
-        Build a subscription callback that coordinates when the background task starts.
+        Build a subscription callback that starts the background task on first subscribe.
 
-        - streams transport: start immediately (events are durable; late subscribers can replay).
-        - pubsub/sharded transport: start on first subscribe, with a short fallback timer so the task
-          still runs if the client never connects.
+        Pub/Sub transports also use a short fallback timer so the task still runs if the
+        client never connects. Streams rely on their prepared delivery boundary instead.
         """
         started = False
         lock = threading.Lock()
@@ -65,18 +66,13 @@ class AppGenerateService:
                 started = True
                 return True
 
-        channel_type = dify_config.PUBSUB_REDIS_CHANNEL_TYPE
-        if channel_type == "streams":
-            # With Redis Streams, we can safely start right away; consumers can read past events.
-            _try_start()
+        if dify_config.PUBSUB_REDIS_CHANNEL_TYPE == "streams":
 
-            # Keep return type Callable[[], None] consistent while allowing an extra (no-op) call.
             def _on_subscribe_streams() -> None:
                 _try_start()
 
             return _on_subscribe_streams
 
-        # Pub/Sub modes (at-most-once): subscribe-gated start with a tiny fallback.
         timer = threading.Timer(SSE_TASK_START_FALLBACK_MS / 1000.0, _try_start)
         timer.daemon = True
         timer.start()

api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_streams_channel.py

@@ -4,6 +4,7 @@ Integration tests for Redis Streams broadcast channel implementation using TestC
 This suite focuses on the semantics that differ from Redis Pub/Sub:
 - Every active subscription should receive each newly published message.
 - Each subscription should only observe messages published after its listener starts.
+- Prepared subscriptions should observe messages published after preparation but before listener startup.
 """
 
 import threading
@@ -16,9 +17,9 @@ import pytest
 import redis
 from testcontainers.redis import RedisContainer
 
-from libs.broadcast_channel.channel import BroadcastChannel, Subscription, Topic
+from libs.broadcast_channel.channel import BroadcastChannel, Subscription, SupportsPreparedSubscription, Topic
 from libs.broadcast_channel.exc import SubscriptionClosedError
-from libs.broadcast_channel.redis.streams_channel import StreamsBroadcastChannel
+from libs.broadcast_channel.redis.streams_channel import StreamsBroadcastChannel, _StreamsSubscription
 
 
 class TestRedisStreamsBroadcastChannelIntegration:
@@ -147,6 +148,39 @@ class TestRedisStreamsBroadcastChannelIntegration:
             first_subscription.close()
             second_subscription.close()
 
+    def test_prepare_subscription_fixes_boundary_before_listener_starts(
+        self,
+        broadcast_channel: BroadcastChannel,
+    ) -> None:
+        topic = broadcast_channel.topic(self._get_test_topic_name())
+        topic.publish(b"before-prepare")
+        subscriber = topic.as_subscriber()
+        assert isinstance(subscriber, SupportsPreparedSubscription)
+
+        subscription = subscriber.prepare_subscription()
+        topic.publish(b"after-prepare-before-start")
+
+        try:
+            assert self._receive_message(subscription) == b"after-prepare-before-start"
+            assert subscription.receive(timeout=0.1) is None
+        finally:
+            subscription.close()
+
+    def test_subscription_with_injected_start_id_reads_only_later_entries(
+        self,
+        redis_client: redis.Redis,
+    ) -> None:
+        stream_key = self._get_test_topic_name()
+        start_id = redis_client.xadd(stream_key, {b"data": b"at-boundary"})
+        redis_client.xadd(stream_key, {b"data": b"after-boundary"})
+        subscription = _StreamsSubscription(redis_client, stream_key, start_id=start_id)
+
+        try:
+            assert self._receive_message(subscription) == b"after-boundary"
+            assert subscription.receive(timeout=0.1) is None
+        finally:
+            subscription.close()
+
     def test_topic_isolation(self, broadcast_channel: BroadcastChannel) -> None:
         """Messages from different topics should remain isolated."""
         topic1 = broadcast_channel.topic(self._get_test_topic_name())

api/tests/unit_tests/core/app/apps/test_message_generator.py

@@ -2,9 +2,33 @@ from unittest.mock import Mock, patch
 
 from core.app.apps.message_generator import MessageGenerator
 from core.app.entities.task_entities import StreamEvent
+from libs.broadcast_channel.channel import SupportsPreparedSubscription
 from models.model import AppMode
 
 
+class _PreparedSubscriber(SupportsPreparedSubscription):
+    def __init__(self) -> None:
+        self.prepare_calls = 0
+
+    def prepare_subscription(self):
+        self.prepare_calls += 1
+        return "prepared-subscription"
+
+    def subscribe(self):
+        return "ordinary-subscription"
+
+
+class _TopicWithSeparateSubscriberView:
+    def __init__(self) -> None:
+        self.subscriber = _PreparedSubscriber()
+
+    def as_subscriber(self) -> _PreparedSubscriber:
+        return self.subscriber
+
+    def subscribe(self):
+        raise AssertionError("event retrieval must use the topic's subscriber view")
+
+
 class TestMessageGenerator:
     def test_get_response_topic(self):
         channel = Mock()
@@ -18,8 +42,11 @@ class TestMessageGenerator:
         channel.topic.assert_called_once_with(expected_key)
 
     def test_retrieve_events_passes_arguments(self):
+        topic = Mock()
+        topic.as_subscriber.return_value = topic
+        topic.subscribe.return_value = "subscription"
         with (
-            patch("core.app.apps.message_generator.MessageGenerator.get_response_topic", return_value="topic"),
+            patch("core.app.apps.message_generator.MessageGenerator.get_response_topic", return_value=topic),
             patch(
                 "core.app.apps.message_generator.stream_topic_events", return_value=iter([{"event": "ping"}])
             ) as mock_stream,
@@ -35,10 +62,30 @@ class TestMessageGenerator:
             )
 
         assert events == [{"event": "ping"}]
+        topic.as_subscriber.assert_called_once_with()
+        topic.subscribe.assert_called_once_with()
         mock_stream.assert_called_once_with(
-            topic="topic",
+            subscription="subscription",
             idle_timeout=1,
             ping_interval=2,
             on_subscribe=None,
             terminal_events=[StreamEvent.WORKFLOW_FINISHED.value],
         )
+
+    def test_retrieve_events_uses_prepared_subscription_capability(self):
+        topic = _TopicWithSeparateSubscriberView()
+        with (
+            patch("core.app.apps.message_generator.MessageGenerator.get_response_topic", return_value=topic),
+            patch("core.app.apps.message_generator.stream_topic_events", return_value=iter([])) as mock_stream,
+        ):
+            events = MessageGenerator.retrieve_events(AppMode.WORKFLOW, "run-1")
+
+        assert topic.subscriber.prepare_calls == 1
+        mock_stream.assert_called_once_with(
+            subscription="prepared-subscription",
+            idle_timeout=300,
+            ping_interval=10.0,
+            on_subscribe=None,
+            terminal_events=None,
+        )
+        assert list(events) == []

api/tests/unit_tests/core/app/apps/test_streaming_utils.py

@@ -8,6 +8,7 @@ import pytest
 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 libs.broadcast_channel.channel import SupportsPreparedSubscription
 from models.model import AppMode
 
 
@@ -42,8 +43,15 @@ class FakeTopic:
     def __init__(self) -> None:
         self._queue: queue.Queue[bytes] = queue.Queue()
         self._state = {"subscribed": False}
+        self.subscribe_calls = 0
+        self.subscriber_view_calls = 0
+
+    def as_subscriber(self):
+        self.subscriber_view_calls += 1
+        return self
 
     def subscribe(self) -> FakeSubscription:
+        self.subscribe_calls += 1
         return FakeSubscription(self._queue, self._state)
 
     def publish(self, payload: bytes) -> None:
@@ -54,7 +62,46 @@ class FakeTopic:
         return self._state["subscribed"]
 
 
-def test_retrieve_events_calls_on_subscribe_after_subscription(monkeypatch: pytest.MonkeyPatch):
+class FakePreparedSubscriber(SupportsPreparedSubscription):
+    def __init__(self, message_queue: queue.Queue[bytes], state: dict[str, bool]) -> None:
+        self._queue = message_queue
+        self._state = state
+        self.prepare_calls = 0
+
+    def prepare_subscription(self) -> FakeSubscription:
+        self.prepare_calls += 1
+        return FakeSubscription(self._queue, self._state)
+
+    def subscribe(self) -> FakeSubscription:
+        return FakeSubscription(self._queue, self._state)
+
+
+class FailingPreparedSubscriber(FakePreparedSubscriber):
+    def prepare_subscription(self) -> FakeSubscription:
+        raise RuntimeError("prepare failed")
+
+
+class FakePreparedTopic(FakeTopic):
+    def __init__(self) -> None:
+        super().__init__()
+        self.subscriber = FakePreparedSubscriber(self._queue, self._state)
+
+    def as_subscriber(self) -> FakePreparedSubscriber:
+        self.subscriber_view_calls += 1
+        return self.subscriber
+
+
+class FailingPreparedTopic(FakeTopic):
+    def __init__(self) -> None:
+        super().__init__()
+        self.subscriber = FailingPreparedSubscriber(self._queue, self._state)
+
+    def as_subscriber(self) -> FailingPreparedSubscriber:
+        self.subscriber_view_calls += 1
+        return self.subscriber
+
+
+def test_retrieve_events_falls_back_to_subscribe_and_invokes_hook_after_entry(monkeypatch: pytest.MonkeyPatch):
     topic = FakeTopic()
 
     def fake_get_response_topic(cls, app_mode, workflow_run_id):
@@ -74,6 +121,9 @@ def test_retrieve_events_calls_on_subscribe_after_subscription(monkeypatch: pyte
         on_subscribe=on_subscribe,
     )
 
+    assert topic.subscriber_view_calls == 1
+    assert topic.subscribe_calls == 1
+    assert topic.subscribed is False
     assert next(generator) == StreamEvent.PING.value
     event = next(generator)
     assert event["event"] == StreamEvent.WORKFLOW_FINISHED.value
@@ -81,6 +131,46 @@ def test_retrieve_events_calls_on_subscribe_after_subscription(monkeypatch: pyte
         next(generator)
 
 
+def test_retrieve_events_prepares_capable_topic_before_generator_iteration(monkeypatch: pytest.MonkeyPatch):
+    topic = FakePreparedTopic()
+
+    def fake_get_response_topic(cls, app_mode, workflow_run_id):
+        return topic
+
+    monkeypatch.setattr(MessageBasedAppGenerator, "get_response_topic", classmethod(fake_get_response_topic))
+
+    generator = MessageBasedAppGenerator.retrieve_events(
+        AppMode.WORKFLOW,
+        "workflow-run-id",
+        idle_timeout=0.5,
+    )
+
+    assert topic.subscriber_view_calls == 1
+    assert topic.subscriber.prepare_calls == 1
+    assert topic.subscribe_calls == 0
+    assert topic.subscribed is False
+
+    topic.publish(json.dumps({"event": StreamEvent.WORKFLOW_FINISHED.value}).encode())
+    assert next(generator) == StreamEvent.PING.value
+    assert next(generator)["event"] == StreamEvent.WORKFLOW_FINISHED.value
+
+
+def test_retrieve_events_propagates_preparation_error_before_generator_iteration(monkeypatch: pytest.MonkeyPatch):
+    topic = FailingPreparedTopic()
+
+    def fake_get_response_topic(cls, app_mode, workflow_run_id):
+        return topic
+
+    monkeypatch.setattr(MessageBasedAppGenerator, "get_response_topic", classmethod(fake_get_response_topic))
+
+    with pytest.raises(RuntimeError, match="prepare failed"):
+        MessageBasedAppGenerator.retrieve_events(
+            AppMode.WORKFLOW,
+            "workflow-run-id",
+            idle_timeout=0.5,
+        )
+
+
 def test_normalize_terminal_events_defaults():
     assert _normalize_terminal_events(None) == {
         StreamEvent.WORKFLOW_FINISHED.value,
@@ -94,6 +184,7 @@ def test_normalize_terminal_events_empty_values():
 
 def test_stream_topic_events_emits_ping_and_idle_timeout(monkeypatch: pytest.MonkeyPatch):
     topic = FakeTopic()
+    subscription = topic.subscribe()
     times = [1000.0, 1000.0, 1001.0, 1001.0, 1002.0]
 
     def fake_time():
@@ -102,7 +193,7 @@ def test_stream_topic_events_emits_ping_and_idle_timeout(monkeypatch: pytest.Mon
     monkeypatch.setattr("core.app.apps.streaming_utils.time.time", fake_time)
 
     generator = stream_topic_events(
-        topic=topic,
+        subscription=subscription,
         idle_timeout=10.0,
         ping_interval=1.0,
     )
@@ -116,9 +207,10 @@ def test_stream_topic_events_can_continue_past_pause():
     topic = FakeTopic()
     topic.publish(json.dumps({"event": StreamEvent.WORKFLOW_PAUSED.value}).encode())
     topic.publish(json.dumps({"event": StreamEvent.WORKFLOW_FINISHED.value}).encode())
+    subscription = topic.subscribe()
 
     generator = stream_topic_events(
-        topic=topic,
+        subscription=subscription,
         idle_timeout=1.0,
         terminal_events=[StreamEvent.WORKFLOW_FINISHED.value],
     )

api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel_unit_tests.py

@@ -1,3 +1,4 @@
+import abc
 import threading
 import time
 from dataclasses import dataclass
@@ -5,6 +6,7 @@ from typing import Any, cast
 
 import pytest
 
+from libs.broadcast_channel import channel as broadcast_channel
 from libs.broadcast_channel.exc import SubscriptionClosedError
 from libs.broadcast_channel.redis.streams_channel import (
     StreamsBroadcastChannel,
@@ -28,6 +30,8 @@ class FakeStreamsRedis:
         self._next_id: dict[str, int] = {}
         self._expire_calls: dict[str, int] = {}
         self._dollar_snapshots: dict[str, int] = {}
+        self._xread_calls = 0
+        self._xrevrange_calls = 0
 
     # Publisher API
     def xadd(self, key: str, fields: dict[str, Any], *, maxlen: int | None = None) -> str:
@@ -44,8 +48,16 @@ class FakeStreamsRedis:
     def expire(self, key: str, seconds: int) -> None:
         self._expire_calls[key] = self._expire_calls.get(key, 0) + 1
 
+    def xrevrange(self, key: str, count: int | None = None):
+        self._xrevrange_calls += 1
+        entries = list(reversed(self._store.get(key, [])))
+        if count is not None:
+            entries = entries[:count]
+        return entries
+
     # Consumer API
     def xread(self, streams: dict[str, Any], block: int | None = None, count: int | None = None):
+        self._xread_calls += 1
         # Expect a single key
         assert len(streams) == 1
         key, last_id = next(iter(streams.items()))
@@ -76,6 +88,11 @@ class FailExpireRedis(FakeStreamsRedis):
         raise RuntimeError("expire failed")
 
 
+class FailXrevrangeRedis(FakeStreamsRedis):
+    def xrevrange(self, key: str, count: int | None = None):
+        raise RuntimeError("xrevrange failed")
+
+
 class BlockingRedis:
     """A Redis mock whose xread blocks until a control event is xadd-ed."""
 
@@ -101,6 +118,12 @@ class BlockingRedis:
             return [(key, entries)]
         return []
 
+    def xrevrange(self, key: str, count: int | None = None):
+        entries = list(reversed(self._store.get(key, [])))
+        if count is not None:
+            entries = entries[:count]
+        return entries
+
     def release(self) -> None:
         self._release.set()
 
@@ -147,6 +170,14 @@ def streams_channel(fake_redis: FakeStreamsRedis) -> StreamsBroadcastChannel:
     return StreamsBroadcastChannel(fake_redis, retention_seconds=60)
 
 
+def test_prepared_subscription_capability_is_an_abstract_base_class():
+    capability = broadcast_channel.SupportsPreparedSubscription
+
+    assert issubclass(capability, abc.ABC)
+    assert broadcast_channel.Subscriber in capability.__mro__
+    assert capability.__abstractmethods__ == {"prepare_subscription", "subscribe"}
+
+
 class TestStreamsBroadcastChannel:
     def test_topic_creation(self, streams_channel: StreamsBroadcastChannel, fake_redis: FakeStreamsRedis):
         topic = streams_channel.topic("alpha")
@@ -183,11 +214,19 @@ class TestStreamsBroadcastChannel:
         assert fake_redis._store["enterprise-a:stream:beta"][0][1] == {b"data": b"hello"}
         assert fake_redis._expire_calls.get("enterprise-a:stream:beta", 0) >= 1
 
-    def test_topic_exposes_self_as_producer_and_subscriber(self, streams_channel: StreamsBroadcastChannel):
+    def test_topic_exposes_producer_and_subscriber_views(self, streams_channel: StreamsBroadcastChannel):
         topic = streams_channel.topic("producer-subscriber")
+        subscriber = topic.as_subscriber()
 
         assert topic.as_producer() is topic
-        assert topic.as_subscriber() is topic
+        assert subscriber is not topic
+        assert not isinstance(topic, broadcast_channel.SupportsPreparedSubscription)
+        assert isinstance(subscriber, broadcast_channel.SupportsPreparedSubscription)
+
+    def test_topic_explicitly_supports_prepared_subscriptions(self, streams_channel: StreamsBroadcastChannel):
+        topic = streams_channel.topic("prepared-capability")
+
+        assert isinstance(topic.as_subscriber(), broadcast_channel.SupportsPreparedSubscription)
 
     def test_publish_logs_warning_when_expire_fails(self, caplog: pytest.LogCaptureFixture):
         channel = StreamsBroadcastChannel(FailExpireRedis(), retention_seconds=60)
@@ -223,6 +262,136 @@ class TestStreamsSubscription:
 
         assert received == [b"after-subscribe-1", b"after-subscribe-2"]
 
+    def test_subscribe_starts_xread_at_latest_id_without_resolving_boundary(self):
+        start_ids: list[Any] = []
+
+        class OneReadRedis:
+            def xread(self, streams: dict[str, Any], block: int | None = None, count: int | None = None):
+                start_ids.extend(streams.values())
+                subscription._closed = True
+                return []
+
+        subscription = _StreamsSubscription(OneReadRedis(), "stream:latest-boundary")
+        subscription._listen()
+
+        assert start_ids == ["$"]
+
+    def test_prepare_subscription_fixes_boundary_without_starting_listener(
+        self,
+        streams_channel: StreamsBroadcastChannel,
+        fake_redis: FakeStreamsRedis,
+    ):
+        topic = streams_channel.topic("prepared-boundary")
+        topic.publish(b"before-prepare")
+        subscriber = topic.as_subscriber()
+        assert isinstance(subscriber, broadcast_channel.SupportsPreparedSubscription)
+
+        sub = subscriber.prepare_subscription()
+
+        assert isinstance(sub, _StreamsSubscription)
+        assert sub._start_id == "1-0"
+        assert sub._listener is None
+        assert fake_redis._xrevrange_calls == 1
+        assert fake_redis._xread_calls == 0
+
+    def test_prepared_subscription_includes_messages_published_before_entry_in_order(
+        self,
+        streams_channel: StreamsBroadcastChannel,
+        fake_redis: FakeStreamsRedis,
+    ):
+        topic = streams_channel.topic("prepared-ordering")
+        topic.publish(b"before-prepare")
+        subscriber = topic.as_subscriber()
+        assert isinstance(subscriber, broadcast_channel.SupportsPreparedSubscription)
+        sub = subscriber.prepare_subscription()
+
+        topic.publish(b"after-prepare-1")
+        topic.publish(b"after-prepare-2")
+
+        received: list[bytes] = []
+        with sub:
+            for _ in range(2):
+                message = sub.receive(timeout=0.1)
+                assert message is not None
+                received.append(message)
+
+        assert received == [b"after-prepare-1", b"after-prepare-2"]
+        assert fake_redis._xrevrange_calls == 1
+
+    def test_prepare_subscription_uses_beginning_boundary_for_empty_stream(
+        self,
+        streams_channel: StreamsBroadcastChannel,
+        fake_redis: FakeStreamsRedis,
+    ):
+        topic = streams_channel.topic("prepared-empty")
+        subscriber = topic.as_subscriber()
+        assert isinstance(subscriber, broadcast_channel.SupportsPreparedSubscription)
+
+        sub = subscriber.prepare_subscription()
+        assert isinstance(sub, _StreamsSubscription)
+        assert sub._start_id == "0-0"
+
+        topic.publish(b"first-event")
+        with sub:
+            assert sub.receive(timeout=0.1) == b"first-event"
+
+        assert fake_redis._xrevrange_calls == 1
+
+    def test_prepare_subscription_propagates_boundary_resolution_error(self):
+        channel = StreamsBroadcastChannel(FailXrevrangeRedis(), retention_seconds=60)
+        topic = channel.topic("broken-checkpoint")
+        subscriber = topic.as_subscriber()
+        assert isinstance(subscriber, broadcast_channel.SupportsPreparedSubscription)
+
+        with pytest.raises(RuntimeError, match="xrevrange failed"):
+            subscriber.prepare_subscription()
+
+    def test_prepared_subscription_receives_messages_published_right_after_entering(
+        self,
+        streams_channel: StreamsBroadcastChannel,
+    ):
+        """Regression test for the `workflow_started`-drop race (dify#40948).
+
+        A prepared subscription fixes the delivery boundary before the background task can
+        publish, so listener-thread scheduling cannot cause the first event to be missed.
+        """
+        topic = streams_channel.topic("race-topic")
+        subscriber = topic.as_subscriber()
+        assert isinstance(subscriber, broadcast_channel.SupportsPreparedSubscription)
+
+        sub = subscriber.prepare_subscription()
+        received: list[bytes] = []
+        with sub:
+            topic.publish(b"workflow_started")
+            for _ in range(5):
+                msg = sub.receive(timeout=0.1)
+                if msg is None:
+                    break
+                received.append(msg)
+
+        assert received == [b"workflow_started"]
+
+    def test_subscribe_does_not_replay_messages_published_before_it_started(
+        self,
+        streams_channel: StreamsBroadcastChannel,
+    ):
+        """`subscribe` preserves the "no stale replay" fix from #34030/#34040."""
+        topic = streams_channel.topic("no-replay-topic")
+        topic.publish(b"stale-event")
+
+        sub = topic.subscribe()
+        received: list[bytes] = []
+        with sub:
+            assert sub.receive(timeout=0.05) is None
+            topic.publish(b"fresh-event")
+            for _ in range(5):
+                msg = sub.receive(timeout=0.1)
+                if msg is None:
+                    break
+                received.append(msg)
+
+        assert received == [b"fresh-event"]
+
     def test_receive_timeout_returns_none(self, streams_channel: StreamsBroadcastChannel):
         topic = streams_channel.topic("delta")
         sub = topic.subscribe()
@@ -267,6 +436,7 @@ class TestStreamsSubscription:
                 return []
 
         subscription = _StreamsSubscription(OneShotRedis(case.fields), "stream:payload-shape")
+        subscription._start_id = "0-0"
         subscription._listen()
 
         received: list[bytes] = []
@@ -300,6 +470,7 @@ class TestStreamsSubscription:
                 return []
 
         subscription = _StreamsSubscription(OneShotRedis(), "stream:close-signal")
+        subscription._start_id = "0-0"
         subscription._listen()
 
         assert subscription._queue.get_nowait() == b"next-event"

api/tests/unit_tests/services/test_app_generate_service.py

@@ -14,8 +14,8 @@ Covers:
 """
 
 import threading
-import time
 import uuid
+from collections.abc import Callable
 from contextlib import contextmanager
 from unittest.mock import MagicMock
 
@@ -95,53 +95,88 @@ def _noop_rate_limit_context(rate_limit, request_id):
 # ---------------------------------------------------------------------------
 # _build_streaming_task_on_subscribe
 # ---------------------------------------------------------------------------
+class _FakeTimer:
+    def __init__(self, interval: float, function: Callable[[], bool]) -> None:
+        self.interval = interval
+        self.function = function
+        self.daemon = False
+        self.started = False
+        self.cancelled = False
+
+    def start(self) -> None:
+        self.started = True
+
+    def cancel(self) -> None:
+        self.cancelled = True
+
+
+def _unexpected_timer(interval: float, function: Callable[[], bool]) -> _FakeTimer:
+    raise AssertionError("streams must not create a fallback timer")
+
+
 class TestBuildStreamingTaskOnSubscribe:
-    """Tests for AppGenerateService._build_streaming_task_on_subscribe."""
-
-    def test_streams_mode_starts_immediately(self, monkeypatch: pytest.MonkeyPatch):
+    def test_streams_starts_only_when_hook_is_invoked_without_creating_timer(self, monkeypatch: pytest.MonkeyPatch):
         monkeypatch.setattr(ags_module.dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", "streams")
-        called = []
-        cb = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1))
-        # task started immediately during build
-        assert called == [1]
-        # calling the returned callback is idempotent
-        cb()
-        assert called == [1]  # not called again
 
-    def test_pubsub_mode_starts_on_subscribe(self, monkeypatch: pytest.MonkeyPatch):
-        monkeypatch.setattr(ags_module.dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", "pubsub")
-        monkeypatch.setattr(ags_module, "SSE_TASK_START_FALLBACK_MS", 60_000)  # large to prevent timer
-        called = []
-        cb = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1))
+        monkeypatch.setattr(ags_module.threading, "Timer", _unexpected_timer)
+        called: list[int] = []
+
+        on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1))
+
         assert called == []
-        cb()
-        assert called == [1]
-        # second call is idempotent
-        cb()
+        on_subscribe()
+        on_subscribe()
         assert called == [1]
 
-    def test_sharded_mode_starts_on_subscribe(self, monkeypatch: pytest.MonkeyPatch):
-        """sharded is treated like pubsub (i.e. not 'streams')."""
-        monkeypatch.setattr(ags_module.dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", "sharded")
-        monkeypatch.setattr(ags_module, "SSE_TASK_START_FALLBACK_MS", 60_000)
-        called = []
-        cb = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1))
+    @pytest.mark.parametrize("channel_type", ["pubsub", "sharded"])
+    def test_pubsub_transports_keep_subscribe_hook_and_fallback_timer(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        channel_type: str,
+    ):
+        monkeypatch.setattr(ags_module.dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", channel_type)
+        timers: list[_FakeTimer] = []
+
+        def build_timer(interval: float, function: Callable[[], bool]) -> _FakeTimer:
+            timer = _FakeTimer(interval, function)
+            timers.append(timer)
+            return timer
+
+        monkeypatch.setattr(ags_module.threading, "Timer", build_timer)
+        called: list[int] = []
+
+        on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1))
+
         assert called == []
-        cb()
-        assert called == [1]
+        assert len(timers) == 1
+        assert timers[0].interval == ags_module.SSE_TASK_START_FALLBACK_MS / 1000.0
+        assert timers[0].started is True
 
-    def test_pubsub_fallback_timer_fires(self, monkeypatch: pytest.MonkeyPatch):
-        """When nobody subscribes fast enough the fallback timer fires."""
+        on_subscribe()
+
+        assert called == [1]
+        assert timers[0].cancelled is True
+
+    def test_pubsub_fallback_starts_task_if_hook_is_never_invoked(self, monkeypatch: pytest.MonkeyPatch):
         monkeypatch.setattr(ags_module.dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", "pubsub")
-        monkeypatch.setattr(ags_module, "SSE_TASK_START_FALLBACK_MS", 50)  # 50 ms
-        called = []
-        _cb = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1))
-        time.sleep(0.2)  # give the timer time to fire
+        timers: list[_FakeTimer] = []
+
+        def build_timer(interval: float, function: Callable[[], bool]) -> _FakeTimer:
+            timer = _FakeTimer(interval, function)
+            timers.append(timer)
+            return timer
+
+        monkeypatch.setattr(ags_module.threading, "Timer", build_timer)
+        called: list[int] = []
+        on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(lambda: called.append(1))
+
+        assert timers[0].function() is True
+        on_subscribe()
         assert called == [1]
 
-    def test_exception_in_start_task_returns_false(self, monkeypatch: pytest.MonkeyPatch):
-        """When start_task raises, _try_start returns False and next call retries."""
+    def test_streams_retries_after_enqueue_failure(self, monkeypatch: pytest.MonkeyPatch):
         monkeypatch.setattr(ags_module.dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", "streams")
+        monkeypatch.setattr(ags_module.threading, "Timer", _unexpected_timer)
         call_count = 0
 
         def _bad():
@@ -150,15 +185,15 @@ class TestBuildStreamingTaskOnSubscribe:
             if call_count == 1:
                 raise RuntimeError("boom")
 
-        cb = AppGenerateService._build_streaming_task_on_subscribe(_bad)
-        # first call inside build raised, but is caught; second call via cb succeeds
+        on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(_bad)
+        on_subscribe()
         assert call_count == 1
-        cb()
+        on_subscribe()
         assert call_count == 2
 
     def test_concurrent_subscribe_only_starts_once(self, monkeypatch: pytest.MonkeyPatch):
-        monkeypatch.setattr(ags_module.dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", "pubsub")
-        monkeypatch.setattr(ags_module, "SSE_TASK_START_FALLBACK_MS", 60_000)
+        monkeypatch.setattr(ags_module.dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", "streams")
+        monkeypatch.setattr(ags_module.threading, "Timer", _unexpected_timer)
         call_count = 0
 
         def _inc():
@@ -368,7 +403,7 @@ class TestGenerate:
         retrieve_spy.assert_not_called()
 
     # -- ADVANCED_CHAT streaming --------------------------------------------
-    def test_advanced_chat_streaming(self, mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    def test_advanced_chat_streaming(self, mocker: MockerFixture):
         workflow = _make_workflow()
         mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow)
         mocker.patch(
@@ -376,9 +411,6 @@ class TestGenerate:
             return_value=MagicMock(workflow_run_id="wfr-1", model_dump_json=MagicMock(return_value="{}")),
         )
         delay_spy = mocker.patch("services.app_generate_service.workflow_based_app_execution_task.delay")
-        # Let _build_streaming_task_on_subscribe call the real on_subscribe
-        # so the inner closure (line 165) actually executes.
-        monkeypatch.setattr(ags_module.dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", "streams")
         gen_instance = MagicMock()
         gen_instance.retrieve_events.return_value = iter([])
         gen_instance.convert_to_event_stream.side_effect = lambda x: x
@@ -397,7 +429,10 @@ class TestGenerate:
         )
         # In streaming mode it should go through retrieve_events, not generate
         gen_instance.retrieve_events.assert_called_once()
-        # The inner on_subscribe closure was invoked by _build_streaming_task_on_subscribe
+        # Dispatch is gated on subscribe; simulate the SSE layer entering the
+        # subscription, which is what actually invokes on_subscribe.
+        on_subscribe = gen_instance.retrieve_events.call_args.kwargs["on_subscribe"]
+        on_subscribe()
         delay_spy.assert_called_once()
 
     # -- WORKFLOW blocking --------------------------------------------------
@@ -428,7 +463,7 @@ class TestGenerate:
         assert call_kwargs["pause_state_config"].state_owner_user_id == "owner-id"
 
     # -- WORKFLOW streaming -------------------------------------------------
-    def test_workflow_streaming(self, mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
+    def test_workflow_streaming(self, mocker: MockerFixture):
         workflow = _make_workflow()
         mocker.patch.object(AppGenerateService, "_get_workflow", return_value=workflow)
         mocker.patch(
@@ -436,9 +471,6 @@ class TestGenerate:
             return_value=MagicMock(workflow_run_id="wfr-2", model_dump_json=MagicMock(return_value="{}")),
         )
         delay_spy = mocker.patch("services.app_generate_service.workflow_based_app_execution_task.delay")
-        # Let _build_streaming_task_on_subscribe invoke the real on_subscribe
-        # so the inner closure (line 216) actually executes.
-        monkeypatch.setattr(ags_module.dify_config, "PUBSUB_REDIS_CHANNEL_TYPE", "streams")
         retrieve_spy = mocker.patch(
             "services.app_generate_service.MessageBasedAppGenerator.retrieve_events",
             return_value=iter([]),
@@ -457,7 +489,10 @@ class TestGenerate:
             session=MagicMock(),
         )
         retrieve_spy.assert_called_once()
-        # The inner on_subscribe closure was invoked by _build_streaming_task_on_subscribe
+        # Dispatch is gated on subscribe; simulate the SSE layer entering the
+        # subscription, which is what actually invokes on_subscribe.
+        on_subscribe = retrieve_spy.call_args.kwargs["on_subscribe"]
+        on_subscribe()
         delay_spy.assert_called_once()
 
     # -- Invalid mode -------------------------------------------------------

api/tests/unit_tests/services/test_app_generate_service_streaming_integration.py

@@ -5,6 +5,7 @@ from typing import Any
 
 import pytest
 
+from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
 from core.app.apps.message_generator import MessageGenerator
 from models.model import AppMode
 from services.app_generate_service import AppGenerateService
@@ -60,6 +61,9 @@ class _FakeStreams:
         # key -> list[(id, {field: value})]
         self._data: dict[str, list[tuple[str, dict]]] = defaultdict(list)
         self._seq: dict[str, int] = defaultdict(int)
+        # Mirrors real Redis: "$" resolves once, to the tail of the stream at the moment
+        # the first xread naming it is issued -- entries added earlier are never seen.
+        self._dollar_snapshots: dict[str, int] = {}
 
     def xadd(self, key: str, fields: dict[str, Any], *, maxlen: int | None = None) -> str:
         # maxlen is accepted for API compatibility with redis-py; ignored in this test double
@@ -72,12 +76,20 @@ class _FakeStreams:
         # no-op for tests
         return None
 
+    def xrevrange(self, key: str, count: int | None = None):
+        entries = list(reversed(self._data.get(key, [])))
+        if count is not None:
+            entries = entries[:count]
+        return entries
+
     def xread(self, streams: dict[str, Any], block: int | None = None, count: int | None = None):
         assert len(streams) == 1
         key, last_id = next(iter(streams.items()))
         entries = self._data.get(key, [])
         start = 0
-        if last_id != "0-0":
+        if last_id == "$":
+            start = self._dollar_snapshots.setdefault(key, len(entries))
+        elif last_id != "0-0":
             for i, (eid, _f) in enumerate(entries):
                 if eid == last_id:
                     start = i + 1
@@ -98,9 +110,10 @@ def _patch_get_channel_streams(monkeypatch: pytest.MonkeyPatch):
     def _get_channel():
         return chan
 
-    # Patch both the source and the imported alias used by MessageGenerator
+    # Patch the source and the imported aliases used by MessageGenerator / MessageBasedAppGenerator
     monkeypatch.setattr("extensions.ext_redis.get_pubsub_broadcast_channel", lambda: chan)
     monkeypatch.setattr("core.app.apps.message_generator.get_pubsub_broadcast_channel", lambda: chan)
+    monkeypatch.setattr("core.app.apps.message_based_app_generator.get_pubsub_broadcast_channel", lambda: chan)
     # Ensure AppGenerateService sees streams mode
     import services.app_generate_service as ags
 
@@ -136,19 +149,27 @@ def _publish_events(app_mode: AppMode, run_id: str, events: list[dict]):
 
 @pytest.mark.usefixtures("_patch_get_channel_streams")
 def test_streams_full_flow_prepublish_and_replay():
+    """Regression test for dify#40948: `workflow_started` dropped for concurrent streaming
+    runs on the streams transport.
+
+    `start_task` is invoked via `on_subscribe` after entering the subscription and publishes
+    both events synchronously, before the listener thread necessarily gets a chance to run its
+    first `xread`. `_StreamsSubscription` fixes its read boundary synchronously before spawning
+    that thread, so neither event is dropped.
+    """
     app_mode = AppMode.WORKFLOW
     run_id = str(uuid.uuid4())
 
-    # Build start_task that publishes two events immediately
     events = [{"event": "workflow_started"}, {"event": "workflow_finished"}]
 
     def start_task():
         _publish_events(app_mode, run_id, events)
 
+    # MessageBasedAppGenerator is what AppGenerateService actually dispatches WORKFLOW /
+    # ADVANCED_CHAT streaming runs through in production.
     on_subscribe = AppGenerateService._build_streaming_task_on_subscribe(start_task)
 
-    # Start retrieving BEFORE subscription is established; in streams mode, we also started immediately
-    gen = MessageGenerator.retrieve_events(app_mode, run_id, idle_timeout=2.0, on_subscribe=on_subscribe)
+    gen = MessageBasedAppGenerator.retrieve_events(app_mode, run_id, idle_timeout=2.0, on_subscribe=on_subscribe)
 
     received = []
     for msg in gen: