feat: event polling can process more than one event at a time (#7858)

* refactor: enhance event polling to retrieve all available events and return in NDJSON format

Updated the event polling logic to collect all available events from the queue instead of just one. The response format has been changed to NDJSON, allowing multiple events to be returned in a single response. Improved error handling for timeouts by returning an empty response instead of raising an error.

* refactor: improve NDJSON event processing in build polling

Enhanced the event polling logic to handle NDJSON responses by processing multiple events in a single fetch. Updated the response handling to read the response as text, split it into individual JSON objects, and process each event accordingly. This change improves the efficiency of event handling and maintains a consistent polling interval.

* fix: change event list type from list[str] to list in get_flow_events_response

* refactor: enhance event stream consumption with timeout and error handling

Updated the consume_and_assert_stream function to include a timeout for processing events and improved error handling for JSON parsing failures. Added logging for better debugging and tracking of processed events. Adjusted the event validation logic to ensure both the first and last expected events are present.

* refactor: enhance event polling logic with improved limits and JSON validation

Updated the event polling mechanism to include configurable limits for total events and empty polls, along with a timeout for requests. Enhanced JSON validation for individual events and added logging for better debugging of polling behavior. This refactor aims to improve robustness and maintainability of the event consumption process.

* refactor: update polling interval and streamline event handling

Modified the polling interval constant to improve responsiveness and adjusted event handling logic in the get_flow_events_response function to simplify the processing of event values. This change enhances the efficiency of event consumption and maintains a consistent approach to handling NDJSON responses.

* refactor: simplify event retrieval in get_flow_events_response

Removed unused event_id and put_time variables from the event retrieval logic in the get_flow_events_response function. This change streamlines the code and enhances readability while maintaining the functionality of event handling.

* refactor: update consume_and_assert_stream to use asyncio.wait_for for timeout handling

Modified the consume_and_assert_stream function to utilize asyncio.wait_for for managing timeouts in event processing. This change improves compatibility with Python 3.10 and enhances the robustness of the event consumption logic while maintaining existing functionality.

src/backend/tests/unit/test_chat_endpoint.py

@@ -1,4 +1,5 @@
 import asyncio
+import json
 import uuid
 from uuid import UUID
 
@@ -6,6 +7,7 @@ import pytest
 from httpx import codes
 from langflow.memory import aget_messages
 from langflow.services.database.models.flow import FlowUpdate
+from loguru import logger
 
 from tests.unit.build_utils import build_flow, consume_and_assert_stream, create_flow, get_build_events
 
@@ -160,33 +162,87 @@ async def test_build_flow_polling(client, json_memory_chatbot_no_llm, logged_in_
             self.job_id = job_id
             self.headers = headers
             self.status_code = codes.OK
+            self.max_total_events = 50  # Limit to prevent infinite loops
+            self.max_empty_polls = 10  # Maximum number of empty polls before giving up
+            self.poll_timeout = 1.0  # Timeout for each polling request
 
         async def aiter_lines(self):
             try:
-                sleeps = 0
-                max_sleeps = 100
-                while True:
-                    response = await self.client.get(
-                        f"api/v1/build/{self.job_id}/events?event_delivery=polling", headers=self.headers
-                    )
-                    assert response.status_code == codes.OK
-                    data = response.json()
+                empty_polls = 0
+                total_events = 0
+                end_event_found = False
 
-                    if data["event"] is None:
-                        # No event available, add delay to prevent tight polling
+                while (
+                    empty_polls < self.max_empty_polls and total_events < self.max_total_events and not end_event_found
+                ):
+                    # Add Accept header for NDJSON
+                    headers = {**self.headers, "Accept": "application/x-ndjson"}
+
+                    # Set a timeout for the request
+                    response = await asyncio.wait_for(
+                        self.client.get(
+                            f"api/v1/build/{self.job_id}/events?event_delivery=polling",
+                            headers=headers,
+                        ),
+                        timeout=self.poll_timeout,
+                    )
+
+                    assert response.status_code == codes.OK
+
+                    # Get the NDJSON response as text
+                    text = response.text
+
+                    # Skip if response is empty
+                    if not text.strip():
+                        empty_polls += 1
                         await asyncio.sleep(0.1)
-                        sleeps += 1
                         continue
 
-                    yield data["event"]
+                    # Reset empty polls counter since we got data
+                    empty_polls = 0
+
+                    # Process each line as an individual JSON object
+                    line_count = 0
+                    for line in text.splitlines():
+                        if not line.strip():
+                            continue
+
+                        line_count += 1
+                        total_events += 1
+
+                        # Check for end event with multiple possible formats
+                        if '"event":"end"' in line or '"event": "end"' in line:
+                            end_event_found = True
+
+                        # Validate it's proper JSON before yielding
+                        try:
+                            json.loads(line)  # Test parse to ensure it's valid JSON
+                            yield line
+                        except json.JSONDecodeError as e:
+                            logger.debug(f"WARNING: Skipping invalid JSON: {line}")
+                            logger.debug(f"Error: {e}")
+                            # Don't yield invalid JSON, but continue processing other lines
+
+                    # If we had no events in this batch, count as empty poll
+                    if line_count == 0:
+                        empty_polls += 1
+
+                    # Add a small delay to prevent tight polling
+                    await asyncio.sleep(0.1)
+
+                # If we hit the limit without finding the end event, log a warning
+                if total_events >= self.max_total_events:
+                    logger.debug(
+                        f"WARNING: Reached maximum event limit ({self.max_total_events}) without finding end event"
+                    )
+
+                if empty_polls >= self.max_empty_polls and not end_event_found:
+                    logger.debug(
+                        f"WARNING: Reached maximum empty polls ({self.max_empty_polls}) without finding end event"
+                    )
 
-                    # If this was the end event, stop polling
-                    if '"end"' in data["event"]:
-                        break
-                    if sleeps > max_sleeps:
-                        msg = "Build event polling timed out."
-                        raise TimeoutError(msg)
             except asyncio.TimeoutError as e:
+                logger.debug(f"ERROR: Polling request timed out after {self.poll_timeout}s")
                 msg = "Build event polling timed out."
                 raise TimeoutError(msg) from e