@lobehub/chat 0.140.1 → 0.141.0

package/src/database/core/sync.ts

@@ -0,0 +1,328 @@
+import Debug from 'debug';
+import { throttle, uniqBy } from 'lodash-es';
+import type { WebrtcProvider } from 'y-webrtc';
+import type { Doc, Transaction } from 'yjs';
+
+import {
+  OnAwarenessChange,
+  OnSyncEvent,
+  OnSyncStatusChange,
+  PeerSyncStatus,
+  StartDataSyncParams,
+} from '@/types/sync';
+
+import { LobeDBSchemaMap, LocalDBInstance } from './db';
+
+const LOG_NAME_SPACE = 'DataSync';
+
+class DataSync {
+  private _ydoc: Doc | null = null;
+  private provider: WebrtcProvider | null = null;
+
+  private syncParams!: StartDataSyncParams;
+  private onAwarenessChange!: OnAwarenessChange;
+
+  private waitForConnecting: any;
+
+  logger = Debug(LOG_NAME_SPACE);
+
+  transact(fn: (transaction: Transaction) => unknown) {
+    this._ydoc?.transact(fn);
+  }
+
+  getYMap = (tableKey: keyof LobeDBSchemaMap) => {
+    return this._ydoc?.getMap(tableKey);
+  };
+
+  startDataSync = async (params: StartDataSyncParams) => {
+    this.syncParams = params;
+    this.onAwarenessChange = params.onAwarenessChange;
+
+    // 开发时由于存在 fast refresh 全局实例会缓存在运行时中
+    // 因此需要在每次重新连接时清理上一次的实例
+    if (window.__ONLY_USE_FOR_CLEANUP_IN_DEV) {
+      await this.cleanConnection(window.__ONLY_USE_FOR_CLEANUP_IN_DEV);
+    }
+
+    await this.connect(params);
+  };
+
+  connect = async (params: StartDataSyncParams) => {
+    const {
+      channel,
+      onSyncEvent,
+      onSyncStatusChange,
+      user,
+      onAwarenessChange,
+      signaling = 'wss://y-webrtc-signaling.lobehub.com',
+    } = params;
+    // ====== 1. init yjs doc ====== //
+
+    await this.initYDoc();
+
+    this.logger('[YJS] start to listen sync event...');
+    this.initYjsObserve(onSyncEvent, onSyncStatusChange);
+
+    // ====== 2. init webrtc provider ====== //
+    this.logger(`[WebRTC] init provider... room: ${channel.name}`);
+    const { WebrtcProvider } = await import('y-webrtc');
+
+    // clients connected to the same room-name share document updates
+    this.provider = new WebrtcProvider(channel.name, this._ydoc!, {
+      password: channel.password,
+      signaling: [signaling],
+    });
+
+    // when fast refresh in dev, the provider will be cached in window
+    // so we need to clean it in destory
+    if (process.env.NODE_ENV === 'development') {
+      window.__ONLY_USE_FOR_CLEANUP_IN_DEV = this.provider;
+    }
+
+    this.logger(`[WebRTC] provider init success`);
+
+    // ====== 3. check signaling server connection  ====== //
+
+    // 当本地设备正确连接到 WebRTC Provider 后，触发 status 事件
+    // 当开始连接，则开始监听事件
+    this.provider.on('status', async ({ connected }) => {
+      this.logger('[WebRTC] peer status:', connected);
+      if (connected) {
+        // this.initObserve(onSyncEvent, onSyncStatusChange);
+        onSyncStatusChange?.(PeerSyncStatus.Connecting);
+      }
+    });
+
+    // check the connection with signaling server
+    let connectionCheckCount = 0;
+
+    this.waitForConnecting = setInterval(() => {
+      const signalingConnection: IWebsocketClient = this.provider!.signalingConns[0];
+
+      if (signalingConnection.connected) {
+        onSyncStatusChange?.(PeerSyncStatus.Ready);
+        clearInterval(this.waitForConnecting);
+        return;
+      }
+
+      connectionCheckCount += 1;
+
+      // check for 5 times, or make it failed
+      if (connectionCheckCount > 5) {
+        onSyncStatusChange?.(PeerSyncStatus.Unconnected);
+        clearInterval(this.waitForConnecting);
+      }
+    }, 2000);
+
+    // ====== 4. handle data sync  ====== //
+
+    // 当各方的数据均完成同步后，YJS 对象之间的数据已经一致时，触发 synced 事件
+    this.provider.on('synced', async ({ synced }) => {
+      this.logger('[WebRTC] peer sync status:', synced);
+      if (synced) {
+        this.logger('[WebRTC] start to init yjs data...');
+        onSyncStatusChange?.(PeerSyncStatus.Syncing);
+        await this.initSync();
+        onSyncStatusChange?.(PeerSyncStatus.Synced);
+        this.logger('[WebRTC] yjs data init success');
+      } else {
+        this.logger('[WebRTC] data not sync, try to reconnect in 1s...');
+        // await this.reconnect(params);
+        setTimeout(() => {
+          onSyncStatusChange?.(PeerSyncStatus.Syncing);
+          this.reconnect(params);
+        }, 1000);
+      }
+    });
+
+    // ====== 5. handle awareness  ====== //
+
+    this.initAwareness({ onAwarenessChange, user });
+
+    return this.provider;
+  };
+
+  reconnect = async (params: StartDataSyncParams) => {
+    await this.cleanConnection(this.provider);
+
+    await this.connect(params);
+  };
+
+  async disconnect() {
+    await this.cleanConnection(this.provider);
+  }
+
+  private initYDoc = async () => {
+    if (typeof window === 'undefined') return;
+
+    this.logger('[YJS] init YDoc...');
+    const { Doc } = await import('yjs');
+    this._ydoc = new Doc();
+  };
+
+  private async cleanConnection(provider: WebrtcProvider | null) {
+    if (provider) {
+      this.logger(`[WebRTC] clean Connection...`);
+      this.logger(`[WebRTC] clean awareness...`);
+      provider.awareness.destroy();
+
+      this.logger(`[WebRTC] clean room...`);
+      provider.room?.disconnect();
+      provider.room?.destroy();
+
+      this.logger(`[WebRTC] clean provider...`);
+      provider.disconnect();
+      provider.destroy();
+
+      this.logger(`[WebRTC] clean yjs doc...`);
+      this._ydoc?.destroy();
+
+      this.logger(`[WebRTC] -------------------`);
+    }
+  }
+
+  private initSync = async () => {
+    await Promise.all(
+      ['sessions', 'sessionGroups', 'topics', 'messages', 'plugins'].map(async (tableKey) =>
+        this.loadDataFromDBtoYjs(tableKey as keyof LobeDBSchemaMap),
+      ),
+    );
+  };
+
+  private initYjsObserve = (onEvent: OnSyncEvent, onSyncStatusChange: OnSyncStatusChange) => {
+    ['sessions', 'sessionGroups', 'topics', 'messages', 'plugins'].forEach((tableKey) => {
+      // listen yjs change
+      this.observeYMapChange(tableKey as keyof LobeDBSchemaMap, onEvent, onSyncStatusChange);
+    });
+  };
+
+  private observeYMapChange = (
+    tableKey: keyof LobeDBSchemaMap,
+    onEvent: OnSyncEvent,
+    onSyncStatusChange: OnSyncStatusChange,
+  ) => {
+    const table = LocalDBInstance[tableKey];
+    const yItemMap = this.getYMap(tableKey);
+    const updateSyncEvent = throttle(onEvent, 1000);
+
+    // 定义一个变量来保存定时器的ID
+    // eslint-disable-next-line no-undef
+    let debounceTimer: NodeJS.Timeout;
+
+    yItemMap?.observe(async (event) => {
+      // abort local change
+      if (event.transaction.local) return;
+
+      // 每次有变更时，都先清除之前的定时器（如果有的话），然后设置新的定时器
+      clearTimeout(debounceTimer);
+
+      onSyncStatusChange(PeerSyncStatus.Syncing);
+
+      this.logger(`[YJS] observe ${tableKey} changes:`, event.keysChanged.size);
+      const pools = Array.from(event.keys).map(async ([id, payload]) => {
+        const item: any = yItemMap.get(id);
+
+        switch (payload.action) {
+          case 'add':
+          case 'update': {
+            await table.put(item, id);
+
+            break;
+          }
+
+          case 'delete': {
+            await table.delete(id);
+            break;
+          }
+        }
+      });
+
+      await Promise.all(pools);
+
+      updateSyncEvent(tableKey);
+
+      // 设置定时器，2000ms 后更新状态为'synced'
+      debounceTimer = setTimeout(() => {
+        onSyncStatusChange(PeerSyncStatus.Synced);
+      }, 2000);
+    });
+  };
+
+  private loadDataFromDBtoYjs = async (tableKey: keyof LobeDBSchemaMap) => {
+    const table = LocalDBInstance[tableKey];
+    const items = await table.toArray();
+    const yItemMap = this.getYMap(tableKey);
+
+    // 定义每批次最多包含的数据条数
+    const batchSize = 50;
+
+    // 计算总批次数
+    const totalBatches = Math.ceil(items.length / batchSize);
+
+    for (let i = 0; i < totalBatches; i++) {
+      // 计算当前批次的起始和结束索引
+      const start = i * batchSize;
+      const end = start + batchSize;
+
+      // 获取当前批次的数据
+      const batchItems = items.slice(start, end);
+
+      // 将当前批次的数据推送到 Yjs 中
+      this._ydoc?.transact(() => {
+        batchItems.forEach((item) => {
+          yItemMap!.set(item.id, item);
+        });
+      });
+    }
+
+    this.logger('[DB]:', tableKey, yItemMap?.size);
+  };
+
+  private initAwareness = ({ user }: Pick<StartDataSyncParams, 'user' | 'onAwarenessChange'>) => {
+    if (!this.provider) return;
+
+    const awareness = this.provider.awareness;
+
+    awareness.setLocalState({ clientID: awareness.clientID, user });
+    this.onAwarenessChange?.([{ ...user, clientID: awareness.clientID, current: true }]);
+
+    awareness.on('change', () => this.syncAwarenessToUI());
+  };
+
+  private syncAwarenessToUI = async () => {
+    const awareness = this.provider?.awareness;
+
+    if (!awareness) return;
+
+    const state = Array.from(awareness.getStates().values()).map((s) => ({
+      ...s.user,
+      clientID: s.clientID,
+      current: s.clientID === awareness.clientID,
+    }));
+
+    this.onAwarenessChange?.(uniqBy(state, 'id'));
+  };
+}
+
+export const dataSync = new DataSync();
+
+interface IWebsocketClient {
+  binaryType: 'arraybuffer' | 'blob' | null;
+  connect(): void;
+  connected: boolean;
+  connecting: boolean;
+  destroy(): void;
+  disconnect(): void;
+  lastMessageReceived: number;
+  send(message: any): void;
+  shouldConnect: boolean;
+  unsuccessfulReconnects: number;
+  url: string;
+  ws: WebSocket;
+}
+
+declare global {
+  interface Window {
+    __ONLY_USE_FOR_CLEANUP_IN_DEV?: WebrtcProvider | null;
+  }
+}

