From 6405228f3fd6a26190e0f747f95299ee429c73be Mon Sep 17 00:00:00 2001
From: Novice <novice12185727@gmail.com>
Date: Wed, 15 Oct 2025 11:45:04 +0800
Subject: [PATCH] feat: improve the refresh token

---
 api/controllers/console/workspace/tool_providers.py | 9 +++++++--
 api/core/mcp/auth/auth_flow.py                      | 3 ++-
 api/core/mcp/error.py                               | 4 ++++
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/api/controllers/console/workspace/tool_providers.py b/api/controllers/console/workspace/tool_providers.py
index 0d8d943c81..0a7fa35a5d 100644
--- a/api/controllers/console/workspace/tool_providers.py
+++ b/api/controllers/console/workspace/tool_providers.py
@@ -18,7 +18,7 @@ from controllers.console.wraps import (
 )
 from core.entities.mcp_provider import MCPAuthentication, MCPConfiguration
 from core.mcp.auth.auth_flow import auth, handle_callback
-from core.mcp.error import MCPAuthError, MCPError
+from core.mcp.error import MCPAuthError, MCPError, MCPRefreshTokenError
 from core.mcp.mcp_client import MCPClient
 from core.model_runtime.utils.encoders import jsonable_encoder
 from core.plugin.impl.oauth import OAuthHandler
@@ -1007,7 +1007,12 @@ class ToolMCPAuthApi(Resource):
                     return {"result": "success"}
             except MCPAuthError as e:
                 service = MCPToolManageService(session=session)
-                return auth(provider_entity, service, args.get("authorization_code"))
+                try:
+                    return auth(provider_entity, service, args.get("authorization_code"))
+                except MCPRefreshTokenError as e:
+                    with session.begin():
+                        service.clear_provider_credentials(provider=db_provider)
+                    raise ValueError(f"Failed to refresh token: {e}") from e
             except MCPError as e:
                 with session.begin():
                     service.clear_provider_credentials(provider=db_provider)
diff --git a/api/core/mcp/auth/auth_flow.py b/api/core/mcp/auth/auth_flow.py
index 4ebf97c7f2..22b0c09ad3 100644
--- a/api/core/mcp/auth/auth_flow.py
+++ b/api/core/mcp/auth/auth_flow.py
@@ -12,6 +12,7 @@ from pydantic import BaseModel, ValidationError
 
 from core.entities.mcp_provider import MCPProviderEntity, MCPSupportGrantType
 from core.helper import ssrf_proxy
+from core.mcp.error import MCPRefreshTokenError
 from core.mcp.types import (
     LATEST_PROTOCOL_VERSION,
     OAuthClientInformation,
@@ -286,7 +287,7 @@ def refresh_authorization(
 
     response = ssrf_proxy.post(token_url, data=params)
     if not response.is_success:
-        raise ValueError(f"Token refresh failed: HTTP {response.status_code}")
+        raise MCPRefreshTokenError(response.text)
     return OAuthTokens.model_validate(response.json())
 
 
diff --git a/api/core/mcp/error.py b/api/core/mcp/error.py
index 92ea7bde09..d4fb8b7674 100644
--- a/api/core/mcp/error.py
+++ b/api/core/mcp/error.py
@@ -8,3 +8,7 @@ class MCPConnectionError(MCPError):
 
 class MCPAuthError(MCPConnectionError):
     pass
+
+
+class MCPRefreshTokenError(MCPError):
+    pass