diff --git a/src/cai/repl/commands/__init__.py b/src/cai/repl/commands/__init__.py index 1ac26681..02fe107f 100644 --- a/src/cai/repl/commands/__init__.py +++ b/src/cai/repl/commands/__init__.py @@ -40,7 +40,8 @@ from cai.repl.commands import ( # pylint: disable=import-error,unused-import,li workspace, virtualization, load, - parallel # Add the new parallel command + parallel, # Add the new parallel command + mcp # Add the MCP command ) # Define helper functions diff --git a/src/cai/repl/commands/mcp.py b/src/cai/repl/commands/mcp.py new file mode 100644 index 00000000..02ee4a3e --- /dev/null +++ b/src/cai/repl/commands/mcp.py @@ -0,0 +1,1066 @@ +""" +MCP (Model Context Protocol) command for CAI CLI + +Provides commands for managing MCP servers and integrating their tools +with agents. + +USAGE EXAMPLES: +============== + +1. Load an SSE (Server-Sent Events) MCP server: + /mcp load http://localhost:9876/sse burp + +2. Load a STDIO MCP server: + /mcp load stdio myserver python mcp_server.py + /mcp load stdio myserver node server.js --port 8080 + +3. List all active MCP connections: + /mcp list + +4. Add MCP tools to an agent: + /mcp add burp redteam_agent # Add by agent name + /mcp add burp 13 # Add by agent number + +5. List tools from a specific server: + /mcp tools burp + +6. Check server connection status: + /mcp status + +7. Remove a server connection: + /mcp remove burp + +8. Show help: + /mcp help + +NOTES: +====== +- Each tool invocation creates a fresh connection to ensure reliability +- SSE servers may show async generator warnings on cleanup (this is normal) +- Use /mcp status to check and reconnect servers if needed +- Tools are added directly to agent.tools for seamless integration + +QUICK START: +=========== +1. Start your MCP server (e.g., Burp Suite MCP extension) +2. Load it: /mcp load http://localhost:9876/sse burp +3. Add to agent: /mcp add burp your_agent +4. Use the tools through the agent +""" + +# Standard library imports +import asyncio +import os +import atexit +import warnings +from typing import Dict, List, Optional + +# Third-party imports +from rich.console import Console +from rich.table import Table +from rich.markdown import Markdown + +# Local imports +from cai.agents import get_available_agents, get_agent_by_name +from cai.repl.commands.base import Command, register_command +from cai.sdk.agents import Agent +from cai.sdk.agents.mcp import ( + MCPServer, + MCPServerSse, + MCPServerSseParams, + MCPServerStdio, + MCPServerStdioParams, + MCPUtil +) +from cai.sdk.agents.tool import FunctionTool +from cai.sdk.agents.tracing import mcp_tools_span +import functools + +console = Console() + +# Global registry for persistent MCP connections +_GLOBAL_MCP_SERVERS: Dict[str, MCPServer] = {} + +# Custom MCPUtil that uses global registry +class GlobalMCPUtil(MCPUtil): + """Custom MCP utility that uses global server registry""" + + @classmethod + def to_function_tool(cls, tool, server_name: str) -> FunctionTool: + """Convert an MCP tool to a CAI function tool using server name instead of object.""" + + # Store the server configuration instead of the server object + server = _GLOBAL_MCP_SERVERS.get(server_name) + if not server: + raise ValueError(f"Server {server_name} not found in registry") + + # Capture server configuration + server_config = { + 'name': server_name, + 'type': type(server).__name__, + 'tool_name': tool.name, + 'tool_schema': tool.inputSchema, + 'tool_description': tool.description + } + + # For SSE servers, capture the URL + if isinstance(server, MCPServerSse): + server_config['url'] = server.params.get('url') + server_config['headers'] = server.params.get('headers') + server_config['timeout'] = server.params.get('timeout', 5) + server_config['sse_read_timeout'] = server.params.get('sse_read_timeout', 60 * 5) + # For STDIO servers, capture the command + elif isinstance(server, MCPServerStdio): + server_config['command'] = server.params.command + server_config['args'] = server.params.args + server_config['env'] = server.params.get('env') + server_config['cwd'] = server.params.get('cwd') + server_config['encoding'] = server.params.get('encoding', 'utf-8') + server_config['encoding_error_handler'] = server.params.get('encoding_error_handler', 'strict') + + # Create a custom invoke function that creates a new connection each time + async def invoke_with_fresh_connection(config, context, input_json): + """Custom invoke function that creates a fresh connection for each invocation""" + import json + import asyncio + import warnings + import sys + from contextlib import suppress + from cai.sdk.agents.exceptions import ModelBehaviorError, AgentsException + from cai.sdk.agents.mcp import MCPServerSse, MCPServerStdio + + # Parse JSON input + try: + json_data = json.loads(input_json) if input_json else {} + except Exception as e: + raise ModelBehaviorError( + f"Invalid JSON input for tool {config['tool_name']}: {input_json}" + ) from e + + # Create a fresh server connection with timeout + server = None + result = None + + # Suppress warnings about async generator cleanup + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=RuntimeWarning) + warnings.filterwarnings("ignore", message=".*asynchronous generator.*") + + try: + if config['type'] == 'MCPServerSse': + # Create new SSE server + params = { + 'url': config['url'], + 'headers': config.get('headers'), + 'timeout': config.get('timeout', 5), + 'sse_read_timeout': config.get('sse_read_timeout', 60 * 5) + } + # Remove None values + params = {k: v for k, v in params.items() if v is not None} + + server = MCPServerSse( + params, + name=config['name'], + cache_tools_list=False # Don't cache since it's temporary + ) + elif config['type'] == 'MCPServerStdio': + # Create new STDIO server + params = { + 'command': config['command'], + 'args': config.get('args', []), + 'env': config.get('env'), + 'cwd': config.get('cwd'), + 'encoding': config.get('encoding', 'utf-8'), + 'encoding_error_handler': config.get('encoding_error_handler', 'strict') + } + # Remove None values + params = {k: v for k, v in params.items() if v is not None} + + server = MCPServerStdio( + params, + name=config['name'], + cache_tools_list=False + ) + else: + raise AgentsException(f"Unknown server type: {config['type']}") + + # Connect to the server with timeout + try: + await asyncio.wait_for(server.connect(), timeout=10.0) + except asyncio.TimeoutError: + raise AgentsException( + f"Timeout connecting to MCP server for tool {config['tool_name']}. " + f"The server may be down or not responding." + ) + + # Call the tool with timeout + try: + result = await asyncio.wait_for( + server.call_tool(config['tool_name'], json_data), + timeout=30.0 + ) + except asyncio.TimeoutError: + raise AgentsException( + f"Timeout calling MCP tool {config['tool_name']}. " + f"The tool took too long to respond." + ) + + except Exception as e: + raise AgentsException( + f"Error invoking MCP tool {config['tool_name']}: {type(e).__name__}: {str(e)}" + ) from e + + finally: + # Cleanup the server - handle SSE cleanup issues + if server: + if config['type'] == 'MCPServerSse': + # For SSE servers, suppress cleanup errors as they're expected + try: + # Don't wait too long for SSE cleanup + await asyncio.wait_for(server.cleanup(), timeout=1.0) + except (asyncio.TimeoutError, RuntimeError, Exception): + # Expected for SSE connections + pass + else: + # For STDIO servers, cleanup normally + try: + await asyncio.wait_for(server.cleanup(), timeout=5.0) + except (asyncio.TimeoutError, Exception): + pass + + # Format the result + if not result: + raise AgentsException(f"No result returned from MCP tool {config['tool_name']}") + + # Convert result to string format + if len(result.content) == 1: + tool_output = result.content[0].model_dump_json() + elif len(result.content) > 1: + tool_output = json.dumps([item.model_dump() for item in result.content]) + else: + tool_output = "Error running tool." + + # Handle tracing if needed + from cai.sdk.agents.tracing import get_current_span, FunctionSpanData + current_span = get_current_span() + if current_span: + if isinstance(current_span.span_data, FunctionSpanData): + current_span.span_data.output = tool_output + current_span.span_data.mcp_data = { + "server": config['name'], + } + + return tool_output + + # Use functools.partial to bind the server config + invoke_func = functools.partial( + invoke_with_fresh_connection, + server_config + ) + + return FunctionTool( + name=tool.name, + description=tool.description or "", + params_json_schema=tool.inputSchema, + on_invoke_tool=invoke_func, + strict_json_schema=False, + ) + +def cleanup_mcp_servers(): + """Cleanup all MCP servers on exit""" + try: + if _GLOBAL_MCP_SERVERS: + # Create new event loop for cleanup if needed + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + async def cleanup_all(): + tasks = [] + for name, server in _GLOBAL_MCP_SERVERS.items(): + try: + tasks.append(server.cleanup()) + except Exception: + pass + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + + loop.run_until_complete(cleanup_all()) + loop.close() + except Exception: + pass + +# Register cleanup on exit +atexit.register(cleanup_mcp_servers) + + +class MCPCommand(Command): + """Command for managing MCP servers and their integration with agents.""" + + def __init__(self): + """Initialize the MCP command.""" + super().__init__( + name="/mcp", + description="Manage MCP servers and add their tools to agents", + aliases=["/m"] + ) + + # Add subcommands manually + self._subcommands = { + "load": "Load an MCP server (SSE or stdio)", + "list": "List active MCP connections", + "add": "Add MCP tools to an agent", + "remove": "Remove an MCP server connection", + "tools": "List tools from an MCP server", + "status": "Check MCP server connection status", + "help": "Show MCP command usage" + } + + def get_subcommands(self) -> List[str]: + """Get list of subcommand names. + + Returns: + List of subcommand names + """ + return list(self._subcommands.keys()) + + def get_subcommand_description(self, subcommand: str) -> str: + """Get description for a subcommand. + + Args: + subcommand: Name of the subcommand + + Returns: + Description of the subcommand + """ + return self._subcommands.get(subcommand, "") + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the MCP command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully + """ + if not args: + return self.handle_list(args) + + subcommand = args[0] + if subcommand in self._subcommands: + handler = getattr(self, f"handle_{subcommand}", None) + if handler: + try: + return handler(args[1:] if len(args) > 1 else None) + except Exception as e: + console.print(f"[red]Error executing command: {e}[/red]") + return False + + console.print(f"[red]Unknown subcommand: {subcommand}[/red]") + self.show_usage() + return False + + def show_usage(self): + """Show usage information for the MCP command.""" + usage_text = """ +# MCP (Model Context Protocol) Command Usage + +The MCP command allows you to manage Model Context Protocol servers and integrate their tools with CAI agents. + +## Commands: + +### Load an MCP Server + +**SSE (Server-Sent Events) Server:** +``` +/mcp load +``` +Example: `/mcp load http://localhost:9876/sse burp` + +**STDIO Server:** +``` +/mcp load stdio [args...] +``` +Example: `/mcp load stdio myserver python mcp_server.py` + +### List Active Connections +``` +/mcp list +``` + +### Add Tools to an Agent +``` +/mcp add +``` +Example: `/mcp add burp redteam_agent` +Example: `/mcp add burp 13` + +### List Tools from a Server +``` +/mcp tools +``` + +### Check Server Status +``` +/mcp status +``` + +### Remove a Server +``` +/mcp remove +``` + +### Show Help +``` +/mcp help +``` + +## Quick Start: + +1. Load an MCP server: + `/mcp load http://localhost:9876/sse burp` + +2. List available tools: + `/mcp tools burp` + +3. Add tools to an agent: + `/mcp add burp redteam_agent` + +4. Switch to the agent and use the tools: + `/agent redteam_agent` +""" + console.print(Markdown(usage_text)) + + def handle_help(self, args: Optional[List[str]] = None) -> bool: + """Handle /mcp help command. + + Args: + args: Optional list of command arguments (not used) + + Returns: + True + """ + self.show_usage() + return True + + def _run_async(self, coro): + """Run async code properly in the CLI context. + + Args: + coro: The coroutine to run + + Returns: + The result of the coroutine + """ + try: + # Try to get existing loop + loop = asyncio.get_running_loop() + # If we're in a loop, we need to use a different approach + import concurrent.futures + import threading + import sys + from io import StringIO + + def run_in_thread(): + # Suppress stderr in the thread too + original_stderr = sys.stderr + try: + sys.stderr = StringIO() + new_loop = asyncio.new_event_loop() + asyncio.set_event_loop(new_loop) + try: + return new_loop.run_until_complete(coro) + finally: + new_loop.close() + finally: + sys.stderr = original_stderr + + with concurrent.futures.ThreadPoolExecutor() as executor: + future = executor.submit(run_in_thread) + return future.result(timeout=30) + + except RuntimeError: + # No running loop, we can use asyncio.run + import sys + from io import StringIO + + # Suppress stderr during asyncio.run + original_stderr = sys.stderr + try: + sys.stderr = StringIO() + return asyncio.run(coro) + finally: + sys.stderr = original_stderr + + def handle_load(self, args: Optional[List[str]] = None) -> bool: + """Handle /mcp load command. + + Usage: + /mcp load - Load SSE server + /mcp load stdio [args...] - Load stdio server + + Args: + args: List of command arguments + + Returns: + True if successful + """ + if not args or len(args) < 2: + console.print("[red]Error: Invalid arguments[/red]") + console.print("Usage:") + console.print(" /mcp load - For SSE servers") + console.print(" /mcp load stdio [args...]") + return False + + # Check if it's a stdio server + if args[0] == "stdio": + if len(args) < 3: + console.print( + "[red]Error: stdio requires name and command[/red]" + ) + return False + + name = args[1] + command = args[2] + cmd_args = args[3:] if len(args) > 3 else [] + + return self._load_stdio_server(name, command, cmd_args) + else: + # SSE server + url = args[0] + name = args[1] + + return self._load_sse_server(url, name) + + def _load_sse_server(self, url: str, name: str) -> bool: + """Load an SSE MCP server. + + Args: + url: URL of the SSE server + name: Name to identify the server + + Returns: + True if successful + """ + if name in _GLOBAL_MCP_SERVERS: + console.print( + f"[yellow]Server '{name}' already exists. " + f"Remove it first.[/yellow]" + ) + return False + + console.print(f"Connecting to SSE server at {url}...") + + async def connect_and_test(): + params: MCPServerSseParams = {"url": url} + server = MCPServerSse( + params, + name=name, + cache_tools_list=True + ) + + # Connect to the server + await server.connect() + + # Test by listing tools + tools = await server.list_tools() + + return server, tools + + try: + # Suppress all stderr output during SSE connection + import sys + from io import StringIO + + # Save the original stderr + original_stderr = sys.stderr + + try: + # Redirect stderr to null + sys.stderr = StringIO() + + # Also suppress warnings + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=RuntimeWarning) + warnings.filterwarnings("ignore", message=".*asynchronous generator.*") + warnings.filterwarnings("ignore", message=".*cancel scope.*") + warnings.filterwarnings("ignore", message=".*didn't stop after athrow.*") + + server, tools = self._run_async(connect_and_test()) + finally: + # Always restore stderr + sys.stderr = original_stderr + + # Store the server globally + _GLOBAL_MCP_SERVERS[name] = server + + console.print( + f"[green]✓ Connected to SSE server '{name}' " + f"at {url}[/green]" + ) + console.print(f"Available tools: {len(tools)}") + + # Show some tool names if available + if tools: + tool_names = [tool.name for tool in tools[:5]] + if len(tools) > 5: + tool_names.append(f"... and {len(tools) - 5} more") + console.print(f"Tools: {', '.join(tool_names)}") + + return True + + except Exception as e: + console.print(f"[red]Error connecting to server: {e}[/red]") + # Clean up if connection failed + if name in _GLOBAL_MCP_SERVERS: + del _GLOBAL_MCP_SERVERS[name] + return False + + def _load_stdio_server( + self, name: str, command: str, cmd_args: List[str] + ) -> bool: + """Load a stdio MCP server. + + Args: + name: Name to identify the server + command: Command to execute + cmd_args: Arguments for the command + + Returns: + True if successful + """ + if name in _GLOBAL_MCP_SERVERS: + console.print( + f"[yellow]Server '{name}' already exists. " + f"Remove it first.[/yellow]" + ) + return False + + console.print( + f"Starting stdio server '{name}' with command: " + f"{command} {' '.join(cmd_args)}" + ) + + async def connect_and_test(): + params: MCPServerStdioParams = { + "command": command, + "args": cmd_args + } + server = MCPServerStdio( + params, + name=name, + cache_tools_list=True + ) + + # Connect to the server + await server.connect() + + # Test by listing tools + tools = await server.list_tools() + + return server, tools + + try: + server, tools = self._run_async(connect_and_test()) + + # Store the server globally + _GLOBAL_MCP_SERVERS[name] = server + + console.print( + f"[green]✓ Started stdio server '{name}'[/green]" + ) + console.print(f"Available tools: {len(tools)}") + + # Show some tool names if available + if tools: + tool_names = [tool.name for tool in tools[:5]] + if len(tools) > 5: + tool_names.append(f"... and {len(tools) - 5} more") + console.print(f"Tools: {', '.join(tool_names)}") + + return True + + except Exception as e: + console.print(f"[red]Error starting server: {e}[/red]") + # Clean up if connection failed + if name in _GLOBAL_MCP_SERVERS: + del _GLOBAL_MCP_SERVERS[name] + return False + + def handle_list(self, args: Optional[List[str]] = None) -> bool: + """Handle /mcp list command. + + Args: + args: Optional list of command arguments (not used) + + Returns: + True + """ + if not _GLOBAL_MCP_SERVERS: + console.print("[yellow]No active MCP connections[/yellow]") + console.print("\nUse `/mcp help` to see how to load servers.") + return True + + table = Table(title="Active MCP Connections") + table.add_column("Name", style="cyan") + table.add_column("Type", style="magenta") + table.add_column("Details", style="green") + table.add_column("Tools", style="yellow") + + for name, server in _GLOBAL_MCP_SERVERS.items(): + server_type = type(server).__name__.replace("MCPServer", "") + + # Get server details + if isinstance(server, MCPServerSse): + details = server.params.get("url", "N/A") + elif isinstance(server, MCPServerStdio): + cmd = server.params.command + args = " ".join(server.params.args) + details = f"{cmd} {args}".strip() + else: + details = "Unknown" + + # Get tool count + try: + async def get_tools(): + return await server.list_tools() + + tools = self._run_async(get_tools()) + tool_count = str(len(tools)) + except Exception: + tool_count = "Error" + + table.add_row(name, server_type, details, tool_count) + + console.print(table) + return True + + def handle_add(self, args: Optional[List[str]] = None) -> bool: + """Handle /mcp add command. + + Usage: /mcp add + + Args: + args: List of command arguments + + Returns: + True if successful + """ + if not args or len(args) < 2: + console.print("[red]Error: Invalid arguments[/red]") + console.print("Usage: /mcp add ") + return False + + server_name = args[0] + agent_identifier = args[1] + + # Check if server exists + if server_name not in _GLOBAL_MCP_SERVERS: + console.print(f"[red]Error: Server '{server_name}' not found[/red]") + console.print("Use /mcp list to see active servers") + return False + + # Get the agent + try: + agent = get_agent_by_name(agent_identifier) + agent_display_name = getattr(agent, "name", agent_identifier) + except ValueError: + # Try by index + try: + agents = get_available_agents() + agent_list = list(agents.items()) + + if agent_identifier.isdigit(): + idx = int(agent_identifier) + if 1 <= idx <= len(agent_list): + agent_key, agent = agent_list[idx - 1] + agent_display_name = getattr(agent, "name", agent_key) + else: + raise ValueError("Invalid index") + else: + raise ValueError("Not found") + except Exception: + console.print( + f"[red]Error: Agent '{agent_identifier}' not found[/red]" + ) + return False + + # Add the MCP server to the agent + server = _GLOBAL_MCP_SERVERS[server_name] + + console.print( + f"Adding tools from MCP server '{server_name}' to agent " + f"'{agent_display_name}'..." + ) + + # Validate the server connection before adding + try: + async def validate_connection(): + try: + # Try to list tools to validate connection + tools = await server.list_tools() + return tools + except Exception as e: + console.print(f"[yellow]Warning: Server connection may be lost, attempting to reconnect...[/yellow]") + # Try to reconnect + await server.connect() + tools = await server.list_tools() + console.print(f"[green]✓ Reconnected to server '{server_name}'[/green]") + return tools + + # Validate the connection and get tools + mcp_tools = self._run_async(validate_connection()) + + except Exception as e: + console.print(f"[red]Error: Cannot connect to server '{server_name}': {e}[/red]") + console.print("Try removing and reloading the server.") + return False + + # Get and display the tools + try: + # Create function tools using GlobalMCPUtil + tools = [] + for mcp_tool in mcp_tools: + # Use GlobalMCPUtil to create tools that use the global registry + function_tool = GlobalMCPUtil.to_function_tool(mcp_tool, server_name) + tools.append(function_tool) + + # Display tools table + table = Table(title=f"Adding tools to {agent_display_name}") + table.add_column("Tool", style="cyan") + table.add_column("Status", style="green") + table.add_column("Details", style="yellow") + + for tool in tools: + table.add_row( + tool.name, + "Added", + f"Available as: {tool.name}" + ) + + console.print(table) + + # Add tools directly to agent.tools + if not hasattr(agent, 'tools'): + agent.tools = [] + + # Remove any existing tools with the same names to avoid duplicates + existing_tool_names = {t.name for t in tools} + agent.tools = [t for t in agent.tools if t.name not in existing_tool_names] + + # Add the new tools + agent.tools.extend(tools) + + console.print( + f"[green]Added {len(tools)} tools from server " + f"'{server_name}' to agent '{agent_display_name}'.[/green]" + ) + + # Test that the tools are accessible + async def test_agent_tools(): + # Get all tools including MCP tools + all_regular_tools = agent.tools if hasattr(agent, 'tools') else [] + all_mcp_tools = await agent.get_mcp_tools() if hasattr(agent, 'mcp_servers') and agent.mcp_servers else [] + return all_regular_tools + all_mcp_tools + + all_tools = self._run_async(test_agent_tools()) + + # Count different types of tools + mcp_server_tools_count = len([t for t in agent.mcp_servers if hasattr(agent, 'mcp_servers')]) if hasattr(agent, 'mcp_servers') else 0 + regular_tools_count = len(agent.tools) if hasattr(agent, 'tools') else 0 + + console.print( + f"[blue]Agent now has {regular_tools_count} tools total[/blue]" + ) + + # Test a simple tool invocation to make sure everything works + console.print("[cyan]Testing MCP tool connectivity...[/cyan]") + try: + if tools: + console.print(f"[green]✓ MCP tools are ready for use![/green]") + else: + console.print("[yellow]Warning: No tools available from server[/yellow]") + except Exception as e: + console.print(f"[yellow]Warning: Tool connectivity test failed: {e}[/yellow]") + + return True + + except Exception as e: + console.print(f"[red]Error adding tools: {e}[/red]") + return False + + def handle_remove(self, args: Optional[List[str]] = None) -> bool: + """Handle /mcp remove command. + + Args: + args: List of command arguments + + Returns: + True if successful + """ + if not args: + console.print("[red]Error: No server name specified[/red]") + console.print("Usage: /mcp remove ") + return False + + server_name = args[0] + + if server_name not in _GLOBAL_MCP_SERVERS: + console.print(f"[red]Error: Server '{server_name}' not found[/red]") + return False + + # Cleanup the server + server = _GLOBAL_MCP_SERVERS[server_name] + + try: + async def cleanup_server(): + await server.cleanup() + + self._run_async(cleanup_server()) + del _GLOBAL_MCP_SERVERS[server_name] + console.print( + f"[green]✓ Removed MCP server '{server_name}'[/green]" + ) + return True + except Exception as e: + console.print(f"[red]Error removing server: {e}[/red]") + # Remove from list anyway + if server_name in _GLOBAL_MCP_SERVERS: + del _GLOBAL_MCP_SERVERS[server_name] + return False + + def handle_status(self, args: Optional[List[str]] = None) -> bool: + """Handle /mcp status command. + + Args: + args: Optional list of command arguments + + Returns: + True if successful + """ + if not _GLOBAL_MCP_SERVERS: + console.print("[yellow]No active MCP connections[/yellow]") + return True + + console.print("[cyan]Checking MCP server connections...[/cyan]") + + table = Table(title="MCP Server Status") + table.add_column("Name", style="cyan") + table.add_column("Type", style="magenta") + table.add_column("Status", style="bold") + table.add_column("Tools", style="yellow") + table.add_column("Details", style="dim") + + healthy_count = 0 + + for name, server in _GLOBAL_MCP_SERVERS.items(): + server_type = type(server).__name__.replace("MCPServer", "") + + # Test server connection + try: + async def test_connection(): + tools = await server.list_tools() + return len(tools), None + + tools_count, error = self._run_async(test_connection()) + status = "[green]✓ Healthy[/green]" + tools_str = str(tools_count) + details = "Connection active" + healthy_count += 1 + + except Exception as e: + status = "[red]✗ Error[/red]" + tools_str = "N/A" + details = f"Error: {str(e)[:50]}..." + + # Try to reconnect + try: + console.print(f"[yellow]Attempting to reconnect to '{name}'...[/yellow]") + + async def reconnect(): + await server.connect() + tools = await server.list_tools() + return len(tools) + + tools_count = self._run_async(reconnect()) + status = "[green]✓ Reconnected[/green]" + tools_str = str(tools_count) + details = "Reconnected successfully" + healthy_count += 1 + + except Exception as reconnect_error: + status = "[red]✗ Failed[/red]" + details = f"Reconnect failed: {str(reconnect_error)[:30]}..." + + table.add_row(name, server_type, status, tools_str, details) + + console.print(table) + + # Summary + total_servers = len(_GLOBAL_MCP_SERVERS) + if healthy_count == total_servers: + console.print(f"[green]✓ All {total_servers} MCP servers are healthy[/green]") + else: + failed_count = total_servers - healthy_count + console.print(f"[yellow]⚠ {healthy_count}/{total_servers} servers healthy, {failed_count} failed[/yellow]") + + return True + + def handle_tools(self, args: Optional[List[str]] = None) -> bool: + """Handle /mcp tools command. + + Args: + args: List of command arguments + + Returns: + True if successful + """ + if not args: + console.print("[red]Error: No server name specified[/red]") + console.print("Usage: /mcp tools ") + return False + + server_name = args[0] + + if server_name not in _GLOBAL_MCP_SERVERS: + console.print(f"[red]Error: Server '{server_name}' not found[/red]") + return False + + server = _GLOBAL_MCP_SERVERS[server_name] + + try: + async def get_tools(): + return await server.list_tools() + + tools = self._run_async(get_tools()) + + if not tools: + console.print( + f"[yellow]No tools available from '{server_name}'[/yellow]" + ) + return True + + table = Table(title=f"Tools from '{server_name}'") + table.add_column("#", style="dim") + table.add_column("Name", style="cyan") + table.add_column("Description", style="green") + + for idx, tool in enumerate(tools, 1): + description = tool.description or "No description" + if len(description) > 60: + description = description[:57] + "..." + table.add_row(str(idx), tool.name, description) + + console.print(table) + return True + + except Exception as e: + console.print(f"[red]Error listing tools: {e}[/red]") + return False + + +# Register the command +register_command(MCPCommand()) \ No newline at end of file diff --git a/src/cai/sdk/agents/mcp/util.py b/src/cai/sdk/agents/mcp/util.py index 496abace..8a28028f 100644 --- a/src/cai/sdk/agents/mcp/util.py +++ b/src/cai/sdk/agents/mcp/util.py @@ -80,10 +80,52 @@ class MCPUtil: logger.debug(f"Invoking MCP tool {tool.name} with input {input_json}") try: + # Check if server session is still valid + if not hasattr(server, 'session') or server.session is None: + logger.warning(f"MCP server session not found for tool {tool.name}, attempting to reconnect...") + # Try to reconnect + try: + await server.connect() + logger.info(f"Successfully reconnected to MCP server for tool {tool.name}") + except Exception as reconnect_error: + logger.error(f"Failed to reconnect to MCP server: {reconnect_error}") + raise AgentsException( + f"MCP server connection lost for tool {tool.name}. " + f"Please remove and re-add the MCP server. " + f"Reconnection error: {str(reconnect_error)}" + ) from reconnect_error + + # Now try to call the tool result = await server.call_tool(tool.name, json_data) + + except AttributeError as ae: + # This often happens when the server object is not properly initialized + logger.error(f"MCP server not properly initialized for tool {tool.name}: {ae}") + logger.error(f"Server type: {type(server)}, has session: {hasattr(server, 'session')}") + raise AgentsException( + f"MCP server not properly initialized for tool {tool.name}. " + f"The server connection may have been lost. " + f"AttributeError: {str(ae)}\n" + f"Try: /mcp remove then /mcp load ... to reconnect." + ) from ae except Exception as e: - logger.error(f"Error invoking MCP tool {tool.name}: {e}") - raise AgentsException(f"Error invoking MCP tool {tool.name}: {e}") from e + # Log the full exception details + logger.error(f"Error invoking MCP tool {tool.name}: {type(e).__name__}: {str(e)}") + logger.error(f"Full exception details: {repr(e)}") + + # Check if it's a connection issue + error_str = str(e).lower() + if "session" in error_str or "connection" in error_str or "closed" in error_str: + raise AgentsException( + f"MCP server connection error for tool {tool.name}. " + f"Error: {type(e).__name__}: {str(e)}\n" + f"Use '/mcp status' to check server health and '/mcp remove' + '/mcp load' to reconnect." + ) from e + else: + # For other errors, include the full error details + raise AgentsException( + f"Error invoking MCP tool {tool.name}: {type(e).__name__}: {str(e)}" + ) from e if _debug.DONT_LOG_TOOL_DATA: logger.debug(f"MCP tool {tool.name} completed.") diff --git a/src/cai/util.py b/src/cai/util.py index ffa87576..45ebf5f4 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -645,9 +645,65 @@ def visualize_agent_graph(start_agent): # Add tools tools_node = node.add("[yellow]Tools[/yellow]") - for tool in getattr(agent, "tools", []): + + # Get regular tools and MCP tools separately + regular_tools = getattr(agent, "tools", []) + mcp_tools = [] + + try: + # Try to get MCP tools specifically + import asyncio + + async def get_mcp_tools_async(): + return await agent.get_mcp_tools() + + # Run async function to get MCP tools + try: + loop = asyncio.get_running_loop() + # If we're in a loop, we need to use a different approach + import concurrent.futures + + def run_in_thread(): + new_loop = asyncio.new_event_loop() + asyncio.set_event_loop(new_loop) + try: + return new_loop.run_until_complete(get_mcp_tools_async()) + finally: + new_loop.close() + + with concurrent.futures.ThreadPoolExecutor() as executor: + future = executor.submit(run_in_thread) + mcp_tools = future.result(timeout=10) + + except RuntimeError: + # No running loop, we can use asyncio.run + mcp_tools = asyncio.run(get_mcp_tools_async()) + + except Exception: + # If MCP tools fetch fails, mcp_tools remains empty list + mcp_tools = [] + + # Show regular tools first + for tool in regular_tools: tool_name = getattr(tool, "name", None) or getattr(tool, "__name__", "") tools_node.add(f"[blue]{tool_name}[/blue]") + + # Show MCP tools with a different color/prefix + if mcp_tools: + for tool in mcp_tools: + tool_name = getattr(tool, "name", None) or getattr(tool, "__name__", "") + tools_node.add(f"[magenta]🔌 {tool_name}[/magenta]") + + # Add a summary line if we have both types + if regular_tools and mcp_tools: + summary_text = f"[dim]({len(regular_tools)} regular, {len(mcp_tools)} MCP tools)[/dim]" + tools_node.add(summary_text) + elif mcp_tools and not regular_tools: + summary_text = f"[dim]({len(mcp_tools)} MCP tools)[/dim]" + tools_node.add(summary_text) + elif regular_tools and not mcp_tools: + summary_text = f"[dim]({len(regular_tools)} regular tools)[/dim]" + tools_node.add(summary_text) # Add handoffs transfers_node = node.add("[magenta]Handoffs[/magenta]")