@lobehub/lobehub 2.0.0-next.194 → 2.0.0-next.196

package/src/store/chat/slices/plugin/action.test.ts

@@ -1429,4 +1429,195 @@ describe('ChatPluginAction', () => {
       });
     });
   });
+
+  describe('Plugin invoke functions use optimisticUpdateToolMessage', () => {
+    const messageId = 'message-id';
+    const payload: ChatToolPayload = {
+      apiName: 'test-api',
+      arguments: '{}',
+      id: 'tool-call-id',
+      identifier: 'test-plugin',
+      type: 'default',
+    };
+
+    describe('invokeMCPTypePlugin', () => {
+      it('should use optimisticUpdateToolMessage for successful result', async () => {
+        const mockResult = {
+          content: 'mcp result content',
+          state: { content: [], isError: false },
+          success: true,
+        };
+
+        // Mock the mcpService
+        const mcpService = await import('@/services/mcp');
+        vi.spyOn(mcpService.mcpService, 'invokeMcpToolCall').mockResolvedValue(mockResult);
+
+        const optimisticUpdateToolMessageMock = vi.fn().mockResolvedValue(undefined);
+
+        act(() => {
+          useChatStore.setState({
+            activeAgentId: 'session-id',
+            messagesMap: { [messageMapKey({ agentId: 'session-id' })]: [] },
+            optimisticUpdateToolMessage: optimisticUpdateToolMessageMock,
+            replaceMessages: vi.fn(),
+            messageOperationMap: {},
+            operations: {},
+          });
+        });
+
+        const { result } = renderHook(() => useChatStore());
+
+        await act(async () => {
+          await result.current.invokeMCPTypePlugin(messageId, payload);
+        });
+
+        expect(optimisticUpdateToolMessageMock).toHaveBeenCalledWith(
+          messageId,
+          {
+            content: mockResult.content,
+            pluginError: undefined,
+            pluginState: mockResult.state,
+          },
+          undefined,
+        );
+      });
+
+      it('should use optimisticUpdateToolMessage for error result', async () => {
+        const mockResult = {
+          content: 'error content',
+          error: { message: 'test error' },
+          state: { content: [], isError: true },
+          success: false,
+        };
+
+        const mcpService = await import('@/services/mcp');
+        vi.spyOn(mcpService.mcpService, 'invokeMcpToolCall').mockResolvedValue(mockResult);
+
+        const optimisticUpdateToolMessageMock = vi.fn().mockResolvedValue(undefined);
+
+        act(() => {
+          useChatStore.setState({
+            activeAgentId: 'session-id',
+            messagesMap: { [messageMapKey({ agentId: 'session-id' })]: [] },
+            optimisticUpdateToolMessage: optimisticUpdateToolMessageMock,
+            replaceMessages: vi.fn(),
+            messageOperationMap: {},
+            operations: {},
+          });
+        });
+
+        const { result } = renderHook(() => useChatStore());
+
+        await act(async () => {
+          await result.current.invokeMCPTypePlugin(messageId, payload);
+        });
+
+        expect(optimisticUpdateToolMessageMock).toHaveBeenCalledWith(
+          messageId,
+          {
+            content: mockResult.content,
+            pluginError: mockResult.error,
+            pluginState: undefined,
+          },
+          undefined,
+        );
+      });
+    });
+
+    describe('invokeKlavisTypePlugin', () => {
+      it('should use optimisticUpdateToolMessage for successful result', async () => {
+        const mockResult = {
+          content: 'klavis result content',
+          state: { data: 'test-data' },
+          success: true,
+        };
+
+        // Mock useToolStore to return a server
+        vi.spyOn(useToolStore, 'getState').mockReturnValue({
+          servers: [{ identifier: 'test-plugin', serverUrl: 'http://test.com' }],
+          callKlavisTool: vi.fn().mockResolvedValue({
+            success: true,
+            data: mockResult,
+          }),
+        } as any);
+
+        const optimisticUpdateToolMessageMock = vi.fn().mockResolvedValue(undefined);
+
+        act(() => {
+          useChatStore.setState({
+            activeAgentId: 'session-id',
+            messagesMap: { [messageMapKey({ agentId: 'session-id' })]: [] },
+            optimisticUpdateToolMessage: optimisticUpdateToolMessageMock,
+            replaceMessages: vi.fn(),
+            messageOperationMap: {},
+            operations: {},
+          });
+        });
+
+        const { result } = renderHook(() => useChatStore());
+
+        await act(async () => {
+          await result.current.invokeKlavisTypePlugin(messageId, payload);
+        });
+
+        expect(optimisticUpdateToolMessageMock).toHaveBeenCalledWith(
+          messageId,
+          {
+            content: mockResult.content,
+            pluginError: undefined,
+            pluginState: mockResult.state,
+          },
+          undefined,
+        );
+      });
+    });
+
+    describe('invokeCloudCodeInterpreterTool', () => {
+      it('should use optimisticUpdateToolMessage for successful result', async () => {
+        const mockResult = {
+          content: 'code interpreter result',
+          state: { output: 'test output' },
+          success: true,
+        };
+
+        // Mock CloudSandboxExecutionRuntime using doMock for dynamic mocking
+        vi.doMock('@lobechat/builtin-tool-cloud-sandbox/executionRuntime', () => ({
+          CloudSandboxExecutionRuntime: class {
+            'test-api' = vi.fn().mockResolvedValue(mockResult);
+          },
+        }));
+
+        const optimisticUpdateToolMessageMock = vi.fn().mockResolvedValue(undefined);
+
+        act(() => {
+          useChatStore.setState({
+            activeAgentId: 'session-id',
+            messagesMap: { [messageMapKey({ agentId: 'session-id' })]: [] },
+            optimisticUpdateToolMessage: optimisticUpdateToolMessageMock,
+            replaceMessages: vi.fn(),
+            messageOperationMap: {},
+            operations: {},
+          });
+        });
+
+        const { result } = renderHook(() => useChatStore());
+
+        await act(async () => {
+          await result.current.invokeCloudCodeInterpreterTool(messageId, payload);
+        });
+
+        expect(optimisticUpdateToolMessageMock).toHaveBeenCalledWith(
+          messageId,
+          {
+            content: mockResult.content,
+            pluginError: undefined,
+            pluginState: mockResult.state,
+          },
+          undefined,
+        );
+
+        vi.doUnmock('@lobechat/builtin-tool-cloud-sandbox/executionRuntime');
+      });
+    });
+  });
 });

