From b270cf246daddd3410c3b44b954e277ef37529e5 Mon Sep 17 00:00:00 2001
From: Fan Huang <117979149+1793649019@users.noreply.github.com>
Date: Tue, 22 Jul 2025 15:36:37 +0800
Subject: [PATCH] Fix MCP command tool assignment and improve CLI error
 messaging (#222)

* fix(mcp): correct server param access and agent tool registration
- Use getattr() instead of .get() for accessing server.params attributes
- Fix tool registration logic to properly add tools to Agent

* fix(cli): improve fallback message for failed or unknown commands
- Replaced misleading "Unknown command" message with "Command failed or unknown"
---
 src/cai/cli.py               |  2 +-
 src/cai/repl/commands/mcp.py | 14 +++++++-------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/cai/cli.py b/src/cai/cli.py
index e1733af4..3c061601 100644
--- a/src/cai/cli.py
+++ b/src/cai/cli.py
@@ -1290,7 +1290,7 @@ def run_cai_cli(
 
                 # If command wasn't recognized, show error (skip for /shell or /s)
                 if command not in ("/shell", "/s"):
-                    console.print(f"[red]Unknown command: {command}[/red]")
+                    console.print(f"[red]Command failed or unknown: {command}[/red]")
                 continue
             from rich.text import Text
 
diff --git a/src/cai/repl/commands/mcp.py b/src/cai/repl/commands/mcp.py
index 57059ba1..a2dc6884 100644
--- a/src/cai/repl/commands/mcp.py
+++ b/src/cai/repl/commands/mcp.py
@@ -116,11 +116,11 @@ class GlobalMCPUtil(MCPUtil):
         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"
+            server_config["env"] = getattr(server.params, "env")
+            server_config["cwd"] = getattr(server.params, "cwd")
+            server_config["encoding"] = getattr(server.params, "encoding", "utf-8")
+            server_config["encoding_error_handler"] = getattr(
+                server.params, "encoding_error_handler", "strict"
             )
 
         # Create a custom invoke function that creates a new connection each time
@@ -810,9 +810,9 @@ Example: `/mcp add burp 13`
 
         # Get the agent
         try:
-            agent = get_agent_by_name(agent_identifier)
+            agent = get_available_agents()[agent_identifier]
             agent_display_name = getattr(agent, "name", agent_identifier)
-        except ValueError:
+        except KeyError:
             # Try by index
             try:
                 agents = get_available_agents()