From 7b3bb1d9606bcaf4800c5e72d274729e931c36fa Mon Sep 17 00:00:00 2001
From: Novice <novice12185727@gmail.com>
Date: Wed, 9 Jul 2025 14:44:51 +0800
Subject: [PATCH] fix: auth api error handle

---
 .../console/workspace/tool_providers.py       | 10 ++++++++--
 api/services/tools/mcp_tools_mange_service.py | 20 ++++++++++++-------
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/api/controllers/console/workspace/tool_providers.py b/api/controllers/console/workspace/tool_providers.py
index 6ac3c4b20b..df50871a38 100644
--- a/api/controllers/console/workspace/tool_providers.py
+++ b/api/controllers/console/workspace/tool_providers.py
@@ -12,7 +12,7 @@ from controllers.console import api
 from controllers.console.wraps import account_initialization_required, enterprise_license_required, setup_required
 from core.mcp.auth.auth_flow import auth, handle_callback
 from core.mcp.auth.auth_provider import OAuthClientProvider
-from core.mcp.error import MCPAuthError
+from core.mcp.error import MCPAuthError, MCPError
 from core.mcp.mcp_client import MCPClient
 from core.model_runtime.utils.encoders import jsonable_encoder
 from extensions.ext_database import db
@@ -733,8 +733,14 @@ class ToolMCPAuthApi(Resource):
 
         except MCPAuthError:
             auth_provider = OAuthClientProvider(provider_id, tenant_id, for_list=True)
-
             return auth(auth_provider, provider.decrypted_server_url, args["authorization_code"])
+        except MCPError as e:
+            MCPToolManageService.update_mcp_provider_credentials(
+                mcp_provider=provider,
+                credentials={},
+                authed=False,
+            )
+            raise ValueError(f"Failed to connect to MCP server: {e}") from e
 
 
 class ToolMCPDetailApi(Resource):
diff --git a/api/services/tools/mcp_tools_mange_service.py b/api/services/tools/mcp_tools_mange_service.py
index b2a88738f6..3b1592230a 100644
--- a/api/services/tools/mcp_tools_mange_service.py
+++ b/api/services/tools/mcp_tools_mange_service.py
@@ -1,12 +1,13 @@
 import hashlib
 import json
 from datetime import datetime
+from typing import Any
 
 from sqlalchemy import or_
 from sqlalchemy.exc import IntegrityError
 
 from core.helper import encrypter
-from core.mcp.error import MCPAuthError, MCPConnectionError
+from core.mcp.error import MCPAuthError, MCPError
 from core.mcp.mcp_client import MCPClient
 from core.tools.entities.api_entities import ToolProviderApiEntity
 from core.tools.entities.common_entities import I18nObject
@@ -119,7 +120,7 @@ class MCPToolManageService:
                 tools = mcp_client.list_tools()
         except MCPAuthError as e:
             raise ValueError("Please auth the tool first")
-        except MCPConnectionError as e:
+        except MCPError as e:
             raise ValueError(f"Failed to connect to MCP server: {e}")
         mcp_provider.tools = json.dumps([tool.model_dump() for tool in tools])
         mcp_provider.authed = True
@@ -173,7 +174,7 @@ class MCPToolManageService:
             server_url_hash = hashlib.sha256(server_url.encode()).hexdigest()
 
             if server_url_hash != mcp_provider.server_url_hash:
-                cls._re_auth_mcp_provider(mcp_provider, provider_id, tenant_id)
+                cls._re_connect_mcp_provider(mcp_provider, provider_id, tenant_id)
                 mcp_provider.server_url_hash = server_url_hash
         try:
             db.session.commit()
@@ -190,7 +191,9 @@ class MCPToolManageService:
                 raise
 
     @classmethod
-    def update_mcp_provider_credentials(cls, mcp_provider: MCPToolProvider, credentials: dict, authed: bool = False):
+    def update_mcp_provider_credentials(
+        cls, mcp_provider: MCPToolProvider, credentials: dict[str, Any], authed: bool = False
+    ):
         provider_controller = MCPToolProviderController._from_db(mcp_provider)
         tool_configuration = ProviderConfigEncrypter(
             tenant_id=mcp_provider.tenant_id,
@@ -202,11 +205,13 @@ class MCPToolManageService:
         mcp_provider.updated_at = datetime.now()
         mcp_provider.encrypted_credentials = json.dumps({**mcp_provider.credentials, **credentials})
         mcp_provider.authed = authed
+        if not authed:
+            mcp_provider.tools = "[]"
         db.session.commit()
 
     @classmethod
-    def _re_auth_mcp_provider(cls, mcp_provider: MCPToolProvider, provider_id: str, tenant_id: str):
-        """re-auth mcp provider"""
+    def _re_connect_mcp_provider(cls, mcp_provider: MCPToolProvider, provider_id: str, tenant_id: str):
+        """re-connect mcp provider"""
         try:
             with MCPClient(
                 mcp_provider.decrypted_server_url,
@@ -221,6 +226,7 @@ class MCPToolManageService:
         except MCPAuthError:
             mcp_provider.authed = False
             mcp_provider.tools = "[]"
-
+        except MCPError as e:
+            raise ValueError(f"Failed to re-connect MCP server: {e}") from e
         # reset credentials
         mcp_provider.encrypted_credentials = "{}"