package/src/store/chat/slices/plugin/actions/internals.ts

@@ -1,6 +1,6 @@
 /* eslint-disable sort-keys-fix/sort-keys-fix, typescript-sort-keys/interface */
 import { ToolNameResolver } from '@lobechat/context-engine';
-import { type ChatToolPayload, type MessageToolCall, type ToolsCallingContext } from '@lobechat/types';
+import { type ChatToolPayload, type MessageToolCall } from '@lobechat/types';
 import { type LobeChatPluginManifest } from '@lobehub/chat-plugin-sdk';
 import { type StateCreator } from 'zustand/vanilla';
 
@@ -9,8 +9,6 @@ import { useToolStore } from '@/store/tool';
 import { klavisStoreSelectors, pluginSelectors } from '@/store/tool/selectors';
 import { builtinTools } from '@/tools';
 
-import { displayMessageSelectors } from '../../message/selectors';
-
 /**
  * Internal utility methods and runtime state management
  * These are building blocks used by other actions
@@ -20,11 +18,6 @@ export interface PluginInternalsAction {
    * Transform tool calls from runtime format to storage format
    */
   internal_transformToolCalls: (toolCalls: MessageToolCall[]) => ChatToolPayload[];
-
-  /**
-   * Construct tools calling context for plugin invocation
-   */
-  internal_constructToolsCallingContext: (id: string) => ToolsCallingContext | undefined;
 }
 
 export const pluginInternals: StateCreator<
@@ -32,7 +25,7 @@ export const pluginInternals: StateCreator<
   [['zustand/devtools', never]],
   [],
   PluginInternalsAction
-> = (set, get) => ({
+> = () => ({
   internal_transformToolCalls: (toolCalls) => {
     const toolNameResolver = new ToolNameResolver();
 
@@ -77,13 +70,4 @@ export const pluginInternals: StateCreator<
       source: sourceMap[payload.identifier],
     }));
   },
-
-  internal_constructToolsCallingContext: (id: string) => {
-    const message = displayMessageSelectors.getDisplayMessageById(id)(get());
-    if (!message) return;
-
-    return {
-      topicId: message.topicId,
-    };
-  },
 });

