@lobehub/chat 0.159.12 → 0.160.1

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

@@ -0,0 +1,44 @@
+import { TRPCError } from '@trpc/server';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { createCallerFactory } from '@/libs/trpc';
+import { AuthContext, createContextInner } from '@/server/context';
+
+import { trpc } from '../init';
+import { userAuth } from './userAuth';
+
+const appRouter = trpc.router({
+  protectedQuery: trpc.procedure.use(userAuth).query(async ({ ctx }) => {
+    return ctx.userId;
+  }),
+});
+
+const createCaller = createCallerFactory(appRouter);
+let ctx: AuthContext;
+let router: ReturnType<typeof createCaller>;
+
+beforeEach(async () => {
+  vi.resetAllMocks();
+});
+
+describe('userAuth middleware', () => {
+  it('should throw UNAUTHORIZED error if userId is not present in context', async () => {
+    ctx = await createContextInner();
+    router = createCaller(ctx);
+
+    try {
+      await router.protectedQuery();
+    } catch (e) {
+      expect(e).toEqual(new TRPCError({ code: 'UNAUTHORIZED' }));
+    }
+  });
+
+  it('should call next with userId in context if userId is present', async () => {
+    ctx = await createContextInner({ userId: 'user-id' });
+    router = createCaller(ctx);
+
+    const data = await router.protectedQuery();
+
+    expect(data).toEqual('user-id');
+  });
+});

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

@@ -0,0 +1,18 @@
+import { TRPCError } from '@trpc/server';
+
+import { trpc } from '../init';
+
+export const userAuth = trpc.middleware(async (opts) => {
+  const { ctx } = opts;
+  // `ctx.user` is nullable
+  if (!ctx.userId) {
+    throw new TRPCError({ code: 'UNAUTHORIZED' });
+  }
+
+  return opts.next({
+    ctx: {
+      // ✅ user value is known to be non-null now
+      userId: ctx.userId,
+    },
+  });
+});

package/src/server/context.ts

@@ -1,7 +1,15 @@
 /* eslint-disable @typescript-eslint/no-unused-vars */
+import { getAuth } from '@clerk/nextjs/server';
 import { NextRequest } from 'next/server';
 
+import { JWTPayload, LOBE_CHAT_AUTH_HEADER, enableClerk } from '@/const/auth';
+
+type ClerkAuth = ReturnType<typeof getAuth>;
+
 export interface AuthContext {
+  auth?: ClerkAuth;
+  authorizationHeader?: string | null;
+  jwtPayload?: JWTPayload | null;
   userId?: string | null;
 }
 
