From cffe1537eccceeb87ff643cb98a1bde8d264c3cd Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 15 Apr 2026 07:32:50 +0000 Subject: [PATCH] fix(stream): preserve ack cb on buffered frames, fail-closed on DB, per-message lock - chatEventHandler queue now stores {event, cb} pairs. Live frames buffered during a replay fence used to lose their Socket.IO ack callback because we only pushed the event. When replayed, callback- style events (confirmation / execute / input) would silently never respond. Now we preserve cb and pass it through to chatEventHandler on flush. - resume_stream ownership check is now fail-CLOSED on DB exception. Previous fail-open skipped the check precisely when infra was unstable. Client still gets a deterministic reply (empty envelopes) so the fence clears; it just can't resume in that specific failure mode. Absent chat_id still falls through to the user-scoped key guarantee. - Added per-message asyncio.Lock via a WeakValueDictionary keyed by (user_id, message_id). Serializes seq-alloc + log-append + emit so concurrent emitters for the same message_id within a worker can't interleave and produce out-of-seq live frames. WeakValueDictionary cleans up automatically once no coroutine holds the lock. Cross-worker concurrency is still unprotected (would need a distributed lock) but that's an even rarer scenario. --- backend/open_webui/socket/main.py | 77 ++++++++++++++++++----------- src/lib/components/chat/Chat.svelte | 11 +++-- 2 files changed, 56 insertions(+), 32 deletions(-) diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index 761b70b524..7a585830a9 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -2,6 +2,7 @@ import asyncio import json import os import random +import weakref import socketio import logging @@ -201,6 +202,26 @@ _HOT_PATH_BREAKER_FAILURE_THRESHOLD = 3 _HOT_PATH_BREAKER_COOLDOWN_SEC = 10.0 _hot_path_breaker = {'failures': 0, 'open_until': 0.0} +# Per-message emitter lock. Serializes the seq-alloc + log-append + emit +# sequence within a worker so two overlapping emitters for the same +# (user_id, message_id) can't interleave and produce out-of-seq live +# frames. Doesn't cover cross-worker concurrency, but concurrent +# emitters for the same message_id on different workers is an even +# rarer scenario. WeakValueDictionary cleans up entries automatically +# once no coroutine holds the lock (i.e., no one is inside the critical +# section and no one is waiting to enter). +_emit_locks = weakref.WeakValueDictionary() + + +def _emit_lock_for(user_id: str, message_id: str) -> asyncio.Lock: + key = f'{user_id}:{message_id}' + lock = _emit_locks.get(key) + if lock is None: + # setdefault is atomic under the GIL; races resolve to a single + # Lock and subsequent calls see it via the first branch. + lock = _emit_locks.setdefault(key, asyncio.Lock()) + return lock + def _breaker_open() -> bool: return time.time() < _hot_path_breaker['open_until'] @@ -732,10 +753,12 @@ async def resume_stream(sid, data): except (TypeError, ValueError): last_seq = 0 # Defense-in-depth: validate chat ownership even though the log - # key is already user-scoped. Rejects only when we can *prove* - # ownership fails — a missing/unknown chat_id falls through to - # the user-scoped key guarantee so we don't regress the "stub - # not yet persisted in DB" refresh case. + # key is already user-scoped. Fails CLOSED on DB errors so an + # infra hiccup can't skip the ownership 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. chat_id = data.get('chat_id') chat_ok = True if chat_id: @@ -744,7 +767,7 @@ async def resume_stream(sid, data): chat_ok = chat is not None except Exception as e: log.warning(f'resume-stream chat ownership check failed: {e}') - chat_ok = True # fail open to not regress legitimate callers + chat_ok = False if chat_ok: envelopes = await _stream_log_read(user_id, message_id, last_seq) @@ -1049,36 +1072,32 @@ async def disconnect(sid): async def get_event_emitter(request_info, update_db=True): - # A single emitter's calls are serial (the streaming loop awaits each - # before the next), so seq ordering is guaranteed within one emitter. - # CONCURRENT emitters for the same (user_id, message_id) — e.g. a - # duplicate request leaking past frontend dedup, or a retry path that - # overlaps with the original — can interleave INCR/XADD/emit across - # tasks and cause live frames to arrive out of seq order. Replay - # reads already sort by seq so resume is safe, but live streaming in - # that edge case can drop a frame via the client dedupe guard. Fixing - # it properly needs either a distributed per-message lock or a - # client-side reorder buffer; neither is worth the complexity for a - # configuration OWUI doesn't normally produce. + # Concurrency note: within one worker the _emit_lock_for serializes + # seq-alloc + log-append + emit per (user_id, message_id), so + # overlapping emitters can't interleave and produce out-of-seq live + # frames. Cross-worker concurrent emitters for the same message_id + # are still unprotected (would need a distributed lock), but that's + # a configuration OWUI doesn't normally produce. async def __event_emitter__(event_data): user_id = request_info['user_id'] chat_id = request_info['chat_id'] message_id = request_info['message_id'] - seq = await _stream_seq_allocate(user_id, message_id) + async with _emit_lock_for(user_id, message_id): + seq = await _stream_seq_allocate(user_id, message_id) - envelope = { - 'chat_id': chat_id, - 'message_id': message_id, - 'data': event_data, - } - # Log before emit so a reconnecting client can't resume-read past - # a frame that hasn't been persisted yet. Client seq guard drops - # duplicates from the inverted race. - if seq is not None: - envelope['seq'] = seq - await _stream_log_append(user_id, message_id, envelope, seq) - await sio.emit('events', envelope, room=f'user:{user_id}') + envelope = { + 'chat_id': chat_id, + 'message_id': message_id, + 'data': event_data, + } + # Log before emit so a reconnecting client can't resume-read past + # a frame that hasn't been persisted yet. Client seq guard drops + # duplicates from the inverted race. + if seq is not None: + envelope['seq'] = seq + await _stream_log_append(user_id, message_id, envelope, seq) + await sio.emit('events', envelope, room=f'user:{user_id}') # Any terminal event shortens TTL so log + seq self-evict together. # Covers normal completion (done:True), explicit cancel, and diff --git a/src/lib/components/chat/Chat.svelte b/src/lib/components/chat/Chat.svelte index 1c6b1e37cd..0e9a331646 100644 --- a/src/lib/components/chat/Chat.svelte +++ b/src/lib/components/chat/Chat.svelte @@ -452,9 +452,12 @@ if (message) { // Buffer live frames during replay; _replayed frames skip the fence. + // Store the ack callback alongside the event so Socket.IO + // call-style events (confirmation/execute/input) don't lose + // their response path when buffered and later replayed. const queue = resumeQueueByMessageId.get(event.message_id); if (queue && !event?._replayed) { - queue.push(event); + queue.push({ event, cb }); return; } @@ -664,12 +667,14 @@ const batch = resumeQueueByMessageId.get(messageId); if (!batch || batch.length === 0) break; resumeQueueByMessageId.set(messageId, []); - for (const event of batch) { + for (const item of batch) { + const event = item?.event; + const cb = item?.cb; if (event && typeof event === 'object') { event._replayed = true; } try { - await chatEventHandler(event); + await chatEventHandler(event, cb); } catch (e) { console.error('resume fence flush error', e); }