refac

backend/open_webui/migrations/versions/9a1b2c3d4e5f_add_current_message_id_to_chat.py

@@ -0,0 +1,139 @@
+"""add current_message_id to chat
+
+Revision ID: 9a1b2c3d4e5f
+Revises: 856c5b02fb54
+Create Date: 2026-07-23 00:00:00.000000
+
+"""
+
+import json
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+revision: str = '9a1b2c3d4e5f'
+down_revision: Union[str, None] = '856c5b02fb54'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    columns = [col['name'] for col in sa.inspect(conn).get_columns('chat')]
+    if 'current_message_id' not in columns:
+        op.add_column('chat', sa.Column('current_message_id', sa.Text(), nullable=True))
+
+    chat = sa.table(
+        'chat',
+        sa.column('id', sa.String()),
+        sa.column('chat', sa.Text()),
+        sa.column('current_message_id', sa.Text()),
+    )
+    chat_message = sa.table(
+        'chat_message',
+        sa.column('id', sa.Text()),
+        sa.column('chat_id', sa.Text()),
+        sa.column('parent_id', sa.Text()),
+        sa.column('created_at', sa.BigInteger()),
+    )
+
+    messages_by_chat: dict[str, dict[str, dict]] = {}
+    for row in conn.execute(
+        sa.select(
+            chat_message.c.chat_id,
+            chat_message.c.id,
+            chat_message.c.parent_id,
+            chat_message.c.created_at,
+        )
+    ):
+        values = row._mapping
+        chat_id = values['chat_id']
+        prefix = f'{chat_id}-'
+        message_id = values['id']
+        if message_id and message_id.startswith(prefix):
+            message_id = message_id[len(prefix) :]
+        if not message_id:
+            continue
+        parent_id = values['parent_id']
+        if parent_id and parent_id.startswith(prefix):
+            parent_id = parent_id[len(prefix) :]
+        messages_by_chat.setdefault(chat_id, {})[message_id] = {
+            'parent_id': parent_id,
+            'created_at': values['created_at'] or 0,
+        }
+
+    for row in conn.execute(sa.select(chat.c.id, chat.c.chat, chat.c.current_message_id)):
+        values = row._mapping
+        chat_id = values['id']
+        prefix = f'{chat_id}-'
+        chat_messages = messages_by_chat.get(chat_id, {})
+
+        chat_data = {}
+        if isinstance(values['chat'], dict):
+            chat_data = values['chat']
+        elif isinstance(values['chat'], str):
+            try:
+                parsed = json.loads(values['chat'])
+                chat_data = parsed if isinstance(parsed, dict) else {}
+            except (TypeError, ValueError, json.JSONDecodeError):
+                chat_data = {}
+
+        history = chat_data.get('history') if isinstance(chat_data.get('history'), dict) else {}
+        messages = history.get('messages') if isinstance(history.get('messages'), dict) else {}
+        if not messages and isinstance(chat_data.get('messages'), list):
+            messages = {
+                message['id']: message
+                for message in chat_data['messages']
+                if isinstance(message, dict) and message.get('id')
+            }
+        json_messages = {
+            message_id: {
+                'parent_id': message.get('parentId') if isinstance(message, dict) else None,
+                'created_at': message.get('timestamp', 0) if isinstance(message, dict) else 0,
+            }
+            for message_id, message in messages.items()
+        }
+        available_messages = chat_messages or json_messages
+        candidates = [
+            values['current_message_id'],
+            history.get('currentId'),
+            chat_data.get('currentId'),
+            chat_data.get('branchPointMessageId'),
+        ]
+
+        current_message_id = next(
+            (
+                candidate[len(prefix) :] if candidate.startswith(prefix) else candidate
+                for candidate in candidates
+                if candidate
+                and (
+                    not available_messages
+                    or (candidate[len(prefix) :] if candidate.startswith(prefix) else candidate) in available_messages
+                )
+            ),
+            None,
+        )
+
+        if not current_message_id and available_messages:
+            parent_ids = {
+                message['parent_id']
+                for message in available_messages.values()
+                if message.get('parent_id') in available_messages
+            }
+            leaf_ids = [message_id for message_id in available_messages if message_id not in parent_ids]
+            current_message_id = max(
+                leaf_ids or list(available_messages),
+                key=lambda message_id: available_messages[message_id].get('created_at') or 0,
+            )
+
+        if current_message_id and current_message_id != values['current_message_id']:
+            conn.execute(
+                sa.update(chat)
+                .where(chat.c.id == chat_id)
+                .values(current_message_id=current_message_id)
+            )
+
+
+def downgrade() -> None:
+    op.drop_column('chat', 'current_message_id')

