alita-sdk 0.3.435__py3-none-any.whl → 0.3.449__py3-none-any.whl

alita_sdk/runtime/toolkits/mcp.py

@@ -5,21 +5,145 @@ Following MCP specification: https://modelcontextprotocol.io/specification/2025-
 """
 
 import logging
+import re
 import requests
-from typing import List, Optional, Any, Dict, Literal, ClassVar
+import asyncio
+from typing import List, Optional, Any, Dict, Literal, ClassVar, Union
 
 from langchain_core.tools import BaseToolkit, BaseTool
 from pydantic import BaseModel, ConfigDict, Field
 
 from ..tools.mcp_server_tool import McpServerTool
+from ..tools.mcp_remote_tool import McpRemoteTool
 from ..tools.mcp_inspect_tool import McpInspectTool
 from ...tools.utils import TOOLKIT_SPLITTER, clean_string
 from ..models.mcp_models import McpConnectionConfig
+from ..utils.mcp_sse_client import McpSseClient
+from ..utils.mcp_oauth import (
+    McpAuthorizationRequired,
+    canonical_resource,
+    extract_resource_metadata_url,
+    fetch_resource_metadata,
+    infer_authorization_servers_from_realm,
+)
 
 logger = logging.getLogger(__name__)
 
 name = "mcp"
 
+def safe_int(value, default):
+    """Convert value to int, handling string inputs."""
+    if value is None:
+        return default
+    try:
+        return int(value)
+    except (ValueError, TypeError):
+        logger.warning(f"Invalid integer value '{value}', using default {default}")
+        return default
+
+def optimize_tool_name(prefix: str, tool_name: str, max_total_length: int = 64) -> str:
+    """
+    Optimize tool name to fit within max_total_length while preserving meaning.
+    
+    Args:
+        prefix: The toolkit prefix (already cleaned)
+        tool_name: The original tool name
+        max_total_length: Maximum total length for the full tool name (default: 64)
+    
+    Returns:
+        Optimized full tool name in format: prefix___tool_name
+    """
+    splitter = TOOLKIT_SPLITTER
+    splitter_len = len(splitter)
+    prefix_len = len(prefix)
+    
+    # Calculate available space for tool name
+    available_space = max_total_length - prefix_len - splitter_len
+    
+    if available_space <= 0:
+        logger.error(f"Prefix '{prefix}' is too long ({prefix_len} chars), cannot create valid tool name")
+        # Fallback: truncate prefix itself
+        prefix = prefix[:max_total_length - splitter_len - 10]  # Leave 10 chars for tool name
+        available_space = max_total_length - len(prefix) - splitter_len
+    
+    # If tool name fits, use it as-is
+    if len(tool_name) <= available_space:
+        return f'{prefix}{splitter}{tool_name}'
+    
+    # Tool name is too long, need to optimize
+    logger.debug(f"Tool name '{tool_name}' is too long ({len(tool_name)} chars), optimizing to fit {available_space} chars")
+    
+    # Split tool name into parts (handle camelCase, snake_case, and mixed)
+    # First, split by underscores and hyphens
+    parts = re.split(r'[_-]', tool_name)
+    
+    # Further split camelCase within each part
+    all_parts = []
+    for part in parts:
+        # Insert underscore before uppercase letters (camelCase to snake_case)
+        snake_case_part = re.sub(r'([a-z0-9])([A-Z])', r'\1_\2', part)
+        all_parts.extend(snake_case_part.split('_'))
+    
+    # Filter out empty parts
+    all_parts = [p for p in all_parts if p]
+    
+    # Remove redundant prefix words (case-insensitive comparison)
+    # Only remove if prefix is meaningful (>= 3 chars) to avoid over-filtering
+    prefix_lower = prefix.lower()
+    filtered_parts = []
+    for part in all_parts:
+        part_lower = part.lower()
+        # Skip if this part contains the prefix or the prefix contains this part
+        # But only if both are meaningful (>= 3 chars)
+        should_remove = False
+        if len(prefix_lower) >= 3 and len(part_lower) >= 3:
+            if part_lower in prefix_lower or prefix_lower in part_lower:
+                should_remove = True
+                logger.debug(f"Removing redundant part '{part}' (matches prefix '{prefix}')")
+        
+        if not should_remove:
+            filtered_parts.append(part)
+    
+    # If we removed all parts, keep the original parts
+    if not filtered_parts:
+        filtered_parts = all_parts
+    
+    # Reconstruct tool name with filtered parts
+    optimized_name = '_'.join(filtered_parts)
+    
+    # If still too long, truncate intelligently
+    if len(optimized_name) > available_space:
+        # Strategy: Keep beginning and end, as they often contain the most important info
+        # For example: "projectalita_github_io_list_branches" -> "projectalita_list_branches"
+        
+        # Try removing middle parts first
+        if len(filtered_parts) > 2:
+            # Keep first and last parts, remove middle
+            kept_parts = [filtered_parts[0], filtered_parts[-1]]
+            optimized_name = '_'.join(kept_parts)
+            
+            # If still too long, add parts from the end until we run out of space
+            if len(optimized_name) <= available_space and len(filtered_parts) > 2:
+                for i in range(len(filtered_parts) - 2, 0, -1):
+                    candidate = '_'.join([filtered_parts[0]] + filtered_parts[i:])
+                    if len(candidate) <= available_space:
+                        optimized_name = candidate
+                        break
+        
+        # If still too long, just truncate
+        if len(optimized_name) > available_space:
+            # Try to truncate at word boundary
+            truncated = optimized_name[:available_space]
+            last_underscore = truncated.rfind('_')
+            if last_underscore > available_space * 0.7:  # Keep if we're not losing too much
+                optimized_name = truncated[:last_underscore]
+            else:
+                optimized_name = truncated
+    
+    full_name = f'{prefix}{splitter}{optimized_name}'
+    logger.info(f"Optimized tool name: '{tool_name}' ({len(tool_name)} chars) -> '{optimized_name}' ({len(optimized_name)} chars), full: '{full_name}' ({len(full_name)} chars)")
+    
+    return full_name
 
 class McpToolkit(BaseToolkit):
     """
@@ -85,11 +209,10 @@ class McpToolkit(BaseToolkit):
                 )
             ),
             timeout=(
-                int,
+                Union[int, str], # TODO: remove one I will figure out why UI sends str
                 Field(
-                    default=60,
-                    ge=1, le=300,
-                    description="Request timeout in seconds"
+                    default=300,
+                    description="Request timeout in seconds (1-3600)"
                 )
             ),
             discovery_mode=(
@@ -103,11 +226,10 @@ class McpToolkit(BaseToolkit):
                 )
             ),
             discovery_interval=(
-                int,
+                Union[int, str],
                 Field(
                     default=300,
-                    ge=60, le=3600,
-                    description="Discovery interval in seconds (for periodic discovery)"
+                    description="Discovery interval in seconds (60-3600, for periodic discovery)"
                 )
             ),
             selected_tools=(
@@ -128,11 +250,10 @@ class McpToolkit(BaseToolkit):
                 )
             ),
             cache_ttl=(
-                int,
+                Union[int, str],
                 Field(
                     default=300,
-                    ge=60, le=3600,
-                    description="Cache TTL in seconds"
+                    description="Cache TTL in seconds (60-3600)"
                 )
             ),
             __config__=ConfigDict(
@@ -194,6 +315,11 @@ class McpToolkit(BaseToolkit):
         if not toolkit_name:
             raise ValueError("toolkit_name is required")
 
+        # Convert numeric parameters that may come as strings from UI
+        timeout = safe_int(timeout, 60)
+        discovery_interval = safe_int(discovery_interval, 300)
+        cache_ttl = safe_int(cache_ttl, 300)
+
         logger.info(f"Creating MCP toolkit: {toolkit_name}")
 
         # Parse headers if they're provided as a JSON string
@@ -216,9 +342,14 @@ class McpToolkit(BaseToolkit):
             logger.error(f"Headers must be a dictionary or JSON string, got: {type(headers)}")
             raise ValueError(f"Headers must be a dictionary or JSON string, got: {type(headers)}")
 
+        # Extract session_id from kwargs if provided
+        session_id = kwargs.get('session_id')
+        if session_id:
+            logger.info(f"[MCP Session] Using provided session ID for toolkit '{toolkit_name}': {session_id}")
+        
         # Create MCP connection configuration
         try:
-            connection_config = McpConnectionConfig(url=url, headers=parsed_headers)
+            connection_config = McpConnectionConfig(url=url, headers=parsed_headers, session_id=session_id)
         except Exception as e:
             logger.error(f"Invalid MCP connection configuration: {e}")
             raise ValueError(f"Invalid MCP connection configuration: {e}")
@@ -258,7 +389,7 @@ class McpToolkit(BaseToolkit):
             logger.info(f"Discovering tools from MCP toolkit '{toolkit_name}' at {connection_config.url}")
 
             # Use synchronous HTTP discovery for toolkit initialization
-            tool_metadata_list = cls._discover_tools_sync(
+            tool_metadata_list, session_id = cls._discover_tools_sync(
                 toolkit_name=toolkit_name,
                 connection_config=connection_config,
                 timeout=timeout
@@ -273,12 +404,18 @@ class McpToolkit(BaseToolkit):
                 ]
 
             # Create BaseTool instances from discovered metadata
+            # Use session_id from frontend (passed via connection_config)
+            if session_id:
+                logger.info(f"[MCP Session] Using session_id from frontend: {session_id}")
+            
             for tool_metadata in tool_metadata_list:
                 server_tool = cls._create_tool_from_dict(
                     tool_dict=tool_metadata,
                     toolkit_name=toolkit_name,
+                    connection_config=connection_config,
                     timeout=timeout,
-                    client=client
+                    client=client,
+                    session_id=session_id  # Use session from discovery
                 )
 
                 if server_tool:
@@ -291,6 +428,9 @@ class McpToolkit(BaseToolkit):
             logger.error(f"Discovery error details - URL: {connection_config.url}, Timeout: {timeout}s")
 
             # Fallback to static mode if available and not already static
+            if isinstance(e, McpAuthorizationRequired):
+                # Authorization is required; surface upstream so the caller can prompt the user
+                raise
             if client and discovery_mode != "static":
                 logger.info(f"Falling back to static discovery for toolkit '{toolkit_name}'")
                 tools = cls._create_tools_static(toolkit_name, selected_tools, timeout, client)
@@ -321,32 +461,76 @@ class McpToolkit(BaseToolkit):
         timeout: int
     ) -> List[Dict[str, Any]]:
         """
-        Synchronously discover tools and prompts from MCP server using HTTP requests.
+        Discover tools and prompts from MCP server using SSE client.
         Returns list of tool/prompt dictionaries with name, description, and inputSchema.
         Prompts are converted to tools that can be invoked.
         """
+        session_id = connection_config.session_id
+        
+        if not session_id:
+            logger.warning(f"[MCP Session] No session_id provided for '{toolkit_name}' - server may require it")
+            logger.warning(f"[MCP Session] Frontend should generate a UUID and include it with mcp_tokens")
+        
+        # Run async discovery in sync context
+        try:
+            all_tools = asyncio.run(
+                cls._discover_tools_async(
+                    toolkit_name=toolkit_name,
+                    connection_config=connection_config,
+                    timeout=timeout
+                )
+            )
+            return all_tools, session_id
+        except Exception as e:
+            logger.error(f"[MCP SSE] Discovery failed for '{toolkit_name}': {e}")
+            raise
+    
+    @classmethod
+    async def _discover_tools_async(
+        cls,
+        toolkit_name: str,
+        connection_config: McpConnectionConfig,
+        timeout: int
+    ) -> List[Dict[str, Any]]:
+        """
+        Async implementation of tool discovery using SSE client.
+        """
         all_tools = []
+        session_id = connection_config.session_id
         
-        # Discover regular tools
-        tools_data = cls._discover_mcp_endpoint(
-            endpoint="tools/list",
-            toolkit_name=toolkit_name,
-            connection_config=connection_config,
+        if not session_id:
+            logger.error(f"[MCP SSE] session_id is required for SSE servers")
+            raise ValueError("session_id is required. Frontend must generate UUID.")
+        
+        logger.info(f"[MCP SSE] Discovering from {connection_config.url} with session {session_id}")
+        
+        # Prepare headers
+        headers = {}
+        if connection_config.headers:
+            headers.update(connection_config.headers)
+        
+        # Create SSE client
+        client = McpSseClient(
+            url=connection_config.url,
+            session_id=session_id,
+            headers=headers,
             timeout=timeout
         )
-        all_tools.extend(tools_data)
-        logger.info(f"Discovered {len(tools_data)} tools from MCP toolkit '{toolkit_name}'")
         
-        # Discover prompts and convert them to tools
+        # Initialize MCP session
+        await client.initialize()
+        logger.info(f"[MCP SSE] Session initialized for '{toolkit_name}'")
+        
+        # Discover tools
+        tools = await client.list_tools()
+        all_tools.extend(tools)
+        logger.info(f"[MCP SSE] Discovered {len(tools)} tools from '{toolkit_name}'")
+        
+        # Discover prompts
         try:
-            prompts_data = cls._discover_mcp_endpoint(
-                endpoint="prompts/list",
-                toolkit_name=toolkit_name,
-                connection_config=connection_config,
-                timeout=timeout
-            )
+            prompts = await client.list_prompts()
             # Convert prompts to tool format
-            for prompt in prompts_data:
+            for prompt in prompts:
                 prompt_tool = {
                     "name": f"prompt_{prompt.get('name', 'unnamed')}",
                     "description": prompt.get('description', f"Execute prompt: {prompt.get('name')}"),
@@ -371,117 +555,22 @@ class McpToolkit(BaseToolkit):
                     "_mcp_prompt_name": prompt.get('name')
                 }
                 all_tools.append(prompt_tool)
-            logger.info(f"Discovered {len(prompts_data)} prompts from MCP toolkit '{toolkit_name}'")
+            logger.info(f"[MCP SSE] Discovered {len(prompts)} prompts from '{toolkit_name}'")
         except Exception as e:
-            logger.warning(f"Failed to discover prompts from MCP toolkit '{toolkit_name}': {e}")
+            logger.warning(f"[MCP SSE] Failed to discover prompts: {e}")
         
-        logger.info(f"Total discovered {len(all_tools)} tools+prompts from MCP toolkit '{toolkit_name}'")
+        logger.info(f"[MCP SSE] Total discovered {len(all_tools)} items from '{toolkit_name}'")
         return all_tools
 
-    @classmethod
-    def _discover_mcp_endpoint(
-        cls,
-        endpoint: str,
-        toolkit_name: str,
-        connection_config: McpConnectionConfig,
-        timeout: int
-    ) -> List[Dict[str, Any]]:
-        """
-        Discover items from a specific MCP endpoint (tools/list or prompts/list).
-        Returns list of dictionaries.
-        """
-        import time
-
-        # MCP protocol request
-        mcp_request = {
-            "jsonrpc": "2.0",
-            "id": f"discover_{endpoint.replace('/', '_')}_{int(time.time())}",
-            "method": endpoint,
-            "params": {}
-        }
-
-        headers = {
-            "Content-Type": "application/json",
-            "Accept": "application/json, text/event-stream"
-        }
-        if connection_config.headers:
-            headers.update(connection_config.headers)
-
-        try:
-            logger.debug(f"Sending MCP {endpoint} request to {connection_config.url}")
-            response = requests.post(
-                connection_config.url,
-                json=mcp_request,
-                headers=headers,
-                timeout=timeout
-            )
-
-            if response.status_code != 200:
-                logger.error(f"MCP server returned non-200 status: {response.status_code}")
-                raise Exception(f"HTTP {response.status_code}: {response.text}")
-
-            # Check content type and parse accordingly
-            content_type = response.headers.get('Content-Type', '')
-            
-            if 'text/event-stream' in content_type:
-                # Parse SSE (Server-Sent Events) format
-                data = cls._parse_sse_response(response.text)
-            elif 'application/json' in content_type:
-                # Parse regular JSON
-                try:
-                    data = response.json()
-                except ValueError as json_err:
-                    raise Exception(f"Invalid JSON response: {json_err}. Response text: {response.text[:200]}")
-            else:
-                raise Exception(f"Unexpected Content-Type: {content_type}. Response: {response.text[:200]}")
-
-            if "error" in data:
-                raise Exception(f"MCP Error: {data['error']}")
-
-            # Parse MCP response - different endpoints return different keys
-            result = data.get("result", {})
-            if endpoint == "tools/list":
-                return result.get("tools", [])
-            elif endpoint == "prompts/list":
-                return result.get("prompts", [])
-            else:
-                return result.get("items", [])
-
-        except Exception as e:
-            logger.error(f"Failed to discover from {endpoint} on MCP toolkit '{toolkit_name}': {e}")
-            raise
-
-    @staticmethod
-    def _parse_sse_response(sse_text: str) -> Dict[str, Any]:
-        """
-        Parse Server-Sent Events (SSE) format response.
-        SSE format: event: message\ndata: {json}\n\n
-        """
-        import json
-        
-        lines = sse_text.strip().split('\n')
-        data_line = None
-        
-        for line in lines:
-            if line.startswith('data:'):
-                data_line = line[5:].strip()  # Remove 'data:' prefix
-                break
-        
-        if not data_line:
-            raise Exception(f"No data found in SSE response: {sse_text[:200]}")
-        
-        try:
-            return json.loads(data_line)
-        except json.JSONDecodeError as e:
-            raise Exception(f"Failed to parse SSE data as JSON: {e}. Data: {data_line[:200]}")
-
     @classmethod
     def _create_tool_from_dict(
         cls,
         tool_dict: Dict[str, Any],
         toolkit_name: str,
+        connection_config: McpConnectionConfig,
         timeout: int,
-        client
+        client,
+        session_id: Optional[str] = None
     ) -> Optional[BaseTool]:
         """Create a BaseTool from a tool/prompt dictionary (from direct HTTP discovery)."""
         try:
@@ -491,23 +580,35 @@ class McpToolkit(BaseToolkit):
             # Clean toolkit name for prefixing
             clean_prefix = clean_string(toolkit_name, max_length_value)
 
-            full_tool_name = f'{clean_prefix}{TOOLKIT_SPLITTER}{tool_dict.get("name", "unknown")}'
+            # Optimize tool name to fit within 64 character limit
+            full_tool_name = optimize_tool_name(clean_prefix, tool_dict.get("name", "unknown"))
             
             # Check if this is a prompt (converted to tool)
             is_prompt = tool_dict.get("_mcp_type") == "prompt"
             item_type = "prompt" if is_prompt else "tool"
 
-            return McpServerTool(
+            # Build description and ensure it doesn't exceed 1000 characters
+            description = f"MCP {item_type} '{tool_dict.get('name')}' from toolkit '{toolkit_name}': {tool_dict.get('description', '')}"
+            if len(description) > 1000:
+                description = description[:997] + "..."
+                logger.debug(f"Trimmed description for tool '{tool_dict.get('name')}' from {len(description)} to 1000 chars")
+
+            # Use McpRemoteTool for remote MCP servers (HTTP/SSE)
+            return McpRemoteTool(
                 name=full_tool_name,
-                description=f"MCP {item_type} '{tool_dict.get('name')}' from toolkit '{toolkit_name}': {tool_dict.get('description', '')}",
+                description=description,
                 args_schema=McpServerTool.create_pydantic_model_from_schema(
                     tool_dict.get("inputSchema", {})
                 ),
                 client=client,
                 server=toolkit_name,
+                server_url=connection_config.url,
+                server_headers=connection_config.headers,
                 tool_timeout_sec=timeout,
                 is_prompt=is_prompt,
-                prompt_name=tool_dict.get("_mcp_prompt_name") if is_prompt else None
+                prompt_name=tool_dict.get("_mcp_prompt_name") if is_prompt else None,
+                original_tool_name=tool_dict.get('name'),  # Store original name for MCP server invocation
+                session_id=session_id  # Pass session ID for stateful SSE servers
             )
         except Exception as e:
             logger.error(f"Failed to create MCP tool '{tool_dict.get('name')}' from toolkit '{toolkit_name}': {e}")
@@ -596,11 +697,18 @@ class McpToolkit(BaseToolkit):
 
             # Clean server name for prefixing (use tool_metadata.server instead of toolkit_name)
             clean_prefix = clean_string(tool_metadata.server, max_length_value)
-            full_tool_name = f'{clean_prefix}{TOOLKIT_SPLITTER}{tool_metadata.name}'
+            # Optimize tool name to fit within 64 character limit
+            full_tool_name = optimize_tool_name(clean_prefix, tool_metadata.name)
+
+            # Build description and ensure it doesn't exceed 1000 characters
+            description = f"MCP tool '{tool_metadata.name}' from server '{tool_metadata.server}': {tool_metadata.description}"
+            if len(description) > 1000:
+                description = description[:997] + "..."
+                logger.debug(f"Trimmed description for tool '{tool_metadata.name}' from {len(description)} to 1000 chars")
 
             return McpServerTool(
                 name=full_tool_name,
-                description=f"MCP tool '{tool_metadata.name}' from server '{tool_metadata.server}': {tool_metadata.description}",
+                description=description,
                 args_schema=McpServerTool.create_pydantic_model_from_schema(tool_metadata.input_schema),
                 client=client,
                 server=tool_metadata.server,
@@ -626,11 +734,18 @@ class McpToolkit(BaseToolkit):
             # Clean toolkit name for prefixing
             clean_prefix = clean_string(toolkit_name, max_length_value)
 
-            full_tool_name = f'{clean_prefix}{TOOLKIT_SPLITTER}{available_tool["name"]}'
+            # Optimize tool name to fit within 64 character limit
+            full_tool_name = optimize_tool_name(clean_prefix, available_tool["name"])
+
+            # Build description and ensure it doesn't exceed 1000 characters
+            description = f"MCP tool '{available_tool['name']}' from toolkit '{toolkit_name}': {available_tool.get('description', '')}"
+            if len(description) > 1000:
+                description = description[:997] + "..."
+                logger.debug(f"Trimmed description for tool '{available_tool['name']}' from {len(description)} to 1000 chars")
 
             return McpServerTool(
                 name=full_tool_name,
-                description=f"MCP tool '{available_tool['name']}' from toolkit '{toolkit_name}': {available_tool.get('description', '')}",
+                description=description,
                 args_schema=McpServerTool.create_pydantic_model_from_schema(
                     available_tool.get("inputSchema", {})
                 ),
@@ -736,16 +851,6 @@ def get_tools(tool_config: dict, alita_client, llm=None, memory_store=None) -> L
         return []
 
     # Type conversion for numeric settings that may come as strings from config
-    def safe_int(value, default):
-        """Convert value to int, handling string inputs."""
-        if value is None:
-            return default
-        try:
-            return int(value)
-        except (ValueError, TypeError):
-            logger.warning(f"Invalid integer value '{value}', using default {default}")
-            return default
-
     return McpToolkit.get_toolkit(
         url=url,
         headers=headers,
@@ -784,4 +889,4 @@ def get_all_discovered_servers():
     """Get status of all discovered servers."""
     from ..clients.mcp_discovery import get_discovery_service
     service = get_discovery_service()
-    return service.get_server_health()
+    return service.get_server_health()

alita_sdk/runtime/toolkits/tools.py

@@ -19,6 +19,7 @@ from ..tools.image_generation import ImageGenerationToolkit
 # Import community tools
 from ...community import get_toolkits as community_toolkits, get_tools as community_tools
 from ...tools.memory import MemoryToolkit
+from ..utils.mcp_oauth import canonical_resource, McpAuthorizationRequired
 from ...tools.utils import TOOLKIT_SPLITTER
 
 logger = logging.getLogger(__name__)
@@ -37,7 +38,7 @@ def get_toolkits():
     return core_toolkits + community_toolkits() + alita_toolkits()
 
 
-def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = None, debug_mode: Optional[bool] = False) -> list:
+def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = None, debug_mode: Optional[bool] = False, mcp_tokens: Optional[dict] = None) -> list:
     prompts = []
     tools = []
 
@@ -109,11 +110,51 @@ def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = Non
                     toolkit_name=tool.get('toolkit_name', ''),
                     **tool['settings']).get_tools())
             elif tool['type'] == 'mcp':
+                settings = dict(tool['settings'])
+                url = settings.get('url')
+                headers = settings.get('headers')
+                token_data = None
+                session_id = None
+                if mcp_tokens and url:
+                    canonical_url = canonical_resource(url)
+                    logger.info(f"[MCP Auth] Looking for token for URL: {url}")
+                    logger.info(f"[MCP Auth] Canonical URL: {canonical_url}")
+                    logger.info(f"[MCP Auth] Available tokens: {list(mcp_tokens.keys())}")
+                    token_data = mcp_tokens.get(canonical_url)
+                    if token_data:
+                        logger.info(f"[MCP Auth] Found token data for {canonical_url}")
+                        # Handle both old format (string) and new format (dict with access_token and session_id)
+                        if isinstance(token_data, dict):
+                            access_token = token_data.get('access_token')
+                            session_id = token_data.get('session_id')
+                            logger.info(f"[MCP Auth] Token data: access_token={'present' if access_token else 'missing'}, session_id={session_id or 'none'}")
+                        else:
+                            # Backward compatibility: treat as plain token string
+                            access_token = token_data
+                            logger.info(f"[MCP Auth] Using legacy token format (string)")
+                    else:
+                        access_token = None
+                        logger.warning(f"[MCP Auth] No token found for {canonical_url}")
+                else:
+                    access_token = None
+                    
+                if access_token:
+                    merged_headers = dict(headers) if headers else {}
+                    merged_headers.setdefault('Authorization', f'Bearer {access_token}')
+                    settings['headers'] = merged_headers
+                    logger.info(f"[MCP Auth] Added Authorization header for {url}")
+                    
+                # Pass session_id to MCP toolkit if available
+                if session_id:
+                    settings['session_id'] = session_id
+                    logger.info(f"[MCP Auth] Passing session_id to toolkit: {session_id}")
                 tools.extend(McpToolkit.get_toolkit(
                     toolkit_name=tool.get('toolkit_name', ''),
                     client=alita_client,
-                    **tool['settings']).get_tools())
+                    **settings).get_tools())
         except Exception as e:
+            if isinstance(e, McpAuthorizationRequired):
+                raise
             logger.error(f"Error initializing toolkit for tool '{tool.get('name', 'unknown')}': {e}", exc_info=True)
             if debug_mode:
                 logger.info("Skipping tool initialization error due to debug mode.")