package/src/store/chat/slices/plugin/actions/pluginTypes.ts

@@ -261,12 +261,6 @@ export const pluginTypes: StateCreator<
   },
 
   invokeCloudCodeInterpreterTool: async (id, payload) => {
-    const {
-      optimisticUpdateMessageContent,
-      optimisticUpdatePluginState,
-      optimisticUpdateMessagePluginError,
-    } = get();
-
     // Get message to extract topicId
     const message = dbMessageSelectors.getDbMessageById(id)(get());
 
@@ -336,16 +330,16 @@ export const pluginTypes: StateCreator<
 
     const context = operationId ? { operationId } : undefined;
 
-    await Promise.all([
-      optimisticUpdateMessageContent(id, data.content, undefined, context),
-      (async () => {
-        if (data.success && data.state) {
-          await optimisticUpdatePluginState(id, data.state, context);
-        } else if (!data.success && data.error) {
-          await optimisticUpdateMessagePluginError(id, data.error, context);
-        }
-      })(),
-    ]);
+    // Use optimisticUpdateToolMessage to update content and state/error in a single call
+    await get().optimisticUpdateToolMessage(
+      id,
+      {
+        content: data.content,
+        pluginError: data.success ? undefined : data.error,
+        pluginState: data.success ? data.state : undefined,
+      },
+      context,
+    );
 
     // Handle exportFile: save exported file and associate with assistant message (parent)
     if (payload.apiName === 'exportFile' && data.success && data.state) {
@@ -422,12 +416,6 @@ export const pluginTypes: StateCreator<
   },
 
   invokeKlavisTypePlugin: async (id, payload) => {
-    const {
-      optimisticUpdateMessageContent,
-      optimisticUpdatePluginState,
-      optimisticUpdateMessagePluginError,
-    } = get();
-
     let data: MCPToolCallResult | undefined;
 
     // Get message to extract sessionId/topicId
@@ -510,13 +498,16 @@ export const pluginTypes: StateCreator<
     // operationId already declared above, reuse it
     const context = operationId ? { operationId } : undefined;
 
-    await Promise.all([
-      optimisticUpdateMessageContent(id, data.content, undefined, context),
-      (async () => {
-        if (data.success) await optimisticUpdatePluginState(id, data.state, context);
-        else await optimisticUpdateMessagePluginError(id, data.error, context);
-      })(),
-    ]);
+    // Use optimisticUpdateToolMessage to update content and state/error in a single call
+    await get().optimisticUpdateToolMessage(
+      id,
+      {
+        content: data.content,
+        pluginError: data.success ? undefined : data.error,
+        pluginState: data.success ? data.state : undefined,
+      },
+      context,
+    );
 
     return data.content;
   },
@@ -561,12 +552,6 @@ export const pluginTypes: StateCreator<
   },
 
   invokeMCPTypePlugin: async (id, payload) => {
-    const {
-      optimisticUpdateMessageContent,
-      internal_constructToolsCallingContext,
-      optimisticUpdatePluginState,
-      optimisticUpdateMessagePluginError,
-    } = get();
     let data: MCPToolCallResult | undefined;
 
     // Get message to extract agentId/topicId
@@ -586,10 +571,9 @@ export const pluginTypes: StateCreator<
     );
 
     try {
-      const context = internal_constructToolsCallingContext(id);
       const result = await mcpService.invokeMcpToolCall(payload, {
         signal: abortController?.signal,
-        topicId: context?.topicId,
+        topicId: message?.topicId,
       });
 
       if (!!result) data = result;
@@ -620,13 +604,16 @@ export const pluginTypes: StateCreator<
     // operationId already declared above, reuse it
     const context = operationId ? { operationId } : undefined;
 
-    await Promise.all([
-      optimisticUpdateMessageContent(id, data.content, undefined, context),
-      (async () => {
-        if (data.success) await optimisticUpdatePluginState(id, data.state, context);
-        else await optimisticUpdateMessagePluginError(id, data.error, context);
-      })(),
-    ]);
+    // Use optimisticUpdateToolMessage to update content and state/error in a single call
+    await get().optimisticUpdateToolMessage(
+      id,
+      {
+        content: data.content,
+        pluginError: data.success ? undefined : data.error,
+        pluginState: data.success ? data.state : undefined,
+      },
+      context,
+    );
 
     return data.content;
   },

