fix: parse and detect sse error event data from dynamo

Author: ajcasagrande

aiperf/clients/http/aiohttp_client.py

@@ -9,7 +9,8 @@
 
 from aiperf.clients.http.defaults import AioHttpDefaults, SocketDefaults
 from aiperf.clients.model_endpoint_info import ModelEndpointInfo
-from aiperf.common.enums import SSEFieldType
+from aiperf.common.enums import SSEEventType, SSEFieldType
+from aiperf.common.exceptions import SSEResponseError
 from aiperf.common.mixins import AIPerfLoggerMixin
 from aiperf.common.models import (
     ErrorDetails,
@@ -129,10 +130,14 @@ async def _request(
                         )
                     record.end_perf_ns = time.perf_counter_ns()
 
+        except SSEResponseError as e:
+            record.end_perf_ns = time.perf_counter_ns()
+            self.error(f"Error in SSE response: {e!r}")
+            record.error = ErrorDetails.from_exception(e)
         except Exception as e:
             record.end_perf_ns = time.perf_counter_ns()
             self.error(f"Error in aiohttp request: {e!r}")
-            record.error = ErrorDetails(type=e.__class__.__name__, message=str(e))
+            record.error = ErrorDetails.from_exception(e)
 
         return record
 
@@ -182,6 +187,27 @@ async def read_complete_stream(self) -> list[SSEMessage]:
         async for raw_message, first_byte_ns in self.__aiter__():
             # Parse the raw SSE message into a SSEMessage object
             message = parse_sse_message(raw_message, first_byte_ns)
+
+            # Check if the message contains an error in the format:
+            # event: error
+            # comment: <error message>
+            # If so, raise an exception with the error message.
+            if (
+                message.packets
+                and message.packets[0].name == SSEFieldType.EVENT
+                and message.packets[0].value == SSEEventType.ERROR
+            ):
+                if (
+                    len(message.packets) > 1
+                    and message.packets[1].name == SSEFieldType.COMMENT
+                ):
+                    error_message = message.packets[1].value
+                else:
+                    error_message = f"Unknown error {message.model_dump_json()}"
+                raise SSEResponseError(
+                    f"Error occurred in SSE response: {error_message}", error_code=502
+                )
+
             messages.append(message)
 
         return messages
@@ -211,18 +237,11 @@ async def __aiter__(self) -> typing.AsyncIterator[tuple[str, int]]:
                 break
             chunk = first_byte + chunk
 
-            try:
-                decoded = chunk.decode("utf-8")
-                for sub_chunk in decoded.split("\n\n"):
-                    if sub_chunk:
-                        yield (sub_chunk, chunk_ns_first_byte)
-                # Use the fastest available decoder
-            except UnicodeDecodeError:
-                # Handle potential encoding issues gracefully
-                yield (
-                    chunk.decode("utf-8", errors="replace").strip(),
-                    chunk_ns_first_byte,
-                )
+            # Replace invalid UTF-8 characters with the Unicode replacement character
+            decoded = chunk.decode("utf-8", errors="replace")
+            for sub_chunk in decoded.split("\n\n"):
+                if sub_chunk:
+                    yield (sub_chunk, chunk_ns_first_byte)
 
 
 def parse_sse_message(raw_message: str, perf_ns: int) -> SSEMessage:

aiperf/common/enums/__init__.py

@@ -93,6 +93,7 @@
     ServiceType,
 )
 from aiperf.common.enums.sse_enums import (
+    SSEEventType,
     SSEFieldType,
 )
 from aiperf.common.enums.system_enums import (
@@ -162,6 +163,7 @@
     "RecordProcessorType",
     "RequestRateMode",
     "ResultsProcessorType",
+    "SSEEventType",
     "SSEFieldType",
     "ServiceRegistrationStatus",
     "ServiceRunType",

aiperf/common/enums/sse_enums.py

@@ -12,3 +12,9 @@ class SSEFieldType(CaseInsensitiveStrEnum):
     ID = "id"
     RETRY = "retry"
     COMMENT = "comment"
+
+
+class SSEEventType(CaseInsensitiveStrEnum):
+    """Event types in an SSE message."""
+
+    ERROR = "error"

aiperf/common/exceptions.py

@@ -155,6 +155,14 @@ class ShutdownError(AIPerfError):
     """Exception raised when a service encounters an error while shutting down."""
 
 
+class SSEResponseError(AIPerfError):
+    """Exception raised when a SSE response contains an error."""
+
+    def __init__(self, message: str, error_code: int = 500) -> None:
+        self.error_code = error_code
+        super().__init__(message)
+
+
 class UnsupportedHookError(AIPerfError):
     """Exception raised when a hook is defined on a class that does not have any base classes that provide that hook type."""
 

aiperf/common/models/error_models.py

@@ -41,10 +41,13 @@ def __hash__(self) -> int:
     @classmethod
     def from_exception(cls, e: BaseException) -> "ErrorDetails":
         """Create an error details object from an exception."""
-        return cls(
+        error_details = cls(
             type=e.__class__.__name__,
             message=str(e),
         )
+        if hasattr(e, "error_code") and isinstance(e.error_code, int):
+            error_details.code = e.error_code
+        return error_details
 
 
 class ExitErrorInfo(AIPerfBaseModel):

tests/clients/http/test_aiohttp_sse.py

@@ -7,12 +7,17 @@
 ################################################################################
 
 import time
+from itertools import count
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
+from pytest import param
 
-from aiperf.clients.http.aiohttp_client import AioHttpSSEStreamReader
+from aiperf.clients.http.aiohttp_client import AioHttpSSEStreamReader, parse_sse_message
+from aiperf.common.enums import SSEFieldType
+from aiperf.common.exceptions import SSEResponseError
 from aiperf.common.models import SSEMessage
+from aiperf.common.models.error_models import ErrorDetails
 from tests.clients.http.conftest import (
     create_sse_chunk_list,
     setup_single_sse_chunk,
@@ -231,3 +236,232 @@ async def test_malformed_sse_stream(self, mock_sse_response: Mock) -> None:
         with pytest.raises(Exception, match="Stream corruption"):
             async for _ in reader:
                 pass
+
+
+@pytest.fixture
+async def stream_data_builder():
+    """Async fixture for building SSE stream data with auto-incrementing timestamps."""
+    timestamp_base = 123456789
+
+    def build(*messages: tuple[str, int | None]) -> list[tuple[str, int]]:
+        ts_gen = count(timestamp_base)
+        return [
+            (msg, ts_override if ts_override is not None else next(ts_gen))
+            for msg, ts_override in messages
+        ]
+
+    return build
+
+
+@pytest.fixture
+async def mock_sse_iterator(mock_sse_response: Mock):
+    """Async fixture that patches AioHttpSSEStreamReader.__aiter__ for testing."""
+
+    async def _read_stream(sse_data: list[tuple[str, int]]) -> list[SSEMessage]:
+        async def mock_aiter():
+            for data in sse_data:
+                yield data
+
+        with patch.object(
+            AioHttpSSEStreamReader, "__aiter__", return_value=mock_aiter()
+        ):
+            reader = AioHttpSSEStreamReader(Mock())
+            return await reader.read_complete_stream()
+
+    return _read_stream
+
+
+@pytest.mark.asyncio
+class TestSSEErrorHandling:
+    """Test suite for SSE error handling in AioHttpSSEStreamReader."""
+
+    TIMESTAMP = 123456789
+
+    @pytest.mark.parametrize(
+        "sse_message,expected_error_match",
+        [
+            param(
+                "event: error\ncomment: Invalid request",
+                "Invalid request",
+                id="error_with_invalid_request",
+            ),
+            param(
+                "event: error\ncomment: Server error",
+                "Server error",
+                id="error_with_server_error",
+            ),
+            param(
+                "event: error\ncomment: Auth failed",
+                "Auth failed",
+                id="error_with_auth_failed",
+            ),
+            param(
+                'event: error\ncomment: {"code": "E001", "msg": "API Error"}',
+                "Error occurred in SSE response",
+                id="error_with_json_message",
+            ),
+        ],
+    )
+    async def test_error_with_comment_message(
+        self,
+        sse_message: str,
+        expected_error_match: str,
+        mock_sse_iterator,
+    ) -> None:
+        """Test SSE error with event: error followed by comment raises error."""
+        with pytest.raises(SSEResponseError, match=expected_error_match):
+            await mock_sse_iterator([(sse_message, self.TIMESTAMP)])
+
+    @pytest.mark.parametrize(
+        "sse_message",
+        [
+            param("event: error", id="error_only"),
+            param("event: error\ndata: error details", id="error_with_data"),
+            param("event: error\nid: err-123", id="error_with_id"),
+        ],
+    )
+    async def test_error_without_comment_field(
+        self, sse_message: str, mock_sse_iterator
+    ) -> None:
+        """Test SSE error without comment field raises with unknown error message."""
+        with pytest.raises(SSEResponseError, match="Unknown error"):
+            await mock_sse_iterator([(sse_message, self.TIMESTAMP)])
+
+    @pytest.mark.parametrize(
+        "error_variant",
+        [
+            param("error", id="lowercase"),
+            param("ERROR", id="uppercase"),
+            param("Error", id="mixed_case"),
+            param("eRrOr", id="random_case"),
+        ],
+    )
+    async def test_error_case_insensitive(
+        self, error_variant: str, mock_sse_iterator
+    ) -> None:
+        """Test that event: error is case-insensitive."""
+        sse_message = f"event: error\ncomment: Test error from {error_variant}"
+        with pytest.raises(SSEResponseError, match="Test error"):
+            await mock_sse_iterator([(sse_message, self.TIMESTAMP)])
+
+    async def test_error_in_middle_of_stream(
+        self, stream_data_builder, mock_sse_iterator
+    ) -> None:
+        """Test that error in the middle of stream raises immediately."""
+        sse_data = stream_data_builder(
+            ("data: First message", None),
+            ("event: error\ncomment: Stream error", None),
+            ("data: This should not be processed", None),
+        )
+
+        with pytest.raises(SSEResponseError, match="Stream error"):
+            await mock_sse_iterator(sse_data)
+
+    async def test_error_code_is_502(self, mock_sse_iterator) -> None:
+        """Test that SSEResponseError has correct error code."""
+        with pytest.raises(SSEResponseError) as exc_info:
+            await mock_sse_iterator([("event: error\ncomment: Test", self.TIMESTAMP)])
+        assert exc_info.value.error_code == 502
+
+    async def test_error_details_preserved_with_502_code(
+        self, mock_sse_iterator
+    ) -> None:
+        """Test that error details are properly preserved when SSEResponseError with 502 is converted."""
+        error_message = "SSE stream error occurred"
+        with pytest.raises(SSEResponseError) as exc_info:
+            await mock_sse_iterator(
+                [("event: error\ncomment: " + error_message, self.TIMESTAMP)]
+            )
+
+        exc = exc_info.value
+        error_details = ErrorDetails.from_exception(exc)
+
+        assert error_details.code == 502
+        assert error_details.type == "SSEResponseError"
+        assert error_message in error_details.message
+
+    @pytest.mark.parametrize(
+        "event_type",
+        [
+            param("message", id="event_message"),
+            param("update", id="event_update"),
+            param("completion", id="event_completion"),
+            param("chunk", id="event_chunk"),
+            param("ping", id="event_ping"),
+        ],
+    )
+    async def test_normal_events_not_treated_as_errors(self, event_type: str) -> None:
+        """Test that normal event types are not treated as errors."""
+        raw_message = f"event: {event_type}\ndata: some data"
+        result = parse_sse_message(raw_message, self.TIMESTAMP)
+
+        assert result.packets[0].name == SSEFieldType.EVENT
+        assert result.packets[0].value == event_type
+
+    async def test_error_with_very_long_message(self, mock_sse_iterator) -> None:
+        """Test SSE error with very long error message (>1000 chars)."""
+        long_message = "x" * 5000
+        with pytest.raises(SSEResponseError, match=long_message[:100]):
+            await mock_sse_iterator(
+                [("event: error\ncomment: " + long_message, self.TIMESTAMP)]
+            )
+
+    @pytest.mark.parametrize(
+        "special_chars",
+        [
+            param("Error with 'quotes'", id="with_single_quotes"),
+            param('Error with "double quotes"', id="with_double_quotes"),
+            param("Error: with: multiple: colons", id="with_colons"),
+            param("Error\twith\ttabs", id="with_tabs"),
+            param("Error with émojis 🚀 💻", id="with_emojis"),
+            param("Error with <html> tags", id="with_html_tags"),
+            param("Error with | pipe | symbols", id="with_pipes"),
+        ],
+    )
+    async def test_error_with_special_characters(
+        self, special_chars: str, mock_sse_iterator
+    ) -> None:
+        """Test SSE error with special characters in message."""
+        with pytest.raises(SSEResponseError):
+            await mock_sse_iterator(
+                [("event: error\ncomment: " + special_chars, self.TIMESTAMP)]
+            )
+
+    async def test_error_with_only_whitespace_comment(self, mock_sse_iterator) -> None:
+        """Test SSE error with only whitespace in comment field."""
+        with pytest.raises(SSEResponseError):
+            await mock_sse_iterator([("event: error\ncomment:   ", self.TIMESTAMP)])
+
+    async def test_multiple_errors_in_sequence(
+        self, stream_data_builder, mock_sse_iterator
+    ) -> None:
+        """Test that first error in sequence is raised immediately."""
+        sse_data = stream_data_builder(
+            ("event: error\ncomment: First error", None),
+            ("event: error\ncomment: Second error", None),
+        )
+
+        with pytest.raises(SSEResponseError, match="First error"):
+            await mock_sse_iterator(sse_data)
+
+    async def test_error_with_mixed_additional_fields(self, mock_sse_iterator) -> None:
+        """Test SSE error with multiple additional fields mixed in."""
+        sse_message = (
+            "event: error\n"
+            "id: err-001\n"
+            "comment: Mixed fields error\n"
+            "retry: 3000\n"
+            "data: extra data"
+        )
+        with pytest.raises(SSEResponseError, match="Mixed fields error"):
+            await mock_sse_iterator([(sse_message, self.TIMESTAMP)])
+
+    async def test_error_with_special_escape_sequences(self) -> None:
+        """Test parsing error message with escape sequences in raw format."""
+        raw_message = "event: error\ncomment: Error\\nwith\\nnewlines"
+        result = parse_sse_message(raw_message, self.TIMESTAMP)
+
+        assert result.packets[0].name == SSEFieldType.EVENT
+        assert result.packets[0].value == "error"
+        assert result.packets[1].name == SSEFieldType.COMMENT
+        assert result.packets[1].value == "Error\\nwith\\nnewlines"