package/src/database/models/__tests__/message.test.ts

@@ -1,6 +1,5 @@
 import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 
-import { DB_Message } from '@/database/schemas/message';
 import { ChatMessage } from '@/types/message';
 
 import { CreateMessageParams, MessageModel } from '../message';

package/src/database/models/__tests__/plugin.test.ts

@@ -59,11 +59,14 @@ describe('PluginModel', () => {
   describe('update', () => {
     it('should update a plugin', async () => {
       await PluginModel.create(pluginData);
-      const updatedPluginData: DB_Plugin = { ...pluginData, type: 'customPlugin' };
+      const updatedPluginData: DB_Plugin = {
+        ...pluginData,
+        type: 'customPlugin',
+      };
       await PluginModel.update(pluginData.identifier, updatedPluginData);
       const plugins = await PluginModel.getList();
       expect(plugins).toHaveLength(1);
-      expect(plugins[0]).toEqual(updatedPluginData);
+      expect(plugins[0]).toEqual({ ...updatedPluginData, updatedAt: expect.any(Number) });
     });
   });
 

package/src/database/models/file.ts

@@ -11,7 +11,7 @@ class _FileModel extends BaseModel<'files'> {
   async create(file: DB_File) {
     const id = nanoid();
 
-    return this._add(file, `file-${id}`);
+    return this._addWithSync(file, `file-${id}`);
   }
 
   async findById(id: string) {

package/src/database/models/message.ts

@@ -107,13 +107,13 @@ class _MessageModel extends BaseModel {
 
     const messageData: DB_Message = this.mapChatMessageToDBMessage(data as ChatMessage);
 
-    return this._add(messageData, id);
+    return this._addWithSync(messageData, id);
   }
 
   async batchCreate(messages: ChatMessage[]) {
     const data: DB_Message[] = messages.map((m) => this.mapChatMessageToDBMessage(m));
 
-    return this._batchAdd(data);
+    return this._batchAdd(data, { withSync: true });
   }
 
   async duplicateMessages(messages: ChatMessage[]): Promise<ChatMessage[]> {
@@ -123,36 +123,14 @@ class _MessageModel extends BaseModel {
     return duplicatedMessages;
   }
 
-  async createDuplicateMessages(messages: ChatMessage[]): Promise<ChatMessage[]> {
-    // 创建一个映射来存储原始消息ID和复制消息ID之间的关系
-    const idMapping = new Map<string, string>();
-
-    // 首先复制所有消息，并为每个复制的消息生成新的ID
-    const duplicatedMessages = messages.map((originalMessage) => {
-      const newId = nanoid();
-      idMapping.set(originalMessage.id, newId);
-
-      return { ...originalMessage, id: newId };
-    });
-
-    // 更新 parentId 为复制后的新ID
-    for (const duplicatedMessage of duplicatedMessages) {
-      if (duplicatedMessage.parentId && idMapping.has(duplicatedMessage.parentId)) {
-        duplicatedMessage.parentId = idMapping.get(duplicatedMessage.parentId);
-      }
-    }
-
-    return duplicatedMessages;
-  }
-
   // **************** Delete *************** //
 
   async delete(id: string) {
-    return this.table.delete(id);
+    return super._deleteWithSync(id);
   }
 
   async clearTable() {
-    return this.table.clear();
+    return this._clearWithSync();
   }
 
   /**
@@ -178,13 +156,32 @@ class _MessageModel extends BaseModel {
     const messageIds = await query.primaryKeys();
 
     // Use the bulkDelete method to delete all selected messages in bulk
-    return this.table.bulkDelete(messageIds);
+    return this._bulkDeleteWithSync(messageIds);
+  }
+
+  async batchDeleteBySessionId(sessionId: string): Promise<void> {
+    // If topicId is specified, use both assistantId and topicId as the filter criteria in the query.
+    // Otherwise, filter by assistantId and require that topicId is undefined.
+    const messageIds = await this.table.where('sessionId').equals(sessionId).primaryKeys();
+
+    // Use the bulkDelete method to delete all selected messages in bulk
+    return this._bulkDeleteWithSync(messageIds);
+  }
+
+  /**
+   * Delete all messages associated with the topicId
+   * @param topicId
+   */
+  async batchDeleteByTopicId(topicId: string): Promise<void> {
+    const messageIds = await this.table.where('topicId').equals(topicId).primaryKeys();
+
+    return this._bulkDeleteWithSync(messageIds);
   }
 
   // **************** Update *************** //
 
   async update(id: string, data: DeepPartial<DB_Message>) {
-    return this._update(id, data);
+    return super._updateWithSync(id, data);
   }
 
   async updatePluginState(id: string, key: string, value: any) {
@@ -202,7 +199,7 @@ class _MessageModel extends BaseModel {
    */
   async batchUpdate(messageIds: string[], updateFields: Partial<DB_Message>): Promise<number> {
     // Retrieve the messages by their IDs
-    const messagesToUpdate = await this.table.where(':id').anyOf(messageIds).toArray();
+    const messagesToUpdate = await this.table.where('id').anyOf(messageIds).toArray();
 
     // Update the specified fields of each message
     const updatedMessages = messagesToUpdate.map((message) => ({
@@ -211,13 +208,35 @@ class _MessageModel extends BaseModel {
     }));
 
     // Use the bulkPut method to update the messages in bulk
-    await this.table.bulkPut(updatedMessages);
+    await this._bulkPutWithSync(updatedMessages);
 
     return updatedMessages.length;
   }
 
   // **************** Helper *************** //
 
+  private async createDuplicateMessages(messages: ChatMessage[]): Promise<ChatMessage[]> {
+    // 创建一个映射来存储原始消息ID和复制消息ID之间的关系
+    const idMapping = new Map<string, string>();
+
+    // 首先复制所有消息，并为每个复制的消息生成新的ID
+    const duplicatedMessages = messages.map((originalMessage) => {
+      const newId = nanoid();
+      idMapping.set(originalMessage.id, newId);
+
+      return { ...originalMessage, id: newId };
+    });
+
+    // 更新 parentId 为复制后的新ID
+    for (const duplicatedMessage of duplicatedMessages) {
+      if (duplicatedMessage.parentId && idMapping.has(duplicatedMessage.parentId)) {
+        duplicatedMessage.parentId = idMapping.get(duplicatedMessage.parentId);
+      }
+    }
+
+    return duplicatedMessages;
+  }
+
   private mapChatMessageToDBMessage(message: ChatMessage): DB_Message {
     const { extra, ...messageData } = message;
 

package/src/database/models/plugin.ts

@@ -21,14 +21,13 @@ class _PluginModel extends BaseModel {
   getList = async (): Promise<DB_Plugin[]> => {
     return this.table.toArray();
   };
-
   // **************** Create *************** //
 
   create = async (plugin: InstallPluginParams) => {
     const old = await this.table.get(plugin.identifier);
     const dbPlugin = this.mapToDBPlugin(plugin);
 
-    return this.table.put(merge(old, dbPlugin), plugin.identifier);
+    return this._putWithSync(merge(old, dbPlugin), plugin.identifier);
   };
 
   batchCreate = async (plugins: LobeTool[]) => {
@@ -39,16 +38,18 @@ class _PluginModel extends BaseModel {
   // **************** Delete *************** //
 
   delete(id: string) {
-    return this.table.delete(id);
+    return this._deleteWithSync(id);
   }
   clear() {
-    return this.table.clear();
+    return this._clearWithSync();
   }
 
   // **************** Update *************** //
 
   update: (id: string, value: Partial<DB_Plugin>) => Promise<number> = async (id, value) => {
-    return this.table.update(id, value);
+    const { success } = await this._updateWithSync(id, value);
+
+    return success;
   };
 
   // **************** Helper *************** //

package/src/database/models/session.ts

@@ -3,7 +3,6 @@ import { DeepPartial } from 'utility-types';
 import { DEFAULT_AGENT_LOBE_SESSION } from '@/const/session';
 import { BaseModel } from '@/database/core';
 import { DBModel } from '@/database/core/types/db';
-import { SessionGroupModel } from '@/database/models/sessionGroup';
 import { DB_Session, DB_SessionSchema } from '@/database/schemas/session';
 import { LobeAgentConfig } from '@/types/agent';
 import {
@@ -16,6 +15,10 @@ import {
 import { merge } from '@/utils/merge';
 import { uuid } from '@/utils/uuid';
 
+import { MessageModel } from './message';
+import { SessionGroupModel } from './sessionGroup';
+import { TopicModel } from './topic';
+
 class _SessionModel extends BaseModel {
   constructor() {
     super('sessions', DB_SessionSchema);
@@ -173,7 +176,7 @@ class _SessionModel extends BaseModel {
   async create(type: 'agent' | 'group', defaultValue: Partial<LobeAgentSession>, id = uuid()) {
     const data = merge(DEFAULT_AGENT_LOBE_SESSION, { type, ...defaultValue });
     const dataDB = this.mapToDB_Session(data);
-    return this._add(dataDB, id);
+    return this._addWithSync(dataDB, id);
   }
 
   async batchCreate(sessions: LobeAgentSession[]) {
@@ -200,7 +203,7 @@ class _SessionModel extends BaseModel {
 
     const newSession = merge(session, { meta: { title: newTitle } });
 
-    return this._add(newSession, uuid());
+    return this._addWithSync(newSession, uuid());
   }
 
   // **************** Delete *************** //
@@ -211,32 +214,28 @@ class _SessionModel extends BaseModel {
   async delete(id: string) {
     return this.db.transaction('rw', [this.table, this.db.topics, this.db.messages], async () => {
       // Delete all topics associated with the session
-      const topics = await this.db.topics.where('sessionId').equals(id).toArray();
-      const topicIds = topics.map((topic) => topic.id);
-      if (topicIds.length > 0) {
-        await this.db.topics.bulkDelete(topicIds);
-      }
+      await TopicModel.batchDeleteBySessionId(id);
 
       // Delete all messages associated with the session
-      const messages = await this.db.messages.where('sessionId').equals(id).toArray();
-      const messageIds = messages.map((message) => message.id);
-      if (messageIds.length > 0) {
-        await this.db.messages.bulkDelete(messageIds);
-      }
+      await MessageModel.batchDeleteBySessionId(id);
 
       // Finally, delete the session itself
-      await this.table.delete(id);
+      await this._deleteWithSync(id);
     });
   }
 
+  async batchDelete(ids: string[]) {
+    return this._bulkDeleteWithSync(ids);
+  }
+
   async clearTable() {
-    return this.table.clear();
+    return this._clearWithSync();
   }
 
   // **************** Update *************** //
 
   async update(id: string, data: Partial<DB_Session>) {
-    return this._update(id, data);
+    return super._updateWithSync(id, data);
   }
 
   async updatePinned(id: string, pinned: boolean) {

package/src/database/models/sessionGroup.ts

@@ -42,7 +42,7 @@ class _SessionGroupModel extends BaseModel {
   // **************** Create *************** //
 
   async create(name: string, sort?: number, id = nanoid()) {
-    return this._add({ name, sort }, id);
+    return this._addWithSync({ name, sort }, id);
   }
 
   async batchCreate(groups: SessionGroups) {
@@ -51,31 +51,37 @@ class _SessionGroupModel extends BaseModel {
 
   // **************** Delete *************** //
   async delete(id: string, removeGroupItem: boolean = false) {
-    this.db.sessions.toCollection().modify((session) => {
+    const { SessionModel } = await import('./session');
+    this.db.sessions.toCollection().modify(async (session) => {
       //  update all session associated with the sessionGroup to default
-      if (session.group === id) session.group = 'default';
+      if (session.group === id) {
+        await SessionModel.update(session.id, { group: 'default' });
+      }
     });
+
     if (!removeGroupItem) {
-      return this.table.delete(id);
+      return this._deleteWithSync(id);
     } else {
-      return this.db.sessions.where('group').equals(id).delete();
+      const sessionIds = await this.db.sessions.where('group').equals(id).primaryKeys();
+
+      return await SessionModel.batchDelete(sessionIds);
     }
   }
 
   async clear() {
-    this.table.clear();
+    await this._clearWithSync();
   }
 
   // **************** Update *************** //
 
   async update(id: string, data: Partial<DB_SessionGroup>) {
-    return super._update(id, data);
+    return super._updateWithSync(id, data);
   }
 
   async updateOrder(sortMap: { id: string; sort: number }[]) {
     return this.db.transaction('rw', this.table, async () => {
       for (const { id, sort } of sortMap) {
-        await this.table.update(id, { sort });
+        await this.update(id, { sort });
       }
     });
   }