backend/open_webui/models/chats.py

@@ -62,6 +62,7 @@ class Chat(Base):  # database table mapping for chat entity
 
     tasks = Column(JSON, nullable=True)
     summary = Column(Text, nullable=True)
+    current_message_id = Column(Text, nullable=True)
 
     last_read_at = Column(BigInteger, nullable=True)
 
@@ -98,6 +99,7 @@ class ChatModel(BaseModel):
 
     tasks: list | None = None
     summary: str | None = None
+    current_message_id: str | None = None
 
     last_read_at: int | None = None
 
@@ -145,6 +147,7 @@ class ChatForm(BaseModel):
 class ChatImportForm(ChatForm):
     meta: dict | None = {}
     pinned: bool | None = False
+    current_message_id: str | None = None
     created_at: int | None = None
     updated_at: int | None = None
 
@@ -177,6 +180,7 @@ class ChatResponse(BaseModel):
 
     tasks: list | None = None
     summary: str | None = None
+    current_message_id: str | None = None
     context_usage: dict | None = None
 
 
@@ -280,6 +284,21 @@ class ChatTable:
         """Recursively remove null bytes from strings in dict/list structures."""
         return sanitize_data_for_db(obj)
 
+    def get_current_message_id(self, chat: dict | None) -> str | None:
+        chat = chat or {}
+        history = chat.get('history') if isinstance(chat.get('history'), dict) else {}
+        current_id = history.get('currentId') or chat.get('currentId') or chat.get('branchPointMessageId')
+        if current_id:
+            return current_id
+
+        messages = chat.get('messages')
+        if isinstance(messages, list):
+            for message in reversed(messages):
+                if isinstance(message, dict) and message.get('id'):
+                    return message['id']
+
+        return None
+
     def _sanitize_chat_row(self, chat_item):
         """
         Clean a Chat SQLAlchemy model's title + chat JSON,
@@ -375,6 +394,7 @@ class ChatTable:
                     'chat': self._clean_null_bytes(form_data.chat),
                     'folder_id': form_data.folder_id,
                     'meta': internal_meta or {},
+                    'current_message_id': self.get_current_message_id(form_data.chat),
                     'created_at': int(time.time()),
                     'updated_at': int(time.time()),
                     'last_read_at': int(time.time()),
@@ -388,8 +408,14 @@ class ChatTable:
 
             # Dual-write initial messages to chat_message table
             try:
-                history = form_data.chat.get('history', {})
+                history = form_data.chat.get('history') if isinstance(form_data.chat.get('history'), dict) else {}
-                messages = history.get('messages', {})
+                messages = history.get('messages') if isinstance(history.get('messages'), dict) else {}
+                if not messages and isinstance(form_data.chat.get('messages'), list):
+                    messages = {
+                        message.get('id'): message
+                        for message in form_data.chat['messages']
+                        if isinstance(message, dict) and message.get('id')
+                    }
                 for message_id, message in messages.items():
                     if isinstance(message, dict) and message.get('role'):
                         await ChatMessages.upsert_message(
@@ -458,6 +484,7 @@ class ChatTable:
                 'meta': form_data.meta,
                 'pinned': form_data.pinned,
                 'folder_id': form_data.folder_id,
+                'current_message_id': form_data.current_message_id or self.get_current_message_id(form_data.chat),
                 'created_at': (form_data.created_at if form_data.created_at else int(time.time())),
                 'updated_at': (form_data.updated_at if form_data.updated_at else int(time.time())),
             }
@@ -497,8 +524,14 @@ class ChatTable:
 
             # Dual-write messages to chat_message table
             for form_data, chat_obj in zip(chat_import_forms, chats):
-                history = form_data.chat.get('history', {})
+                history = form_data.chat.get('history') if isinstance(form_data.chat.get('history'), dict) else {}
-                messages = history.get('messages', {})
+                messages = history.get('messages') if isinstance(history.get('messages'), dict) else {}
+                if not messages and isinstance(form_data.chat.get('messages'), list):
+                    messages = {
+                        message.get('id'): message
+                        for message in form_data.chat['messages']
+                        if isinstance(message, dict) and message.get('id')
+                    }
                 for message_id, message in messages.items():
                     if isinstance(message, dict) and message.get('role'):
                         try:
@@ -530,6 +563,8 @@ class ChatTable:
 
                 chat_item.chat = self._clean_null_bytes(chat)
                 chat_item.title = self._clean_null_bytes(chat['title']) if 'title' in chat else 'New Chat'
+                if any(key in chat for key in ('history', 'messages', 'currentId', 'branchPointMessageId')):
+                    chat_item.current_message_id = self.get_current_message_id(chat)
 
                 if touch:
                     chat_item.updated_at = int(time.time())

backend/open_webui/routers/chats.py

@@ -38,6 +38,7 @@ from open_webui.tasks import has_active_tasks, stop_item_tasks
 from open_webui.utils.access_control import filter_allowed_access_grants, has_permission
 from open_webui.utils.access_control.folders import has_folder_access
 from open_webui.utils.auth import get_admin_user, get_verified_user
+from open_webui.utils.chat_fork import build_fork_history
 from open_webui.utils.context_compaction import compact_chat_branch, get_chat_context_usage
 from open_webui.utils.misc import get_message_list
 from open_webui.utils.models import get_all_models
@@ -1167,7 +1168,8 @@ async def compact_chat_by_id(
 
     history = (chat.chat or {}).get('history') or {}
     messages_map = await Chats.get_messages_map_by_chat_id(id)
-    message_list = get_message_list(messages_map or history.get('messages') or {}, history.get('currentId'))
+    current_message_id = chat.current_message_id or history.get('currentId')
+    message_list = get_message_list(messages_map or history.get('messages') or {}, current_message_id)
     model_id = (form_data.model if form_data else None) or next(
         (message.get('model') for message in reversed(message_list) if message.get('model')),
         None,
@@ -1564,6 +1566,96 @@ class CloneForm(BaseModel):
     title: str | None = None
 
 
+class ForkForm(BaseModel):
+    message_id: str | None = None
+
+
+@router.post('/{id}/fork', response_model=ChatResponse | None)
+async def fork_chat_by_id(
+    request: Request,
+    id: str,
+    form_data: ForkForm | None = None,
+    user=Depends(get_verified_user),
+    db: AsyncSession = Depends(get_async_session),
+):
+    await require_chat_import_permission(request, user, db)
+
+    chat = await Chats.get_chat_by_id_and_user_id(id, user.id, db=db)
+    if not chat:
+        raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.DEFAULT())
+
+    if await has_active_tasks(request.app.state.redis, id):
+        raise HTTPException(
+            status_code=status.HTTP_409_CONFLICT,
+            detail='Wait for the current response to finish before forking.',
+        )
+
+    history = (chat.chat or {}).get('history') or {}
+    messages_map = await Chats.get_messages_map_by_chat_id(id) or history.get('messages') or {}
+    if any(
+        message.get('role') == 'assistant' and message.get('done') is False
+        for message in messages_map.values()
+        if isinstance(message, dict)
+    ):
+        raise HTTPException(
+            status_code=status.HTTP_409_CONFLICT,
+            detail='Wait for the current response to finish before forking.',
+        )
+
+    source_message_id = (form_data.message_id if form_data else None) or chat.current_message_id or history.get('currentId')
+    if not source_message_id:
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='chat has no messages to fork')
+
+    try:
+        fork_history, fork_messages = build_fork_history(messages_map, source_message_id)
+    except ValueError as exc:
+        detail = str(exc)
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND if detail == 'message not found' else status.HTTP_400_BAD_REQUEST,
+            detail=detail,
+        ) from exc
+
+    updated_chat = {**(chat.chat or {})}
+    updated_chat.pop('currentId', None)
+    updated_chat.update(
+        {
+            'originalChatId': chat.id,
+            'branchPointMessageId': source_message_id,
+            'title': f'{chat.title} (fork)',
+            'history': fork_history,
+            'messages': fork_messages,
+        }
+    )
+    meta = {
+        **(chat.meta or {}),
+        'forked_from': chat.id,
+        'forked_from_message_id': source_message_id,
+    }
+
+    fork = await Chats.insert_new_chat(
+        str(uuid4()),
+        user.id,
+        ChatForm(chat=updated_chat, folder_id=chat.folder_id),
+        db=db,
+        internal_meta=meta,
+    )
+
+    if not fork:
+        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=ERROR_MESSAGES.DEFAULT())
+
+    if chat.pinned:
+        fork = await Chats.toggle_chat_pinned_by_id(fork.id, db=db) or fork
+
+    await publish_event(
+        request,
+        EVENTS.CHAT_CLONED,
+        actor=user,
+        subject_id=fork.id,
+        data={'original_chat_id': id, 'forked_from_message_id': source_message_id},
+    )
+    return ChatResponse(**fork.model_dump())
+
+
 @router.post('/{id}/clone', response_model=ChatResponse | None)
 async def clone_chat_by_id(
     request: Request,

backend/open_webui/utils/chat_fork.py

@@ -0,0 +1,41 @@
+from copy import deepcopy
+
+
+def build_fork_history(messages_map: dict, source_message_id: str) -> tuple[dict, list[dict]]:
+    if not messages_map:
+        raise ValueError('chat has no messages to fork')
+
+    branch: list[tuple[str, dict]] = []
+    seen: set[str] = set()
+    message_id = source_message_id
+
+    while message_id:
+        if message_id in seen:
+            raise ValueError('message branch contains a cycle')
+        seen.add(message_id)
+
+        message = messages_map.get(message_id)
+        if not isinstance(message, dict):
+            raise ValueError('message not found')
+
+        branch.append((message_id, message))
+        message_id = message.get('parentId')
+
+    fork_messages: dict[str, dict] = {}
+    ordered_messages: list[dict] = []
+    parent_id = None
+
+    for message_id, message in reversed(branch):
+        copied = deepcopy(message)
+        copied['id'] = message_id
+        copied['parentId'] = parent_id
+        copied['childrenIds'] = []
+
+        if parent_id:
+            fork_messages[parent_id]['childrenIds'] = [message_id]
+
+        fork_messages[message_id] = copied
+        ordered_messages.append(copied)
+        parent_id = message_id
+
+    return {'messages': fork_messages, 'currentId': source_message_id}, ordered_messages

backend/open_webui/utils/context_compaction.py

@@ -148,8 +148,13 @@ async def compact_chat_branch(request, user, chat: Any, model_id: str, models: d
     if not config['enable']:
         return {'ok': True, 'compacted': False, 'reason': 'disabled'}
 
-    history = (chat.chat or {}).get('history') or {}
+    chat_data = chat.chat or {}
-    current_id = history.get('currentId')
+    history = chat_data.get('history') or {}
+    current_id = getattr(chat, 'current_message_id', None) or history.get('currentId')
+    if not current_id:
+        current_id = chat_data.get('currentId') or chat_data.get('branchPointMessageId')
+    if not current_id and isinstance(chat_data.get('messages'), list) and chat_data['messages']:
+        current_id = chat_data['messages'][-1].get('id')
     if not current_id:
         return {'ok': True, 'compacted': False, 'reason': 'empty'}
 
@@ -216,8 +221,13 @@ def _resolve_token_threshold(global_threshold: int, global_cap: int, metadata: d
 
 
 async def get_chat_context_usage(chat: Any, model_id: str | None = None) -> dict | None:
-    history = (chat.chat or {}).get('history') or {}
+    chat_data = chat.chat or {}
-    current_id = history.get('currentId')
+    history = chat_data.get('history') or {}
+    current_id = getattr(chat, 'current_message_id', None) or history.get('currentId')
+    if not current_id:
+        current_id = chat_data.get('currentId') or chat_data.get('branchPointMessageId')
+    if not current_id and isinstance(chat_data.get('messages'), list) and chat_data['messages']:
+        current_id = chat_data['messages'][-1].get('id')
     if not current_id:
         return None