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

package/src/services/mcp.test.ts

@@ -36,6 +36,11 @@ vi.mock('@lobechat/utils', () => ({
 
 vi.mock('@/libs/trpc/client', () => ({
   toolsClient: {
+    market: {
+      callCloudMcpEndpoint: {
+        mutate: vi.fn(),
+      },
+    },
     mcp: {
       callTool: {
         mutate: vi.fn(),
@@ -53,6 +58,7 @@ vi.mock('@/utils/electron/ipc', () => ({
 
 vi.mock('./discover', () => ({
   discoverService: {
+    injectMPToken: vi.fn().mockResolvedValue(undefined),
     reportPluginCall: vi.fn().mockResolvedValue(undefined),
   },
 }));
@@ -129,22 +135,20 @@ describe('MCPService', () => {
             name: 'test-plugin',
             env: { API_KEY: 'test-key' },
           },
+          meta: {
+            customPluginInfo: undefined,
+            isCustomPlugin: false,
+            sessionId: 'topic-1',
+            version: '1.0.0',
+          },
           toolName: 'testMethod',
         },
         { signal: undefined },
       );
 
-      // Wait for async reporting to complete
+      // For SSE type, reporting is handled by server-side, frontend should NOT call reportPluginCall
       await new Promise((resolve) => setTimeout(resolve, 10));
-
-      expect(discoverService.reportPluginCall).toHaveBeenCalled();
-      const reportCall = vi.mocked(discoverService.reportPluginCall).mock.calls[0][0];
-      expect(reportCall).toMatchObject({
-        identifier: 'test-plugin',
-        methodName: 'testMethod',
-        success: true,
-        isCustomPlugin: false,
-      });
+      expect(discoverService.reportPluginCall).not.toHaveBeenCalled();
     });
 
     it('should invoke tool call with custom plugin', async () => {
@@ -274,28 +278,22 @@ describe('MCPService', () => {
 
       await expect(mcpService.invokeMcpToolCall(payload, {})).rejects.toThrow('Tool call failed');
 
-      // Wait for async reporting to complete
+      // For SSE type, reporting is handled by server-side, frontend should NOT call reportPluginCall
       await new Promise((resolve) => setTimeout(resolve, 10));
-
-      expect(discoverService.reportPluginCall).toHaveBeenCalled();
-      const reportCall = vi.mocked(discoverService.reportPluginCall).mock.calls[0][0];
-      expect(reportCall).toMatchObject({
-        success: false,
-        errorCode: 'CALL_FAILED',
-        errorMessage: 'Tool call failed',
-      });
+      expect(discoverService.reportPluginCall).not.toHaveBeenCalled();
     });
 
-    it('should calculate request and response sizes correctly', async () => {
-      const { toolsClient } = await import('@/libs/trpc/client');
+    it('should call toolsClient.market.callCloudMcpEndpoint for cloud type and not report from frontend', async () => {
       const { discoverService } = await import('./discover');
+      const { toolsClient } = await import('@/libs/trpc/client');
 
+      // Use cloud type which now reports from server-side
       const mockPlugin = {
         customParams: {
-          mcp: { type: 'sse' },
+          mcp: { type: 'cloud' },
         },
         manifest: {
-          meta: { title: 'Size Test Plugin' },
+          meta: { title: 'Cloud Plugin' },
           version: '1.0.0',
         },
       };
@@ -303,32 +301,44 @@ describe('MCPService', () => {
       mockPluginSelectors.getInstalledPluginById.mockReturnValue(() => mockPlugin);
       mockPluginSelectors.getCustomPluginById.mockReturnValue(() => null);
 
+      // Mock the toolsClient for cloud type
       const mockResult = {
         content: 'response data',
         state: {
-          content: [{ text: 'response data', type: 'text' }],
+          content: [{ text: 'response data', type: 'text' as const }],
         },
         success: true,
       };
-      vi.mocked(toolsClient.mcp.callTool.mutate).mockResolvedValue(mockResult);
+      vi.mocked(toolsClient.market.callCloudMcpEndpoint.mutate).mockResolvedValue(mockResult);
 
       const payload: ChatToolPayload = {
         id: 'tool-call-6',
-        identifier: 'size-plugin',
-        apiName: 'sizeMethod',
+        identifier: 'cloud-plugin',
+        apiName: 'cloudMethod',
         arguments: '{"key": "value"}',
         type: 'standalone',
       };
 
-      await mcpService.invokeMcpToolCall(payload, {});
+      const result = await mcpService.invokeMcpToolCall(payload, { topicId: 'topic-123' });
+
+      expect(result).toEqual(mockResult);
+      expect(toolsClient.market.callCloudMcpEndpoint.mutate).toHaveBeenCalledWith({
+        apiParams: { key: 'value' },
+        identifier: 'cloud-plugin',
+        meta: {
+          customPluginInfo: undefined,
+          isCustomPlugin: false,
+          sessionId: 'topic-123',
+          version: '1.0.0',
+        },
+        toolName: 'cloudMethod',
+      });
 
       // Wait for async reporting
       await new Promise((resolve) => setTimeout(resolve, 10));
 
-      expect(discoverService.reportPluginCall).toHaveBeenCalled();
-      const reportCall = vi.mocked(discoverService.reportPluginCall).mock.calls[0][0];
-      expect(reportCall.requestSizeBytes).toBeGreaterThan(0);
-      expect(reportCall.responseSizeBytes).toBeGreaterThan(0);
+      // Cloud type should NOT report from frontend (handled server-side)
+      expect(discoverService.reportPluginCall).not.toHaveBeenCalled();
     });
 
     it('should handle abort signal', async () => {
@@ -372,9 +382,8 @@ describe('MCPService', () => {
       expect(result).toEqual(mockResult);
     });
 
-    it('should report custom plugin info correctly', async () => {
+    it('should pass meta to server for custom plugin', async () => {
       const { toolsClient } = await import('@/libs/trpc/client');
-      const { discoverService } = await import('./discover');
 
       const mockCustomPlugin = {
         customParams: {
@@ -406,19 +415,24 @@ describe('MCPService', () => {
         type: 'standalone',
       };
 
-      await mcpService.invokeMcpToolCall(payload, {});
-
-      // Wait for async reporting
-      await new Promise((resolve) => setTimeout(resolve, 10));
+      await mcpService.invokeMcpToolCall(payload, { topicId: 'topic-123' });
 
-      expect(discoverService.reportPluginCall).toHaveBeenCalled();
-      const reportCall = vi.mocked(discoverService.reportPluginCall).mock.calls[0][0];
-      expect(reportCall.isCustomPlugin).toBe(true);
-      expect(reportCall.customPluginInfo).toEqual({
-        avatar: '🔧',
-        description: 'Custom tool description',
-        name: 'Custom Tool',
-      });
+      // Verify meta is passed to server
+      expect(toolsClient.mcp.callTool.mutate).toHaveBeenCalledWith(
+        expect.objectContaining({
+          meta: {
+            customPluginInfo: {
+              avatar: '🔧',
+              description: 'Custom tool description',
+              name: 'Custom Tool',
+            },
+            isCustomPlugin: true,
+            sessionId: 'topic-123',
+            version: '3.0.0',
+          },
+        }),
+        expect.anything(),
+      );
     });
   });
 

package/src/services/mcp.ts

@@ -1,12 +1,16 @@
 import { CURRENT_VERSION, isDesktop } from '@lobechat/const';
-import { type ChatToolPayload, type CheckMcpInstallResult, type CustomPluginMetadata } from '@lobechat/types';
+import {
+  type ChatToolPayload,
+  type CheckMcpInstallResult,
+  type CustomPluginMetadata,
+} from '@lobechat/types';
 import { isLocalOrPrivateUrl, safeParseJSON } from '@lobechat/utils';
 import { type PluginManifest } from '@lobehub/market-sdk';
 import { type CallReportRequest } from '@lobehub/market-types';
 import superjson from 'superjson';
 
 import { type MCPToolCallResult } from '@/libs/mcp';
-import { lambdaClient, toolsClient } from '@/libs/trpc/client';
+import { toolsClient } from '@/libs/trpc/client';
 import { ensureElectronIpc } from '@/utils/electron/ipc';
 
 import { discoverService } from './discover';
@@ -31,6 +35,8 @@ class MCPService {
     payload: ChatToolPayload,
     { signal, topicId }: { signal?: AbortSignal; topicId?: string },
   ) {
+    await discoverService.injectMPToken();
+
     const { pluginSelectors } = await import('@/store/tool/selectors');
     const { getToolStoreState } = await import('@/store/tool/store');
 
@@ -80,12 +86,28 @@ class MCPService {
 
     const isStdio = plugin?.customParams?.mcp?.type === 'stdio';
     const isCloud = plugin?.customParams?.mcp?.type === 'cloud';
+    const isCustomPlugin = !!customPlugin;
+
+    // Build meta for server-side reporting
+    const meta = {
+      customPluginInfo: isCustomPlugin
+        ? {
+            avatar: plugin.manifest?.meta.avatar,
+            description: plugin.manifest?.meta.description,
+            name: plugin.manifest?.meta.title,
+          }
+        : undefined,
+      isCustomPlugin,
+      sessionId: topicId,
+      version: plugin.manifest?.version || 'unknown',
+    };
 
     const data = {
       // For desktop IPC, always pass a record/object for tool "arguments"
       // (IPC layer will superjson serialize the whole payload).
       args: isDesktop && isStdio ? (safeParseJSON(args) ?? {}) : args,
       env: connection?.type === 'stdio' ? params.env : (pluginSettings ?? connection?.env),
+      meta,
       params,
       toolName: apiName,
     };
@@ -103,13 +125,14 @@ class MCPService {
         // Parse args
         const apiParams = safeParseJSON(args) || {};
 
-        // Call cloud gateway via lambda market endpoint
+        // Call cloud gateway via tools market endpoint
         // Server will automatically get user access token from database
         // and format the result to MCPToolCallResult
-        // @ts-ignore tsgo 误报错误
-        result = await lambdaClient.market.callCloudMcpEndpoint.mutate({
+        // Server-side also handles telemetry reporting
+        result = await toolsClient.market.callCloudMcpEndpoint.mutate({
           apiParams,
           identifier,
+          meta,
           toolName: apiName,
         });
       } else if (isDesktop && isStdio) {
@@ -134,50 +157,46 @@ class MCPService {
       // Rethrow error, maintain original error handling logic
       throw error;
     } finally {
-      // Asynchronously report call result without affecting main flow
-      const callEndTime = Date.now();
-      const callDurationMs = callEndTime - callStartTime;
-
-      // Calculate request size
-      const inputParams = safeParseJSON(args) || args;
-
-      const requestSizeBytes = calculateObjectSizeBytes(inputParams);
-      // Calculate response size
-      const responseSizeBytes = success && result ? calculateObjectSizeBytes(result.state) : 0;
-
-      const isCustomPlugin = !!customPlugin;
-      // Construct report data
-      const reportData: CallReportRequest = {
-        callDurationMs,
-        customPluginInfo: isCustomPlugin
-          ? {
-              avatar: plugin.manifest?.meta.avatar,
-              description: plugin.manifest?.meta.description,
-              name: plugin.manifest?.meta.title,
-            }
-          : undefined,
-        errorCode,
-        errorMessage,
-        identifier,
-        isCustomPlugin,
-        metadata: {
-          appVersion: CURRENT_VERSION,
-          command: plugin.customParams?.mcp?.command,
-          mcpType: plugin.customParams?.mcp?.type,
-        },
-        methodName: apiName,
-        methodType: 'tool' as const,
-        requestSizeBytes,
-        responseSizeBytes,
-        sessionId: topicId,
-        success,
-        version: plugin.manifest!.version || 'unknown',
-      };
+      // HTTP/SSE/streamable types: reporting is handled by server-side via mcp.callTool
+      // Cloud type: reporting is handled by server-side via market.callCloudMcpEndpoint
+      // Only report from frontend for stdio types (desktop only)
+      if (isStdio) {
+        const callEndTime = Date.now();
+        const callDurationMs = callEndTime - callStartTime;
+
+        // Calculate request size
+        const inputParams = safeParseJSON(args) || args;
+        const requestSizeBytes = calculateObjectSizeBytes(inputParams);
+        // Calculate response size
+        const responseSizeBytes = success && result ? calculateObjectSizeBytes(result.state) : 0;
 
-      // Asynchronously report without affecting main flow
-      discoverService.reportPluginCall(reportData).catch((reportError) => {
-        console.warn('Failed to report MCP tool call:', reportError);
-      });
+        // Construct report data
+        const reportData: CallReportRequest = {
+          callDurationMs,
+          customPluginInfo: meta.customPluginInfo,
+          errorCode,
+          errorMessage,
+          identifier,
+          isCustomPlugin,
+          metadata: {
+            appVersion: CURRENT_VERSION,
+            command: plugin.customParams?.mcp?.command,
+            mcpType: plugin.customParams?.mcp?.type,
+          },
+          methodName: apiName,
+          methodType: 'tool' as const,
+          requestSizeBytes,
+          responseSizeBytes,
+          sessionId: topicId,
+          success,
+          version: meta.version,
+        };
+
+        // Asynchronously report without affecting main flow
+        discoverService.reportPluginCall(reportData).catch((reportError) => {
+          console.warn('Failed to report MCP tool call:', reportError);
+        });
+      }
     }
   }
 

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