@lobehub/chat 1.98.2 → 1.99.0

package/src/server/services/generation/index.ts

@@ -0,0 +1,239 @@
+import debug from 'debug';
+import { sha256 } from 'js-sha256';
+import mime from 'mime';
+import { nanoid } from 'nanoid';
+import sharp from 'sharp';
+
+import { LobeChatDatabase } from '@/database/type';
+import { FileService } from '@/server/services/file';
+import { getYYYYmmddHHMMss } from '@/utils/time';
+
+const log = debug('lobe-image:generation-service');
+
+interface ImageForGeneration {
+  buffer: Buffer;
+  extension: string;
+  hash: string;
+  height: number;
+  mime: string;
+  size: number;
+  width: number;
+}
+
+/**
+ * 图片生成服务
+ * 负责处理AI生成图片的转换、上传和封面创建
+ */
+export class GenerationService {
+  private fileService: FileService;
+
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.fileService = new FileService(db, userId);
+  }
+
+  /**
+   * Generate width 512px image as thumbnail when width > 512, end with _512.webp
+   */
+  async transformImageForGeneration(url: string): Promise<{
+    image: ImageForGeneration;
+    thumbnailImage: ImageForGeneration;
+  }> {
+    log('Starting image transformation for:', url.startsWith('data:') ? 'base64 data' : url);
+
+    // If the url is in base64 format, extract the Buffer directly; otherwise, use fetch to get the Buffer
+    let originalImageBuffer: Buffer;
+    let originalMimeType: string;
+
+    if (url.startsWith('data:')) {
+      log('Processing base64 image data');
+      // Extract the MIME type and base64 data part
+      const [mimeTypePart, base64Data] = url.split(',');
+      originalMimeType = mimeTypePart.split(':')[1].split(';')[0];
+      originalImageBuffer = Buffer.from(base64Data, 'base64');
+    } else {
+      log('Fetching image from URL:', url);
+      const response = await fetch(url);
+      if (!response.ok) {
+        throw new Error(
+          `Failed to fetch image from ${url}: ${response.status} ${response.statusText}`,
+        );
+      }
+      const arrayBuffer = await response.arrayBuffer();
+      originalImageBuffer = Buffer.from(arrayBuffer);
+      originalMimeType = response.headers.get('content-type') || 'application/octet-stream';
+      log('Successfully fetched image, buffer size:', originalImageBuffer.length);
+    }
+
+    // Calculate hash for original image
+    const originalHash = sha256(originalImageBuffer);
+    log('Original image hash calculated:', originalHash);
+
+    const sharpInstance = sharp(originalImageBuffer);
+    const { format, width, height } = await sharpInstance.metadata();
+    log('Image metadata:', { format, height, width });
+
+    if (!width || !height) {
+      throw new Error(`Invalid image format: ${format}, url: ${url}`);
+    }
+
+    const shouldResize = format !== 'webp' || width > 512 || height > 512;
+    const thumbnailWidth = shouldResize
+      ? width > height
+        ? 512
+        : Math.round((width * 512) / height)
+      : width;
+    const thumbnailHeight = shouldResize
+      ? height > width
+        ? 512
+        : Math.round((height * 512) / width)
+      : height;
+
+    log('Thumbnail dimensions calculated:', { shouldResize, thumbnailHeight, thumbnailWidth });
+
+    const thumbnailBuffer = shouldResize
+      ? await sharpInstance.resize(thumbnailWidth, thumbnailHeight).webp().toBuffer()
+      : originalImageBuffer;
+
+    // Calculate hash for thumbnail
+    const thumbnailHash = sha256(thumbnailBuffer);
+    log('Thumbnail image hash calculated:', thumbnailHash);
+
+    log('Image transformation completed successfully');
+
+    // Determine extension without fallback
+    let extension: string;
+    if (url.startsWith('data:')) {
+      const mimeExtension = mime.getExtension(originalMimeType);
+      if (!mimeExtension) {
+        throw new Error(`Unable to determine file extension for MIME type: ${originalMimeType}`);
+      }
+      extension = mimeExtension;
+    } else {
+      const urlExtension = url.split('.').pop();
+      if (!urlExtension) {
+        throw new Error(`Unable to determine file extension from URL: ${url}`);
+      }
+      extension = urlExtension;
+    }
+
+    return {
+      image: {
+        buffer: originalImageBuffer,
+        extension,
+        hash: originalHash,
+        height,
+        mime: originalMimeType,
+        size: originalImageBuffer.length,
+        width,
+      },
+      thumbnailImage: {
+        buffer: thumbnailBuffer,
+        extension: 'webp',
+        hash: thumbnailHash,
+        height: thumbnailHeight,
+        mime: 'image/webp',
+        size: thumbnailBuffer.length,
+        width: thumbnailWidth,
+      },
+    };
+  }
+
+  async uploadImageForGeneration(image: ImageForGeneration, thumbnail: ImageForGeneration) {
+    log('Starting image upload for generation');
+
+    const generationImagesFolder = 'generations/images';
+    const uuid = nanoid();
+    const dateTime = getYYYYmmddHHMMss(new Date());
+    const pathPrefix = `${generationImagesFolder}/${uuid}_${image.width}x${image.height}_${dateTime}`;
+    const imageKey = `${pathPrefix}_raw.${image.extension}`;
+    const thumbnailKey = `${pathPrefix}_thumb.${thumbnail.extension}`;
+
+    log('Generated paths:', { imagePath: imageKey, thumbnailPath: thumbnailKey });
+
+    // Check if image and thumbnail buffers are identical
+    const isIdenticalBuffer = image.buffer.equals(thumbnail.buffer);
+    log('Buffer comparison:', {
+      imageSize: image.buffer.length,
+      isIdenticalBuffer,
+      thumbnailSize: thumbnail.buffer.length,
+    });
+
+    if (isIdenticalBuffer) {
+      log('Buffers are identical, uploading single image');
+      // If buffers are identical, only upload once
+      const result = await this.fileService.uploadMedia(imageKey, image.buffer);
+      log('Single image uploaded successfully:', result.key);
+      // Use the same key for both image and thumbnail
+      return {
+        imageUrl: result.key,
+        thumbnailImageUrl: result.key,
+      };
+    } else {
+      log('Buffers are different, uploading both images');
+      // If buffers are different, upload both
+      const [imageResult, thumbnailResult] = await Promise.all([
+        this.fileService.uploadMedia(imageKey, image.buffer),
+        this.fileService.uploadMedia(thumbnailKey, thumbnail.buffer),
+      ]);
+
+      log('Both images uploaded successfully:', {
+        imageUrl: imageResult.key,
+        thumbnailImageUrl: thumbnailResult.key,
+      });
+
+      return {
+        imageUrl: imageResult.key,
+        thumbnailImageUrl: thumbnailResult.key,
+      };
+    }
+  }
+
+  /**
+   * Create a 256x256 cover image from a given URL and upload
+   * @param coverUrl - The source image URL (can be base64 or HTTP URL)
+   * @returns The key of the uploaded cover image
+   */
+  async createCoverFromUrl(coverUrl: string): Promise<string> {
+    log('Creating cover image from URL:', coverUrl.startsWith('data:') ? 'base64 data' : coverUrl);
+
+    // Download the original image
+    let originalImageBuffer: Buffer;
+    if (coverUrl.startsWith('data:')) {
+      log('Processing base64 cover image data');
+      // Extract base64 data part
+      const [, base64Data] = coverUrl.split(',');
+      originalImageBuffer = Buffer.from(base64Data, 'base64');
+    } else {
+      log('Fetching cover image from URL:', coverUrl);
+      const response = await fetch(coverUrl);
+      if (!response.ok) {
+        throw new Error(
+          `Failed to fetch cover image from ${coverUrl}: ${response.status} ${response.statusText}`,
+        );
+      }
+      const arrayBuffer = await response.arrayBuffer();
+      originalImageBuffer = Buffer.from(arrayBuffer);
+      log('Successfully fetched cover image, buffer size:', originalImageBuffer.length);
+    }
+
+    // Process image to 256x256 cover with webp format
+    log('Processing cover image to 256x256 webp format');
+    const coverBuffer = await sharp(originalImageBuffer)
+      .resize(256, 256, { fit: 'cover', position: 'center' })
+      .toBuffer();
+
+    log('Cover image processed, final size:', coverBuffer.length);
+
+    // Upload using FileService
+    const coverFolder = 'generations/covers';
+    const uuid = nanoid();
+    const dateTime = getYYYYmmddHHMMss(new Date());
+    const coverKey = `${coverFolder}/${uuid}_cover_${dateTime}.webp`;
+
+    log('Uploading cover image:', coverKey);
+    const result = await this.fileService.uploadMedia(coverKey, coverBuffer);
+
+    log('Cover image uploaded successfully:', result.key);
+    return result.key;
+  }
+}

