♻️ refactor: refactor upload router into lambda and decide to remove it in V2 (#9766)

refactor edge router and decide to remove it in V2

src/app/(backend)/trpc/edge/[trpc]/route.ts

@@ -5,8 +5,6 @@ import { pino } from '@/libs/logger';
 import { createEdgeContext } from '@/libs/trpc/edge/context';
 import { edgeRouter } from '@/server/routers/edge';
 
-export const runtime = 'edge';
-
 const handler = (req: NextRequest) =>
   fetchRequestHandler({
     /**

src/server/routers/edge/index.ts

@@ -1,5 +1,6 @@
 /**
- * This file contains the edge router of Lobe Chat tRPC-backend
+ * @deprecated
+ * TODO: it will be remove in V2.0
  */
 import { publicProcedure, router } from '@/libs/trpc/edge';
 

src/server/routers/lambda/index.ts

@@ -28,6 +28,7 @@ import { sessionRouter } from './session';
 import { sessionGroupRouter } from './sessionGroup';
 import { threadRouter } from './thread';
 import { topicRouter } from './topic';
+import { uploadRouter } from './upload';
 import { userRouter } from './user';
 
 export const lambdaRouter = router({
@@ -57,6 +58,7 @@ export const lambdaRouter = router({
   sessionGroup: sessionGroupRouter,
   thread: threadRouter,
   topic: topicRouter,
+  upload: uploadRouter,
   user: userRouter,
 });
 

src/server/routers/lambda/upload.ts

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

src/services/__tests__/upload.test.ts

@@ -1,15 +1,28 @@
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { fileEnv } from '@/envs/file';
-import { edgeClient } from '@/libs/trpc/client';
+import { lambdaClient } from '@/libs/trpc/client';
 import { API_ENDPOINTS } from '@/services/_url';
 import { clientS3Storage } from '@/services/file/ClientS3';
 
 import { UPLOAD_NETWORK_ERROR, uploadService } from '../upload';
 
 // Mock dependencies
+vi.mock('@lobechat/const', () => ({
+  isDesktop: false,
+  isServerMode: false,
+}));
+
+vi.mock('@lobechat/model-runtime', () => ({
+  parseDataUri: vi.fn(),
+}));
+
+vi.mock('@lobechat/utils', () => ({
+  uuid: () => 'mock-uuid',
+}));
+
 vi.mock('@/libs/trpc/client', () => ({
-  edgeClient: {
+  lambdaClient: {
     upload: {
       createS3PreSignedUrl: {
         mutate: vi.fn(),
@@ -24,8 +37,24 @@ vi.mock('@/services/file/ClientS3', () => ({
   },
 }));
 
-vi.mock('@/utils/uuid', () => ({
-  uuid: () => 'mock-uuid',
+vi.mock('@/store/electron', () => ({
+  getElectronStoreState: vi.fn(() => ({})),
+}));
+
+vi.mock('@/store/electron/selectors', () => ({
+  electronSyncSelectors: {
+    isSyncActive: vi.fn(() => false),
+  },
+}));
+
+vi.mock('@/services/electron/file', () => ({
+  desktopFileAPI: {
+    uploadFile: vi.fn(),
+  },
+}));
+
+vi.mock('js-sha256', () => ({
+  sha256: vi.fn((data) => 'mock-hash-' + data.byteLength),
 }));
 
 describe('UploadService', () => {
@@ -38,23 +67,174 @@ describe('UploadService', () => {
     vi.spyOn(Date, 'now').mockImplementation(() => 3600000); // 1 hour in milliseconds
   });
 
-  describe('uploadWithProgress', () => {
+  describe('uploadFileToS3', () => {
+    it('should upload to client S3 for non-server mode with image file', async () => {
+      const { sha256 } = await import('js-sha256');
+      vi.mocked(sha256).mockReturnValue('test-hash');
+      vi.mocked(clientS3Storage.putObject).mockResolvedValue(undefined);
+
+      const result = await uploadService.uploadFileToS3(mockFile, {});
+
+      expect(result.success).toBe(true);
+      expect(result.data).toEqual({
+        date: '1',
+        dirname: '',
+        filename: mockFile.name,
+        path: 'client-s3://test-hash',
+      });
+      expect(clientS3Storage.putObject).toHaveBeenCalledWith('test-hash', mockFile);
+    });
+
+    it('should call onNotSupported for non-image/video files', async () => {
+      const nonImageFile = new File(['test'], 'test.txt', { type: 'text/plain' });
+      const onNotSupported = vi.fn();
+
+      const result = await uploadService.uploadFileToS3(nonImageFile, {
+        onNotSupported,
+      });
+
+      expect(result.success).toBe(false);
+      expect(onNotSupported).toHaveBeenCalled();
+    });
+
+    it('should skip file type check when skipCheckFileType is true', async () => {
+      const nonImageFile = new File(['test'], 'test.txt', { type: 'text/plain' });
+      const { sha256 } = await import('js-sha256');
+      vi.mocked(sha256).mockReturnValue('test-hash');
+      vi.mocked(clientS3Storage.putObject).mockResolvedValue(undefined);
+
+      const result = await uploadService.uploadFileToS3(nonImageFile, {
+        skipCheckFileType: true,
+      });
+
+      expect(result.success).toBe(true);
+      expect(clientS3Storage.putObject).toHaveBeenCalled();
+    });
+
+    it('should upload video files', async () => {
+      const videoFile = new File(['test'], 'test.mp4', { type: 'video/mp4' });
+      const { sha256 } = await import('js-sha256');
+      vi.mocked(sha256).mockReturnValue('video-hash');
+      vi.mocked(clientS3Storage.putObject).mockResolvedValue(undefined);
+
+      const result = await uploadService.uploadFileToS3(videoFile, {});
+
+      expect(result.success).toBe(true);
+      expect(clientS3Storage.putObject).toHaveBeenCalledWith('video-hash', videoFile);
+    });
+  });
+
+  describe('uploadBase64ToS3', () => {
+    it('should upload base64 data successfully', async () => {
+      const { parseDataUri } = await import('@lobechat/model-runtime');
+      vi.mocked(parseDataUri).mockReturnValueOnce({
+        base64: 'dGVzdA==', // "test" in base64
+        mimeType: 'image/png',
+        type: 'base64',
+      });
+
+      const { sha256 } = await import('js-sha256');
+      vi.mocked(sha256).mockReturnValue('base64-hash');
+      vi.mocked(clientS3Storage.putObject).mockResolvedValue(undefined);
+
+      const base64Data = 'data:image/png;base64,dGVzdA==';
+      const result = await uploadService.uploadBase64ToS3(base64Data);
+
+      expect(result).toMatchObject({
+        fileType: 'image/png',
+        hash: expect.any(String),
+        metadata: expect.objectContaining({
+          path: expect.stringContaining('client-s3://'),
+        }),
+        size: expect.any(Number),
+      });
+    });
+
+    it('should throw error for invalid base64 data', async () => {
+      const { parseDataUri } = await import('@lobechat/model-runtime');
+      vi.mocked(parseDataUri).mockReturnValueOnce({
+        base64: null,
+        mimeType: null,
+        type: 'url',
+      });
+
+      const invalidBase64 = 'not-a-base64-string';
+
+      await expect(uploadService.uploadBase64ToS3(invalidBase64)).rejects.toThrow(
+        'Invalid base64 data for image',
+      );
+    });
+
+    it('should use custom filename when provided', async () => {
+      const { parseDataUri } = await import('@lobechat/model-runtime');
+      vi.mocked(parseDataUri).mockReturnValueOnce({
+        base64: 'dGVzdA==',
+        mimeType: 'image/png',
+        type: 'base64',
+      });
+
+      const { sha256 } = await import('js-sha256');
+      vi.mocked(sha256).mockReturnValue('custom-hash');
+      vi.mocked(clientS3Storage.putObject).mockResolvedValue(undefined);
+
+      const base64Data = 'data:image/png;base64,dGVzdA==';
+      const result = await uploadService.uploadBase64ToS3(base64Data, {
+        filename: 'custom-image',
+      });
+
+      expect(result.metadata.filename).toContain('custom-image');
+    });
+  });
+
+  describe('uploadDataToS3', () => {
+    it('should upload JSON data successfully', async () => {
+      const { sha256 } = await import('js-sha256');
+      vi.mocked(sha256).mockReturnValue('json-hash');
+      vi.mocked(clientS3Storage.putObject).mockResolvedValue(undefined);
+
+      const data = { key: 'value', number: 123 };
+      // uploadDataToS3 internally calls uploadFileToS3, which needs skipCheckFileType for JSON
+      const result = await uploadService.uploadDataToS3(data, {
+        skipCheckFileType: true,
+      });
+
+      expect(result.success).toBe(true);
+      expect(clientS3Storage.putObject).toHaveBeenCalled();
+    });
+
+    it('should use custom filename when provided', async () => {
+      const { sha256 } = await import('js-sha256');
+      vi.mocked(sha256).mockReturnValue('custom-json-hash');
+      vi.mocked(clientS3Storage.putObject).mockResolvedValue(undefined);
+
+      const data = { test: true };
+      const result = await uploadService.uploadDataToS3(data, {
+        filename: 'custom.json',
+        skipCheckFileType: true,
+      });
+
+      expect(result.success).toBe(true);
+      expect(result.data.filename).toBe('custom.json');
+    });
+  });
+
+  describe('uploadToServerS3', () => {
     beforeEach(() => {
       // Mock XMLHttpRequest
       const xhrMock = {
-        upload: {
-          addEventListener: vi.fn(),
-        },
+        addEventListener: vi.fn(),
         open: vi.fn(),
         send: vi.fn(),
         setRequestHeader: vi.fn(),
-        addEventListener: vi.fn(),
         status: 200,
+        upload: {
+          addEventListener: vi.fn(),
+        },
       };
       global.XMLHttpRequest = vi.fn(() => xhrMock) as any;
 
       // Mock createS3PreSignedUrl
-      (edgeClient.upload.createS3PreSignedUrl.mutate as any).mockResolvedValue(mockPreSignUrl);
+      vi.mocked(lambdaClient.upload.createS3PreSignedUrl.mutate).mockResolvedValue(mockPreSignUrl);
     });
 
     it('should upload file successfully with progress', async () => {
@@ -64,7 +244,7 @@ describe('UploadService', () => {
       // Simulate successful upload
       vi.spyOn(xhr, 'addEventListener').mockImplementation((event, handler) => {
         if (event === 'load') {
-          // @ts-ignore
+          // @ts-expect-error - mock implementation
           handler({ target: { status: 200 } });
         }
       });
@@ -79,6 +259,41 @@ describe('UploadService', () => {
       });
     });
 
+    it('should report progress during upload', async () => {
+      const onProgress = vi.fn();
+      const xhr = new XMLHttpRequest();
+
+      // Simulate progress events
+      vi.spyOn(xhr.upload, 'addEventListener').mockImplementation((event, handler) => {
+        if (event === 'progress') {
+          // @ts-expect-error - mock implementation
+          handler({
+            lengthComputable: true,
+            loaded: 500,
+            total: 1000,
+          });
+        }
+      });
+
+      vi.spyOn(xhr, 'addEventListener').mockImplementation((event, handler) => {
+        if (event === 'load') {
+          // @ts-expect-error - mock implementation
+          handler({ target: { status: 200 } });
+        }
+      });
+
+      await uploadService.uploadToServerS3(mockFile, { onProgress });
+
+      expect(onProgress).toHaveBeenCalledWith(
+        'uploading',
+        expect.objectContaining({
+          progress: expect.any(Number),
+          restTime: expect.any(Number),
+          speed: expect.any(Number),
+        }),
+      );
+    });
+
     it('should handle network error', async () => {
       const xhr = new XMLHttpRequest();
 
@@ -86,7 +301,7 @@ describe('UploadService', () => {
       vi.spyOn(xhr, 'addEventListener').mockImplementation((event, handler) => {
         if (event === 'error') {
           Object.assign(xhr, { status: 0 });
-          // @ts-ignore
+          // @ts-expect-error - mock implementation
           handler({});
         }
       });
@@ -102,13 +317,46 @@ describe('UploadService', () => {
         if (event === 'load') {
           Object.assign(xhr, { status: 400, statusText: 'Bad Request' });
 
-          // @ts-ignore
+          // @ts-expect-error - mock implementation
           handler({});
         }
       });
 
       await expect(uploadService.uploadToServerS3(mockFile, {})).rejects.toBe('Bad Request');
     });
+
+    it('should use custom directory when provided', async () => {
+      const xhr = new XMLHttpRequest();
+      vi.spyOn(xhr, 'addEventListener').mockImplementation((event, handler) => {
+        if (event === 'load') {
+          // @ts-expect-error - mock implementation
+          handler({ target: { status: 200 } });
+        }
+      });
+
+      const result = await uploadService.uploadToServerS3(mockFile, {
+        directory: 'custom/dir',
+      });
+
+      expect(result.dirname).toContain('custom/dir');
+    });
+
+    it('should use custom pathname when provided', async () => {
+      const xhr = new XMLHttpRequest();
+      vi.spyOn(xhr, 'addEventListener').mockImplementation((event, handler) => {
+        if (event === 'load') {
+          // @ts-expect-error - mock implementation
+          handler({ target: { status: 200 } });
+        }
+      });
+
+      const customPath = 'custom/path/file.png';
+      const result = await uploadService.uploadToServerS3(mockFile, {
+        pathname: customPath,
+      });
+
+      expect(result.path).toBe(customPath);
+    });
   });
 
   describe('uploadToClientS3', () => {
@@ -121,7 +369,7 @@ describe('UploadService', () => {
         path: `client-s3://${hash}`,
       };
 
-      (clientS3Storage.putObject as any).mockResolvedValue(undefined);
+      vi.mocked(clientS3Storage.putObject).mockResolvedValue(undefined);
 
       const result = await uploadService['uploadToClientS3'](hash, mockFile);
 
@@ -140,9 +388,9 @@ describe('UploadService', () => {
       const filename = 'test.png';
       const mockArrayBuffer = new ArrayBuffer(8);
 
-      (global.fetch as any).mockResolvedValue({
+      vi.mocked(global.fetch).mockResolvedValue({
         arrayBuffer: () => Promise.resolve(mockArrayBuffer),
-      });
+      } as Response);
 
       const result = await uploadService.getImageFileByUrlWithCORS(url, filename);
 
@@ -161,9 +409,9 @@ describe('UploadService', () => {
       const fileType = 'image/jpeg';
       const mockArrayBuffer = new ArrayBuffer(8);
 
-      (global.fetch as any).mockResolvedValue({
+      vi.mocked(global.fetch).mockResolvedValue({
         arrayBuffer: () => Promise.resolve(mockArrayBuffer),
-      });
+      } as Response);
 
       const result = await uploadService.getImageFileByUrlWithCORS(url, filename, fileType);
 

src/services/upload.ts

@@ -5,7 +5,7 @@ import dayjs from 'dayjs';
 import { sha256 } from 'js-sha256';
 
 import { fileEnv } from '@/envs/file';
-import { edgeClient } from '@/libs/trpc/client';
+import { lambdaClient } from '@/libs/trpc/client';
 import { API_ENDPOINTS } from '@/services/_url';
 import { clientS3Storage } from '@/services/file/ClientS3';
 import { FileMetadata, UploadBase64ToS3Result } from '@/types/files';
@@ -264,7 +264,7 @@ class UploadService {
     // 生成文件路径元数据
     const { date, dirname, filename, pathname } = generateFilePathMetadata(file.name, options);
 
-    const preSignUrl = await edgeClient.upload.createS3PreSignedUrl.mutate({ pathname });
+    const preSignUrl = await lambdaClient.upload.createS3PreSignedUrl.mutate({ pathname });
 
     return {
       date,