♻️ refactor: refactor the create message flow to fix some bugs (lobehub#2521)

* ♻️ refactor: refactor the create message flow

* ✅ test: fix test

* ✅ test: add test

* ✅ test: add more test

* 💄 style: improve loading state

Author: arvinxx

src/app/(main)/chat/@session/features/SessionHydration.tsx

@@ -13,6 +13,7 @@ import { useSessionStore } from '@/store/session';
 const SessionHydration = memo(() => {
   const useStoreUpdater = createStoreUpdater(useSessionStore);
   const useAgentStoreUpdater = createStoreUpdater(useAgentStore);
+  const useChatStoreUpdater = createStoreUpdater(useChatStore);
   const [switchTopic] = useChatStore((s) => [s.switchTopic]);
 
   // two-way bindings the url and session store
@@ -22,6 +23,7 @@ const SessionHydration = memo(() => {
   );
   useStoreUpdater('activeId', session);
   useAgentStoreUpdater('activeId', session);
+  useChatStoreUpdater('activeId', session);
 
   useEffect(() => {
     const unsubscribe = useSessionStore.subscribe(

src/features/Conversation/Error/index.tsx

@@ -1,6 +1,8 @@
 import { IPluginErrorType, PluginErrorType } from '@lobehub/chat-plugin-sdk';
 import type { AlertProps } from '@lobehub/ui';
-import { memo } from 'react';
+import { Skeleton } from 'antd';
+import dynamic from 'next/dynamic';
+import { Suspense, memo } from 'react';
 
 import { AgentRuntimeErrorType, ILobeAgentRuntimeErrorType } from '@/libs/agent-runtime';
 import { ChatErrorType, ErrorType } from '@/types/fetch';
@@ -10,9 +12,12 @@ import ClerkLogin from './ClerkLogin';
 import ErrorJsonViewer from './ErrorJsonViewer';
 import InvalidAPIKey from './InvalidAPIKey';
 import InvalidAccessCode from './InvalidAccessCode';
-import OllamaBizError from './OllamaBizError';
 import OpenAiBizError from './OpenAiBizError';
-import PluginSettings from './PluginSettings';
+
+const loading = () => <Skeleton active />;
+
+const OllamaBizError = dynamic(() => import('./OllamaBizError'), { loading, ssr: false });
+const PluginSettings = dynamic(() => import('./PluginSettings'), { loading, ssr: false });
 
 // Config for the errorMessage display
 export const getErrorAlertConfig = (
@@ -95,4 +100,8 @@ const ErrorMessageExtra = memo<{ data: ChatMessage }>(({ data }) => {
   }
 });
 
-export default ErrorMessageExtra;
+export default memo<{ data: ChatMessage }>(({ data }) => (
+  <Suspense fallback={<Skeleton active style={{ width: '100%' }} />}>
+    <ErrorMessageExtra data={data} />
+  </Suspense>
+));

src/features/Conversation/components/VirtualizedList/index.tsx

@@ -29,14 +29,17 @@ const VirtualizedList = memo<VirtualizedListProps>(({ mobile }) => {
 
   const [id] = useChatStore((s) => [chatSelectors.currentChatKey(s)]);
 
-  const [activeTopicId, useFetchMessages, isFirstLoading] = useChatStore((s) => [
-    s.activeTopicId,
-    s.useFetchMessages,
-    chatSelectors.currentChatLoadingState(s),
-  ]);
+  const [activeTopicId, useFetchMessages, isFirstLoading, isCurrentChatLoaded] = useChatStore(
+    (s) => [
+      s.activeTopicId,
+      s.useFetchMessages,
+      chatSelectors.currentChatLoadingState(s),
+      chatSelectors.isCurrentChatLoaded(s),
+    ],
+  );
 
   const [sessionId] = useSessionStore((s) => [s.activeId]);
-  const { isLoading } = useFetchMessages(sessionId, activeTopicId);
+  useFetchMessages(sessionId, activeTopicId);
 
   const data = useChatStore((s) => {
     const showInboxWelcome = chatSelectors.showInboxWelcome(s);
@@ -77,23 +80,26 @@ const VirtualizedList = memo<VirtualizedListProps>(({ mobile }) => {
     [mobile],
   );
 
-  // first time loading
+  // first time loading or not loaded
   if (isFirstLoading) return <SkeletonList mobile={mobile} />;
 
-  // in server mode and switch page
-  if (isServerMode && isLoading) return <SkeletonList mobile={mobile} />;
-
-  // in client mode using the center loading for more
-  return isLoading ? (
-    <Center height={'100%'} width={'100%'}>
-      <Icon
-        icon={Loader2Icon}
-        size={{ fontSize: 32 }}
-        spin
-        style={{ color: theme.colorTextTertiary }}
-      />
-    </Center>
-  ) : (
+  if (!isCurrentChatLoaded)
+    // use skeleton list when not loaded in server mode due to the loading duration is much longer than client mode
+    return isServerMode ? (
+      <SkeletonList mobile={mobile} />
+    ) : (
+      // in client mode and switch page, using the center loading for smooth transition
+      <Center height={'100%'} width={'100%'}>
+        <Icon
+          icon={Loader2Icon}
+          size={{ fontSize: 32 }}
+          spin
+          style={{ color: theme.colorTextTertiary }}
+        />
+      </Center>
+    );
+
+  return (
     <Flexbox height={'100%'}>
       <Virtuoso
         atBottomStateChange={setAtBottom}

src/services/message/type.ts

@@ -1,15 +1,18 @@
 import { DB_Message } from '@/database/client/schemas/message';
-import { ChatMessage, ChatMessageError, ChatPluginPayload } from '@/types/message';
+import { ChatMessage, ChatMessageError, ChatPluginPayload, MessageRoleType } from '@/types/message';
 
 /* eslint-disable typescript-sort-keys/interface */
 
 export interface CreateMessageParams
-  extends Partial<Omit<ChatMessage, 'content' | 'role'>>,
-    Pick<ChatMessage, 'content' | 'role'> {
+  extends Partial<Omit<ChatMessage, 'content' | 'role' | 'topicId'>> {
   fromModel?: string;
   fromProvider?: string;
   sessionId: string;
   traceId?: string;
+  topicId?: string;
+  content: string;
+  error?: ChatMessageError;
+  role: MessageRoleType;
 }
 
 export interface IMessageService {

src/store/chat/slices/enchance/action.test.ts

@@ -5,6 +5,7 @@ import { chainLangDetect } from '@/chains/langDetect';
 import { chainTranslate } from '@/chains/translate';
 import { chatService } from '@/services/chat';
 import { messageService } from '@/services/message';
+import { messageMapKey } from '@/store/chat/slices/message/utils';
 
 import { useChatStore } from '../../store';
 
@@ -73,18 +74,21 @@ describe('ChatEnhanceAction', () => {
 
       act(() => {
         useChatStore.setState({
-          messages: [
-            {
-              id: messageId,
-              content: messageContent,
-              createdAt: Date.now(),
-              updatedAt: Date.now(),
-              role: 'user',
-              sessionId: 'test',
-              topicId: 'test',
-              meta: {},
-            },
-          ],
+          activeId: 'session',
+          messagesMap: {
+            [messageMapKey('session')]: [
+              {
+                id: messageId,
+                content: messageContent,
+                createdAt: Date.now(),
+                updatedAt: Date.now(),
+                role: 'user',
+                sessionId: 'test',
+                topicId: 'test',
+                meta: {},
+              },
+            ],
+          },
         });
       });