@lobehub/chat 1.98.2 → 1.99.0

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

@@ -0,0 +1,248 @@
+import debug from 'debug';
+import { and, eq } from 'drizzle-orm';
+import { z } from 'zod';
+
+import { AsyncTaskModel } from '@/database/models/asyncTask';
+import {
+  NewGeneration,
+  NewGenerationBatch,
+  asyncTasks,
+  generationBatches,
+  generations,
+} from '@/database/schemas';
+import { authedProcedure, router } from '@/libs/trpc/lambda';
+import { keyVaults, serverDatabase } from '@/libs/trpc/lambda/middleware';
+import { createAsyncCaller } from '@/server/routers/async/caller';
+import { FileService } from '@/server/services/file';
+import {
+  AsyncTaskError,
+  AsyncTaskErrorType,
+  AsyncTaskStatus,
+  AsyncTaskType,
+} from '@/types/asyncTask';
+import { generateUniqueSeeds } from '@/utils/number';
+
+const log = debug('lobe-image:lambda');
+
+const imageProcedure = authedProcedure
+  .use(keyVaults)
+  .use(serverDatabase)
+  .use(async (opts) => {
+    const { ctx } = opts;
+
+    const { apiKey } = ctx.jwtPayload;
+    if (apiKey) {
+      log('API key found in jwtPayload: %s', apiKey);
+    } else {
+      log('No API key found in jwtPayload');
+    }
+
+    return opts.next({
+      ctx: {
+        asyncTaskModel: new AsyncTaskModel(ctx.serverDB, ctx.userId),
+        fileService: new FileService(ctx.serverDB, ctx.userId),
+      },
+    });
+  });
+
+const createImageInputSchema = z.object({
+  generationTopicId: z.string(),
+  imageNum: z.number(),
+  model: z.string(),
+  params: z
+    .object({
+      cfg: z.number().optional(),
+      height: z.number().optional(),
+      imageUrls: z.array(z.string()).optional(),
+      prompt: z.string(),
+      seed: z.number().nullable().optional(),
+      steps: z.number().optional(),
+      width: z.number().optional(),
+    })
+    .passthrough(),
+  provider: z.string(),
+});
+export type CreateImageServicePayload = z.infer<typeof createImageInputSchema>;
+
+export const imageRouter = router({
+  createImage: imageProcedure.input(createImageInputSchema).mutation(async ({ input, ctx }) => {
+    const { userId, serverDB, asyncTaskModel, fileService } = ctx;
+    const { generationTopicId, provider, model, imageNum, params } = input;
+
+    log('Starting image creation process, input: %O', input);
+
+    // 如果 params 中包含 imageUrls，将它们转换为 S3 keys 用于数据库存储
+    let configForDatabase = { ...params };
+    if (Array.isArray(params.imageUrls) && params.imageUrls.length > 0) {
+      log('Converting imageUrls to S3 keys for database storage: %O', params.imageUrls);
+      try {
+        const imageKeys = params.imageUrls.map((url) => {
+          const key = fileService.getKeyFromFullUrl(url);
+          log('Converted URL %s to key %s', url, key);
+          return key;
+        });
+
+        // 将转换后的 keys 存储为数据库配置
+        configForDatabase = {
+          ...params,
+          imageUrls: imageKeys,
+        };
+        log('Successfully converted imageUrls to keys for database: %O', imageKeys);
+      } catch (error) {
+        log('Error converting imageUrls to keys: %O', error);
+        // 如果转换失败，保持原始 URLs（可能是本地文件或其他格式）
+        log('Keeping original imageUrls due to conversion error');
+      }
+    }
+
+    // 步骤 1: 在事务中原子性地创建所有数据库记录
+    const { batch: createdBatch, generationsWithTasks } = await serverDB.transaction(async (tx) => {
+      log('Starting database transaction for image generation');
+
+      // 1. 创建 generationBatch
+      const newBatch: NewGenerationBatch = {
+        config: configForDatabase,
+        generationTopicId,
+        height: params.height,
+        model,
+        prompt: params.prompt,
+        provider,
+        userId,
+        width: params.width, // 使用转换后的配置存储到数据库
+      };
+      log('Creating generation batch: %O', newBatch);
+      const [batch] = await tx.insert(generationBatches).values(newBatch).returning();
+      log('Generation batch created successfully: %s', batch.id);
+
+      // 2. 创建 4 个 generation（一期固定生成 4 张）
+      const seeds =
+        'seed' in params
+          ? generateUniqueSeeds(imageNum)
+          : Array.from({ length: imageNum }, () => null);
+      const newGenerations: NewGeneration[] = Array.from({ length: imageNum }, (_, index) => {
+        return {
+          generationBatchId: batch.id,
+          seed: seeds[index],
+          userId,
+        };
+      });
+
+      log('Creating %d generations for batch: %s', newGenerations.length, batch.id);
+      const createdGenerations = await tx.insert(generations).values(newGenerations).returning();
+      log(
+        'Generations created successfully: %O',
+        createdGenerations.map((g) => g.id),
+      );
+
+      // 3. 并发为每个 generation 创建 asyncTask（在事务中）
+      log('Creating async tasks for generations');
+      const generationsWithTasks = await Promise.all(
+        createdGenerations.map(async (generation) => {
+          // 在事务中直接创建 asyncTask
+          const [createdAsyncTask] = await tx
+            .insert(asyncTasks)
+            .values({
+              status: AsyncTaskStatus.Pending,
+              type: AsyncTaskType.ImageGeneration,
+              userId,
+            })
+            .returning();
+
+          const asyncTaskId = createdAsyncTask.id;
+          log('Created async task %s for generation %s', asyncTaskId, generation.id);
+
+          // 更新 generation 的 asyncTaskId
+          await tx
+            .update(generations)
+            .set({ asyncTaskId })
+            .where(and(eq(generations.id, generation.id), eq(generations.userId, userId)));
+
+          return { asyncTaskId, generation };
+        }),
+      );
+      log('All async tasks created in transaction');
+
+      return {
+        batch,
+        generationsWithTasks,
+      };
+    });
+
+    log('Database transaction completed successfully. Starting async task triggers directly.');
+
+    // 步骤 2: 直接执行所有生图任务（去掉 after 包装）
+    log('Starting async image generation tasks directly');
+
+    try {
+      log('Creating unified async caller for userId: %s', userId);
+      log(
+        'Lambda context - userId: %s, jwtPayload keys: %O',
+        ctx.userId,
+        Object.keys(ctx.jwtPayload || {}),
+      );
+
+      // 使用统一的 caller 工厂创建 caller
+      const asyncCaller = await createAsyncCaller({
+        jwtPayload: ctx.jwtPayload,
+        userId: ctx.userId,
+      });
+
+      log('Unified async caller created successfully for userId: %s', ctx.userId);
+      log('Processing %d async image generation tasks', generationsWithTasks.length);
+
+      // 启动所有图像生成任务（不等待完成，真正的后台任务）
+      generationsWithTasks.forEach(({ generation, asyncTaskId }) => {
+        log('Starting background async task %s for generation %s', asyncTaskId, generation.id);
+
+        // 不使用 await，让任务在后台异步执行
+        // 这里不应该 await 也不应该 .then.catch，让 runtime 早点释放计算资源
+        asyncCaller.image.createImage({
+          generationId: generation.id,
+          model,
+          params,
+          provider,
+          taskId: asyncTaskId, // 使用原始参数
+        });
+      });
+
+      log('All %d background async image generation tasks started', generationsWithTasks.length);
+    } catch (e) {
+      console.error('[createImage] Failed to process async tasks:', e);
+      log('Failed to process async tasks: %O', e);
+
+      // 如果整体失败，更新所有任务状态为失败
+      try {
+        await Promise.allSettled(
+          generationsWithTasks.map(({ asyncTaskId }) =>
+            asyncTaskModel.update(asyncTaskId, {
+              error: new AsyncTaskError(
+                AsyncTaskErrorType.ServerError,
+                'start async task error: ' + (e instanceof Error ? e.message : 'Unknown error'),
+              ),
+              status: AsyncTaskStatus.Error,
+            }),
+          ),
+        );
+      } catch (batchUpdateError) {
+        console.error('Failed to update batch task statuses:', batchUpdateError);
+      }
+    }
+
+    const createdGenerations = generationsWithTasks.map((item) => item.generation);
+    log('Image creation process completed successfully: %O', {
+      batchId: createdBatch.id,
+      generationCount: createdGenerations.length,
+      generationIds: createdGenerations.map((g) => g.id),
+    });
+
+    return {
+      data: {
+        batch: createdBatch,
+        generations: createdGenerations,
+      },
+      success: true,
+    };
+  }),
+});
+
+export type ImageRouter = typeof imageRouter;

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

