alita-sdk 0.3.257__py3-none-any.whl → 0.3.562__py3-none-any.whl

alita_sdk/runtime/clients/client.py

@@ -13,6 +13,7 @@ from langchain_core.messages import (
 from langchain_core.tools import ToolException
 from langgraph.store.base import BaseStore
 from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from langchain_anthropic import ChatAnthropic
 
 from ..langchain.assistant import Assistant as LangChainAssistant
 # from ..llamaindex.assistant import Assistant as LLamaAssistant
@@ -20,7 +21,9 @@ from .prompt import AlitaPrompt
 from .datasource import AlitaDataSource
 from .artifact import Artifact
 from ..langchain.chat_message_template import Jinja2TemplatedChatMessagesTemplate
-from ...tools import get_available_toolkit_models
+from ..utils.mcp_oauth import McpAuthorizationRequired
+from ...tools import get_available_toolkit_models, instantiate_toolkit
+from ...tools.base_indexer_toolkit import IndexTools
 
 logger = logging.getLogger(__name__)
 
@@ -41,6 +44,7 @@ class AlitaClient:
         self.base_url = base_url.rstrip('/')
         self.api_path = '/api/v1'
         self.llm_path = '/llm/v1'
+        self.allm_path = '/llm'
         self.project_id = project_id
         self.auth_token = auth_token
         self.headers = {
@@ -67,7 +71,15 @@ class AlitaClient:
         self.bucket_url = f"{self.base_url}{self.api_path}/artifacts/buckets/{self.project_id}"
         self.configurations_url = f'{self.base_url}{self.api_path}/integrations/integrations/default/{self.project_id}?section=configurations&unsecret=true'
         self.ai_section_url = f'{self.base_url}{self.api_path}/integrations/integrations/default/{self.project_id}?section=ai'
+        self.models_url = f'{self.base_url}{self.api_path}/configurations/models/{self.project_id}?include_shared=true'
+        self.image_generation_url = f"{self.base_url}{self.llm_path}/images/generations"
         self.configurations: list = configurations or []
+        self.model_timeout = kwargs.get('model_timeout', 120)
+        self.model_image_generation = kwargs.get('model_image_generation')
+        
+        # Cache for generated images to avoid token consumption
+        # This is used by image_generation and artifact toolkits to pass data via reference
+        self._generated_images_cache: Dict[str, Dict[str, Any]] = {}
 
     def get_mcp_toolkits(self):
         if user_id := self._get_real_user_id():
@@ -80,8 +92,23 @@ class AlitaClient:
     def mcp_tool_call(self, params: dict[str, Any]):
         if user_id := self._get_real_user_id():
             url = f"{self.mcp_tools_call}/{user_id}"
-            data = requests.post(url, headers=self.headers, json=params, verify=False).json()
-            return data
+            #
+            # This loop iterates over each key-value pair in the arguments dictionary,
+            # and if a value is a Pydantic object, it replaces it with its dictionary representation using .dict().
+            for arg_name, arg_value in params.get('params', {}).get('arguments', {}).items():
+                if isinstance(arg_value, list):
+                    params['params']['arguments'][arg_name] = [
+                        item.dict() if hasattr(item, "dict") and callable(item.dict) else item
+                        for item in arg_value
+                    ]
+                elif hasattr(arg_value, "dict") and callable(arg_value.dict):
+                    params['params']['arguments'][arg_name] = arg_value.dict()
+            #
+            response = requests.post(url, headers=self.headers, json=params, verify=False)
+            try:
+                return response.json()
+            except (ValueError, TypeError):
+                return response.text
         else:
             return f"Error: Could not determine user ID for MCP tool call"
 
@@ -124,6 +151,19 @@ class AlitaClient:
         data = requests.get(url, headers=self.headers, verify=False).json()
         return data
 
+    def toolkit(self, toolkit_id: int):
+        url = f"{self.base_url}{self.api_path}/tool/prompt_lib/{self.project_id}/{toolkit_id}"
+        response = requests.get(url, headers=self.headers, verify=False)
+        if not response.ok:
+            raise ValueError(f"Failed to fetch toolkit {toolkit_id}: {response.text}")
+        
+        tool_data = response.json()
+        if 'settings' not in tool_data:
+            tool_data['settings'] = {}
+        tool_data['settings']['alita'] = self
+        
+        return instantiate_toolkit(tool_data)
+
     def get_list_of_apps(self):
         apps = []
         limit = 10
@@ -156,6 +196,20 @@ class AlitaClient:
             return resp.json()
         return []
 
+    def get_available_models(self):
+        """Get list of available models from the configurations API.
+
+        Returns:
+            List of model dictionaries with 'name' and other properties,
+            or empty list if request fails.
+        """
+        resp = requests.get(self.models_url, headers=self.headers, verify=False)
+        if resp.ok:
+            data = resp.json()
+            # API returns {"items": [...], ...}
+            return data.get('items', [])
+        return []
+
     def get_embeddings(self, embedding_model: str) -> OpenAIEmbeddings:
         """
         Get an instance of OpenAIEmbeddings configured with the project ID and auth token.
@@ -168,38 +222,160 @@ class AlitaClient:
             model=embedding_model,
             api_key=self.auth_token,
             openai_organization=str(self.project_id),
+            request_timeout=self.model_timeout
         )
 
-    def get_llm(self, model_name: str, model_config: dict) -> ChatOpenAI:
+    def get_llm(self, model_name: str, model_config: dict):
         """
-        Get a ChatOpenAI model instance based on the model name and configuration.
+        Get a ChatOpenAI or ChatAnthropic model instance based on the model name and configuration.
 
         Args:
             model_name: Name of the model to retrieve
             model_config: Configuration parameters for the model
 
         Returns:
-            An instance of ChatOpenAI configured with the provided parameters.
+            An instance of ChatOpenAI or ChatAnthropic configured with the provided parameters.
         """
         if not model_name:
             raise ValueError("Model name must be provided")
 
-        logger.info(f"Creating ChatOpenAI model: {model_name} with config: {model_config}")
+        # Determine if this is an Anthropic model
+        model_name_lower = model_name.lower()
+        is_anthropic = "anthropic" in model_name_lower or "claude" in model_name_lower
 
-        return ChatOpenAI(
-            base_url=f"{self.base_url}{self.llm_path}",
-            model=model_name,
-            api_key=self.auth_token,
-            streaming=model_config.get("streaming", True),
-            stream_usage=model_config.get("stream_usage", True),
-            max_tokens=model_config.get("max_tokens", None),
-            top_p=model_config.get("top_p"),
-            temperature=model_config.get("temperature"),
-            max_retries=model_config.get("max_retries", 3),
-            seed=model_config.get("seed", None),
-            openai_organization=str(self.project_id),
-        )
+        logger.info(f"Creating {'ChatAnthropic' if is_anthropic else 'ChatOpenAI'} model: {model_name} with config: {model_config}")
+
+        try:
+            from tools import this  # pylint: disable=E0401,C0415
+            worker_config = this.for_module("indexer_worker").descriptor.config
+        except:  # pylint: disable=W0702
+            worker_config = {}
+
+        use_responses_api = False
+
+        if worker_config and isinstance(worker_config, dict):
+            for target_name_tag in worker_config.get("use_responses_api_for", []):
+                if target_name_tag in model_name:
+                    use_responses_api = True
+                    break
+
+        # handle case when max_tokens are auto-configurable == -1 or None
+        llm_max_tokens = model_config.get("max_tokens", None)
+        if llm_max_tokens is None or llm_max_tokens == -1:
+            logger.warning(f'User selected `MAX COMPLETION TOKENS` as `auto` or value is None/missing')
+            # default number for a case when auto is selected for an agent
+            llm_max_tokens = 4000
+
+        if is_anthropic:
+            # ChatAnthropic configuration
+            # Anthropic requires max_tokens to be an integer, never None
+            target_kwargs = {
+                "base_url": f"{self.base_url}{self.allm_path}",
+                "model": model_name,
+                "api_key": self.auth_token,
+                "streaming": model_config.get("streaming", True),
+                "max_tokens": llm_max_tokens,  # Always an integer now
+                "temperature": model_config.get("temperature"),
+                "max_retries": model_config.get("max_retries", 3),
+                "default_headers": {"openai-organization": str(self.project_id),
+                                    "Authorization": f"Bearer {self.auth_token}"},
+            }
+            
+            # TODO": Check on ChatAnthropic client when they get "effort" support back
+            if model_config.get("reasoning_effort"):
+                if model_config["reasoning_effort"].lower() == "low":
+                    target_kwargs['thinking'] = {"type": "enabled", "budget_tokens": 2048}
+                    target_kwargs['temperature'] = 1
+                    target_kwargs["max_tokens"] = 2048 + target_kwargs["max_tokens"]
+                elif model_config["reasoning_effort"].lower() == "medium":
+                    target_kwargs['thinking'] = {"type": "enabled", "budget_tokens": 4096}
+                    target_kwargs['temperature'] = 1
+                    target_kwargs["max_tokens"] = 4096 + target_kwargs["max_tokens"]
+                elif model_config["reasoning_effort"].lower() == "high":
+                    target_kwargs['thinking'] = {"type": "enabled", "budget_tokens": 9092}
+                    target_kwargs['temperature'] = 1
+                    target_kwargs["max_tokens"] = 9092 + target_kwargs["max_tokens"]
+                    
+            # Add http_client if provided
+            if "http_client" in model_config:
+                target_kwargs["http_client"] = model_config["http_client"]
+            
+            llm = ChatAnthropic(**target_kwargs)
+        else:
+            # ChatOpenAI configuration
+            target_kwargs = {
+                "base_url": f"{self.base_url}{self.llm_path}",
+                "model": model_name,
+                "api_key": self.auth_token,
+                "streaming": model_config.get("streaming", True),
+                "stream_usage": model_config.get("stream_usage", True),
+                "max_tokens": llm_max_tokens,
+                "temperature": model_config.get("temperature"),
+                "reasoning_effort": model_config.get("reasoning_effort"),
+                "max_retries": model_config.get("max_retries", 3),
+                "seed": model_config.get("seed", None),
+                "openai_organization": str(self.project_id),
+            }
+
+            if use_responses_api:
+                target_kwargs["use_responses_api"] = True
+            
+            llm = ChatOpenAI(**target_kwargs)
+        return llm
+        
+    def generate_image(self,
+                       prompt: str,
+                       n: int = 1,
+                       size: str = "auto",
+                       quality: str = "auto",
+                       response_format: str = "b64_json",
+                       style: Optional[str] = None) -> dict:
+
+        if not self.model_image_generation:
+            raise ValueError("Image generation model is not configured for this client")
+
+        image_generation_data = {
+            "prompt": prompt,
+            "model": self.model_image_generation,
+            "n": n,
+            "response_format": response_format,
+        }
+
+        # Only add optional parameters if they have meaningful values
+        if size and size.lower() != "auto":
+            image_generation_data["size"] = size
+
+        if quality and quality.lower() != "auto":
+            image_generation_data["quality"] = quality
 
+        if style:
+            image_generation_data["style"] = style
+
+        # Standard headers for image generation
+        image_headers = self.headers.copy()
+        image_headers.update({
+            "Content-Type": "application/json",
+        })
+
+        logger.info(f"Generating image with model: {self.model_image_generation}, prompt: {prompt[:50]}...")
+
+        try:
+            response = requests.post(
+                self.image_generation_url,
+                headers=image_headers,
+                json=image_generation_data,
+                verify=False,
+                timeout=self.model_timeout
+            )
+            response.raise_for_status()
+            return response.json()
+
+        except requests.exceptions.HTTPError as e:
+            logger.error(f"Image generation failed: {e.response.status_code} - {e.response.text}")
+            raise
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Image generation request failed: {e}")
+            raise
 
     def get_app_version_details(self, application_id: int, application_version_id: int) -> dict:
         url = f"{self.application_versions}/{application_id}/{application_version_id}"
@@ -231,7 +407,8 @@ class AlitaClient:
                     app_type=None, memory=None, runtime='langchain',
                     application_variables: Optional[dict] = None,
                     version_details: Optional[dict] = None, store: Optional[BaseStore] = None,
-                    llm: Optional[ChatOpenAI] = None):
+                    llm: Optional[ChatOpenAI] = None, mcp_tokens: Optional[dict] = None,
+                    conversation_id: Optional[str] = None, ignored_mcp_servers: Optional[list] = None):
         if tools is None:
             tools = []
         if chat_history is None:
@@ -251,11 +428,15 @@ class AlitaClient:
                 if var['name'] in application_variables:
                     var.update(application_variables[var['name']])
         if llm is None:
+            max_tokens = data['llm_settings'].get('max_tokens', 4000)
+            if max_tokens == -1:
+                # default nuber for case when auto is selected for agent
+                max_tokens = 4000
             llm = self.get_llm(
                 model_name=data['llm_settings']['model_name'],
                 model_config={
-                    "max_tokens": data['llm_settings']['max_tokens'],
-                    "top_p": data['llm_settings']['top_p'],
+                    "max_tokens": max_tokens,
+                    "reasoning_effort": data['llm_settings'].get('reasoning_effort'),
                     "temperature": data['llm_settings']['temperature'],
                     "model_project_id": data['llm_settings'].get('model_project_id'),
                 }
@@ -270,13 +451,18 @@ class AlitaClient:
             app_type = "react"
         elif app_type == 'autogen':
             app_type = "react"
+
+        # LangChainAssistant constructor calls get_tools() which may raise McpAuthorizationRequired
+        # The exception will propagate naturally to the indexer worker's outer handler
         if runtime == 'nonrunnable':
             return LangChainAssistant(self, data, llm, chat_history, app_type,
-                                      tools=tools, memory=memory, store=store)
+                                      tools=tools, memory=memory, store=store, mcp_tokens=mcp_tokens,
+                                      conversation_id=conversation_id, ignored_mcp_servers=ignored_mcp_servers)
         if runtime == 'langchain':
             return LangChainAssistant(self, data, llm,
                                       chat_history, app_type,
-                                      tools=tools, memory=memory, store=store).runnable()
+                                      tools=tools, memory=memory, store=store, mcp_tokens=mcp_tokens,
+                                      conversation_id=conversation_id, ignored_mcp_servers=ignored_mcp_servers).runnable()
         elif runtime == 'llama':
             raise NotImplementedError("LLama runtime is not supported")
 
@@ -338,16 +524,50 @@ class AlitaClient:
         return self._process_requst(resp)
 
     def list_artifacts(self, bucket_name: str):
-        url = f'{self.artifacts_url}/{bucket_name}'
+        # Ensure bucket name is lowercase as required by the API
+        url = f'{self.artifacts_url}/{bucket_name.lower()}'
         data = requests.get(url, headers=self.headers, verify=False)
         return self._process_requst(data)
 
     def create_artifact(self, bucket_name, artifact_name, artifact_data):
+        # Sanitize filename to prevent regex errors during indexing
+        sanitized_name, was_modified = self._sanitize_artifact_name(artifact_name)
+        if was_modified:
+            logger.warning(f"Artifact filename sanitized: '{artifact_name}' -> '{sanitized_name}'")
+        
         url = f'{self.artifacts_url}/{bucket_name.lower()}'
         data = requests.post(url, headers=self.headers, files={
-            'file': (artifact_name, artifact_data)
+            'file': (sanitized_name, artifact_data)
         }, verify=False)
         return self._process_requst(data)
+    
+    @staticmethod
+    def _sanitize_artifact_name(filename: str) -> tuple:
+        """Sanitize filename for safe storage and regex pattern matching."""
+        import re
+        from pathlib import Path
+        
+        if not filename or not filename.strip():
+            return "unnamed_file", True
+        
+        original = filename
+        path_obj = Path(filename)
+        name = path_obj.stem
+        extension = path_obj.suffix
+        
+        # Whitelist: alphanumeric, underscore, hyphen, space, Unicode letters/digits
+        sanitized_name = re.sub(r'[^\w\s-]', '', name, flags=re.UNICODE)
+        sanitized_name = re.sub(r'[-\s]+', '-', sanitized_name)
+        sanitized_name = sanitized_name.strip('-').strip()
+        
+        if not sanitized_name:
+            sanitized_name = "file"
+        
+        if extension:
+            extension = re.sub(r'[^\w.-]', '', extension, flags=re.UNICODE)
+        
+        sanitized = sanitized_name + extension
+        return sanitized, (sanitized != original)
 
     def download_artifact(self, bucket_name, artifact_name):
         url = f'{self.artifact_url}/{bucket_name.lower()}/{artifact_name}'
@@ -489,25 +709,32 @@ class AlitaClient:
             monitoring_meta = tasknode_task.meta.get("monitoring", {})
             return monitoring_meta["user_id"]
         except Exception as e:
-            logger.warning(f"Error: Could not determine user ID for MCP tool: {e}")
+            logger.debug(f"Error: Could not determine user ID for MCP tool: {e}")
             return None
 
     def predict_agent(self, llm: ChatOpenAI, instructions: str = "You are a helpful assistant.",
                       tools: Optional[list] = None, chat_history: Optional[List[Any]] = None,
                       memory=None, runtime='langchain', variables: Optional[list] = None,
-                      store: Optional[BaseStore] = None):
+                      store: Optional[BaseStore] = None, debug_mode: Optional[bool] = False,
+                      mcp_tokens: Optional[dict] = None, conversation_id: Optional[str] = None,
+                      ignored_mcp_servers: Optional[list] = None, persona: Optional[str] = "generic"):
         """
         Create a predict-type agent with minimal configuration.
 
         Args:
             llm: The LLM to use
             instructions: System instructions for the agent
-            tools: Optional list of tools to provide to the agent
+            tools: Optional list of tool configurations (not tool instances) to provide to the agent.
+                   Tool configs will be processed through get_tools() to create tool instances.
+                   Each tool config should have 'type', 'settings', etc.
             chat_history: Optional chat history
             memory: Optional memory/checkpointer
             runtime: Runtime type (default: 'langchain')
             variables: Optional list of variables for the agent
             store: Optional store for memory
+            debug_mode: Enable debug mode for cases when assistant can be initialized without tools
+            ignored_mcp_servers: Optional list of MCP server URLs to ignore (user chose to continue without auth)
+            persona: Default persona for chat: 'generic' or 'qa' (default: 'generic')
 
         Returns:
             Runnable agent ready for execution
@@ -521,17 +748,34 @@ class AlitaClient:
 
         # Create a minimal data structure for predict agent
         # All LLM settings are taken from the passed client instance
+        # Note: 'tools' here are tool CONFIGURATIONS, not tool instances
+        # They will be converted to tool instances by LangChainAssistant via get_tools()
         agent_data = {
             'instructions': instructions,
-            'tools': tools,  # Tools are handled separately in predict agents
+            'tools': tools,  # Tool configs that will be processed by get_tools()
             'variables': variables
         }
-        return LangChainAssistant(self, agent_data, llm,
-                                  chat_history, "predict", memory=memory, store=store).runnable()
+
+        # LangChainAssistant constructor calls get_tools() which may raise McpAuthorizationRequired
+        # The exception will propagate naturally to the indexer worker's outer handler
+        return LangChainAssistant(
+            self,
+            agent_data,
+            llm,
+            chat_history,
+            "predict",
+            memory=memory,
+            store=store,
+            debug_mode=debug_mode,
+            mcp_tokens=mcp_tokens,
+            conversation_id=conversation_id,
+            ignored_mcp_servers=ignored_mcp_servers,
+            persona=persona
+        ).runnable()
 
     def test_toolkit_tool(self, toolkit_config: dict, tool_name: str, tool_params: dict = None,
                           runtime_config: dict = None, llm_model: str = None,
-                          llm_config: dict = None) -> dict:
+                          llm_config: dict = None, mcp_tokens: dict = None) -> dict:
         """
         Test a single tool from a toolkit with given parameters and runtime callbacks.
 
@@ -550,6 +794,7 @@ class AlitaClient:
                 - configurable: Additional configuration parameters
                 - tags: Tags for the execution
             llm_model: Name of the LLM model to use (default: 'gpt-4o-mini')
+            mcp_tokens: Optional dictionary of MCP OAuth tokens by server URL
             llm_config: Configuration for the LLM containing:
                 - max_tokens: Maximum tokens for response (default: 1000)
                 - temperature: Temperature for response generation (default: 0.1)
@@ -597,7 +842,6 @@ class AlitaClient:
             llm_config = {
                 'max_tokens': 1024,
                 'temperature': 0.1,
-                'top_p': 1.0
             }
         import logging
         logger = logging.getLogger(__name__)
@@ -668,7 +912,23 @@ class AlitaClient:
                 }
 
             # Instantiate the toolkit with client and LLM support
-            tools = instantiate_toolkit_with_client(toolkit_config, llm, self)
+            try:
+                tools = instantiate_toolkit_with_client(toolkit_config, llm, self, mcp_tokens=mcp_tokens, use_prefix=False)
+            except McpAuthorizationRequired:
+                # Re-raise McpAuthorizationRequired to allow proper handling upstream
+                logger.info(f"McpAuthorizationRequired detected, re-raising")
+                raise
+            except Exception as toolkit_error:
+                # For other errors, return error response
+                return {
+                    "success": False,
+                    "error": f"Failed to instantiate toolkit '{toolkit_config.get('toolkit_name')}': {str(toolkit_error)}",
+                    "tool_name": tool_name,
+                    "toolkit_config": toolkit_config_parsed_json,
+                    "llm_model": llm_model,
+                    "events_dispatched": events_dispatched,
+                    "execution_time_seconds": 0.0
+                }
 
             if not tools:
                 return {
@@ -744,7 +1004,6 @@ class AlitaClient:
             if target_tool is None:
                 available_tools = []
                 base_available_tools = []
-                full_available_tools = []
 
                 for tool in tools:
                     tool_name_attr = None
@@ -761,16 +1020,14 @@ class AlitaClient:
                         if base_name not in base_available_tools:
                             base_available_tools.append(base_name)
 
-                        # Track full names separately
-                        if '___' in tool_name_attr:
-                            full_available_tools.append(tool_name_attr)
-
                 # Create comprehensive error message
-                error_msg = f"Tool '{tool_name}' not found in toolkit '{toolkit_config.get('toolkit_name')}'."
+                error_msg = f"Tool '{tool_name}' not found in toolkit '{toolkit_config.get('toolkit_name')}'.\n"
 
-                if base_available_tools and full_available_tools:
-                    error_msg += f" Available tools: {base_available_tools} (base names) or {full_available_tools} (full names)"
-                elif base_available_tools:
+                # Custom error for index tools
+                if toolkit_name in [tool.value for tool in IndexTools]:
+                    error_msg += f" Please make sure proper PGVector configuration and embedding model are set in the platform.\n"
+
+                if base_available_tools:
                     error_msg += f" Available tools: {base_available_tools}"
                 elif available_tools:
                     error_msg += f" Available tools: {available_tools}"
@@ -779,10 +1036,7 @@ class AlitaClient:
 
                 # Add helpful hint about naming conventions
                 if '___' in tool_name:
-                    error_msg += f" Note: You provided a full name '{tool_name}'. Try using just the base name '{extract_base_tool_name(tool_name)}'."
-                elif full_available_tools:
-                    possible_full_name = create_full_tool_name(tool_name, toolkit_name)
-                    error_msg += f" Note: You provided a base name '{tool_name}'. The full name might be '{possible_full_name}'."
+                    error_msg += f" Note: Tool names no longer use '___' prefixes. Try using just the base name '{extract_base_tool_name(tool_name)}'."
 
                 return {
                     "success": False,
@@ -888,6 +1142,9 @@ class AlitaClient:
                 }
 
         except Exception as e:
+            # Re-raise McpAuthorizationRequired to allow proper handling upstream
+            if isinstance(e, McpAuthorizationRequired):
+                raise
             logger = logging.getLogger(__name__)
             logger.error(f"Error in test_toolkit_tool: {str(e)}")
             return {
@@ -900,23 +1157,157 @@ class AlitaClient:
                 "execution_time_seconds": 0.0
             }
 
-    def _get_real_user_id(self) -> str:
-        """Extract the real user ID from the auth token for MCP tool calls."""
+    def test_mcp_connection(self, toolkit_config: dict, mcp_tokens: dict = None) -> dict:
+        """
+        Test MCP server connection using protocol-level list_tools.
+
+        This method verifies MCP server connectivity and authentication by calling
+        the protocol-level tools/list JSON-RPC method (NOT executing a tool).
+        This is ideal for auth checks as it validates the connection without
+        requiring any tool execution.
+
+        Args:
+            toolkit_config: Configuration dictionary for the MCP toolkit containing:
+                - toolkit_name: Name of the toolkit
+                - settings: Dictionary with 'url', optional 'headers', 'session_id'
+            mcp_tokens: Optional dictionary of MCP OAuth tokens by server URL
+                Format: {canonical_url: {access_token: str, session_id: str}}
+
+        Returns:
+            Dictionary containing:
+                - success: Boolean indicating if the connection was successful
+                - tools: List of tool names available on the MCP server (if successful)
+                - tools_count: Number of tools discovered
+                - server_session_id: Session ID provided by the server (if any)
+                - error: Error message (if unsuccessful)
+                - toolkit_config: Original toolkit configuration
+
+        Raises:
+            McpAuthorizationRequired: If MCP server requires OAuth authorization
+
+        Example:
+            >>> config = {
+            ...     'toolkit_name': 'my-mcp-server',
+            ...     'type': 'mcp',
+            ...     'settings': {
+            ...         'url': 'https://mcp-server.example.com/mcp',
+            ...         'headers': {'X-Custom': 'value'}
+            ...     }
+            ... }
+            >>> result = client.test_mcp_connection(config)
+            >>> if result['success']:
+            ...     print(f"Connected! Found {result['tools_count']} tools")
+        """
+        import asyncio
+        import time
+        from ..utils.mcp_client import McpClient
+        from ..utils.mcp_oauth import canonical_resource
+
+        toolkit_name = toolkit_config.get('toolkit_name', 'unknown')
+        settings = toolkit_config.get('settings', {})
+
+        # Extract connection parameters
+        url = settings.get('url')
+        if not url:
+            return {
+                "success": False,
+                "error": "MCP toolkit configuration missing 'url' in settings",
+                "toolkit_config": toolkit_config,
+                "tools": [],
+                "tools_count": 0
+            }
+
+        headers = settings.get('headers') or {}
+        session_id = settings.get('session_id')
+
+        # Apply OAuth token if available
+        if mcp_tokens and url:
+            canonical_url = canonical_resource(url)
+            token_data = mcp_tokens.get(canonical_url)
+            if token_data:
+                if isinstance(token_data, dict):
+                    access_token = token_data.get('access_token')
+                    if not session_id:
+                        session_id = token_data.get('session_id')
+                else:
+                    # Backward compatibility: plain token string
+                    access_token = token_data
+
+                if access_token:
+                    headers = dict(headers)  # Copy to avoid mutating original
+                    headers.setdefault('Authorization', f'Bearer {access_token}')
+                    logger.info(f"[MCP Auth Check] Applied OAuth token for {canonical_url}")
+
+        logger.info(f"Testing MCP connection to '{toolkit_name}' at {url}")
+
+        start_time = time.time()
+
+        async def _test_connection():
+            client = McpClient(
+                url=url,
+                session_id=session_id,
+                headers=headers,
+                timeout=60  # Reasonable timeout for connection test
+            )
+
+            async with client:
+                # Initialize MCP protocol session
+                await client.initialize()
+                logger.info(f"[MCP Auth Check] Session initialized (transport={client.detected_transport})")
+
+                # Call protocol-level list_tools (tools/list JSON-RPC method)
+                tools = await client.list_tools()
+
+                return {
+                    "tools": tools,
+                    "server_session_id": client.server_session_id,
+                    "transport": client.detected_transport
+                }
+
         try:
-            import base64
-            import json
-            # Assuming JWT token, extract user ID from payload
-            # This is a basic implementation - adjust based on your token format
-            token_parts = self.auth_token.split('.')
-            if len(token_parts) >= 2:
-                payload_part = token_parts[1]
-                # Add padding if needed
-                padding = len(payload_part) % 4
-                if padding:
-                    payload_part += '=' * (4 - padding)
-                payload = json.loads(base64.b64decode(payload_part))
-                return payload.get('user_id') or payload.get('sub') or payload.get('uid')
-        except Exception as e:
-            logger.error(f"Error extracting user ID from token: {e}")
-        return None
+            # Run async operation
+            try:
+                loop = asyncio.get_event_loop()
+                if loop.is_running():
+                    # If we're already in an async context, create a new task
+                    import concurrent.futures
+                    with concurrent.futures.ThreadPoolExecutor() as executor:
+                        future = executor.submit(asyncio.run, _test_connection())
+                        result = future.result(timeout=120)
+                else:
+                    result = loop.run_until_complete(_test_connection())
+            except RuntimeError:
+                # No event loop, create one
+                result = asyncio.run(_test_connection())
 
+            execution_time = time.time() - start_time
+
+            # Extract tool names for the response
+            tool_names = [tool.get('name', 'unknown') for tool in result.get('tools', [])]
+
+            logger.info(f"[MCP Auth Check] Connection successful to '{toolkit_name}': {len(tool_names)} tools in {execution_time:.3f}s")
+
+            return {
+                "success": True,
+                "tools": tool_names,
+                "tools_count": len(tool_names),
+                "server_session_id": result.get('server_session_id'),
+                "transport": result.get('transport'),
+                "toolkit_config": toolkit_config,
+                "execution_time_seconds": execution_time
+            }
+
+        except McpAuthorizationRequired:
+            # Re-raise to allow proper handling upstream
+            raise
+        except Exception as e:
+            execution_time = time.time() - start_time
+            logger.error(f"[MCP Auth Check] Connection failed to '{toolkit_name}': {str(e)}")
+            return {
+                "success": False,
+                "error": f"MCP connection failed: {str(e)}",
+                "toolkit_config": toolkit_config,
+                "tools": [],
+                "tools_count": 0,
+                "execution_time_seconds": execution_time
+            }