From dd91a61e958fbc2afe6826be36d7b38b2b29c779 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 4 Aug 2026 13:19:52 -0700 Subject: [PATCH] fix: parse Server-Sent Events in Agent Engine streaming responses Agent Engine streaming responses are newline-delimited JSON. A response may instead arrive as Server-Sent Events, with each JSON object wrapped in a `data:` frame. The streaming readers now remove that framing before parsing, so both shapes yield the same parsed objects rather than raw `data: {...}` strings, and the `HttpBody` reader parses a `text/event-stream` response rather than yielding the unparsed message. Newline-delimited JSON responses are unaffected. PiperOrigin-RevId: 959194089 --- agentplatform/_genai/_agent_engines_utils.py | 45 +++++++++++++++- .../agentplatform/genai/test_agent_engines.py | 54 +++++++++++++++++++ vertexai/_genai/_agent_engines_utils.py | 36 +++++++++++++ vertexai/agent_engines/_utils.py | 43 ++++++++++++++- vertexai/reasoning_engines/_utils.py | 43 ++++++++++++++- 5 files changed, 218 insertions(+), 3 deletions(-) diff --git a/agentplatform/_genai/_agent_engines_utils.py b/agentplatform/_genai/_agent_engines_utils.py index 6cb42423c4..cd23c35e51 100644 --- a/agentplatform/_genai/_agent_engines_utils.py +++ b/agentplatform/_genai/_agent_engines_utils.py @@ -1888,6 +1888,39 @@ async def _method(self, **kwargs) -> Any: # type: ignore[no-untyped-def] return _method # type: ignore[return-value] +_SSE_DATA_PREFIX = "data:" +_STREAMABLE_CONTENT_TYPES = ("application/json", "text/event-stream") + + +def _strip_sse_framing(line: str) -> str: + """Returns the payload of a Server-Sent Events `data:` line. + + Streaming responses are newline-delimited JSON. A response may instead + arrive as Server-Sent Events, in which case each JSON object is wrapped in a + `data:` frame; removing that framing here lets both shapes be parsed the + same way + (https://github.com/googleapis/python-aiplatform/issues/5586). + + A serialized JSON value never begins with `data:` -- it begins with `{`, + `[`, `"`, a digit, or one of `true`/`false`/`null` -- so stripping the + prefix unconditionally cannot corrupt a newline-delimited JSON response. A + chunk whose value is the string `data: hello` is serialized as + `"data: hello"`, with the quote first. + + Args: + line: A single line of the response body. + + Returns: + The line with any SSE `data:` framing removed. + """ + line = line.rstrip("\r") + if not line.startswith(_SSE_DATA_PREFIX): + return line + # The single space after the colon is optional per the SSE specification. + payload = line[len(_SSE_DATA_PREFIX) :] + return payload[1:] if payload.startswith(" ") else payload + + def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterator[Any]: """Converts the body of the HTTP Response message to JSON format. @@ -1904,6 +1937,9 @@ def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterat # Handle the case of multiple dictionaries delimited by newlines. for line in http_response.body.split("\n"): + # Strip before the emptiness check so the blank line that terminates an + # SSE frame, and a `data:` line with an empty payload, are both skipped. + line = _strip_sse_framing(line) if line: try: line = json.loads(line) @@ -1931,7 +1967,11 @@ def _yield_parsed_json_from_httpbody(body: httpbody_pb2.HttpBody) -> Iterator[An content_type = getattr(body, "content_type", None) data = getattr(body, "data", None) - if content_type is None or data is None or "application/json" not in content_type: + if ( + content_type is None + or data is None + or not any(t in content_type for t in _STREAMABLE_CONTENT_TYPES) + ): yield body return @@ -1948,6 +1988,9 @@ def _yield_parsed_json_from_httpbody(body: httpbody_pb2.HttpBody) -> Iterator[An # Handle the case of multiple dictionaries delimited by newlines. for line in utf8_data.split("\n"): + # Strip before the emptiness check so the blank line that terminates an + # SSE frame, and a `data:` line with an empty payload, are both skipped. + line = _strip_sse_framing(line) if line: try: line = json.loads(line) diff --git a/tests/unit/agentplatform/genai/test_agent_engines.py b/tests/unit/agentplatform/genai/test_agent_engines.py index e92f121e8d..c64094379f 100644 --- a/tests/unit/agentplatform/genai/test_agent_engines.py +++ b/tests/unit/agentplatform/genai/test_agent_engines.py @@ -1805,6 +1805,60 @@ def test_yield_parsed_json_from_httpbody(self, obj, expected): got = list(_agent_engines_utils._yield_parsed_json_from_httpbody(obj)) assert got == expected + + # pytest does not allow absl.testing.parameterized.named_parameters. + @pytest.mark.parametrize( + "obj, expected", + [ + ( + # "sse_single_event", + genai_types.HttpResponse(body='data: {"a": 1}\n\n'), + [{"a": 1}], + ), + ( + # "sse_multiple_events", + genai_types.HttpResponse( + body='data: {"a": 1}\n\ndata: {"a": 2}\n\n' + ), + [{"a": 1}, {"a": 2}], + ), + ( + # "sse_no_space_after_colon", + genai_types.HttpResponse(body='data:{"a": 1}\n\n'), + [{"a": 1}], + ), + ( + # "sse_crlf_line_endings", + genai_types.HttpResponse(body='data: {"a": 1}\r\n\r\n'), + [{"a": 1}], + ), + ( + # "sse_empty_data_line_is_skipped", + genai_types.HttpResponse(body='data:\n\ndata: {"a": 1}\n\n'), + [{"a": 1}], + ), + ( + # "json_string_value_beginning_with_data_is_untouched", + genai_types.HttpResponse(body='"data: hello"'), + ["data: hello"], + ), + ], + ) + def test_to_parsed_json_server_sent_events(self, obj, expected): + """An SSE-framed response is parsed into the same objects as NDJSON.""" + assert list(_agent_engines_utils._yield_parsed_json(obj)) == expected + + def test_yield_parsed_json_from_httpbody_event_stream(self): + """The gRPC path parses SSE instead of yielding the raw proto.""" + body = httpbody_pb2.HttpBody( + content_type="text/event-stream", + data=b'data: {"a": 1}\n\ndata: {"a": 2}\n\n', + ) + assert list(_agent_engines_utils._yield_parsed_json_from_httpbody(body)) == [ + {"a": 1}, + {"a": 2}, + ] + def test_yield_parsed_json_from_httpbody_non_json_content_type(self): body = httpbody_pb2.HttpBody(content_type="text/plain", data=b"hello") assert list(_agent_engines_utils._yield_parsed_json_from_httpbody(body)) == [ diff --git a/vertexai/_genai/_agent_engines_utils.py b/vertexai/_genai/_agent_engines_utils.py index e3a2fb8b68..52c9c58e3a 100644 --- a/vertexai/_genai/_agent_engines_utils.py +++ b/vertexai/_genai/_agent_engines_utils.py @@ -2007,6 +2007,39 @@ async def _method(self, **kwargs) -> Any: # type: ignore[no-untyped-def] _wrap_a2a_operation = _wrap_a2a_operation_v03 +_SSE_DATA_PREFIX = "data:" +_STREAMABLE_CONTENT_TYPES = ("application/json", "text/event-stream") + + +def _strip_sse_framing(line: str) -> str: + """Returns the payload of a Server-Sent Events `data:` line. + + Streaming responses are newline-delimited JSON. A response may instead + arrive as Server-Sent Events, in which case each JSON object is wrapped in a + `data:` frame; removing that framing here lets both shapes be parsed the + same way + (https://github.com/googleapis/python-aiplatform/issues/5586). + + A serialized JSON value never begins with `data:` -- it begins with `{`, + `[`, `"`, a digit, or one of `true`/`false`/`null` -- so stripping the + prefix unconditionally cannot corrupt a newline-delimited JSON response. A + chunk whose value is the string `data: hello` is serialized as + `"data: hello"`, with the quote first. + + Args: + line: A single line of the response body. + + Returns: + The line with any SSE `data:` framing removed. + """ + line = line.rstrip("\r") + if not line.startswith(_SSE_DATA_PREFIX): + return line + # The single space after the colon is optional per the SSE specification. + payload = line[len(_SSE_DATA_PREFIX) :] + return payload[1:] if payload.startswith(" ") else payload + + def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterator[Any]: """Converts the body of the HTTP Response message to JSON format. @@ -2023,6 +2056,9 @@ def _yield_parsed_json(http_response: google_genai_types.HttpResponse) -> Iterat # Handle the case of multiple dictionaries delimited by newlines. for line in http_response.body.split("\n"): + # Strip before the emptiness check so the blank line that terminates an + # SSE frame, and a `data:` line with an empty payload, are both skipped. + line = _strip_sse_framing(line) if line: try: line = json.loads(line) diff --git a/vertexai/agent_engines/_utils.py b/vertexai/agent_engines/_utils.py index f7c359c93d..5126c957d3 100644 --- a/vertexai/agent_engines/_utils.py +++ b/vertexai/agent_engines/_utils.py @@ -305,6 +305,40 @@ def to_json_serializable_autogen_object( return _autogen_run_response_protocol_to_dict(obj) + +_SSE_DATA_PREFIX = "data:" +_STREAMABLE_CONTENT_TYPES = ("application/json", "text/event-stream") + + +def _strip_sse_framing(line: str) -> str: + """Returns the payload of a Server-Sent Events `data:` line. + + Streaming responses are newline-delimited JSON. A response may instead + arrive as Server-Sent Events, in which case each JSON object is wrapped in a + `data:` frame; removing that framing here lets both shapes be parsed the + same way + (https://github.com/googleapis/python-aiplatform/issues/5586). + + A serialized JSON value never begins with `data:` -- it begins with `{`, + `[`, `"`, a digit, or one of `true`/`false`/`null` -- so stripping the + prefix unconditionally cannot corrupt a newline-delimited JSON response. A + chunk whose value is the string `data: hello` is serialized as + `"data: hello"`, with the quote first. + + Args: + line: A single line of the response body. + + Returns: + The line with any SSE `data:` framing removed. + """ + line = line.rstrip("\r") + if not line.startswith(_SSE_DATA_PREFIX): + return line + # The single space after the colon is optional per the SSE specification. + payload = line[len(_SSE_DATA_PREFIX) :] + return payload[1:] if payload.startswith(" ") else payload + + def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]: """Converts the contents of the httpbody message to JSON format. @@ -318,7 +352,11 @@ def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]: content_type = getattr(body, "content_type", None) data = getattr(body, "data", None) - if content_type is None or data is None or "application/json" not in content_type: + if ( + content_type is None + or data is None + or not any(t in content_type for t in _STREAMABLE_CONTENT_TYPES) + ): yield body return @@ -335,6 +373,9 @@ def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]: # Handle the case of multiple dictionaries delimited by newlines. for line in utf8_data.split("\n"): + # Strip before the emptiness check so the blank line that terminates an + # SSE frame, and a `data:` line with an empty payload, are both skipped. + line = _strip_sse_framing(line) if line: try: line = json.loads(line) diff --git a/vertexai/reasoning_engines/_utils.py b/vertexai/reasoning_engines/_utils.py index dbb0938748..eea0121b8d 100644 --- a/vertexai/reasoning_engines/_utils.py +++ b/vertexai/reasoning_engines/_utils.py @@ -162,6 +162,40 @@ def to_json_serializable_llama_index_object( return str(obj) + +_SSE_DATA_PREFIX = "data:" +_STREAMABLE_CONTENT_TYPES = ("application/json", "text/event-stream") + + +def _strip_sse_framing(line: str) -> str: + """Returns the payload of a Server-Sent Events `data:` line. + + Streaming responses are newline-delimited JSON. A response may instead + arrive as Server-Sent Events, in which case each JSON object is wrapped in a + `data:` frame; removing that framing here lets both shapes be parsed the + same way + (https://github.com/googleapis/python-aiplatform/issues/5586). + + A serialized JSON value never begins with `data:` -- it begins with `{`, + `[`, `"`, a digit, or one of `true`/`false`/`null` -- so stripping the + prefix unconditionally cannot corrupt a newline-delimited JSON response. A + chunk whose value is the string `data: hello` is serialized as + `"data: hello"`, with the quote first. + + Args: + line: A single line of the response body. + + Returns: + The line with any SSE `data:` framing removed. + """ + line = line.rstrip("\r") + if not line.startswith(_SSE_DATA_PREFIX): + return line + # The single space after the colon is optional per the SSE specification. + payload = line[len(_SSE_DATA_PREFIX) :] + return payload[1:] if payload.startswith(" ") else payload + + def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]: """Converts the contents of the httpbody message to JSON format. @@ -175,7 +209,11 @@ def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]: content_type = getattr(body, "content_type", None) data = getattr(body, "data", None) - if content_type is None or data is None or "application/json" not in content_type: + if ( + content_type is None + or data is None + or not any(t in content_type for t in _STREAMABLE_CONTENT_TYPES) + ): yield body return @@ -192,6 +230,9 @@ def yield_parsed_json(body: httpbody_pb2.HttpBody) -> Iterable[Any]: # Handle the case of multiple dictionaries delimited by newlines. for line in utf8_data.split("\n"): + # Strip before the emptiness check so the blank line that terminates an + # SSE frame, and a `data:` line with an empty payload, are both skipped. + line = _strip_sse_framing(line) if line: try: line = json.loads(line)