fix(agent): preserve Unicode separators in SSE streams (#39008)

This commit is contained in:
zyssyz123
2026-07-15 16:38:08 +08:00
committed by GitHub
parent d914ece3fb
commit 71d5c35e8f
4 changed files with 136 additions and 4 deletions
+56 -2
View File
@@ -94,6 +94,48 @@ class _ReconnectableStreamError(Exception):
super().__init__(str(error))
class _SSELineDecoder:
"""Split SSE text using only the line endings defined by the SSE specification."""
_buffer: str
def __init__(self) -> None:
self._buffer = ""
def decode(self, text: str) -> list[str]:
data = self._buffer + text
self._buffer = ""
lines: list[str] = []
start = 0
index = 0
while index < len(data):
char = data[index]
if char == "\n":
lines.append(data[start:index])
index += 1
start = index
continue
if char == "\r":
if index + 1 == len(data):
break
lines.append(data[start:index])
index += 2 if data[index + 1] == "\n" else 1
start = index
continue
index += 1
self._buffer = data[start:]
return lines
def flush(self) -> list[str]:
if not self._buffer:
return []
line = self._buffer[:-1] if self._buffer.endswith("\r") else self._buffer
self._buffer = ""
return [line]
class _SSEDecoder:
"""Incrementally decode SSE lines into typed run events.
@@ -624,7 +666,13 @@ class Client:
_ = await response.aread()
_raise_for_stream_status(response)
decoder = _SSEDecoder()
async for line in response.aiter_lines():
line_decoder = _SSELineDecoder()
async for text in response.aiter_text():
for line in line_decoder.decode(text):
event = decoder.feed_line(line)
if event is not None:
yield event
for line in line_decoder.flush():
event = decoder.feed_line(line)
if event is not None:
yield event
@@ -653,7 +701,13 @@ class Client:
_ = response.read()
_raise_for_stream_status(response)
decoder = _SSEDecoder()
for line in response.iter_lines():
line_decoder = _SSELineDecoder()
for text in response.iter_text():
for line in line_decoder.decode(text):
event = decoder.feed_line(line)
if event is not None:
yield event
for line in line_decoder.flush():
event = decoder.feed_line(line)
if event is not None:
yield event
+10 -1
View File
@@ -9,6 +9,14 @@ from collections.abc import AsyncIterable, AsyncIterator
from dify_agent.protocol.schemas import RUN_EVENT_ADAPTER, RunEvent
_SSE_UNSAFE_LINE_SEPARATORS = str.maketrans(
{
"\x85": "\\u0085",
"\u2028": "\\u2028",
"\u2029": "\\u2029",
}
)
def format_sse_event(event: RunEvent) -> str:
"""Serialize one event as an SSE frame."""
@@ -16,7 +24,8 @@ def format_sse_event(event: RunEvent) -> str:
if event.id is not None:
lines.append(f"id: {event.id}")
lines.append(f"event: {event.type}")
lines.append(f"data: {RUN_EVENT_ADAPTER.dump_json(event).decode()}")
payload = RUN_EVENT_ADAPTER.dump_json(event).decode().translate(_SSE_UNSAFE_LINE_SEPARATORS)
lines.append(f"data: {payload}")
return "\n".join(lines) + "\n\n"
@@ -28,6 +28,8 @@ from dify_agent.protocol import (
RunCancelledEvent,
RunEvent,
RunEventsResponse,
RunFailedEvent,
RunFailedEventData,
RunStartedEvent,
RunSucceededEvent,
RunSucceededEventData,
@@ -64,6 +66,14 @@ def _run_succeeded_event(*, event_id: str = "2-0", run_id: str = "run-1") -> Run
)
def _run_failed_event(error: str, *, event_id: str = "2-0", run_id: str = "run-1") -> RunFailedEvent:
return RunFailedEvent(
id=event_id,
run_id=run_id,
data=RunFailedEventData(error=error),
)
def _run_status_json(status: str) -> dict[str, object]:
now = datetime(2026, 5, 11, tzinfo=UTC).isoformat()
return {"run_id": "run-1", "status": status, "created_at": now, "updated_at": now, "error": None}
@@ -448,6 +458,26 @@ def test_sync_sse_parser_supports_comments_multiline_data_and_id_fill() -> None:
assert [event.type for event in events] == ["run_started"]
@pytest.mark.parametrize("separator", ["\x85", "\u2028", "\u2029"])
def test_sync_sse_parser_preserves_unicode_line_separators(separator: str) -> None:
error = f"before{separator}after"
body = _event_frame(_run_failed_event(error))
def handler(_request: httpx.Request) -> httpx.Response:
return httpx.Response(200, content=body)
client = Client(
base_url="http://testserver",
sync_http_client=httpx.Client(transport=httpx.MockTransport(handler)),
)
events = list(client.stream_events_sync("run-1", reconnect=False))
assert len(events) == 1
assert isinstance(events[0], RunFailedEvent)
assert events[0].data.error == error
def test_stream_events_stops_after_terminal_event() -> None:
calls = 0
body = "".join(
@@ -595,6 +625,27 @@ def test_async_stream_events_yields_terminal_event() -> None:
asyncio.run(scenario())
def test_async_sse_parser_preserves_unicode_line_separators() -> None:
error = "next-line:\x85line-separator:\u2028paragraph-separator:\u2029done"
body = _event_frame(_run_failed_event(error))
def handler(_request: httpx.Request) -> httpx.Response:
return httpx.Response(200, content=body)
async def scenario() -> None:
http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
client = Client(base_url="http://testserver", async_http_client=http_client)
events = [event async for event in client.stream_events("run-1")]
assert len(events) == 1
assert isinstance(events[0], RunFailedEvent)
assert events[0].data.error == error
await http_client.aclose()
asyncio.run(scenario())
def test_async_stream_events_reconnects_after_http_5xx_response() -> None:
seen_after: list[str] = []
@@ -1,4 +1,6 @@
from dify_agent.protocol.schemas import RunStartedEvent
import json
from dify_agent.protocol.schemas import RunFailedEvent, RunFailedEventData, RunStartedEvent
from dify_agent.server.sse import format_sse_event
@@ -10,3 +12,19 @@ def test_format_sse_event_uses_id_event_and_json_data() -> None:
assert frame.startswith("id: 7-0\nevent: run_started\ndata: ")
assert '"run_id":"run-1"' in frame
assert frame.endswith("\n\n")
def test_format_sse_event_escapes_unicode_line_separators() -> None:
error = "next-line:\x85line-separator:\u2028paragraph-separator:\u2029done"
event = RunFailedEvent(id="8-0", run_id="run-1", data=RunFailedEventData(error=error))
frame = format_sse_event(event)
data = next(line.removeprefix("data: ") for line in frame.splitlines() if line.startswith("data: "))
assert "\x85" not in frame
assert "\u2028" not in frame
assert "\u2029" not in frame
assert "\\u0085" in frame
assert "\\u2028" in frame
assert "\\u2029" in frame
assert json.loads(data)["data"]["error"] == error