@lobehub/lobehub 2.0.0-next.1 → 2.0.0-next.10

package/src/features/Conversation/Messages/User/index.tsx

@@ -15,7 +15,7 @@ import { useUserAvatar } from '@/hooks/useUserAvatar';
 import { useAgentStore } from '@/store/agent';
 import { agentChatConfigSelectors } from '@/store/agent/selectors';
 import { useChatStore } from '@/store/chat';
-import { chatSelectors } from '@/store/chat/selectors';
+import { messageStateSelectors } from '@/store/chat/selectors';
 import { useSessionStore } from '@/store/session';
 import { sessionSelectors } from '@/store/session/selectors';
 import { useUserStore } from '@/store/user';
@@ -55,9 +55,9 @@ const UserMessage = memo<UserMessageProps>((props) => {
   const displayMode = useAgentStore(agentChatConfigSelectors.displayMode);
 
   const [editing, generating, isInRAGFlow] = useChatStore((s) => [
-    chatSelectors.isMessageEditing(id)(s),
-    chatSelectors.isMessageGenerating(id)(s),
-    chatSelectors.isMessageInRAGFlow(id)(s),
+    messageStateSelectors.isMessageEditing(id)(s),
+    messageStateSelectors.isMessageGenerating(id)(s),
+    messageStateSelectors.isMessageInRAGFlow(id)(s),
   ]);
 
   const loading = isInRAGFlow || generating;

package/src/features/Conversation/Messages/index.tsx

@@ -1,17 +1,17 @@
 'use client';
 
+import { isDesktop } from '@lobechat/const';
 import { createStyles } from 'antd-style';
 import isEqual from 'fast-deep-equal';
 import { ReactNode, memo, useCallback, useEffect, useMemo, useRef } from 'react';
 import { Flexbox } from 'react-layout-kit';
 
-import { isDesktop } from '@/const/version';
 import {
   removeVirtuosoVisibleItem,
   upsertVirtuosoVisibleItem,
 } from '@/features/Conversation/components/VirtualizedList/VirtuosoContext';
 import { useChatStore } from '@/store/chat';
-import { chatSelectors } from '@/store/chat/selectors';
+import { chatSelectors, messageStateSelectors } from '@/store/chat/selectors';
 
 import History from '../components/History';
 import { InPortalThreadContext } from '../context/InPortalThreadContext';
@@ -57,7 +57,7 @@ const Item = memo<ChatListItemProps>(
 
     const item = useChatStore(chatSelectors.getMessageById(id), isEqual);
 
-    const [isMessageLoading] = useChatStore((s) => [chatSelectors.isMessageLoading(id)(s)]);
+    const [isMessageLoading] = useChatStore((s) => [messageStateSelectors.isMessageLoading(id)(s)]);
 
     // ======================= Performance Optimization ======================= //
     // these useMemo/useCallback are all for the performance optimization

package/src/features/Conversation/components/AutoScroll.tsx

@@ -1,7 +1,7 @@
 import { memo, useEffect } from 'react';
 
 import { useChatStore } from '@/store/chat';
-import { chatSelectors } from '@/store/chat/selectors';
+import { chatSelectors, messageStateSelectors } from '@/store/chat/selectors';
 
 import BackBottom from './BackBottom';
 
@@ -11,7 +11,7 @@ interface AutoScrollProps {
   onScrollToBottom: (type: 'auto' | 'click') => void;
 }
 const AutoScroll = memo<AutoScrollProps>(({ atBottom, isScrolling, onScrollToBottom }) => {
-  const trackVisibility = useChatStore(chatSelectors.isAIGenerating);
+  const trackVisibility = useChatStore(messageStateSelectors.isAIGenerating);
   const str = useChatStore(chatSelectors.mainAIChatsMessageString);
   const reasoningStr = useChatStore(chatSelectors.mainAILatestMessageReasoningContent);
 

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/libs/next-auth/auth.config.ts

@@ -58,7 +58,7 @@ export default {
     signIn: '/next-auth/signin',
   },
   providers: initSSOProviders(),
-  secret: NEXT_AUTH_SECRET,
+  secret: NEXT_AUTH_SECRET ?? process.env.AUTH_SECRET,
   session: {
     // Force use JWT if server service is disabled
     strategy: NEXT_PUBLIC_ENABLED_SERVER_SERVICE ? NEXT_AUTH_SSO_SESSION_STRATEGY : 'jwt',

package/src/libs/next-auth/sso-providers/auth0.ts

@@ -1,7 +1,5 @@
 import Auth0 from 'next-auth/providers/auth0';
 
-import { authEnv } from '@/envs/auth';
-
 import { CommonProviderConfig } from './sso.config';
 
 const provider = {
@@ -11,11 +9,6 @@ const provider = {
     // Specify auth scope, at least include 'openid email'
     // all scopes in Auth0 ref: https://auth0.com/docs/get-started/apis/scopes/openid-connect-scopes#standard-claims
     authorization: { params: { scope: 'openid email profile' } },
-    // TODO(NextAuth ENVs Migration): Remove once nextauth envs migration time end
-    clientId: authEnv.AUTH0_CLIENT_ID ?? process.env.AUTH_AUTH0_ID,
-    clientSecret: authEnv.AUTH0_CLIENT_SECRET ?? process.env.AUTH_AUTH0_SECRET,
-    issuer: authEnv.AUTH0_ISSUER ?? process.env.AUTH_AUTH0_ISSUER,
-    // Remove End
     profile(profile) {
       return {
         email: profile.email,

package/src/libs/next-auth/sso-providers/authelia.ts

@@ -1,7 +1,5 @@
 import type { OIDCConfig } from '@auth/core/providers';
 
-import { authEnv } from '@/envs/auth';
-
 import { CommonProviderConfig } from './sso.config';
 
 export type AutheliaProfile = {
@@ -21,10 +19,10 @@ const provider = {
     ...CommonProviderConfig,
     authorization: { params: { scope: 'openid email profile' } },
     checks: ['state', 'pkce'],
-    clientId: authEnv.AUTHELIA_CLIENT_ID ?? process.env.AUTH_AUTHELIA_ID,
-    clientSecret: authEnv.AUTHELIA_CLIENT_SECRET ?? process.env.AUTH_AUTHELIA_SECRET,
+    clientId: process.env.AUTH_AUTHELIA_ID,
+    clientSecret: process.env.AUTH_AUTHELIA_SECRET,
     id: 'authelia',
-    issuer: authEnv.AUTHELIA_ISSUER ?? process.env.AUTH_AUTHELIA_ISSUER,
+    issuer: process.env.AUTH_AUTHELIA_ISSUER,
     name: 'Authelia',
     profile(profile) {
       return {

package/src/libs/next-auth/sso-providers/authentik.ts

@@ -1,7 +1,5 @@
 import Authentik from 'next-auth/providers/authentik';
 
-import { authEnv } from '@/envs/auth';
-
 import { CommonProviderConfig } from './sso.config';
 
 const provider = {
@@ -11,11 +9,6 @@ const provider = {
     // Specify auth scope, at least include 'openid email'
     // all scopes in Authentik ref: https://goauthentik.io/docs/providers/oauth2
     authorization: { params: { scope: 'openid email profile' } },
-    // TODO(NextAuth ENVs Migration): Remove once nextauth envs migration time end
-    clientId: authEnv.AUTHENTIK_CLIENT_ID ?? process.env.AUTH_AUTHENTIK_ID,
-    clientSecret: authEnv.AUTHENTIK_CLIENT_SECRET ?? process.env.AUTH_AUTHENTIK_SECRET,
-    issuer: authEnv.AUTHENTIK_ISSUER ?? process.env.AUTH_AUTHENTIK_ISSUER,
-    // Remove end
     // TODO(NextAuth): map unique user id to `providerAccountId` field
     //  profile(profile) {
     //   return {

package/src/libs/next-auth/sso-providers/cloudflare-zero-trust.ts

@@ -1,7 +1,5 @@
 import type { OIDCConfig } from '@auth/core/providers';
 
-import { authEnv } from '@/envs/auth';
-
 import { CommonProviderConfig } from './sso.config';
 
 export type CloudflareZeroTrustProfile = {
@@ -16,11 +14,10 @@ const provider = {
     ...CommonProviderConfig,
     authorization: { params: { scope: 'openid email profile' } },
     checks: ['state', 'pkce'],
-    clientId: authEnv.CLOUDFLARE_ZERO_TRUST_CLIENT_ID ?? process.env.AUTH_CLOUDFLARE_ZERO_TRUST_ID,
-    clientSecret:
-      authEnv.CLOUDFLARE_ZERO_TRUST_CLIENT_SECRET ?? process.env.AUTH_CLOUDFLARE_ZERO_TRUST_SECRET,
+    clientId: process.env.AUTH_CLOUDFLARE_ZERO_TRUST_ID,
+    clientSecret: process.env.AUTH_CLOUDFLARE_ZERO_TRUST_SECRET,
     id: 'cloudflare-zero-trust',
-    issuer: authEnv.CLOUDFLARE_ZERO_TRUST_ISSUER ?? process.env.AUTH_CLOUDFLARE_ZERO_TRUST_ISSUER,
+    issuer: process.env.AUTH_CLOUDFLARE_ZERO_TRUST_ISSUER,
     name: 'Cloudflare Zero Trust',
     profile(profile) {
       return {

package/src/libs/next-auth/sso-providers/cognito.ts

@@ -2,11 +2,7 @@ import Cognito from 'next-auth/providers/cognito';
 
 const provider = {
   id: 'cognito',
-  provider: Cognito({
-    clientId: process.env.AUTH_COGNITO_ID,
-    clientSecret: process.env.AUTH_COGNITO_SECRET,
-    issuer: process.env.AUTH_COGNITO_ISSUER,
-  }),
+  provider: Cognito({}),
 };
 
 export default provider;

package/src/libs/next-auth/sso-providers/generic-oidc.ts

@@ -1,7 +1,5 @@
 import type { OIDCConfig } from '@auth/core/providers';
 
-import { authEnv } from '@/envs/auth';
-
 import { CommonProviderConfig } from './sso.config';
 
 export type GenericOIDCProfile = {
@@ -19,10 +17,10 @@ const provider = {
     ...CommonProviderConfig,
     authorization: { params: { scope: 'email openid profile' } },
     checks: ['state', 'pkce'],
-    clientId: authEnv.GENERIC_OIDC_CLIENT_ID ?? process.env.AUTH_GENERIC_OIDC_ID,
-    clientSecret: authEnv.GENERIC_OIDC_CLIENT_SECRET ?? process.env.AUTH_GENERIC_OIDC_SECRET,
+    clientId: process.env.AUTH_GENERIC_OIDC_ID,
+    clientSecret: process.env.AUTH_GENERIC_OIDC_SECRET,
     id: 'generic-oidc',
-    issuer: authEnv.GENERIC_OIDC_ISSUER ?? process.env.AUTH_GENERIC_OIDC_ISSUER,
+    issuer: process.env.AUTH_GENERIC_OIDC_ISSUER,
     name: 'Generic OIDC',
     profile(profile) {
       return {

package/src/libs/next-auth/sso-providers/github.ts

@@ -1,7 +1,5 @@
 import GitHub from 'next-auth/providers/github';
 
-import { authEnv } from '@/envs/auth';
-
 import { CommonProviderConfig } from './sso.config';
 
 const provider = {
@@ -10,10 +8,6 @@ const provider = {
     ...CommonProviderConfig,
     // Specify auth scope, at least include 'openid email'
     authorization: { params: { scope: 'read:user user:email' } },
-    // TODO(NextAuth ENVs Migration): Remove once nextauth envs migration time end
-    clientId: authEnv.GITHUB_CLIENT_ID ?? process.env.AUTH_GITHUB_ID,
-    clientSecret: authEnv.GITHUB_CLIENT_SECRET ?? process.env.AUTH_GITHUB_SECRET,
-    // Remove end
     profile: (profile) => {
       return {
         email: profile.email,

package/src/libs/next-auth/sso-providers/google.ts

@@ -12,8 +12,6 @@ const provider = {
           'openid https://www.googleapis.com/auth/userinfo.profile https://www.googleapis.com/auth/userinfo.email openid',
       },
     },
-    clientId: process.env.AUTH_GOOGLE_CLIENT_ID,
-    clientSecret: process.env.AUTH_GOOGLE_CLIENT_SECRET,
   }),
 };
 

package/src/libs/next-auth/sso-providers/index.ts

@@ -1,7 +1,6 @@
 import Auth0 from './auth0';
 import Authelia from './authelia';
 import Authentik from './authentik';
-import AzureAD from './azure-ad';
 import Casdoor from './casdoor';
 import CloudflareZeroTrust from './cloudflare-zero-trust';
 import Cognito from './cognito';
@@ -19,7 +18,6 @@ import Zitadel from './zitadel';
 export const ssoProviders = [
   Auth0,
   Authentik,
-  AzureAD,
   GenericOIDC,
   Github,
   Zitadel,

package/src/libs/next-auth/sso-providers/keycloak.ts

@@ -8,9 +8,6 @@ const provider = {
     ...CommonProviderConfig,
     // Specify auth scope, at least include 'openid email'
     authorization: { params: { scope: 'openid email profile' } },
-    clientId: process.env.AUTH_KEYCLOAK_ID,
-    clientSecret: process.env.AUTH_KEYCLOAK_SECRET,
-    issuer: process.env.AUTH_KEYCLOAK_ISSUER,
     profile(profile) {
       return {
         email: profile.email,

package/src/libs/next-auth/sso-providers/logto.ts

@@ -1,7 +1,5 @@
 import { OIDCConfig, OIDCUserConfig } from '@auth/core/providers';
 
-import { authEnv } from '@/envs/auth';
-
 import { CommonProviderConfig } from './sso.config';
 
 interface LogtoProfile extends Record<string, any> {
@@ -41,9 +39,9 @@ const provider = {
     },
     // You can get the issuer value from the Logto Application Details page,
     // in the field "Issuer endpoint"
-    clientId: authEnv.LOGTO_CLIENT_ID ?? process.env.AUTH_LOGTO_ID,
-    clientSecret: authEnv.LOGTO_CLIENT_SECRET ?? process.env.AUTH_LOGTO_SECRET,
-    issuer: authEnv.LOGTO_ISSUER ?? process.env.AUTH_LOGTO_ISSUER,
+    clientId: process.env.AUTH_LOGTO_ID,
+    clientSecret: process.env.AUTH_LOGTO_SECRET,
+    issuer: process.env.AUTH_LOGTO_ISSUER,
   }),
 };
 

package/src/libs/next-auth/sso-providers/okta.ts

@@ -7,10 +7,6 @@ const provider = {
   provider: Okta({
     ...CommonProviderConfig,
     authorization: { params: { scope: 'openid email profile' } },
-    clientId: process.env.AUTH_OKTA_ID,
-    clientSecret: process.env.AUTH_OKTA_SECRET,
-    issuer: process.env.AUTH_OKTA_ISSUER,
-    // Remove End
     profile(profile) {
       return {
         email: profile.email,

package/src/libs/next-auth/sso-providers/zitadel.ts

@@ -1,17 +1,10 @@
 import Zitadel from 'next-auth/providers/zitadel';
 
-import { authEnv } from '@/envs/auth';
-
 const provider = {
   id: 'zitadel',
   provider: Zitadel({
     // Available scopes in ZITADEL: https://zitadel.com/docs/apis/openidoauth/scopes
     authorization: { params: { scope: 'openid email profile' } },
-    // TODO(NextAuth ENVs Migration): Remove once nextauth envs migration time end
-    clientId: authEnv.ZITADEL_CLIENT_ID ?? process.env.AUTH_ZITADEL_ID,
-    clientSecret: authEnv.ZITADEL_CLIENT_SECRET ?? process.env.AUTH_ZITADEL_SECRET,
-    issuer: authEnv.ZITADEL_ISSUER ?? process.env.AUTH_ZITADEL_ISSUER,
-    // Remove end
     // TODO(NextAuth): map unique user id to `providerAccountId` field
     // profile(profile) {
     //   return {

package/src/libs/oidc-provider/provider.ts

@@ -340,7 +340,7 @@ export const createOIDCProvider = async (db: LobeChatDatabase): Promise<Provider
 
     // 8. 令牌有效期
     ttl: {
-      AccessToken: 25 * 3600, // 25 hour
+      AccessToken: 7 * 24 * 3600, // 7 days
       AuthorizationCode: 600, // 10 minutes
       DeviceCode: 600, // 10 minutes (if enabled)
 

package/src/server/modules/AssistantStore/index.ts

@@ -93,7 +93,7 @@ export class AssistantStore {
       },
     });
     if (!res.ok) {
-      res = await fetch(this.getAgentUrl(DEFAULT_LANG), {
+      res = await fetch(this.getAgentUrl(identifier, DEFAULT_LANG), {
         cache: 'force-cache',
         next: {
           revalidate: CacheRevalidate.Details,

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/server/routers/lambda/session.ts

@@ -97,14 +97,17 @@ export const sessionRouter = router({
     }),
 
   getGroupedSessions: publicProcedure.query(async ({ ctx }): Promise<ChatSessionList> => {
-    if (!ctx.userId) return { sessionGroups: [], sessions: [] };
+    const userId = ctx.userId;
+    if (!userId) return { sessionGroups: [], sessions: [] };
 
     const serverDB = await getServerDB();
-    const sessionModel = new SessionModel(serverDB, ctx.userId!);
-    const chatGroupModel = new ChatGroupModel(serverDB, ctx.userId!);
+    const sessionModel = new SessionModel(serverDB, userId);
+    const chatGroupModel = new ChatGroupModel(serverDB, userId);
 
-    const { sessions, sessionGroups } = await sessionModel.queryWithGroups();
-    const chatGroups = await chatGroupModel.queryWithMemberDetails();
+    const [{ sessions, sessionGroups }, chatGroups] = await Promise.all([
+      sessionModel.queryWithGroups(),
+      chatGroupModel.queryWithMemberDetails(),
+    ]);
 
     const groupSessions: LobeGroupSession[] = chatGroups.map((group) => {
       const { title, description, avatar, backgroundColor, groupId, ...rest } = group;