fix: assistant message editing and continuation

Author: Algorithm5838

backend/open_webui/routers/chats.py

@@ -6,7 +6,7 @@
 from fastapi.responses import StreamingResponse
 
 
-from open_webui.utils.misc import get_message_list
+from open_webui.utils.misc import get_message_list, reconcile_output_with_content
 from open_webui.socket.main import get_event_emitter
 from open_webui.models.chats import (
     ChatForm,
@@ -888,6 +888,15 @@ async def update_chat_by_id(
     chat = Chats.get_chat_by_id_and_user_id(id, user.id, db=db)
     if chat:
         updated_chat = {**chat.chat, **form_data.chat}
+
+        # Reconcile output with content for any edited assistant messages so
+        # the DB is always consistent and process_messages_with_output can use
+        # output as the source of truth without re-parsing content at read time.
+        messages = updated_chat.get('history', {}).get('messages', {})
+        for mid, msg in messages.items():
+            if msg.get('role') == 'assistant' and msg.get('output'):
+                messages[mid] = reconcile_output_with_content(msg)
+
         chat = Chats.update_chat_by_id(id, updated_chat, db=db)
         return ChatResponse(**chat.model_dump())
     else:

backend/open_webui/utils/middleware.py

@@ -481,9 +481,9 @@ def serialize_output(output: list) -> str:
             )
 
             if status == 'completed' or duration is not None or not is_last_item:
-                content = f'{content}<details type="reasoning" done="true" duration="{duration or 0}">\n<summary>Thought for {duration or 0} seconds</summary>\n{display}\n</details>\n'
+                content = f'{content}<details type="reasoning" done="true" id="{item.get("id", "")}" duration="{duration or 0}">\n<summary>Thought for {duration or 0} seconds</summary>\n{display}\n</details>\n'
             else:
-                content = f'{content}<details type="reasoning" done="false">\n<summary>Thinking…</summary>\n{display}\n</details>\n'
+                content = f'{content}<details type="reasoning" done="false" id="{item.get("id", "")}">\n<summary>Thinking…</summary>\n{display}\n</details>\n'
 
         elif item_type == 'open_webui:code_interpreter':
             content_stripped, original_whitespace = split_content_and_whitespace(content)
@@ -519,9 +519,9 @@ def serialize_output(output: list) -> str:
                 output_attr = f' output="{html.escape(output_json)}"'
 
             if status == 'completed' or duration is not None or not is_last_item:
-                content += f'<details type="code_interpreter" done="true" duration="{duration or 0}"{output_attr}>\n<summary>Analyzed</summary>\n{display}\n</details>\n'
+                content += f'<details type="code_interpreter" done="true" id="{item.get("id", "")}" duration="{duration or 0}"{output_attr}>\n<summary>Analyzed</summary>\n{display}\n</details>\n'
             else:
-                content += f'<details type="code_interpreter" done="false"{output_attr}>\n<summary>Analyzing…</summary>\n{display}\n</details>\n'
+                content += f'<details type="code_interpreter" done="false" id="{item.get("id", "")}"{output_attr}>\n<summary>Analyzing…</summary>\n{display}\n</details>\n'
 
     return content.strip()
 
@@ -2066,15 +2066,23 @@ async def convert_url_images_to_base64(form_data):
     return form_data
 
 
-def load_messages_from_db(chat_id: str, message_id: str) -> Optional[list[dict]]:
+def load_messages_from_db(
+    chat_id: str, message_id: str, require_done: bool = False
+) -> Optional[list[dict]]:
     """
     Load the message chain from DB up to message_id,
     keeping only LLM-relevant fields (role, content, output).
+
+    If require_done is True, returns None when the anchor message is not
+    marked done — prevents loading a partial mid-stream output as context.
     """
     messages_map = Chats.get_messages_map_by_chat_id(chat_id)
     if not messages_map:
         return None
 
+    if require_done and not messages_map.get(message_id, {}).get('done'):
+        return None
+
     db_messages = get_message_list(messages_map, message_id)
     if not db_messages:
         return None
@@ -2088,6 +2096,8 @@ def process_messages_with_output(messages: list[dict]) -> list[dict]:
 
     For assistant messages with 'output' field, produces properly formatted
     OpenAI-style messages (tool_calls + tool results). Strips 'output' before LLM.
+    The DB is kept consistent at write-time (update_chat_by_id), so output can
+    be used directly as the source of truth here.
     """
     processed = []
 
@@ -2146,37 +2156,47 @@ async def process_chat_payload(request, form_data, user, metadata, model):
     # Load messages from DB when available — DB preserves structured 'output' items
     # which the frontend strips, causing tool calls to be merged into content.
     chat_id = metadata.get('chat_id')
+    message_id = metadata.get('message_id')
     parent_message_id = metadata.get('parent_message_id')
-
-    if chat_id and parent_message_id and not chat_id.startswith('local:'):
-        db_messages = load_messages_from_db(chat_id, parent_message_id)
-        if db_messages:
-            system_message = get_system_message(form_data.get('messages', []))
-            form_data['messages'] = [system_message, *db_messages] if system_message else db_messages
-
-            # Inject image files into content as image_url parts (mirrors frontend logic)
-            for message in form_data['messages']:
-                image_files = [
-                    f
-                    for f in message.get('files', [])
-                    if f.get('type') == 'image' or (f.get('content_type') or '').startswith('image/')
-                ]
-                if message.get('role') == 'user' and image_files:
-                    text_content = message.get('content', '')
-                    if isinstance(text_content, str):
-                        message['content'] = [
-                            {'type': 'text', 'text': text_content},
-                            *[
-                                {
-                                    'type': 'image_url',
-                                    'image_url': {'url': f['url']},
-                                }
-                                for f in image_files
-                                if f.get('url')
-                            ],
-                        ]
-                # Strip files field — it's been incorporated into content
-                message.pop('files', None)
+    db_messages = None
+
+    if chat_id and not chat_id.startswith('local:'):
+        # For "Continue", message_id is the existing assistant message already in the
+        # DB — load from it so the reconciled output is included in context.
+        # For a new generation, message_id is a fresh UUID not yet saved, so
+        # load_messages_from_db returns None and we fall through to parent_message_id.
+        if message_id:
+            db_messages = load_messages_from_db(chat_id, message_id, require_done=True)
+        if db_messages is None and parent_message_id:
+            db_messages = load_messages_from_db(chat_id, parent_message_id)
+
+    if db_messages:
+        system_message = get_system_message(form_data.get('messages', []))
+        form_data['messages'] = [system_message, *db_messages] if system_message else db_messages
+
+        # Inject image files into content as image_url parts (mirrors frontend logic)
+        for message in form_data['messages']:
+            image_files = [
+                f
+                for f in message.get('files', [])
+                if f.get('type') == 'image' or (f.get('content_type') or '').startswith('image/')
+            ]
+            if message.get('role') == 'user' and image_files:
+                text_content = message.get('content', '')
+                if isinstance(text_content, str):
+                    message['content'] = [
+                        {'type': 'text', 'text': text_content},
+                        *[
+                            {
+                                'type': 'image_url',
+                                'image_url': {'url': f['url']},
+                            }
+                            for f in image_files
+                            if f.get('url')
+                        ],
+                    ]
+            # Strip files field — it's been incorporated into content
+            message.pop('files', None)
 
     # Process messages with OR-aligned output items for clean LLM messages
     form_data['messages'] = process_messages_with_output(form_data.get('messages', []))

backend/open_webui/utils/misc.py

@@ -276,6 +276,85 @@ def flush_pending():
     return messages
 
 
+def filter_output_by_content(output: list, content: str) -> list:
+    """
+    Drop output items whose <details> block was removed from content.
+    Matches by id attribute. Falls back to keeping all items for types
+    that don't yet carry id= (legacy content compatibility).
+    """
+    present_ids = set(re.findall(r'<details\b[^>]*\bid="([^"]+)"', content))
+    types_present = set(re.findall(r'<details\b[^>]*\btype="([^"]+)"', content))
+
+    # Types that carry id= in content (vs legacy types without id=).
+    types_with_ids: set[str] = set()
+    for m in re.finditer(r'<details\b(?=[^>]*\btype="([^"]+)")(?=[^>]*\bid=)', content):
+        types_with_ids.add(m.group(1))
+
+    # Map output item type → <details> type attribute value
+    DETAILS_TYPE = {
+        'function_call': 'tool_calls',
+        'function_call_output': 'tool_calls',
+        'reasoning': 'reasoning',
+        'open_webui:code_interpreter': 'code_interpreter',
+    }
+
+    filtered = []
+    for item in output:
+        details_type = DETAILS_TYPE.get(item.get('type', ''))
+        if details_type is None:
+            filtered.append(item)
+        elif details_type not in types_present:
+            pass  # type removed from content
+        elif details_type not in types_with_ids:
+            filtered.append(item)
+        else:
+            item_id = item.get('call_id') or item.get('id', '')
+            if item_id in present_ids:
+                filtered.append(item)
+            # else: user removed this block
+
+    return filtered
+
+
+def reconcile_output_with_content(message: dict) -> dict:
+    """Align the 'output' list of an assistant message with a user-edited 'content'.
+
+    Called at write-time (update_chat_by_id) so the DB always holds a consistent
+    output list. process_messages_with_output then needs no special-casing.
+    """
+    content = message.get('content', '')
+    output = message.get('output', [])
+
+    # Extract plain text by stripping <details> blocks that serialize_output wrote
+    edited_text = re.sub(r'<details\b[^>]*>.*?</details>', '', content, flags=re.S).strip()
+
+    # Drop output items for <details> blocks the user removed from content
+    filtered = filter_output_by_content(output, content)
+
+    # Update the 'message' output item's text to match the edit
+    reconciled_output = []
+    text_applied = False
+    for item in filtered:
+        if item.get('type') == 'message' and not text_applied:
+            reconciled_output.append({
+                **item,
+                'content': [{'type': 'output_text', 'text': edited_text}],
+            })
+            text_applied = True
+        else:
+            reconciled_output.append(item)
+
+    # If there was no 'message' item, append one so the text is not lost
+    if not text_applied and edited_text:
+        reconciled_output.append({
+            'type': 'message',
+            'role': 'assistant',
+            'content': [{'type': 'output_text', 'text': edited_text}],
+        })
+
+    return {**message, 'output': reconciled_output}
+
+
 def get_last_user_message(messages: list[dict]) -> Optional[str]:
     message = get_last_user_message_item(messages)
     if message is None: