@lobehub/chat 1.77.5 → 1.77.7

package/src/database/core/db-adaptor.ts

@@ -1,6 +1,24 @@
-// import { isDesktop } from '@/const/version';
 import { getDBInstance } from '@/database/core/web-server';
+import { LobeChatDatabase } from '@/database/type';
 
-// import { getPgliteInstance } from './electron';
+/**
+ * 懒加载数据库实例
+ * 避免每次模块导入时都初始化数据库
+ */
+let cachedDB: LobeChatDatabase | null = null;
+
+export const getServerDB = async (): Promise<LobeChatDatabase> => {
+  // 如果已经有缓存的实例，直接返回
+  if (cachedDB) return cachedDB;
+
+  try {
+    // 根据环境选择合适的数据库实例
+    cachedDB = getDBInstance();
+    return cachedDB;
+  } catch (error) {
+    console.error('❌ Failed to initialize database:', error);
+    throw error;
+  }
+};
 
 export const serverDB = getDBInstance();

package/src/database/models/__tests__/aiProvider.test.ts

@@ -4,6 +4,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { LobeChatDatabase } from '@/database/type';
 import { ModelProvider } from '@/libs/agent-runtime';
+import { sleep } from '@/utils/sleep';
 
 import { aiProviders, users } from '../../schemas';
 import { AiProviderModel } from '../aiProvider';
