From 1d5b5a88ca835f3527c605ee723fbdb09cc1f916 Mon Sep 17 00:00:00 2001 From: vaibhavsrv Date: Thu, 17 Sep 2026 10:04:12 +0530 Subject: [PATCH] fix(pixel-edge): accumulate trailing buffered content and synthesize fallback finish --- .github/workflows/test-linux.yml | 1 + .../services/pixel-edge/pixel_edge.py | 5 +- ods/tests/test_pixel_edge_buffered_content.py | 91 +++++++++++++++++++ 3 files changed, 96 insertions(+), 1 deletion(-) create mode 100644 ods/tests/test_pixel_edge_buffered_content.py diff --git a/.github/workflows/test-linux.yml b/.github/workflows/test-linux.yml index 65a2a7a515..331743b152 100644 --- a/.github/workflows/test-linux.yml +++ b/.github/workflows/test-linux.yml @@ -414,6 +414,7 @@ jobs: python3 -m unittest -v tests/test_pixel_native_search.py bash tests/test-pixel-compose-wiring.sh python3 tests/test_pixel_timeout_contract.py + python3 -m unittest -v tests/test_pixel_edge_buffered_content.py python3 -m unittest -v extensions/services/pixel-agent/tests/test_artifact_promoter.py node --test extensions/services/pixel-agent/tests/*.test.mjs diff --git a/ods/extensions/services/pixel-edge/pixel_edge.py b/ods/extensions/services/pixel-edge/pixel_edge.py index aa35d69cb3..b87f6952c9 100644 --- a/ods/extensions/services/pixel-edge/pixel_edge.py +++ b/ods/extensions/services/pixel-edge/pixel_edge.py @@ -1360,6 +1360,8 @@ async def replace_pending(template: dict, *, synthesize_finish: bool): else: event, content, finish_reason = _sse_event(line) queue_pending(line, event, content, finish_reason) + if content is not None: + pending_text += content if not passthrough and pending: normalized = pending_text.strip() if not normalized or normalized in _RESERVED_ASSISTANT_REPLIES: @@ -1367,7 +1369,8 @@ async def replace_pending(template: dict, *, synthesize_finish: bool): (item[1] for item in reversed(pending) if isinstance(item[1], dict)), {"model": _PIXEL_REWRITE}, ) - await replace_pending(template, synthesize_finish=False) + has_finish = any(item[3] is not None for item in pending) + await replace_pending(template, synthesize_finish=not has_finish) else: await flush_pending() except (ConnectionError, OSError, asyncio.TimeoutError) as exc: diff --git a/ods/tests/test_pixel_edge_buffered_content.py b/ods/tests/test_pixel_edge_buffered_content.py new file mode 100644 index 0000000000..d01c9f8fe2 --- /dev/null +++ b/ods/tests/test_pixel_edge_buffered_content.py @@ -0,0 +1,91 @@ +"""Regression test for un-terminated trailing buffered SSE in Pixel Edge.""" + +from __future__ import annotations + +import asyncio +import os +import sys +import unittest +from pathlib import Path +from unittest.mock import MagicMock + +os.environ.setdefault("PIXEL_OPENWEBUI_KEY", "test-pixel-openwebui-key-01234567890123456789") +os.environ.setdefault("PIXEL_PREVIEW_PROXY_KEY", "test-pixel-preview-proxy-key-01234567890123456789") + +EDGE_DIR = Path(__file__).resolve().parents[1] / "extensions" / "services" / "pixel-edge" +edge = None + + +class _FakeStreamContent: + def __init__(self, chunks: list[bytes]) -> None: + self._chunks = chunks + + async def iter_any(self): + for chunk in self._chunks: + yield chunk + + +class _FakeUpstreamResponse: + def __init__(self, chunks: list[bytes], status: int = 200) -> None: + self.status = status + self.content = _FakeStreamContent(chunks) + + +class _CaptureResponse: + def __init__(self) -> None: + self.status = 200 + self.written: list[bytes] = [] + + async def prepare(self, _request) -> None: + pass + + async def write(self, data: bytes) -> None: + self.written.append(data) + + +class PixelEdgeBufferedContentTests(unittest.IsolatedAsyncioTestCase): + @classmethod + def setUpClass(cls) -> None: + global edge + try: + import aiohttp # noqa: F401 + sys.path.insert(0, str(EDGE_DIR)) + import pixel_edge as edge_module + edge = edge_module + except ImportError: + raise unittest.SkipTest("aiohttp is required for pixel edge tests") + def setUp(self) -> None: + self.orig_stream_response = edge.web.StreamResponse + + def tearDown(self) -> None: + edge.web.StreamResponse = self.orig_stream_response + + async def test_trailing_buffered_content_is_flushed_not_replaced(self) -> None: + # Single chunk without trailing newline at EOF + chunks = [b'data: {"model":"m","choices":[{"delta":{"content":"Tokyo"}}]}'] + resp = _FakeUpstreamResponse(chunks) + capture = _CaptureResponse() + edge.web.StreamResponse = lambda **kw: capture + + await edge._stream_upstream(MagicMock(), resp, "fallback reply text") + output = b"".join(capture.written).decode("utf-8", errors="replace") + + self.assertIn("Tokyo", output) + self.assertNotIn("fallback reply text", output) + + async def test_empty_trailing_buffer_synthesizes_finished_fallback(self) -> None: + # Empty chunk without trailing newline at EOF triggers fallback with synthesized finish + chunks = [b'data: {"model":"m","choices":[{"delta":{}}]}'] + resp = _FakeUpstreamResponse(chunks) + capture = _CaptureResponse() + edge.web.StreamResponse = lambda **kw: capture + + await edge._stream_upstream(MagicMock(), resp, "fallback reply text") + output = b"".join(capture.written).decode("utf-8", errors="replace") + + self.assertIn("fallback reply text", output) + self.assertIn('"finish_reason": "stop"', output) + + +if __name__ == "__main__": + unittest.main()