refactor(mcp): clean the client service code

api/core/mcp/auth_client.py

@@ -0,0 +1,204 @@
+"""
+MCP Client with Authentication Retry Support
+
+This module provides a wrapper around MCPClient that automatically handles
+authentication failures and retries operations after refreshing tokens.
+"""
+
+import logging
+from collections.abc import Callable
+from types import TracebackType
+from typing import Any, Optional
+
+from sqlalchemy.orm import Session
+
+from core.entities.mcp_provider import MCPProviderEntity
+from core.mcp.error import MCPAuthError
+from core.mcp.mcp_client import MCPClient
+from core.mcp.types import CallToolResult, Tool
+from extensions.ext_database import db
+
+logger = logging.getLogger(__name__)
+
+
+class MCPClientWithAuthRetry:
+    """
+    A wrapper around MCPClient that provides automatic authentication retry.
+
+    This class intercepts MCPAuthError exceptions and attempts to refresh
+    authentication before retrying the failed operation.
+    """
+
+    def __init__(
+        self,
+        server_url: str,
+        headers: dict[str, str] | None = None,
+        timeout: float | None = None,
+        sse_read_timeout: float | None = None,
+        provider_entity: MCPProviderEntity | None = None,
+        auth_callback: Callable[[MCPProviderEntity, Optional[str]], dict[str, str]] | None = None,
+        authorization_code: Optional[str] = None,
+        by_server_id: bool = False,
+    ):
+        """
+        Initialize the MCP client with auth retry capability.
+
+        Args:
+            server_url: The MCP server URL
+            headers: Optional headers for requests
+            timeout: Request timeout
+            sse_read_timeout: SSE read timeout
+            provider_entity: Provider entity for authentication
+            auth_callback: Authentication callback function
+            authorization_code: Optional authorization code for initial auth
+        """
+        self.server_url = server_url
+        self.headers = headers or {}
+        self.timeout = timeout
+        self.sse_read_timeout = sse_read_timeout
+        self.provider_entity = provider_entity
+        self.auth_callback = auth_callback
+        self.authorization_code = authorization_code
+        self._has_retried = False
+        self._client: MCPClient | None = None
+        self.by_server_id = by_server_id
+
+    def _create_client(self) -> MCPClient:
+        """Create a new MCPClient instance with current headers."""
+        return MCPClient(
+            server_url=self.server_url,
+            headers=self.headers,
+            timeout=self.timeout,
+            sse_read_timeout=self.sse_read_timeout,
+        )
+
+    def _handle_auth_error(self, error: MCPAuthError) -> None:
+        """
+        Handle authentication error by refreshing tokens.
+
+        Args:
+            error: The authentication error
+
+        Raises:
+            MCPAuthError: If authentication fails or max retries reached
+        """
+        from services.tools.mcp_oauth_service import MCPOAuthService
+
+        if not self.provider_entity or not self.auth_callback:
+            raise error
+
+        if self._has_retried:
+            raise error
+
+        self._has_retried = True
+
+        try:
+            # Perform authentication
+            self.auth_callback(self.provider_entity, self.authorization_code)
+
+            # Retrieve new tokens
+            with Session(db.engine) as session:
+                oauth_service = MCPOAuthService(session=session)
+                self.provider_entity = oauth_service.get_provider_entity(
+                    self.provider_entity.id, self.provider_entity.tenant_id, by_server_id=self.by_server_id
+                )
+            token = self.provider_entity.retrieve_tokens()
+            if not token:
+                raise MCPAuthError("Authentication failed - no token received")
+
+            # Update headers with new token
+            self.headers["Authorization"] = f"{token.token_type.capitalize()} {token.access_token}"
+
+            # Clear authorization code after first use
+            self.authorization_code = None
+
+        except Exception as e:
+            logger.exception("Authentication retry failed")
+            raise MCPAuthError(f"Authentication retry failed: {e}") from e
+
+    def _execute_with_retry(self, func: Callable[..., Any], *args, **kwargs) -> Any:
+        """
+        Execute a function with authentication retry logic.
+
+        Args:
+            func: The function to execute
+            *args: Positional arguments for the function
+            **kwargs: Keyword arguments for the function
+
+        Returns:
+            The result of the function call
+
+        Raises:
+            MCPAuthError: If authentication fails after retries
+            Any other exceptions from the function
+        """
+        try:
+            return func(*args, **kwargs)
+        except MCPAuthError as e:
+            self._handle_auth_error(e)
+            # Recreate client with new headers
+            if self._client:
+                self._client.cleanup()
+                self._client = self._create_client()
+                self._client.__enter__()
+            return func(*args, **kwargs)
+        finally:
+            # Reset retry flag after operation completes
+            self._has_retried = False
+
+    def __enter__(self):
+        """Enter the context manager."""
+        self._client = self._create_client()
+
+        # Try to initialize with retry
+        def initialize():
+            if self._client is None:
+                raise ValueError("Client not created")
+            self._client.__enter__()
+            return self
+
+        return self._execute_with_retry(initialize)
+
+    def __exit__(self, exc_type: type | None, exc_value: BaseException | None, traceback: TracebackType | None):
+        """Exit the context manager."""
+        if self._client:
+            self._client.__exit__(exc_type, exc_value, traceback)
+            self._client = None
+
+    def list_tools(self) -> list[Tool]:
+        """
+        List available tools from the MCP server.
+
+        Returns:
+            List of available tools
+
+        Raises:
+            MCPAuthError: If authentication fails after retries
+        """
+        if not self._client:
+            raise ValueError("Client not initialized. Use within a context manager.")
+        return self._execute_with_retry(self._client.list_tools)
+
+    def invoke_tool(self, tool_name: str, tool_args: dict[str, Any]) -> CallToolResult:
+        """
+        Invoke a tool on the MCP server.
+
+        Args:
+            tool_name: Name of the tool to invoke
+            tool_args: Arguments for the tool
+
+        Returns:
+            Result of the tool invocation
+
+        Raises:
+            MCPAuthError: If authentication fails after retries
+        """
+        if not self._client:
+            raise ValueError("Client not initialized. Use within a context manager.")
+        return self._execute_with_retry(self._client.invoke_tool, tool_name, tool_args)
+
+    def cleanup(self):
+        """Clean up resources."""
+        if self._client:
+            self._client.cleanup()
+            self._client = None