From 3b802a0b857c318c1e4ac3e294022efb32f3a3eb Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 15 Apr 2026 08:11:54 +0000 Subject: [PATCH] fix(stream): drop message-in-chat check, refresh seq TTL every append MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Critical: the message-in-chat check introduced in 72429ea would reject resume for exactly the flow this feature exists for. Traced the frontend submitPrompt → sendMessage → sendMessageSocket path: the assistant stub is added to in-memory history before the stream starts, but saveChatHandler/initChatHandler that persists it to DB runs inside the completion event handler, not before. During active streaming, the stub isn't in the DB yet, so `message_id in chat.history.messages` would return False and block replay. Dropped that portion. Kept the chat ownership check since that's checking a pre-existing chat — which IS in the DB — and the user-scoped log key already enforces cross-user isolation even when message binding isn't validated. Warning: seq-key TTL was only refreshed every 64 appends. A sparse stream emitting fewer than 64 frames per RESUME_STREAM_TTL_SEC (1h) could let the counter expire mid-stream and INCR would restart at 1, corrupting replay. Refresh on every append instead — pipelined with the XADD so no extra round-trip. Removed the unused refresh cadence constant. Deferred: prune resumeSeqByMessageId on per-message terminal. This is the third time the bot has flip-flopped on this; the current design (no per-message prune, cleared at chat/navigation boundaries) was chosen in response to its earlier finding that pruning caused continuation-reuses-message_id to replay duplicate content. Memory is int-per-message bounded by chat size. --- backend/open_webui/socket/main.py | 39 ++++++++++++------------------- 1 file changed, 15 insertions(+), 24 deletions(-) diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index 99579fa9ab..93fc44b241 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -180,7 +180,6 @@ YDOC_MANAGER = YdocManager( RESUME_STREAM_MAXLEN = 2000 RESUME_STREAM_TTL_SEC = 3600 RESUME_STREAM_DONE_TTL_SEC = 30 -RESUME_STREAM_TTL_REFRESH_EVERY = 64 # Hot-path timeout: tight so a slow Redis can't stall live tokens. # Replay read timeout: looser since a resume is user-blocking anyway # and silent timeout here is worse than a brief extra wait. Both @@ -304,7 +303,6 @@ async def _stream_log_append(user_id: str, message_id: str, envelope: dict, seq: if _breaker_open(): return try: - refresh_ttl = (seq == 1) or (seq % RESUME_STREAM_TTL_REFRESH_EVERY == 0) key = _stream_key(user_id, message_id) seq_key = _stream_seq_key(user_id, message_id) pipe = REDIS.pipeline(transaction=False) @@ -314,12 +312,12 @@ async def _stream_log_append(user_id: str, message_id: str, envelope: dict, seq: maxlen=RESUME_STREAM_MAXLEN, approximate=True, ) - if refresh_ttl: - # Refresh both keys together — without this the seq counter - # can expire mid-stream on responses longer than the TTL and - # INCR restarts at 1, corrupting replay ordering. - pipe.expire(key, RESUME_STREAM_TTL_SEC) - pipe.expire(seq_key, RESUME_STREAM_TTL_SEC) + # Refresh both keys every append. Pipelined with XADD so no extra + # round-trip; guarantees sparse / slow streams (where the seq + # counter might otherwise idle past the TTL and INCR would + # restart at 1) stay alive as long as the stream is active. + pipe.expire(key, RESUME_STREAM_TTL_SEC) + pipe.expire(seq_key, RESUME_STREAM_TTL_SEC) await asyncio.wait_for( pipe.execute(), timeout=RESUME_STREAM_REDIS_TIMEOUT_SEC ) @@ -765,13 +763,14 @@ async def resume_stream(sid, data): last_seq = int(data.get('last_seq') or 0) except (TypeError, ValueError): last_seq = 0 - # Defense-in-depth: validate chat ownership AND message-in-chat - # binding even though the log key is already user-scoped. Fails - # CLOSED on DB errors so an infra hiccup can't skip the check. - # When chat_id is absent we fall through to the user-scoped key - # guarantee so legacy clients that don't send chat_id still work. - # Either way we reply with (at worst) empty envelopes below so - # the client fence clears deterministically. + # Defense-in-depth: validate chat ownership when chat_id is + # present. The message-in-chat lookup can legitimately fail for + # in-flight assistants that haven't been persisted to DB yet + # (ENABLE_REALTIME_CHAT_SAVE=False flow), which is exactly the + # case this feature is meant to recover — so we don't gate + # replay on message presence. The user-scoped log key already + # prevents cross-user access regardless. Fails CLOSED on DB + # errors or when the user doesn't own the supplied chat_id. chat_id = data.get('chat_id') chat_ok = True if chat_id: @@ -779,16 +778,8 @@ async def resume_stream(sid, data): chat = await Chats.get_chat_by_id_and_user_id(chat_id, user_id) if not chat: chat_ok = False - else: - # Reuse the chat record for the message-in-chat check - # so we don't pay a second DB round-trip. - messages = ( - getattr(chat, 'chat', {}) or {} - ).get('history', {}).get('messages', {}) or {} - if message_id not in messages: - chat_ok = False except Exception as e: - log.warning(f'resume-stream chat/message check failed: {e}') + log.warning(f'resume-stream chat ownership check failed: {e}') chat_ok = False if chat_ok: envelopes = await _stream_log_read(user_id, message_id, last_seq)