package/packages/database/src/client/db.test.ts

@@ -1,52 +0,0 @@
-import { PGlite } from '@electric-sql/pglite';
-import { beforeEach, describe, expect, it, vi } from 'vitest';
-
-vi.mock('@electric-sql/pglite', () => ({
-  PGlite: vi.fn(() => ({})),
-}));
-
-vi.mock('@electric-sql/pglite/vector', () => ({
-  vector: vi.fn(),
-}));
-
-vi.mock('drizzle-orm/pglite', () => ({
-  drizzle: vi.fn(() => ({
-    dialect: {
-      migrate: vi.fn().mockResolvedValue(undefined),
-    },
-  })),
-}));
-
-beforeEach(() => {
-  vi.clearAllMocks();
-  vi.resetModules();
-});
-
-describe('DatabaseManager', () => {
-  describe('initializeDB', () => {
-    it('should initialize database with PGlite', async () => {
-      const { initializeDB } = await import('./db');
-      await initializeDB();
-
-      expect(PGlite).toHaveBeenCalledWith('idb://lobechat', {
-        extensions: { vector: expect.any(Function) },
-        relaxedDurability: true,
-      });
-    });
-
-    it('should only initialize once when called multiple times', async () => {
-      const { initializeDB } = await import('./db');
-      await Promise.all([initializeDB(), initializeDB()]);
-
-      expect(PGlite).toHaveBeenCalledTimes(1);
-    });
-  });
-
-  describe('clientDB proxy', () => {
-    it('should provide access to database after initialization', async () => {
-      const { clientDB, initializeDB } = await import('./db');
-      await initializeDB();
-      expect(clientDB).toBeDefined();
-    });
-  });
-});

package/packages/database/src/client/db.ts