@@ -11,6 +11,10 @@ import { configRouter } from './config';
 import { documentRouter } from './document';
 import { exporterRouter } from './exporter';
 import { fileRouter } from './file';
+import { generationRouter } from './generation';
+import { generationBatchRouter } from './generationBatch';
+import { generationTopicRouter } from './generationTopic';
+import { imageRouter } from './image';
 import { importerRouter } from './importer';
 import { knowledgeBaseRouter } from './knowledgeBase';
 import { marketRouter } from './market';
@@ -32,7 +36,11 @@ export const lambdaRouter = router({
   document: documentRouter,
   exporter: exporterRouter,
   file: fileRouter,
+  generation: generationRouter,
+  generationBatch: generationBatchRouter,
+  generationTopic: generationTopicRouter,
   healthcheck: publicProcedure.query(() => "i'm live!"),
+  image: imageRouter,
   importer: importerRouter,
   knowledgeBase: knowledgeBaseRouter,
   market: marketRouter,

package/src/server/routers/tools/mcp.ts

@@ -6,8 +6,19 @@ import { passwordProcedure } from '@/libs/trpc/edge';
 import { authedProcedure, router } from '@/libs/trpc/lambda';
 import { mcpService } from '@/server/services/mcp';
 
+const StreamableHTTPAuthSchema = z
+  .object({
+    // Bearer Token
+    accessToken: z.string().optional(),
+    token: z.string().optional(),
+    type: z.enum(['none', 'bearer', 'oauth2']), // OAuth2 Access Token
+  })
+  .optional();
+
 // Define Zod schemas for MCP Client parameters
 const httpParamsSchema = z.object({
+  auth: StreamableHTTPAuthSchema,
+  headers: z.record(z.string()).optional(),
   name: z.string().min(1),
   type: z.literal('http'),
   url: z.string().url(),
@@ -38,6 +49,8 @@ export const mcpRouter = router({
   getStreamableMcpServerManifest: mcpProcedure
     .input(
       z.object({
+        auth: StreamableHTTPAuthSchema,
+        headers: z.record(z.string()).optional(),
         identifier: z.string(),
         metadata: z
           .object({
@@ -53,6 +66,8 @@ export const mcpRouter = router({
         input.identifier,
         input.url,
         input.metadata,
+        input.auth,
+        input.headers,
       );
     }),
   /* eslint-disable sort-keys-fix/sort-keys-fix */

package/src/server/services/file/__tests__/index.test.ts

@@ -22,6 +22,8 @@ vi.mock('../impls', () => ({
     createPreSignedUrlForPreview: vi.fn(),
     uploadContent: vi.fn(),
     getFullFileUrl: vi.fn(),
+    getKeyFromFullUrl: vi.fn(),
+    uploadMedia: vi.fn(),
   }),
 }));
 
@@ -39,6 +41,7 @@ describe('FileService', () => {
   const mockUserId = 'test-user';
   let mockFileModel: any;
   let mockTempManager: any;
+  let consoleErrorSpy: any;
 
   beforeEach(() => {
     mockFileModel = {
@@ -51,11 +54,16 @@ describe('FileService', () => {
     };
     vi.mocked(FileModel).mockImplementation(() => mockFileModel);
     vi.mocked(TempFileManager).mockImplementation(() => mockTempManager);
+
+    // Mock console.error to test error logging
+    consoleErrorSpy = vi.spyOn(console, 'error').mockImplementation(() => {});
+
     service = new FileService(mockDb, mockUserId);
   });
 
   afterEach(() => {
     vi.clearAllMocks();
+    consoleErrorSpy?.mockRestore();
   });
 
   describe('downloadFileToLocal', () => {
@@ -93,6 +101,30 @@ describe('FileService', () => {
       expect(mockFileModel.delete).toHaveBeenCalledWith('test-file-id', false);
     });
 
+    it('should log error and rethrow for non-NoSuchKey errors', async () => {
+      const originalError = new Error('Network error');
+      mockFileModel.findById.mockResolvedValue(mockFile);
+      vi.mocked(service['impl'].getFileByteArray).mockRejectedValue(originalError);
+
+      await expect(service.downloadFileToLocal('test-file-id')).rejects.toThrow(
+        new TRPCError({ code: 'BAD_REQUEST', message: 'File content is empty' }),
+      );
+
+      // 验证错误被记录到控制台
+      expect(consoleErrorSpy).toHaveBeenCalledWith(originalError);
+      // 验证没有调用删除操作（因为不是NoSuchKey错误）
+      expect(mockFileModel.delete).not.toHaveBeenCalled();
+    });
+
+    it('should handle getFileByteArray returning null content', async () => {
+      mockFileModel.findById.mockResolvedValue(mockFile);
+      vi.mocked(service['impl'].getFileByteArray).mockResolvedValue(null as any);
+
+      await expect(service.downloadFileToLocal('test-file-id')).rejects.toThrow(
+        new TRPCError({ code: 'BAD_REQUEST', message: 'File content is empty' }),
+      );
+    });
+
     it('should successfully download file to local', async () => {
       const mockContent = new Uint8Array([1, 2, 3]);
       const mockFilePath = '/tmp/test.txt';
@@ -112,4 +144,107 @@ describe('FileService', () => {
       expect(mockTempManager.writeTempFile).toHaveBeenCalledWith(mockContent, mockFile.name);
     });
   });
+
+  it('should delegate deleteFile to implementation', async () => {
+    const testKey = 'test-key';
+    await service.deleteFile(testKey);
+
+    expect(service['impl'].deleteFile).toHaveBeenCalledWith(testKey);
+  });
+
+  it('should delegate deleteFiles to implementation', async () => {
+    const testKeys = ['key1', 'key2'];
+    await service.deleteFiles(testKeys);
+
+    expect(service['impl'].deleteFiles).toHaveBeenCalledWith(testKeys);
+  });
+
+  it('should delegate getFileContent to implementation', async () => {
+    const testKey = 'test-key';
+    const expectedContent = 'file content';
+    vi.mocked(service['impl'].getFileContent).mockResolvedValue(expectedContent);
+
+    const result = await service.getFileContent(testKey);
+
+    expect(service['impl'].getFileContent).toHaveBeenCalledWith(testKey);
+    expect(result).toBe(expectedContent);
+  });
+
+  it('should delegate getFileByteArray to implementation', async () => {
+    const testKey = 'test-key';
+    const expectedBytes = new Uint8Array([1, 2, 3]);
+    vi.mocked(service['impl'].getFileByteArray).mockResolvedValue(expectedBytes);
+
+    const result = await service.getFileByteArray(testKey);
+
+    expect(service['impl'].getFileByteArray).toHaveBeenCalledWith(testKey);
+    expect(result).toBe(expectedBytes);
+  });
+
+  it('should delegate createPreSignedUrl to implementation', async () => {
+    const testKey = 'test-key';
+    const expectedUrl = 'https://example.com/signed-url';
+    vi.mocked(service['impl'].createPreSignedUrl).mockResolvedValue(expectedUrl);
+
+    const result = await service.createPreSignedUrl(testKey);
+
+    expect(service['impl'].createPreSignedUrl).toHaveBeenCalledWith(testKey);
+    expect(result).toBe(expectedUrl);
+  });
+
+  it('should delegate createPreSignedUrlForPreview to implementation', async () => {
+    const testKey = 'test-key';
+    const expiresIn = 3600;
+    const expectedUrl = 'https://example.com/preview-url';
+    vi.mocked(service['impl'].createPreSignedUrlForPreview).mockResolvedValue(expectedUrl);
+
+    const result = await service.createPreSignedUrlForPreview(testKey, expiresIn);
+
+    expect(service['impl'].createPreSignedUrlForPreview).toHaveBeenCalledWith(testKey, expiresIn);
+    expect(result).toBe(expectedUrl);
+  });
+
+  it('should delegate uploadContent to implementation', async () => {
+    const testPath = 'test-path';
+    const testContent = 'test content';
+
+    await service.uploadContent(testPath, testContent);
+
+    expect(service['impl'].uploadContent).toHaveBeenCalledWith(testPath, testContent);
+  });
+
+  it('should delegate getFullFileUrl to implementation', async () => {
+    const testUrl = 'test-url';
+    const expiresIn = 3600;
+    const expectedUrl = 'https://example.com/full-url';
+    vi.mocked(service['impl'].getFullFileUrl).mockResolvedValue(expectedUrl);
+
+    const result = await service.getFullFileUrl(testUrl, expiresIn);
+
+    expect(service['impl'].getFullFileUrl).toHaveBeenCalledWith(testUrl, expiresIn);
+    expect(result).toBe(expectedUrl);
+  });
+
+  it('should delegate getKeyFromFullUrl to implementation', () => {
+    const testUrl = 'https://example.com/path/to/file.jpg';
+    const expectedKey = 'path/to/file.jpg';
+    vi.mocked(service['impl'].getKeyFromFullUrl).mockReturnValue(expectedKey);
+
+    const result = service.getKeyFromFullUrl(testUrl);
+
+    expect(service['impl'].getKeyFromFullUrl).toHaveBeenCalledWith(testUrl);
+    expect(result).toBe(expectedKey);
+  });
+
+  it('should delegate uploadMedia to implementation', async () => {
+    const testKey = 'test-key';
+    const testBuffer = Buffer.from('test content');
+    const expectedResult = { key: testKey };
+    vi.mocked(service['impl'].uploadMedia).mockResolvedValue(expectedResult);
+
+    const result = await service.uploadMedia(testKey, testBuffer);
+
+    expect(service['impl'].uploadMedia).toHaveBeenCalledWith(testKey, testBuffer);
+    expect(result).toBe(expectedResult);
+  });
 });