package/src/server/services/mcp/index.ts

@@ -216,7 +216,7 @@ class MCPService {
       log(`New client initialized and cached for key: ${key.slice(0, 20)}`);
       return client;
     } catch (error) {
-      console.error(`Failed to initialize MCP client for key ${key}:`, error);
+      console.error(`Failed to initialize MCP client:`, error);
 
       // 保留完整的错误信息，特别是详细的 stderr 输出
       const errorMessage = error instanceof Error ? error.message : String(error);
@@ -266,8 +266,26 @@ class MCPService {
     identifier: string,
     url: string,
     metadata?: CustomPluginMetadata,
+    auth?: {
+      accessToken?: string;
+      token?: string;
+      type: 'none' | 'bearer' | 'oauth2';
+    },
+    headers?: Record<string, string>,
   ): Promise<LobeChatPluginManifest> {
-    const tools = await this.listTools({ name: identifier, type: 'http', url }); // Get client using params
+    const mcpParams = { name: identifier, type: 'http' as const, url };
+
+    // 如果有认证信息，添加到参数中
+    if (auth) {
+      (mcpParams as any).auth = auth;
+    }
+
+    // 如果有 headers 信息，添加到参数中
+    if (headers) {
+      (mcpParams as any).headers = headers;
+    }
+
+    const tools = await this.listTools(mcpParams);
 
     return {
       api: tools,

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

@@ -0,0 +1,40 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { lambdaClient } from '@/libs/trpc/client';
+
+import { generationService } from '../generation';
+
+vi.mock('@/libs/trpc/client', () => ({
+  lambdaClient: {
+    generation: {
+      getGenerationStatus: { query: vi.fn() },
+      deleteGeneration: { mutate: vi.fn() },
+    },
+  },
+}));
+
+describe('GenerationService', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it('getGenerationStatus should call lambdaClient with correct params', async () => {
+    const generationId = 'test-generation-id';
+    const asyncTaskId = 'test-async-task-id';
+
+    await generationService.getGenerationStatus(generationId, asyncTaskId);
+
+    expect(lambdaClient.generation.getGenerationStatus.query).toBeCalledWith({
+      generationId,
+      asyncTaskId,
+    });
+  });
+
+  it('deleteGeneration should call lambdaClient with correct params', async () => {
+    const generationId = 'test-generation-id';
+
+    await generationService.deleteGeneration(generationId);
+
+    expect(lambdaClient.generation.deleteGeneration.mutate).toBeCalledWith({ generationId });
+  });
+});

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

@@ -0,0 +1,36 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { lambdaClient } from '@/libs/trpc/client';
+
+import { generationBatchService } from '../generationBatch';
+
+vi.mock('@/libs/trpc/client', () => ({
+  lambdaClient: {
+    generationBatch: {
+      getGenerationBatches: { query: vi.fn() },
+      deleteGenerationBatch: { mutate: vi.fn() },
+    },
+  },
+}));
+
+describe('GenerationBatchService', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it('getGenerationBatches should call lambdaClient with correct params', async () => {
+    const topicId = 'test-topic-id';
+
+    await generationBatchService.getGenerationBatches(topicId);
+
+    expect(lambdaClient.generationBatch.getGenerationBatches.query).toBeCalledWith({ topicId });
+  });
+
+  it('deleteGenerationBatch should call lambdaClient with correct params', async () => {
+    const batchId = 'test-batch-id';
+
+    await generationBatchService.deleteGenerationBatch(batchId);
+
+    expect(lambdaClient.generationBatch.deleteGenerationBatch.mutate).toBeCalledWith({ batchId });
+  });
+});

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

@@ -0,0 +1,72 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { lambdaClient } from '@/libs/trpc/client';
+import { UpdateTopicValue } from '@/server/routers/lambda/generationTopic';
+
+import { ServerService } from '../generationTopic';
+
+vi.mock('@/libs/trpc/client', () => ({
+  lambdaClient: {
+    generationTopic: {
+      getAllGenerationTopics: { query: vi.fn() },
+      createTopic: { mutate: vi.fn() },
+      updateTopic: { mutate: vi.fn() },
+      updateTopicCover: { mutate: vi.fn() },
+      deleteTopic: { mutate: vi.fn() },
+    },
+  },
+}));
+
+describe('GenerationTopic ServerService', () => {
+  let service: ServerService;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    service = new ServerService();
+  });
+
+  it('getAllGenerationTopics should call lambdaClient', async () => {
+    await service.getAllGenerationTopics();
+    expect(lambdaClient.generationTopic.getAllGenerationTopics.query).toBeCalled();
+  });
+
+  it('createTopic should call lambdaClient with undefined', async () => {
+    await service.createTopic();
+    expect(lambdaClient.generationTopic.createTopic.mutate).toBeCalledWith(undefined);
+  });
+
+  it('updateTopic should call lambdaClient with correct params', async () => {
+    const id = 'test-topic-id';
+    const data: UpdateTopicValue = {
+      title: 'Updated Topic',
+      coverUrl: 'https://example.com/cover.jpg',
+    };
+
+    await service.updateTopic(id, data);
+
+    expect(lambdaClient.generationTopic.updateTopic.mutate).toBeCalledWith({
+      id,
+      value: data,
+    });
+  });
+
+  it('updateTopicCover should call lambdaClient with correct params', async () => {
+    const id = 'test-topic-id';
+    const coverUrl = 'https://example.com/cover.jpg';
+
+    await service.updateTopicCover(id, coverUrl);
+
+    expect(lambdaClient.generationTopic.updateTopicCover.mutate).toBeCalledWith({
+      id,
+      coverUrl,
+    });
+  });
+
+  it('deleteTopic should call lambdaClient with correct params', async () => {
+    const id = 'test-topic-id';
+
+    await service.deleteTopic(id);
+
+    expect(lambdaClient.generationTopic.deleteTopic.mutate).toBeCalledWith({ id });
+  });
+});

package/src/services/electron/file.ts

@@ -10,11 +10,13 @@ class DesktopFileAPI {
    * 上传文件到桌面应用
    * @param file 文件对象
    * @param hash 文件哈希
+   * @param path 文件存储路径
    * @returns 上传结果
    */
   async uploadFile(
     file: File,
     hash: string,
+    path: string,
   ): Promise<{ metadata: FileMetadata; success: boolean }> {
     const arrayBuffer = await file.arrayBuffer();
 
@@ -22,7 +24,7 @@ class DesktopFileAPI {
       content: arrayBuffer,
       filename: file.name,
       hash,
-      path: file.name,
+      path,
       type: file.type,
     });
   }

package/src/services/generation.ts

@@ -0,0 +1,16 @@
+import { lambdaClient } from '@/libs/trpc/client';
+
+class GenerationService {
+  async getGenerationStatus(generationId: string, asyncTaskId: string) {
+    return lambdaClient.generation.getGenerationStatus.query({ asyncTaskId, generationId });
+  }
+
+  /**
+   * Delete a single generation
+   */
+  async deleteGeneration(generationId: string) {
+    return lambdaClient.generation.deleteGeneration.mutate({ generationId });
+  }
+}
+
+export const generationService = new GenerationService();

package/src/services/generationBatch.ts

@@ -0,0 +1,25 @@
+import { GenerationBatchItem } from '@/database/schemas';
+import { lambdaClient } from '@/libs/trpc/client';
+import { Generation, GenerationBatch } from '@/types/generation';
+
+type GenerationBatchWithAsyncTaskId = GenerationBatch & {
+  generations: (Generation & { asyncTaskId?: string | null })[];
+};
+
+class GenerationBatchService {
+  /**
+   * Get generation batches for a specific topic
+   */
+  async getGenerationBatches(topicId: string): Promise<GenerationBatchWithAsyncTaskId[]> {
+    return lambdaClient.generationBatch.getGenerationBatches.query({ topicId });
+  }
+
+  /**
+   * Delete a generation batch
+   */
+  async deleteGenerationBatch(batchId: string): Promise<GenerationBatchItem | undefined> {
+    return lambdaClient.generationBatch.deleteGenerationBatch.mutate({ batchId });
+  }
+}
+
+export const generationBatchService = new GenerationBatchService();

package/src/services/generationTopic.ts

@@ -0,0 +1,28 @@
+import { GenerationTopicItem } from '@/database/schemas';
+import { lambdaClient } from '@/libs/trpc/client';
+import { UpdateTopicValue } from '@/server/routers/lambda/generationTopic';
+import { ImageGenerationTopic } from '@/types/generation';
+
+export class ServerService {
+  async getAllGenerationTopics(): Promise<ImageGenerationTopic[]> {
+    return lambdaClient.generationTopic.getAllGenerationTopics.query();
+  }
+
+  async createTopic(): Promise<string> {
+    return lambdaClient.generationTopic.createTopic.mutate(undefined);
+  }
+
+  async updateTopic(id: string, data: UpdateTopicValue): Promise<GenerationTopicItem | undefined> {
+    return lambdaClient.generationTopic.updateTopic.mutate({ id, value: data });
+  }
+
+  async updateTopicCover(id: string, coverUrl: string): Promise<GenerationTopicItem | undefined> {
+    return lambdaClient.generationTopic.updateTopicCover.mutate({ coverUrl, id });
+  }
+
+  async deleteTopic(id: string): Promise<GenerationTopicItem | undefined> {
+    return lambdaClient.generationTopic.deleteTopic.mutate({ id });
+  }
+}
+
+export const generationTopicService = new ServerService();

package/src/services/image.ts

@@ -0,0 +1,33 @@
+import debug from 'debug';
+
+import { lambdaClient } from '@/libs/trpc/client';
+import { CreateImageServicePayload } from '@/server/routers/lambda/image';
+
+// Create debug logger
+const log = debug('lobe-image:service');
+
+export class AiImageService {
+  async createImage(payload: CreateImageServicePayload) {
+    log('Creating image with payload: %O', payload);
+
+    try {
+      const result = await lambdaClient.image.createImage.mutate(payload);
+      log('Image creation service call completed successfully: %O', {
+        batchId: result.data?.batch?.id,
+        generationCount: result.data?.generations?.length,
+        success: result.success,
+      });
+
+      return result;
+    } catch (error) {
+      log('Image creation service call failed: %O', {
+        error: (error as Error).message,
+        payload,
+      });
+
+      throw error;
+    }
+  }
+}
+
+export const imageService = new AiImageService();

package/src/services/mcp.ts

@@ -128,15 +128,20 @@ class MCPService {
   }
 
   async getStreamableMcpServerManifest(
-    identifier: string,
-    url: string,
-    metadata?: CustomPluginMetadata,
+    params: {
+      auth?: {
+        accessToken?: string;
+        token?: string;
+        type: 'none' | 'bearer' | 'oauth2';
+      };
+      headers?: Record<string, string>;
+      identifier: string;
+      metadata?: CustomPluginMetadata;
+      url: string;
+    },
     signal?: AbortSignal,
   ) {
-    return toolsClient.mcp.getStreamableMcpServerManifest.query(
-      { identifier, metadata, url },
-      { signal },
-    );
+    return toolsClient.mcp.getStreamableMcpServerManifest.query(params, { signal });
   }
 
   async getStdioMcpServerManifest(

package/src/services/upload.ts

@@ -13,6 +13,38 @@ import { uuid } from '@/utils/uuid';
 
 export const UPLOAD_NETWORK_ERROR = 'NetWorkError';
 
+/**
+ * Generate file storage path metadata for S3-compatible storage
+ * @param originalFilename - Original filename
+ * @param options - Path generation options
+ * @returns Path metadata including date, dirname, filename, and pathname
+ */
+const generateFilePathMetadata = (
+  originalFilename: string,
+  options: { directory?: string; pathname?: string } = {},
+): {
+  date: string;
+  dirname: string;
+  filename: string;
+  pathname: string;
+} => {
+  // Generate unique filename with UUID prefix and original extension
+  const extension = originalFilename.split('.').at(-1);
+  const filename = `${uuid()}.${extension}`;
+
+  // Generate timestamp-based directory path
+  const date = (Date.now() / 1000 / 60 / 60).toFixed(0);
+  const dirname = `${options.directory || fileEnv.NEXT_PUBLIC_S3_FILE_PATH}/${date}`;
+  const pathname = options.pathname ?? `${dirname}/${filename}`;
+
+  return {
+    date,
+    dirname,
+    filename,
+    pathname,
+  };
+};
+
 interface UploadFileToS3Options {
   directory?: string;
   filename?: string;
@@ -38,7 +70,7 @@ class UploadService {
 
     // 桌面端上传逻辑（并且没开启 sync 同步）
     if (isDesktop && !isSyncActive) {
-      const data = await this.uploadToDesktopS3(file);
+      const data = await this.uploadToDesktopS3(file, { directory, pathname });
       return { data, success: true };
     }
 
@@ -182,12 +214,18 @@ class UploadService {
     return result;
   };
 
-  private uploadToDesktopS3 = async (file: File) => {
+  private uploadToDesktopS3 = async (
+    file: File,
+    options: { directory?: string; pathname?: string } = {},
+  ) => {
     const fileArrayBuffer = await file.arrayBuffer();
     const hash = sha256(fileArrayBuffer);
 
+    // 生成文件路径元数据
+    const { pathname } = generateFilePathMetadata(file.name, options);
+
     const { desktopFileAPI } = await import('@/services/electron/file');
-    const { metadata } = await desktopFileAPI.uploadFile(file, hash);
+    const { metadata } = await desktopFileAPI.uploadFile(file, hash, pathname);
     return metadata;
   };
 
@@ -223,12 +261,8 @@ class UploadService {
       preSignUrl: string;
     }
   > => {
-    const filename = `${uuid()}.${file.name.split('.').at(-1)}`;
-
-    // 精确到以 h 为单位的 path
-    const date = (Date.now() / 1000 / 60 / 60).toFixed(0);
-    const dirname = `${options.directory || fileEnv.NEXT_PUBLIC_S3_FILE_PATH}/${date}`;
-    const pathname = options.pathname ?? `${dirname}/${filename}`;
+    // 生成文件路径元数据
+    const { date, dirname, filename, pathname } = generateFilePathMetadata(file.name, options);
 
     const preSignUrl = await edgeClient.upload.createS3PreSignedUrl.mutate({ pathname });