🐛 fix: fix agent config not load correctly (#6252)

* fix agent config issue * improve loading for agent * revert chat change
lobehub · Feb 17, 2025 · fe9bc16 · fe9bc16
1 parent 977301b
commit fe9bc16
Show file tree

Hide file tree

Showing 15 changed files with 80 additions and 39 deletions.
diff --git a/src/app/[variants]/(main)/chat/(workspace)/@conversation/features/ChatInput/Mobile/index.tsx b/src/app/[variants]/(main)/chat/(workspace)/@conversation/features/ChatInput/Mobile/index.tsx
@@ -11,10 +11,10 @@ import { ActionKeys } from '@/features/ChatInput/ActionBar/config';
 import STT from '@/features/ChatInput/STT';
 import SaveTopic from '@/features/ChatInput/Topic';
 import { useSendMessage } from '@/features/ChatInput/useSend';
+import { useInitAgentConfig } from '@/hooks/useInitAgentConfig';
 import { useChatStore } from '@/store/chat';
 import { chatSelectors } from '@/store/chat/selectors';
 
-import { useInitAgentConfig } from '../../../../_layout/useInitAgentConfig';
 import Files from './Files';
 import InputArea from './InputArea';
 import SendButton from './Send';

diff --git a/src/app/[variants]/(main)/chat/(workspace)/@topic/features/SystemRole/SystemRoleContent.tsx b/src/app/[variants]/(main)/chat/(workspace)/@topic/features/SystemRole/SystemRoleContent.tsx
@@ -31,7 +31,8 @@ const SystemRole = memo(() => {
     sessionMetaSelectors.currentAgentMeta(s),
   ]);
 
-  const [systemRole, updateAgentConfig] = useAgentStore((s) => [
+  const [isAgentConfigLoading, systemRole, updateAgentConfig] = useAgentStore((s) => [
+    agentSelectors.isAgentConfigLoading(s),
     agentSelectors.currentAgentSystemRole(s),
     s.updateAgentConfig,
   ]);
@@ -49,14 +50,16 @@ const SystemRole = memo(() => {
 
   const { t } = useTranslation('common');
 
+  const isLoading = !init || isAgentConfigLoading;
+
   const handleOpenWithEdit = () => {
-    if (!init) return;
+    if (isLoading) return;
     setEditing(true);
     setOpen(true);
   };
 
   const handleOpen = () => {
-    if (!init) return;
+    if (isLoading) return;
 
     setOpen(true);
   };
@@ -77,7 +80,7 @@ const SystemRole = memo(() => {
           if (e.altKey) handleOpenWithEdit();
         }}
       >
-        {!init ? (
+        {isLoading ? (
           <Skeleton
             active
             avatar={false}

diff --git a/src/app/[variants]/(main)/chat/(workspace)/_layout/Desktop/ChatHeader/Main.tsx b/src/app/[variants]/(main)/chat/(workspace)/_layout/Desktop/ChatHeader/Main.tsx
@@ -10,13 +10,13 @@ import { useTranslation } from 'react-i18next';
 import { Flexbox } from 'react-layout-kit';
 
 import { DESKTOP_HEADER_ICON_SIZE } from '@/const/layoutTokens';
+import { useInitAgentConfig } from '@/hooks/useInitAgentConfig';
 import { useOpenChatSettings } from '@/hooks/useInterceptingRoutes';
 import { useGlobalStore } from '@/store/global';
 import { systemStatusSelectors } from '@/store/global/selectors';
 import { useSessionStore } from '@/store/session';
 import { sessionMetaSelectors, sessionSelectors } from '@/store/session/selectors';
 
-import { useInitAgentConfig } from '../../useInitAgentConfig';
 import Tags from './Tags';
 
 const Main = memo(() => {

diff --git a/src/app/[variants]/(main)/chat/(workspace)/_layout/Desktop/ChatHeader/Tags.tsx b/src/app/[variants]/(main)/chat/(workspace)/_layout/Desktop/ChatHeader/Tags.tsx
@@ -1,4 +1,5 @@
 import { ModelTag } from '@lobehub/icons';
+import { Skeleton } from 'antd';
 import isEqual from 'fast-deep-equal';
 import { memo } from 'react';
 import { Flexbox } from 'react-layout-kit';
@@ -21,8 +22,11 @@ const TitleTags = memo(() => {
   const enabledKnowledge = useAgentStore(agentSelectors.currentEnabledKnowledge, isEqual);
 
   const showPlugin = useModelSupportToolUse(model, provider);
+  const isLoading = useAgentStore(agentSelectors.isAgentConfigLoading);
 
-  return (
+  return isLoading ? (
+    <Skeleton.Button active size={'small'} style={{ height: 20 }} />
+  ) : (
     <Flexbox align={'center'} horizontal>
       <ModelSwitchPanel>
         <ModelTag model={model} />

diff --git a/src/app/[variants]/(main)/chat/(workspace)/features/AgentSettings/index.tsx b/src/app/[variants]/(main)/chat/(workspace)/features/AgentSettings/index.tsx
@@ -14,10 +14,10 @@ import AgentMeta from '@/features/AgentSetting/AgentMeta';
 import AgentModal from '@/features/AgentSetting/AgentModal';
 import AgentPlugin from '@/features/AgentSetting/AgentPlugin';
 import AgentPrompt from '@/features/AgentSetting/AgentPrompt';
+import { AgentSettingsProvider } from '@/features/AgentSetting/AgentSettingsProvider';
 import AgentTTS from '@/features/AgentSetting/AgentTTS';
-import StoreUpdater from '@/features/AgentSetting/StoreUpdater';
-import { Provider, createStore } from '@/features/AgentSetting/store';
 import Footer from '@/features/Setting/Footer';
+import { useInitAgentConfig } from '@/hooks/useInitAgentConfig';
 import { useAgentStore } from '@/store/agent';
 import { agentSelectors } from '@/store/agent/slices/chat';
 import { ChatSettingsTabs } from '@/store/global/initialState';
@@ -31,6 +31,8 @@ const AgentSettings = memo(() => {
   const id = useSessionStore((s) => s.activeId);
   const config = useAgentStore(agentSelectors.currentAgentConfig, isEqual);
   const meta = useSessionStore(sessionMetaSelectors.currentAgentMeta, isEqual);
+
+  const { isLoading } = useInitAgentConfig();
   const [showAgentSetting, updateAgentConfig] = useAgentStore((s) => [
     s.showAgentSetting,
     s.updateAgentConfig,
@@ -49,14 +51,14 @@ const AgentSettings = memo(() => {
 
   const category = <CategoryContent setTab={setTab} tab={tab} />;
   return (
-    <Provider createStore={createStore}>
-      <StoreUpdater
-        config={config}
-        id={id}
-        meta={meta}
-        onConfigChange={updateAgentConfig}
-        onMetaChange={updateAgentMeta}
-      />
+    <AgentSettingsProvider
+      config={config}
+      id={id}
+      loading={isLoading}
+      meta={meta}
+      onConfigChange={updateAgentConfig}
+      onMetaChange={updateAgentMeta}
+    >
       <Drawer
         height={'100vh'}
         onClose={() => {
@@ -109,7 +111,7 @@ const AgentSettings = memo(() => {
           </Flexbox>
         </Flexbox>
       </Drawer>
-    </Provider>
+    </AgentSettingsProvider>
   );
 });
 

diff --git a/src/app/[variants]/(main)/chat/settings/page.tsx b/src/app/[variants]/(main)/chat/settings/page.tsx
@@ -14,6 +14,7 @@ import AgentModal from '@/features/AgentSetting/AgentModal';
 import AgentPlugin from '@/features/AgentSetting/AgentPlugin';
 import AgentPrompt from '@/features/AgentSetting/AgentPrompt';
 import AgentTTS from '@/features/AgentSetting/AgentTTS';
+import { useInitAgentConfig } from '@/hooks/useInitAgentConfig';
 import { useAgentStore } from '@/store/agent';
 import { agentSelectors } from '@/store/agent/selectors';
 import { ChatSettingsTabs } from '@/store/global/initialState';
@@ -30,15 +31,14 @@ const EditPage = memo(() => {
     s.updateSessionMeta,
     sessionMetaSelectors.currentAgentTitle(s),
   ]);
-  const [useFetchAgentConfig, updateAgentConfig] = useAgentStore((s) => [
-    s.useFetchAgentConfig,
-    s.updateAgentConfig,
-  ]);
+
+  const [updateAgentConfig] = useAgentStore((s) => [s.updateAgentConfig]);
 
   const config = useAgentStore(agentSelectors.currentAgentConfig, isEqual);
   const meta = useSessionStore(sessionMetaSelectors.currentAgentMeta, isEqual);
 
-  const { isLoading } = useFetchAgentConfig(id);
+  const { isLoading } = useInitAgentConfig();
+
   const { enablePlugins } = useServerConfigStore(featureFlagsSelectors);
 
   return (

diff --git a/src/features/AgentSetting/AgentPrompt/index.tsx b/src/features/AgentSetting/AgentPrompt/index.tsx
@@ -47,6 +47,26 @@ const AgentPrompt = memo<{ modal?: boolean }>(({ modal }) => {
   ]);
 
   if (loading) {
+    if (modal)
+      return (
+        <Form
+          items={[
+            {
+              children: (
+                <>
+                  <div style={{ height: 24 }} />
+                  <Skeleton active title={false} />
+                </>
+              ),
+              title: t('settingAgent.prompt.title'),
+            },
+          ]}
+          itemsType={'group'}
+          variant={'pure'}
+          {...FORM_STYLE}
+        />
+      );
+
     return (
       <div className={styles.wrapper}>
         <Flexbox className={styles.container} padding={4}>

diff --git a/src/features/AgentSetting/store/initialState.ts b/src/features/AgentSetting/store/initialState.ts
@@ -25,5 +25,6 @@ export const initialState: State = {
     title: false,
   },
   config: DEFAULT_AGENT_CONFIG,
+  loading: true,
   meta: DEFAULT_AGENT_META,
 };
diff --git a/...(workspace)/_layout/useInitAgentConfig.ts → src/hooks/useInitAgentConfig.ts b/...(workspace)/_layout/useInitAgentConfig.ts → src/hooks/useInitAgentConfig.ts
@@ -1,10 +1,16 @@
 import { useAgentStore } from '@/store/agent';
 import { useSessionStore } from '@/store/session';
+import { useUserStore } from '@/store/user';
+import { authSelectors } from '@/store/user/selectors';
 
 export const useInitAgentConfig = () => {
   const [useFetchAgentConfig] = useAgentStore((s) => [s.useFetchAgentConfig]);
 
+  const isLogin = useUserStore(authSelectors.isLogin);
+
   const [sessionId] = useSessionStore((s) => [s.activeId]);
 
-  return useFetchAgentConfig(sessionId);
+  const data = useFetchAgentConfig(isLogin, sessionId);
+
+  return { ...data, isLoading: data.isLoading && isLogin };
 };
diff --git a/src/services/message/client.test.ts b/src/services/message/client.test.ts
@@ -1,8 +1,6 @@
-import dayjs from 'dayjs';
 import { and, eq } from 'drizzle-orm';
-import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
-import { MessageModel } from '@/database/_deprecated/models/message';
 import { clientDB, initializeDB } from '@/database/client/db';
 import {
   files,

diff --git a/src/services/session/server.ts b/src/services/session/server.ts
@@ -1,8 +1,5 @@
 /* eslint-disable @typescript-eslint/no-unused-vars */
-import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
 import { lambdaClient } from '@/libs/trpc/client';
-import { useUserStore } from '@/store/user';
-import { authSelectors } from '@/store/user/selectors';
 
 import { ISessionService } from './type';
 
@@ -55,9 +52,6 @@ export class ServerService implements ISessionService {
   // TODO: Need to be fixed
   // @ts-ignore
   getSessionConfig: ISessionService['getSessionConfig'] = async (id) => {
-    const isLogin = authSelectors.isLogin(useUserStore.getState());
-    if (!isLogin) return DEFAULT_AGENT_CONFIG;
-
     return lambdaClient.agent.getAgentConfig.query({ sessionId: id });
   };
 

diff --git a/src/store/agent/slices/chat/action.test.ts b/src/store/agent/slices/chat/action.test.ts
@@ -175,7 +175,7 @@ describe('AgentSlice', () => {
 
       vi.spyOn(sessionService, 'getSessionConfig').mockResolvedValueOnce({ model: 'gpt-4' } as any);
 
-      renderHook(() => result.current.useFetchAgentConfig('test-session-id'));
+      renderHook(() => result.current.useFetchAgentConfig(true, 'test-session-id'));
 
       await waitFor(() => {
         expect(result.current.agentMap['test-session-id']).toEqual({ model: 'gpt-4' });
@@ -199,7 +199,7 @@ describe('AgentSlice', () => {
         model: 'gpt-3.5-turbo',
       } as any);
 
-      renderHook(() => result.current.useFetchAgentConfig('test-session-id'));
+      renderHook(() => result.current.useFetchAgentConfig(true, 'test-session-id'));
 
       await waitFor(() => {
         expect(result.current.agentMap['test-session-id']).toEqual({ model: 'gpt-3.5-turbo' });

diff --git a/src/store/agent/slices/chat/action.ts b/src/store/agent/slices/chat/action.ts
@@ -48,7 +48,7 @@ export interface AgentChatAction {
   togglePlugin: (id: string, open?: boolean) => Promise<void>;
   updateAgentChatConfig: (config: Partial<LobeAgentChatConfig>) => Promise<void>;
   updateAgentConfig: (config: DeepPartial<LobeAgentConfig>) => Promise<void>;
-  useFetchAgentConfig: (id: string) => SWRResponse<LobeAgentConfig>;
+  useFetchAgentConfig: (isLogin: boolean | undefined, id: string) => SWRResponse<LobeAgentConfig>;
   useFetchFilesAndKnowledgeBases: () => SWRResponse<KnowledgeItem[]>;
   useInitInboxAgentStore: (
     isLogin: boolean | undefined,
@@ -158,14 +158,22 @@ export const createChatSlice: StateCreator<
 
     await get().internal_updateAgentConfig(activeId, config, controller.signal);
   },
-  useFetchAgentConfig: (sessionId) =>
+  useFetchAgentConfig: (isLogin, sessionId) =>
     useClientDataSWR<LobeAgentConfig>(
-      [FETCH_AGENT_CONFIG_KEY, sessionId],
+      isLogin ? [FETCH_AGENT_CONFIG_KEY, sessionId] : null,
       ([, id]: string[]) => sessionService.getSessionConfig(id),
       {
         onSuccess: (data) => {
           get().internal_dispatchAgentMap(sessionId, data, 'fetch');
-          set({ activeAgentId: data.id }, false, 'updateActiveAgentId');
+
+          set(
+            {
+              activeAgentId: data.id,
+              agentConfigInitMap: { ...get().agentConfigInitMap, [sessionId]: true },
+            },
+            false,
+            'fetchAgentConfig',
+          );
         },
       },
     ),

diff --git a/src/store/agent/slices/chat/initialState.ts b/src/store/agent/slices/chat/initialState.ts
@@ -7,6 +7,7 @@ import { LobeAgentConfig } from '@/types/agent';
 export interface AgentState {
   activeAgentId?: string;
   activeId: string;
+  agentConfigInitMap: Record<string, boolean>;
   agentMap: Record<string, DeepPartial<LobeAgentConfig>>;
   agentSettingInstance?: AgentSettingsInstance | null;
   defaultAgentConfig: LobeAgentConfig;
@@ -18,6 +19,7 @@ export interface AgentState {
 
 export const initialAgentChatState: AgentState = {
   activeId: 'inbox',
+  agentConfigInitMap: {},
   agentMap: {},
   defaultAgentConfig: DEFAULT_AGENT_CONFIG,
   isInboxAgentConfigInit: false,

diff --git a/src/store/agent/slices/chat/selectors.ts b/src/store/agent/slices/chat/selectors.ts
@@ -138,6 +138,8 @@ const currentKnowledgeIds = (s: AgentStore) => {
   };
 };
 
+const isAgentConfigLoading = (s: AgentStore) => !s.agentConfigInitMap[s.activeId];
+
 export const agentSelectors = {
   currentAgentChatConfig,
   currentAgentConfig,
@@ -157,5 +159,6 @@ export const agentSelectors = {
   hasSystemRole,
   inboxAgentConfig,
   inboxAgentModel,
+  isAgentConfigLoading,
   isInboxSession,
 };