alita-sdk 0.3.486__py3-none-any.whl → 0.3.515__py3-none-any.whl

alita_sdk/runtime/toolkits/mcp.py

@@ -15,9 +15,8 @@ from pydantic import BaseModel, ConfigDict, Field, SecretStr
 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_client import McpClient
 from ..utils.mcp_oauth import (
     McpAuthorizationRequired,
     canonical_resource,
@@ -40,110 +39,6 @@ def safe_int(value, default):
         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):
     """
     MCP Toolkit for connecting to a single remote MCP server and exposing its tools.
@@ -153,9 +48,6 @@ class McpToolkit(BaseToolkit):
     tools: List[BaseTool] = []
     toolkit_name: Optional[str] = None
 
-    # Class variable (not Pydantic field) for tool name length limit
-    toolkit_max_length: ClassVar[int] = 0  # No limit for MCP tool names
-
     def __getstate__(self):
         """Custom serialization for pickle compatibility."""
         state = self.__dict__.copy()
@@ -426,11 +318,6 @@ class McpToolkit(BaseToolkit):
         except Exception as e:
             logger.error(f"Direct discovery failed for MCP toolkit '{toolkit_name}': {e}", exc_info=True)
             logger.error(f"Discovery error details - URL: {connection_config.url}, Timeout: {timeout}s")
-
-            # Check if the exception wraps McpAuthorizationRequired (can happen with asyncio)
-            if hasattr(e, '__cause__') and isinstance(e.__cause__, McpAuthorizationRequired):
-                logger.info(f"Found wrapped McpAuthorizationRequired, re-raising")
-                raise e.__cause__
             
             # For new MCP toolkits (no client), don't silently return empty - surface the error
             # This helps users understand why tool discovery failed
@@ -464,37 +351,36 @@ class McpToolkit(BaseToolkit):
         toolkit_name: str,
         connection_config: McpConnectionConfig,
         timeout: int
-    ) -> List[Dict[str, Any]]:
+    ) -> tuple[List[Dict[str, Any]], Optional[str]]:
         """
         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.
+        
+        Returns:
+            Tuple of (tool_list, server_session_id) - session_id may be server-provided
         """
-        session_id = connection_config.session_id
+        initial_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")
+        if not initial_session_id:
+            logger.warning(f"[MCP Session] No session_id provided for '{toolkit_name}' - will generate one")
         
         # Run async discovery in sync context
         try:
-            all_tools = asyncio.run(
+            all_tools, server_session_id = asyncio.run(
                 cls._discover_tools_async(
                     toolkit_name=toolkit_name,
                     connection_config=connection_config,
                     timeout=timeout
                 )
             )
-            return all_tools, session_id
+            # Return tools and the session_id (server-provided or generated)
+            logger.info(f"[MCP Session] Final session_id for '{toolkit_name}': {server_session_id}")
+            return all_tools, server_session_id
         except McpAuthorizationRequired:
             # Re-raise auth required exceptions directly
             logger.info(f"[MCP SSE] Authorization required for '{toolkit_name}'")
             raise
         except Exception as e:
             logger.error(f"[MCP SSE] Discovery failed for '{toolkit_name}': {e}")
-            # Check if the exception wraps McpAuthorizationRequired
-            if hasattr(e, '__cause__') and isinstance(e.__cause__, McpAuthorizationRequired):
-                raise e.__cause__
             raise
     
     @classmethod
@@ -503,9 +389,12 @@ class McpToolkit(BaseToolkit):
         toolkit_name: str,
         connection_config: McpConnectionConfig,
         timeout: int
-    ) -> List[Dict[str, Any]]:
+    ) -> tuple[List[Dict[str, Any]], Optional[str]]:
         """
         Async implementation of tool discovery using SSE client.