@@ -9,8 +17,14 @@ export interface AuthContext {
  * Inner function for `createContext` where we create the context.
  * This is useful for testing when we don't want to mock Next.js' request/response
  */
-export const createContextInner = async (): Promise<AuthContext> => ({
-  userId: null,
+export const createContextInner = async (params?: {
+  auth?: ClerkAuth;
+  authorizationHeader?: string | null;
+  userId?: string | null;
+}): Promise<AuthContext> => ({
+  auth: params?.auth,
+  authorizationHeader: params?.authorizationHeader,
+  userId: params?.userId,
 });
 
 export type Context = Awaited<ReturnType<typeof createContextInner>>;
@@ -22,5 +36,16 @@ export type Context = Awaited<ReturnType<typeof createContextInner>>;
 export const createContext = async (request: NextRequest): Promise<Context> => {
   // for API-response caching see https://trpc.io/docs/v11/caching
 
-  return createContextInner();
+  const authorization = request.headers.get(LOBE_CHAT_AUTH_HEADER);
+
+  let userId;
+  let auth;
+
+  if (enableClerk) {
+    auth = getAuth(request);
+
+    userId = auth.userId;
+  }
+
+  return createContextInner({ auth, authorizationHeader: authorization, userId });
 };

package/src/server/files/s3.ts

@@ -0,0 +1,58 @@
+import { ListObjectsCommand, PutObjectCommand, S3Client } from '@aws-sdk/client-s3';
+import { getSignedUrl } from '@aws-sdk/s3-request-presigner';
+import { z } from 'zod';
+
+import { fileEnv } from '@/config/file';
+
+export const fileSchema = z.object({
+  Key: z.string(),
+  LastModified: z.date(),
+  Size: z.number(),
+});
+
+export const listFileSchema = z.array(fileSchema);
+
+export type FileType = z.infer<typeof fileSchema>;
+
+const DEFAULT_S3_REGION = 'us-east-1';
+
+export class S3 {
+  private readonly client: S3Client;
+
+  private readonly bucket: string;
+
+  constructor() {
+    if (!fileEnv.S3_ACCESS_KEY_ID || !fileEnv.S3_SECRET_ACCESS_KEY || !fileEnv.S3_BUCKET)
+      throw new Error('S3 environment variables are not set completely, please check your env');
+
+    this.bucket = fileEnv.S3_BUCKET;
+
+    this.client = new S3Client({
+      credentials: {
+        accessKeyId: fileEnv.S3_ACCESS_KEY_ID,
+        secretAccessKey: fileEnv.S3_SECRET_ACCESS_KEY,
+      },
+      endpoint: fileEnv.S3_ENDPOINT,
+      region: fileEnv.S3_REGION || DEFAULT_S3_REGION,
+    });
+  }
+
+  public async getImages(): Promise<FileType[]> {
+    const command = new ListObjectsCommand({
+      Bucket: this.bucket,
+    });
+
+    const res = await this.client.send(command);
+    return listFileSchema.parse(res.Contents);
+  }
+
+  public async createPreSignedUrl(key: string): Promise<string> {
+    const command = new PutObjectCommand({
+      ACL: 'public-read',
+      Bucket: this.bucket,
+      Key: key,
+    });
+
+    return getSignedUrl(this.client, command, { expiresIn: 3600 });
+  }
+}

package/src/server/globalConfig/index.ts

@@ -1,3 +1,4 @@
+import { fileEnv } from '@/config/file';
 import {
   OllamaProviderCard,
   OpenAIProviderCard,
@@ -50,6 +51,7 @@ export const getServerGlobalConfig = () => {
       config: parseAgentConfig(DEFAULT_AGENT_CONFIG),
     },
 
+    enableUploadFileToServer: !!fileEnv.S3_SECRET_ACCESS_KEY,
     enabledAccessCode: ACCESS_CODES?.length > 0,
     enabledOAuthSSO: enableNextAuth,
     languageModel: {

package/src/server/mock.ts

@@ -3,6 +3,6 @@
  */
 import { createCallerFactory } from '@/libs/trpc';
 
-import { appRouter } from './routers';
+import { edgeRouter } from './routers';
 
-export const createCaller = createCallerFactory(appRouter);
+export const createCaller = createCallerFactory(edgeRouter);

package/src/server/routers/{config → edge/config}/index.test.ts

@@ -1,3 +1,4 @@
+// @vitest-environment node
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 
 /**

package/src/server/routers/edge/upload.ts

@@ -0,0 +1,16 @@
+import { z } from 'zod';
+
+import { passwordProcedure, router } from '@/libs/trpc';
+import { S3 } from '@/server/files/s3';
+
+export const uploadRouter = router({
+  createS3PreSignedUrl: passwordProcedure
+    .input(z.object({ pathname: z.string() }))
+    .mutation(async ({ input }) => {
+      const s3 = new S3();
+
+      return await s3.createPreSignedUrl(input.pathname);
+    }),
+});
+
+export type FileRouter = typeof uploadRouter;

package/src/server/routers/index.ts

@@ -3,11 +3,13 @@
  */
 import { publicProcedure, router } from '@/libs/trpc';
 
-import { configRouter } from './config';
+import { configRouter } from './edge/config';
+import { uploadRouter } from './edge/upload';
 
-export const appRouter = router({
+export const edgeRouter = router({
   config: configRouter,
   healthcheck: publicProcedure.query(() => "i'm live!"),
+  upload: uploadRouter,
 });
 
-export type AppRouter = typeof appRouter;
+export type EdgeRouter = typeof edgeRouter;

package/src/services/__tests__/global.test.ts

@@ -1,7 +1,6 @@
 import { Mock, beforeEach, describe, expect, it, vi } from 'vitest';
 
-import { trpcClient } from '@/libs/trpc/client';
-import { LobeAgentConfig } from '@/types/agent';
+import { edgeClient } from '@/libs/trpc/client';
 import { GlobalServerConfig } from '@/types/serverConfig';
 
 import { globalService } from '../global';
@@ -14,7 +13,7 @@ beforeEach(() => {
 
 vi.mock('@/libs/trpc/client', () => {
   return {
-    trpcClient: {
+    edgeClient: {
       config: {
         getGlobalConfig: { query: vi.fn() },
         getDefaultAgentConfig: { query: vi.fn() },
@@ -76,7 +75,7 @@ describe('GlobalService', () => {
     it('should return the serverConfig when fetch is successful', async () => {
       // Arrange
       const mockConfig = { enabledOAuthSSO: true } as GlobalServerConfig;
-      vi.spyOn(trpcClient.config.getGlobalConfig, 'query').mockResolvedValue(mockConfig);
+      vi.spyOn(edgeClient.config.getGlobalConfig, 'query').mockResolvedValue(mockConfig);
 
       // Act
       const config = await globalService.getGlobalConfig();
@@ -87,7 +86,7 @@ describe('GlobalService', () => {
 
     it('should return the defaultAgentConfig when fetch is successful', async () => {
       // Arrange
-      vi.spyOn(trpcClient.config.getDefaultAgentConfig, 'query').mockResolvedValue({
+      vi.spyOn(edgeClient.config.getDefaultAgentConfig, 'query').mockResolvedValue({
         model: 'gemini-pro',
       });
 

package/src/services/__tests__/sync.test.ts

@@ -0,0 +1,56 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { dataSync } from '@/database/client/core';
+import { StartDataSyncParams } from '@/types/sync';
+
+import { syncService } from '../sync';
+
+vi.mock('@/database/client/core', () => ({
+  dataSync: {
+    startDataSync: vi.fn(),
+    disconnect: vi.fn(),
+  },
+}));
+
+describe('SyncService', () => {
+  afterEach(() => {
+    vi.resetAllMocks();
+  });
+
+  describe('enabledSync', () => {
+    it('should return false when running on server side', async () => {
+      const params = { user: { id: '123' }, authToken: 'abc' } as unknown as StartDataSyncParams;
+
+      const origin = global.window;
+
+      // @ts-ignore
+      global.window = undefined;
+
+      const result = await syncService.enabledSync(params);
+
+      expect(result).toBe(false);
+      expect(dataSync.startDataSync).not.toHaveBeenCalled();
+
+      // reset
+      global.window = origin;
+    });
+
+    it('should start data sync and return true when running on client side', async () => {
+      const params = { user: { id: '123' }, authToken: 'abc' } as unknown as StartDataSyncParams;
+
+      const result = await syncService.enabledSync(params);
+
+      expect(result).toBe(true);
+      expect(dataSync.startDataSync).toHaveBeenCalledWith(params);
+    });
+  });
+
+  describe('disableSync', () => {
+    it('should disconnect data sync and return false', async () => {
+      const result = await syncService.disableSync();
+
+      expect(result).toBe(false);
+      expect(dataSync.disconnect).toHaveBeenCalled();
+    });
+  });
+});

package/src/services/__tests__/upload.test.ts

@@ -0,0 +1,72 @@
+import { beforeAll, describe, expect, it, vi } from 'vitest';
+
+import { DB_File } from '@/database/client/schemas/files';
+import { edgeClient } from '@/libs/trpc/client';
+import { API_ENDPOINTS } from '@/services/_url';
+import { serverConfigSelectors } from '@/store/serverConfig/selectors';
+import { createServerConfigStore } from '@/store/serverConfig/store';
+
+import { uploadService } from '../upload';
+
+vi.mock('@/store/serverConfig/selectors');
+vi.mock('@/libs/trpc/client', () => {
+  return {
+    edgeClient: {
+      upload: {
+        createS3PreSignedUrl: { mutate: vi.fn() },
+      },
+    },
+  };
+});
+
+beforeAll(() => {
+  createServerConfigStore();
+});
+
+describe('UploadService', () => {
+  describe('uploadFile', () => {
+    it('should upload file to server when enableServer is true', async () => {
+      // Arrange
+      const file: DB_File = {
+        data: new ArrayBuffer(10),
+        fileType: 'text/plain',
+        metadata: {},
+        name: 'test.txt',
+        saveMode: 'local',
+        size: 10,
+      };
+      const mockCreateS3Url = vi.fn().mockResolvedValue('https://example.com');
+      vi.mocked(edgeClient.upload.createS3PreSignedUrl.mutate).mockImplementation(mockCreateS3Url);
+      vi.spyOn(serverConfigSelectors, 'enableUploadFileToServer').mockReturnValue(true);
+      global.fetch = vi.fn().mockResolvedValue({ ok: true } as Response);
+
+      // Act
+      const result = await uploadService.uploadFile(file);
+
+      // Assert
+      expect(mockCreateS3Url).toHaveBeenCalledTimes(1);
+      expect(fetch).toHaveBeenCalledTimes(1);
+      expect(result.url).toMatch(/\/\d+\/[a-f0-9-]+\.txt/);
+      expect(result.saveMode).toBe('url');
+    });
+
+    it('should save file locally when enableServer is false', async () => {
+      // Arrange
+      const file: DB_File = {
+        data: new ArrayBuffer(10),
+        fileType: 'text/plain',
+        metadata: {},
+        name: 'test.txt',
+        saveMode: 'local',
+        size: 10,
+      };
+      vi.spyOn(serverConfigSelectors, 'enableUploadFileToServer').mockReturnValue(false);
+
+      // Act
+      const result = await uploadService.uploadFile(file);
+
+      // Assert
+      expect(result).toEqual(file);
+    });
+  });
+});

package/src/services/_url.ts

@@ -1,3 +1,5 @@
+// TODO: 未来所有路由需要全部迁移到 trpc
+
 /* eslint-disable sort-keys-fix/sort-keys-fix */
 import { transform } from 'lodash-es';
 

package/src/services/file/client.test.ts

@@ -1,18 +1,35 @@
-import { Mock, beforeEach, describe, expect, it, vi } from 'vitest';
+import { Mock, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
 
+import { fileEnv } from '@/config/file';
 import { FileModel } from '@/database/client/models/file';
 import { DB_File } from '@/database/client/schemas/files';
+import { serverConfigSelectors } from '@/store/serverConfig/selectors';
+import { createServerConfigStore } from '@/store/serverConfig/store';
 
 import { ClientService } from './client';
 
 const fileService = new ClientService();
 
+beforeAll(() => {
+  createServerConfigStore();
+});
 // Mocks for the FileModel
 vi.mock('@/database/client/models/file', () => ({
   FileModel: {
     create: vi.fn(),
     delete: vi.fn(),
     findById: vi.fn(),
+    clear: vi.fn(),
+  },
+}));
+
+let s3Domain: string;
+
+vi.mock('@/config/file', () => ({
+  fileEnv: {
+    get NEXT_PUBLIC_S3_DOMAIN() {
+      return s3Domain;
+    },
   },
 }));
 
@@ -20,13 +37,14 @@ vi.mock('@/database/client/models/file', () => ({
 global.URL.createObjectURL = vi.fn();
 global.Blob = vi.fn();
 
-describe('FileService', () => {
-  beforeEach(() => {
-    // Reset all mocks before each test
-    vi.resetAllMocks();
-  });
+beforeEach(() => {
+  // Reset all mocks before each test
+  vi.resetAllMocks();
+  s3Domain = '';
+});
 
-  it('uploadFile should save the file to the database', async () => {
+describe('FileService', () => {
+  it('createFile should save the file to the database', async () => {
     const localFile: DB_File = {
       name: 'test',
       data: new ArrayBuffer(1),
@@ -37,7 +55,7 @@ describe('FileService', () => {
 
     (FileModel.create as Mock).mockResolvedValue(localFile);
 
-    const result = await fileService.uploadFile(localFile);
+    const result = await fileService.createFile(localFile);
 
     expect(FileModel.create).toHaveBeenCalledWith(localFile);
     expect(result).toEqual(localFile);
@@ -53,38 +71,88 @@ describe('FileService', () => {
     expect(result).toBe(true);
   });
 
-  it('getFile should retrieve and convert file info to FilePreview', async () => {
-    const fileId = '1';
-    const fileData: DB_File = {
-      name: 'test',
-      data: new ArrayBuffer(1),
-      fileType: 'image/png',
-      saveMode: 'local',
-      size: 1,
-    };
+  describe('getFile', () => {
+    it('should retrieve and convert local file info to FilePreview', async () => {
+      const fileId = '1';
+      const fileData: DB_File = {
+        name: 'test',
+        data: new ArrayBuffer(1),
+        fileType: 'image/png',
+        saveMode: 'local',
+        size: 1,
+      };
+
+      (FileModel.findById as Mock).mockResolvedValue(fileData);
+      (global.URL.createObjectURL as Mock).mockReturnValue('blob:test');
+      (global.Blob as Mock).mockImplementation(() => ['test']);
+
+      const result = await fileService.getFile(fileId);
+
+      expect(FileModel.findById).toHaveBeenCalledWith(fileId);
+      expect(result).toEqual({
+        base64Url: 'data:image/png;base64,AA==',
+        fileType: 'image/png',
+        name: 'test',
+        saveMode: 'local',
+        url: 'blob:test',
+      });
+    });
 
-    (FileModel.findById as Mock).mockResolvedValue(fileData);
-    (global.URL.createObjectURL as Mock).mockReturnValue('blob:test');
-    (global.Blob as Mock).mockImplementation(() => ['test']);
+    it('should retrieve and convert URL file info to FilePreview when enableServer is true', async () => {
+      const fileId = '1';
+      const fileData: DB_File = {
+        name: 'test',
+        fileType: 'image/png',
+        saveMode: 'url',
+        metadata: {
+          filename: 'test.png',
+        },
+        size: 1,
+        url: '/test.png',
+      };
+
+      (FileModel.findById as Mock).mockResolvedValue(fileData);
+      vi.spyOn(serverConfigSelectors, 'enableUploadFileToServer').mockReturnValue(true);
+      s3Domain = 'https://example.com';
+
+      const result = await fileService.getFile(fileId);
+
+      expect(FileModel.findById).toHaveBeenCalledWith(fileId);
+      expect(result).toEqual({
+        fileType: 'image/png',
+        name: 'test.png',
+        saveMode: 'url',
+        url: 'https://example.com/test.png',
+      });
+    });
 
-    const result = await fileService.getFile(fileId);
+    it('should throw an error when enableServer is true but NEXT_PUBLIC_S3_DOMAIN is not set', async () => {
+      const fileId = '1';
+      const fileData: DB_File = {
+        name: 'test',
+        fileType: 'image/png',
+        saveMode: 'url',
+        size: 1,
+        url: '/test.png',
+      };
 
-    expect(FileModel.findById).toHaveBeenCalledWith(fileId);
-    expect(result).toEqual({
-      base64Url: 'data:image/png;base64,AA==',
-      fileType: 'image/png',
-      name: 'test',
-      saveMode: 'local',
-      url: 'blob:test',
+      (FileModel.findById as Mock).mockResolvedValue(fileData);
+      vi.spyOn(serverConfigSelectors, 'enableUploadFileToServer').mockReturnValue(true);
+
+      const getFilePromise = fileService.getFile(fileId);
+
+      await expect(getFilePromise).rejects.toThrow(
+        'fileEnv.NEXT_PUBLIC_S3_DOMAIN is not set while enable server upload',
+      );
     });
-  });
 
-  it('getFile should throw an error when the file is not found', async () => {
-    const fileId = 'non-existent';
-    (FileModel.findById as Mock).mockResolvedValue(null);
+    it('should throw an error when the file is not found', async () => {
+      const fileId = 'non-existent';
+      (FileModel.findById as Mock).mockResolvedValue(null);
 
-    const getFilePromise = fileService.getFile(fileId);
+      const getFilePromise = fileService.getFile(fileId);
 
-    await expect(getFilePromise).rejects.toThrow('file not found');
+      await expect(getFilePromise).rejects.toThrow('file not found');
+    });
   });
 });