@lobehub/lobehub 2.0.0-next.4 → 2.0.0-next.40

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'));
+  },
+});

package/src/store/chat/slices/message/initialState.ts

@@ -15,6 +15,14 @@ export interface ChatMessageState {
    * Derived from session.type, used for caching to avoid repeated lookups
    */
   activeSessionType?: 'agent' | 'group';
+  /**
+   * is the message is continuing generation (used for disable continue button)
+   */
+  continuingIds: string[];
+  /**
+   * Raw messages from database (flat structure)
+   */
+  dbMessagesMap: Record<string, UIChatMessage[]>;
   /**
    * Group agents maps by group ID
    */
@@ -40,7 +48,14 @@ export interface ChatMessageState {
    * whether messages have fetched
    */
   messagesInit: boolean;
+  /**
+   * Parsed messages for display (includes assistantGroup from conversation-flow)
+   */
   messagesMap: Record<string, UIChatMessage[]>;
+  /**
+   * is the message is regenerating (used for disable regenerate button)
+   */
+  regeneratingIds: string[];
   /**
    * Supervisor decision debounce timers by group ID
    */
@@ -62,6 +77,8 @@ export interface ChatMessageState {
 export const initialMessageState: ChatMessageState = {
   activeId: 'inbox',
   activeSessionType: undefined,
+  continuingIds: [],
+  dbMessagesMap: {},
   groupAgentMaps: {},
   groupMaps: {},
   groupsInit: false,
@@ -70,6 +87,7 @@ export const initialMessageState: ChatMessageState = {
   messageLoadingIds: [],
   messagesInit: false,
   messagesMap: {},
+  regeneratingIds: [],
   supervisorDebounceTimers: {},
   supervisorDecisionAbortControllers: {},
   supervisorDecisionLoading: [],

package/src/store/chat/slices/message/reducer.test.ts

@@ -56,18 +56,6 @@ describe('messagesReducer', () => {
 
       expect(newState).toEqual(initialState);
     });
-
-    it('should not modify the state if the specified message does not exist', () => {
-      const payload: MessageDispatch = {
-        type: 'updateMessage',
-        id: 'nonexistentMessage',
-        value: { content: 'Updated Message' },
-      };
-
-      const newState = messagesReducer(initialState, payload);
-
-      expect(newState).toEqual(initialState);
-    });
   });
 
   describe('unimplemented type', () => {
@@ -139,6 +127,54 @@ describe('messagesReducer', () => {
     });
   });
 
+  describe('updateMessageMetadata', () => {
+    it('should merge update the metadata field of a message', () => {
+      const payload: MessageDispatch = {
+        type: 'updateMessageMetadata',
+        id: 'message1',
+        value: { activeBranchIndex: 1 },
+      };
+
+      const newState = messagesReducer(initialState, payload);
+      const updatedMessage = newState.find((m) => m.id === 'message1');
+
+      expect(updatedMessage?.metadata).toEqual({ activeBranchIndex: 1 });
+      expect(updatedMessage?.updatedAt).toBeGreaterThan(initialState[0].updatedAt);
+    });
+
+    it('should merge update the metadata field if metadata already exists', () => {
+      const state = [
+        {
+          ...initialState[0],
+          metadata: { activeBranchIndex: 0 },
+        },
+      ];
+
+      const payload: MessageDispatch = {
+        type: 'updateMessageMetadata',
+        id: 'message1',
+        value: { activeBranchIndex: 2 },
+      };
+
+      const newState = messagesReducer(state, payload);
+      const updatedMessage = newState.find((m) => m.id === 'message1');
+
+      expect(updatedMessage?.metadata).toEqual({ activeBranchIndex: 2 });
+      expect(updatedMessage?.updatedAt).toBeGreaterThan(initialState[0].updatedAt);
+    });
+
+    it('should not modify state if message is not found', () => {
+      const payload: MessageDispatch = {
+        type: 'updateMessageMetadata',
+        id: 'nonexistent',
+        value: { activeBranchIndex: 1 },
+      };
+
+      const newState = messagesReducer(initialState, payload);
+      expect(newState).toEqual(initialState);
+    });
+  });
+
   describe('updatePluginState', () => {
     it('should update the plugin state of a message', () => {
       const payload: MessageDispatch = {

package/src/store/chat/slices/message/reducer.ts

@@ -75,12 +75,19 @@ interface UpdateMessageExtra {
   value: any;
 }
 
+interface UpdateMessageMetadata {
+  id: string;
+  type: 'updateMessageMetadata';
+  value: Partial<UIChatMessage['metadata']>;
+}
+
 export type MessageDispatch =
   | CreateMessage
   | UpdateMessage
   | UpdateMessages
   | UpdatePluginState
   | UpdateMessageExtra
+  | UpdateMessageMetadata
   | DeleteMessage
   | UpdateMessagePlugin
   | UpdateMessageTools
@@ -96,10 +103,11 @@ export const messagesReducer = (
     case 'updateMessage': {
       return produce(state, (draftState) => {
         const { id, value } = payload;
-        const index = draftState.findIndex((i) => i.id === id);
-        if (index < 0) return;
 
-        draftState[index] = merge(draftState[index], { ...value, updatedAt: Date.now() });
+        const index = draftState.findIndex((i) => i.id === id);
+        if (index >= 0) {
+          draftState[index] = merge(draftState[index], { ...value, updatedAt: Date.now() });
+        }
       });
     }
 
@@ -119,6 +127,17 @@ export const messagesReducer = (
       });
     }
 
+    case 'updateMessageMetadata': {
+      return produce(state, (draftState) => {
+        const { id, value } = payload;
+        const message = draftState.find((i) => i.id === id);
+        if (!message) return;
+
+        message.metadata = merge(message.metadata, value);
+        message.updatedAt = Date.now();
+      });
+    }
+
     case 'updatePluginState': {
       return produce(state, (draftState) => {
         const { id, key, value } = payload;
@@ -226,6 +245,7 @@ export const messagesReducer = (
         });
       });
     }
+
     default: {
       throw new Error('暂未实现的 type，请检查 reducer');
     }

package/src/store/chat/slices/message/selectors/chat.test.ts

@@ -0,0 +1,92 @@
+import { UIChatMessage } from '@lobechat/types';
+import { describe, expect, it } from 'vitest';
+
+import { ChatStore } from '@/store/chat';
+import { initialState } from '@/store/chat/initialState';
+import { messageMapKey } from '@/store/chat/utils/messageMapKey';
+import { createServerConfigStore } from '@/store/serverConfig/store';
+import { merge } from '@/utils/merge';
+
+import { chatSelectors } from './chat';
+
+vi.mock('i18next', () => ({
+  t: vi.fn((key) => key),
+}));
+
+const initialStore = initialState as ChatStore;
+
+const mockMessages = [
+  {
+    id: 'msg1',
+    content: 'Hello World',
+    role: 'user',
+  },
+  {
+    id: 'msg2',
+    content: 'Goodbye World',
+    role: 'user',
+  },
+] as UIChatMessage[];
+
+beforeAll(() => {
+  createServerConfigStore();
+});
+
+describe('chatSelectors - Backward Compatibility Layer', () => {
+  describe('getMessageById', () => {
+    it('should work as backward compatibility alias for getDisplayMessageById', () => {
+      const state = merge(initialStore, {
+        messagesMap: {
+          [messageMapKey('abc')]: mockMessages,
+        },
+        activeId: 'abc',
+      });
+      const message = chatSelectors.getMessageById('msg1')(state);
+      expect(message?.id).toBe('msg1');
+      expect(message?.content).toBe('Hello World');
+    });
+  });
+
+  describe('currentChatKey', () => {
+    it('should work as backward compatibility alias for currentDisplayChatKey', () => {
+      const state: Partial<ChatStore> = {
+        activeId: 'testId',
+        activeTopicId: 'topicId',
+      };
+      const result = chatSelectors.currentChatKey(state as ChatStore);
+      expect(result).toBe(messageMapKey('testId', 'topicId'));
+    });
+  });
+
+  describe('activeBaseChats', () => {
+    it('should work as backward compatibility alias for activeDisplayMessages', () => {
+      const state = merge(initialStore, {
+        messagesMap: {
+          [messageMapKey('abc')]: mockMessages,
+        },
+        activeId: 'abc',
+      });
+      const chats = chatSelectors.activeBaseChats(state);
+      expect(chats).toHaveLength(2);
+      expect(chats[0].id).toBe('msg1');
+    });
+  });
+
+  describe('currentToolMessages', () => {
+    it('should work as backward compatibility alias for dbToolMessages', () => {
+      const messages = [
+        { id: '1', role: 'user', content: 'Hello' },
+        { id: '2', role: 'tool', content: 'Tool message' },
+      ] as UIChatMessage[];
+      const state: Partial<ChatStore> = {
+        activeId: 'test-id',
+        dbMessagesMap: {
+          [messageMapKey('test-id')]: messages,
+        },
+      };
+      const result = chatSelectors.currentToolMessages(state as ChatStore);
+      expect(result).toHaveLength(1);
+      expect(result[0].role).toBe('tool');
+    });
+  });
+});

package/src/store/chat/slices/message/selectors/chat.ts

@@ -0,0 +1,106 @@
+/**
+ * Chat Selectors - Backward Compatibility Layer
+ *
+ * This file provides backward compatibility by re-exporting selectors from:
+ * - dbMessage.ts: Raw database message selectors (for data operations)
+ * - displayMessage.ts: Display message selectors (for UI rendering)
+ *
+ * MIGRATION GUIDE:
+ * ================
+ *
+ * For Data Operations (create, update, delete):
+ * ----------------------------------------------
+ * Use dbMessageSelectors from './dbMessage'
+ *
+ * Before: chatSelectors.getMessageById(id)(state)
+ * After:  dbMessageSelectors.getDbMessageById(id)(state)
+ *
+ * Before: chatSelectors.latestMessage(state)
+ * After:  dbMessageSelectors.latestDbMessage(state)
+ *
+ * Before: chatSelectors.currentToolMessages(state)
+ * After:  dbMessageSelectors.dbToolMessages(state)
+ *
+ * For UI Rendering:
+ * -----------------
+ * Use displayMessageSelectors from './displayMessage'
+ *
+ * Before: chatSelectors.mainDisplayChats(state)
+ * After:  displayMessageSelectors.mainDisplayChats(state)
+ *
+ * Before: chatSelectors.getMessageById(id)(state)  // For display
+ * After:  displayMessageSelectors.getDisplayMessageById(id)(state)
+ */
+import { ChatFileItem } from '@lobechat/types';
+
+import type { ChatStoreState } from '../../../initialState';
+import { dbMessageSelectors } from './dbMessage';
+import { displayMessageSelectors } from './displayMessage';
+
+// Re-export new selector groups
+export { dbMessageSelectors } from './dbMessage';
+export { displayMessageSelectors } from './displayMessage';
+
+// ============= Backward Compatibility Exports ========== //
+// These maintain the old API for gradual migration
+// Prefer using dbMessageSelectors or displayMessageSelectors directly
+
+/**
+ * @deprecated Use displayMessageSelectors.currentDisplayChatKey instead
+ */
+export const currentChatKey = displayMessageSelectors.currentDisplayChatKey;
+
+/**
+ * @deprecated Use displayMessageSelectors.mainDisplayChatIDs instead
+ */
+export const mainDisplayChatIDs = displayMessageSelectors.mainDisplayChatIDs;
+
+/**
+ * Legacy chatSelectors object for backward compatibility
+ * @deprecated Import dbMessageSelectors or displayMessageSelectors directly
+ */
+export const chatSelectors = {
+  // Display message selectors (from messagesMap)
+  activeBaseChats: displayMessageSelectors.activeDisplayMessages,
+  // DB message selectors (from dbMessagesMap)
+  countMessagesByThreadId: dbMessageSelectors.countDbMessagesByThreadId,
+
+  currentChatKey: displayMessageSelectors.currentDisplayChatKey,
+
+  currentChatLoadingState: displayMessageSelectors.currentChatLoadingState,
+
+  currentToolMessages: dbMessageSelectors.dbToolMessages,
+
+  currentUserFiles: (s: ChatStoreState) =>
+    dbMessageSelectors.dbUserFiles(s) as unknown as ChatFileItem[],
+
+  getBaseChatsByKey: displayMessageSelectors.getDisplayMessagesByKey,
+
+  getMessageById: displayMessageSelectors.getDisplayMessageById,
+
+  getSupervisorTodos: displayMessageSelectors.getSupervisorTodos,
+
+  getThreadMessageIDs: displayMessageSelectors.getThreadMessageIDs,
+
+  getThreadMessages: displayMessageSelectors.getThreadMessages,
+
+  getTraceIdByMessageId: dbMessageSelectors.getTraceIdByDbMessageId,
+
+  inboxActiveTopicMessages: displayMessageSelectors.inboxActiveTopicDisplayMessages,
+
+  isCurrentChatLoaded: displayMessageSelectors.isCurrentDisplayChatLoaded,
+
+  isSupervisorLoading: displayMessageSelectors.isSupervisorLoading,
+
+  latestMessage: dbMessageSelectors.latestDbMessage,
+
+  mainAIChats: displayMessageSelectors.mainAIChats,
+
+  mainAIChatsMessageString: displayMessageSelectors.mainAIChatsMessageString,
+
+  mainAIChatsWithHistoryConfig: displayMessageSelectors.mainAIChatsWithHistoryConfig,
+  mainAILatestMessageReasoningContent: displayMessageSelectors.mainAILatestMessageReasoningContent,
+  mainDisplayChatIDs: displayMessageSelectors.mainDisplayChatIDs,
+  mainDisplayChats: displayMessageSelectors.mainDisplayChats,
+  showInboxWelcome: displayMessageSelectors.showInboxWelcome,
+};

package/src/store/chat/slices/message/selectors/dbMessage.test.ts

@@ -0,0 +1,51 @@
+import { UIChatMessage } from '@lobechat/types';
+import { describe, expect, it } from 'vitest';
+
+import { ChatStore } from '@/store/chat';
+import { messageMapKey } from '@/store/chat/utils/messageMapKey';
+
+import { dbMessageSelectors } from './dbMessage';
+
+describe('dbMessageSelectors', () => {
+  describe('dbToolMessages', () => {
+    it('should return only tool messages', () => {
+      const messages = [
+        { id: '1', role: 'user', content: 'Hello' },
+        { id: '2', role: 'assistant', content: 'Hi' },
+        { id: '3', role: 'tool', content: 'Tool message 1' },
+        { id: '4', role: 'user', content: 'Query' },
+        { id: '5', role: 'tool', tools: [] },
+      ] as UIChatMessage[];
+      const state: Partial<ChatStore> = {
+        activeId: 'test-id',
+        dbMessagesMap: {
+          [messageMapKey('test-id')]: messages,
+        },
+        messagesMap: {
+          [messageMapKey('test-id')]: messages,
+        },
+      };
+      const result = dbMessageSelectors.dbToolMessages(state as ChatStore);
+      expect(result).toHaveLength(2);
+      expect(result.every((msg) => msg.role === 'tool')).toBe(true);
+    });
+
+    it('should return an empty array when no tool messages exist', () => {
+      const messages = [
+        { id: '1', role: 'user', content: 'Hello' },
+        { id: '2', role: 'assistant', content: 'Hi' },
+      ] as UIChatMessage[];
+      const state: Partial<ChatStore> = {
+        activeId: 'test-id',
+        dbMessagesMap: {
+          [messageMapKey('test-id')]: messages,
+        },
+        messagesMap: {
+          [messageMapKey('test-id')]: messages,
+        },
+      };
+      const result = dbMessageSelectors.dbToolMessages(state as ChatStore);
+      expect(result).toHaveLength(0);
+    });
+  });
+});

package/src/store/chat/slices/message/selectors/dbMessage.ts

@@ -0,0 +1,140 @@
+import { UIChatMessage } from '@lobechat/types';
+
+import { chatHelpers } from '../../../helpers';
+import type { ChatStoreState } from '../../../initialState';
+import { messageMapKey } from '../../../utils/messageMapKey';
+
+/**
+ * DB Message Selectors
+ *
+ * These selectors access raw messages from dbMessagesMap (database source).
+ * Use these selectors when you need to:
+ * - Operate on raw message data (create, update, delete)
+ * - Access original message structure before processing
+ * - Perform data mutations or service calls
+ *
+ * DO NOT use these for UI rendering - use displayMessage.ts selectors instead.
+ */
+
+// ============= Basic DB Message Access ========== //
+
+/**
+ * Get the current chat key for accessing dbMessagesMap
+ */
+export const currentDbChatKey = (s: ChatStoreState) => messageMapKey(s.activeId, s.activeTopicId);
+
+/**
+ * Get raw messages from database by key
+ */
+const getDbMessagesByKey =
+  (key: string) =>
+  (s: ChatStoreState): UIChatMessage[] => {
+    return s.dbMessagesMap[key] || [];
+  };
+
+/**
+ * Get current active session's raw messages from database
+ */
+const activeDbMessages = (s: ChatStoreState): UIChatMessage[] => {
+  if (!s.activeId) return [];
+  return getDbMessagesByKey(currentDbChatKey(s))(s);
+};
+
+// ============= DB Message Queries ========== //
+
+/**
+ * Get raw message by ID from database
+ * This searches in dbMessagesMap, which contains flat message structure
+ */
+const getDbMessageById = (id: string) => (s: ChatStoreState) =>
+  chatHelpers.getMessageById(activeDbMessages(s), id);
+
+/**
+ * Get raw message by tool_call_id from database
+ */
+export const getDbMessageByToolCallId = (id: string) => (s: ChatStoreState) => {
+  const messages = activeDbMessages(s);
+  return messages.find((m) => m.tool_call_id === id);
+};
+
+/**
+ * Get traceId from a message by ID
+ */
+const getTraceIdByDbMessageId = (id: string) => (s: ChatStoreState) =>
+  getDbMessageById(id)(s)?.traceId;
+
+/**
+ * Get latest raw message from database
+ */
+const latestDbMessage = (s: ChatStoreState) => activeDbMessages(s).at(-1);
+
+// ============= DB Message Filtering ========== //
+
+/**
+ * Get all tool messages from database
+ */
+const dbToolMessages = (s: ChatStoreState) => {
+  const messages = activeDbMessages(s);
+  return messages.filter((m) => m.role === 'tool');
+};
+
+/**
+ * Get all user messages from database
+ */
+const dbUserMessages = (s: ChatStoreState) => {
+  const messages = activeDbMessages(s);
+  return messages.filter((m) => m.role === 'user');
+};
+
+/**
+ * Get all file attachments from user messages
+ */
+const dbUserFiles = (s: ChatStoreState) => {
+  const userMessages = dbUserMessages(s);
+  return userMessages
+    .filter((m) => m.fileList && m.fileList.length > 0)
+    .flatMap((m) => m.fileList)
+    .filter(Boolean);
+};
+
+// ============= DB Message Counting ========== //
+
+/**
+ * Count messages in a specific thread
+ */
+const countDbMessagesByThreadId = (id: string) => (s: ChatStoreState) => {
+  const messages = activeDbMessages(s).filter((m) => m.threadId === id);
+  return messages.length;
+};
+
+// ============= DB Message Status ========== //
+
+/**
+ * Check if current chat is loaded in dbMessagesMap
+ */
+const isCurrentDbChatLoaded = (s: ChatStoreState) => !!s.dbMessagesMap[currentDbChatKey(s)];
+
+/**
+ * Get inbox active topic raw messages from database
+ */
+const inboxActiveTopicDbMessages = (state: ChatStoreState) => {
+  const activeTopicId = state.activeTopicId;
+  const key = messageMapKey('inbox', activeTopicId);
+  return state.dbMessagesMap[key] || [];
+};
+
+export const dbMessageSelectors = {
+  activeDbMessages,
+  countDbMessagesByThreadId,
+  currentDbChatKey,
+  dbToolMessages,
+  dbUserFiles,
+  dbUserMessages,
+  getDbMessageById,
+  getDbMessageByToolCallId,
+  getDbMessagesByKey,
+  getTraceIdByDbMessageId,
+  inboxActiveTopicDbMessages,
+  isCurrentDbChatLoaded,
+  latestDbMessage,
+};