@lobehub/lobehub 2.0.0-next.8 → 2.0.0-next.9

package/src/features/Conversation/components/Extras/Usage/UsageDetail/index.tsx

@@ -11,7 +11,7 @@ import InfoTooltip from '@/components/InfoTooltip';
 import { aiModelSelectors, useAiInfraStore } from '@/store/aiInfra';
 import { useGlobalStore } from '@/store/global';
 import { systemStatusSelectors } from '@/store/global/selectors';
-import { formatNumber } from '@/utils/format';
+import { formatNumber, formatShortenNumber } from '@/utils/format';
 
 import ModelCard from './ModelCard';
 import TokenProgress, { TokenProgressItem } from './TokenProgress';
@@ -111,10 +111,13 @@ const TokenDetail = memo<TokenDetailProps>(({ meta, model, provider }) => {
     },
   ].filter(Boolean) as TokenProgressItem[];
 
-  const displayTotal =
+  const totalCount =
     isShowCredit && !!detailTokens.totalTokens
-      ? formatNumber(detailTokens.totalTokens.credit)
-      : formatNumber(detailTokens.totalTokens!.token);
+      ? detailTokens.totalTokens.credit
+      : detailTokens.totalTokens!.token;
+
+  const shortTotal = (formatShortenNumber(totalCount) as string).toLowerCase?.();
+  const detailTotal = formatNumber(totalCount);
 
   const averagePricing = formatNumber(
     detailTokens.totalTokens!.credit / detailTokens.totalTokens!.token,
@@ -171,7 +174,7 @@ const TokenDetail = memo<TokenDetailProps>(({ meta, model, provider }) => {
                 <div style={{ color: theme.colorTextSecondary }}>
                   {t('messages.tokenDetails.total')}
                 </div>
-                <div style={{ fontWeight: 500 }}>{displayTotal}</div>
+                <div style={{ fontWeight: 500 }}>{detailTotal}</div>
               </Flexbox>
               {isShowCredit && (
                 <Flexbox align={'center'} gap={4} horizontal justify={'space-between'}>
@@ -212,7 +215,7 @@ const TokenDetail = memo<TokenDetailProps>(({ meta, model, provider }) => {
     >
       <Center gap={2} horizontal style={{ cursor: 'default' }}>
         <Icon icon={isShowCredit ? BadgeCent : CoinsIcon} />
-        {displayTotal}
+        {shortTotal}
       </Center>
     </Popover>
   );

package/src/features/PluginTag/index.tsx

@@ -7,7 +7,6 @@ import { memo } from 'react';
 import { Center } from 'react-layout-kit';
 
 import Avatar from '@/components/Plugins/PluginAvatar';
-import { featureFlagsSelectors, useServerConfigStore } from '@/store/serverConfig';
 import { pluginHelpers, useToolStore } from '@/store/tool';
 import { toolSelectors } from '@/store/tool/selectors';
 
@@ -18,8 +17,7 @@ export interface PluginTagProps {
 }
 
 const PluginTag = memo<PluginTagProps>(({ plugins }) => {
-  const { showDalle } = useServerConfigStore(featureFlagsSelectors);
-  const list = useToolStore(toolSelectors.metaList(showDalle), isEqual);
+  const list = useToolStore(toolSelectors.metaList, isEqual);
 
   const displayPlugin = useToolStore(toolSelectors.getMetaById(plugins[0]), isEqual);
 

package/src/features/PluginsUI/Render/BuiltinType/index.test.tsx

@@ -1,22 +1,19 @@
 import { render, screen } from '@testing-library/react';
 import { describe, expect, it, vi } from 'vitest';
 
-import { DalleManifest } from '@/tools/dalle';
-import { BuiltinToolsRenders } from '@/tools/renders';
-
 import BuiltinType from './index';
 
-// Mock Render component and useParseContent hook
+// Mock renders module
 vi.mock('@/tools/renders', () => ({
   BuiltinToolsRenders: {
-    dalle3: vi.fn(() => <div>Test Renderer</div>),
-    [DalleManifest.identifier]: vi.fn(() => <div>{DalleManifest.identifier}</div>),
+    'lobe-web-browsing': vi.fn(({ content }) => <div>WebBrowsingRender: {content}</div>),
+    'lobe-code-interpreter': vi.fn(({ content }) => <div>CodeInterpreterRender: {content}</div>),
   },
 }));
 
-// Mock Loading component
-vi.mock('../Loading', () => ({
-  default: vi.fn(() => <div>Loading...</div>),
+// Mock useParseContent hook
+vi.mock('../useParseContent', () => ({
+  useParseContent: vi.fn((content) => ({ data: content })),
 }));
 
 describe('BuiltinType', () => {
@@ -25,29 +22,41 @@ describe('BuiltinType', () => {
     expect(container).toBeEmptyDOMElement();
   });
 
-  it('should not render anything if content is not JSON and no identifier', () => {
-    const { container } = render(<BuiltinType content="..." id="123" />);
-    expect(container).toBeEmptyDOMElement();
-  });
-
   it('should not render anything if identifier is unknown', () => {
     const { container } = render(<BuiltinType content="{}" id="123" identifier="unknown" />);
     expect(container).toBeEmptyDOMElement();
   });
 
-  describe('DALL·E', () => {
-    it('should render the correct renderer if identifier is dalle3', () => {
-      render(<BuiltinType content='{"some":"data"}' id="123" identifier="dalle3" />);
-      expect(BuiltinToolsRenders.dalle3).toHaveBeenCalled();
-      expect(screen.getByText('Test Renderer')).toBeInTheDocument();
-    });
-
-    it('should render the correct renderer if is DALL·E ', () => {
-      render(
-        <BuiltinType content='{"some":"data"}' id="123" identifier={DalleManifest.identifier} />,
-      );
-      expect(BuiltinToolsRenders.dalle3).toHaveBeenCalled();
-      expect(screen.getByText(DalleManifest.identifier)).toBeInTheDocument();
-    });
+  it('should render the correct renderer for web browsing', () => {
+    const content = '{"query":"test"}';
+    render(<BuiltinType content={content} id="123" identifier="lobe-web-browsing" />);
+    expect(screen.getByText(`WebBrowsingRender: ${content}`)).toBeInTheDocument();
+  });
+
+  it('should render the correct renderer for code interpreter', () => {
+    const content = '{"code":"print(1)"}';
+    render(<BuiltinType content={content} id="123" identifier="lobe-code-interpreter" />);
+    expect(screen.getByText(`CodeInterpreterRender: ${content}`)).toBeInTheDocument();
+  });
+
+  it('should pass correct props to renderer', () => {
+    const content = '{"test":"data"}';
+    const args = '{"arg":"value"}';
+    const pluginState = { state: 'value' };
+    const pluginError = { error: 'test' };
+
+    render(
+      <BuiltinType
+        content={content}
+        id="msg-123"
+        identifier="lobe-web-browsing"
+        arguments={args}
+        pluginState={pluginState}
+        pluginError={pluginError}
+        apiName="testApi"
+      />,
+    );
+
+    expect(screen.getByText(`WebBrowsingRender: ${content}`)).toBeInTheDocument();
   });
 });

package/src/features/Portal/Artifacts/Body/index.tsx

@@ -3,7 +3,7 @@ import { memo, useEffect, useMemo } from 'react';
 import { Flexbox } from 'react-layout-kit';
 
 import { useChatStore } from '@/store/chat';
-import { chatPortalSelectors, chatSelectors } from '@/store/chat/selectors';
+import { chatPortalSelectors, messageStateSelectors } from '@/store/chat/selectors';
 import { ArtifactDisplayMode } from '@/store/chat/slices/portal/initialState';
 import { ArtifactType } from '@/types/artifact';
 
@@ -24,7 +24,7 @@ const ArtifactsUI = memo(() => {
     return [
       messageId,
       s.portalArtifactDisplayMode,
-      chatSelectors.isMessageGenerating(messageId)(s),
+      messageStateSelectors.isMessageGenerating(messageId)(s),
       chatPortalSelectors.artifactType(s),
       chatPortalSelectors.artifactCode(messageId)(s),
       chatPortalSelectors.artifactCodeLanguage(s),

package/src/server/modules/ModelRuntime/trace.ts

@@ -42,6 +42,16 @@ export const createTraceOptions = (
     startTime: new Date(),
   });
 
+  const headers = new Headers();
+
+  if (trace?.id) {
+    headers.set(LOBE_CHAT_TRACE_ID, trace.id);
+  }
+
+  if (generation?.id) {
+    headers.set(LOBE_CHAT_OBSERVATION_ID, generation.id);
+  }
+
   return {
     callback: {
       onCompletion: async ({ text, thinking, usage, grounding, toolsCalling }) => {
@@ -94,9 +104,6 @@ export const createTraceOptions = (
         });
       },
     } as ChatStreamCallbacks,
-    headers: {
-      [LOBE_CHAT_OBSERVATION_ID]: generation?.id,
-      [LOBE_CHAT_TRACE_ID]: trace?.id,
-    },
+    headers: headers,
   };
 };

package/src/server/routers/lambda/message.ts

@@ -1,4 +1,9 @@
-import { BatchTaskResult, UIChatMessage, UpdateMessageRAGParamsSchema } from '@lobechat/types';
+import {
+  BatchTaskResult,
+  UIChatMessage,
+  UpdateMessageParamsSchema,
+  UpdateMessageRAGParamsSchema,
+} from '@lobechat/types';
 import { z } from 'zod';
 
 import { MessageModel } from '@/database/models/message';
@@ -180,11 +185,17 @@ export const messageRouter = router({
     .input(
       z.object({
         id: z.string(),
-        value: z.object({}).passthrough().partial(),
+        sessionId: z.string().nullable().optional(),
+        topicId: z.string().nullable().optional(),
+        value: UpdateMessageParamsSchema,
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      return ctx.messageModel.update(input.id, input.value);
+      return ctx.messageModel.update(input.id, input.value as any, {
+        postProcessUrl: (path) => ctx.fileService.getFullFileUrl(path),
+        sessionId: input.sessionId,
+        topicId: input.topicId,
+      });
     }),
 
   updateMessagePlugin: messageProcedure

package/src/services/chat/chat.test.ts

@@ -8,18 +8,12 @@ import { type Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vite
 
 import { DEFAULT_USER_AVATAR } from '@/const/meta';
 import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
-import * as isCanUseFCModule from '@/helpers/isCanUseFC';
 import * as toolEngineeringModule from '@/helpers/toolEngineering';
-import { agentChatConfigSelectors, agentSelectors } from '@/store/agent/selectors';
+import { agentChatConfigSelectors } from '@/store/agent/selectors';
 import { aiModelSelectors } from '@/store/aiInfra';
 import { useToolStore } from '@/store/tool';
-import { toolSelectors } from '@/store/tool/selectors';
-import { modelProviderSelectors } from '@/store/user/selectors';
-import { DalleManifest } from '@/tools/dalle';
 import { WebBrowsingManifest } from '@/tools/web-browsing';
 
-import { API_ENDPOINTS } from '../_url';
-import * as helpers from './helper';
 import { chatService } from './index';
 
 // Mocking external dependencies
@@ -831,39 +825,6 @@ describe('ChatService', () => {
           undefined,
         );
       });
-
-      it('work with dalle3', async () => {
-        const getChatCompletionSpy = vi.spyOn(chatService, 'getChatCompletion');
-        const messages = [
-          {
-            role: 'user',
-            content: 'https://vercel.com/ 请分析 chatGPT 关键词\n\n',
-            sessionId: 'inbox',
-            createdAt: 1702723964330,
-            id: 'vyQvEw6V',
-            updatedAt: 1702723964330,
-            extra: {},
-            meta: {
-              avatar: DEFAULT_USER_AVATAR,
-            },
-          },
-        ] as UIChatMessage[];
-
-        await chatService.createAssistantMessage({
-          messages,
-          model: 'gpt-3.5-turbo-1106',
-          top_p: 1,
-          plugins: [DalleManifest.identifier],
-        });
-
-        // Assert that getChatCompletionSpy was called with the expected arguments
-        expect(getChatCompletionSpy).toHaveBeenCalled();
-
-        const calls = getChatCompletionSpy.mock.lastCall;
-        // Take a snapshot of the first call's first argument
-        expect(calls![0]).toMatchSnapshot();
-        expect(calls![1]).toBeUndefined();
-      });
     });
 
     describe('search functionality', () => {

package/src/services/chat/contextEngineering.test.ts

@@ -248,36 +248,6 @@ describe('contextEngineering', () => {
     ]);
   });
 
-  it('should inject INBOX_GUIDE_SYSTEM_ROLE for welcome questions in inbox session', async () => {
-    // Don't mock INBOX_GUIDE_SYSTEMROLE, use the real one
-    const messages: UIChatMessage[] = [
-      {
-        role: 'user',
-        content: 'Hello, this is my first question',
-        createdAt: Date.now(),
-        id: 'test-welcome',
-        meta: {},
-        updatedAt: Date.now(),
-      },
-    ];
-
-    const result = await contextEngineering({
-      messages,
-      model: 'gpt-4',
-      provider: 'openai',
-      isWelcomeQuestion: true,
-      sessionId: 'inbox',
-    });
-
-    // Should have system message with inbox guide content
-    const systemMessage = result.find((msg) => msg.role === 'system');
-    expect(systemMessage).toBeDefined();
-    // Check for characteristic content of the actual INBOX_GUIDE_SYSTEMROLE
-    expect(systemMessage!.content).toContain('LobeChat Support Assistant');
-    expect(systemMessage!.content).toContain('LobeHub');
-    expect(Object.keys(systemMessage!).length).toEqual(2);
-  });
-
   it('should inject historySummary into system message when provided', async () => {
     const historySummary = 'Previous conversation summary: User discussed AI topics.';
 

package/src/services/chat/contextEngineering.ts

@@ -1,9 +1,8 @@
-import { INBOX_GUIDE_SYSTEMROLE, INBOX_SESSION_ID, isDesktop, isServerMode } from '@lobechat/const';
+import { isDesktop, isServerMode } from '@lobechat/const';
 import {
   ContextEngine,
   HistorySummaryProvider,
   HistoryTruncateProcessor,
-  InboxGuideProvider,
   InputTemplateProcessor,
   MessageCleanupProcessor,
   MessageContentProcessor,
@@ -48,8 +47,6 @@ export const contextEngineering = async ({
   enableHistoryCount,
   historyCount,
   historySummary,
-  sessionId,
-  isWelcomeQuestion,
 }: ContextEngineeringContext): Promise<OpenAIChatMessage[]> => {
   const toolNameResolver = new ToolNameResolver();
 
@@ -63,14 +60,6 @@ export const contextEngineering = async ({
       // 2. System role injection (agent's system role)
       new SystemRoleInjector({ systemRole }),
 
-      // 3. Inbox guide system role injection
-      new InboxGuideProvider({
-        inboxGuideSystemRole: INBOX_GUIDE_SYSTEMROLE,
-        inboxSessionId: INBOX_SESSION_ID,
-        isWelcomeQuestion: isWelcomeQuestion,
-        sessionId: sessionId,
-      }),
-
       // 4. Tool system role injection
       new ToolSystemRoleProvider({
         getToolSystemRoles: (tools) => toolSelectors.enabledSystemRoles(tools)(getToolStoreState()),

package/src/services/chat/index.ts

@@ -66,7 +66,6 @@ interface FetchAITaskResultParams extends FetchSSEOptions {
 interface CreateAssistantMessageStream extends FetchSSEOptions {
   abortController?: AbortController;
   historySummary?: string;
-  isWelcomeQuestion?: boolean;
   params: GetChatCompletionPayload;
   trace?: TracePayload;
 }
@@ -113,9 +112,7 @@ class ChatService {
       enableHistoryCount: agentChatConfigSelectors.enableHistoryCount(agentStoreState),
       // include user messages
       historyCount: agentChatConfigSelectors.historyCount(agentStoreState) + 2,
-      historySummary: options?.historySummary,
       inputTemplate: chatConfig.inputTemplate,
-      isWelcomeQuestion: options?.isWelcomeQuestion,
       messages,
       model: payload.model,
       provider: payload.provider!,
@@ -217,12 +214,10 @@ class ChatService {
     onErrorHandle,
     onFinish,
     trace,
-    isWelcomeQuestion,
     historySummary,
   }: CreateAssistantMessageStream) => {
     await this.createAssistantMessage(params, {
       historySummary,
-      isWelcomeQuestion,
       onAbort,
       onErrorHandle,
       onFinish,
@@ -404,7 +399,7 @@ class ChatService {
     onLoadingChange?.(true);
 
     try {
-      const oaiMessages = await contextEngineering({
+      const llmMessages = await contextEngineering({
         messages: params.messages as any,
         model: params.model!,
         provider: params.provider!,
@@ -421,7 +416,7 @@ class ChatService {
       // remove plugins
       delete params.plugins;
       await this.getChatCompletion(
-        { ...params, messages: oaiMessages, tools },
+        { ...params, messages: llmMessages, tools },
         {
           onErrorHandle: (error) => {
             errorHandle(new Error(error.message), error);

package/src/services/chat/types.ts

@@ -1,9 +1,9 @@
 import { TracePayload } from '@lobechat/types';
+
 import { FetchSSEOptions } from '@/utils/fetch';
 
 export interface FetchOptions extends FetchSSEOptions {
   historySummary?: string;
-  isWelcomeQuestion?: boolean;
   signal?: AbortSignal | undefined;
   trace?: TracePayload;
 }

package/src/services/message/_deprecated.ts

@@ -102,7 +102,7 @@ export class ClientService implements IMessageService {
   }
 
   // @ts-ignore
-  async updateMessage(id: string, message: Partial<DB_Message>) {
+  async updateMessage(id: string, message: Partial<DB_Message>): Promise<any> {
     return MessageModel.update(id, message);
   }
 

package/src/services/message/client.ts

@@ -45,6 +45,7 @@ export class ClientService extends BaseClientService implements IMessageService
         topicId,
       },
       {
+        groupAssistantMessages: false,
         postProcessUrl: this.postProcessUrl,
       },
     );
@@ -60,6 +61,7 @@ export class ClientService extends BaseClientService implements IMessageService
         topicId,
       },
       {
+        groupAssistantMessages: false,
         postProcessUrl: this.postProcessUrl,
       },
     );
@@ -99,8 +101,12 @@ export class ClientService extends BaseClientService implements IMessageService
     return this.messageModel.update(id, { error });
   };
 
-  updateMessage: IMessageService['updateMessage'] = async (id, message) => {
-    return this.messageModel.update(id, message);
+  updateMessage: IMessageService['updateMessage'] = async (id, message, options) => {
+    return this.messageModel.update(id, message, {
+      postProcessUrl: this.postProcessUrl,
+      sessionId: options?.sessionId,
+      topicId: options?.topicId,
+    });
   };
 
   updateMessageTTS: IMessageService['updateMessageTTS'] = async (id, tts) => {

package/src/services/message/server.ts

@@ -81,8 +81,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

@@ -10,6 +10,7 @@ import {
   UIChatMessage,
   UpdateMessageParams,
   UpdateMessageRAGParams,
+  UpdateMessageResult,
 } from '@lobechat/types';
 import type { HeatmapsProps } from '@lobehub/charts';
 
@@ -37,7 +38,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/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';