Merge pull request #1220 from escapade-mckv/bug-fixes

Bug fixes

Author: escapade-mckv

backend/agent/prompt.py

@@ -202,6 +202,7 @@ async def register_mcp_tools(self, agent_config: dict) -> Optional[MCPToolWrappe
                         "schema": schema
                     }
 
+            logger.info(f"⚡ Registered {len(updated_schemas)} MCP tools (Redis cache enabled)")
             return mcp_wrapper_instance
         except Exception as e:
             logger.error(f"Failed to initialize MCP tools: {e}")

backend/agent/run.py

@@ -314,8 +314,16 @@ async def get_current_agent_config(self) -> ToolResult:
                 "updated_at": agent_data.get("updated_at"),
                 "current_version": agent_config.get("version_name", "v1") if version_data else "No version data"
             }
-            
-            tools_count = len([t for t, cfg in config_summary["agentpress_tools"].items() if cfg.get("enabled")])
+
+            enabled_tools = []
+            for tool_name, tool_config in config_summary["agentpress_tools"].items():
+                if isinstance(tool_config, bool):
+                    if tool_config:
+                        enabled_tools.append(tool_name)
+                elif isinstance(tool_config, dict):
+                    if tool_config.get("enabled", False):
+                        enabled_tools.append(tool_name)
+            tools_count = len(enabled_tools)
             mcps_count = len(config_summary["configured_mcps"])
             custom_mcps_count = len(config_summary["custom_mcps"])
 

backend/agent/tools/agent_builder_tools/agent_config_tool.py

@@ -140,7 +140,15 @@ async def create_credential_profile(
 
             if result.connected_account.redirect_url:
                 response_data["connection_link"] = result.connected_account.redirect_url
-                response_data["instructions"] = f"🔗 **IMPORTANT: Please visit the connection link to authenticate your {result.toolkit.name} account with this profile. After connecting, you'll be able to use {result.toolkit.name} tools in your agent.**"
+                # Include both the toolkit name and slug in a parseable format
+                # Format: [toolkit:slug:name] to help frontend identify the service accurately
+                response_data["instructions"] = f"""🔗 **{result.toolkit.name} Authentication Required**
+
+Please authenticate your {result.toolkit.name} account by clicking the link below:
+
+[toolkit:{toolkit_slug}:{result.toolkit.name}] Authentication: {result.connected_account.redirect_url}
+
+After connecting, you'll be able to use {result.toolkit.name} tools in your agent."""
             else:
                 response_data["instructions"] = f"This {result.toolkit.name} profile has been created and is ready to use."
 
@@ -256,12 +264,48 @@ async def configure_profile_for_agent(
                 change_description=f"Configured {display_name or profile.display_name} with {len(enabled_tools)} tools"
             )
 
+            # Dynamically register the MCP tools in the current runtime
+            try:
+                from agent.tools.mcp_tool_wrapper import MCPToolWrapper
+                
+                mcp_config_for_wrapper = {
+                    'name': profile.toolkit_name,
+                    'qualifiedName': f"composio.{profile.toolkit_slug}",
+                    'config': {
+                        'profile_id': profile_id,
+                        'toolkit_slug': profile.toolkit_slug,
+                        'mcp_qualified_name': profile.mcp_qualified_name
+                    },
+                    'enabledTools': enabled_tools,
+                    'instructions': '',
+                    'isCustom': True,
+                    'customType': 'composio'
+                }
+                
+                mcp_wrapper_instance = MCPToolWrapper(mcp_configs=[mcp_config_for_wrapper])
+                await mcp_wrapper_instance.initialize_and_register_tools()
+                updated_schemas = mcp_wrapper_instance.get_schemas()
+                
+                for method_name, schema_list in updated_schemas.items():
+                    for schema in schema_list:
+                        self.thread_manager.tool_registry.tools[method_name] = {
+                            "instance": mcp_wrapper_instance,
+                            "schema": schema
+                        }
+                        logger.info(f"Dynamically registered MCP tool: {method_name}")
+                
+                logger.info(f"Successfully registered {len(updated_schemas)} MCP tools dynamically for {profile.toolkit_name}")
+                
+            except Exception as e:
+                logger.warning(f"Could not dynamically register MCP tools in current runtime: {str(e)}. Tools will be available on next agent run.")
+
             return self.success_response({
-                "message": f"Profile '{profile.profile_name}' updated with {len(enabled_tools)} tools",
+                "message": f"Profile '{profile.profile_name}' configured with {len(enabled_tools)} tools and registered in current runtime",
                 "enabled_tools": enabled_tools,
                 "total_tools": len(enabled_tools),
                 "version_id": new_version.version_id,
-                "version_name": new_version.version_name
+                "version_name": new_version.version_name,
+                "runtime_registration": "success"
             })
 
         except Exception as e:

backend/agent/tools/agent_builder_tools/credential_profile_tool.py

@@ -52,8 +52,13 @@ async def _get_available_tools_for_agent(self) -> List[str]:
 
             agentpress_tools = agent_config.get('agentpress_tools', {})
             for tool_key, tool_names in tool_mapping.items():
-                if agentpress_tools.get(tool_key, {}).get('enabled', False):
-                    available_tools.extend(tool_names)
+                tool_config = agentpress_tools.get(tool_key, False)
+                if isinstance(tool_config, bool):
+                    if tool_config:
+                        available_tools.extend(tool_names)
+                elif isinstance(tool_config, dict):
+                    if tool_config.get('enabled', False):
+                        available_tools.extend(tool_names)
 
             configured_mcps = agent_config.get('configured_mcps', [])
             for mcp in configured_mcps:

backend/agent/tools/agent_builder_tools/workflow_tool.py

@@ -3,20 +3,120 @@
 from mcp_module import mcp_service
 from utils.logger import logger
 import inspect
+import asyncio
+import time
+import hashlib
+import json
 from agent.tools.utils.mcp_connection_manager import MCPConnectionManager
 from agent.tools.utils.custom_mcp_handler import CustomMCPHandler
 from agent.tools.utils.dynamic_tool_builder import DynamicToolBuilder
 from agent.tools.utils.mcp_tool_executor import MCPToolExecutor
+from services import redis as redis_service
 
 
+class MCPSchemaRedisCache:
+    def __init__(self, ttl_seconds: int = 3600, key_prefix: str = "mcp_schema:"):
+        self._ttl = ttl_seconds
+        self._key_prefix = key_prefix
+        self._redis_client = None
+    
+    async def _ensure_redis(self):
+        if not self._redis_client:
+            try:
+                self._redis_client = await redis_service.get_client()
+            except Exception as e:
+                logger.warning(f"Redis not available for MCP cache: {e}")
+                return False
+        return True
+    
+    def _get_cache_key(self, config: Dict[str, Any]) -> str:
+        config_str = json.dumps(config, sort_keys=True)
+        config_hash = hashlib.md5(config_str.encode()).hexdigest()
+        return f"{self._key_prefix}{config_hash}"
+    
+    async def get(self, config: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        if not await self._ensure_redis():
+            return None
+            
+        try:
+            key = self._get_cache_key(config)
+            cached_data = await self._redis_client.get(key)
+            
+            if cached_data:
+                logger.debug(f"⚡ Redis cache hit for MCP: {config.get('name', config.get('qualifiedName', 'Unknown'))}")
+                return json.loads(cached_data)
+            else:
+                logger.debug(f"Redis cache miss for MCP: {config.get('name', config.get('qualifiedName', 'Unknown'))}")
+                return None
+                
+        except Exception as e:
+            logger.warning(f"Error reading from Redis cache: {e}")
+            return None
+    
+    async def set(self, config: Dict[str, Any], data: Dict[str, Any]):
+        if not await self._ensure_redis():
+            return
+            
+        try:
+            key = self._get_cache_key(config)
+            serialized_data = json.dumps(data)
+            
+            await self._redis_client.setex(key, self._ttl, serialized_data)
+            logger.debug(f"✅ Cached MCP schema in Redis for {config.get('name', config.get('qualifiedName', 'Unknown'))} (TTL: {self._ttl}s)")
+            
+        except Exception as e:
+            logger.warning(f"Error writing to Redis cache: {e}")
+    
+    async def clear_pattern(self, pattern: Optional[str] = None):
+        if not await self._ensure_redis():
+            return
+        try:
+            if pattern:
+                search_pattern = f"{self._key_prefix}{pattern}*"
+            else:
+                search_pattern = f"{self._key_prefix}*"
+            
+            keys = []
+            async for key in self._redis_client.scan_iter(match=search_pattern):
+                keys.append(key)
+            
+            if keys:
+                await self._redis_client.delete(*keys)
+                logger.info(f"Cleared {len(keys)} MCP schema cache entries from Redis")
+            
+        except Exception as e:
+            logger.warning(f"Error clearing Redis cache: {e}")
+    
+    async def get_stats(self) -> Dict[str, Any]:
+        if not await self._ensure_redis():
+            return {"available": False}
+        try:
+            count = 0
+            async for _ in self._redis_client.scan_iter(match=f"{self._key_prefix}*"):
+                count += 1
+            
+            return {
+                "available": True,
+                "cached_schemas": count,
+                "ttl_seconds": self._ttl,
+                "key_prefix": self._key_prefix
+            }
+        except Exception as e:
+            logger.warning(f"Error getting cache stats: {e}")
+            return {"available": False, "error": str(e)}
+
+
+_redis_cache = MCPSchemaRedisCache(ttl_seconds=3600)
+
 class MCPToolWrapper(Tool):
-    def __init__(self, mcp_configs: Optional[List[Dict[str, Any]]] = None):
+    def __init__(self, mcp_configs: Optional[List[Dict[str, Any]]] = None, use_cache: bool = True):
         self.mcp_manager = mcp_service
         self.mcp_configs = mcp_configs or []
         self._initialized = False
         self._schemas: Dict[str, List[ToolSchema]] = {}
         self._dynamic_tools = {}
         self._custom_tools = {}
+        self.use_cache = use_cache
 
         self.connection_manager = MCPConnectionManager()
         self.custom_handler = CustomMCPHandler(self.connection_manager)
@@ -32,23 +132,109 @@ async def _ensure_initialized(self):
             self._initialized = True
 
     async def _initialize_servers(self):
+        start_time = time.time()
+        
         standard_configs = [cfg for cfg in self.mcp_configs if not cfg.get('isCustom', False)]
         custom_configs = [cfg for cfg in self.mcp_configs if cfg.get('isCustom', False)]
 
+        cached_configs = []
+        cached_tools_data = []
+        
+        initialization_tasks = []
+        
         if standard_configs:
-            await self._initialize_standard_servers(standard_configs)
+            for config in standard_configs:
+                if self.use_cache:
+                    cached_data = await _redis_cache.get(config)
+                    if cached_data:
+                        cached_configs.append(config.get('qualifiedName', 'Unknown'))
+                        cached_tools_data.append(cached_data)
+                        continue
+                
+                task = self._initialize_single_standard_server(config)
+                initialization_tasks.append(('standard', config, task))
 
         if custom_configs:
-            await self.custom_handler.initialize_custom_mcps(custom_configs)
+            for config in custom_configs:
+                if self.use_cache:
+                    cached_data = await _redis_cache.get(config)
+                    if cached_data:
+                        cached_configs.append(config.get('name', 'Unknown'))
+                        cached_tools_data.append(cached_data)
+                        continue
+                
+                task = self._initialize_single_custom_mcp(config)
+                initialization_tasks.append(('custom', config, task))
+        
+        if cached_tools_data:
+            logger.info(f"⚡ Loaded {len(cached_configs)} MCP schemas from Redis cache: {', '.join(cached_configs)}")
+            for cached_data in cached_tools_data:
+                try:
+                    if cached_data.get('type') == 'standard':
+                        logger.debug("Standard MCP tools found in cache but require connection to restore")
+                    elif cached_data.get('type') == 'custom':
+                        custom_tools = cached_data.get('tools', {})
+                        if custom_tools:
+                            self.custom_handler.custom_tools.update(custom_tools)
+                            logger.debug(f"Restored {len(custom_tools)} custom tools from cache")
+                except Exception as e:
+                    logger.warning(f"Failed to restore cached tools: {e}")
+        
+        if initialization_tasks:
+            logger.info(f"🚀 Initializing {len(initialization_tasks)} MCP servers in parallel (cache enabled: {self.use_cache})...")
+            
+            tasks = [task for _, _, task in initialization_tasks]
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            successful = 0
+            failed = 0
+            
+            for i, result in enumerate(results):
+                task_type, config, _ = initialization_tasks[i]
+                if isinstance(result, Exception):
+                    failed += 1
+                    config_name = config.get('name', config.get('qualifiedName', 'Unknown'))
+                    logger.error(f"Failed to initialize MCP server '{config_name}': {result}")
+                else:
+                    successful += 1
+                    if self.use_cache and result:
+                        await _redis_cache.set(config, result)
+            
+            elapsed_time = time.time() - start_time
+            logger.info(f"⚡ MCP initialization completed in {elapsed_time:.2f}s - {successful} successful, {failed} failed, {len(cached_configs)} from cache")
+        else:
+            if cached_configs:
+                elapsed_time = time.time() - start_time
+                logger.info(f"⚡ All {len(cached_configs)} MCP schemas loaded from Redis cache in {elapsed_time:.2f}s - instant startup!")
+            else:
+                logger.info("No MCP servers to initialize")
+    
+    async def _initialize_single_standard_server(self, config: Dict[str, Any]):
+        try:
+            logger.debug(f"Connecting to standard MCP server: {config['qualifiedName']}")
+            await self.mcp_manager.connect_server(config)
+            logger.debug(f"✓ Connected to MCP server: {config['qualifiedName']}")
+            
+            tools_info = self.mcp_manager.get_all_tools_openapi()
+            return {'tools': tools_info, 'type': 'standard', 'timestamp': time.time()}
+        except Exception as e:
+            logger.error(f"✗ Failed to connect to MCP server {config['qualifiedName']}: {e}")
+            raise e
+    
+    async def _initialize_single_custom_mcp(self, config: Dict[str, Any]):
+        try:
+            logger.debug(f"Initializing custom MCP: {config.get('name', 'Unknown')}")
+            await self.custom_handler._initialize_single_custom_mcp(config)
+            logger.debug(f"✓ Initialized custom MCP: {config.get('name', 'Unknown')}")
+            
+            custom_tools = self.custom_handler.get_custom_tools()
+            return {'tools': custom_tools, 'type': 'custom', 'timestamp': time.time()}
+        except Exception as e:
+            logger.error(f"✗ Failed to initialize custom MCP {config.get('name', 'Unknown')}: {e}")
+            raise e
 
     async def _initialize_standard_servers(self, standard_configs: List[Dict[str, Any]]):
-        for config in standard_configs:
-            try:
-                logger.info(f"Attempting to connect to MCP server: {config['qualifiedName']}")
-                await self.mcp_manager.connect_server(config)
-                logger.info(f"Successfully connected to MCP server: {config['qualifiedName']}")
-            except Exception as e:
-                logger.error(f"Failed to connect to MCP server {config['qualifiedName']}: {e}")
+        pass
 
     async def _create_dynamic_tools(self):
         try:
@@ -77,7 +263,6 @@ async def _create_dynamic_tools(self):
 
             logger.info(f"Created {len(self._dynamic_tools)} dynamic MCP tool methods")
 
-            # Re-register schemas to pick up the dynamic methods
             self._register_schemas()
             logger.info(f"Re-registered schemas after creating dynamic tools - total: {len(self._schemas)}")
 

backend/agent/tools/mcp_tool_wrapper.py

@@ -12,18 +12,33 @@
 class CustomMCPHandler:
     def __init__(self, connection_manager: MCPConnectionManager):
         self.connection_manager = connection_manager
-        self.custom_tools: Dict[str, Dict[str, Any]] = {}
+        self.custom_tools = {}
 
     async def initialize_custom_mcps(self, custom_configs: List[Dict[str, Any]]) -> Dict[str, Dict[str, Any]]:
+        initialization_tasks = []
+        
         for config in custom_configs:
-            try:
-                await self._initialize_single_custom_mcp(config)
-            except Exception as e:
-                logger.error(f"Failed to initialize custom MCP {config.get('name', 'Unknown')}: {e}")
-                continue
+            task = self._initialize_single_custom_mcp_safe(config)
+            initialization_tasks.append(task)
+        
+        if initialization_tasks:
+            logger.info(f"Initializing {len(initialization_tasks)} custom MCPs in parallel...")
+            results = await asyncio.gather(*initialization_tasks, return_exceptions=True)
+            
+            for i, result in enumerate(results):
+                if isinstance(result, Exception):
+                    config_name = custom_configs[i].get('name', 'Unknown')
+                    logger.error(f"Failed to initialize custom MCP {config_name}: {result}")
 
         return self.custom_tools
 
+    async def _initialize_single_custom_mcp_safe(self, config: Dict[str, Any]):
+        try:
+            await self._initialize_single_custom_mcp(config)
+        except Exception as e:
+            logger.error(f"Failed to initialize custom MCP {config.get('name', 'Unknown')}: {e}")
+            return e
+    
     async def _initialize_single_custom_mcp(self, config: Dict[str, Any]):
         custom_type = config.get('customType', 'sse')
         server_config = config.get('config', {})