+        
+        Returns:
+            Tuple of (tool_list, server_session_id)
         """
         all_tools = []
         session_id = connection_config.session_id
@@ -517,65 +406,74 @@ class McpToolkit(BaseToolkit):
             session_id = str(uuid.uuid4())
             logger.info(f"[MCP SSE] Generated temporary session_id for OAuth: {session_id}")
         
-        logger.info(f"[MCP SSE] Discovering from {connection_config.url} with session {session_id}")
+        logger.info(f"[MCP] 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(
+        # Create unified MCP client (auto-detects SSE vs Streamable HTTP)
+        client = McpClient(
             url=connection_config.url,
             session_id=session_id,
             headers=headers,
             timeout=timeout
         )
         
-        # 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 = await client.list_prompts()
-            # Convert prompts to tool format
-            for prompt in prompts:
-                prompt_tool = {
-                    "name": f"prompt_{prompt.get('name', 'unnamed')}",
-                    "description": prompt.get('description', f"Execute prompt: {prompt.get('name')}"),
-                    "inputSchema": {
-                        "type": "object",
-                        "properties": {
-                            "arguments": {
-                                "type": "object",
-                                "description": "Arguments for the prompt template",
-                                "properties": {
-                                    arg.get("name"): {
-                                        "type": "string",
-                                        "description": arg.get("description", ""),
-                                        "required": arg.get("required", False)
+        server_session_id = None
+        async with client:
+            # Initialize MCP session
+            await client.initialize()
+            logger.info(f"[MCP] Session initialized for '{toolkit_name}' (transport={client.detected_transport})")
+            
+            # Capture server-provided session_id (from mcp-session-id header)
+            server_session_id = client.server_session_id
+            if server_session_id:
+                logger.info(f"[MCP] Server provided session_id: {server_session_id}")
+            
+            # Discover tools
+            tools = await client.list_tools()
+            all_tools.extend(tools)
+            logger.info(f"[MCP] Discovered {len(tools)} tools from '{toolkit_name}'")
+            
+            # Discover prompts
+            try:
+                prompts = await client.list_prompts()
+                # Convert prompts to tool format
+                for prompt in prompts:
+                    prompt_tool = {
+                        "name": f"prompt_{prompt.get('name', 'unnamed')}",
+                        "description": prompt.get('description', f"Execute prompt: {prompt.get('name')}"),
+                        "inputSchema": {
+                            "type": "object",
+                            "properties": {
+                                "arguments": {
+                                    "type": "object",
+                                    "description": "Arguments for the prompt template",
+                                    "properties": {
+                                        arg.get("name"): {
+                                            "type": "string",
+                                            "description": arg.get("description", ""),
+                                            "required": arg.get("required", False)
+                                        }
+                                        for arg in prompt.get("arguments", [])
                                     }
-                                    for arg in prompt.get("arguments", [])
                                 }
                             }
-                        }
-                    },
-                    "_mcp_type": "prompt",
-                    "_mcp_prompt_name": prompt.get('name')
-                }
-                all_tools.append(prompt_tool)
-            logger.info(f"[MCP SSE] Discovered {len(prompts)} prompts from '{toolkit_name}'")
-        except Exception as e:
-            logger.warning(f"[MCP SSE] Failed to discover prompts: {e}")
+                        },
+                        "_mcp_type": "prompt",
+                        "_mcp_prompt_name": prompt.get('name')
+                    }
+                    all_tools.append(prompt_tool)
+                logger.info(f"[MCP] Discovered {len(prompts)} prompts from '{toolkit_name}'")
+            except Exception as e:
+                logger.warning(f"[MCP] Failed to discover prompts: {e}")
         
-        logger.info(f"[MCP SSE] Total discovered {len(all_tools)} items from '{toolkit_name}'")
-        return all_tools
+        logger.info(f"[MCP] Total discovered {len(all_tools)} items from '{toolkit_name}'")
+        # Return tools and server-provided session_id (use server's if available, else the one we sent)
+        final_session_id = server_session_id or session_id
+        return all_tools, final_session_id
 
     @classmethod
     def _create_tool_from_dict(
@@ -589,28 +487,23 @@ class McpToolkit(BaseToolkit):
     ) -> Optional[BaseTool]:
         """Create a BaseTool from a tool/prompt dictionary (from direct HTTP discovery)."""
         try:
-            # Store toolkit_max_length in local variable to avoid contextual access issues
-            max_length_value = cls.toolkit_max_length
-
-            # Clean toolkit name for prefixing
-            clean_prefix = clean_string(toolkit_name, max_length_value)
-
-            # Optimize tool name to fit within 64 character limit
-            full_tool_name = optimize_tool_name(clean_prefix, tool_dict.get("name", "unknown"))
+            # Use original tool name directly
+            tool_name = 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"
 
-            # 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', '')}"
+            # Build description with toolkit context and ensure it doesn't exceed 1000 characters
+            base_description = tool_dict.get('description', '')
+            description = f"{base_description}\nToolkit: {toolkit_name} ({connection_config.url})"
             if len(description) > 1000:
                 description = description[:997] + "..."
-                logger.debug(f"Trimmed description for tool '{tool_dict.get('name')}' from {len(description)} to 1000 chars")
+                logger.debug(f"Trimmed description for tool '{tool_name}' to 1000 chars")
 
             # Use McpRemoteTool for remote MCP servers (HTTP/SSE)
             return McpRemoteTool(
-                name=full_tool_name,
+                name=tool_name,
                 description=description,
                 args_schema=McpServerTool.create_pydantic_model_from_schema(
                     tool_dict.get("inputSchema", {})
@@ -622,11 +515,11 @@ class McpToolkit(BaseToolkit):
                 tool_timeout_sec=timeout,
                 is_prompt=is_prompt,
                 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
+                original_tool_name=tool_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}")
+            logger.error(f"Failed to create MCP tool '{tool_name}' from toolkit '{toolkit_name}': {e}")
             return None
 
     @classmethod
@@ -685,7 +578,7 @@ class McpToolkit(BaseToolkit):
             # We don't have full connection config in static mode, so create a basic one
             # The inspection tool will work as long as the server is accessible
             inspection_tool = McpInspectTool(
-                name=f"{clean_string(toolkit_name, 50)}{TOOLKIT_SPLITTER}mcp_inspect",
+                name="mcp_inspect",
                 server_name=toolkit_name,
                 server_url="",  # Will be populated by the client if available
                 description=f"Inspect available tools, prompts, and resources from MCP toolkit '{toolkit_name}'"
@@ -707,22 +600,17 @@ class McpToolkit(BaseToolkit):
     ) -> Optional[BaseTool]:
         """Create a BaseTool from discovered metadata."""
         try:
-            # Store toolkit_max_length in local variable to avoid contextual access issues
-            max_length_value = cls.toolkit_max_length
+            # Use original tool name directly
+            tool_name = tool_metadata.name
 
-            # Clean server name for prefixing (use tool_metadata.server instead of toolkit_name)
-            clean_prefix = clean_string(tool_metadata.server, max_length_value)
-            # 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}"
+            # Build description with toolkit context and ensure it doesn't exceed 1000 characters
+            description = f"{tool_metadata.description}\nToolkit: {toolkit_name}"
             if len(description) > 1000:
                 description = description[:997] + "..."
-                logger.debug(f"Trimmed description for tool '{tool_metadata.name}' from {len(description)} to 1000 chars")
+                logger.debug(f"Trimmed description for tool '{tool_name}' to 1000 chars")
 
             return McpServerTool(
-                name=full_tool_name,
+                name=tool_name,
                 description=description,
                 args_schema=McpServerTool.create_pydantic_model_from_schema(tool_metadata.input_schema),
                 client=client,
@@ -730,7 +618,7 @@ class McpToolkit(BaseToolkit):
                 tool_timeout_sec=timeout
             )
         except Exception as e:
-            logger.error(f"Failed to create MCP tool '{tool_metadata.name}' from server '{tool_metadata.server}': {e}")
+            logger.error(f"Failed to create MCP tool '{tool_name}' from server '{tool_metadata.server}': {e}")
             return None
 
     @classmethod
@@ -743,23 +631,18 @@ class McpToolkit(BaseToolkit):
     ) -> Optional[BaseTool]:
         """Create a single MCP tool."""
         try:
-            # Store toolkit_max_length in local variable to avoid contextual access issues
-            max_length_value = cls.toolkit_max_length
-
-            # Clean toolkit name for prefixing
-            clean_prefix = clean_string(toolkit_name, max_length_value)
-
-            # Optimize tool name to fit within 64 character limit
-            full_tool_name = optimize_tool_name(clean_prefix, available_tool["name"])
+            # Use original tool name directly
+            tool_name = 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', '')}"
+            # Build description with toolkit context and ensure it doesn't exceed 1000 characters
+            base_description = available_tool.get('description', '')
+            description = f"{base_description}\nToolkit: {toolkit_name}"
             if len(description) > 1000:
                 description = description[:997] + "..."
-                logger.debug(f"Trimmed description for tool '{available_tool['name']}' from {len(description)} to 1000 chars")
+                logger.debug(f"Trimmed description for tool '{tool_name}' to 1000 chars")
 
             return McpServerTool(
-                name=full_tool_name,
+                name=tool_name,
                 description=description,
                 args_schema=McpServerTool.create_pydantic_model_from_schema(
                     available_tool.get("inputSchema", {})
@@ -769,7 +652,7 @@ class McpToolkit(BaseToolkit):
                 tool_timeout_sec=timeout
             )
         except Exception as e:
-            logger.error(f"Failed to create MCP tool '{available_tool.get('name')}' from toolkit '{toolkit_name}': {e}")
+            logger.error(f"Failed to create MCP tool '{tool_name}' from toolkit '{toolkit_name}': {e}")
             return None
 
     @classmethod
@@ -780,16 +663,8 @@ class McpToolkit(BaseToolkit):
     ) -> Optional[BaseTool]:
         """Create the inspection tool for the MCP toolkit."""
         try:
-            # Store toolkit_max_length in local variable to avoid contextual access issues
-            max_length_value = cls.toolkit_max_length
-
-            # Clean toolkit name for prefixing
-            clean_prefix = clean_string(toolkit_name, max_length_value)
-
-            full_tool_name = f'{clean_prefix}{TOOLKIT_SPLITTER}mcp_inspect'
-
             return McpInspectTool(
-                name=full_tool_name,
+                name="mcp_inspect",
                 server_name=toolkit_name,
                 server_url=connection_config.url,
                 server_headers=connection_config.headers,

alita_sdk/runtime/toolkits/planning.py

@@ -16,7 +16,7 @@ from pydantic.fields import FieldInfo
 
 from ..tools.planning import PlanningWrapper
 from ...tools.base.tool import BaseAction
-from ...tools.utils import clean_string, TOOLKIT_SPLITTER, get_max_toolkit_length
+from ...tools.utils import clean_string, get_max_toolkit_length
 
 
 class PlanningToolkit(BaseToolkit):
@@ -94,7 +94,9 @@ class PlanningToolkit(BaseToolkit):
                         "label": "Planning",
                         "description": "Tools for managing multi-step execution plans with progress tracking. Uses PostgreSQL when configured, filesystem otherwise.",
                         "icon_url": None,
-                        "max_length": PlanningToolkit._toolkit_max_length
+                        "max_length": PlanningToolkit._toolkit_max_length,
+                        "categories": ["planning", "internal_tool"],
+                        "extra_categories": ["task management", "todo", "progress tracking"]
                     }
                 }
             )
@@ -148,8 +150,8 @@ class PlanningToolkit(BaseToolkit):
             plan_callback=plan_callback,
         )
         
-        # Build tool name prefix
-        prefix = clean_string(toolkit_name, cls._toolkit_max_length) + TOOLKIT_SPLITTER if toolkit_name else ''
+        # Use clean toolkit name for context (max 1000 chars in description)
+        toolkit_context = f" [Toolkit: {clean_string(toolkit_name, 0)}]" if toolkit_name else ''
         
         # Create tools from wrapper
         available_tools = wrapper.get_available_tools()
@@ -157,10 +159,15 @@ class PlanningToolkit(BaseToolkit):
             if tool["name"] not in selected_tools:
                 continue
             
+            # Add toolkit context to description with character limit
+            description = tool["description"]
+            if toolkit_context and len(description + toolkit_context) <= 1000:
+                description = description + toolkit_context
+            
             tools.append(BaseAction(
                 api_wrapper=wrapper,
-                name=prefix + tool["name"],
-                description=tool["description"],
+                name=tool["name"],
+                description=description,
                 args_schema=tool["args_schema"]
             ))