@lobehub/lobehub 2.0.0-next.35 → 2.0.0-next.36

package/src/services/message/__tests__/metadata-race-condition.test.ts

@@ -0,0 +1,157 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { lambdaClient } from '@/libs/trpc/client';
+
+import { MessageService } from '../index';
+
+vi.mock('@/libs/trpc/client', () => ({
+  lambdaClient: {
+    message: {
+      updateMetadata: {
+        mutate: vi.fn(),
+      },
+    },
+  },
+}));
+
+describe('MessageService - Race Condition Control', () => {
+  let messageService: MessageService;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    messageService = new MessageService();
+  });
+
+  describe('updateMessageMetadata race condition', () => {
+    it('should cancel previous request when new update is triggered for same message', async () => {
+      const messageId = 'test-message-id';
+      let firstRequestAborted = false;
+      let secondRequestCompleted = false;
+
+      // Mock first request (slow)
+      vi.mocked(lambdaClient.message.updateMetadata.mutate).mockImplementationOnce(
+        (_params, options) =>
+          new Promise((resolve, reject) => {
+            const signal = options?.signal;
+            if (signal) {
+              signal.addEventListener('abort', () => {
+                firstRequestAborted = true;
+                reject(new Error('Aborted'));
+              });
+            }
+            // Simulate slow request
+            setTimeout(() => resolve({ success: true, messages: [] }), 200);
+          }),
+      );
+
+      // Mock second request (fast)
+      vi.mocked(lambdaClient.message.updateMetadata.mutate).mockImplementationOnce(
+        async (_params, _options) => {
+          secondRequestCompleted = true;
+          return { success: true, messages: [] };
+        },
+      );
+
+      // Start first update
+      const firstPromise = messageService.updateMessageMetadata(messageId, { compare: true });
+
+      // Wait a bit then start second update
+      await new Promise((resolve) => setTimeout(resolve, 10));
+      const secondPromise = messageService.updateMessageMetadata(messageId, { compare: false });
+
+      // First should be aborted
+      await expect(firstPromise).rejects.toThrow('Aborted');
+      expect(firstRequestAborted).toBe(true);
+
+      // Second should complete successfully
+      await expect(secondPromise).resolves.toEqual({ success: true, messages: [] });
+      expect(secondRequestCompleted).toBe(true);
+    });
+
+    it('should allow concurrent updates for different messages', async () => {
+      const message1Id = 'message-1';
+      const message2Id = 'message-2';
+
+      vi.mocked(lambdaClient.message.updateMetadata.mutate).mockResolvedValue({
+        success: true,
+        messages: [],
+      });
+
+      const [result1, result2] = await Promise.all([
+        messageService.updateMessageMetadata(message1Id, { cost: 0.001 }),
+        messageService.updateMessageMetadata(message2Id, { cost: 0.002 }),
+      ]);
+
+      expect(result1).toEqual({ success: true, messages: [] });
+      expect(result2).toEqual({ success: true, messages: [] });
+      expect(lambdaClient.message.updateMetadata.mutate).toHaveBeenCalledTimes(2);
+    });
+
+    it('should handle rapid successive updates correctly', async () => {
+      const messageId = 'test-message-id';
+      let completedUpdates = 0;
+      const abortedUpdates: number[] = [];
+
+      // All but the last request should be aborted
+      let callIndex = 0;
+      vi.mocked(lambdaClient.message.updateMetadata.mutate).mockImplementation(
+        (_params, options) => {
+          const currentIndex = callIndex++;
+          return new Promise((resolve, reject) => {
+            const signal = options?.signal;
+            let isAborted = false;
+
+            if (signal) {
+              signal.addEventListener('abort', () => {
+                isAborted = true;
+                abortedUpdates.push(currentIndex);
+                reject(new Error('Aborted'));
+              });
+            }
+
+            setTimeout(() => {
+              if (!isAborted) {
+                completedUpdates++;
+                resolve({ success: true, messages: [] });
+              }
+            }, 50);
+          });
+        },
+      );
+
+      // Trigger 5 rapid updates sequentially with catch to prevent unhandled rejections
+      const promise1 = messageService
+        .updateMessageMetadata(messageId, { cost: 0.001 })
+        .catch((e) => e);
+      await new Promise((resolve) => setTimeout(resolve, 5));
+      const promise2 = messageService
+        .updateMessageMetadata(messageId, { cost: 0.002 })
+        .catch((e) => e);
+      await new Promise((resolve) => setTimeout(resolve, 5));
+      const promise3 = messageService.updateMessageMetadata(messageId, { tps: 10 }).catch((e) => e);
+      await new Promise((resolve) => setTimeout(resolve, 5));
+      const promise4 = messageService.updateMessageMetadata(messageId, { tps: 20 }).catch((e) => e);
+      await new Promise((resolve) => setTimeout(resolve, 5));
+      const promise5 = messageService
+        .updateMessageMetadata(messageId, { compare: true })
+        .catch((e) => e);
+
+      // Wait for all to settle
+      const results = await Promise.all([promise1, promise2, promise3, promise4, promise5]);
+
+      // First 4 should be errors (aborted), last should succeed
+      expect(results[0]).toBeInstanceOf(Error);
+      expect(results[1]).toBeInstanceOf(Error);
+      expect(results[2]).toBeInstanceOf(Error);
+      expect(results[3]).toBeInstanceOf(Error);
+      expect(results[4]).toEqual({ success: true, messages: [] });
+
+      // 4 requests should have been aborted
+      expect(abortedUpdates.length).toBe(4);
+      expect(abortedUpdates).toEqual([0, 1, 2, 3]);
+
+      // Only the last request should complete
+      expect(completedUpdates).toBe(1);
+    });
+  });
+});

package/src/services/message/index.ts

@@ -5,6 +5,7 @@ import {
   ChatTranslate,
   CreateMessageParams,
   CreateMessageResult,
+  MessageMetadata,
   ModelRankItem,
   UIChatMessage,
   UpdateMessageParams,
@@ -15,14 +16,10 @@ import type { HeatmapsProps } from '@lobehub/charts';
 
 import { INBOX_SESSION_ID } from '@/const/session';
 import { lambdaClient } from '@/libs/trpc/client';
-import { useUserStore } from '@/store/user';
-import { labPreferSelectors } from '@/store/user/selectors';
 
-export class MessageService {
-  private get useGroup() {
-    return labPreferSelectors.enableAssistantMessageGroup(useUserStore.getState());
-  }
+import { abortableRequest } from '../utils/abortableRequest';
 
+export class MessageService {
   createMessage = async ({
     sessionId,
     ...params
@@ -30,7 +27,6 @@ export class MessageService {
     return lambdaClient.message.createMessage.mutate({
       ...params,
       sessionId: sessionId ? this.toDbSessionId(sessionId) : undefined,
-      useGroup: this.useGroup,
     });
   };
 
@@ -43,7 +39,6 @@ export class MessageService {
       groupId,
       sessionId: this.toDbSessionId(sessionId),
       topicId,
-      useGroup: this.useGroup,
     });
 
     return data as unknown as UIChatMessage[];
@@ -53,7 +48,6 @@ export class MessageService {
     const data = await lambdaClient.message.getMessages.query({
       groupId,
       topicId,
-      useGroup: this.useGroup,
     });
     return data as unknown as UIChatMessage[];
   };
@@ -100,7 +94,6 @@ export class MessageService {
       id,
       sessionId: options?.sessionId,
       topicId: options?.topicId,
-      useGroup: this.useGroup,
       value,
     });
   };
@@ -113,6 +106,24 @@ export class MessageService {
     return lambdaClient.message.updateTTS.mutate({ id, value: tts });
   };
 
+  updateMessageMetadata = async (
+    id: string,
+    value: Partial<MessageMetadata>,
+    options?: { sessionId?: string | null; topicId?: string | null },
+  ): Promise<UpdateMessageResult> => {
+    return abortableRequest.execute(`message-metadata-${id}`, (signal) =>
+      lambdaClient.message.updateMetadata.mutate(
+        {
+          id,
+          sessionId: options?.sessionId,
+          topicId: options?.topicId,
+          value,
+        },
+        { signal },
+      ),
+    );
+  };
+
   updateMessagePluginState = async (
     id: string,
     value: Record<string, any>,
@@ -122,7 +133,6 @@ export class MessageService {
       id,
       sessionId: options?.sessionId,
       topicId: options?.topicId,
-      useGroup: this.useGroup,
       value,
     });
   };
@@ -136,7 +146,6 @@ export class MessageService {
       id,
       sessionId: options?.sessionId,
       topicId: options?.topicId,
-      useGroup: this.useGroup,
       value: error as any,
     });
   };
@@ -150,7 +159,6 @@ export class MessageService {
       id,
       sessionId: options?.sessionId,
       topicId: options?.topicId,
-      useGroup: this.useGroup,
       value: data,
     });
   };
@@ -163,7 +171,6 @@ export class MessageService {
       id,
       sessionId: options?.sessionId,
       topicId: options?.topicId,
-      useGroup: this.useGroup,
     });
   };
 
@@ -175,7 +182,6 @@ export class MessageService {
       ids,
       sessionId: options?.sessionId,
       topicId: options?.topicId,
-      useGroup: this.useGroup,
     });
   };
 

package/src/services/upload.ts

@@ -68,13 +68,13 @@ class UploadService {
     const state = getElectronStoreState();
     const isSyncActive = electronSyncSelectors.isSyncActive(state);
 
-    // 桌面端上传逻辑（并且没开启 sync 同步）
+    // Desktop upload logic (when sync is not enabled)
     if (isDesktop && !isSyncActive) {
       const data = await this.uploadToDesktopS3(file, { directory, pathname });
       return { data, success: true };
     }
 
-    // 服务端上传逻辑
+    // Server-side upload logic
     if (isServerMode) {
       // if is server mode, upload to server s3,
 
@@ -83,7 +83,7 @@ class UploadService {
     }
 
     // upload to client s3
-    // 客户端上传逻辑
+    // Client-side upload logic
     if (!skipCheckFileType && !file.type.startsWith('image') && !file.type.startsWith('video')) {
       onNotSupported?.();
       return { data: undefined as unknown as FileMetadata, success: false };
@@ -103,18 +103,18 @@ class UploadService {
     base64Data: string,
     options: UploadFileToS3Options = {},
   ): Promise<UploadBase64ToS3Result> => {
-    // 解析 base64 数据
+    // Parse base64 data
     const { base64, mimeType, type } = parseDataUri(base64Data);
 
     if (!base64 || !mimeType || type !== 'base64') {
       throw new Error('Invalid base64 data for image');
     }
 
-    // 将 base64 转换为 Blob
+    // Convert base64 to Blob
     const byteCharacters = atob(base64);
     const byteArrays = [];
 
-    // 分块处理以避免内存问题
+    // Process in chunks to avoid memory issues
     for (let offset = 0; offset < byteCharacters.length; offset += 1024) {
       const slice = byteCharacters.slice(offset, offset + 1024);
 
@@ -129,14 +129,14 @@ class UploadService {
 
     const blob = new Blob(byteArrays, { type: mimeType });
 
-    // 确定文件扩展名
+    // Determine file extension
     const fileExtension = mimeType.split('/')[1] || 'png';
     const fileName = `${options.filename || `image_${dayjs().format('YYYY-MM-DD-hh-mm-ss')}`}.${fileExtension}`;
 
-    // 创建文件对象
+    // Create file object
     const file = new File([blob], fileName, { type: mimeType });
 
-    // 使用统一的上传方法
+    // Use unified upload method
     const { data: metadata } = await this.uploadFileToS3(file, options);
     const hash = sha256(await file.arrayBuffer());
 
@@ -221,7 +221,7 @@ class UploadService {
     const fileArrayBuffer = await file.arrayBuffer();
     const hash = sha256(fileArrayBuffer);
 
-    // 生成文件路径元数据
+    // Generate file path metadata
     const { pathname } = generateFilePathMetadata(file.name, options);
 
     const { desktopFileAPI } = await import('@/services/electron/file');
@@ -261,7 +261,7 @@ class UploadService {
       preSignUrl: string;
     }
   > => {
-    // 生成文件路径元数据
+    // Generate file path metadata
     const { date, dirname, filename, pathname } = generateFilePathMetadata(file.name, options);
 
     const preSignUrl = await lambdaClient.upload.createS3PreSignedUrl.mutate({ pathname });

package/src/services/utils/abortableRequest.test.ts

@@ -0,0 +1,161 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { abortableRequest } from './abortableRequest';
+
+describe('AbortableRequestManager', () => {
+  beforeEach(() => {
+    abortableRequest.cancelAll();
+  });
+
+  describe('execute', () => {
+    it('should execute request successfully', async () => {
+      const mockFetcher = vi.fn(async (signal: AbortSignal) => {
+        return 'result';
+      });
+
+      const result = await abortableRequest.execute('test-key', mockFetcher);
+
+      expect(result).toBe('result');
+      expect(mockFetcher).toHaveBeenCalledTimes(1);
+      expect(mockFetcher).toHaveBeenCalledWith(expect.any(AbortSignal));
+    });
+
+    it('should cancel previous request when new request with same key is triggered', async () => {
+      let firstRequestAborted = false;
+      let secondRequestAborted = false;
+
+      const firstFetcher = vi.fn(
+        async (signal: AbortSignal) =>
+          new Promise((resolve, reject) => {
+            signal.addEventListener('abort', () => {
+              firstRequestAborted = true;
+              reject(new Error('Aborted'));
+            });
+            setTimeout(() => resolve('first'), 100);
+          }),
+      );
+
+      const secondFetcher = vi.fn(
+        async (signal: AbortSignal) =>
+          new Promise((resolve, reject) => {
+            signal.addEventListener('abort', () => {
+              secondRequestAborted = true;
+              reject(new Error('Aborted'));
+            });
+            setTimeout(() => resolve('second'), 100);
+          }),
+      );
+
+      // Start first request
+      const firstPromise = abortableRequest.execute('same-key', firstFetcher);
+
+      // Start second request with same key (should cancel first)
+      await new Promise((resolve) => setTimeout(resolve, 10));
+      const secondPromise = abortableRequest.execute('same-key', secondFetcher);
+
+      // First should be aborted
+      await expect(firstPromise).rejects.toThrow('Aborted');
+      expect(firstRequestAborted).toBe(true);
+
+      // Second should succeed
+      const result = await secondPromise;
+      expect(result).toBe('second');
+      expect(secondRequestAborted).toBe(false);
+    });
+
+    it('should allow concurrent requests with different keys', async () => {
+      const fetcher1 = vi.fn(async () => 'result1');
+      const fetcher2 = vi.fn(async () => 'result2');
+
+      const [result1, result2] = await Promise.all([
+        abortableRequest.execute('key1', fetcher1),
+        abortableRequest.execute('key2', fetcher2),
+      ]);
+
+      expect(result1).toBe('result1');
+      expect(result2).toBe('result2');
+      expect(fetcher1).toHaveBeenCalledTimes(1);
+      expect(fetcher2).toHaveBeenCalledTimes(1);
+    });
+
+    it('should clean up controller after request completes', async () => {
+      const fetcher = vi.fn(async () => 'result');
+
+      await abortableRequest.execute('cleanup-test', fetcher);
+
+      // Manually check that controller is cleaned up by starting a new request
+      // and verifying it doesn't abort anything (since map should be empty)
+      let aborted = false;
+      const fetcher2 = vi.fn(
+        async (signal: AbortSignal) =>
+          new Promise((resolve) => {
+            signal.addEventListener('abort', () => {
+              aborted = true;
+            });
+            setTimeout(() => resolve('result2'), 50);
+          }),
+      );
+
+      await abortableRequest.execute('cleanup-test', fetcher2);
+      expect(aborted).toBe(false);
+    });
+  });
+
+  describe('cancel', () => {
+    it('should cancel specific request by key', async () => {
+      let aborted = false;
+      const fetcher = vi.fn(
+        async (signal: AbortSignal) =>
+          new Promise((resolve, reject) => {
+            signal.addEventListener('abort', () => {
+              aborted = true;
+              reject(new Error('Cancelled'));
+            });
+            setTimeout(() => resolve('result'), 100);
+          }),
+      );
+
+      const promise = abortableRequest.execute('cancel-key', fetcher);
+
+      await new Promise((resolve) => setTimeout(resolve, 10));
+      abortableRequest.cancel('cancel-key');
+
+      await expect(promise).rejects.toThrow('Cancelled');
+      expect(aborted).toBe(true);
+    });
+
+    it('should do nothing when canceling non-existent key', () => {
+      expect(() => abortableRequest.cancel('non-existent')).not.toThrow();
+    });
+  });
+
+  describe('cancelAll', () => {
+    it('should cancel all pending requests', async () => {
+      const results = { req1: false, req2: false, req3: false };
+
+      const createFetcher = (key: keyof typeof results) =>
+        vi.fn(
+          async (signal: AbortSignal) =>
+            new Promise((resolve, reject) => {
+              signal.addEventListener('abort', () => {
+                results[key] = true;
+                reject(new Error('Cancelled'));
+              });
+              setTimeout(() => resolve(`result-${key}`), 100);
+            }),
+        );
+
+      const promise1 = abortableRequest.execute('key1', createFetcher('req1'));
+      const promise2 = abortableRequest.execute('key2', createFetcher('req2'));
+      const promise3 = abortableRequest.execute('key3', createFetcher('req3'));
+
+      await new Promise((resolve) => setTimeout(resolve, 10));
+      abortableRequest.cancelAll();
+
+      await expect(Promise.all([promise1, promise2, promise3])).rejects.toThrow();
+      expect(results.req1).toBe(true);
+      expect(results.req2).toBe(true);
+      expect(results.req3).toBe(true);
+    });
+  });
+});

package/src/services/utils/abortableRequest.ts

@@ -0,0 +1,67 @@
+/**
+ * Abortable Request Manager
+ *
+ * Provides race condition control for async requests by canceling previous
+ * requests when a new one with the same key is triggered.
+ *
+ * @example
+ * ```ts
+ * const result = await abortableRequest.execute(
+ *   'update-user-profile',
+ *   (signal) => api.updateProfile(data, { signal })
+ * );
+ * ```
+ */
+class AbortableRequestManager {
+  private controllers = new Map<string, AbortController>();
+
+  /**
+   * Execute a request with race condition control
+   * @param key - Unique key to identify the request group
+   * @param fetcher - Request function that accepts AbortSignal
+   * @returns Promise with the request result
+   */
+  async execute<T>(key: string, fetcher: (signal: AbortSignal) => Promise<T>): Promise<T> {
+    // Cancel previous request with same key
+    const existingController = this.controllers.get(key);
+    if (existingController) {
+      existingController.abort('New request triggered');
+    }
+
+    const controller = new AbortController();
+    this.controllers.set(key, controller);
+
+    try {
+      return await fetcher(controller.signal);
+    } finally {
+      // Clean up controller if it's still the active one
+      if (this.controllers.get(key) === controller) {
+        this.controllers.delete(key);
+      }
+    }
+  }
+
+  /**
+   * Manually cancel a request by key
+   * @param key - Request key to cancel
+   */
+  cancel(key: string): void {
+    const controller = this.controllers.get(key);
+    if (controller) {
+      controller.abort('Manually cancelled');
+      this.controllers.delete(key);
+    }
+  }
+
+  /**
+   * Cancel all pending requests
+   */
+  cancelAll(): void {
+    for (const controller of this.controllers.values()) {
+      controller.abort('All requests cancelled');
+    }
+    this.controllers.clear();
+  }
+}
+
+export const abortableRequest = new AbortableRequestManager();