fix: invalidate team vector count cache after vector changes (#6902)

packages/global/common/system/utils.ts

@@ -19,6 +19,27 @@ export const retryFn = async <T>(fn: () => Promise<T>, attempts = 3): Promise<T>
   }
 };
 
+export const withTimeout = async <T>(
+  promise: Promise<T>,
+  timeoutMs: number,
+  timeoutMessage = `Operation timed out after ${timeoutMs}ms`
+): Promise<T> => {
+  let timer: ReturnType<typeof setTimeout> | undefined;
+
+  try {
+    return await Promise.race([
+      promise,
+      new Promise<never>((_, reject) => {
+        timer = setTimeout(() => {
+          reject(new Error(timeoutMessage));
+        }, timeoutMs);
+      })
+    ]);
+  } finally {
+    if (timer) clearTimeout(timer);
+  }
+};
+
 export const batchRun = async <T, R>(
   arr: T[],
   fn: (item: T, index: number) => Promise<R>,

packages/global/test/common/system/utils.test.ts

-import { describe, expect, it, vi } from 'vitest';
-import { delay, retryFn, batchRun } from '@fastgpt/global/common/system/utils';
+import { afterEach, describe, expect, it, vi } from 'vitest';
+import { delay, retryFn, batchRun, withTimeout } from '@fastgpt/global/common/system/utils';
 
 describe('system utils', () => {
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
   describe('delay', () => {
     it('should resolve after specified milliseconds', async () => {
       const start = Date.now();
@@ -89,6 +93,35 @@ describe('system utils', () => {
     });
   });
 
+  describe('withTimeout', () => {
+    it('should resolve when promise settles before timeout', async () => {
+      vi.useFakeTimers();
+
+      const resultPromise = withTimeout(Promise.resolve('success'), 1000);
+
+      await expect(resultPromise).resolves.toBe('success');
+    });
+
+    it('should reject when promise times out', async () => {
+      vi.useFakeTimers();
+
+      const resultPromise = withTimeout(new Promise(() => {}), 1000, 'custom timeout');
+      const assertion = expect(resultPromise).rejects.toThrow('custom timeout');
+
+      await vi.advanceTimersByTimeAsync(1000);
+
+      await assertion;
+    });
+
+    it('should reject with source error when promise rejects before timeout', async () => {
+      vi.useFakeTimers();
+
+      const resultPromise = withTimeout(Promise.reject(new Error('source failure')), 1000);
+
+      await expect(resultPromise).rejects.toThrow('source failure');
+    });
+  });
+
   describe('batchRun', () => {
     it('should process all items', async () => {
       const arr = [1, 2, 3, 4, 5];

packages/service/common/vectorDB/controller.ts

@@ -18,12 +18,37 @@ import {
   setRedisCache,
   getRedisCache,
   delRedisCache,
-  incrValueToCache,
   CacheKeyEnum,
   CacheKeyEnumTime
 } from '../redis/cache';
-import { throttle } from 'lodash';
-import { retryFn } from '@fastgpt/global/common/system/utils';
+import { retryFn, withTimeout } from '@fastgpt/global/common/system/utils';
+import { getLogger, LogCategories } from '../logger';
+
+const logger = getLogger(LogCategories.INFRA.REDIS);
+const TEAM_VECTOR_CACHE_OPERATION_TIMEOUT_MS = 3000;
+
+const runTeamVectorCacheOperation = async <T>({
+  teamId,
+  operation,
+  warnMessage,
+  action
+}: {
+  teamId: string;
+  operation: string;
+  warnMessage: string;
+  action: () => Promise<T>;
+}) => {
+  try {
+    return await withTimeout(
+      action(),
+      TEAM_VECTOR_CACHE_OPERATION_TIMEOUT_MS,
+      `${operation} timed out after ${TEAM_VECTOR_CACHE_OPERATION_TIMEOUT_MS}ms`
+    );
+  } catch (error) {
+    logger.warn(warnMessage, { teamId, error });
+    return undefined;
+  }
+};
 
 const getVectorObj = (): VectorControllerType => {
   if (SEEKDB_ADDRESS) return new SeekVectorCtrl({ type: 'seekdb' });
@@ -40,29 +65,35 @@ const teamVectorCache = {
     return `${CacheKeyEnum.team_vector_count}:${teamId}`;
   },
   get: async function (teamId: string) {
-    const countStr = await getRedisCache(teamVectorCache.getKey(teamId));
+    const countStr = await runTeamVectorCacheOperation({
+      teamId,
+      operation: 'Get team vector count cache',
+      warnMessage: 'Failed to get team vector count cache',
+      action: () => getRedisCache(teamVectorCache.getKey(teamId))
+    });
     if (countStr) {
       return Number(countStr);
     }
     return undefined;
   },
   set: function ({ teamId, count }: { teamId: string; count: number }) {
-    retryFn(() =>
-      setRedisCache(teamVectorCache.getKey(teamId), count, CacheKeyEnumTime.team_vector_count)
-    ).catch();
+    void runTeamVectorCacheOperation({
+      teamId,
+      operation: 'Set team vector count cache',
+      warnMessage: 'Failed to set team vector count cache',
+      action: () =>
+        retryFn(() =>
+          setRedisCache(teamVectorCache.getKey(teamId), count, CacheKeyEnumTime.team_vector_count)
+        )
+    });
   },
-  delete: throttle(
-    function (teamId: string) {
-      return retryFn(() => delRedisCache(teamVectorCache.getKey(teamId))).catch();
-    },
-    30000,
-    {
-      leading: true,
-      trailing: true
-    }
-  ),
-  incr: function (teamId: string, count: number) {
-    retryFn(() => incrValueToCache(teamVectorCache.getKey(teamId), count)).catch();
+  invalidate: async function (teamId: string) {
+    await runTeamVectorCacheOperation({
+      teamId,
+      operation: 'Invalidate team vector count cache',
+      warnMessage: 'Failed to invalidate team vector count cache',
+      action: () => delRedisCache(teamVectorCache.getKey(teamId))
+    });
   }
 };
 
@@ -92,7 +123,7 @@ export const insertDatasetDataVector = async ({
     })
   );
 
-  teamVectorCache.incr(props.teamId, insertIds.length);
+  await teamVectorCache.invalidate(props.teamId);
 
   return {
     tokens,
@@ -102,7 +133,7 @@ export const insertDatasetDataVector = async ({
 
 export const deleteDatasetDataVector: VectorControllerType['delete'] = async (props) => {
   const result = await retryFn(() => Vector.delete(props));
-  teamVectorCache.delete(props.teamId);
+  await teamVectorCache.invalidate(props.teamId);
   return result;
 };
 

packages/service/test/common/vectorDB/controller.test.ts

@@ -26,13 +26,12 @@ import {
 const mockGetRedisCache = vi.fn();
 const mockSetRedisCache = vi.fn();
 const mockDelRedisCache = vi.fn();
-const mockIncrValueToCache = vi.fn();
+const mockLoggerWarn = vi.fn();
 
 vi.mock('@fastgpt/service/common/redis/cache', () => ({
   setRedisCache: (...args: any[]) => mockSetRedisCache(...args),
   getRedisCache: (...args: any[]) => mockGetRedisCache(...args),
   delRedisCache: (...args: any[]) => mockDelRedisCache(...args),
-  incrValueToCache: (...args: any[]) => mockIncrValueToCache(...args),
   CacheKeyEnum: {
     team_vector_count: 'team_vector_count',
     team_point_surplus: 'team_point_surplus',
@@ -45,16 +44,34 @@ vi.mock('@fastgpt/service/common/redis/cache', () => ({
   }
 }));
 
+vi.mock('@fastgpt/service/common/logger', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('@fastgpt/service/common/logger')>();
+
+  return {
+    ...actual,
+    getLogger: () => ({
+      debug: vi.fn(),
+      info: vi.fn(),
+      warn: (...args: any[]) => mockLoggerWarn(...args),
+      error: vi.fn()
+    })
+  };
+});
+
 describe('VectorDB Controller', () => {
   beforeEach(() => {
     resetVectorMocks();
     mockGetRedisCache.mockReset();
     mockSetRedisCache.mockReset();
     mockDelRedisCache.mockReset();
-    mockIncrValueToCache.mockReset();
+    mockLoggerWarn.mockReset();
     mockGetVectorsByText.mockClear();
   });
 
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
   describe('initVectorStore', () => {
     it('should call Vector.init', async () => {
       await initVectorStore();
@@ -143,10 +160,67 @@ describe('VectorDB Controller', () => {
       expect(result).toBe(50);
       expect(mockGetVectorCount).toHaveBeenCalledWith({ teamId: 'team_789' });
     });
+
+    it('should fallback to Vector count when cache read fails', async () => {
+      mockGetRedisCache.mockRejectedValueOnce(new Error('redis down'));
+      mockGetVectorCount.mockResolvedValue(200);
+
+      const result = await getVectorCountByTeamId('team_456');
+
+      expect(result).toBe(200);
+      expect(mockGetVectorCount).toHaveBeenCalledWith({ teamId: 'team_456' });
+      expect(mockLoggerWarn).toHaveBeenCalledWith('Failed to get team vector count cache', {
+        teamId: 'team_456',
+        error: expect.any(Error)
+      });
+    });
+
+    it('should fallback to Vector count when cache read times out', async () => {
+      vi.useFakeTimers();
+      mockGetRedisCache.mockReturnValueOnce(new Promise(() => {}));
+      mockGetVectorCount.mockResolvedValue(120);
+
+      const resultPromise = getVectorCountByTeamId('team_timeout');
+
+      await vi.advanceTimersByTimeAsync(3000);
+
+      await expect(resultPromise).resolves.toBe(120);
+      expect(mockGetVectorCount).toHaveBeenCalledWith({ teamId: 'team_timeout' });
+      expect(mockLoggerWarn).toHaveBeenCalledWith('Failed to get team vector count cache', {
+        teamId: 'team_timeout',
+        error: expect.any(Error)
+      });
+    });
+
+    it('should not block count result when cache write times out', async () => {
+      vi.useFakeTimers();
+      mockGetRedisCache.mockResolvedValue(null);
+      mockGetVectorCount.mockResolvedValue(300);
+      mockSetRedisCache.mockReturnValueOnce(new Promise(() => {}));
+
+      const result = await getVectorCountByTeamId('team_set_timeout');
+
+      expect(result).toBe(300);
+      expect(mockGetVectorCount).toHaveBeenCalledWith({ teamId: 'team_set_timeout' });
+      expect(mockSetRedisCache).toHaveBeenCalledWith(
+        'team_vector_count:team_set_timeout',
+        300,
+        1800
+      );
+
+      await vi.advanceTimersByTimeAsync(3000);
+
+      expect(mockLoggerWarn).toHaveBeenCalledWith('Failed to set team vector count cache', {
+        teamId: 'team_set_timeout',
+        error: expect.any(Error)
+      });
+    });
   });
 
   describe('getVectorCount', () => {
     it('should call Vector.getVectorCount', async () => {
+      mockGetVectorCount.mockResolvedValue(50);
+
       const result = await getVectorCount({ teamId: 'team_1', datasetId: 'dataset_1' });
 
       expect(mockGetVectorCount).toHaveBeenCalledWith({
@@ -207,7 +281,7 @@ describe('VectorDB Controller', () => {
       });
     });
 
-    it('should increment team vector cache', async () => {
+    it('should invalidate team vector cache after insert', async () => {
       mockGetVectorsByText.mockResolvedValue({
         tokens: 50,
         vectors: [[0.1]]
@@ -224,9 +298,68 @@ describe('VectorDB Controller', () => {
         model: mockModel as any
       });
 
-      // Cache increment is called asynchronously
-      await new Promise((resolve) => setTimeout(resolve, 10));
-      expect(mockIncrValueToCache).toHaveBeenCalled();
+      expect(mockDelRedisCache).toHaveBeenCalledWith('team_vector_count:team_abc');
+    });
+
+    it('should return insert result when team vector cache invalidation fails', async () => {
+      mockGetVectorsByText.mockResolvedValue({
+        tokens: 50,
+        vectors: [[0.1]]
+      });
+      mockVectorInsert.mockResolvedValue({
+        insertIds: ['id_1']
+      });
+      mockDelRedisCache.mockRejectedValueOnce(new Error('redis down'));
+
+      const result = await insertDatasetDataVector({
+        teamId: 'team_abc',
+        datasetId: 'dataset_def',
+        collectionId: 'col_ghi',
+        inputs: ['single input'],
+        model: mockModel as any
+      });
+
+      expect(result).toEqual({
+        tokens: 50,
+        insertIds: ['id_1']
+      });
+      expect(mockDelRedisCache).toHaveBeenCalledWith('team_vector_count:team_abc');
+      expect(mockLoggerWarn).toHaveBeenCalledWith('Failed to invalidate team vector count cache', {
+        teamId: 'team_abc',
+        error: expect.any(Error)
+      });
+    });
+
+    it('should return insert result when team vector cache invalidation times out', async () => {
+      vi.useFakeTimers();
+      mockGetVectorsByText.mockResolvedValue({
+        tokens: 50,
+        vectors: [[0.1]]
+      });
+      mockVectorInsert.mockResolvedValue({
+        insertIds: ['id_1']
+      });
+      mockDelRedisCache.mockReturnValueOnce(new Promise(() => {}));
+
+      const resultPromise = insertDatasetDataVector({
+        teamId: 'team_abc',
+        datasetId: 'dataset_def',
+        collectionId: 'col_ghi',
+        inputs: ['single input'],
+        model: mockModel as any
+      });
+
+      await vi.advanceTimersByTimeAsync(3000);
+
+      await expect(resultPromise).resolves.toEqual({
+        tokens: 50,
+        insertIds: ['id_1']
+      });
+      expect(mockDelRedisCache).toHaveBeenCalledWith('team_vector_count:team_abc');
+      expect(mockLoggerWarn).toHaveBeenCalledWith('Failed to invalidate team vector count cache', {
+        teamId: 'team_abc',
+        error: expect.any(Error)
+      });
     });
 
     it('should handle empty inputs', async () => {
@@ -311,6 +444,50 @@ describe('VectorDB Controller', () => {
 
       expect(mockVectorDelete).toHaveBeenCalledWith(props);
       expect(result).toEqual({ deletedCount: 5 });
+      expect(mockDelRedisCache).toHaveBeenCalledWith('team_vector_count:team_cache_test');
+    });
+
+    it('should return delete result when team vector cache invalidation fails', async () => {
+      mockVectorDelete.mockResolvedValue({ deletedCount: 5 });
+      mockDelRedisCache.mockRejectedValueOnce(new Error('redis down'));
+
+      const props = {
+        teamId: 'team_cache_test',
+        id: 'some_id'
+      };
+
+      const result = await deleteDatasetDataVector(props);
+
+      expect(mockVectorDelete).toHaveBeenCalledWith(props);
+      expect(result).toEqual({ deletedCount: 5 });
+      expect(mockDelRedisCache).toHaveBeenCalledWith('team_vector_count:team_cache_test');
+      expect(mockLoggerWarn).toHaveBeenCalledWith('Failed to invalidate team vector count cache', {
+        teamId: 'team_cache_test',
+        error: expect.any(Error)
+      });
+    });
+
+    it('should return delete result when team vector cache invalidation times out', async () => {
+      vi.useFakeTimers();
+      mockVectorDelete.mockResolvedValue({ deletedCount: 5 });
+      mockDelRedisCache.mockReturnValueOnce(new Promise(() => {}));
+
+      const props = {
+        teamId: 'team_cache_test',
+        id: 'some_id'
+      };
+
+      const resultPromise = deleteDatasetDataVector(props);
+
+      await vi.advanceTimersByTimeAsync(3000);
+
+      await expect(resultPromise).resolves.toEqual({ deletedCount: 5 });
+      expect(mockVectorDelete).toHaveBeenCalledWith(props);
+      expect(mockDelRedisCache).toHaveBeenCalledWith('team_vector_count:team_cache_test');
+      expect(mockLoggerWarn).toHaveBeenCalledWith('Failed to invalidate team vector count cache', {
+        teamId: 'team_cache_test',
+        error: expect.any(Error)
+      });
     });
   });
 });