Added MCP Server Integration

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/<server_id>/<tool_name>`
+- **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. 
 

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/<id>/<tool_name>
+  #
+  # 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: []

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:
         """

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/<server_id>/<tool_name>
+            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>")
+            _, 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

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/<server_id>/<tool_name>
+    """
+
+    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 {})
+

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+

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)
+