From 15883e5229ca5c07d85ebe312dbbdd469d8a83a2 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Wed, 1 Apr 2026 06:00:53 -0500 Subject: [PATCH] refac --- backend/open_webui/models/chat_messages.py | 5 ++++- backend/open_webui/utils/middleware.py | 1 + 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/backend/open_webui/models/chat_messages.py b/backend/open_webui/models/chat_messages.py index 97490c1602..ac75fcf973 100644 --- a/backend/open_webui/models/chat_messages.py +++ b/backend/open_webui/models/chat_messages.py @@ -169,7 +169,10 @@ class ChatMessageTable: info = data.get('info', {}) usage = info.get('usage') if info else None if usage: - existing.usage = usage + # Deep-merge: preserve existing keys not present in new data + # This prevents background tasks (follow-ups, title, tags) + # from accidentally clearing the primary response's token counts + existing.usage = {**(existing.usage or {}), **usage} existing.updated_at = now db.commit() db.refresh(existing) diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 05e18e555f..fc8c2fc909 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -4626,6 +4626,7 @@ async def streaming_chat_response_handler(response, ctx): 'content': serialize_output(output), 'output': output, 'title': title, + **({'usage': usage} if usage else {}), } if not ENABLE_REALTIME_CHAT_SAVE: