mirror of
https://github.com/open-webui/open-webui.git
synced 2026-08-28 05:27:35 +00:00
fix: long streamed lines no longer abort the response (#28114)
Some providers send one very large piece of a streamed answer in a single go: a long reasoning trace, a code execution result, a turn with many tool calls, or a response echo carrying a big tool list. Anything past 128 KB in one line killed the chat mid-answer with a misleading `400, message: Got more than 131072 bytes when reading`. Nothing was rejected upstream, that is our own reader giving up on an oversized line. Open WebUI already had code that assembles lines itself with no such limit, but it only ran when CHAT_STREAM_RESPONSE_CHUNK_MAX_BUFFER_SIZE was set. Unset is the default, and in that case the raw capped reader was used instead, so a default install always broke. That path now always assembles lines, and the setting goes back to being what its name says: an optional cap, off by default. It applies to the Ollama stream as well, since both now share the same reader. The assembly loop only splits once a line actually completes, because the old one re-concatenated and re-split the whole buffer on every network chunk. Without that, allowing long lines would have traded an error for multi-second event loop stalls. | | 20 MB in one line | 200k small lines | | --- | --- | --- | | before | 4249 ms | 27.3 ms | | after | 37 ms | 25.2 ms |
This commit is contained in:
parent
beb3c114d3
commit
2d2bcb5332
3 changed files with 31 additions and 51 deletions
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue