diff --git a/backend/open_webui/routers/openai.py b/backend/open_webui/routers/openai.py index 3bc56ebcf6..0f21defdd5 100644 --- a/backend/open_webui/routers/openai.py +++ b/backend/open_webui/routers/openai.py @@ -43,10 +43,7 @@ from open_webui.utils.anthropic import ANTHROPIC_VERSION, get_anthropic_models, from open_webui.utils.auth import get_admin_user, get_verified_user from open_webui.utils.headers import get_custom_headers, include_user_info_headers from open_webui.utils.json_codec import JSONCodec -from open_webui.utils.misc import ( - convert_logit_bias_input_to_json, - stream_chunks_handler, -) +from open_webui.utils.misc import convert_logit_bias_input_to_json from open_webui.utils.model_ids import strip_provider_model_prefix from open_webui.utils.payload import ( apply_model_params_to_body_openai, @@ -1679,7 +1676,7 @@ async def generate_chat_completion( streaming = True return StreamingResponse( - stream_wrapper(r, content_handler=stream_chunks_handler), + stream_wrapper(r), status_code=r.status, headers=_clean_proxy_headers(r.headers), ) diff --git a/backend/open_webui/utils/misc.py b/backend/open_webui/utils/misc.py index 96fedc3cd9..2761adc191 100644 --- a/backend/open_webui/utils/misc.py +++ b/backend/open_webui/utils/misc.py @@ -1241,64 +1241,48 @@ async def stream_wrapper(response, session, content_handler=None): def stream_chunks_handler(stream: aiohttp.StreamReader): """ - Handle stream response chunks without using aiohttp's line reader. - When configured and a single line exceeds max_buffer_size, returns an empty - JSON string {} and skips subsequent data until encountering normally sized data. + Assemble lines from raw chunks, so a line over aiohttp's reader limit no longer aborts the stream. + When CHAT_STREAM_RESPONSE_CHUNK_MAX_BUFFER_SIZE is set, a line exceeding it is dropped. :param stream: The stream reader to handle. - :return: An async generator that yields the stream data. + :return: An async generator that yields the stream one line at a time. """ max_buffer_size = CHAT_STREAM_RESPONSE_CHUNK_MAX_BUFFER_SIZE - if max_buffer_size is not None and max_buffer_size <= 0: - max_buffer_size = None + if max_buffer_size is None or max_buffer_size <= 0: + max_buffer_size = float('inf') # unset: no line is too long async def yield_safe_stream_chunks(): - buffer = b'' - skip_mode = False + buffer = bytearray() # bytearray, not bytes: `+=` on bytes reallocates, quadratic on long lines + dropping_line_tail = False async for data, _ in stream.iter_chunks(): if not data: continue - # In skip_mode, if buffer already exceeds the limit, clear it (it's part of an oversized line) - if max_buffer_size is not None and skip_mode and len(buffer) > max_buffer_size: - buffer = b'' + buffer += data - lines = (buffer + data).split(b'\n') + # Only split once a line completed: splitting every chunk re-copies the buffer, quadratic + if b'\n' in data: + *lines, rest = bytes(buffer).split(b'\n') + buffer = bytearray(rest) - # Process complete lines (except the last possibly incomplete fragment) - for i in range(len(lines) - 1): - line = lines[i] - - if skip_mode: - # Skip mode: check if current line is small enough to exit skip mode - if max_buffer_size is None or len(line) <= max_buffer_size: - skip_mode = False - yield line - else: - yield b'data: {}\n' - else: - # Normal mode: check if line exceeds limit - if max_buffer_size is not None and len(line) > max_buffer_size: - skip_mode = True - yield b'data: {}\n' - log.info('Skip mode triggered, line size: %s', len(line)) + for line in lines: + if dropping_line_tail: + dropping_line_tail = False + elif len(line) > max_buffer_size: + log.info('Dropped line over max buffer size: %s bytes', len(line)) else: yield line + b'\n' - # Save the last incomplete fragment - buffer = lines[-1] + # Oversized line still arriving: drop it instead of buffering the rest + if len(buffer) > max_buffer_size: + if not dropping_line_tail: + log.info('Dropping line over max buffer size, buffered so far: %s bytes', len(buffer)) + dropping_line_tail = True + buffer.clear() - # Check if buffer exceeds limit - if max_buffer_size is not None and not skip_mode and len(buffer) > max_buffer_size: - skip_mode = True - log.info('Skip mode triggered, buffer size: %s', len(buffer)) - # Clear oversized buffer to prevent unlimited growth - buffer = b'' - - # Process remaining buffer data - if buffer and not skip_mode: - yield buffer + b'\n' + if buffer and not dropping_line_tail: + yield bytes(buffer) return yield_safe_stream_chunks() diff --git a/backend/open_webui/utils/session_pool.py b/backend/open_webui/utils/session_pool.py index fb66c3e2cf..bc9a3a8f74 100644 --- a/backend/open_webui/utils/session_pool.py +++ b/backend/open_webui/utils/session_pool.py @@ -34,6 +34,7 @@ from open_webui.env import ( AIOHTTP_POOL_CONNECTIONS_PER_HOST, AIOHTTP_POOL_DNS_TTL, ) +from open_webui.utils.misc import stream_chunks_handler log = logging.getLogger(__name__) @@ -115,7 +116,7 @@ async def cleanup_response( await result -async def stream_wrapper(response, session=None, content_handler=None, passthrough=False): +async def stream_wrapper(response, session=None, passthrough=False): """Wrap a stream to ensure cleanup happens even if streaming is interrupted. This is more reliable than BackgroundTask which may not run if the client @@ -126,12 +127,10 @@ async def stream_wrapper(response, session=None, content_handler=None, passthrou line. Only for streams no internal consumer parses line-by-line. """ try: - if content_handler: - stream = content_handler(response.content) - elif passthrough: + if passthrough: stream = response.content.iter_any() else: - stream = response.content + stream = stream_chunks_handler(response.content) async for chunk in stream: yield chunk finally: