@lobehub/lobehub 2.0.0-next.35 → 2.0.0-next.36

package/src/store/chat/slices/message/actions/optimisticUpdate.ts

@@ -0,0 +1,260 @@
+import {
+  ChatErrorType,
+  ChatImageItem,
+  ChatMessageError,
+  ChatMessagePluginError,
+  CreateMessageParams,
+  GroundingSearch,
+  MessageMetadata,
+  MessageToolCall,
+  ModelReasoning,
+  UIChatMessage,
+  UpdateMessageRAGParams,
+} from '@lobechat/types';
+import { nanoid } from '@lobechat/utils';
+import { StateCreator } from 'zustand/vanilla';
+
+import { messageService } from '@/services/message';
+import { ChatStore } from '@/store/chat/store';
+
+/**
+ * Optimistic update operations
+ * All methods follow the pattern: update frontend first, then persist to database
+ */
+export interface MessageOptimisticUpdateAction {
+  /**
+   * create a message with optimistic update
+   * returns the created message ID and updated message list
+   */
+  optimisticCreateMessage: (
+    params: CreateMessageParams,
+    context?: { groupMessageId?: string; skipRefresh?: boolean; tempMessageId?: string },
+  ) => Promise<{ id: string; messages: UIChatMessage[] } | undefined>;
+
+  /**
+   * create a temp message for optimistic update
+   * otherwise the message will be too slow to show
+   */
+  optimisticCreateTmpMessage: (params: CreateMessageParams) => string;
+
+  /**
+   * delete the message content with optimistic update
+   */
+  optimisticDeleteMessage: (id: string) => Promise<void>;
+
+  /**
+   * update the message content with optimistic update
+   * a method used by other action
+   */
+  optimisticUpdateMessageContent: (
+    id: string,
+    content: string,
+    extra?: {
+      imageList?: ChatImageItem[];
+      metadata?: MessageMetadata;
+      model?: string;
+      provider?: string;
+      reasoning?: ModelReasoning;
+      search?: GroundingSearch;
+      toolCalls?: MessageToolCall[];
+    },
+  ) => Promise<void>;
+
+  /**
+   * update the message error with optimistic update
+   */
+  optimisticUpdateMessageError: (id: string, error: ChatMessageError | null) => Promise<void>;
+
+  /**
+   * update the message metadata with optimistic update
+   */
+  optimisticUpdateMessageMetadata: (
+    id: string,
+    metadata: Partial<MessageMetadata>,
+  ) => Promise<void>;
+
+  /**
+   * update the message plugin error with optimistic update
+   */
+  optimisticUpdateMessagePluginError: (
+    id: string,
+    error: ChatMessagePluginError | null,
+  ) => Promise<void>;
+
+  /**
+   * update message RAG with optimistic update
+   */
+  optimisticUpdateMessageRAG: (id: string, input: UpdateMessageRAGParams) => Promise<void>;
+}
+
+export const messageOptimisticUpdate: StateCreator<
+  ChatStore,
+  [['zustand/devtools', never]],
+  [],
+  MessageOptimisticUpdateAction
+> = (set, get) => ({
+  optimisticCreateMessage: async (message, context) => {
+    const {
+      optimisticCreateTmpMessage,
+      internal_toggleMessageLoading,
+      internal_dispatchMessage,
+      replaceMessages,
+    } = get();
+
+    let tempId = context?.tempMessageId;
+    if (!tempId) {
+      tempId = optimisticCreateTmpMessage(message as any);
+      internal_toggleMessageLoading(true, tempId);
+    }
+
+    try {
+      const result = await messageService.createMessage(message);
+
+      if (!context?.skipRefresh) {
+        // Use the messages returned from createMessage (already grouped)
+        replaceMessages(result.messages);
+      }
+
+      internal_toggleMessageLoading(false, tempId);
+      return result;
+    } catch (e) {
+      internal_toggleMessageLoading(false, tempId);
+      internal_dispatchMessage({
+        id: tempId,
+        type: 'updateMessage',
+        value: {
+          error: {
+            body: e,
+            message: (e as Error).message,
+            type: ChatErrorType.CreateMessageError,
+          },
+        },
+      });
+    }
+  },
+
+  optimisticCreateTmpMessage: (message) => {
+    const { internal_dispatchMessage } = get();
+
+    // use optimistic update to avoid the slow waiting
+    const tempId = 'tmp_' + nanoid();
+    internal_dispatchMessage({ id: tempId, type: 'createMessage', value: message });
+
+    return tempId;
+  },
+
+  optimisticDeleteMessage: async (id: string) => {
+    get().internal_dispatchMessage({ id, type: 'deleteMessage' });
+    const result = await messageService.removeMessage(id, {
+      sessionId: get().activeId,
+      topicId: get().activeTopicId,
+    });
+    if (result?.success && result.messages) {
+      get().replaceMessages(result.messages);
+    }
+  },
+
+  optimisticUpdateMessageContent: async (id, content, extra) => {
+    const {
+      internal_dispatchMessage,
+      refreshMessages,
+      internal_transformToolCalls,
+      replaceMessages,
+    } = get();
+
+    // Due to the async update method and refresh need about 100ms
+    // we need to update the message content at the frontend to avoid the update flick
+    // refs: https://medium.com/@kyledeguzmanx/what-are-optimistic-updates-483662c3e171
+    if (extra?.toolCalls) {
+      internal_dispatchMessage({
+        id,
+        type: 'updateMessage',
+        value: { tools: internal_transformToolCalls(extra?.toolCalls) },
+      });
+    } else {
+      internal_dispatchMessage({
+        id,
+        type: 'updateMessage',
+        value: { content },
+      });
+    }
+
+    const result = await messageService.updateMessage(
+      id,
+      {
+        content,
+        imageList: extra?.imageList,
+        metadata: extra?.metadata,
+        model: extra?.model,
+        provider: extra?.provider,
+        reasoning: extra?.reasoning,
+        search: extra?.search,
+        tools: extra?.toolCalls ? internal_transformToolCalls(extra?.toolCalls) : undefined,
+      },
+      { sessionId: get().activeId, topicId: get().activeTopicId },
+    );
+
+    if (result && result.success && result.messages) {
+      replaceMessages(result.messages);
+    } else {
+      await refreshMessages();
+    }
+  },
+
+  optimisticUpdateMessageError: async (id, error) => {
+    get().internal_dispatchMessage({ id, type: 'updateMessage', value: { error } });
+    const result = await messageService.updateMessage(
+      id,
+      { error },
+      { sessionId: get().activeId, topicId: get().activeTopicId },
+    );
+    if (result?.success && result.messages) {
+      get().replaceMessages(result.messages);
+    } else {
+      await get().refreshMessages();
+    }
+  },
+
+  optimisticUpdateMessageMetadata: async (id, metadata) => {
+    const { internal_dispatchMessage, refreshMessages, replaceMessages } = get();
+
+    // Optimistic update: update the frontend immediately
+    internal_dispatchMessage({
+      id,
+      type: 'updateMessageMetadata',
+      value: metadata,
+    });
+
+    // Persist to database
+    const result = await messageService.updateMessageMetadata(id, metadata, {
+      sessionId: get().activeId,
+      topicId: get().activeTopicId,
+    });
+
+    if (result?.success && result.messages) {
+      replaceMessages(result.messages);
+    } else {
+      await refreshMessages();
+    }
+  },
+
+  optimisticUpdateMessagePluginError: async (id, error) => {
+    const result = await messageService.updateMessagePluginError(id, error, {
+      sessionId: get().activeId,
+      topicId: get().activeTopicId,
+    });
+    if (result?.success && result.messages) {
+      get().replaceMessages(result.messages);
+    }
+  },
+
+  optimisticUpdateMessageRAG: async (id, data) => {
+    const result = await messageService.updateMessageRAG(id, data, {
+      sessionId: get().activeId,
+      topicId: get().activeTopicId,
+    });
+    if (result?.success && result.messages) {
+      get().replaceMessages(result.messages);
+    }
+  },
+});

package/src/store/chat/slices/message/actions/publicApi.ts

@@ -0,0 +1,224 @@
+/* eslint-disable sort-keys-fix/sort-keys-fix,typescript-sort-keys/interface */
+import { TraceEventType } from '@lobechat/types';
+import { copyToClipboard } from '@lobehub/ui';
+import isEqual from 'fast-deep-equal';
+import { StateCreator } from 'zustand/vanilla';
+
+import { messageService } from '@/services/message';
+import { topicService } from '@/services/topic';
+import { ChatStore } from '@/store/chat/store';
+import { useSessionStore } from '@/store/session';
+import { sessionSelectors } from '@/store/session/selectors';
+import { setNamespace } from '@/utils/storeDebug';
+
+import { dbMessageSelectors, displayMessageSelectors } from '../../../selectors';
+import { toggleBooleanList } from '../../../utils';
+
+const n = setNamespace('m');
+
+/**
+ * Public API for components
+ * These methods are directly called by UI components
+ */
+export interface MessagePublicApiAction {
+  // ===== Create ===== //
+  addAIMessage: () => Promise<void>;
+  addUserMessage: (params: { message: string; fileList?: string[] }) => Promise<void>;
+
+  // ===== Delete ===== //
+  /**
+   * clear message on the active session
+   */
+  clearMessage: () => Promise<void>;
+  deleteMessage: (id: string) => Promise<void>;
+  deleteToolMessage: (id: string) => Promise<void>;
+  clearAllMessages: () => Promise<void>;
+
+  // ===== Update ===== //
+  /**
+   * Update message input box content
+   */
+  updateMessageInput: (message: string) => void;
+  modifyMessageContent: (id: string, content: string) => Promise<void>;
+  toggleMessageEditing: (id: string, editing: boolean) => void;
+
+  // ===== Others ===== //
+  copyMessage: (id: string, content: string) => Promise<void>;
+}
+
+export const messagePublicApi: StateCreator<
+  ChatStore,
+  [['zustand/devtools', never]],
+  [],
+  MessagePublicApiAction
+> = (set, get) => ({
+  addAIMessage: async () => {
+    const { optimisticCreateMessage, updateMessageInput, activeTopicId, activeId, inputMessage } =
+      get();
+    if (!activeId) return;
+
+    const parentId = displayMessageSelectors.lastDisplayMessageId(get());
+
+    const result = await optimisticCreateMessage({
+      content: inputMessage,
+      role: 'assistant',
+      sessionId: activeId,
+      // if there is activeTopicId，then add topicId to message
+      topicId: activeTopicId,
+      parentId,
+    });
+
+    if (result) {
+      updateMessageInput('');
+    }
+  },
+
+  addUserMessage: async ({ message, fileList }) => {
+    const { optimisticCreateMessage, updateMessageInput, activeTopicId, activeId, activeThreadId } =
+      get();
+    if (!activeId) return;
+
+    const parentId = displayMessageSelectors.lastDisplayMessageId(get());
+
+    const result = await optimisticCreateMessage({
+      content: message,
+      files: fileList,
+      role: 'user',
+      sessionId: activeId,
+      // if there is activeTopicId，then add topicId to message
+      topicId: activeTopicId,
+      threadId: activeThreadId,
+      parentId,
+    });
+
+    if (result) {
+      updateMessageInput('');
+    }
+  },
+
+  deleteMessage: async (id) => {
+    const message = displayMessageSelectors.getDisplayMessageById(id)(get());
+    if (!message) return;
+
+    let ids = [message.id];
+    const allMessages = displayMessageSelectors.activeDisplayMessages(get());
+
+    // Handle assistantGroup messages: delete all child blocks and tool results
+    if (message.role === 'assistantGroup' && message.children) {
+      // Collect all child block IDs
+      const childIds = message.children.map((child) => child.id);
+      ids = ids.concat(childIds);
+
+      // Collect all tool result IDs from children
+      const toolResultIds = message.children.flatMap((child) => {
+        if (!child.tools) return [];
+        return child.tools.filter((tool) => tool.result?.id).map((tool) => tool.result!.id);
+      });
+      ids = ids.concat(toolResultIds);
+    }
+    // Handle regular messages with tools: find and delete related tool messages
+    else if (message.tools) {
+      const toolMessageIds = message.tools.flatMap((tool) => {
+        const messages = allMessages.filter((m) => m.tool_call_id === tool.id);
+        return messages.map((m) => m.id);
+      });
+      ids = ids.concat(toolMessageIds);
+    }
+
+    get().internal_dispatchMessage({ type: 'deleteMessages', ids });
+    const result = await messageService.removeMessages(ids, {
+      sessionId: get().activeId,
+      topicId: get().activeTopicId,
+    });
+    if (result?.success && result.messages) {
+      get().replaceMessages(result.messages);
+    }
+  },
+
+  deleteToolMessage: async (id) => {
+    const message = dbMessageSelectors.getDbMessageById(id)(get());
+    if (!message || message.role !== 'tool') return;
+
+    const removeToolInAssistantMessage = async () => {
+      if (!message.parentId) return;
+      await get().internal_removeToolToAssistantMessage(message.parentId, message.tool_call_id);
+    };
+
+    await Promise.all([
+      // 1. remove tool message
+      get().optimisticDeleteMessage(id),
+      // 2. remove the tool item in the assistant tools
+      removeToolInAssistantMessage(),
+    ]);
+  },
+
+  clearMessage: async () => {
+    const { activeId, activeTopicId, refreshTopic, switchTopic, activeSessionType } = get();
+
+    // Check if this is a group session - use activeSessionType if available, otherwise check session store
+    let isGroupSession = activeSessionType === 'group';
+    if (activeSessionType === undefined) {
+      // Fallback: check session store directly
+      const sessionStore = useSessionStore.getState();
+      isGroupSession = sessionSelectors.isCurrentSessionGroupSession(sessionStore);
+    }
+
+    // For group sessions, we need to clear group messages using groupId
+    // For regular sessions, we clear session messages using sessionId
+    if (isGroupSession) {
+      // For group chat, activeId is the groupId
+      await messageService.removeMessagesByGroup(activeId, activeTopicId);
+    } else {
+      // For regular session, activeId is the sessionId
+      await messageService.removeMessagesByAssistant(activeId, activeTopicId);
+    }
+
+    if (activeTopicId) {
+      await topicService.removeTopic(activeTopicId);
+    }
+    await refreshTopic();
+
+    // Clear messages directly since all messages are deleted
+    get().replaceMessages([]);
+
+    // after remove topic , go back to default topic
+    switchTopic();
+  },
+
+  clearAllMessages: async () => {
+    await messageService.removeAllMessages();
+    // Clear messages directly since all messages are deleted
+    get().replaceMessages([]);
+  },
+
+  copyMessage: async (id, content) => {
+    await copyToClipboard(content);
+
+    get().internal_traceMessage(id, { eventType: TraceEventType.CopyMessage });
+  },
+
+  toggleMessageEditing: (id, editing) => {
+    set(
+      { messageEditingIds: toggleBooleanList(get().messageEditingIds, id, editing) },
+      false,
+      'toggleMessageEditing',
+    );
+  },
+
+  updateMessageInput: (message) => {
+    if (isEqual(message, get().inputMessage)) return;
+
+    set({ inputMessage: message }, false, n('updateMessageInput', message));
+  },
+
+  modifyMessageContent: async (id, content) => {
+    // tracing the diff of update
+    // due to message content will change, so we need send trace before update,or will get wrong data
+    get().internal_traceMessage(id, {
+      eventType: TraceEventType.ModifyMessage,
+      nextContent: content,
+    });
+
+    await get().optimisticUpdateMessageContent(id, content);
+  },
+});

package/src/store/chat/slices/message/actions/query.ts

@@ -0,0 +1,120 @@
+import { parse } from '@lobechat/conversation-flow';
+import { UIChatMessage } from '@lobechat/types';
+import isEqual from 'fast-deep-equal';
+import { SWRResponse, mutate } from 'swr';
+import { StateCreator } from 'zustand/vanilla';
+
+import { useClientDataSWR } from '@/libs/swr';
+import { messageService } from '@/services/message';
+import { ChatStore } from '@/store/chat/store';
+import { setNamespace } from '@/utils/storeDebug';
+
+import { messageMapKey } from '../../../utils/messageMapKey';
+
+const n = setNamespace('m');
+
+const SWR_USE_FETCH_MESSAGES = 'SWR_USE_FETCH_MESSAGES';
+
+/**
+ * Data query and synchronization actions
+ * Handles fetching, refreshing, and replacing message data
+ */
+export interface MessageQueryAction {
+  /**
+   * Manually refresh messages from server
+   */
+  refreshMessages: () => Promise<void>;
+
+  /**
+   * Replace current messages with new data
+   */
+  replaceMessages: (
+    messages: UIChatMessage[],
+    params?: {
+      action?: any;
+      sessionId?: string;
+      topicId?: string;
+    },
+  ) => void;
+
+  /**
+   * Fetch messages using SWR
+   * @param enable - whether to enable the fetch
+   * @param messageContextId - Can be sessionId or groupId
+   */
+  useFetchMessages: (
+    enable: boolean,
+    messageContextId: string,
+    activeTopicId?: string,
+    type?: 'session' | 'group',
+  ) => SWRResponse<UIChatMessage[]>;
+}
+
+export const messageQuery: StateCreator<
+  ChatStore,
+  [['zustand/devtools', never]],
+  [],
+  MessageQueryAction
+> = (set, get) => ({
+  // TODO: The mutate should only be called once, but since we haven't merge session and group,
+  // we need to call it twice
+  refreshMessages: async () => {
+    await mutate([SWR_USE_FETCH_MESSAGES, get().activeId, get().activeTopicId, 'session']);
+    await mutate([SWR_USE_FETCH_MESSAGES, get().activeId, get().activeTopicId, 'group']);
+  },
+
+  replaceMessages: (messages, params) => {
+    const messagesKey = messageMapKey(
+      params?.sessionId ?? get().activeId,
+      params?.topicId ?? get().activeTopicId,
+    );
+
+    // Get raw messages from dbMessagesMap and apply reducer
+    const nextDbMap = { ...get().dbMessagesMap, [messagesKey]: messages };
+
+    if (isEqual(nextDbMap, get().dbMessagesMap)) return;
+
+    // Parse messages using conversation-flow
+    const { flatList } = parse(messages);
+
+    set(
+      {
+        // Store raw messages from backend
+        dbMessagesMap: nextDbMap,
+        // Store parsed messages for display
+        messagesMap: { ...get().messagesMap, [messagesKey]: flatList },
+      },
+      false,
+      params?.action ?? 'replaceMessages',
+    );
+  },
+
+  useFetchMessages: (enable, messageContextId, activeTopicId, type = 'session') =>
+    useClientDataSWR<UIChatMessage[]>(
+      enable ? [SWR_USE_FETCH_MESSAGES, messageContextId, activeTopicId, type] : null,
+      async ([, sessionId, topicId, type]: [string, string, string | undefined, string]) =>
+        type === 'session'
+          ? messageService.getMessages(sessionId, topicId)
+          : messageService.getGroupMessages(sessionId, topicId),
+      {
+        onSuccess: (messages, key) => {
+          const nextMap = {
+            ...get().dbMessagesMap,
+            [messageMapKey(messageContextId || '', activeTopicId)]: messages,
+          };
+
+          // no need to update map if the messages have been init and the map is the same
+          if (get().messagesInit && isEqual(nextMap, get().dbMessagesMap)) return;
+
+          set(
+            { messagesInit: true },
+            false,
+            n('useFetchMessages(success)', { messages, queryKey: key }),
+          );
+          get().replaceMessages(messages, {
+            action: n('useFetchMessages/updateMessages'),
+          });
+        },
+      },
+    ),
+});