@@ -1,195 +0,0 @@
-import { PGlite } from '@electric-sql/pglite';
-import { vector } from '@electric-sql/pglite/vector';
-import { sql } from 'drizzle-orm';
-import { PgliteDatabase, drizzle } from 'drizzle-orm/pglite';
-import { Md5 } from 'ts-md5';
-
-import migrations from '../core/migrations.json';
-import { DrizzleMigrationModel } from '../models/drizzleMigration';
-import * as schema from '../schemas';
-
-const pgliteSchemaHashCache = 'LOBE_CHAT_PGLITE_SCHEMA_HASH';
-const DB_NAME = 'lobechat';
-
-type DrizzleInstance = PgliteDatabase<typeof schema>;
-
-class DatabaseManager {
-  private static instance: DatabaseManager;
-  private dbInstance: DrizzleInstance | null = null;
-  private initPromise: Promise<DrizzleInstance> | null = null;
-  private isLocalDBSchemaSynced = false;
-
-  private constructor() {}
-
-  static getInstance() {
-    if (!DatabaseManager.instance) {
-      DatabaseManager.instance = new DatabaseManager();
-    }
-    return DatabaseManager.instance;
-  }
-
-  private async migrate(): Promise<DrizzleInstance> {
-    if (this.isLocalDBSchemaSynced) return this.db;
-
-    let hash: string | undefined;
-    if (typeof localStorage !== 'undefined') {
-      const cacheHash = localStorage.getItem(pgliteSchemaHashCache);
-      hash = Md5.hashStr(JSON.stringify(migrations));
-      // if hash is the same, no need to migrate
-      if (hash === cacheHash) {
-        try {
-          const drizzleMigration = new DrizzleMigrationModel(this.db as any);
-
-          // Check if tables exist in database
-          const tableCount = await drizzleMigration.getTableCounts();
-
-          // If table count > 0, consider database properly initialized
-          if (tableCount > 0) {
-            this.isLocalDBSchemaSynced = true;
-            return this.db;
-          }
-        } catch (error) {
-          console.warn('Error checking table existence, proceeding with migration', error);
-        }
-      }
-    }
-
-    const start = Date.now();
-    try {
-      // @ts-expect-error - migrate internal API
-      await this.db.dialect.migrate(migrations, this.db.session, {});
-
-      if (typeof localStorage !== 'undefined' && hash) {
-        localStorage.setItem(pgliteSchemaHashCache, hash);
-      }
-
-      this.isLocalDBSchemaSynced = true;
-      console.info(`🗂 Migration success, take ${Date.now() - start}ms`);
-    } catch (cause) {
-      console.error('❌ Local database schema migration failed', cause);
-      throw cause;
-    }
-
-    return this.db;
-  }
-
-  async initialize(): Promise<DrizzleInstance> {
-    if (this.initPromise) return this.initPromise;
-
-    this.initPromise = (async () => {
-      if (this.dbInstance) return this.dbInstance;
-
-      const time = Date.now();
-
-      // 直接使用 pglite，自动处理 wasm 加载
-      const pglite = new PGlite(`idb://${DB_NAME}`, {
-        extensions: { vector },
-        relaxedDurability: true,
-      });
-
-      this.dbInstance = drizzle({ client: pglite, schema });
-
-      await this.migrate();
-
-      console.log(`✅ Database initialized in ${Date.now() - time}ms`);
-
-      return this.dbInstance;
-    })();
-
-    return this.initPromise;
-  }
-
-  get db(): DrizzleInstance {
-    if (!this.dbInstance) {
-      throw new Error('Database not initialized. Please call initialize() first.');
-    }
-    return this.dbInstance;
-  }
-
-  createProxy(): DrizzleInstance {
-    return new Proxy({} as DrizzleInstance, {
-      get: (target, prop) => {
-        return this.db[prop as keyof DrizzleInstance];
-      },
-    });
-  }
-
-  async resetDatabase(): Promise<void> {
-    // 1. Close existing PGlite connection
-    if (this.dbInstance) {
-      try {
-        // @ts-ignore
-        await (this.dbInstance.session as any).client.close();
-        console.log('PGlite instance closed successfully.');
-      } catch (e) {
-        console.error('Error closing PGlite instance:', e);
-      }
-    }
-
-    // 2. Reset database instance and initialization state
-    this.dbInstance = null;
-    this.initPromise = null;
-    this.isLocalDBSchemaSynced = false;
-
-    // 3. Delete IndexedDB database
-    return new Promise<void>((resolve, reject) => {
-      if (typeof indexedDB === 'undefined') {
-        console.warn('IndexedDB is not available, cannot delete database');
-        resolve();
-        return;
-      }
-
-      const dbName = `/pglite/${DB_NAME}`;
-      const request = indexedDB.deleteDatabase(dbName);
-
-      request.onsuccess = () => {
-        console.log(`✅ Database '${dbName}' reset successfully`);
-
-        if (typeof localStorage !== 'undefined') {
-          localStorage.removeItem(pgliteSchemaHashCache);
-        }
-
-        resolve();
-      };
-
-      // eslint-disable-next-line unicorn/prefer-add-event-listener
-      request.onerror = (event) => {
-        const error = (event.target as IDBOpenDBRequest)?.error;
-        console.error(`❌ Error resetting database '${dbName}':`, error);
-        reject(
-          new Error(
-            `Failed to reset database '${dbName}'. Error: ${error?.message || 'Unknown error'}`,
-          ),
-        );
-      };
-
-      request.onblocked = (event) => {
-        console.warn(`Deletion of database '${dbName}' is blocked.`, event);
-        reject(
-          new Error(
-            `Failed to reset database '${dbName}' because it is blocked by other open connections.`,
-          ),
-        );
-      };
-    });
-  }
-}
-
-// Export singleton
-const dbManager = DatabaseManager.getInstance();
-
-export const clientDB = dbManager.createProxy();
-
-export const initializeDB = () => dbManager.initialize();
-
-export const resetClientDatabase = async () => {
-  await dbManager.resetDatabase();
-};
-
-export const updateMigrationRecord = async (migrationHash: string) => {
-  await clientDB.execute(
-    sql`INSERT INTO "drizzle"."__drizzle_migrations" ("hash", "created_at") VALUES (${migrationHash}, ${Date.now()});`,
-  );
-
-  await initializeDB();
-};

package/packages/database/src/client/type.ts

@@ -1,6 +0,0 @@
-export interface InitMeta {
-  dbName: string;
-  fsBundle: Blob;
-  vectorBundlePath: string;
-  wasmModule: WebAssembly.Module;
-}