From 87e5a11ef3fbb206f55ae12ed25d06dee9d3775b Mon Sep 17 00:00:00 2001 From: YMPgit Date: Fri, 8 May 2026 18:19:05 +0530 Subject: [PATCH] Added MCP Server Integration --- README.md | 27 ++++++ aios/config/config.yaml.example | 22 +++++ aios/config/config_manager.py | 28 +++++- aios/tool/manager.py | 15 +++- aios/tool/mcp_clients.py | 147 +++++++++++++++++++++++++++++++ requirements.txt | 1 + tests/test_mcp_clients_config.py | 17 ++++ 7 files changed, 254 insertions(+), 3 deletions(-) create mode 100644 aios/tool/mcp_clients.py create mode 100644 tests/test_mcp_clients_config.py diff --git a/README.md b/README.md index 5354f64f7..b6624ff14 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,33 @@ Below shows how agents utilize AIOS SDK to interact with AIOS kernel and how AIO ### Computer-use Specialized Architecture For computer-use agent, the architecture extends the AIOS Kernel with significant enhancements focused on computer contextualization. While preserving essential components like LLM Core(s), Context Manager, and Memory Manager, the Tool Manager module has been fundamentally redesigned to incorporate a VM (Virtual Machine) Controller and MCP Server. + +### MCP server integration (external tools) + +AIOS can also **connect to external MCP servers** and expose their tools through the existing Tool Manager. Configure MCP clients in `aios/config/config.yaml` under `tool.mcp_clients`. + +- **Tool name format**: `mcp//` +- **Supported transports**: + - **stdio**: spawn a subprocess and communicate via stdio + - **streamable_http**: connect to an MCP server via Streamable HTTP + +Example config snippet: + +```yaml +tool: + enable_mcp_server: true + mcp_server_script_path: "aios/tool/mcp_server.py" + + mcp_clients: + - id: "playwright" + transport: "stdio" + command: "npx" + args: ["-y", "@microsoft/playwright-mcp"] + + - id: "docs" + transport: "streamable_http" + url: "http://localhost:8000/mcp" +``` This redesign creates a sandboxed environment that allows agents to safely interact with computer systems while maintaining a consistent semantic mapping between agent intentions and computer operations. diff --git a/aios/config/config.yaml.example b/aios/config/config.yaml.example index 743bb532a..a8e1c621e 100644 --- a/aios/config/config.yaml.example +++ b/aios/config/config.yaml.example @@ -113,3 +113,25 @@ agent_factory: server: host: "localhost" port: 8000 + +# Tool / MCP Configuration +tool: + # Internal MCP server (Computer-as-MCP-Server style). Can be disabled if not needed. + enable_mcp_server: true + # Path is relative to the repository root. + mcp_server_script_path: "aios/tool/mcp_server.py" + + # External MCP clients (connect to other MCP servers and expose their tools as: + # mcp// + # + # Example: Playwright MCP (Node) + # - id: "playwright" + # transport: "stdio" + # command: "npx" + # args: ["-y", "@microsoft/playwright-mcp"] + # + # Example: Streamable HTTP server + # - id: "docs" + # transport: "streamable_http" + # url: "http://localhost:8000/mcp" + mcp_clients: [] diff --git a/aios/config/config_manager.py b/aios/config/config_manager.py index 9351146d0..8f7de6e7c 100644 --- a/aios/config/config_manager.py +++ b/aios/config/config_manager.py @@ -1,6 +1,7 @@ import os import yaml from typing import Any, Optional +import shutil class ConfigManager: """ @@ -45,6 +46,10 @@ def __init__(self): os.path.dirname(__file__), # aios/config directory 'config.yaml' ) + self.example_config_path = os.path.join( + os.path.dirname(__file__), + 'config.yaml.example', + ) self.load_config() def load_config(self): @@ -55,7 +60,15 @@ def load_config(self): FileNotFoundError: If the config file doesn't exist """ if not os.path.exists(self.config_path): - raise FileNotFoundError(f"Config file not found at {self.config_path}") + # Bootstrapping: if a user cloned the repo and didn't create config.yaml yet, + # initialize it from config.yaml.example. + if os.path.exists(self.example_config_path): + os.makedirs(os.path.dirname(self.config_path), exist_ok=True) + shutil.copyfile(self.example_config_path, self.config_path) + else: + raise FileNotFoundError( + f"Config file not found at {self.config_path} and no example config at {self.example_config_path}" + ) with open(self.config_path, 'r') as f: self.config = yaml.safe_load(f) @@ -214,7 +227,18 @@ def get_mcp_server_script_path(self) -> str: """ Retrieves the path to the MCP server script. """ - return os.path.join(os.getcwd(), self.config.get("tool", {}).get("mcp_server_script_path")) + tool_cfg = self.config.get("tool", {}) or {} + # Default to the in-repo MCP server. + rel = tool_cfg.get("mcp_server_script_path") or os.path.join("aios", "tool", "mcp_server.py") + return os.path.join(os.getcwd(), rel) + + def get_mcp_clients_config(self) -> list[dict]: + """Returns configured external MCP clients.""" + tool_cfg = self.config.get("tool", {}) or {} + clients = tool_cfg.get("mcp_clients") or [] + if not isinstance(clients, list): + return [] + return clients def get_scheduler_config(self) -> dict: """ diff --git a/aios/tool/manager.py b/aios/tool/manager.py index 56051fba7..5f17f77af 100644 --- a/aios/tool/manager.py +++ b/aios/tool/manager.py @@ -13,6 +13,8 @@ from threading import Lock +from aios.tool.mcp_clients import MCPClientManager, MCPTool + class ToolManager: def __init__( self, @@ -22,10 +24,14 @@ def __init__( self.tool_conflict_map = {} self.tool_conflict_map_lock = Lock() self.mcp_server_process = None # To store the mcp_server subprocess - self._start_mcp_server() # Start mcp_server on initialization + self.mcp_clients = MCPClientManager(config=config) + self._start_mcp_server() # Start mcp_server on initialization if enabled def _start_mcp_server(self): """Starts the mcp_server.py script as a background subprocess.""" + tool_cfg = config.get_tool_config() or {} + if tool_cfg.get("enable_mcp_server", True) is False: + return if self.mcp_server_process is None or self.mcp_server_process.poll() is not None: try: # Assuming mcp_server.py is in the same directory as manager.py @@ -115,5 +121,12 @@ def address_request(self, syscall) -> None: ) def load_tool_instance(self, tool_org_and_name): + if tool_org_and_name.startswith("mcp/"): + # Format: mcp// + parts = tool_org_and_name.split("/", 2) + if len(parts) < 3: + raise ValueError("MCP tool name must be in form mcp//") + _, server_id, tool_name = parts + return MCPTool(mcp_manager=self.mcp_clients, server_id=server_id, tool_name=tool_name) tool_instance = AutoTool.from_preloaded(tool_org_and_name) return tool_instance diff --git a/aios/tool/mcp_clients.py b/aios/tool/mcp_clients.py new file mode 100644 index 000000000..dc1a85a2a --- /dev/null +++ b/aios/tool/mcp_clients.py @@ -0,0 +1,147 @@ +import asyncio +import os +from dataclasses import dataclass +from typing import Any, Optional + + +@dataclass(frozen=True) +class MCPServerSpec: + id: str + transport: str # "stdio" | "streamable_http" + command: Optional[str] = None + args: Optional[list[str]] = None + env: Optional[dict[str, str]] = None + cwd: Optional[str] = None + url: Optional[str] = None + + +class MCPClientManager: + """ + Lightweight MCP client registry. + + Exposes configured external MCP servers as tools addressable via: + mcp// + """ + + def __init__(self, config): + self._config = config + self._specs: dict[str, MCPServerSpec] = {} + self._load_specs() + + def _load_specs(self) -> None: + self._specs = {} + for raw in self._config.get_mcp_clients_config(): + if not isinstance(raw, dict): + continue + server_id = raw.get("id") + if not server_id or not isinstance(server_id, str): + continue + transport = raw.get("transport", "stdio") + spec = MCPServerSpec( + id=server_id, + transport=transport, + command=raw.get("command"), + args=raw.get("args") or None, + env=raw.get("env") or None, + cwd=raw.get("cwd") or None, + url=raw.get("url") or None, + ) + self._specs[server_id] = spec + + def refresh(self) -> None: + self._load_specs() + + def list_server_ids(self) -> list[str]: + return sorted(self._specs.keys()) + + def get_server_spec(self, server_id: str) -> MCPServerSpec: + if server_id not in self._specs: + raise KeyError(f"Unknown MCP server_id '{server_id}'. Configured: {self.list_server_ids()}") + return self._specs[server_id] + + async def _call_tool_async(self, server_id: str, tool_name: str, arguments: dict[str, Any]) -> Any: + # Import lazily so AIOS can still run without MCP dependencies installed until needed. + from mcp import ClientSession, StdioServerParameters + from mcp.client.stdio import stdio_client + from mcp.client.streamable_http import streamable_http_client + from mcp import types + + spec = self.get_server_spec(server_id) + + if spec.transport == "streamable_http": + if not spec.url: + raise ValueError(f"MCP server '{server_id}' transport streamable_http requires 'url'") + async with streamable_http_client(spec.url) as (read_stream, write_stream, _): + async with ClientSession(read_stream, write_stream) as session: + await session.initialize() + result = await session.call_tool(tool_name, arguments=arguments or {}) + return self._normalize_call_result(result, types) + + # Default: stdio + if not spec.command: + raise ValueError(f"MCP server '{server_id}' transport stdio requires 'command'") + + env = dict(os.environ) + if spec.env: + env.update({k: str(v) for k, v in spec.env.items()}) + + server_params = StdioServerParameters( + command=spec.command, + args=spec.args or [], + env=env, + cwd=spec.cwd, + ) + + async with stdio_client(server_params) as (read, write): + async with ClientSession(read, write) as session: + await session.initialize() + result = await session.call_tool(tool_name, arguments=arguments or {}) + return self._normalize_call_result(result, types) + + def call_tool(self, server_id: str, tool_name: str, arguments: dict[str, Any]) -> Any: + """ + Synchronous wrapper used by the current ToolManager pipeline. + """ + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + + if loop and loop.is_running(): + # Avoid nested event loops; run in a new task and wait. + return asyncio.run_coroutine_threadsafe( + self._call_tool_async(server_id, tool_name, arguments), + loop, + ).result() + + return asyncio.run(self._call_tool_async(server_id, tool_name, arguments)) + + @staticmethod + def _normalize_call_result(result, types_module) -> Any: + # Prefer structuredContent if present. + if getattr(result, "structuredContent", None) is not None: + return result.structuredContent + + # Otherwise join text blocks. + texts: list[str] = [] + for c in getattr(result, "content", []) or []: + if isinstance(c, types_module.TextContent): + texts.append(c.text) + if texts: + return "\n".join(texts) + return str(result) + + +class MCPTool: + """ + Tool wrapper compatible with the existing ToolManager interface (expects .run(params=...)). + """ + + def __init__(self, mcp_manager: MCPClientManager, server_id: str, tool_name: str): + self._mcp = mcp_manager + self._server_id = server_id + self._tool_name = tool_name + + def run(self, params: dict[str, Any]) -> Any: + return self._mcp.call_tool(self._server_id, self._tool_name, params or {}) + diff --git a/requirements.txt b/requirements.txt index 67c9299b7..eef55e7b4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -25,6 +25,7 @@ gdown typing-extensions qdrant-client fastembed +mcp # Pin opentelemetry packages — chromadb pulls in old versions # whose generated protobuf code is incompatible with protobuf v5+ diff --git a/tests/test_mcp_clients_config.py b/tests/test_mcp_clients_config.py new file mode 100644 index 000000000..445afebae --- /dev/null +++ b/tests/test_mcp_clients_config.py @@ -0,0 +1,17 @@ +from aios.config.config_manager import ConfigManager + + +def test_config_bootstrap_has_tool_section_or_defaults(tmp_path, monkeypatch): + """ + Minimal smoke test: ConfigManager should be able to load config via example bootstrap + and provide tool defaults without crashing. + """ + # Point config_manager at a temp config directory by monkeypatching __file__-relative resolution. + # We do this by copying the example config next to a temp config_manager module path simulation + # is out of scope; instead we validate the default accessors are defensive. + cfg = ConfigManager() + tool_cfg = cfg.get_tool_config() + assert isinstance(tool_cfg, dict) + assert isinstance(cfg.get_mcp_server_script_path(), str) + assert isinstance(cfg.get_mcp_clients_config(), list) +