@lobehub/lobehub 2.0.0-next.2 → 2.0.0-next.21

package/src/services/message/index.ts

@@ -1,13 +1,3 @@
-import { isDesktop } from '@/const/version';
-
-import { ClientService as DeprecatedService } from './_deprecated';
-import { ClientService } from './client';
 import { ServerService } from './server';
 
-const clientService =
-  process.env.NEXT_PUBLIC_CLIENT_DB === 'pglite' ? new ClientService() : new DeprecatedService();
-
-export const messageService =
-  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' || isDesktop
-    ? new ServerService()
-    : clientService;
+export const messageService = new ServerService();

package/src/services/message/server.ts

@@ -21,10 +21,6 @@ export class ServerService implements IMessageService {
     });
   };
 
-  batchCreateMessages: IMessageService['batchCreateMessages'] = async (messages) => {
-    return lambdaClient.message.batchCreateMessages.mutate(messages);
-  };
-
   getMessages: IMessageService['getMessages'] = async (sessionId, topicId, groupId) => {
     const data = await lambdaClient.message.getMessages.query({
       groupId,
@@ -81,8 +77,13 @@ export class ServerService implements IMessageService {
     return lambdaClient.message.updateMessagePlugin.mutate({ id, value: { arguments: args } });
   };
 
-  updateMessage: IMessageService['updateMessage'] = async (id, value) => {
-    return lambdaClient.message.update.mutate({ id, value });
+  updateMessage: IMessageService['updateMessage'] = async (id, value, options) => {
+    return lambdaClient.message.update.mutate({
+      id,
+      sessionId: options?.sessionId,
+      topicId: options?.topicId,
+      value,
+    });
   };
 
   updateMessageTranslate: IMessageService['updateMessageTranslate'] = async (id, translate) => {

package/src/services/message/type.ts

@@ -5,11 +5,11 @@ import {
   ChatTranslate,
   CreateMessageParams,
   CreateMessageResult,
-  DBMessageItem,
   ModelRankItem,
   UIChatMessage,
   UpdateMessageParams,
   UpdateMessageRAGParams,
+  UpdateMessageResult,
 } from '@lobechat/types';
 import type { HeatmapsProps } from '@lobehub/charts';
 
@@ -18,7 +18,6 @@ import type { HeatmapsProps } from '@lobehub/charts';
 export interface IMessageService {
   createMessage(data: CreateMessageParams): Promise<string>;
   createNewMessage(data: CreateMessageParams): Promise<CreateMessageResult>;
-  batchCreateMessages(messages: DBMessageItem[]): Promise<any>;
 
   getMessages(sessionId: string, topicId?: string, groupId?: string): Promise<UIChatMessage[]>;
   getGroupMessages(groupId: string, topicId?: string): Promise<UIChatMessage[]>;
@@ -37,7 +36,11 @@ export interface IMessageService {
   rankModels(): Promise<ModelRankItem[]>;
   getHeatmaps(): Promise<HeatmapsProps['data']>;
   updateMessageError(id: string, error: ChatMessageError): Promise<any>;
-  updateMessage(id: string, message: Partial<UpdateMessageParams>): Promise<any>;
+  updateMessage(
+    id: string,
+    message: Partial<UpdateMessageParams>,
+    options?: { sessionId?: string | null; topicId?: string | null },
+  ): Promise<UpdateMessageResult>;
   updateMessageTTS(id: string, tts: Partial<ChatTTS> | false): Promise<any>;
   updateMessageTranslate(id: string, translate: Partial<ChatTranslate> | false): Promise<any>;
   updateMessagePluginState(id: string, value: Record<string, any>): Promise<any>;

package/src/services/models.ts

@@ -1,8 +1,5 @@
-import { isDeprecatedEdition } from '@/const/version';
 import { createHeaderWithAuth } from '@/services/_auth';
 import { aiProviderSelectors, getAiInfraStoreState } from '@/store/aiInfra';
-import { useUserStore } from '@/store/user';
-import { modelConfigSelectors } from '@/store/user/selectors';
 import { ChatModelCard } from '@/types/llm';
 import { getMessageError } from '@/utils/fetch';
 
@@ -10,14 +7,8 @@ import { API_ENDPOINTS } from './_url';
 import { initializeWithClientStore } from './chat/clientModelRuntime';
 import { resolveRuntimeProvider } from './chat/helper';
 
-const isEnableFetchOnClient = (provider: string) => {
-  // TODO: remove this condition in V2.0
-  if (isDeprecatedEdition) {
-    return modelConfigSelectors.isProviderFetchOnClient(provider)(useUserStore.getState());
-  } else {
-    return aiProviderSelectors.isProviderFetchOnClient(provider)(getAiInfraStoreState());
-  }
-};
+const isEnableFetchOnClient = (provider: string) =>
+  aiProviderSelectors.isProviderFetchOnClient(provider)(getAiInfraStoreState());
 
 // 进度信息接口
 export interface ModelProgressInfo {

package/src/services/plugin/index.ts

@@ -1,13 +1,3 @@
-import { isDesktop } from '@/const/version';
-
-import { ClientService as DeprecatedService } from './_deprecated';
-import { ClientService } from './client';
 import { ServerService } from './server';
 
-const clientService =
-  process.env.NEXT_PUBLIC_CLIENT_DB === 'pglite' ? new ClientService() : new DeprecatedService();
-
-export const pluginService =
-  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' || isDesktop
-    ? new ServerService()
-    : clientService;
+export const pluginService = new ServerService();

package/src/services/session/index.ts

@@ -1,13 +1,3 @@
-import { isDesktop } from '@/const/version';
-
-import { ClientService as DeprecatedService } from './_deprecated';
-import { ClientService } from './client';
 import { ServerService } from './server';
 
-const clientService =
-  process.env.NEXT_PUBLIC_CLIENT_DB === 'pglite' ? new ClientService() : new DeprecatedService();
-
-export const sessionService =
-  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' || isDesktop
-    ? new ServerService()
-    : clientService;
+export const sessionService = new ServerService();

package/src/services/tableViewer/client.ts

@@ -1,16 +1,13 @@
-import { clientDB } from '@/database/client/db';
-import { TableViewerRepo } from '@/database/repositories/tableViewer';
-import { BaseClientService } from '@/services/baseClientService';
-
-export class ClientService extends BaseClientService {
-  private get tableViewerRepo(): TableViewerRepo {
-    return new TableViewerRepo(clientDB as any, this.userId);
-  }
-
-  getAllTables = async () => this.tableViewerRepo.getAllTables();
-
-  getTableDetails = async (tableName: string) => this.tableViewerRepo.getTableDetails(tableName);
-
-  getTableData = async (tableName: string) =>
-    this.tableViewerRepo.getTableData(tableName, { page: 1, pageSize: 300 });
+export class ClientService {
+  getAllTables = async () => [];
+
+  getTableDetails = async (tableName: string) => {
+    console.log('getTableDetails:', tableName);
+    return [];
+  };
+
+  getTableData = async (tableName: string) => {
+    console.log('getTableData:', tableName);
+    return [];
+  };
 }

package/src/services/thread/index.ts

@@ -1,9 +1,3 @@
-import { isDesktop } from '@/const/version';
-
-import { ClientService } from './client';
 import { ServerService } from './server';
 
-export const threadService =
-  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' || isDesktop
-    ? new ServerService()
-    : new ClientService();
+export const threadService = new ServerService();

package/src/services/topic/index.ts

@@ -1,13 +1,3 @@
-import { isDesktop } from '@/const/version';
-
-import { ClientService as DeprecatedService } from './_deprecated';
-import { ClientService } from './client';
 import { ServerService } from './server';
 
-const clientService =
-  process.env.NEXT_PUBLIC_CLIENT_DB === 'pglite' ? new ClientService() : new DeprecatedService();
-
-export const topicService =
-  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' || isDesktop
-    ? new ServerService()
-    : clientService;
+export const topicService = new ServerService();

package/src/services/user/index.ts

@@ -1,15 +1,3 @@
-import { isDesktop } from '@/const/version';
-
-import { ClientService as DeprecatedService } from './_deprecated';
-import { ClientService } from './client';
 import { ServerService } from './server';
 
-const clientService =
-  process.env.NEXT_PUBLIC_CLIENT_DB === 'pglite' ? new ClientService() : new DeprecatedService();
-
-export const userService =
-  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' || isDesktop
-    ? new ServerService()
-    : clientService;
-
-export const userClientService = clientService;
+export const userService = new ServerService();

package/src/store/chat/helpers.test.ts

@@ -64,6 +64,105 @@ describe('chatHelpers', () => {
       const message = chatHelpers.getMessageById([], '1');
       expect(message).toBeUndefined();
     });
+
+    it('finds a block within a group message', () => {
+      const messagesWithGroup = [
+        { id: '1', content: 'Hello' },
+        {
+          id: 'group1',
+          role: 'group',
+          content: '',
+          children: [
+            { id: 'block1', content: 'First block' },
+            { id: 'block2', content: 'Second block' },
+          ],
+        },
+      ] as UIChatMessage[];
+
+      const block = chatHelpers.getMessageById(messagesWithGroup, 'block1');
+      expect(block).toBeDefined();
+      expect(block?.id).toBe('block1');
+      expect(block?.content).toBe('First block');
+    });
+
+    it('returns block with parentId set to group message id', () => {
+      const messagesWithGroup = [
+        {
+          id: 'group1',
+          role: 'group',
+          content: '',
+          children: [{ id: 'block1', content: 'Block content' }],
+        },
+      ] as UIChatMessage[];
+
+      const block = chatHelpers.getMessageById(messagesWithGroup, 'block1');
+      expect(block).toBeDefined();
+      expect(block?.parentId).toBe('group1');
+    });
+
+    it('searches across multiple group messages', () => {
+      const messagesWithGroups = [
+        {
+          id: 'group1',
+          role: 'group',
+          content: '',
+          children: [{ id: 'block1', content: 'First group block' }],
+        },
+        {
+          id: 'group2',
+          role: 'group',
+          content: '',
+          children: [{ id: 'block2', content: 'Second group block' }],
+        },
+      ] as UIChatMessage[];
+
+      const block = chatHelpers.getMessageById(messagesWithGroups, 'block2');
+      expect(block).toBeDefined();
+      expect(block?.id).toBe('block2');
+      expect(block?.parentId).toBe('group2');
+      expect(block?.content).toBe('Second group block');
+    });
+
+    it('prioritizes top-level message over block with same id', () => {
+      const messagesWithConflict = [
+        { id: 'duplicate', content: 'Top-level message', role: 'user' },
+        {
+          id: 'group1',
+          role: 'group',
+          content: '',
+          children: [{ id: 'duplicate', content: 'Block message' }],
+        },
+      ] as UIChatMessage[];
+
+      const message = chatHelpers.getMessageById(messagesWithConflict, 'duplicate');
+      expect(message).toBeDefined();
+      expect(message?.content).toBe('Top-level message');
+      expect(message?.role).toBe('user');
+      expect(message?.parentId).toBeUndefined();
+    });
+
+    it('returns undefined when block is not found in any group', () => {
+      const messagesWithGroup = [
+        {
+          id: 'group1',
+          role: 'group',
+          content: '',
+          children: [{ id: 'block1', content: 'Block content' }],
+        },
+      ] as UIChatMessage[];
+
+      const block = chatHelpers.getMessageById(messagesWithGroup, 'nonexistent');
+      expect(block).toBeUndefined();
+    });
+
+    it('handles group message without children', () => {
+      const messagesWithEmptyGroup = [
+        { id: 'group1', role: 'group', content: '' },
+      ] as UIChatMessage[];
+
+      const block = chatHelpers.getMessageById(messagesWithEmptyGroup, 'block1');
+      expect(block).toBeUndefined();
+    });
   });
 
   describe('getSlicedMessages', () => {

package/src/store/chat/helpers.ts

@@ -5,8 +5,27 @@ import { encodeAsync } from '@/utils/tokenizer';
 export const getMessagesTokenCount = async (messages: OpenAIChatMessage[]) =>
   encodeAsync(messages.map((m) => m.content).join(''));
 
-export const getMessageById = (messages: UIChatMessage[], id: string) =>
-  messages.find((m) => m.id === id);
+export const getMessageById = (
+  messages: UIChatMessage[],
+  id: string,
+): UIChatMessage | undefined => {
+  // First try to find in top-level messages
+  const directMatch = messages.find((m) => m.id === id);
+  if (directMatch) return directMatch;
+
+  // If not found, search in group message children (blocks)
+  for (const message of messages) {
+    if (message.role === 'group' && message.children) {
+      const blockMatch = message.children.find((block) => block.id === id);
+      if (blockMatch) {
+        // Return the block with parentId set to group message ID
+        return { ...blockMatch, parentId: message.id } as UIChatMessage;
+      }
+    }
+  }
+
+  return undefined;
+};
 
 const getSlicedMessages = (
   messages: UIChatMessage[],

package/src/store/chat/selectors.ts

@@ -1,6 +1,6 @@
 export { aiChatSelectors } from './slices/aiChat/selectors';
 export { chatToolSelectors } from './slices/builtinTool/selectors';
-export { chatSelectors } from './slices/message/selectors';
+export * from './slices/message/selectors';
 export * from './slices/portal/selectors';
 export { threadSelectors } from './slices/thread/selectors';
 export { topicSelectors } from './slices/topic/selectors';

package/src/store/chat/slices/aiChat/actions/__tests__/generateAIChat.test.ts

@@ -87,57 +87,6 @@ describe('chatMessage actions', () => {
     });
 
     describe('message creation', () => {
-      it('should create user message and trigger AI processing', async () => {
-        const { result } = renderHook(() => useChatStore());
-
-        await act(async () => {
-          await result.current.sendMessage({ message: TEST_CONTENT.USER_MESSAGE });
-        });
-
-        expect(messageService.createMessage).toHaveBeenCalledWith({
-          content: TEST_CONTENT.USER_MESSAGE,
-          files: undefined,
-          role: 'user',
-          sessionId: TEST_IDS.SESSION_ID,
-          topicId: TEST_IDS.TOPIC_ID,
-        });
-        expect(result.current.internal_coreProcessMessage).toHaveBeenCalled();
-      });
-
-      it('should send message with files attached', async () => {
-        const { result } = renderHook(() => useChatStore());
-        const files = [{ id: TEST_IDS.FILE_ID } as UploadFileItem];
-
-        await act(async () => {
-          await result.current.sendMessage({ message: TEST_CONTENT.USER_MESSAGE, files });
-        });
-
-        expect(messageService.createMessage).toHaveBeenCalledWith({
-          content: TEST_CONTENT.USER_MESSAGE,
-          files: [TEST_IDS.FILE_ID],
-          role: 'user',
-          sessionId: TEST_IDS.SESSION_ID,
-          topicId: TEST_IDS.TOPIC_ID,
-        });
-      });
-
-      it('should send files without message content', async () => {
-        const { result } = renderHook(() => useChatStore());
-        const files = [{ id: TEST_IDS.FILE_ID } as UploadFileItem];
-
-        await act(async () => {
-          await result.current.sendMessage({ message: TEST_CONTENT.EMPTY, files });
-        });
-
-        expect(messageService.createMessage).toHaveBeenCalledWith({
-          content: TEST_CONTENT.EMPTY,
-          files: [TEST_IDS.FILE_ID],
-          role: 'user',
-          sessionId: TEST_IDS.SESSION_ID,
-          topicId: TEST_IDS.TOPIC_ID,
-        });
-      });
-
       it('should not process AI when onlyAddUserMessage is true', async () => {
         const { result } = renderHook(() => useChatStore());
 
@@ -169,196 +118,6 @@ describe('chatMessage actions', () => {
         expect(result.current.internal_coreProcessMessage).not.toHaveBeenCalled();
       });
     });
-
-    describe('auto-create topic', () => {
-      const TOPIC_THRESHOLD = 5;
-
-      it('should create topic when threshold is reached and feature is enabled', async () => {
-        const { result } = renderHook(() => useChatStore());
-
-        const createTopicMock = vi.fn(() => Promise.resolve(TEST_IDS.NEW_TOPIC_ID));
-        const switchTopicMock = vi.fn();
-
-        act(() => {
-          setupMockSelectors({
-            agentConfig: {
-              enableAutoCreateTopic: true,
-              autoCreateTopicThreshold: TOPIC_THRESHOLD,
-            },
-          });
-
-          useChatStore.setState({
-            activeTopicId: undefined,
-            messagesMap: {
-              [messageMapKey(TEST_IDS.SESSION_ID)]: createMockMessages(TOPIC_THRESHOLD),
-            },
-            createTopic: createTopicMock,
-            switchTopic: switchTopicMock,
-          });
-        });
-
-        await act(async () => {
-          await result.current.sendMessage({ message: TEST_CONTENT.USER_MESSAGE });
-        });
-
-        expect(createTopicMock).toHaveBeenCalled();
-        expect(switchTopicMock).toHaveBeenCalledWith(TEST_IDS.NEW_TOPIC_ID, true);
-      });
-    });
-
-    describe('RAG integration', () => {
-      it('should include RAG query when RAG is enabled', async () => {
-        const { result } = renderHook(() => useChatStore());
-        vi.spyOn(result.current, 'internal_shouldUseRAG').mockReturnValue(true);
-
-        await act(async () => {
-          await result.current.sendMessage({ message: TEST_CONTENT.RAG_QUERY });
-        });
-
-        expect(result.current.internal_coreProcessMessage).toHaveBeenCalledWith(
-          expect.any(Array),
-          expect.any(String),
-          expect.objectContaining({
-            ragQuery: TEST_CONTENT.RAG_QUERY,
-          }),
-        );
-      });
-
-      it('should not use RAG when feature is disabled', async () => {
-        const { result } = renderHook(() => useChatStore());
-        vi.spyOn(result.current, 'internal_shouldUseRAG').mockReturnValue(false);
-        const retrieveChunksSpy = vi.spyOn(result.current, 'internal_retrieveChunks');
-
-        await act(async () => {
-          await result.current.sendMessage({ message: TEST_CONTENT.USER_MESSAGE });
-        });
-
-        expect(retrieveChunksSpy).not.toHaveBeenCalled();
-        expect(result.current.internal_coreProcessMessage).toHaveBeenCalledWith(
-          expect.any(Array),
-          expect.any(String),
-          expect.not.objectContaining({
-            ragQuery: expect.anything(),
-          }),
-        );
-      });
-    });
-
-    describe('special flags', () => {
-      it('should pass isWelcomeQuestion flag to processing', async () => {
-        const { result } = renderHook(() => useChatStore());
-
-        await act(async () => {
-          await result.current.sendMessage({
-            message: TEST_CONTENT.USER_MESSAGE,
-            isWelcomeQuestion: true,
-          });
-        });
-
-        expect(result.current.internal_coreProcessMessage).toHaveBeenCalledWith(
-          expect.anything(),
-          expect.anything(),
-          { isWelcomeQuestion: true },
-        );
-      });
-
-      it('should return early when onlyAddUserMessage is true', async () => {
-        const { result } = renderHook(() => useChatStore());
-
-        await act(async () => {
-          await result.current.sendMessage({
-            message: TEST_CONTENT.USER_MESSAGE,
-            onlyAddUserMessage: true,
-          });
-        });
-
-        expect(result.current.internal_coreProcessMessage).not.toHaveBeenCalled();
-      });
-    });
-
-    describe('topic creation flow', () => {
-      it('should handle tempMessage during topic creation', async () => {
-        setupMockSelectors({
-          chatConfig: { enableAutoCreateTopic: true, autoCreateTopicThreshold: 2 },
-        });
-
-        act(() => {
-          useChatStore.setState({ activeTopicId: undefined });
-          setupStoreWithMessages(createMockMessages(5));
-        });
-
-        const { result } = renderHook(() => useChatStore());
-        const createTopicMock = vi
-          .spyOn(result.current, 'createTopic')
-          .mockResolvedValue(TEST_IDS.NEW_TOPIC_ID);
-        const toggleMessageLoadingSpy = vi.spyOn(result.current, 'internal_toggleMessageLoading');
-        const createTmpMessageSpy = vi
-          .spyOn(result.current, 'internal_createTmpMessage')
-          .mockReturnValue('temp-id');
-        vi.spyOn(result.current, 'internal_fetchMessages').mockResolvedValue();
-        vi.spyOn(result.current, 'switchTopic').mockResolvedValue();
-
-        await act(async () => {
-          await result.current.sendMessage({ message: TEST_CONTENT.USER_MESSAGE });
-        });
-
-        expect(createTmpMessageSpy).toHaveBeenCalled();
-        expect(toggleMessageLoadingSpy).toHaveBeenCalledWith(true, 'temp-id');
-        expect(createTopicMock).toHaveBeenCalled();
-      });
-
-      it('should call summaryTopicTitle after processing when new topic created', async () => {
-        setupMockSelectors({
-          chatConfig: { enableAutoCreateTopic: true, autoCreateTopicThreshold: 2 },
-        });
-
-        act(() => {
-          useChatStore.setState({ activeTopicId: undefined });
-          setupStoreWithMessages(createMockMessages(5));
-        });
-
-        const { result } = renderHook(() => useChatStore());
-        vi.spyOn(result.current, 'createTopic').mockResolvedValue(TEST_IDS.NEW_TOPIC_ID);
-        vi.spyOn(result.current, 'internal_createTmpMessage').mockReturnValue('temp-id');
-        vi.spyOn(result.current, 'internal_fetchMessages').mockResolvedValue();
-        vi.spyOn(result.current, 'switchTopic').mockResolvedValue();
-
-        const summaryTopicTitleSpy = vi
-          .spyOn(result.current, 'summaryTopicTitle')
-          .mockResolvedValue();
-
-        await act(async () => {
-          await result.current.sendMessage({ message: TEST_CONTENT.USER_MESSAGE });
-        });
-
-        expect(summaryTopicTitleSpy).toHaveBeenCalledWith(TEST_IDS.NEW_TOPIC_ID, expect.any(Array));
-      });
-
-      it('should handle topic creation failure gracefully', async () => {
-        setupMockSelectors({
-          chatConfig: { enableAutoCreateTopic: true, autoCreateTopicThreshold: 2 },
-        });
-
-        act(() => {
-          useChatStore.setState({ activeTopicId: undefined });
-          setupStoreWithMessages(createMockMessages(5));
-        });
-
-        const { result } = renderHook(() => useChatStore());
-        vi.spyOn(result.current, 'createTopic').mockResolvedValue(undefined);
-        vi.spyOn(result.current, 'internal_createTmpMessage').mockReturnValue('temp-id');
-        const toggleLoadingSpy = vi.spyOn(result.current, 'internal_toggleMessageLoading');
-        const updateTopicLoadingSpy = vi.spyOn(result.current, 'internal_updateTopicLoading');
-
-        await act(async () => {
-          await result.current.sendMessage({ message: TEST_CONTENT.USER_MESSAGE });
-        });
-
-        // Should still call the AI processing even if topic creation fails
-        expect(result.current.internal_coreProcessMessage).toHaveBeenCalled();
-        expect(updateTopicLoadingSpy).not.toHaveBeenCalled();
-      });
-    });
   });
 
   describe('regenerateMessage', () => {

package/src/store/chat/slices/aiChat/actions/__tests__/generateAIChatV2.test.ts

@@ -19,8 +19,33 @@ import { resetTestEnvironment, setupMockSelectors, spyOnMessageService } from '.
 // Keep zustand mock as it's needed globally
 vi.mock('zustand/traditional');
 
+// Mock AntdStaticMethods
+vi.mock('@/components/AntdStaticMethods', () => ({
+  notification: {
+    error: vi.fn(),
+    success: vi.fn(),
+    info: vi.fn(),
+    warning: vi.fn(),
+  },
+  message: {
+    error: vi.fn(),
+    success: vi.fn(),
+    info: vi.fn(),
+    warning: vi.fn(),
+  },
+}));
+
+// Mock sessionService to prevent TRPC requests
+vi.mock('@/services/session', () => ({
+  sessionService: {
+    updateSession: vi.fn(),
+    updateSessionConfig: vi.fn(),
+    updateSessionChatConfig: vi.fn(),
+  },
+}));
+
 // Mock server mode for V2 tests
-vi.mock('@/const/version', async (importOriginal) => {
+vi.mock('@lobechat/const', async (importOriginal) => {
   const module = await importOriginal();
   return {
     ...(module as any),

package/src/store/chat/slices/aiChat/actions/__tests__/helpers.ts

@@ -60,7 +60,9 @@ export const spyOnMessageService = () => {
   const createMessageSpy = vi
     .spyOn(messageService, 'createMessage')
     .mockResolvedValue(TEST_IDS.NEW_MESSAGE_ID);
-  const updateMessageSpy = vi.spyOn(messageService, 'updateMessage').mockResolvedValue(undefined);
+  const updateMessageSpy = vi
+    .spyOn(messageService, 'updateMessage')
+    .mockResolvedValue({ messages: [], success: true });
   const removeMessageSpy = vi.spyOn(messageService, 'removeMessage').mockResolvedValue(undefined);
   const updateMessageErrorSpy = vi
     .spyOn(messageService, 'updateMessageError')