From 4cea6aa27748a171741837edd4cbf8db74b15a44 Mon Sep 17 00:00:00 2001 From: luijait Date: Thu, 29 May 2025 11:28:44 +0200 Subject: [PATCH 1/6] Add artifacts to .gitignore --- .gitignore | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index e881af55..7a69a7e8 100644 --- a/.gitignore +++ b/.gitignore @@ -155,4 +155,7 @@ workspaces/ benchmarks/outputs # other -nohup.out \ No newline at end of file +nohup.out +venv* +caiextensions-* +pentestperf From e22c79092f73402610ee3136af405e1f56589f59 Mon Sep 17 00:00:00 2001 From: luijait Date: Thu, 29 May 2025 12:42:02 +0200 Subject: [PATCH 2/6] Solve duplicate agents in handoffs, add retesting agent, add bug bounty triage handoff (bug bounty agent + retesting agent), solve red team thought pattern, rename pattenrs Signed-off-by: luijait --- src/cai/agents/patterns/bb_triage.py | 61 ++++++++++++++++++++ src/cai/agents/patterns/red_team.py | 31 ++++++---- src/cai/agents/retester.py | 46 +++++++++++++++ src/cai/prompts/system_triage_agent.md | 78 ++++++++++++++++++++++++++ 4 files changed, 206 insertions(+), 10 deletions(-) create mode 100644 src/cai/agents/patterns/bb_triage.py create mode 100644 src/cai/agents/retester.py create mode 100644 src/cai/prompts/system_triage_agent.md diff --git a/src/cai/agents/patterns/bb_triage.py b/src/cai/agents/patterns/bb_triage.py new file mode 100644 index 00000000..96e075e8 --- /dev/null +++ b/src/cai/agents/patterns/bb_triage.py @@ -0,0 +1,61 @@ +""" +Implementation of a Cyclic Swarm Pattern for Bug Bounty Triage Operations + +This module establishes a coordinated multi-agent system where specialized agents +collaborate on vulnerability discovery and verification tasks. The pattern +implements a directed graph of agent relationships, where each agent can transfer +context (message history) to another agent through handoff functions, creating a +complete communication network for comprehensive bug bounty and triage analysis. +""" +from cai.agents.retester import retester_agent +from cai.agents.bug_bounter import bug_bounter_agent +from cai.sdk.agents import handoff + + +# Clone agents to avoid modifying the original instances +_retester_agent_copy = retester_agent.clone() +_bug_bounter_agent_copy = bug_bounter_agent.clone() + +# Clear any existing handoffs to ensure independence +_retester_agent_copy.handoffs = [] +_bug_bounter_agent_copy.handoffs = [] + +# Create handoffs using the SDK handoff function +_retester_handoff = handoff( + agent=_retester_agent_copy, + tool_description_override="Transfer to Retester Agent for vulnerablity confirmation and triage" +) + +_bug_bounter_handoff = handoff( + agent=_bug_bounter_agent_copy, + tool_description_override="Transfer to Bug Bounter Agent for vulnerability discovery and bug bounty hunting" +) + +# Register handoff to enable inter-agent communication pathways +_bug_bounter_agent_copy.handoffs.append(_retester_handoff) +_retester_agent_copy.handoffs.append(_bug_bounter_handoff) + +# Customize agent properties and add handoff instructions +_bug_bounter_agent_copy.name = "Bug bounty Triage Agent" +_bug_bounter_agent_copy.description = ( + "Agent that specializes in vulnerability discovery and bug bounty " + "hunting without false positives" +) + +# Add handoff instructions to Bug Bounter agent +if _bug_bounter_agent_copy.instructions: + _bug_bounter_agent_copy.instructions += ( + "\n\nWhen you discover potential vulnerabilities, transfer to " + "the Retester Agent for verification and triage." + ) + +# Add handoff instructions to Retester agent +if _retester_agent_copy.instructions: + _retester_agent_copy.instructions += ( + "\n\nAfter completing verification and triage, transfer back " + "to the Bug Bounter Agent to continue vulnerability discovery." + ) + +# Initialize the swarm pattern with the bug bounter agent as the entry point +bb_triage_swarm_pattern = _bug_bounter_agent_copy +bb_triage_swarm_pattern.pattern = "swarm" diff --git a/src/cai/agents/patterns/red_team.py b/src/cai/agents/patterns/red_team.py index bf950f67..86dd1e6f 100644 --- a/src/cai/agents/patterns/red_team.py +++ b/src/cai/agents/patterns/red_team.py @@ -13,27 +13,38 @@ from cai.agents.mail import dns_smtp_agent from cai.sdk.agents import handoff +# Clone agents to avoid modifying the original instances +_redteam_agent_copy = redteam_agent.clone() +_thought_agent_copy = thought_agent.clone() +_dns_smtp_agent_copy = dns_smtp_agent.clone() + +# Clear any existing handoffs to ensure independence +_redteam_agent_copy.handoffs = [] +_thought_agent_copy.handoffs = [] +_dns_smtp_agent_copy.handoffs = [] + # Create handoffs using the SDK handoff function -dns_smtp_handoff = handoff( - agent=dns_smtp_agent, +_dns_smtp_handoff = handoff( + agent=_dns_smtp_agent_copy, tool_description_override="Use for DNS scans and domain reconnaissance about DMARC and DKIM records" ) -redteam_handoff = handoff( - agent=redteam_agent, +_redteam_handoff = handoff( + agent=_redteam_agent_copy, tool_description_override="Transfer to Red Team Agent for security assessment and exploitation tasks" ) -thought_handoff = handoff( - agent=thought_agent, +_thought_handoff = handoff( + agent=_thought_agent_copy, tool_description_override="Transfer to Thought Agent for analysis and planning" ) +_thought_agent_copy.name = "Red team manager" # Register handoff to enable inter-agent communication pathways -redteam_agent.handoffs.append(dns_smtp_handoff) -dns_smtp_agent.handoffs.append(redteam_handoff) -thought_agent.handoffs.append(redteam_handoff) +_redteam_agent_copy.handoffs.append(_dns_smtp_handoff) +_dns_smtp_agent_copy.handoffs.append(_redteam_handoff) +_thought_agent_copy.handoffs.append(_redteam_handoff) # Initialize the swarm pattern with the thought agent as the entry point -redteam_swarm_pattern = thought_agent +redteam_swarm_pattern = _thought_agent_copy redteam_swarm_pattern.pattern = "swarm" \ No newline at end of file diff --git a/src/cai/agents/retester.py b/src/cai/agents/retester.py new file mode 100644 index 00000000..1fafd31b --- /dev/null +++ b/src/cai/agents/retester.py @@ -0,0 +1,46 @@ +"""Retester Agent for vulnerability verification and triage""" +import os +from dotenv import load_dotenv +from cai.sdk.agents import Agent, OpenAIChatCompletionsModel +from openai import AsyncOpenAI +from cai.util import load_prompt_template +from cai.tools.reconnaissance.generic_linux_command import ( # pylint: disable=import-error # noqa: E501 + generic_linux_command +) +from cai.tools.web.search_web import ( # pylint: disable=import-error # noqa: E501 + make_google_search +) +from cai.tools.reconnaissance.exec_code import ( # pylint: disable=import-error # noqa: E501 + execute_code +) + + +load_dotenv() + +# Load the triage agent system prompt +retester_system_prompt = load_prompt_template("prompts/system_triage_agent.md") + +tools = [ + generic_linux_command, + execute_code +] + +if os.getenv('GOOGLE_SEARCH_API_KEY') and os.getenv('GOOGLE_SEARCH_CX'): + tools.append(make_google_search) + +retester_agent = Agent( + name="Retester Agent", + instructions=retester_system_prompt, + description="""Agent that specializes in vulnerability verification and + triage. Expert in determining exploitability and + eliminating false positives.""", + tools=tools, + model=OpenAIChatCompletionsModel( + model=os.getenv('CAI_MODEL', "qwen2.5:14b"), + openai_client=AsyncOpenAI(), + ) +) + + + + diff --git a/src/cai/prompts/system_triage_agent.md b/src/cai/prompts/system_triage_agent.md new file mode 100644 index 00000000..936a7f78 --- /dev/null +++ b/src/cai/prompts/system_triage_agent.md @@ -0,0 +1,78 @@ +# Cybersecurity Triage Agent + +You are an autonomous cybersecurity agent specializing in vulnerability +verification and false positive elimination. Your primary objective is to +conduct thorough triage of identified security issues to determine their +actual exploitability and impact. + +## Core Objectives + +1. **Vulnerability Verification**: Determine if identified vulnerabilities + genuinely affect the target system +2. **Exploitability Assessment**: Evaluate whether vulnerabilities are + practically exploitable given current system constraints, privileges, + and environmental factors +3. **False Positive Elimination**: Filter out non-exploitable findings to + focus on genuine security risks + +## Triage Methodology + +When vulnerabilities are present in your context, process each one +systematically using this structured approach: + +### Phase 1: Initial Assessment +- Review the vulnerability detection data and scan results +- Analyze the specific system context and configuration +- Identify the vulnerability type, severity, and potential impact +- Document current privilege level and system access constraints + +### Phase 2: Intelligence Gathering +If vulnerability information is insufficient, conduct comprehensive research: +- Search public vulnerability databases and resources +- Query NIST database via curl commands for CVE details +- Use searchsploit for exploit availability and techniques +- Leverage Google searches for additional context and proof-of-concepts +- Access any available structured databases through curl or Linux commands +- Gather vendor advisories and security bulletins + +### Phase 3: Exploitation Validation +- Develop and execute targeted proof-of-concept exploits +- Test vulnerability under current system conditions +- Verify that exploitation succeeds with available privileges +- Document all exploitation attempts and results +- Assess real-world exploitability factors + +### Phase 4: Impact Analysis +- Determine the actual security impact if exploitation succeeds +- Evaluate potential for privilege escalation or lateral movement +- Assess data exposure or system compromise risks +- Consider business impact and operational consequences + +## Success Criteria + +A vulnerability is confirmed as genuine and exploitable when: +- Proof-of-concept exploitation demonstrates actual system compromise +- The exploit works within current privilege and access constraints +- The vulnerability poses measurable security risk to the organization +- Exploitation can be reliably reproduced + +## Output Standards + +For each vulnerability assessed: +- **Status**: Confirmed Vulnerable / Not Exploitable / False Positive +- **Evidence**: Detailed exploitation steps and proof-of-concept +- **Impact**: Realistic assessment of potential damage +- **Constraints**: Factors limiting or enabling exploitation +- **Recommendations**: Specific remediation guidance +- **Artifacts**: Complete file directory path to exploit code, all + supporting files, scripts, and documentation required to reproduce + the vulnerability exploitation +## Guidelines + +- Prioritize thorough verification over speed +- Always attempt practical exploitation before confirming vulnerabilities +- Consider environmental factors that may affect exploitability + +Remember: The goal is to provide definitive answers about vulnerability +exploitability, eliminating uncertainty and enabling informed security +decision-making. From 0c9544a76ba0071bc62d1313fe391eaa905cb49a Mon Sep 17 00:00:00 2001 From: luijait Date: Thu, 29 May 2025 16:13:05 +0200 Subject: [PATCH 3/6] Add deepseek r1 reasoning traces suppport --- .../agents/models/openai_chatcompletions.py | 19 ++++- src/cai/util.py | 79 ++++++++++++++----- 2 files changed, 78 insertions(+), 20 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 0df96012..9d9bc01a 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -1141,7 +1141,9 @@ class OpenAIChatCompletionsModel(Model): else: # Non-streaming mode: Use simple text output from cai.util import print_claude_reasoning_simple, detect_claude_thinking_in_stream - if detect_claude_thinking_in_stream(str(self.model)): + # Check if model supports reasoning (Claude or DeepSeek) + model_str_lower = str(self.model).lower() + if detect_claude_thinking_in_stream(str(self.model)) or "deepseek" in model_str_lower: print_claude_reasoning_simple(reasoning_content, self.agent_name, str(self.model)) @@ -2153,6 +2155,14 @@ class OpenAIChatCompletionsModel(Model): # Remove tool_choice if no tools are specified if not converted_tools: kwargs.pop("tool_choice", None) + + # Add reasoning support for DeepSeek + # DeepSeek supports reasoning_effort parameter + if hasattr(model_settings, "reasoning_effort") and model_settings.reasoning_effort: + kwargs["reasoning_effort"] = model_settings.reasoning_effort + else: + # Default to "low" reasoning effort if model supports it + kwargs["reasoning_effort"] = "high" elif provider == "claude" or "claude" in model_str: litellm.drop_params = True kwargs.pop("store", None) @@ -2355,6 +2365,13 @@ class OpenAIChatCompletionsModel(Model): provider_kwargs["custom_llm_provider"] = "deepseek" provider_kwargs.pop("store", None) # DeepSeek doesn't support store parameter provider_kwargs.pop("parallel_tool_calls", None) # DeepSeek doesn't support parallel tool calls + + # Add reasoning support for DeepSeek + if hasattr(model_settings, "reasoning_effort") and model_settings.reasoning_effort: + provider_kwargs["reasoning_effort"] = model_settings.reasoning_effort + else: + # Default to "low" reasoning effort + provider_kwargs["reasoning_effort"] = "high" elif provider == "claude" or "claude" in model_str: provider_kwargs["custom_llm_provider"] = "anthropic" provider_kwargs.pop("store", None) # Claude doesn't support store parameter diff --git a/src/cai/util.py b/src/cai/util.py index ffa87576..04e54f5e 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -3165,8 +3165,8 @@ def setup_ctf(): def create_claude_thinking_context(agent_name, counter, model): """ - Create a streaming context for Claude thinking/reasoning display. - This creates a dedicated panel that shows Claude's internal reasoning process. + Create a streaming context for AI thinking/reasoning display. + This creates a dedicated panel that shows the model's internal reasoning process. Args: agent_name: The name of the agent @@ -3198,10 +3198,19 @@ def create_claude_thinking_context(agent_name, counter, model): terminal_width, _ = shutil.get_terminal_size((100, 24)) panel_width = min(terminal_width - 4, 120) + # Determine model type for display + model_str = str(model).lower() + if "claude" in model_str: + model_display = "Claude" + elif "deepseek" in model_str: + model_display = "DeepSeek" + else: + model_display = "AI" + # Create the thinking panel header header = Text() header.append("šŸ§  ", style="bold yellow") - header.append(f"Claude Reasoning [{counter}]", style="bold yellow") + header.append(f"{model_display} Reasoning [{counter}]", style="bold yellow") header.append(f" | {agent_name}", style="bold cyan") header.append(f" | {timestamp}", style="dim") @@ -3211,7 +3220,7 @@ def create_claude_thinking_context(agent_name, counter, model): # Create the panel for thinking panel = Panel( Group(header, Text("\n"), thinking_content), - title="[bold yellow]šŸ§  Thinking Process[/bold yellow]", + title=f"[bold yellow]šŸ§  {model_display} Thinking Process[/bold yellow]", border_style="yellow", box=ROUNDED, padding=(1, 2), @@ -3230,6 +3239,7 @@ def create_claude_thinking_context(agent_name, counter, model): "thinking_content": thinking_content, "timestamp": timestamp, "model": model, + "model_display": model_display, "agent_name": agent_name, "panel_width": panel_width, "is_started": False, @@ -3242,12 +3252,12 @@ def create_claude_thinking_context(agent_name, counter, model): return context except Exception as e: - print(f"Error creating Claude thinking context: {e}") + print(f"Error creating {model_display} thinking context: {e}") return None def update_claude_thinking_content(context, thinking_delta): """ - Update the Claude thinking content with new reasoning text. + Update the AI thinking content with new reasoning text. Args: context: The thinking context created by create_claude_thinking_context @@ -3283,6 +3293,9 @@ def update_claude_thinking_content(context, thinking_delta): # For short thinking, use regular text with styling thinking_display = Text(thinking_text, style="white") + # Get model display name from context + model_display = context.get("model_display", "AI") + # Update the panel content updated_panel = Panel( Group( @@ -3290,7 +3303,7 @@ def update_claude_thinking_content(context, thinking_delta): Text("\n"), thinking_display ), - title="[bold yellow]šŸ§  Thinking Process[/bold yellow]", + title=f"[bold yellow]šŸ§  {model_display} Thinking Process[/bold yellow]", border_style="yellow", box=ROUNDED, padding=(1, 2), @@ -3304,7 +3317,8 @@ def update_claude_thinking_content(context, thinking_delta): context["live"].start() context["is_started"] = True except Exception as e: - print(f"Error starting Claude thinking display: {e}") + model_display = context.get("model_display", "AI") + print(f"Error starting {model_display} thinking display: {e}") return False # Update the live display @@ -3315,12 +3329,13 @@ def update_claude_thinking_content(context, thinking_delta): return True except Exception as e: - print(f"Error updating Claude thinking content: {e}") + model_display = context.get("model_display", "AI") + print(f"Error updating {model_display} thinking content: {e}") return False def finish_claude_thinking_display(context): """ - Finish the Claude thinking display session. + Finish the AI thinking display session. Args: context: The thinking context to finish @@ -3339,10 +3354,13 @@ def finish_claude_thinking_display(context): from rich.syntax import Syntax from rich.console import Group + # Get model display name + model_display = context.get("model_display", "AI") + # Add final formatting to show completion final_header = Text() final_header.append("šŸ§  ", style="bold green") - final_header.append(f"Claude Reasoning Complete", style="bold green") + final_header.append(f"{model_display} Reasoning Complete", style="bold green") final_header.append(f" | {context['agent_name']}", style="bold cyan") final_header.append(f" | {context['timestamp']}", style="dim") @@ -3368,7 +3386,7 @@ def finish_claude_thinking_display(context): Text("\n"), final_thinking_display ), - title="[bold green]šŸ§  Thinking Complete[/bold green]", + title=f"[bold green]šŸ§  {model_display} Thinking Complete[/bold green]", border_style="green", box=ROUNDED, padding=(1, 2), @@ -3390,13 +3408,14 @@ def finish_claude_thinking_display(context): return True except Exception as e: - print(f"Error finishing Claude thinking display: {e}") + model_display = context.get("model_display", "AI") + print(f"Error finishing {model_display} thinking display: {e}") return False def detect_claude_thinking_in_stream(model_name): """ Detect if a model should show thinking/reasoning display. - Only applies to Claude models with reasoning capability. + Applies to Claude and DeepSeek models with reasoning capability. Args: model_name: The model name to check @@ -3412,7 +3431,7 @@ def detect_claude_thinking_in_stream(model_name): # Check for Claude models with reasoning capability # Claude 4 models (like claude-sonnet-4-20250514) support reasoning # Also check for explicit "thinking" in model name - has_reasoning = ( + has_claude_reasoning = ( "claude" in model_str and ( # Claude 4 models (sonnet-4, haiku-4, opus-4) "-4-" in model_str or @@ -3425,11 +3444,23 @@ def detect_claude_thinking_in_stream(model_name): ) ) - return has_reasoning + # Check for DeepSeek models with reasoning capability + has_deepseek_reasoning = ( + "deepseek" in model_str and ( + # DeepSeek reasoner models + "reasoner" in model_str or + # DeepSeek chat models also support reasoning + "chat" in model_str or + # Generic deepseek models likely support it + "/" in model_str # e.g., deepseek/deepseek-chat + ) + ) + + return has_claude_reasoning or has_deepseek_reasoning def print_claude_reasoning_simple(reasoning_content, agent_name, model_name): """ - Print Claude reasoning content in simple mode (no Rich panels). + Print AI reasoning content in simple mode (no Rich panels). Used when CAI_STREAM=False. Args: @@ -3440,16 +3471,26 @@ def print_claude_reasoning_simple(reasoning_content, agent_name, model_name): if not reasoning_content or not reasoning_content.strip(): return + # Determine model type for display + model_str = str(model_name).lower() + if "claude" in model_str: + model_display = "Claude" + elif "deepseek" in model_str: + model_display = "DeepSeek" + else: + model_display = "AI" + # Simple text output without Rich formatting timestamp = datetime.now().strftime("%H:%M:%S") - print(f"\nšŸ§  Reasoning | {agent_name} | {model_name} | {timestamp}") + print(f"\nšŸ§  {model_display} Reasoning | {agent_name} | {model_name} | {timestamp}") print("=" * 60) print(reasoning_content) print("=" * 60 + "\n") def start_claude_thinking_if_applicable(model_name, agent_name, counter): """ - Start Claude thinking display if the model supports it AND streaming is enabled. + Start AI thinking display if the model supports it AND streaming is enabled. + Supports Claude and DeepSeek models with reasoning capabilities. Args: model_name: The model name From b53547ba4678320ed071d5f682bacdf43de21823 Mon Sep 17 00:00:00 2001 From: luijait Date: Thu, 29 May 2025 16:32:56 +0200 Subject: [PATCH 4/6] Add deepseek resilence --- .../agents/models/openai_chatcompletions.py | 11 ++- src/cai/util.py | 78 ++++++++++++------- 2 files changed, 58 insertions(+), 31 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 9d9bc01a..e3fe3772 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -2162,7 +2162,7 @@ class OpenAIChatCompletionsModel(Model): kwargs["reasoning_effort"] = model_settings.reasoning_effort else: # Default to "low" reasoning effort if model supports it - kwargs["reasoning_effort"] = "high" + kwargs["reasoning_effort"] = "low" elif provider == "claude" or "claude" in model_str: litellm.drop_params = True kwargs.pop("store", None) @@ -2371,7 +2371,7 @@ class OpenAIChatCompletionsModel(Model): provider_kwargs["reasoning_effort"] = model_settings.reasoning_effort else: # Default to "low" reasoning effort - provider_kwargs["reasoning_effort"] = "high" + provider_kwargs["reasoning_effort"] = "low" elif provider == "claude" or "claude" in model_str: provider_kwargs["custom_llm_provider"] = "anthropic" provider_kwargs.pop("store", None) # Claude doesn't support store parameter @@ -3251,6 +3251,9 @@ class _Converter: call_id = func_output["call_id"] output_content = func_output["output"] + # IMPORTANT: Truncate call_id to 40 characters for consistency + truncated_call_id = call_id[:40] if call_id else call_id + # Update execution timing if we have the start time if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls: tool_call_details = cls.recent_tool_calls[call_id] # Renamed for clarity @@ -3371,10 +3374,10 @@ class _Converter: # The responsibility for ensuring a preceding assistant message # is now fully deferred to fix_message_list, called later. - # Now add the tool message + # Now add the tool message with truncated call_id msg: ChatCompletionToolMessageParam = { "role": "tool", - "tool_call_id": func_output["call_id"], + "tool_call_id": truncated_call_id, "content": func_output["output"], } result.append(msg) diff --git a/src/cai/util.py b/src/cai/util.py index 04e54f5e..b69a4043 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -741,6 +741,7 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 a tool_result block (tool message with matching tool_call_id). 6. Each 'tool' message must be immediately preceded by an 'assistant' message with matching tool_call_id in its tool_calls. + 7. Tool call IDs are truncated to 40 characters for API compatibility. Args: messages (List[dict]): List of message dictionaries containing @@ -754,22 +755,45 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 # Deep-copy to ensure we don't modify the input sanitized_messages = [] - # First pass - identify tool_call_ids from assistant messages and tool messages + # First, truncate all tool call IDs to 40 characters throughout the messages + # This ensures consistency for providers like DeepSeek that have strict ID matching + for msg in messages: + msg_copy = msg.copy() + + # Truncate tool_call_id in tool messages + if msg_copy.get("role") == "tool" and msg_copy.get("tool_call_id"): + if len(msg_copy["tool_call_id"]) > 40: + msg_copy["tool_call_id"] = msg_copy["tool_call_id"][:40] + + # Truncate IDs in assistant tool_calls + if msg_copy.get("role") == "assistant" and msg_copy.get("tool_calls"): + tool_calls_copy = [] + for tc in msg_copy["tool_calls"]: + tc_copy = tc.copy() + if tc_copy.get("id") and len(tc_copy["id"]) > 40: + tc_copy["id"] = tc_copy["id"][:40] + tool_calls_copy.append(tc_copy) + msg_copy["tool_calls"] = tool_calls_copy + + sanitized_messages.append(msg_copy) + + # Now process the messages with truncated IDs + processed_messages = [] tool_call_map = {} # Map from tool_call_id to (assistant_idx, tool_idx) - for i, msg in enumerate(messages): + for i, msg in enumerate(sanitized_messages): # Skip empty messages (considered empty if 'content' is None or only whitespace) if msg.get("role") in ["user", "system"] and (msg.get("content") is None or not str(msg.get("content", "")).strip()): # Special case: if it's a system message, set content to empty string instead of skipping if msg.get("role") == "system": # Replace None with empty string msg["content"] = "" - sanitized_messages.append(msg) + processed_messages.append(msg) # Skip empty user messages entirely continue - # Add valid messages to our sanitized list first - sanitized_messages.append(msg) + # Add valid messages to our processed list first + processed_messages.append(msg) # Now track tool calls and tool messages for pairing if msg.get("role") == "assistant" and msg.get("tool_calls"): @@ -777,12 +801,12 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 if tc.get("id"): tool_id = tc.get("id") if tool_id not in tool_call_map: - tool_call_map[tool_id] = {"assistant_idx": len(sanitized_messages) - 1, "tool_idx": None} + tool_call_map[tool_id] = {"assistant_idx": len(processed_messages) - 1, "tool_idx": None} if msg.get("role") == "tool" and msg.get("tool_call_id"): tool_id = msg.get("tool_call_id") if tool_id in tool_call_map: - tool_call_map[tool_id]["tool_idx"] = len(sanitized_messages) - 1 + tool_call_map[tool_id]["tool_idx"] = len(processed_messages) - 1 else: # Tool response without a matching tool call - create a synthetic pair # by adding a dummy assistant message with a tool_call @@ -799,15 +823,15 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 }] } # Insert the assistant message *before* the tool message - sanitized_messages.insert(len(sanitized_messages) - 1, assistant_msg) + processed_messages.insert(len(processed_messages) - 1, assistant_msg) # Update mapping - tool_call_map[tool_id] = {"assistant_idx": len(sanitized_messages) - 2, "tool_idx": len(sanitized_messages) - 1} + tool_call_map[tool_id] = {"assistant_idx": len(processed_messages) - 2, "tool_idx": len(processed_messages) - 1} # Second pass - ensure correct sequence (tool messages must directly follow their assistant messages) # This fixes the error "messages with role 'tool' must be a response to a preceeding message with 'tool_calls'" i = 0 - while i < len(sanitized_messages): - msg = sanitized_messages[i] + while i < len(processed_messages): + msg = processed_messages[i] # Check if this is a tool message that might be out of sequence if msg.get("role") == "tool" and msg.get("tool_call_id"): @@ -815,7 +839,7 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 # If this isn't the first message, check if the previous message is a matching assistant message if i > 0: - prev_msg = sanitized_messages[i-1] + prev_msg = processed_messages[i-1] # Check if the previous message is an assistant message with matching tool_call_id is_valid_sequence = ( @@ -827,7 +851,7 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 if not is_valid_sequence: # Find the assistant message with this tool_call_id assistant_idx = None - for j, assistant_msg in enumerate(sanitized_messages): + for j, assistant_msg in enumerate(processed_messages): if (assistant_msg.get("role") == "assistant" and assistant_msg.get("tool_calls") and any(tc.get("id") == tool_id for tc in assistant_msg.get("tool_calls", []))): @@ -837,10 +861,10 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 # If we found a matching assistant message, move this tool message right after it if assistant_idx is not None: # Remember to save the tool message - tool_msg = sanitized_messages.pop(i) + tool_msg = processed_messages.pop(i) # Insert right after the assistant message - sanitized_messages.insert(assistant_idx + 1, tool_msg) + processed_messages.insert(assistant_idx + 1, tool_msg) # Adjust i to account for the move if assistant_idx < i: @@ -867,7 +891,7 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 } # Insert the assistant message before the tool message - sanitized_messages.insert(i, assistant_msg) + processed_messages.insert(i, assistant_msg) # Skip past both messages i += 2 @@ -889,7 +913,7 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 } # Insert the assistant message before the tool message - sanitized_messages.insert(0, assistant_msg) + processed_messages.insert(0, assistant_msg) # Skip past both messages i += 2 @@ -903,7 +927,7 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 if indices["tool_idx"] is None: # Tool call without a response - create a synthetic tool message assistant_idx = indices["assistant_idx"] - assistant_msg = sanitized_messages[assistant_idx] + assistant_msg = processed_messages[assistant_idx] # Find the relevant tool call tool_name = "unknown_function" @@ -921,18 +945,18 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 } # Insert immediately after the assistant message - if assistant_idx + 1 < len(sanitized_messages): + if assistant_idx + 1 < len(processed_messages): # Insert at the position after assistant - sanitized_messages.insert(assistant_idx + 1, tool_msg) + processed_messages.insert(assistant_idx + 1, tool_msg) else: # Just append if we're at the end - sanitized_messages.append(tool_msg) + processed_messages.append(tool_msg) # Update the map to note that this tool call now has a response tool_call_map[tool_id]["tool_idx"] = assistant_idx + 1 # Ensure messages have non-null content (required by some providers) - for msg in sanitized_messages: + for msg in processed_messages: # For assistant messages with tool_calls, content can be None if msg.get("role") == "assistant" and msg.get("tool_calls"): # Assistant messages with tool calls can have None content - this is valid @@ -949,9 +973,9 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 # Special case for Claude: ensure strict alternating pattern between assistant tool_calls and tool results # If multiple consecutive assistant messages with tool_calls exist, interleave them with tool responses i = 0 - while i < len(sanitized_messages) - 1: - current_msg = sanitized_messages[i] - next_msg = sanitized_messages[i + 1] + while i < len(processed_messages) - 1: + current_msg = processed_messages[i] + next_msg = processed_messages[i + 1] # When current message is assistant with tool_calls and next message is NOT a tool response if (current_msg.get("role") == "assistant" and @@ -972,13 +996,13 @@ def fix_message_list(messages): # pylint: disable=R0914,R0915,R0912 } # Insert the tool message after the current assistant message - sanitized_messages.insert(i + 1, tool_msg) + processed_messages.insert(i + 1, tool_msg) # Skip over the newly inserted message i += 2 else: i += 1 - return sanitized_messages + return processed_messages def cli_print_tool_call(tool_name="", args="", output="", prefix=" "): """Print a tool call with pretty formatting""" From e5a00afeefe3b6c48fae1f9e127ae8bf10d484e3 Mon Sep 17 00:00:00 2001 From: luijait Date: Thu, 29 May 2025 23:58:35 +0200 Subject: [PATCH 5/6] Add MCP command --- src/cai/repl/commands/__init__.py | 3 +- src/cai/repl/commands/mcp.py | 1066 +++++++++++++++++++++++++++++ src/cai/sdk/agents/mcp/util.py | 46 +- src/cai/util.py | 58 +- 4 files changed, 1169 insertions(+), 4 deletions(-) create mode 100644 src/cai/repl/commands/mcp.py 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]") From 1bde081d6ff3a9895800df33e0877f5b4cee1f28 Mon Sep 17 00:00:00 2001 From: luijait Date: Fri, 30 May 2025 12:00:58 +0200 Subject: [PATCH 6/6] Fix parallel input bug --- src/cai/repl/commands/parallel.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/cai/repl/commands/parallel.py b/src/cai/repl/commands/parallel.py index 5582045d..e78cd821 100644 --- a/src/cai/repl/commands/parallel.py +++ b/src/cai/repl/commands/parallel.py @@ -98,8 +98,9 @@ class ParallelCommand(Command): model = args[i + 1] i += 2 elif args[i] == "--prompt" and i + 1 < len(args): - prompt = args[i + 1] - i += 2 + # Capture all remaining arguments as the prompt + prompt = " ".join(args[i + 1:]) + break # Stop parsing after --prompt since we take everything after it else: i += 1