fix: resolve race condition in SSE streaming causing empty message re…

api/core/app/apps/base_app_queue_manager.py

@@ -52,9 +52,16 @@ def __init__(self, task_id: str, user_id: str, invoke_from: InvokeFrom):
         self._stopped_cache: TTLCache[tuple, bool] = TTLCache(maxsize=1, ttl=1)
         self._cache_lock = threading.Lock()
 
+        # Graceful shutdown mechanism to prevent race condition (fixes #31611)
+        self._should_stop = threading.Event()
+        self._stop_lock = threading.Lock()
+
     def listen(self):
         """
-        Listen to queue
+        Listen to queue with graceful shutdown support.
+
+        This method implements a graceful shutdown mechanism to prevent race conditions
+        where messages could be lost during queue shutdown (fixes #31611).
         :return:
         """
         # wait for APP_MAX_EXECUTION_TIME seconds to stop listen
@@ -63,12 +70,22 @@ def listen(self):
         last_ping_time: int | float = 0
         while True:
             try:
-                message = self._q.get(timeout=1)
+                # Use shorter timeout to be more responsive to stop signals
+                message = self._q.get(timeout=0.5)
                 if message is None:
+                    # Drain remaining messages before exiting to prevent message loss
+                    yield from self._drain_remaining_messages()
                     break
 
                 yield message
             except queue.Empty:
+                # Check graceful stop signal
+                if self._should_stop.is_set():
+                    # Give a small window for any pending messages
+                    time.sleep(0.05)
+                    # Drain any remaining messages
+                    yield from self._drain_remaining_messages()
+                    break
                 continue
             finally:
                 elapsed_time = time.time() - start_time
@@ -83,14 +100,51 @@ def listen(self):
                     self.publish(QueuePingEvent(), PublishFrom.TASK_PIPELINE)
                     last_ping_time = elapsed_time // 10
 
+    def _drain_remaining_messages(self):
+        """
+        Drain all remaining messages from the queue.
+
+        This ensures no messages are lost during shutdown (fixes #31611).
+        """
+        drained_count = 0
+        max_drain = 1000  # Safety limit to prevent infinite loop
+
+        while drained_count < max_drain:
+            try:
+                remaining = self._q.get_nowait()
+                if remaining is not None:
+                    drained_count += 1
+                    yield remaining
+                else:
+                    # Another None, we're done
+                    break
+            except queue.Empty:
+                break
+
+        if drained_count > 0:
+            logger.debug("Drained %d remaining messages from queue", drained_count)
+
     def stop_listen(self):
         """
-        Stop listen to queue
+        Stop listen to queue gracefully.
+
+        This method implements graceful shutdown to prevent race conditions
+        where messages could be lost (fixes #31611).
         :return:
         """
-        self._clear_task_belong_cache()
-        self._q.put(None)
-        self._graph_runtime_state = None  # Release reference to allow GC to reclaim memory
+        with self._stop_lock:
+            self._clear_task_belong_cache()
+
+            # Set stop signal first to allow pending publishes to complete
+            self._should_stop.set()
+
+            # Small delay to allow pending publishes to complete
+            # This helps prevent the race condition where messages are
+            # published after we've started stopping but before we put None
+            time.sleep(0.02)
+
+            self._q.put(None)
+            self._graph_runtime_state = None  # Release reference to allow GC to reclaim memory
 
     def _clear_task_belong_cache(self) -> None:
         """

api/core/app/apps/message_based_app_queue_manager.py

@@ -1,3 +1,5 @@
+import time
+
 from core.app.apps.base_app_queue_manager import AppQueueManager, PublishFrom
 from core.app.apps.exc import GenerateTaskStoppedError
 from core.app.entities.app_invoke_entities import InvokeFrom
@@ -21,9 +23,15 @@ def __init__(
         self._app_mode = app_mode
         self._message_id = str(message_id)
 
+        # Terminal event delay to prevent race condition (fixes #31611)
+        self._terminal_event_delay = 0.05
+
     def _publish(self, event: AppQueueEvent, pub_from: PublishFrom):
         """
-        Publish event to queue
+        Publish event to queue with improved synchronization.
+
+        This method includes a small delay before stopping to prevent race conditions
+        where messages could be lost during queue shutdown (fixes #31611).
         :param event:
         :param pub_from:
         :return:
@@ -41,7 +49,32 @@ def _publish(self, event: AppQueueEvent, pub_from: PublishFrom):
         if isinstance(
             event, QueueStopEvent | QueueErrorEvent | QueueMessageEndEvent | QueueAdvancedChatMessageEndEvent
         ):
+            # Add delay to allow concurrent publishes to complete
+            # This prevents the race condition where:
+            # 1. Thread A is about to publish a LLMChunkEvent
+            # 2. Thread B publishes MessageEndEvent and calls stop_listen()
+            # 3. Thread A's message never gets processed
+            time.sleep(self._terminal_event_delay)
+
+            # Wait for queue to be reasonably empty before stopping
+            self._wait_for_queue_flush(timeout=1.0)
+
             self.stop_listen()
 
         if pub_from == PublishFrom.APPLICATION_MANAGER and self._is_stopped():
             raise GenerateTaskStoppedError()
+
+    def _wait_for_queue_flush(self, timeout: float = 1.0):
+        """
+        Wait for the queue to be flushed (or timeout).
+
+        This gives the consumer a chance to process pending messages (fixes #31611).
+        """
+        start_time = time.time()
+        check_interval = 0.01  # 10ms
+
+        while time.time() - start_time < timeout:
+            # Check if queue is empty or nearly empty
+            if self._q.qsize() <= 1:  # Allow for the terminal event itself
+                break
+            time.sleep(check_interval)