@@ -96,6 +97,7 @@ describe('AiProviderModel', () => {
   describe('query', () => {
     it('should query ai providers for the user', async () => {
       await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix' });
+      await sleep(10);
       await aiProviderModel.create({ name: 'AiHubMix', source: 'custom', id: 'aihubmix-2' });
 
       const userGroups = await aiProviderModel.query();

package/src/database/models/__tests__/message.test.ts

@@ -1660,7 +1660,15 @@ describe('MessageModel', () => {
 
   describe('getHeatmaps', () => {
     it('should return heatmap data for the last year', async () => {
-      const today = dayjs();
+      // 使用固定日期进行测试
+      vi.useFakeTimers();
+      const fixedDate = new Date('2023-04-07T13:00:00Z');
+      vi.setSystemTime(fixedDate);
+
+      const today = dayjs(fixedDate);
+      const twoDaysAgoDate = today.subtract(2, 'day').format('YYYY-MM-DD');
+      const oneDayAgoDate = today.subtract(1, 'day').format('YYYY-MM-DD');
+      const todayDate = today.format('YYYY-MM-DD');
 
       // 创建测试数据
       await serverDB.insert(messages).values([
@@ -1695,31 +1703,39 @@ describe('MessageModel', () => {
       expect(result.length).toBeLessThan(368);
 
       // 检查两天前的数据
-      const twoDaysAgo = result.find(
-        (item) => item.date === today.subtract(2, 'day').format('YYYY-MM-DD'),
-      );
+      const twoDaysAgo = result.find((item) => item.date === twoDaysAgoDate);
       expect(twoDaysAgo?.count).toBe(2);
       expect(twoDaysAgo?.level).toBe(1);
 
       // 检查一天前的数据
-      const oneDayAgo = result.find(
-        (item) => item.date === today.subtract(1, 'day').format('YYYY-MM-DD'),
-      );
+      const oneDayAgo = result.find((item) => item.date === oneDayAgoDate);
       expect(oneDayAgo?.count).toBe(1);
       expect(oneDayAgo?.level).toBe(1);
 
       // 检查今天的数据
-      const todayData = result.find((item) => item.date === today.format('YYYY-MM-DD'));
+      const todayData = result.find((item) => item.date === todayDate);
       expect(todayData?.count).toBe(0);
       expect(todayData?.level).toBe(0);
+
+      vi.useRealTimers();
     });
 
     it('should calculate correct levels based on message count', async () => {
-      const today = dayjs();
+      // 使用固定日期进行测试
+      vi.useFakeTimers();
+      const fixedDate = new Date('2023-05-15T12:00:00Z');
+      vi.setSystemTime(fixedDate);
+
+      const today = dayjs(fixedDate);
+      const fourDaysAgoDate = today.subtract(4, 'day').format('YYYY-MM-DD');
+      const threeDaysAgoDate = today.subtract(3, 'day').format('YYYY-MM-DD');
+      const twoDaysAgoDate = today.subtract(2, 'day').format('YYYY-MM-DD');
+      const oneDayAgoDate = today.subtract(1, 'day').format('YYYY-MM-DD');
+      const todayDate = today.format('YYYY-MM-DD');
 
       // 创建测试数据 - 不同数量的消息以测试不同的等级
       await serverDB.insert(messages).values([
-        // 1 message - level 0
+        // 1 message - level 1
         {
           id: '1',
           userId,
@@ -1727,7 +1743,7 @@ describe('MessageModel', () => {
           content: 'message 1',
           createdAt: today.subtract(4, 'day').toDate(),
         },
-        // 6 messages - level 1
+        // 6 messages - level 2
         ...Array(6)
           .fill(0)
           .map((_, i) => ({
@@ -1737,7 +1753,7 @@ describe('MessageModel', () => {
             content: `message 2-${i}`,
             createdAt: today.subtract(3, 'day').toDate(),
           })),
-        // 11 messages - level 2
+        // 11 messages - level 3
         ...Array(11)
           .fill(0)
           .map((_, i) => ({
@@ -1747,7 +1763,7 @@ describe('MessageModel', () => {
             content: `message 3-${i}`,
             createdAt: today.subtract(2, 'day').toDate(),
           })),
-        // 16 messages - level 3
+        // 16 messages - level 4
         ...Array(16)
           .fill(0)
           .map((_, i) => ({
@@ -1773,33 +1789,88 @@ describe('MessageModel', () => {
       const result = await messageModel.getHeatmaps();
 
       // 检查不同天数的等级
-      const fourDaysAgo = result.find(
-        (item) => item.date === today.subtract(4, 'day').format('YYYY-MM-DD'),
-      );
+      const fourDaysAgo = result.find((item) => item.date === fourDaysAgoDate);
       expect(fourDaysAgo?.count).toBe(1);
       expect(fourDaysAgo?.level).toBe(1);
 
-      const threeDaysAgo = result.find(
-        (item) => item.date === today.subtract(3, 'day').format('YYYY-MM-DD'),
-      );
+      const threeDaysAgo = result.find((item) => item.date === threeDaysAgoDate);
       expect(threeDaysAgo?.count).toBe(6);
       expect(threeDaysAgo?.level).toBe(2);
 
-      const twoDaysAgo = result.find(
-        (item) => item.date === today.subtract(2, 'day').format('YYYY-MM-DD'),
-      );
+      const twoDaysAgo = result.find((item) => item.date === twoDaysAgoDate);
       expect(twoDaysAgo?.count).toBe(11);
       expect(twoDaysAgo?.level).toBe(3);
 
-      const oneDayAgo = result.find(
-        (item) => item.date === today.subtract(1, 'day').format('YYYY-MM-DD'),
-      );
+      const oneDayAgo = result.find((item) => item.date === oneDayAgoDate);
       expect(oneDayAgo?.count).toBe(16);
       expect(oneDayAgo?.level).toBe(4);
 
-      const todayData = result.find((item) => item.date === today.format('YYYY-MM-DD'));
+      const todayData = result.find((item) => item.date === todayDate);
       expect(todayData?.count).toBe(21);
       expect(todayData?.level).toBe(4);
+
+      vi.useRealTimers();
+    });
+
+    it.skip('should return time count correctly when 19:00 time', async () => {
+      // 使用固定日期进行测试
+      vi.useFakeTimers();
+      const fixedDate = new Date('2025-04-02T19:00:00Z');
+      vi.setSystemTime(fixedDate);
+
+      const today = dayjs(fixedDate);
+      const twoDaysAgoDate = today.subtract(2, 'day').format('YYYY-MM-DD');
+      const oneDayAgoDate = today.subtract(1, 'day').format('YYYY-MM-DD');
+      const todayDate = today.format('YYYY-MM-DD');
+
+      // 创建测试数据
+      await serverDB.insert(messages).values([
+        {
+          id: '1',
+          userId,
+          role: 'user',
+          content: 'message 1',
+          createdAt: today.subtract(2, 'day').toDate(),
+        },
+        {
+          id: '2',
+          userId,
+          role: 'user',
+          content: 'message 2',
+          createdAt: today.subtract(2, 'day').toDate(),
+        },
+        {
+          id: '3',
+          userId,
+          role: 'user',
+          content: 'message 3',
+          createdAt: today.subtract(1, 'day').toDate(),
+        },
+      ]);
+
+      // 调用 getHeatmaps 方法
+      const result = await messageModel.getHeatmaps();
+
+      // 断言结果
+      expect(result.length).toBeGreaterThanOrEqual(366);
+      expect(result.length).toBeLessThan(368);
+
+      // 检查两天前的数据
+      const twoDaysAgo = result.find((item) => item.date === twoDaysAgoDate);
+      expect(twoDaysAgo?.count).toBe(2);
+      expect(twoDaysAgo?.level).toBe(1);
+
+      // 检查一天前的数据
+      const oneDayAgo = result.find((item) => item.date === oneDayAgoDate);
+      expect(oneDayAgo?.count).toBe(1);
+      expect(oneDayAgo?.level).toBe(1);
+
+      // 检查今天的数据
+      const todayData = result.find((item) => item.date === todayDate);
+      expect(todayData?.count).toBe(0);
+      expect(todayData?.level).toBe(0);
+
+      vi.useRealTimers();
     });
 
     it('should handle empty data', async () => {

package/src/database/models/__tests__/session.test.ts

@@ -894,7 +894,9 @@ describe('SessionModel', () => {
         await trx.insert(topics).values([
           { id: 't1', sessionId: '1', userId },
           { id: 't2', sessionId: '1', userId },
+          { id: 't6', sessionId: '1', userId },
           { id: 't3', sessionId: '2', userId },
+          { id: 't8', sessionId: '2', userId },
           { id: 't4', sessionId: '3', userId },
         ]);
       });

package/src/database/models/drizzleMigration.ts

@@ -1,3 +1,5 @@
+import { sql } from 'drizzle-orm';
+
 import { LobeChatDatabase } from '@/database/type';
 import { MigrationTableItem } from '@/types/clientDB';
 
@@ -8,6 +10,19 @@ export class DrizzleMigrationModel {
     this.db = db;
   }
 
+  getTableCounts = async () => {
+    // 这里使用 pg_tables 系统表查询用户表数量
+    const result = await this.db.execute(
+      sql`
+        SELECT COUNT(*) as table_count
+        FROM information_schema.tables
+        WHERE table_schema = 'public'
+      `,
+    );
+
+    return parseInt((result.rows[0] as any).table_count || '0');
+  };
+
   getMigrationList = async () => {
     const res = await this.db.execute(
       'SELECT * FROM "drizzle"."__drizzle_migrations" ORDER BY "created_at" DESC;',

package/src/database/models/message.ts

@@ -382,15 +382,20 @@ export class MessageModel {
       .orderBy(desc(sql`heatmaps_date`));
 
     const heatmapData: HeatmapsProps['data'] = [];
-    let currentDate = startDate;
+    let currentDate = startDate.clone();
+
+    const dateCountMap = new Map<string, number>();
+    for (const item of result) {
+      if (item?.date) {
+        const dateStr = dayjs(item.date as string).format('YYYY-MM-DD');
+        dateCountMap.set(dateStr, Number(item.count) || 0);
+      }
+    }
 
     while (currentDate.isBefore(endDate) || currentDate.isSame(endDate, 'day')) {
       const formattedDate = currentDate.format('YYYY-MM-DD');
-      const matchingResult = result.find(
-        (r) => r?.date && dayjs(r.date as string).format('YYYY-MM-DD') === formattedDate,
-      );
+      const count = dateCountMap.get(formattedDate) || 0;
 
-      const count = matchingResult ? matchingResult.count : 0;
       const levelCount = count > 0 ? Math.ceil(count / 5) : 0;
       const level = levelCount > 4 ? 4 : levelCount;
 

package/src/database/models/session.ts

@@ -193,6 +193,14 @@ export class SessionModel {
     type: 'agent' | 'group';
   }): Promise<SessionItem> => {
     return this.db.transaction(async (trx) => {
+      if (slug) {
+        const existResult = await trx.query.sessions.findFirst({
+          where: and(eq(sessions.slug, slug), eq(sessions.userId, this.userId)),
+        });
+
+        if (existResult) return existResult;
+      }
+
       const newAgents = await trx
         .insert(agents)
         .values({

package/src/database/models/user.ts

@@ -85,7 +85,13 @@ export class UserModel {
     const state = result[0];
 
     // Decrypt keyVaults
-    const decryptKeyVaults = await decryptor(state.settingsKeyVaults, this.userId);
+    let decryptKeyVaults = {};
+
+    try {
+      decryptKeyVaults = await decryptor(state.settingsKeyVaults, this.userId);
+    } catch {
+      /* empty */
+    }
 
     const settings: DeepPartial<UserSettings> = {
       defaultAgent: state.settingsDefaultAgent || {},
@@ -171,6 +177,9 @@ export class UserModel {
   };
 
   // Static method
+  static makeSureUserExist = async (db: LobeChatDatabase, userId: string) => {
+    await db.insert(users).values({ id: userId }).onConflictDoNothing();
+  };
 
   static createUser = async (db: LobeChatDatabase, params: NewUser) => {
     // if user already exists, skip creation

package/src/database/server/index.ts

@@ -1 +1 @@
-export { serverDB } from '../core/db-adaptor';
+export { getServerDB, serverDB } from '../core/db-adaptor';

package/src/features/DevPanel/features/FloatPanel.tsx

@@ -4,14 +4,16 @@ import { ActionIcon, FluentEmoji, Icon, SideNav } from '@lobehub/ui';
 import { FloatButton } from 'antd';
 import { createStyles } from 'antd-style';
 import { BugIcon, BugOff, XIcon } from 'lucide-react';
+import { usePathname } from 'next/navigation';
 import { ReactNode, memo, useEffect, useState } from 'react';
 import { Flexbox } from 'react-layout-kit';
 import { Rnd } from 'react-rnd';
 
 import { BRANDING_NAME } from '@/const/branding';
+import { isDesktop } from '@/const/version';
 
 // 定义样式
-const useStyles = createStyles(({ token, css, prefixCls }) => {
+export const useStyles = createStyles(({ token, css, prefixCls }) => {
   return {
     collapsed: css`
       pointer-events: none;
@@ -86,6 +88,7 @@ const CollapsibleFloatPanel = memo<CollapsibleFloatPanelProps>(({ items }) => {
   const [position, setPosition] = useState({ x: 100, y: 100 });
   const [size, setSize] = useState({ height: minHeight, width: minWidth });
 
+  const pathname = usePathname();
   useEffect(() => {
     try {
       const localStoragePosition = localStorage.getItem('debug-panel-position');
@@ -108,11 +111,25 @@ const CollapsibleFloatPanel = memo<CollapsibleFloatPanelProps>(({ items }) => {
 
   return (
     <>
-      <FloatButton
-        className={styles.floatButton}
-        icon={<Icon icon={isExpanded ? BugOff : BugIcon} />}
-        onClick={() => setIsExpanded(!isExpanded)}
-      />
+      {
+        // desktop devtools 下隐藏
+        pathname !== '/desktop/devtools' && (
+          <FloatButton
+            className={styles.floatButton}
+            icon={<Icon icon={isExpanded ? BugOff : BugIcon} />}
+            onClick={async () => {
+              if (isDesktop) {
+                const { electronDevtoolsService } = await import('@/services/electron/devtools');
+
+                await electronDevtoolsService.openDevtools();
+
+                return;
+              }
+              setIsExpanded(!isExpanded);
+            }}
+          />
+        )
+      }
       {isExpanded && (
         <Rnd
           bounds="window"

package/src/features/User/UserPanel/index.tsx

@@ -4,16 +4,20 @@ import { Popover } from 'antd';
 import { createStyles } from 'antd-style';
 import { PropsWithChildren, memo, useState } from 'react';
 
+import { isDesktop } from '@/const/version';
+
 import PanelContent from './PanelContent';
 import UpgradeBadge from './UpgradeBadge';
 import { useNewVersion } from './useNewVersion';
 
-const useStyles = createStyles(({ css }) => ({
-  popover: css`
-    inset-block-start: 8px !important;
-    inset-inline-start: 8px !important;
-  `,
-}));
+const useStyles = createStyles(({ css }) => {
+  return {
+    popover: css`
+      inset-block-start: ${isDesktop ? 24 : 8}px !important;
+      inset-inline-start: 8px !important;
+    `,
+  };
+});
 
 const UserPanel = memo<PropsWithChildren>(({ children }) => {
   const hasNewVersion = useNewVersion();

package/src/libs/trpc/async/index.ts

@@ -1,3 +1,5 @@
+import { getServerDB } from '@/database/core/db-adaptor';
+
 import { asyncAuth } from './asyncAuth';
 import { asyncTrpc } from './init';
 
@@ -5,6 +7,14 @@ export const publicProcedure = asyncTrpc.procedure;
 
 export const asyncRouter = asyncTrpc.router;
 
-export const asyncAuthedProcedure = asyncTrpc.procedure.use(asyncAuth);
+export const asyncAuthedProcedure = asyncTrpc.procedure.use(asyncAuth).use(
+  asyncTrpc.middleware(async (opts) => {
+    const serverDB = await getServerDB();
+
+    return opts.next({
+      ctx: { serverDB },
+    });
+  }),
+);
 
 export const createAsyncCallerFactory = asyncTrpc.createCallerFactory;

package/src/libs/trpc/lambda/index.ts

@@ -0,0 +1 @@
+export * from './serverDatabase';

package/src/libs/trpc/lambda/serverDatabase.ts

@@ -0,0 +1,10 @@
+import { getServerDB } from '@/database/core/db-adaptor';
+import { trpc } from '@/libs/trpc/init';
+
+export const serverDatabase = trpc.middleware(async (opts) => {
+  const serverDB = await getServerDB();
+
+  return opts.next({
+    ctx: { serverDB },
+  });
+});

package/src/libs/trpc/middleware/userAuth.ts

@@ -1,11 +1,21 @@
 import { TRPCError } from '@trpc/server';
 
 import { enableClerk } from '@/const/auth';
+import { DESKTOP_USER_ID } from '@/const/desktop';
+import { isDesktop } from '@/const/version';
 
 import { trpc } from '../init';
 
 export const userAuth = trpc.middleware(async (opts) => {
   const { ctx } = opts;
+
+  if (isDesktop) {
+    return opts.next({
+      ctx: {
+        userId: DESKTOP_USER_ID,
+      },
+    });
+  }
   // `ctx.user` is nullable
   if (!ctx.userId) {
     if (enableClerk) {

package/src/server/routers/async/file.ts

@@ -11,7 +11,6 @@ import { ChunkModel } from '@/database/models/chunk';
 import { EmbeddingModel } from '@/database/models/embedding';
 import { FileModel } from '@/database/models/file';
 import { NewChunkItem, NewEmbeddingsItem } from '@/database/schemas';
-import { serverDB } from '@/database/server';
 import { asyncAuthedProcedure, asyncRouter as router } from '@/libs/trpc/async';
 import { getServerDefaultFilesConfig } from '@/server/globalConfig';
 import { initAgentRuntimeWithUserPayload } from '@/server/modules/AgentRuntime';
@@ -31,11 +30,11 @@ const fileProcedure = asyncAuthedProcedure.use(async (opts) => {
 
   return opts.next({
     ctx: {
-      asyncTaskModel: new AsyncTaskModel(serverDB, ctx.userId),
-      chunkModel: new ChunkModel(serverDB, ctx.userId),
+      asyncTaskModel: new AsyncTaskModel(ctx.serverDB, ctx.userId),
+      chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
       chunkService: new ChunkService(ctx.userId),
-      embeddingModel: new EmbeddingModel(serverDB, ctx.userId),
-      fileModel: new FileModel(serverDB, ctx.userId),
+      embeddingModel: new EmbeddingModel(ctx.serverDB, ctx.userId),
+      fileModel: new FileModel(ctx.serverDB, ctx.userId),
     },
   });
 });

package/src/server/routers/async/ragEval.ts

@@ -7,7 +7,6 @@ import { DEFAULT_EMBEDDING_MODEL, DEFAULT_MODEL } from '@/const/settings';
 import { ChunkModel } from '@/database/models/chunk';
 import { EmbeddingModel } from '@/database/models/embedding';
 import { FileModel } from '@/database/models/file';
-import { serverDB } from '@/database/server';
 import {
   EvalDatasetRecordModel,
   EvalEvaluationModel,
@@ -25,13 +24,13 @@ const ragEvalProcedure = asyncAuthedProcedure.use(async (opts) => {
 
   return opts.next({
     ctx: {
-      chunkModel: new ChunkModel(serverDB, ctx.userId),
+      chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
       chunkService: new ChunkService(ctx.userId),
       datasetRecordModel: new EvalDatasetRecordModel(ctx.userId),
-      embeddingModel: new EmbeddingModel(serverDB, ctx.userId),
+      embeddingModel: new EmbeddingModel(ctx.serverDB, ctx.userId),
       evalRecordModel: new EvaluationRecordModel(ctx.userId),
       evaluationModel: new EvalEvaluationModel(ctx.userId),
-      fileModel: new FileModel(serverDB, ctx.userId),
+      fileModel: new FileModel(ctx.serverDB, ctx.userId),
     },
   });
 });

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

@@ -2,16 +2,16 @@ import { z } from 'zod';
 
 import { SessionGroupModel } from '@/database/models/sessionGroup';
 import { insertSessionGroupSchema } from '@/database/schemas';
-import { serverDB } from '@/database/server';
 import { authedProcedure, router } from '@/libs/trpc';
+import { serverDatabase } from '@/libs/trpc/lambda';
 import { SessionGroupItem } from '@/types/session';
 
-const sessionProcedure = authedProcedure.use(async (opts) => {
+const sessionProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
   const { ctx } = opts;
 
   return opts.next({
     ctx: {
-      sessionGroupModel: new SessionGroupModel(serverDB, ctx.userId),
+      sessionGroupModel: new SessionGroupModel(ctx.serverDB, ctx.userId),
     },
   });
 });
@@ -69,8 +69,6 @@ export const sessionGroupRouter = router({
       }),
     )
     .mutation(async ({ input, ctx }) => {
-      console.log('sortMap:', input.sortMap);
-
       return ctx.sessionGroupModel.updateOrder(input.sortMap);
     }),
 });

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

@@ -7,22 +7,22 @@ import { FileModel } from '@/database/models/file';
 import { KnowledgeBaseModel } from '@/database/models/knowledgeBase';
 import { SessionModel } from '@/database/models/session';
 import { UserModel } from '@/database/models/user';
-import { serverDB } from '@/database/server';
 import { pino } from '@/libs/logger';
 import { authedProcedure, router } from '@/libs/trpc';
+import { serverDatabase } from '@/libs/trpc/lambda';
 import { AgentService } from '@/server/services/agent';
 import { KnowledgeItem, KnowledgeType } from '@/types/knowledgeBase';
 
-const agentProcedure = authedProcedure.use(async (opts) => {
+const agentProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
   const { ctx } = opts;
 
   return opts.next({
     ctx: {
-      agentModel: new AgentModel(serverDB, ctx.userId),
-      agentService: new AgentService(serverDB, ctx.userId),
-      fileModel: new FileModel(serverDB, ctx.userId),
-      knowledgeBaseModel: new KnowledgeBaseModel(serverDB, ctx.userId),
-      sessionModel: new SessionModel(serverDB, ctx.userId),
+      agentModel: new AgentModel(ctx.serverDB, ctx.userId),
+      agentService: new AgentService(ctx.serverDB, ctx.userId),
+      fileModel: new FileModel(ctx.serverDB, ctx.userId),
+      knowledgeBaseModel: new KnowledgeBaseModel(ctx.serverDB, ctx.userId),
+      sessionModel: new SessionModel(ctx.serverDB, ctx.userId),
     },
   });
 });
@@ -90,7 +90,7 @@ export const agentRouter = router({
         // if there is no session for user, create one
         if (!item) {
           // if there is no user, return default config
-          const user = await UserModel.findById(serverDB, ctx.userId);
+          const user = await UserModel.findById(ctx.serverDB, ctx.userId);
           if (!user) return DEFAULT_AGENT_CONFIG;
 
           const res = await ctx.agentService.createInbox();

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

@@ -3,8 +3,8 @@ import { z } from 'zod';
 import { AiModelModel } from '@/database/models/aiModel';
 import { UserModel } from '@/database/models/user';
 import { AiInfraRepos } from '@/database/repositories/aiInfra';
-import { serverDB } from '@/database/server';
 import { authedProcedure, router } from '@/libs/trpc';
+import { serverDatabase } from '@/libs/trpc/lambda';
 import { getServerGlobalConfig } from '@/server/globalConfig';
 import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
 import {
@@ -15,7 +15,7 @@ import {
 } from '@/types/aiModel';
 import { ProviderConfig } from '@/types/user/settings';
 
-const aiModelProcedure = authedProcedure.use(async (opts) => {
+const aiModelProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
   const { ctx } = opts;
 
   const gateKeeper = await KeyVaultsGateKeeper.initWithEnvKey();
@@ -24,13 +24,13 @@ const aiModelProcedure = authedProcedure.use(async (opts) => {
   return opts.next({
     ctx: {
       aiInfraRepos: new AiInfraRepos(
-        serverDB,
+        ctx.serverDB,
         ctx.userId,
         aiProvider as Record<string, ProviderConfig>,
       ),
-      aiModelModel: new AiModelModel(serverDB, ctx.userId),
+      aiModelModel: new AiModelModel(ctx.serverDB, ctx.userId),
       gateKeeper,
-      userModel: new UserModel(serverDB, ctx.userId),
+      userModel: new UserModel(ctx.serverDB, ctx.userId),
     },
   });
 });

package/src/server/routers/lambda/aiProvider.test.ts

@@ -1,9 +1,7 @@
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { AiProviderModel } from '@/database/models/aiProvider';
-import { UserModel } from '@/database/models/user';
 import { AiInfraRepos } from '@/database/repositories/aiInfra';
-import { serverDB } from '@/database/server';
 import { getServerGlobalConfig } from '@/server/globalConfig';
 import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
 import { AiProviderDetailItem, AiProviderRuntimeState } from '@/types/aiProvider';