From 610e0b8fab7ad5065cbb37bc9f5dc2f9e4b1ee87 Mon Sep 17 00:00:00 2001 From: Classic298 <27028174+Classic298@users.noreply.github.com> Date: Fri, 15 May 2026 22:30:59 +0200 Subject: [PATCH] fix: enforce graph invariants on message upsert, stop clobbering existing nodes Addresses the automated review on the lost-placeholder hardening: - chats.py: enforce structural invariants (id, parentId key present, childrenIds list, role, timestamp) AFTER merge, so dict-spread order can no longer let a partial payload persist a malformed node. - chats.py: a partial update no longer degrades structural fields an existing node already has; an absent/empty childrenIds no longer wipes real children. - chats.py: history.setdefault('messages', {}) so a missing messages map does not raise; warn when synthesizing a missing node so unexpected callers are visible. - middleware.py: build_assistant_message_update carries only placeholder-repair fields; drop childrenIds/timestamp (these clobbered existing nodes via the merge), require metadata['message_id'], and set parentId only when known so a missing user_message_id cannot null out a valid parent. Co-Authored-By: Claude Opus 4.7 (1M context) --- backend/open_webui/models/chats.py | 64 +++++++++++++++++++------- backend/open_webui/utils/middleware.py | 16 +++++-- 2 files changed, 58 insertions(+), 22 deletions(-) diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py index 571a3d8378..e7e2b3037e 100644 --- a/backend/open_webui/models/chats.py +++ b/backend/open_webui/models/chats.py @@ -563,24 +563,54 @@ class ChatTable: chat = chat.chat history = chat.get('history', {}) - if message_id in history.get('messages', {}): - history['messages'][message_id] = { - **history['messages'][message_id], - **message, - } + now = int(time.time()) + messages = history.setdefault('messages', {}) + + def enforce_node_invariants(node: dict) -> dict: + # Structural fields are graph invariants, not free-form payload. + # A partial streaming/final update must never be able to leave a + # node without them, regardless of caller or merge order. + node['id'] = message_id + # The key must always exist (explicit null), otherwise the + # frontend's parent walk treats the node as having an unknown + # parent and the conversation gets stuck loading. + node.setdefault('parentId', None) + if not isinstance(node.get('childrenIds'), list): + node['childrenIds'] = [] + if not node.get('role'): + node['role'] = 'assistant' + if not node.get('timestamp'): + node['timestamp'] = now + return node + + existing = messages.get(message_id) + if existing is not None: + merged = {**existing, **message} + # A partial update (done/content/output/usage) must not degrade + # structural fields the node already has: an empty or null + # incoming value never wins over a valid existing one. + for key in ('parentId', 'role', 'timestamp'): + if not merged.get(key) and existing.get(key): + merged[key] = existing[key] + # Only let childrenIds change when the caller actually sends one; + # an absent/empty incoming list must not wipe real children. + if not message.get('childrenIds') and existing.get('childrenIds'): + merged['childrenIds'] = existing['childrenIds'] + messages[message_id] = enforce_node_invariants(merged) else: - now = int(time.time()) - # This upsert is also used for partial streaming/final updates. - # If a concurrent whole-chat write dropped the assistant placeholder, - # never persist the partial payload as a malformed history node. - history['messages'][message_id] = { - 'id': message_id, - 'parentId': message.get('parentId'), - 'childrenIds': message.get('childrenIds', []), - 'role': message.get('role', 'assistant'), - 'timestamp': message.get('timestamp', now), - **message, - } + # The target node does not exist. This upsert is also used for + # partial streaming/final updates that assume an assistant + # placeholder already exists; if a concurrent whole-chat write + # dropped it, synthesize a structurally valid node instead of + # persisting the partial payload as a malformed history node. + incoming_role = message.get('role') + log.warning( + 'upsert_message_to_chat_by_id_and_message_id: creating missing ' + f'message node {message_id} from a partial payload ' + f'(role={incoming_role!r}); an assistant placeholder was ' + 'likely lost by a concurrent chat write' + ) + messages[message_id] = enforce_node_invariants({**message}) history['currentId'] = message_id diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 617f6983ec..a0e346309e 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -3583,15 +3583,21 @@ async def streaming_chat_response_handler(response, ctx): model_id = form_data.get('model', '') def build_assistant_message_update(**fields): - return { - 'id': metadata.get('message_id'), - 'parentId': metadata.get('user_message_id'), - 'childrenIds': [], + # Only carry the fields that repair a lost assistant placeholder. + # childrenIds and timestamp are deliberately omitted: forcing them + # here would clobber a real children list or the original creation + # time on an already-existing node through the upsert merge. + # chats.py synthesizes those only when the node is actually missing. + update = { + 'id': metadata['message_id'], 'role': 'assistant', 'model': model_id, - 'timestamp': int(time.time()), **fields, } + parent_id = metadata.get('user_message_id') + if parent_id: + update['parentId'] = parent_id + return update # Handle as a background task async def response_handler(response, events):