From 9d8dacd0c8749ad66b4b4bcaf5d2bb5322b30fbc Mon Sep 17 00:00:00 2001 From: zl86790 Date: Fri, 21 Aug 2026 03:46:50 +0000 Subject: [PATCH] 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 --- .../app/apps/message_based_app_generator.py | 10 +- api/core/app/apps/message_generator.py | 10 +- api/core/app/apps/streaming_utils.py | 6 +- api/libs/broadcast_channel/channel.py | 10 + .../redis/streams_channel.py | 38 +++- api/services/app_generate_service.py | 18 +- .../redis/test_streams_channel.py | 38 +++- .../core/app/apps/test_message_generator.py | 51 ++++- .../core/app/apps/test_streaming_utils.py | 98 +++++++++- .../redis/test_streams_channel_unit_tests.py | 175 +++++++++++++++++- .../services/test_app_generate_service.py | 137 +++++++++----- ..._generate_service_streaming_integration.py | 31 +++- 12 files changed, 529 insertions(+), 93 deletions(-) diff --git a/api/core/app/apps/message_based_app_generator.py b/api/core/app/apps/message_based_app_generator.py index 6988ef124ea..acf590d9ddc 100644 --- a/api/core/app/apps/message_based_app_generator.py +++ b/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, ) diff --git a/api/core/app/apps/message_generator.py b/api/core/app/apps/message_generator.py index c04f20c7965..be3c4d0e07e 100644 --- a/api/core/app/apps/message_generator.py +++ b/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, diff --git a/api/core/app/apps/streaming_utils.py b/api/core/app/apps/streaming_utils.py index 5743bad4b6a..9bb31ca4103 100644 --- a/api/core/app/apps/streaming_utils.py +++ b/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: diff --git a/api/libs/broadcast_channel/channel.py b/api/libs/broadcast_channel/channel.py index a4957c97715..adf3e048149 100644 --- a/api/libs/broadcast_channel/channel.py +++ b/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. diff --git a/api/libs/broadcast_channel/redis/streams_channel.py b/api/libs/broadcast_channel/redis/streams_channel.py index c86d77ada6b..bdf73792187 100644 --- a/api/libs/broadcast_channel/redis/streams_channel.py +++ b/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}", diff --git a/api/services/app_generate_service.py b/api/services/app_generate_service.py index 1fc9d7be2da..4b271c6e94e 100644 --- a/api/services/app_generate_service.py +++ b/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() diff --git a/api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_streams_channel.py b/api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_streams_channel.py index a79208f649d..b859bed03df 100644 --- a/api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_streams_channel.py +++ b/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()) diff --git a/api/tests/unit_tests/core/app/apps/test_message_generator.py b/api/tests/unit_tests/core/app/apps/test_message_generator.py index 90c9abf35c0..a7ecf79b1e0 100644 --- a/api/tests/unit_tests/core/app/apps/test_message_generator.py +++ b/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) == [] 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 12f3ed9f078..66afebb3e60 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 @@ -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], ) diff --git a/api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel_unit_tests.py b/api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel_unit_tests.py index 73a65d3b5fc..e822bfc5bbd 100644 --- a/api/tests/unit_tests/libs/broadcast_channel/redis/test_streams_channel_unit_tests.py +++ b/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" diff --git a/api/tests/unit_tests/services/test_app_generate_service.py b/api/tests/unit_tests/services/test_app_generate_service.py index ff8eceed664..db80881b187 100644 --- a/api/tests/unit_tests/services/test_app_generate_service.py +++ b/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 ------------------------------------------------------- diff --git a/api/tests/unit_tests/services/test_app_generate_service_streaming_integration.py b/api/tests/unit_tests/services/test_app_generate_service_streaming_integration.py index 4fc08516d3d..fe667ef5771 100644 --- a/api/tests/unit_tests/services/test_app_generate_service_streaming_integration.py +++ b/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: