@lobehub/chat 1.98.2 → 1.99.0

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

@@ -0,0 +1,267 @@
+import { TRPCError } from '@trpc/server';
+import { describe, expect, it, vi } from 'vitest';
+
+import { AsyncTaskModel } from '@/database/models/asyncTask';
+import { GenerationModel } from '@/database/models/generation';
+import { FileService } from '@/server/services/file';
+import { AsyncTaskStatus } from '@/types/asyncTask';
+
+import { generationRouter } from './generation';
+
+vi.mock('@/database/models/asyncTask');
+vi.mock('@/database/models/generation');
+vi.mock('@/server/services/file');
+
+describe('generationRouter', () => {
+  const mockCtx = {
+    userId: 'test-user',
+  };
+
+  describe('getGenerationStatus', () => {
+    it('should return generation status when task is successful', async () => {
+      const mockGeneration = {
+        id: 'gen-1',
+        asset: { url: 'https://example.com/image.jpg' },
+      };
+      const mockAsyncTask = {
+        id: 'task-1',
+        status: AsyncTaskStatus.Success,
+        error: null,
+      };
+      const mockCheckTimeoutTasks = vi.fn().mockResolvedValue(undefined);
+      const mockFindById = vi.fn().mockResolvedValue(mockAsyncTask);
+      const mockFindByIdAndTransform = vi.fn().mockResolvedValue(mockGeneration);
+
+      vi.mocked(AsyncTaskModel).mockImplementation(
+        () =>
+          ({
+            checkTimeoutTasks: mockCheckTimeoutTasks,
+            findById: mockFindById,
+          }) as any,
+      );
+      vi.mocked(GenerationModel).mockImplementation(
+        () =>
+          ({
+            findByIdAndTransform: mockFindByIdAndTransform,
+          }) as any,
+      );
+
+      const caller = generationRouter.createCaller(mockCtx);
+
+      const result = await caller.getGenerationStatus({
+        generationId: 'gen-1',
+        asyncTaskId: 'task-1',
+      });
+
+      expect(result.status).toBe(AsyncTaskStatus.Success);
+      expect(result.generation).toEqual(mockGeneration);
+      expect(result.error).toBeNull();
+      expect(mockCheckTimeoutTasks).toHaveBeenCalledWith(['task-1']);
+      expect(mockFindById).toHaveBeenCalledWith('task-1');
+      expect(mockFindByIdAndTransform).toHaveBeenCalledWith('gen-1');
+    });
+
+    it('should return error when task failed', async () => {
+      const mockError = { code: 'GENERATION_ERROR', message: 'Generation failed' };
+      const mockAsyncTask = {
+        id: 'task-1',
+        status: AsyncTaskStatus.Error,
+        error: mockError,
+      };
+      const mockCheckTimeoutTasks = vi.fn().mockResolvedValue(undefined);
+      const mockFindById = vi.fn().mockResolvedValue(mockAsyncTask);
+
+      vi.mocked(AsyncTaskModel).mockImplementation(
+        () =>
+          ({
+            checkTimeoutTasks: mockCheckTimeoutTasks,
+            findById: mockFindById,
+          }) as any,
+      );
+
+      const caller = generationRouter.createCaller(mockCtx);
+
+      const result = await caller.getGenerationStatus({
+        generationId: 'gen-1',
+        asyncTaskId: 'task-1',
+      });
+
+      expect(result.status).toBe(AsyncTaskStatus.Error);
+      expect(result.generation).toBeNull();
+      expect(result.error).toEqual(mockError);
+    });
+
+    it('should return pending status when task is running', async () => {
+      const mockAsyncTask = {
+        id: 'task-1',
+        status: AsyncTaskStatus.Pending,
+        error: null,
+      };
+      const mockCheckTimeoutTasks = vi.fn().mockResolvedValue(undefined);
+      const mockFindById = vi.fn().mockResolvedValue(mockAsyncTask);
+
+      vi.mocked(AsyncTaskModel).mockImplementation(
+        () =>
+          ({
+            checkTimeoutTasks: mockCheckTimeoutTasks,
+            findById: mockFindById,
+          }) as any,
+      );
+
+      const caller = generationRouter.createCaller(mockCtx);
+
+      const result = await caller.getGenerationStatus({
+        generationId: 'gen-1',
+        asyncTaskId: 'task-1',
+      });
+
+      expect(result.status).toBe(AsyncTaskStatus.Pending);
+      expect(result.generation).toBeNull();
+      expect(result.error).toBeNull();
+    });
+
+    it('should throw error when async task not found', async () => {
+      const mockCheckTimeoutTasks = vi.fn().mockResolvedValue(undefined);
+      const mockFindById = vi.fn().mockResolvedValue(null);
+
+      vi.mocked(AsyncTaskModel).mockImplementation(
+        () =>
+          ({
+            checkTimeoutTasks: mockCheckTimeoutTasks,
+            findById: mockFindById,
+          }) as any,
+      );
+
+      const caller = generationRouter.createCaller(mockCtx);
+
+      await expect(
+        caller.getGenerationStatus({
+          generationId: 'gen-1',
+          asyncTaskId: 'task-1',
+        }),
+      ).rejects.toThrow(TRPCError);
+    });
+
+    it('should throw error when generation not found for successful task', async () => {
+      const mockAsyncTask = {
+        id: 'task-1',
+        status: AsyncTaskStatus.Success,
+        error: null,
+      };
+      const mockCheckTimeoutTasks = vi.fn().mockResolvedValue(undefined);
+      const mockFindById = vi.fn().mockResolvedValue(mockAsyncTask);
+      const mockFindByIdAndTransform = vi.fn().mockResolvedValue(null);
+
+      vi.mocked(AsyncTaskModel).mockImplementation(
+        () =>
+          ({
+            checkTimeoutTasks: mockCheckTimeoutTasks,
+            findById: mockFindById,
+          }) as any,
+      );
+      vi.mocked(GenerationModel).mockImplementation(
+        () =>
+          ({
+            findByIdAndTransform: mockFindByIdAndTransform,
+          }) as any,
+      );
+
+      const caller = generationRouter.createCaller(mockCtx);
+
+      await expect(
+        caller.getGenerationStatus({
+          generationId: 'gen-1',
+          asyncTaskId: 'task-1',
+        }),
+      ).rejects.toThrow(TRPCError);
+    });
+  });
+
+  describe('deleteGeneration', () => {
+    it('should delete generation with thumbnail', async () => {
+      const mockDeletedGeneration = {
+        id: 'gen-1',
+        asset: { thumbnailUrl: 'thumb-key' },
+      };
+      const mockDelete = vi.fn().mockResolvedValue(mockDeletedGeneration);
+      const mockDeleteFile = vi.fn().mockResolvedValue(true);
+
+      vi.mocked(GenerationModel).mockImplementation(
+        () =>
+          ({
+            delete: mockDelete,
+          }) as any,
+      );
+      vi.mocked(FileService).mockImplementation(
+        () =>
+          ({
+            deleteFile: mockDeleteFile,
+          }) as any,
+      );
+
+      const caller = generationRouter.createCaller(mockCtx);
+
+      const result = await caller.deleteGeneration({ generationId: 'gen-1' });
+
+      expect(result).toEqual(mockDeletedGeneration);
+      expect(mockDelete).toHaveBeenCalledWith('gen-1');
+      expect(mockDeleteFile).toHaveBeenCalledWith('thumb-key');
+    });
+
+    it('should delete generation without thumbnail', async () => {
+      const mockDeletedGeneration = {
+        id: 'gen-1',
+        asset: { url: 'main-url' },
+      };
+      const mockDelete = vi.fn().mockResolvedValue(mockDeletedGeneration);
+      const mockDeleteFile = vi.fn().mockResolvedValue(true);
+
+      vi.mocked(GenerationModel).mockImplementation(
+        () =>
+          ({
+            delete: mockDelete,
+          }) as any,
+      );
+      vi.mocked(FileService).mockImplementation(
+        () =>
+          ({
+            deleteFile: mockDeleteFile,
+          }) as any,
+      );
+
+      const caller = generationRouter.createCaller(mockCtx);
+
+      const result = await caller.deleteGeneration({ generationId: 'gen-1' });
+
+      expect(result).toEqual(mockDeletedGeneration);
+      expect(mockDelete).toHaveBeenCalledWith('gen-1');
+      expect(mockDeleteFile).not.toHaveBeenCalled();
+    });
+
+    it('should handle when generation not found', async () => {
+      const mockDelete = vi.fn().mockResolvedValue(null);
+      const mockDeleteFile = vi.fn().mockResolvedValue(true);
+
+      vi.mocked(GenerationModel).mockImplementation(
+        () =>
+          ({
+            delete: mockDelete,
+          }) as any,
+      );
+      vi.mocked(FileService).mockImplementation(
+        () =>
+          ({
+            deleteFile: mockDeleteFile,
+          }) as any,
+      );
+
+      const caller = generationRouter.createCaller(mockCtx);
+
+      const result = await caller.deleteGeneration({ generationId: 'gen-1' });
+
+      expect(result).toBeUndefined();
+      expect(mockDelete).toHaveBeenCalledWith('gen-1');
+      expect(mockDeleteFile).not.toHaveBeenCalled();
+    });
+  });
+});

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

@@ -0,0 +1,86 @@
+import { TRPCError } from '@trpc/server';
+import { z } from 'zod';
+
+import { AsyncTaskModel } from '@/database/models/asyncTask';
+import { GenerationModel } from '@/database/models/generation';
+import { authedProcedure, router } from '@/libs/trpc/lambda';
+import { serverDatabase } from '@/libs/trpc/lambda/middleware';
+import { FileService } from '@/server/services/file';
+import { AsyncTaskError, AsyncTaskStatus } from '@/types/asyncTask';
+import { Generation } from '@/types/generation';
+
+const generationProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
+  const { ctx } = opts;
+
+  return opts.next({
+    ctx: {
+      asyncTaskModel: new AsyncTaskModel(ctx.serverDB, ctx.userId),
+      fileService: new FileService(ctx.serverDB, ctx.userId),
+      generationModel: new GenerationModel(ctx.serverDB, ctx.userId),
+    },
+  });
+});
+
+export type GetGenerationStatusResult = {
+  error: AsyncTaskError | null;
+  generation: Generation | null;
+  status: AsyncTaskStatus;
+};
+
+export const generationRouter = router({
+  deleteGeneration: generationProcedure
+    .input(z.object({ generationId: z.string() }))
+    .mutation(async ({ ctx, input }) => {
+      // Delete the generation record from database and get the deleted data
+      const deletedGeneration = await ctx.generationModel.delete(input.generationId);
+
+      if (!deletedGeneration) return;
+
+      // Note: Based on new requirements, don't delete main file (fileId), only delete thumbnail
+      // If generation has a thumbnail, delete it from S3
+      if (deletedGeneration.asset) {
+        const asset = deletedGeneration.asset as any;
+
+        // Only delete thumbnail URL if exists
+        if (asset.thumbnailUrl) {
+          await ctx.fileService.deleteFile(asset.thumbnailUrl);
+        }
+      }
+
+      return deletedGeneration;
+    }),
+
+  getGenerationStatus: generationProcedure
+    .input(z.object({ asyncTaskId: z.string(), generationId: z.string() }))
+    .query(async ({ ctx, input }) => {
+      // Check for timeout tasks before querying
+      await ctx.asyncTaskModel.checkTimeoutTasks([input.asyncTaskId]);
+
+      const asyncTask = await ctx.asyncTaskModel.findById(input.asyncTaskId);
+      if (!asyncTask) {
+        throw new TRPCError({ code: 'NOT_FOUND', message: 'Async task not found' });
+      }
+
+      const { status, error } = asyncTask;
+      const result: GetGenerationStatusResult = {
+        error: null,
+        generation: null,
+        status: status as AsyncTaskStatus,
+      };
+
+      if (asyncTask.status === AsyncTaskStatus.Success) {
+        const generation = await ctx.generationModel.findByIdAndTransform(input.generationId);
+        if (!generation) {
+          throw new TRPCError({ code: 'NOT_FOUND', message: 'Generation not found' });
+        }
+
+        result.generation = generation;
+      } else if (asyncTask.status === AsyncTaskStatus.Error) {
+        result.error = error as AsyncTaskError;
+      }
+
+      return result;
+    }),
+});
+
+export type GenerationRouter = typeof generationRouter;