@lobehub/lobehub 2.0.0-next.5 → 2.0.0-next.50

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

@@ -1,14 +1,17 @@
-import { BatchTaskResult, UIChatMessage, UpdateMessageRAGParamsSchema } from '@lobechat/types';
+import {
+  CreateNewMessageParamsSchema,
+  UpdateMessageParamsSchema,
+  UpdateMessagePluginSchema,
+  UpdateMessageRAGParamsSchema,
+} from '@lobechat/types';
 import { z } from 'zod';
 
 import { MessageModel } from '@/database/models/message';
-import { updateMessagePluginSchema } from '@/database/schemas';
 import { getServerDB } from '@/database/server';
 import { authedProcedure, publicProcedure, router } from '@/libs/trpc/lambda';
 import { serverDatabase } from '@/libs/trpc/lambda/middleware';
 import { FileService } from '@/server/services/file';
-
-type ChatMessageList = UIChatMessage[];
+import { MessageService } from '@/server/services/message';
 
 const messageProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
   const { ctx } = opts;
@@ -17,19 +20,12 @@ const messageProcedure = authedProcedure.use(serverDatabase).use(async (opts) =>
     ctx: {
       fileService: new FileService(ctx.serverDB, ctx.userId),
       messageModel: new MessageModel(ctx.serverDB, ctx.userId),
+      messageService: new MessageService(ctx.serverDB, ctx.userId),
     },
   });
 });
 
 export const messageRouter = router({
-  batchCreateMessages: messageProcedure
-    .input(z.array(z.any()))
-    .mutation(async ({ input, ctx }): Promise<BatchTaskResult> => {
-      const data = await ctx.messageModel.batchCreate(input);
-
-      return { added: data.rowCount as number, ids: [], skips: [], success: true };
-    }),
-
   count: messageProcedure
     .input(
       z
@@ -59,35 +55,9 @@ export const messageRouter = router({
     }),
 
   createMessage: messageProcedure
-    .input(z.object({}).passthrough().partial())
+    .input(CreateNewMessageParamsSchema)
     .mutation(async ({ input, ctx }) => {
-      const data = await ctx.messageModel.create(input as any);
-
-      return data.id;
-    }),
-
-  createNewMessage: messageProcedure
-    .input(z.object({}).passthrough().partial())
-    .mutation(async ({ input, ctx }) => {
-      return ctx.messageModel.createNewMessage(input as any, {
-        postProcessUrl: (path) => ctx.fileService.getFullFileUrl(path),
-      });
-    }),
-
-  // TODO: it will be removed in V2
-  getAllMessages: messageProcedure.query(async ({ ctx }): Promise<ChatMessageList> => {
-    return ctx.messageModel.queryAll() as any;
-  }),
-
-  // TODO: it will be removed in V2
-  getAllMessagesInSession: messageProcedure
-    .input(
-      z.object({
-        sessionId: z.string().nullable().optional(),
-      }),
-    )
-    .query(async ({ ctx, input }): Promise<ChatMessageList> => {
-      return ctx.messageModel.queryBySessionId(input.sessionId) as any;
+      return ctx.messageService.createMessage(input as any);
     }),
 
   getHeatmaps: messageProcedure.query(async ({ ctx }) => {
@@ -113,6 +83,7 @@ export const messageRouter = router({
       const fileService = new FileService(serverDB, ctx.userId);
 
       return messageModel.query(input, {
+        groupAssistantMessages: false,
         postProcessUrl: (path) => fileService.getFullFileUrl(path),
       });
     }),
@@ -126,9 +97,16 @@ export const messageRouter = router({
   }),
 
   removeMessage: messageProcedure
-    .input(z.object({ id: z.string() }))
+    .input(
+      z.object({
+        id: z.string(),
+        sessionId: z.string().nullable().optional(),
+        topicId: z.string().nullable().optional(),
+      }),
+    )
     .mutation(async ({ input, ctx }) => {
-      return ctx.messageModel.deleteMessage(input.id);
+      const { id, ...options } = input;
+      return ctx.messageService.removeMessage(id, options);
     }),
 
   removeMessageQuery: messageProcedure
@@ -138,9 +116,16 @@ export const messageRouter = router({
     }),
 
   removeMessages: messageProcedure
-    .input(z.object({ ids: z.array(z.string()) }))
+    .input(
+      z.object({
+        ids: z.array(z.string()),
+        sessionId: z.string().nullable().optional(),
+        topicId: z.string().nullable().optional(),
+      }),
+    )
     .mutation(async ({ input, ctx }) => {
-      return ctx.messageModel.deleteMessages(input.ids);
+      const { ids, ...options } = input;
+      return ctx.messageService.removeMessages(ids, options);
     }),
 
   removeMessagesByAssistant: messageProcedure
@@ -180,61 +165,82 @@ 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);
+      const { id, value, ...options } = input;
+      return ctx.messageService.updateMessage(id, value as any, options);
     }),
 
   updateMessagePlugin: messageProcedure
     .input(
       z.object({
         id: z.string(),
-        value: updateMessagePluginSchema.partial(),
+        sessionId: z.string().nullable().optional(),
+        topicId: z.string().nullable().optional(),
+        value: UpdateMessagePluginSchema.partial(),
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      return ctx.messageModel.updateMessagePlugin(input.id, input.value);
+      const { id, value, ...options } = input;
+      return ctx.messageService.updateMessagePlugin(id, value, options);
     }),
 
   updateMessageRAG: messageProcedure
-    .input(UpdateMessageRAGParamsSchema)
+    .input(
+      UpdateMessageRAGParamsSchema.extend({
+        sessionId: z.string().nullable().optional(),
+        topicId: z.string().nullable().optional(),
+      }),
+    )
     .mutation(async ({ input, ctx }) => {
-      await ctx.messageModel.updateMessageRAG(input.id, input.value);
+      const { id, value, ...options } = input;
+      return ctx.messageService.updateMessageRAG(id, value, options);
     }),
 
   updateMetadata: messageProcedure
     .input(
       z.object({
         id: z.string(),
+        sessionId: z.string().nullable().optional(),
+        topicId: z.string().nullable().optional(),
         value: z.object({}).passthrough(),
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      return ctx.messageModel.updateMetadata(input.id, input.value);
+      const { id, value, ...options } = input;
+      return ctx.messageService.updateMetadata(id, value, options);
     }),
 
   updatePluginError: messageProcedure
     .input(
       z.object({
         id: z.string(),
+        sessionId: z.string().nullable().optional(),
+        topicId: z.string().nullable().optional(),
         value: z.object({}).passthrough().nullable(),
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      return ctx.messageModel.updateMessagePlugin(input.id, { error: input.value });
+      const { id, value, ...options } = input;
+      return ctx.messageService.updatePluginError(id, value, options);
     }),
 
   updatePluginState: messageProcedure
     .input(
       z.object({
         id: z.string(),
+        sessionId: z.string().nullable().optional(),
+        topicId: z.string().nullable().optional(),
         value: z.object({}).passthrough(),
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      return ctx.messageModel.updatePluginState(input.id, input.value);
+      const { id, value, ...options } = input;
+      return ctx.messageService.updatePluginState(id, value, options);
     }),
 
   updateTTS: messageProcedure
@@ -257,7 +263,6 @@ export const messageRouter = router({
 
       return ctx.messageModel.updateTTS(input.id, input.value);
     }),
-
   updateTranslate: messageProcedure
     .input(
       z.object({
@@ -279,5 +284,3 @@ export const messageRouter = router({
       return ctx.messageModel.updateTranslate(input.id, input.value);
     }),
 });
-
-export type MessageRouter = typeof messageRouter;

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

@@ -128,7 +128,13 @@ export const topicRouter = router({
     }),
 
   searchTopics: topicProcedure
-    .input(z.object({ keywords: z.string(), sessionId: z.string().nullable().optional() }))
+    .input(
+      z.object({
+        groupId: z.string().nullable().optional(),
+        keywords: z.string(),
+        sessionId: z.string().nullable().optional(),
+      }),
+    )
     .query(async ({ input, ctx }) => {
       return ctx.topicModel.queryByKeyword(input.keywords, input.sessionId);
     }),

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

@@ -0,0 +1,36 @@
+import { z } from 'zod';
+
+import { authedProcedure, router } from '@/libs/trpc/lambda';
+import { serverDatabase } from '@/libs/trpc/lambda/middleware';
+import { UsageRecordService } from '@/server/services/usage';
+
+const usageProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
+  const { ctx } = opts;
+  return opts.next({
+    ctx: {
+      usageRecordService: new UsageRecordService(ctx.serverDB, ctx.userId),
+    },
+  });
+});
+
+export const usageRouter = router({
+  findAndGroupByDay: usageProcedure
+    .input(
+      z.object({
+        mo: z.string().optional(),
+      }),
+    )
+    .query(async ({ ctx, input }) => {
+      return await ctx.usageRecordService.findAndGroupByDay(input.mo);
+    }),
+
+  findByMonth: usageProcedure
+    .input(
+      z.object({
+        mo: z.string().optional(),
+      }),
+    )
+    .query(async ({ ctx, input }) => {
+      return await ctx.usageRecordService.findByMonth(input.mo);
+    }),
+});

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

@@ -1,9 +1,17 @@
 import { UserJSON } from '@clerk/backend';
+import { enableClerk, isDesktop } from '@lobechat/const';
+import {
+  NextAuthAccountSchame,
+  UserGuideSchema,
+  UserInitializationState,
+  UserPreference,
+  UserPreferenceSchema,
+  UserSettings,
+  UserSettingsSchema,
+} from '@lobechat/types';
 import { v4 as uuidv4 } from 'uuid';
 import { z } from 'zod';
 
-import { enableClerk } from '@/const/auth';
-import { isDesktop } from '@/const/version';
 import { MessageModel } from '@/database/models/message';
 import { SessionModel } from '@/database/models/session';
 import { UserModel, UserNotFoundError } from '@/database/models/user';
@@ -16,20 +24,15 @@ import { S3 } from '@/server/modules/S3';
 import { FileService } from '@/server/services/file';
 import { NextAuthUserService } from '@/server/services/nextAuthUser';
 import { UserService } from '@/server/services/user';
-import {
-  NextAuthAccountSchame,
-  UserGuideSchema,
-  UserInitializationState,
-  UserPreference,
-} from '@/types/user';
-import { UserSettings } from '@/types/user/settings';
 
 const userProcedure = authedProcedure.use(serverDatabase).use(async ({ ctx, next }) => {
   return next({
     ctx: {
       clerkAuth: new ClerkAuth(),
       fileService: new FileService(ctx.serverDB, ctx.userId),
+      messageModel: new MessageModel(ctx.serverDB, ctx.userId),
       nextAuthUserService: new NextAuthUserService(ctx.serverDB),
+      sessionModel: new SessionModel(ctx.serverDB, ctx.userId),
       userModel: new UserModel(ctx.serverDB, ctx.userId),
     },
   });
@@ -96,12 +99,12 @@ export const userRouter = router({
       }
     }
 
-    const messageModel = new MessageModel(ctx.serverDB, ctx.userId);
-    const hasMoreThan4Messages = await messageModel.hasMoreThanN(4);
-
-    const sessionModel = new SessionModel(ctx.serverDB, ctx.userId);
-    const hasAnyMessages = await messageModel.hasMoreThanN(0);
-    const hasExtraSession = await sessionModel.hasMoreThanN(1);
+    // Run all count queries in parallel
+    const [hasMoreThan4Messages, hasAnyMessages, hasExtraSession] = await Promise.all([
+      ctx.messageModel.hasMoreThanN(4),
+      ctx.messageModel.hasMoreThanN(0),
+      ctx.sessionModel.hasMoreThanN(1),
+    ]);
 
     return {
       avatar: state.avatar,
@@ -199,30 +202,28 @@ export const userRouter = router({
     return ctx.userModel.updateGuide(input);
   }),
 
-  updatePreference: userProcedure.input(z.any()).mutation(async ({ ctx, input }) => {
+  updatePreference: userProcedure.input(UserPreferenceSchema).mutation(async ({ ctx, input }) => {
     return ctx.userModel.updatePreference(input);
   }),
 
-  updateSettings: userProcedure
-    .input(z.object({}).passthrough())
-    .mutation(async ({ ctx, input }) => {
-      const { keyVaults, ...res } = input as Partial<UserSettings>;
+  updateSettings: userProcedure.input(UserSettingsSchema).mutation(async ({ ctx, input }) => {
+    const { keyVaults, ...res } = input as Partial<UserSettings>;
 
-      // Encrypt keyVaults
-      let encryptedKeyVaults: string | null = null;
+    // Encrypt keyVaults
+    let encryptedKeyVaults: string | null = null;
 
-      if (keyVaults) {
-        // TODO: better to add a validation
-        const data = JSON.stringify(keyVaults);
-        const gateKeeper = await KeyVaultsGateKeeper.initWithEnvKey();
+    if (keyVaults) {
+      // TODO: better to add a validation
+      const data = JSON.stringify(keyVaults);
+      const gateKeeper = await KeyVaultsGateKeeper.initWithEnvKey();
 
-        encryptedKeyVaults = await gateKeeper.encrypt(data);
-      }
+      encryptedKeyVaults = await gateKeeper.encrypt(data);
+    }
 
-      const nextValue = { ...res, keyVaults: encryptedKeyVaults };
+    const nextValue = { ...res, keyVaults: encryptedKeyVaults };
 
-      return ctx.userModel.updateSetting(nextValue);
-    }),
+    return ctx.userModel.updateSetting(nextValue);
+  }),
 });
 
 export type UserRouter = typeof userRouter;

package/src/server/routers/mobile/index.ts

@@ -8,23 +8,35 @@ import { agentRouter } from '../lambda/agent';
 import { aiChatRouter } from '../lambda/aiChat';
 import { aiModelRouter } from '../lambda/aiModel';
 import { aiProviderRouter } from '../lambda/aiProvider';
+import { chunkRouter } from '../lambda/chunk';
+import { configRouter } from '../lambda/config';
+import { documentRouter } from '../lambda/document';
+import { fileRouter } from '../lambda/file';
+import { knowledgeBaseRouter } from '../lambda/knowledgeBase';
 import { marketRouter } from '../lambda/market';
 import { messageRouter } from '../lambda/message';
 import { sessionRouter } from '../lambda/session';
 import { sessionGroupRouter } from '../lambda/sessionGroup';
 import { topicRouter } from '../lambda/topic';
+import { uploadRouter } from '../lambda/upload';
+import { userRouter } from '../lambda/user';
 
 export const mobileRouter = router({
   agent: agentRouter,
   aiChat: aiChatRouter,
   aiModel: aiModelRouter,
   aiProvider: aiProviderRouter,
+  chunk: chunkRouter,
+  config: configRouter,
+  document: documentRouter,
+  file: fileRouter,
   healthcheck: publicProcedure.query(() => "i'm live!"),
+  knowledgeBase: knowledgeBaseRouter,
   market: marketRouter,
   message: messageRouter,
   session: sessionRouter,
   sessionGroup: sessionGroupRouter,
   topic: topicRouter,
+  upload: uploadRouter,
+  user: userRouter,
 });
-
-export type MobileRouter = typeof mobileRouter;

package/src/server/routers/tools/mcp.ts

@@ -1,3 +1,4 @@
+import { isDesktop } from '@lobechat/const';
 import {
   GetStreamableMcpServerManifestInputSchema,
   StreamableHTTPAuthSchema,
@@ -5,10 +6,12 @@ import {
 import { TRPCError } from '@trpc/server';
 import { z } from 'zod';
 
-import { isDesktop, isServerMode } from '@/const/version';
-import { passwordProcedure } from '@/libs/trpc/edge';
+import { ToolCallContent } from '@/libs/mcp';
 import { authedProcedure, router } from '@/libs/trpc/lambda';
+import { serverDatabase } from '@/libs/trpc/lambda/middleware';
+import { FileService } from '@/server/services/file';
 import { mcpService } from '@/server/services/mcp';
+import { processContentBlocks } from '@/server/services/mcp/contentProcessor';
 
 // Define Zod schemas for MCP Client parameters
 const httpParamsSchema = z.object({
@@ -38,7 +41,13 @@ const checkStdioEnvironment = (params: z.infer<typeof mcpClientParamsSchema>) =>
   }
 };
 
-const mcpProcedure = isServerMode ? authedProcedure : passwordProcedure;
+const mcpProcedure = authedProcedure.use(serverDatabase).use(async ({ ctx, next }) => {
+  return next({
+    ctx: {
+      fileService: new FileService(ctx.serverDB, ctx.userId),
+    },
+  });
+});
 
 export const mcpRouter = router({
   getStreamableMcpServerManifest: mcpProcedure
@@ -96,13 +105,21 @@ export const mcpRouter = router({
         toolName: z.string(),
       }),
     )
-    .mutation(async ({ input }) => {
+    .mutation(async ({ input, ctx }) => {
       // Stdio check can be done here or rely on the service/client layer
       checkStdioEnvironment(input.params);
 
-      // Pass the validated params, toolName, and args to the service
-      const data = await mcpService.callTool(input.params, input.toolName, input.args);
+      // Create a closure that binds fileService and userId to processContentBlocks
+      const boundProcessContentBlocks = async (blocks: ToolCallContent[]) => {
+        return processContentBlocks(blocks, ctx.fileService);
+      };
 
-      return JSON.stringify(data);
+      // Pass the validated params, toolName, args, and bound processContentBlocks to the service
+      return await mcpService.callTool({
+        clientParams: input.params,
+        toolName: input.toolName,
+        argsStr: input.args,
+        processContentBlocks: boundProcessContentBlocks,
+      });
     }),
 });

package/src/server/routers/tools/search.test.ts

@@ -23,13 +23,7 @@ vi.mock('@/server/services/search/impls/searxng/client');
 
 describe('searchRouter', () => {
   const mockContext = {
-    req: {
-      headers: {
-        authorization: 'Bearer mock-token',
-      },
-    },
-    authorizationHeader: 'Bearer mock-token',
-    jwtPayload: { userId: '1' },
+    userId: 'test-user-id',
   };
 
   beforeEach(() => {

package/src/server/routers/tools/search.ts

@@ -1,12 +1,9 @@
-import { isServerMode } from '@lobechat/const';
 import { z } from 'zod';
 
-import { passwordProcedure } from '@/libs/trpc/edge';
 import { authedProcedure, router } from '@/libs/trpc/lambda';
 import { searchService } from '@/server/services/search';
 
-// TODO: password procedure 未来的处理方式可能要思考下
-const searchProcedure = isServerMode ? authedProcedure : passwordProcedure;
+const searchProcedure = authedProcedure;
 
 export const searchRouter = router({
   crawlPages: searchProcedure