♻️ refactor: refactor trpc request to use zod schema (#10016)

arvinxx · web-flow · commit 1a84f2cb003a · 2025-11-03T22:16:37.000+08:00
refactor request api
diff --git a/packages/types/src/message/ui/params.ts b/packages/types/src/message/ui/params.ts
@@ -1,9 +1,14 @@
 /* eslint-disable sort-keys-fix/sort-keys-fix , typescript-sort-keys/interface */
+import { z } from 'zod';
+
 import { UploadFileItem } from '../../files';
 import { MessageSemanticSearchChunk } from '../../rag';
-import { ChatMessageError } from '../common/base';
+import { ChatMessageError, ChatMessageErrorSchema } from '../common/base';
 import { ChatPluginPayload } from '../common/tools';
-import { UIChatMessage, UIMessageRoleType } from './chat';
+import { UIChatMessage } from './chat';
+import { SemanticSearchChunkSchema } from './rag';
+
+export type CreateMessageRoleType = 'user' | 'assistant' | 'tool' | 'supervisor';
 
 export interface CreateMessageParams
   extends Partial<Omit<UIChatMessage, 'content' | 'role' | 'topicId' | 'chunksList'>> {
@@ -14,7 +19,7 @@ export interface CreateMessageParams
   fromModel?: string;
   fromProvider?: string;
   groupId?: string;
-  role: UIMessageRoleType;
+  role: CreateMessageRoleType;
   sessionId: string;
   targetId?: string | null;
   threadId?: string | null;
@@ -28,7 +33,7 @@ export interface CreateMessageParams
  */
 export interface CreateNewMessageParams {
   // ========== Required fields ==========
-  role: UIMessageRoleType;
+  role: CreateMessageRoleType;
   content: string;
   sessionId: string;
 
@@ -103,3 +108,92 @@ export interface SendGroupMessageParams {
    */
   targetMemberId?: string | null;
 }
+
+// ========== Zod Schemas ========== //
+
+const UIMessageRoleTypeSchema = z.enum(['user', 'assistant', 'tool', 'supervisor']);
+
+const ChatPluginPayloadSchema = z.object({
+  apiName: z.string(),
+  arguments: z.string(),
+  identifier: z.string(),
+  type: z.string(),
+});
+
+export const CreateMessageParamsSchema = z
+  .object({
+    content: z.string(),
+    role: UIMessageRoleTypeSchema,
+    sessionId: z.string().nullable().optional(),
+    error: ChatMessageErrorSchema.nullable().optional(),
+    fileChunks: z.array(SemanticSearchChunkSchema).optional(),
+    files: z.array(z.string()).optional(),
+    fromModel: z.string().optional(),
+    fromProvider: z.string().optional(),
+    groupId: z.string().optional(),
+    targetId: z.string().nullable().optional(),
+    threadId: z.string().nullable().optional(),
+    topicId: z.string().optional(),
+    traceId: z.string().optional(),
+    // Allow additional fields from UIChatMessage (many can be null)
+    agentId: z.string().optional(),
+    children: z.any().optional(),
+    chunksList: z.any().optional(),
+    createdAt: z.number().optional(),
+    extra: z.any().optional(),
+    favorite: z.boolean().optional(),
+    fileList: z.any().optional(),
+    id: z.string().optional(),
+    imageList: z.any().optional(),
+    meta: z.any().optional(),
+    metadata: z.any().nullable().optional(),
+    model: z.string().nullable().optional(),
+    observationId: z.string().optional(),
+    parentId: z.string().optional(),
+    performance: z.any().optional(),
+    plugin: z.any().optional(),
+    pluginError: z.any().optional(),
+    pluginState: z.any().optional(),
+    provider: z.string().nullable().optional(),
+    quotaId: z.string().optional(),
+    ragQuery: z.string().nullable().optional(),
+    ragQueryId: z.string().nullable().optional(),
+    reasoning: z.any().optional(),
+    search: z.any().optional(),
+    tool_call_id: z.string().optional(),
+    toolCalls: z.any().optional(),
+    tools: z.any().optional(),
+    translate: z.any().optional(),
+    tts: z.any().optional(),
+    updatedAt: z.number().optional(),
+  })
+  .passthrough();
+
+export const CreateNewMessageParamsSchema = z
+  .object({
+    // Required fields
+    role: UIMessageRoleTypeSchema,
+    content: z.string(),
+    sessionId: z.string().nullable().optional(),
+    // Tool related
+    tool_call_id: z.string().optional(),
+    plugin: ChatPluginPayloadSchema.optional(),
+    // Grouping
+    parentId: z.string().optional(),
+    groupId: z.string().optional(),
+    // Context
+    topicId: z.string().optional(),
+    threadId: z.string().nullable().optional(),
+    targetId: z.string().nullable().optional(),
+    // Model info
+    model: z.string().nullable().optional(),
+    provider: z.string().nullable().optional(),
+    // Content
+    files: z.array(z.string()).optional(),
+    // Error handling
+    error: ChatMessageErrorSchema.nullable().optional(),
+    // Metadata
+    traceId: z.string().optional(),
+    fileChunks: z.array(SemanticSearchChunkSchema).optional(),
+  })
+  .passthrough();
diff --git a/packages/types/src/user/index.ts b/packages/types/src/user/index.ts
@@ -3,7 +3,7 @@ import { z } from 'zod';
 
 import { Plans } from '../subscription';
 import { TopicDisplayMode } from '../topic';
-import { UserSettings } from '../user/settings';
+import { UserSettings } from './settings';
 
 export interface LobeUser {
   avatar?: string;
@@ -74,3 +74,15 @@ export const NextAuthAccountSchame = z.object({
   provider: z.string(),
   providerAccountId: z.string(),
 });
+
+export const UserPreferenceSchema = z
+  .object({
+    disableInputMarkdownRender: z.boolean().optional(),
+    enableGroupChat: z.boolean().optional(),
+    guide: UserGuideSchema.optional(),
+    hideSyncAlert: z.boolean().optional(),
+    telemetry: z.boolean().nullable(),
+    topicDisplayMode: z.nativeEnum(TopicDisplayMode).optional(),
+    useCmdEnterToSend: z.boolean().optional(),
+  })
+  .partial();
diff --git a/packages/types/src/user/settings/index.ts b/packages/types/src/user/settings/index.ts
@@ -1,3 +1,5 @@
+import { z } from 'zod';
+
 import type { LobeAgentSettings } from '../../session';
 import { UserGeneralConfig } from './general';
 import { UserHotkeyConfig } from './hotkey';
@@ -18,6 +20,7 @@ export * from './keyVaults';
 export * from './modelProvider';
 export * from './sync';
 export * from './systemAgent';
+export * from './tool';
 export * from './tts';
 
 /**
@@ -34,3 +37,22 @@ export interface UserSettings {
   tool: UserToolConfig;
   tts: UserTTSConfig;
 }
+
+/**
+ * Zod schema for partial UserSettings updates
+ * Uses passthrough to allow any nested settings fields
+ */
+export const UserSettingsSchema = z
+  .object({
+    defaultAgent: z.any().optional(),
+    general: z.any().optional(),
+    hotkey: z.any().optional(),
+    image: z.any().optional(),
+    keyVaults: z.any().optional(),
+    languageModel: z.any().optional(),
+    systemAgent: z.any().optional(),
+    tool: z.any().optional(),
+    tts: z.any().optional(),
+  })
+  .passthrough()
+  .partial();
diff --git a/src/server/routers/lambda/__tests__/integration/message.integration.test.ts b/src/server/routers/lambda/__tests__/integration/message.integration.test.ts
@@ -318,47 +318,6 @@ describe('Message Router Integration Tests', () => {
     });
   });
 
-  describe('batchCreateMessages', () => {
-    it('should create multiple messages in batch', async () => {
-      const caller = messageRouter.createCaller(createTestContext(userId));
-
-      const messagesToCreate = [
-        {
-          content: 'Batch message 1',
-          role: 'user' as const,
-          sessionId: testSessionId,
-        },
-        {
-          content: 'Batch message 2',
-          role: 'assistant' as const,
-          sessionId: testSessionId,
-        },
-        {
-          content: 'Batch message 3',
-          role: 'user' as const,
-          sessionId: testSessionId,
-          topicId: testTopicId,
-        },
-      ];
-
-      const result = await caller.batchCreateMessages(messagesToCreate);
-
-      expect(result.success).toBe(true);
-      // Note: rowCount might be undefined in PGlite, so we skip this check
-      // expect(result.added).toBe(3);
-
-      // 验证数据库中的消息
-      const dbMessages = await serverDB
-        .select()
-        .from(messages)
-        .where(eq(messages.sessionId, testSessionId));
-
-      expect(dbMessages.length).toBeGreaterThanOrEqual(3);
-      const topicMessage = dbMessages.find((m) => m.content === 'Batch message 3');
-      expect(topicMessage?.topicId).toBe(testTopicId);
-    });
-  });
-
   describe('removeMessages', () => {
     it('should remove multiple messages', async () => {
       const caller = messageRouter.createCaller(createTestContext(userId));
diff --git a/src/server/routers/lambda/message.ts b/src/server/routers/lambda/message.ts
@@ -1,5 +1,6 @@
 import {
-  BatchTaskResult,
+  CreateMessageParamsSchema,
+  CreateNewMessageParamsSchema,
   UIChatMessage,
   UpdateMessageParamsSchema,
   UpdateMessageRAGParamsSchema,
@@ -27,14 +28,6 @@ const messageProcedure = authedProcedure.use(serverDatabase).use(async (opts) =>
 });
 
 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
@@ -64,15 +57,15 @@ export const messageRouter = router({
     }),
 
   createMessage: messageProcedure
-    .input(z.object({}).passthrough().partial())
+    .input(CreateMessageParamsSchema)
     .mutation(async ({ input, ctx }) => {
       const data = await ctx.messageModel.create(input as any);
 
       return data.id;
     }),
 
   createNewMessage: messageProcedure
-    .input(z.object({}).passthrough().partial())
+    .input(CreateNewMessageParamsSchema)
     .mutation(async ({ input, ctx }) => {
       return ctx.messageModel.createNewMessage(input as any, {
         postProcessUrl: (path) => ctx.fileService.getFullFileUrl(path),
diff --git a/src/server/routers/lambda/user.ts b/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,13 +24,6 @@ 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({
@@ -199,30 +200,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;
diff --git a/src/services/message/server.ts b/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,
diff --git a/src/services/message/type.ts b/src/services/message/type.ts
@@ -5,7 +5,6 @@ import {
   ChatTranslate,
   CreateMessageParams,
   CreateMessageResult,
-  DBMessageItem,
   ModelRankItem,
   UIChatMessage,
   UpdateMessageParams,
@@ -19,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[]>;
