perf: mcp server status checking logic

2025-11-25 11:29:39 +08:00 · 2025-11-04 17:32:05 +08:00
parent 3ee7736361
commit 1afecf01e4
11 changed files with 182 additions and 90 deletions
--- a/pkg/api/http/service/mcp.py
+++ b/pkg/api/http/service/mcp.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import sqlalchemy
 import uuid
 import traceback
+import asyncio

 from ....core import app
 from ....entity.persistence import mcp as persistence_mcp
@@ -124,7 +125,6 @@ class MCPService:

    async def create_mcp_server(self, server_data: dict) -> str:
        server_data['uuid'] = str(uuid.uuid4())
-        print('server_data:', server_data)
        await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_mcp.MCPServer).values(server_data))

        result = await self.ap.persistence_mgr.execute_async(
@@ -134,7 +134,8 @@ class MCPService:
        if server_entity:
            server_config = self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server_entity)
            if self.ap.tool_mgr.mcp_tool_loader:
-                await self.ap.tool_mgr.mcp_tool_loader.load_mcp_server(server_config)
+                task = asyncio.create_task(self.ap.tool_mgr.mcp_tool_loader.host_mcp_server(server_config))
+                self.ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks.append(task)

        return server_data['uuid']

@@ -175,7 +176,9 @@ class MCPService:
            if updated_server:
                # convert entity to config dict
                server_config = self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, updated_server)
-                await self.ap.tool_mgr.mcp_tool_loader.load_mcp_server(server_config)
+                # await self.ap.tool_mgr.mcp_tool_loader.host_mcp_server(server_config)
+                task = asyncio.create_task(self.ap.tool_mgr.mcp_tool_loader.host_mcp_server(server_config))
+                self.ap.tool_mgr.mcp_tool_loader._hosted_mcp_tasks.append(task)

    async def delete_mcp_server(self, server_uuid: str) -> None:
        result = await self.ap.persistence_mgr.execute_async(
--- a/pkg/provider/tools/loaders/mcp.py
+++ b/pkg/provider/tools/loaders/mcp.py
@@ -1,5 +1,6 @@
 from __future__ import annotations

+import enum
 import typing
 from contextlib import AsyncExitStack
 import traceback
@@ -16,6 +17,12 @@ import langbot_plugin.api.entities.builtin.resource.tool as resource_tool
 from ....entity.persistence import mcp as persistence_mcp


+class MCPSessionStatus(enum.Enum):
+    CONNECTING = 'connecting'
+    CONNECTED = 'connected'
+    ERROR = 'error'
+
+
 class RuntimeMCPSession:
    """运行时 MCP 会话"""

@@ -33,7 +40,8 @@ class RuntimeMCPSession:

    enable: bool

-    connected: bool
+    # connected: bool
+    status: MCPSessionStatus

    last_test_error_message: str

@@ -47,7 +55,7 @@ class RuntimeMCPSession:
        self.exit_stack = AsyncExitStack()
        self.functions = []

-        self.connected = False
+        self.status = MCPSessionStatus.CONNECTING
        self.last_test_error_message = ''

    async def _init_stdio_python_server(self):
@@ -117,10 +125,10 @@ class RuntimeMCPSession:
                    )
                )

-            self.connected = True
+            self.status = MCPSessionStatus.CONNECTED
            self.last_test_error_message = ''
        except Exception as e:
-            self.connected = False
+            self.status = MCPSessionStatus.ERROR
            self.last_test_error_message = str(e)
            raise e

@@ -129,7 +137,7 @@ class RuntimeMCPSession:

    def get_runtime_info_dict(self) -> dict:
        return {
-            'connected': self.connected,
+            'status': self.status.value,
            'error_message': self.last_test_error_message,
            'tool_count': len(self.get_tools()),
            'tools': [
@@ -163,13 +171,13 @@ class MCPLoader(loader.ToolLoader):

    _last_listed_functions: list[resource_tool.LLMTool]

-    _startup_load_tasks: list[asyncio.Task]
+    _hosted_mcp_tasks: list[asyncio.Task]

    def __init__(self, ap: app.Application):
        super().__init__(ap)
        self.sessions = {}
        self._last_listed_functions = []
-        self._startup_load_tasks = []
+        self._hosted_mcp_tasks = []

    async def initialize(self):
        await self.load_mcp_servers_from_db()
@@ -185,30 +193,30 @@ class MCPLoader(loader.ToolLoader):
        for server in servers:
            config = self.ap.persistence_mgr.serialize_model(persistence_mcp.MCPServer, server)

-            async def load_mcp_server_task(server_config: dict):
-                self.ap.logger.debug(f'Loading MCP server {server_config}')
-                try:
-                    session = await self.load_mcp_server(server_config)
-                    self.sessions[server_config['name']] = session
-                except Exception as e:
-                    self.ap.logger.error(
-                        f'Failed to load MCP server from db: {server_config["name"]}({server_config["uuid"]}): {e}\n{traceback.format_exc()}'
-                    )
-                    return
+            task = asyncio.create_task(self.host_mcp_server(config))
+            self._hosted_mcp_tasks.append(task)

-                self.ap.logger.debug(f'Starting MCP server {server_config["name"]}({server_config["uuid"]})')
-                try:
-                    await session.start()
-                except Exception as e:
-                    self.ap.logger.error(
-                        f'Failed to start MCP server {server_config["name"]}({server_config["uuid"]}): {e}\n{traceback.format_exc()}'
-                    )
-                    return
+    async def host_mcp_server(self, server_config: dict):
+        self.ap.logger.debug(f'Loading MCP server {server_config}')
+        try:
+            session = await self.load_mcp_server(server_config)
+            self.sessions[server_config['name']] = session
+        except Exception as e:
+            self.ap.logger.error(
+                f'Failed to load MCP server from db: {server_config["name"]}({server_config["uuid"]}): {e}\n{traceback.format_exc()}'
+            )
+            return

-                self.ap.logger.debug(f'Started MCP server {server_config["name"]}({server_config["uuid"]})')
+        self.ap.logger.debug(f'Starting MCP server {server_config["name"]}({server_config["uuid"]})')
+        try:
+            await session.start()
+        except Exception as e:
+            self.ap.logger.error(
+                f'Failed to start MCP server {server_config["name"]}({server_config["uuid"]}): {e}\n{traceback.format_exc()}'
+            )
+            return

-            task = asyncio.create_task(load_mcp_server_task(config))
-            self._startup_load_tasks.append(task)
+        self.ap.logger.debug(f'Started MCP server {server_config["name"]}({server_config["uuid"]})')

    async def load_mcp_server(self, server_config: dict) -> RuntimeMCPSession:
        """加载 MCP 服务器到运行时
@@ -271,20 +279,6 @@ class MCPLoader(loader.ToolLoader):

        raise ValueError(f'Tool not found: {name}')

-    async def reload_mcp_server(self, server_config: dict):
-        """重新加载 MCP 服务器（先移除再加载）
-
-        Args:
-            server_config: 服务器配置字典，必须包含 name 字段
-        """
-        server_name = server_config['name']
-
-        if server_name in self.sessions:
-            await self.remove_mcp_server(server_name)
-
-        # 重新加载
-        await self.load_mcp_server(server_config)
-
    async def remove_mcp_server(self, server_name: str):
        """移除 MCP 服务器"""
        if server_name not in self.sessions:
--- a/web/src/app/home/plugins/mcp-server/MCPCardVO.ts
+++ b/web/src/app/home/plugins/mcp-server/MCPCardVO.ts
@@ -1,10 +1,10 @@
-import { MCPServer } from '@/app/infra/entities/api';
+import { MCPServer, MCPSessionStatus } from '@/app/infra/entities/api';

 export class MCPCardVO {
  name: string;
  mode: 'stdio' | 'sse';
  enable: boolean;
-  status: 'connected' | 'disconnected' | 'error' | 'disabled';
+  status: MCPSessionStatus;
  tools: number;
  error?: string;

@@ -15,45 +15,15 @@ export class MCPCardVO {

    // Determine status from runtime_info
    if (!data.runtime_info) {
-      this.status = 'disconnected';
+      this.status = MCPSessionStatus.ERROR;
      this.tools = 0;
-    } else if (data.runtime_info.connected) {
-      this.status = 'connected';
+    } else if (data.runtime_info.status === MCPSessionStatus.CONNECTED) {
+      this.status = data.runtime_info.status;
      this.tools = data.runtime_info.tool_count || 0;
    } else {
-      this.status = 'error';
+      this.status = data.runtime_info.status;
      this.tools = 0;
      this.error = data.runtime_info.error_message;
    }
  }
-
-  getStatusColor(): string {
-    switch (this.status) {
-      case 'connected':
-        return 'text-green-600';
-      case 'disconnected':
-        return 'text-gray-500';
-      case 'error':
-        return 'text-red-600';
-      case 'disabled':
-        return 'text-gray-400';
-      default:
-        return 'text-gray-500';
-    }
-  }
-
-  getStatusIcon(): string {
-    switch (this.status) {
-      case 'connected':
-        return 'M9 12l2 2 4-4m6 2a9 9 0 11-18 0 9 9 0 0118 0z';
-      case 'disconnected':
-        return 'M10 14l2-2m0 0l2-2m-2 2l-2-2m2 2l2 2m7-2a9 9 0 11-18 0 9 9 0 0118 0z';
-      case 'error':
-        return 'M12 8v4m0 4h.01M21 12a9 9 0 11-18 0 9 9 0 0118 0z';
-      case 'disabled':
-        return 'M18.364 18.364A9 9 0 005.636 5.636m12.728 12.728A9 9 0 015.636 5.636m12.728 12.728L5.636 5.636';
-      default:
-        return 'M10 14l2-2m0 0l2-2m-2 2l-2-2m2 2l2 2m7-2a9 9 0 11-18 0 9 9 0 0118 0z';
-    }
-  }
 }
--- a/web/src/app/home/plugins/mcp-server/MCPServerComponent.tsx
+++ b/web/src/app/home/plugins/mcp-server/MCPServerComponent.tsx
@@ -1,9 +1,10 @@
 'use client';

-import { useEffect, useState } from 'react';
+import { useEffect, useState, useRef } from 'react';
 import MCPCardComponent from '@/app/home/plugins/mcp-server/mcp-card/MCPCardComponent';
 import { MCPCardVO } from '@/app/home/plugins/mcp-server/MCPCardVO';
 import { useTranslation } from 'react-i18next';
+import { MCPSessionStatus } from '@/app/infra/entities/api';

 import { httpClient } from '@/app/infra/http/HttpClient';

@@ -16,11 +17,44 @@ export default function MCPComponent({
  const { t } = useTranslation();
  const [installedServers, setInstalledServers] = useState<MCPCardVO[]>([]);
  const [loading, setLoading] = useState(false);
+  const pollingIntervalRef = useRef<NodeJS.Timeout | null>(null);

  useEffect(() => {
    fetchInstalledServers();
+
+    return () => {
+      // Cleanup: clear polling interval when component unmounts
+      if (pollingIntervalRef.current) {
+        clearInterval(pollingIntervalRef.current);
+      }
+    };
  }, []);

+  // Check if any server is connecting and start/stop polling accordingly
+  useEffect(() => {
+    const hasConnecting = installedServers.some(
+      (server) => server.status === MCPSessionStatus.CONNECTING,
+    );
+
+    if (hasConnecting && !pollingIntervalRef.current) {
+      // Start polling every 3 seconds
+      pollingIntervalRef.current = setInterval(() => {
+        fetchInstalledServers();
+      }, 3000);
+    } else if (!hasConnecting && pollingIntervalRef.current) {
+      // Stop polling when no server is connecting
+      clearInterval(pollingIntervalRef.current);
+      pollingIntervalRef.current = null;
+    }
+
+    return () => {
+      if (pollingIntervalRef.current) {
+        clearInterval(pollingIntervalRef.current);
+        pollingIntervalRef.current = null;
+      }
+    };
+  }, [installedServers]);
+
  function fetchInstalledServers() {
    setLoading(true);
    httpClient
--- a/web/src/app/home/plugins/mcp-server/mcp-card/MCPCardComponent.tsx
+++ b/web/src/app/home/plugins/mcp-server/mcp-card/MCPCardComponent.tsx
@@ -5,7 +5,8 @@ import { Switch } from '@/components/ui/switch';
 import { Button } from '@/components/ui/button';
 import { toast } from 'sonner';
 import { useTranslation } from 'react-i18next';
-import { RefreshCcw, Wrench, Ban, AlertCircle } from 'lucide-react';
+import { RefreshCcw, Wrench, Ban, AlertCircle, Loader2 } from 'lucide-react';
+import { MCPSessionStatus } from '@/app/infra/entities/api';

 export default function MCPCardComponent({
  cardVO,
@@ -110,7 +111,7 @@ export default function MCPCardComponent({
                  {t('mcp.statusDisabled')}
                </div>
              </div>
-            ) : status === 'connected' ? (
+            ) : status === MCPSessionStatus.CONNECTED ? (
              // 连接成功 - 显示工具数量
              <div className="flex h-full flex-row items-center justify-center gap-[0.4rem]">
                <Wrench className="w-5 h-5" />
@@ -118,6 +119,14 @@ export default function MCPCardComponent({
                  {t('mcp.toolCount', { count: toolsCount })}
                </div>
              </div>
+            ) : status === MCPSessionStatus.CONNECTING ? (
+              // 连接中 - 蓝色加载
+              <div className="flex flex-row items-center gap-[0.4rem]">
+                <Loader2 className="w-4 h-4 text-blue-500 dark:text-blue-400 animate-spin" />
+                <div className="text-sm text-blue-500 dark:text-blue-400 font-medium">
+                  {t('mcp.connecting')}
+                </div>
+              </div>
            ) : (
              // 连接失败 - 红色
              <div className="flex flex-row items-center gap-[0.4rem]">
--- a/web/src/app/home/plugins/mcp-server/mcp-form/MCPFormDialog.tsx
+++ b/web/src/app/home/plugins/mcp-server/mcp-form/MCPFormDialog.tsx
@@ -1,6 +1,6 @@
 'use client';

-import React, { useState, useEffect } from 'react';
+import React, { useState, useEffect, useRef } from 'react';
 import { useTranslation } from 'react-i18next';
 import { Resolver, useForm } from 'react-hook-form';
 import { zodResolver } from '@hookform/resolvers/zod';
@@ -42,9 +42,10 @@ import {
  MCPServerRuntimeInfo,
  MCPTool,
  MCPServer,
+  MCPSessionStatus,
 } from '@/app/infra/entities/api';

-// Status Display Component - 只在测试中或连接失败时使用
+// Status Display Component - 在测试中、连接中或连接失败时使用
 function StatusDisplay({
  testing,
  runtimeInfo,
@@ -82,6 +83,35 @@ function StatusDisplay({
    );
  }

+  // 连接中
+  if (runtimeInfo.status === MCPSessionStatus.CONNECTING) {
+    return (
+      <div className="flex items-center gap-2 text-blue-600">
+        <svg
+          className="w-5 h-5 animate-spin"
+          xmlns="http://www.w3.org/2000/svg"
+          fill="none"
+          viewBox="0 0 24 24"
+        >
+          <circle
+            className="opacity-25"
+            cx="12"
+            cy="12"
+            r="10"
+            stroke="currentColor"
+            strokeWidth="4"
+          />
+          <path
+            className="opacity-75"
+            fill="currentColor"
+            d="M4 12a8 8 0 018-8V0C5.373 0 0 5.373 0 12h4zm2 5.291A7.962 7.962 0 014 12H0c0 3.042 1.135 5.824 3 7.938l3-2.647z"
+          />
+        </svg>
+        <span className="font-medium">{t('mcp.connecting')}</span>
+      </div>
+    );
+  }
+
  // 连接失败
  return (
    <div className="space-y-1">
@@ -201,6 +231,7 @@ export default function MCPFormDialog({
  const [runtimeInfo, setRuntimeInfo] = useState<MCPServerRuntimeInfo | null>(
    null,
  );
+  const pollingIntervalRef = useRef<NodeJS.Timeout | null>(null);

  // Load server data when editing
  useEffect(() => {
@@ -212,8 +243,48 @@ export default function MCPFormDialog({
      setExtraArgs([]);
      setRuntimeInfo(null);
    }
+
+    // Cleanup polling interval when dialog closes
+    return () => {
+      if (pollingIntervalRef.current) {
+        clearInterval(pollingIntervalRef.current);
+        pollingIntervalRef.current = null;
+      }
+    };
  }, [open, isEditMode, serverName]);

+  // Poll for updates when runtime_info status is CONNECTING
+  useEffect(() => {
+    if (
+      !open ||
+      !isEditMode ||
+      !serverName ||
+      !runtimeInfo ||
+      runtimeInfo.status !== MCPSessionStatus.CONNECTING
+    ) {
+      // Stop polling if conditions are not met
+      if (pollingIntervalRef.current) {
+        clearInterval(pollingIntervalRef.current);
+        pollingIntervalRef.current = null;
+      }
+      return;
+    }
+
+    // Start polling if not already running
+    if (!pollingIntervalRef.current) {
+      pollingIntervalRef.current = setInterval(() => {
+        loadServerForEdit(serverName);
+      }, 3000);
+    }
+
+    return () => {
+      if (pollingIntervalRef.current) {
+        clearInterval(pollingIntervalRef.current);
+        pollingIntervalRef.current = null;
+      }
+    };
+  }, [open, isEditMode, serverName, runtimeInfo?.status]);
+
  async function loadServerForEdit(serverName: string) {
    try {
      const resp = await httpClient.getMCPServer(serverName);
@@ -312,7 +383,7 @@ export default function MCPFormDialog({
                taskResp.runtime.exception || t('mcp.unknownError');
              toast.error(`${t('mcp.testError')}: ${errorMsg}`);
              setRuntimeInfo({
-                connected: false,
+                status: MCPSessionStatus.ERROR,
                error_message: errorMsg,
                tool_count: 0,
                tools: [],
@@ -387,7 +458,8 @@ export default function MCPFormDialog({
        {isEditMode && runtimeInfo && (
          <div className="mb-4 space-y-3">
            {/* 测试中或连接失败时显示状态 */}
-            {(mcpTesting || !runtimeInfo.connected) && (
+            {(mcpTesting ||
+              runtimeInfo.status !== MCPSessionStatus.CONNECTED) && (
              <div className="p-3 rounded-lg border">
                <StatusDisplay
                  testing={mcpTesting}
@@ -399,7 +471,7 @@ export default function MCPFormDialog({

            {/* 连接成功时只显示工具列表 */}
            {!mcpTesting &&
-              runtimeInfo.connected &&
+              runtimeInfo.status === MCPSessionStatus.CONNECTED &&
              runtimeInfo.tools?.length > 0 && (
                <ToolsList tools={runtimeInfo.tools} />
              )}
--- a/web/src/app/infra/entities/api/index.ts
+++ b/web/src/app/infra/entities/api/index.ts
@@ -325,8 +325,14 @@ export interface MCPServerExtraArgsSSE {
  ssereadtimeout: number;
 }

+export enum MCPSessionStatus {
+  CONNECTING = 'connecting',
+  CONNECTED = 'connected',
+  ERROR = 'error',
+}
+
 export interface MCPServerRuntimeInfo {
-  connected: boolean;
+  status: MCPSessionStatus;
  error_message: string;
  tool_count: number;
  tools: MCPTool[];
--- a/web/src/i18n/locales/en-US.ts
+++ b/web/src/i18n/locales/en-US.ts
@@ -311,6 +311,7 @@ const enUS = {
    keyName: 'Key Name',
    value: 'Value',
    testing: 'Testing...',
+    connecting: 'Connecting...',
    testSuccess: 'Connection test successful',
    testFailed: 'Connection test failed: ',
    testError: 'Connection test error',
--- a/web/src/i18n/locales/ja-JP.ts
+++ b/web/src/i18n/locales/ja-JP.ts
@@ -313,6 +313,7 @@ const jaJP = {
    keyName: 'キー名',
    value: '値',
    testing: 'テスト中...',
+    connecting: '接続中...',
    testSuccess: '接続テストに成功しました',
    testFailed: '接続テストに失敗しました：',
    testError: '接続テストエラー',
--- a/web/src/i18n/locales/zh-Hans.ts
+++ b/web/src/i18n/locales/zh-Hans.ts
@@ -299,6 +299,7 @@ const zhHans = {
    keyName: '键名',
    value: '值',
    testing: '测试中...',
+    connecting: '连接中...',
    testSuccess: '连接测试成功',
    testFailed: '连接测试失败：',
    testError: '连接测试出错',
--- a/web/src/i18n/locales/zh-Hant.ts
+++ b/web/src/i18n/locales/zh-Hant.ts
@@ -297,6 +297,7 @@ const zhHant = {
    keyName: '鍵名',
    value: '值',
    testing: '測試中...',
+    connecting: '連接中...',
    testSuccess: '連接測試成功',
    testFailed: '連接測試失敗：',
    testError: '連接測試出錯',