package/src/store/chat/slices/message/actions/runtimeState.ts

@@ -0,0 +1,108 @@
+import { StateCreator } from 'zustand/vanilla';
+
+import { ChatStore } from '@/store/chat/store';
+import { Action, setNamespace } from '@/utils/storeDebug';
+
+import type { ChatStoreState } from '../../../initialState';
+import { preventLeavingFn, toggleBooleanList } from '../../../utils';
+
+const n = setNamespace('m');
+
+/**
+ * Runtime state management for message-related states
+ * Handles loading states, active session tracking, etc.
+ */
+export interface MessageRuntimeStateAction {
+  /**
+   * helper to toggle the loading state of the array,used by these three toggleXXXLoading
+   */
+  internal_toggleLoadingArrays: (
+    key: keyof ChatStoreState,
+    loading: boolean,
+    id?: string,
+    action?: Action,
+  ) => AbortController | undefined;
+
+  /**
+   * method to toggle message create loading state
+   * the AI message status is creating -> generating
+   * other message role like user and tool , only this method need to be called
+   */
+  internal_toggleMessageLoading: (loading: boolean, id: string) => void;
+
+  /**
+   * Update active session ID with cleanup of pending operations
+   */
+  internal_updateActiveId: (activeId: string) => void;
+
+  /**
+   * Update active session type
+   */
+  internal_updateActiveSessionType: (sessionType?: 'agent' | 'group') => void;
+}
+
+export const messageRuntimeState: StateCreator<
+  ChatStore,
+  [['zustand/devtools', never]],
+  [],
+  MessageRuntimeStateAction
+> = (set, get) => ({
+  internal_toggleLoadingArrays: (key, loading, id, action) => {
+    const abortControllerKey = `${key}AbortController`;
+    if (loading) {
+      window.addEventListener('beforeunload', preventLeavingFn);
+
+      const abortController = new AbortController();
+      set(
+        {
+          [abortControllerKey]: abortController,
+          [key]: toggleBooleanList(get()[key] as string[], id!, loading),
+        },
+        false,
+        action,
+      );
+
+      return abortController;
+    } else {
+      if (!id) {
+        set({ [abortControllerKey]: undefined, [key]: [] }, false, action);
+      } else
+        set(
+          {
+            [abortControllerKey]: undefined,
+            [key]: toggleBooleanList(get()[key] as string[], id, loading),
+          },
+          false,
+          action,
+        );
+
+      window.removeEventListener('beforeunload', preventLeavingFn);
+    }
+  },
+
+  internal_toggleMessageLoading: (loading, id) => {
+    set(
+      {
+        messageLoadingIds: toggleBooleanList(get().messageLoadingIds, id, loading),
+      },
+      false,
+      `internal_toggleMessageLoading/${loading ? 'start' : 'end'}`,
+    );
+  },
+
+  internal_updateActiveId: (activeId: string) => {
+    const currentActiveId = get().activeId;
+    if (currentActiveId === activeId) return;
+
+    // Before switching sessions, cancel all pending supervisor decisions
+    get().internal_cancelAllSupervisorDecisions();
+
+    set({ activeId }, false, n(`updateActiveId/${activeId}`));
+  },
+
+  internal_updateActiveSessionType: (sessionType?: 'agent' | 'group') => {
+    if (get().activeSessionType === sessionType) return;
+
+    set({ activeSessionType: sessionType }, false, n('updateActiveSessionType'));
+  },
+});