mirror of https://github.com/aliasrobotics/cai.git
Merge branch 'main' into agent2pwnbench
This commit is contained in:
cris
commit
cd1103a1db
37
README.md
37
README.md
|
|
@ -22,7 +22,8 @@
|
|||
[](https://arxiv.org/abs/2506.23592)
|
||||
[](https://arxiv.org/abs/2508.13588)
|
||||
[](https://arxiv.org/abs/2508.21669)
|
||||
|
||||
[](https://arxiv.org/abs/2509.14096)
|
||||
[](https://arxiv.org/abs/2509.14139)
|
||||
|
||||
|
||||
</div>
|
||||
|
|
@ -124,22 +125,29 @@ Cybersecurity AI (CAI) is a lightweight, open-source framework that empowers sec
|
|||
[](https://lu.ma/roboticshack?tk=RuryKF)
|
||||
|
||||
### 📊 Research Impact
|
||||
- Pioneered LLM-powered AI Security with PentestGPT, establishing the foundation for the `Cybersecurity AI` research domain [](https://arxiv.org/pdf/2308.06782)
|
||||
- Established the `Cybersecurity AI` research line with **4 papers and technical reports** and active research collaborations [](https://arxiv.org/pdf/2504.06017) [](https://arxiv.org/abs/2506.23592) [](https://arxiv.org/abs/2508.13588) [](https://arxiv.org/abs/2508.21669)
|
||||
- Demonstrated **3,600× performance improvement** over human penetration testers in standardized CTF benchmark evaluations [](https://arxiv.org/pdf/2504.06017)
|
||||
- Identified **CVSS 4.3-7.5 severity vulnerabilities** in production systems through automated security assessment [](https://arxiv.org/pdf/2504.06017)
|
||||
- **Democratization of AI-empowered vulnerability research**: CAI enables both non-security domain experts and experienced researchers to conduct more efficient vulnerability discovery, expanding the security research community while empowering small and medium enterprises to conduct autonomous security assessments [](https://arxiv.org/pdf/2504.06017)
|
||||
- **Systematic evaluation of large language models** across both proprietary and open-weight architectures, revealing <u>substantial gaps</u> between vendor-reported capabilities and empirical cybersecurity performance metrics [](https://arxiv.org/pdf/2504.06017)
|
||||
- Established the **autonomy levels in cybersecurity** and argued about autonomy vs automation in the field [](https://arxiv.org/abs/2506.23592)
|
||||
- **Collaborative research initiatives** with international academic institutions focused on developing cybersecurity education curricula and training methodologies [](https://arxiv.org/abs/2508.13588)
|
||||
- **Contributed a comprehensive defense framework against prompt injection in AI security agents**: developed and empirically validated a multi-layered defense system that addresses the identified prompt injection issues [](https://arxiv.org/abs/2508.21669)
|
||||
- Pioneered LLM-powered AI Security with PentestGPT, establishing the foundation for the `Cybersecurity AI` research domain [](https://arxiv.org/pdf/2308.06782)
|
||||
- Established the `Cybersecurity AI` research line with **6 papers and technical reports**, with active research collaborations [](https://arxiv.org/pdf/2504.06017) [](https://arxiv.org/abs/2506.23592) [](https://arxiv.org/abs/2508.13588) [](https://arxiv.org/abs/2508.21669) [](https://arxiv.org/abs/2509.14096) [](https://arxiv.org/abs/2509.14139)
|
||||
|
||||
- Demonstrated **3,600× performance improvement** over human penetration testers in standardized CTF benchmark evaluations [](https://arxiv.org/pdf/2504.06017)
|
||||
- Identified **CVSS 4.3-7.5 severity vulnerabilities** in production systems through automated security assessment [](https://arxiv.org/pdf/2504.06017)
|
||||
- **Democratization of AI-empowered vulnerability research**: CAI enables both non-security domain experts and experienced researchers to conduct more efficient vulnerability discovery, expanding the security research community while empowering small and medium enterprises to conduct autonomous security assessments [](https://arxiv.org/pdf/2504.06017)
|
||||
- **Systematic evaluation of large language models** across both proprietary and open-weight architectures, revealing <u>substantial gaps</u> between vendor-reported capabilities and empirical cybersecurity performance metrics [](https://arxiv.org/pdf/2504.06017)
|
||||
- Established the **autonomy levels in cybersecurity** and argued about autonomy vs automation in the field [](https://arxiv.org/abs/2506.23592)
|
||||
- **Collaborative research initiatives** with international academic institutions focused on developing cybersecurity education curricula and training methodologies [](https://arxiv.org/abs/2508.13588)
|
||||
- **Contributed a comprehensive defense framework against prompt injection in AI security agents**: developed and empirically validated a multi-layered defense system that addresses the identified prompt injection issues [](https://arxiv.org/abs/2508.21669)
|
||||
- Explord the Cybersecurity of Humanoid Robots with CAI and identified new attack vectors showing how it `(a)` operates simultaneously as a covert surveillance node and `(b)` can be purposed as an active cyber operations platform [](https://arxiv.org/abs/2509.14096) [](https://arxiv.org/abs/2509.14139)
|
||||
|
||||
|
||||
### 📚 Research products: `Cybersecurity AI`
|
||||
|
||||
| CAI, An Open, Bug Bounty-Ready Cybersecurity AI | The Dangerous Gap Between Automation and Autonomy | CAI Fluency, A Framework for Cybersecurity AI Fluency | Hacking the AI Hackers via Prompt Injection |
|
||||
|---|---|---|---|
|
||||
| [<img src="https://aliasrobotics.com/img/paper-cai.png" width="350">](https://arxiv.org/pdf/2508.13588) | [<img src="https://aliasrobotics.com/img/cai_automation_vs_autonomy.png" width="350">](https://www.arxiv.org/pdf/2506.23592) | [<img src="https://aliasrobotics.com/img/cai_fluency_cover.png" width="350">](https://arxiv.org/pdf/2504.06017) | [<img src="https://aliasrobotics.com/img/aihackers.jpeg" width="500">](https://arxiv.org/pdf/2508.21669) |
|
||||
| CAI, An Open, Bug Bounty-Ready Cybersecurity AI [](https://arxiv.org/pdf/2504.06017) | The Dangerous Gap Between Automation and Autonomy [](https://arxiv.org/abs/2506.23592) | CAI Fluency, A Framework for Cybersecurity AI Fluency [](https://arxiv.org/abs/2508.13588) |
|
||||
|---|---|---|
|
||||
| [<img src="https://aliasrobotics.com/img/paper-cai.png" width="350">](https://arxiv.org/pdf/2504.06017) | [<img src="https://aliasrobotics.com/img/cai_automation_vs_autonomy.png" width="350">](https://www.arxiv.org/pdf/2506.23592) | [<img src="https://aliasrobotics.com/img/cai_fluency_cover.png" width="350">](https://arxiv.org/pdf/2508.13588) |
|
||||
|
||||
|
||||
| Hacking the AI Hackers via Prompt Injection [](https://arxiv.org/abs/2508.21669) | Humanoid Robots as Attack Vectors [](https://arxiv.org/abs/2509.14139) | The Cybersecurity of a Humanoid Robot [](https://arxiv.org/abs/2509.14096) |
|
||||
|---|---|---|
|
||||
| [<img src="https://aliasrobotics.com/img/aihackers.jpeg" width="350">](https://arxiv.org/pdf/2508.21669) | [<img src="https://aliasrobotics.com/img/humanoids-cover.png" width="350">](https://arxiv.org/pdf/2509.14139) | [<img src="https://aliasrobotics.com/img/humanoid.png" width="350">](https://arxiv.org/pdf/2509.14096) |
|
||||
|
||||
|
||||
|
||||
|
|
@ -192,6 +200,7 @@ Cybersecurity AI is a critical field, yet many groups are misguidedly pursuing i
|
|||
- [CrackenAGI](https://cracken.ai/)
|
||||
- [ETHIACK](https://ethiack.com/)
|
||||
- [Horizon3](https://horizon3.ai/)
|
||||
- [Irregular](https://www.irregular.com/)
|
||||
- [Kindo](https://www.kindo.ai/)
|
||||
- [Lakera](https://lakera.ai)
|
||||
- [Mindfort](www.mindfort.ai)
|
||||
|
|
@ -235,7 +244,7 @@ Cybersecurity AI is a critical field, yet many groups are misguidedly pursuing i
|
|||
| **Episode 0**: What is CAI? | Cybersecurity AI (`CAI`) explained | [](https://www.youtube.com/watch?v=nBdTxbKM4oo) | [](https://www.youtube.com/watch?v=FaUL9HXrQ5k) |
|
||||
| **Episode 1**: The `CAI` Framework | Vision & Ethics - Explore the core motivation behind CAI and delve into the crucial ethical principles guiding its development. Understand the motivation behind CAI and how you can actively contribute to the future of cybersecurity and the CAI framework. | [](https://www.youtube.com/watch?v=QEiGdsMf29M&list=PLLc16OUiZWd4RuFdN5_Wx9xwjCVVbopzr&index=3) | |
|
||||
| **Episode 2**: From Zero to Cyber Hero | Breaking into Cybersecurity with AI - A comprehensive guide for complete beginners to become cybersecurity practitioners using CAI and AI tools. Learn how to leverage artificial intelligence to accelerate your cybersecurity learning journey, from understanding basic security concepts to performing real-world security assessments, all without requiring prior cybersecurity experience. | [](https://www.youtube.com/watch?v=hSTLHOOcQoY&list=PLLc16OUiZWd4RuFdN5_Wx9xwjCVVbopzr&index=14) | |
|
||||
| **Episode 3**: Vibe-Hacking Tutorial | "My first Hack" - A Vibe-Hacking guide for newbies. We demonstrate a simple web security hack using a default agent and show how to leverage tools and interpret CIA output with the help of the CAI Python API. You'll also learn to compare different LLM models to find the best fit for your hacking endeavors. | [](https://www.youtube.com/watch?v=9vZ_Iyex7uI&list=PLLc16OUiZWd4RuFdN5_Wx9xwjCVVbopzr&index=1) | [](https://www.youtube.com/watch?v=iAOMaI1ftiA&list=PLLc16OUiZWd4RuFdN5_Wx9xwjCVVbopzr&index=2) |
|
||||
| **Episode 3**: Vibe-Hacking Tutorial | "My first Hack" - A Vibe-Hacking guide for newbies. We demonstrate a simple web security hack using a default agent and show how to leverage tools and interpret CAI output with the help of the CAI Python API. You'll also learn to compare different LLM models to find the best fit for your hacking endeavors. | [](https://www.youtube.com/watch?v=9vZ_Iyex7uI&list=PLLc16OUiZWd4RuFdN5_Wx9xwjCVVbopzr&index=1) | [](https://www.youtube.com/watch?v=iAOMaI1ftiA&list=PLLc16OUiZWd4RuFdN5_Wx9xwjCVVbopzr&index=2) |
|
||||
| **Episode 4**: Intro ReAct | The Evolution of LLMs - Learn how LLMs evolved from basic language models to advanced multiagency AI systems. From basic LLMs to Chain-of-Thought and Reasoning LLMs towards ReAct and Multi-Agent Architectures. Get to know the basic terms | [](https://www.youtube.com/watch?v=tLdFO1flj_o&list=PLLc16OUiZWd4RuFdN5_Wx9xwjCVVbopzr&index=13) | |
|
||||
| **Episode 5**: CAI on CTF challenges | Dive into Capture The Flag (CTF) competitions using CAI. Learn how to leverage AI agents to solve various cybersecurity challenges including web exploitation, cryptography, reverse engineering, and forensics. Discover how to configure CAI for competitive hacking scenarios and maximize your CTF performance with intelligent automation. | [](https://www.youtube.com/watch?v=MrXTQ0e2to4&list=PLLc16OUiZWd4RuFdN5_Wx9xwjCVVbopzr&index=13) | [](https://www.youtube.com/watch?v=r9US_JZa9_c&list=PLLc16OUiZWd4RuFdN5_Wx9xwjCVVbopzr&index=12) |
|
||||
| | | | |
|
||||
|
|
|
|||
|
|
@ -79,6 +79,9 @@ console = Console()
|
|||
# Global registry for persistent MCP connections
|
||||
_GLOBAL_MCP_SERVERS: Dict[str, MCPServer] = {}
|
||||
|
||||
# Per-server locks to serialize tool invocations for persistent connections
|
||||
_SERVER_INVOCATION_LOCKS: Dict[str, asyncio.Lock] = {}
|
||||
|
||||
# Global registry for agent-MCP associations
|
||||
# Maps agent name to list of MCP server names
|
||||
_AGENT_MCP_ASSOCIATIONS: Dict[str, List[str]] = {}
|
||||
|
|
@ -104,6 +107,7 @@ class GlobalMCPUtil(MCPUtil):
|
|||
"tool_name": tool.name,
|
||||
"tool_schema": tool.inputSchema,
|
||||
"tool_description": tool.description,
|
||||
"persistent": isinstance(server, MCPServerStdio),
|
||||
}
|
||||
|
||||
# For SSE servers, capture the URL
|
||||
|
|
@ -123,9 +127,9 @@ class GlobalMCPUtil(MCPUtil):
|
|||
server.params, "encoding_error_handler", "strict"
|
||||
)
|
||||
|
||||
# Create a custom invoke function that creates a new connection each time
|
||||
# Create a custom invoke function that manages the server lifecycle per invocation
|
||||
async def invoke_with_fresh_connection(config, context, input_json):
|
||||
"""Custom invoke function that creates a fresh connection for each invocation"""
|
||||
"""Invoke an MCP tool, keeping STDIO transports persistent."""
|
||||
import asyncio
|
||||
import json
|
||||
import warnings
|
||||
|
|
@ -133,7 +137,6 @@ class GlobalMCPUtil(MCPUtil):
|
|||
from cai.sdk.agents.exceptions import AgentsException, ModelBehaviorError
|
||||
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:
|
||||
|
|
@ -141,130 +144,167 @@ class GlobalMCPUtil(MCPUtil):
|
|||
f"Invalid JSON input for tool {config['tool_name']}: {input_json}"
|
||||
) from e
|
||||
|
||||
# Create a fresh server connection with timeout
|
||||
server = None
|
||||
result = None
|
||||
max_retries = 2
|
||||
retry_count = 0
|
||||
server = None
|
||||
should_cleanup = False
|
||||
persistent = bool(config.get("persistent"))
|
||||
|
||||
# Suppress warnings about async generator cleanup
|
||||
with warnings.catch_warnings():
|
||||
warnings.filterwarnings("ignore", category=RuntimeWarning)
|
||||
warnings.filterwarnings("ignore", message=".*asynchronous generator.*")
|
||||
warnings.filterwarnings("ignore", message=".*ClosedResourceError.*")
|
||||
|
||||
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}
|
||||
if persistent:
|
||||
server_name = config["name"]
|
||||
server = _GLOBAL_MCP_SERVERS.get(server_name)
|
||||
if not server or not isinstance(server, MCPServerStdio):
|
||||
raise AgentsException(
|
||||
f"MCP server '{server_name}' is unavailable. Use /mcp status to verify it is loaded."
|
||||
)
|
||||
|
||||
server = MCPServerSse(
|
||||
params,
|
||||
name=config["name"],
|
||||
cache_tools_list=False, # Don't cache since it's temporary
|
||||
lock = _SERVER_INVOCATION_LOCKS.setdefault(
|
||||
server_name, asyncio.Lock()
|
||||
)
|
||||
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)
|
||||
async with lock:
|
||||
while retry_count < max_retries:
|
||||
try:
|
||||
if not getattr(server, "session", None):
|
||||
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']}. "
|
||||
"The server may be down or not responding."
|
||||
)
|
||||
|
||||
result = await asyncio.wait_for(
|
||||
server.call_tool(config["tool_name"], json_data),
|
||||
timeout=30.0,
|
||||
)
|
||||
break
|
||||
except asyncio.TimeoutError:
|
||||
raise AgentsException(
|
||||
f"Timeout calling MCP tool {config['tool_name']}. "
|
||||
f"The tool took too long to respond."
|
||||
)
|
||||
except Exception:
|
||||
retry_count += 1
|
||||
if retry_count >= max_retries:
|
||||
raise
|
||||
import logging
|
||||
|
||||
logging.debug(
|
||||
f"Retrying MCP tool {config['tool_name']} (attempt {retry_count}/{max_retries})"
|
||||
)
|
||||
try:
|
||||
await server.cleanup()
|
||||
except Exception:
|
||||
pass
|
||||
server.session = None
|
||||
await asyncio.sleep(0.5)
|
||||
else:
|
||||
raise AgentsException(f"Unknown server type: {config['type']}")
|
||||
if config["type"] == "MCPServerSse":
|
||||
params = {
|
||||
"url": config["url"],
|
||||
"headers": config.get("headers"),
|
||||
"timeout": config.get("timeout", 5),
|
||||
"sse_read_timeout": config.get("sse_read_timeout", 60 * 5),
|
||||
}
|
||||
params = {k: v for k, v in params.items() if v is not None}
|
||||
|
||||
# Retry logic for connection and tool calls
|
||||
while retry_count < max_retries:
|
||||
try:
|
||||
# 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."
|
||||
)
|
||||
server = MCPServerSse(
|
||||
params,
|
||||
name=config["name"],
|
||||
cache_tools_list=False,
|
||||
)
|
||||
elif config["type"] == "MCPServerStdio":
|
||||
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"
|
||||
),
|
||||
}
|
||||
params = {k: v for k, v in params.items() if v is not None}
|
||||
|
||||
# Call the tool with timeout
|
||||
server = MCPServerStdio(
|
||||
params, name=config["name"], cache_tools_list=False
|
||||
)
|
||||
else:
|
||||
raise AgentsException(f"Unknown server type: {config['type']}")
|
||||
|
||||
should_cleanup = True
|
||||
|
||||
while retry_count < max_retries:
|
||||
try:
|
||||
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."
|
||||
)
|
||||
|
||||
result = await asyncio.wait_for(
|
||||
server.call_tool(config["tool_name"], json_data), timeout=30.0
|
||||
server.call_tool(config["tool_name"], json_data),
|
||||
timeout=30.0,
|
||||
)
|
||||
break # Success, exit retry loop
|
||||
break
|
||||
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:
|
||||
retry_count += 1
|
||||
if retry_count >= max_retries:
|
||||
raise
|
||||
# Log retry attempt
|
||||
import logging
|
||||
logging.debug(f"Retrying MCP tool {config['tool_name']} (attempt {retry_count}/{max_retries})")
|
||||
# Clear session for SSE servers
|
||||
if config["type"] == "MCPServerSse" and hasattr(server, 'session'):
|
||||
server.session = None
|
||||
await asyncio.sleep(0.5) # Brief delay before retry
|
||||
except Exception:
|
||||
retry_count += 1
|
||||
if retry_count >= max_retries:
|
||||
raise
|
||||
import logging
|
||||
|
||||
logging.debug(
|
||||
f"Retrying MCP tool {config['tool_name']} (attempt {retry_count}/{max_retries})"
|
||||
)
|
||||
if isinstance(server, MCPServerSse) and hasattr(server, "session"):
|
||||
server.session = None
|
||||
await asyncio.sleep(0.5)
|
||||
except Exception as e:
|
||||
# Handle ClosedResourceError and connection issues
|
||||
error_type = type(e).__name__
|
||||
error_str = str(e).lower()
|
||||
|
||||
# Improved error messages for common issues
|
||||
if (error_type in ("ClosedResourceError", "ExceptionGroup") or
|
||||
"closedresourceerror" in error_str or
|
||||
"closed" in error_str or
|
||||
"connection" in error_str):
|
||||
|
||||
if (
|
||||
error_type in ("ClosedResourceError", "ExceptionGroup")
|
||||
or "closedresourceerror" in error_str
|
||||
or "closed" in error_str
|
||||
or "connection" in error_str
|
||||
):
|
||||
raise AgentsException(
|
||||
f"Connection lost to MCP server for tool {config['tool_name']}. "
|
||||
f"This is normal for SSE servers. The tool will reconnect automatically "
|
||||
f"on the next invocation."
|
||||
"Use /mcp status to reconnect if the issue persists."
|
||||
) from e
|
||||
else:
|
||||
raise AgentsException(
|
||||
f"Error invoking MCP tool {config['tool_name']}: {type(e).__name__}: {str(e)}"
|
||||
) from 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
|
||||
if should_cleanup and server:
|
||||
if isinstance(server, MCPServerSse):
|
||||
try:
|
||||
# Don't wait too long for SSE cleanup
|
||||
await asyncio.wait_for(server.cleanup(), timeout=0.5)
|
||||
except (asyncio.TimeoutError, RuntimeError, Exception):
|
||||
# Expected for SSE connections - they close abruptly
|
||||
pass
|
||||
# Explicitly clear the session to force reconnection next time
|
||||
server.session = None
|
||||
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']}")
|
||||
|
||||
|
|
@ -336,6 +376,7 @@ def cleanup_mcp_servers():
|
|||
# Only close the loop if it's not running
|
||||
if not loop.is_running():
|
||||
loop.close()
|
||||
_SERVER_INVOCATION_LOCKS.clear()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
|
@ -987,6 +1028,7 @@ Example: `/mcp add burp 13`
|
|||
|
||||
self._run_async(cleanup_server())
|
||||
del _GLOBAL_MCP_SERVERS[server_name]
|
||||
_SERVER_INVOCATION_LOCKS.pop(server_name, None)
|
||||
console.print(f"[green]✓ Removed MCP server '{server_name}'[/green]")
|
||||
return True
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -69,6 +69,8 @@ class _MCPServerWithClientSession(MCPServer, abc.ABC):
|
|||
self.session: ClientSession | None = None
|
||||
self.exit_stack: AsyncExitStack = AsyncExitStack()
|
||||
self._cleanup_lock: asyncio.Lock = asyncio.Lock()
|
||||
self._connect_lock: asyncio.Lock = asyncio.Lock()
|
||||
self._call_lock: asyncio.Lock = asyncio.Lock()
|
||||
self.cache_tools_list = cache_tools_list
|
||||
|
||||
# The cache is always dirty at startup, so that we fetch tools at least once
|
||||
|
|
@ -100,25 +102,32 @@ class _MCPServerWithClientSession(MCPServer, abc.ABC):
|
|||
|
||||
async def connect(self):
|
||||
"""Connect to the server."""
|
||||
try:
|
||||
transport = await self.exit_stack.enter_async_context(self.create_streams())
|
||||
read, write = transport
|
||||
session = await self.exit_stack.enter_async_context(ClientSession(read, write))
|
||||
await session.initialize()
|
||||
self.session = session
|
||||
except Exception as e:
|
||||
# Only log connection errors at debug level
|
||||
error_str = str(e).lower()
|
||||
error_type = type(e).__name__
|
||||
if ("connection" in error_str or
|
||||
"refused" in error_str or
|
||||
"taskgroup" in error_str or
|
||||
error_type == "ExceptionGroup"):
|
||||
logger.debug(f"Expected connection error during MCP server init: {e}")
|
||||
else:
|
||||
logger.error(f"Error initializing MCP server: {e}")
|
||||
await self.cleanup()
|
||||
raise
|
||||
if self.session is not None:
|
||||
return
|
||||
|
||||
async with self._connect_lock:
|
||||
if self.session is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
transport = await self.exit_stack.enter_async_context(self.create_streams())
|
||||
read, write = transport
|
||||
session = await self.exit_stack.enter_async_context(ClientSession(read, write))
|
||||
await session.initialize()
|
||||
self.session = session
|
||||
except Exception as e:
|
||||
# Only log connection errors at debug level
|
||||
error_str = str(e).lower()
|
||||
error_type = type(e).__name__
|
||||
if ("connection" in error_str or
|
||||
"refused" in error_str or
|
||||
"taskgroup" in error_str or
|
||||
error_type == "ExceptionGroup"):
|
||||
logger.debug(f"Expected connection error during MCP server init: {e}")
|
||||
else:
|
||||
logger.error(f"Error initializing MCP server: {e}")
|
||||
await self.cleanup()
|
||||
raise
|
||||
|
||||
async def list_tools(self) -> list[MCPTool]:
|
||||
"""List the tools available on the server."""
|
||||
|
|
@ -141,7 +150,16 @@ class _MCPServerWithClientSession(MCPServer, abc.ABC):
|
|||
if not self.session:
|
||||
raise UserError("Server not initialized. Make sure you call `connect()` first.")
|
||||
|
||||
return await self.session.call_tool(tool_name, arguments)
|
||||
async with self._call_lock:
|
||||
if self.session is None:
|
||||
await self.connect()
|
||||
|
||||
try:
|
||||
return await self.session.call_tool(tool_name, arguments)
|
||||
except Exception:
|
||||
# Ensure resources are released and mark session stale so callers can reconnect
|
||||
await self.cleanup()
|
||||
raise
|
||||
|
||||
async def cleanup(self):
|
||||
"""Cleanup the server."""
|
||||
|
|
|
|||
Loading…
Reference in New Issue