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

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;
-}

package/src/server/routers/tools/codeInterpreter.ts

@@ -1,255 +0,0 @@
-import { type CodeInterpreterToolName, MarketSDK } from '@lobehub/market-sdk';
-import debug from 'debug';
-import { z } from 'zod';
-
-import { DocumentModel } from '@/database/models/document';
-import { FileModel } from '@/database/models/file';
-import { authedProcedure, router } from '@/libs/trpc/lambda';
-import { marketUserInfo, serverDatabase } from '@/libs/trpc/lambda/middleware';
-import { generateTrustedClientToken } from '@/libs/trusted-client';
-import { FileS3 } from '@/server/modules/S3';
-
-const log = debug('lobe-server:tools:code-interpreter');
-
-const codeInterpreterProcedure = authedProcedure.use(serverDatabase).use(marketUserInfo);
-
-// Schema for tool call request
-const callToolSchema = z.object({
-  /** Market access token from OIDC (stored in user settings) */
-  marketAccessToken: z.string().optional(),
-  params: z.record(z.any()),
-  toolName: z.string(),
-  // Session context for isolation
-  topicId: z.string(),
-  userId: z.string(),
-});
-
-// Schema for getting export file upload URL
-const getExportFileUploadUrlSchema = z.object({
-  /** Original filename from sandbox */
-  filename: z.string(),
-  /** Topic ID for organizing files */
-  topicId: z.string(),
-});
-
-// Schema for saving exported file content to document
-const saveExportedFileContentSchema = z.object({
-  /** File content (text content from code-interpreter export) */
-  content: z.string(),
-  /** File ID to associate with the document */
-  fileId: z.string(),
-  /** File MIME type */
-  fileType: z.string(),
-  /** Filename */
-  filename: z.string(),
-  /** File URL */
-  url: z.string(),
-});
-
-export type CallToolInput = z.infer<typeof callToolSchema>;
-export type GetExportFileUploadUrlInput = z.infer<typeof getExportFileUploadUrlSchema>;
-export type SaveExportedFileContentInput = z.infer<typeof saveExportedFileContentSchema>;
-
-export interface CallToolResult {
-  error?: {
-    message: string;
-    name?: string;
-  };
-  result: any;
-  sessionExpiredAndRecreated?: boolean;
-  success: boolean;
-}
-
-export interface GetExportFileUploadUrlResult {
-  /** The download URL after file is uploaded */
-  downloadUrl: string;
-  error?: {
-    message: string;
-  };
-  /** The S3 key where file will be stored */
-  key: string;
-  success: boolean;
-  /** Pre-signed upload URL */
-  uploadUrl: string;
-}
-
-export interface SaveExportedFileContentResult {
-  /** Created document ID */
-  documentId?: string;
-  error?: {
-    message: string;
-  };
-  success: boolean;
-}
-
-export const codeInterpreterRouter = router({
-  callTool: codeInterpreterProcedure.input(callToolSchema).mutation(async ({ input, ctx }) => {
-    const { toolName, params, userId, topicId, marketAccessToken } = input;
-
-    log('Calling cloud code interpreter tool: %s with params: %O', toolName, {
-      params,
-      topicId,
-      userId,
-    });
-    log('Market access token available: %s', marketAccessToken ? 'yes' : 'no');
-
-    // Generate trusted client token if user info is available
-    const trustedClientToken = ctx.marketUserInfo
-      ? generateTrustedClientToken(ctx.marketUserInfo)
-      : undefined;
-
-    try {
-      // Initialize MarketSDK with market access token and trusted client token
-      const market = new MarketSDK({
-        accessToken: marketAccessToken,
-        baseURL: process.env.NEXT_PUBLIC_MARKET_BASE_URL,
-        trustedClientToken,
-      });
-
-      // Call market-sdk's runBuildInTool
-      // API signature: runBuildInTool(toolName, params, context, options?)
-      const response = await market.plugins.runBuildInTool(
-        toolName as CodeInterpreterToolName,
-        params as any,
-        { topicId, userId },
-      );
-
-      log('Cloud code interpreter tool %s response: %O', toolName, response);
-
-      if (!response.success) {
-        return {
-          error: {
-            message: response.error?.message || 'Unknown error',
-            name: response.error?.code,
-          },
-          result: null,
-          sessionExpiredAndRecreated: false,
-          success: false,
-        } as CallToolResult;
-      }
-
-      return {
-        result: response.data?.result,
-        sessionExpiredAndRecreated: response.data?.sessionExpiredAndRecreated || false,
-        success: true,
-      } as CallToolResult;
-    } catch (error) {
-      log('Error calling cloud code interpreter tool %s: %O', toolName, error);
-
-      return {
-        error: {
-          message: (error as Error).message,
-          name: (error as Error).name,
-        },
-        result: null,
-        sessionExpiredAndRecreated: false,
-        success: false,
-      } as CallToolResult;
-    }
-  }),
-
-  /**
-   * Generate a pre-signed upload URL for exporting files from sandbox
-   * The URL can be used by the sandbox to upload the file directly to S3
-   */
-  getExportFileUploadUrl: codeInterpreterProcedure
-    .input(getExportFileUploadUrlSchema)
-
-    // TODO if upload success, should add it path to files db
-    // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    .mutation(async ({ ctx, input }) => {
-      const { filename, topicId } = input;
-
-      log('Generating export file upload URL for: %s in topic: %s', filename, topicId);
-
-      try {
-        const s3 = new FileS3();
-
-        // Generate a unique key for the exported file
-        // Format: code-interpreter-exports/{topicId}/{filename}
-        const key = `code-interpreter-exports/${topicId}/${filename}`;
-
-        // Generate pre-signed upload URL
-        const uploadUrl = await s3.createPreSignedUrl(key);
-
-        // Generate download URL (pre-signed for preview)
-        const downloadUrl = await s3.createPreSignedUrlForPreview(key);
-
-        log('Generated upload URL for key: %s', key);
-
-        return {
-          downloadUrl,
-          key,
-          success: true,
-          uploadUrl,
-        } as GetExportFileUploadUrlResult;
-      } catch (error) {
-        log('Error generating export file upload URL: %O', error);
-
-        return {
-          downloadUrl: '',
-          error: {
-            message: (error as Error).message,
-          },
-          key: '',
-          success: false,
-          uploadUrl: '',
-        } as GetExportFileUploadUrlResult;
-      }
-    }),
-
-  /**
-   * Save exported file content to documents table
-   * This creates a document record linked to the file, allowing content to be retrieved
-   * when querying messages with file attachments
-   */
-  saveExportedFileContent: codeInterpreterProcedure
-    .input(saveExportedFileContentSchema)
-    .use(serverDatabase)
-    .mutation(async ({ ctx, input }) => {
-      const { content, fileId, fileType, filename, url } = input;
-
-      log('Saving exported file content: fileId=%s, filename=%s', fileId, filename);
-
-      try {
-        const documentModel = new DocumentModel(ctx.serverDB, ctx.userId);
-        const fileModel = new FileModel(ctx.serverDB, ctx.userId);
-
-        // Verify the file exists
-        const file = await fileModel.findById(fileId);
-        if (!file) {
-          return {
-            error: { message: 'File not found' },
-            success: false,
-          } as SaveExportedFileContentResult;
-        }
-
-        // Create document record with the file content
-        const document = await documentModel.create({
-          content,
-          fileId,
-          fileType,
-          filename,
-          source: url,
-          sourceType: 'file',
-          title: filename,
-          totalCharCount: content.length,
-          totalLineCount: content.split('\n').length,
-        });
-
-        log('Created document for exported file: documentId=%s, fileId=%s', document.id, fileId);
-
-        return {
-          documentId: document.id,
-          success: true,
-        } as SaveExportedFileContentResult;
-      } catch (error) {
-        log('Error saving exported file content: %O', error);
-
-        return {
-          error: { message: (error as Error).message },
-          success: false,
-        } as SaveExportedFileContentResult;
-      }
-    }),
-});