a working mcp integration

agent/README.md

@@ -19,6 +19,5 @@ context_manager/
 └── manager.py         # Message history management
 
 config.py              # Config with model_name + tools
-tools/                 # Tool implementations (todo)
 utils/                 # Logging, etc
 ```

agent/config.py

@@ -1,14 +1,18 @@
+import json
+import os
+import re
+from typing import Any, Union
+
+from dotenv import load_dotenv
+from fastmcp.mcp_config import (
+    RemoteMCPServer,
+    StdioMCPServer,
+)
 from litellm import Tool
 from pydantic import BaseModel
 
-
-class MCPServerConfig(BaseModel):
-    """Configuration for an MCP server"""
-
-    name: str
-    command: str
-    args: list[str] = []
-    env: dict[str, str] | None = None
+# These two are the canonical server config types for MCP servers.
+MCPServerConfig = Union[StdioMCPServer, RemoteMCPServer]
 
 
 class Config(BaseModel):
@@ -17,10 +21,58 @@ class Config(BaseModel):
     model_name: str
     tools: list[Tool] = []
     system_prompt_path: str = ""
-    mcp_servers: list[MCPServerConfig] = []
+    mcpServers: dict[str, MCPServerConfig] = {}
+
+
+def substitute_env_vars(obj: Any) -> Any:
+    """
+    Recursively substitute environment variables in any data structure.
+
+    Supports ${VAR_NAME} syntax for required variables and ${VAR_NAME:-default} for optional.
+    """
+    if isinstance(obj, str):
+        pattern = r"\$\{([^}:]+)(?::(-)?([^}]*))?\}"
+
+        def replacer(match):
+            var_name = match.group(1)
+            has_default = match.group(2) is not None
+            default_value = match.group(3) if has_default else None
+
+            env_value = os.environ.get(var_name)
+
+            if env_value is not None:
+                return env_value
+            elif has_default:
+                return default_value or ""
+            else:
+                raise ValueError(
+                    f"Environment variable '{var_name}' is not set. "
+                    f"Add it to your .env file."
+                )
+
+        return re.sub(pattern, replacer, obj)
+
+    elif isinstance(obj, dict):
+        return {key: substitute_env_vars(value) for key, value in obj.items()}
+
+    elif isinstance(obj, list):
+        return [substitute_env_vars(item) for item in obj]
+
+    return obj
 
 
 def load_config(config_path: str = "config.json") -> Config:
-    """Load configuration from file"""
+    """
+    Load configuration with environment variable substitution.
+
+    Use ${VAR_NAME} in your JSON for any secret.
+    Automatically loads from .env file.
+    """
+    # Load environment variables from .env file
+    load_dotenv()
+
     with open(config_path, "r") as f:
-        return Config.model_validate_json(f.read())
+        raw_config = json.load(f)
+
+    config_with_env = substitute_env_vars(raw_config)
+    return Config.model_validate(config_with_env)

agent/config_mcp_example.json

@@ -2,12 +2,20 @@
   "model_name": "anthropic/claude-sonnet-4-5-20250929",
   "tools": [],
   "system_prompt_path": "",
-  "mcp_servers": [
-    {
-      "name": "weather",
-      "command": "python",
-      "args": ["path/to/your/weather_server.py"],
-      "env": null
+  "mcpServers": {
+    "hf-mcp-server": {
+      "transport": "http",
+      "url": "https://huggingface.co/mcp?login",
+      "headers": {
+        "Authorization": "Bearer ${HF_TOKEN}"
+      }
+    },
+    "playwright": {
+      "transport": "stdio",
+      "command": "npx",
+      "args": [
+        "@playwright/mcp@latest"
+      ]
     }
-  ]
-}
+  }
+}

agent/core/__init__.py

@@ -3,12 +3,9 @@ Core agent implementation
 Contains the main agent logic, decision-making, and orchestration
 """
 
-from agent.core.mcp_client import McpClient, McpConnectionManager
 from agent.core.tools import ToolRouter, ToolSpec, create_builtin_tools
 
 __all__ = [
-    "McpClient",
-    "McpConnectionManager",
     "ToolRouter",
     "ToolSpec",
     "create_builtin_tools",

agent/core/agent_loop.py

@@ -9,6 +9,7 @@ from litellm import ChatCompletionMessageToolCall, Message, ModelResponse, acomp
 
 from agent.config import Config
 from agent.core.session import Event, OpType, Session
+from agent.core.tools import ToolRouter
 
 ToolCall = ChatCompletionMessageToolCall
 
@@ -54,13 +55,18 @@ class Handlers:
                         assistant_msg = Message(role="assistant", content=content)
                         session.context_manager.add_message(assistant_msg)
                         await session.send_event(
-                            Event(event_type="assistant_message", data={"content": content})
+                            Event(
+                                event_type="assistant_message",
+                                data={"content": content},
+                            )
                         )
                     break
 
                 # Add assistant message with tool calls to history
                 # LiteLLM will format this correctly for the provider
-                assistant_msg = Message(role="assistant", content=content, tool_calls=tool_calls)
+                assistant_msg = Message(
+                    role="assistant", content=content, tool_calls=tool_calls
+                )
                 session.context_manager.add_message(assistant_msg)
 
                 if content:
@@ -80,7 +86,7 @@ class Handlers:
                         )
                     )
 
-                    output, success = await session.tool_router.execute_tool(
+                    output, success = await session.tool_router.call_tool(
                         tool_name, tool_args
                     )
 
@@ -98,8 +104,7 @@ class Handlers:
                             event_type="tool_output",
                             data={
                                 "tool": tool_name,
-                                "output": output[:200]
-                                + ("..." if len(output) > 200 else ""),
+                                "output": output,
                                 "success": success,
                             },
                         )
@@ -108,8 +113,13 @@ class Handlers:
                 iteration += 1
 
             except Exception as e:
+                import traceback
+
                 await session.send_event(
-                    Event(event_type="error", data={"error": str(e)})
+                    Event(
+                        event_type="error",
+                        data={"error": str(e + "\n" + traceback.format_exc())},
+                    )
                 )
                 break
 
@@ -195,58 +205,26 @@ async def process_submission(session: Session, submission) -> bool:
 async def submission_loop(
     submission_queue: asyncio.Queue,
     event_queue: asyncio.Queue,
-    tool_router=None,
     config: Config | None = None,
+    tool_router: ToolRouter | None = None,
 ) -> None:
     """
     Main agent loop - processes submissions and dispatches to handlers.
     This is the core of the agent (like submission_loop in codex.rs:1259-1340)
     """
-    # Import here to avoid circular imports
-    from agent.core.mcp_client import McpConnectionManager
-    from agent.core.tools import ToolRouter, create_builtin_tools
-
-    # Initialize MCP and tools
-    if tool_router is None:
-        mcp_manager = McpConnectionManager()
-
-        # Add MCP servers from config
-        if config and config.mcp_servers:
-            print("🔌 Initializing MCP connections...")
-            for server_config in config.mcp_servers:
-                try:
-                    await mcp_manager.add_server(
-                        server_name=server_config.name,
-                        command=server_config.command,
-                        args=server_config.args,
-                        env=server_config.env,
-                    )
-                except Exception as e:
-                    print(
-                        f"⚠️  Failed to connect to MCP server {server_config.name}: {e}"
-                    )
-
-        # Create tool router
-        tool_router = ToolRouter(mcp_manager)
-
-        # Register built-in tools
-        for tool in create_builtin_tools():
-            tool_router.register_tool(tool)
-
-        # Register MCP tools
-        tool_router.register_mcp_tools()
-
-        print(f"📦 Registered {len(tool_router.tools)} tools:")
-        for tool_name in tool_router.tools.keys():
-            print(f"   - {tool_name}")
 
     # Create session and assign tool router
     session = Session(event_queue, config=config)
     session.tool_router = tool_router
     print("🤖 Agent loop started")
 
-    try:
-        # Main processing loop
+    # Main processing loop
+    async with tool_router:
+        # Emit ready event after initialization
+        await session.send_event(
+            Event(event_type="ready", data={"message": "Agent initialized"})
+        )
+
         while session.is_running:
             submission = await submission_queue.get()
 
@@ -261,9 +239,5 @@ async def submission_loop(
                 await session.send_event(
                     Event(event_type="error", data={"error": str(e)})
                 )
-    finally:
-        # Cleanup MCP connections
-        if hasattr(tool_router, "mcp_manager") and tool_router.mcp_manager:
-            await tool_router.mcp_manager.shutdown_all()
 
     print("🛑 Agent loop exited")

agent/core/mcp_client.py

@@ -1,164 +0,0 @@
-"""
-MCP (Model Context Protocol) client integration for the agent
-Based on the official MCP SDK implementation
-"""
-
-import os
-from contextlib import AsyncExitStack
-from typing import Any, Optional
-
-from mcp import ClientSession, StdioServerParameters, types
-from mcp.client.stdio import stdio_client
-
-
-class McpClient:
-    """
-    Client for connecting to MCP servers using the official MCP SDK.
-    Based on codex-rs/core/src/mcp_connection_manager.rs
-    """
-
-    def __init__(
-        self,
-        server_name: str,
-        command: str,
-        args: list[str] | None = None,
-        env: dict[str, str] | None = None,
-    ):
-        self.server_name = server_name
-        self.command = command
-        self.args = args or []
-        self.env = env or {}
-        self.session: Optional[ClientSession] = None
-        self.tools: dict[str, dict[str, Any]] = {}
-        self.exit_stack = AsyncExitStack()
-
-    async def start(self) -> None:
-        """Start the MCP server connection using official SDK"""
-        # Merge environment variables
-        full_env = {**dict(os.environ), **self.env} if self.env else None
-
-        # Create server parameters
-        server_params = StdioServerParameters(
-            command=self.command,
-            args=self.args,
-            env=full_env,
-        )
-
-        # Connect using stdio_client
-        read, write = await self.exit_stack.enter_async_context(
-            stdio_client(server_params)
-        )
-
-        # Create session
-        self.session = await self.exit_stack.enter_async_context(
-            ClientSession(read, write)
-        )
-
-        # Initialize
-        await self.session.initialize()
-
-        # List available tools
-        tools_result = await self.session.list_tools()
-        for tool in tools_result.tools:
-            qualified_name = f"mcp__{self.server_name}__{tool.name}"
-            self.tools[qualified_name] = {
-                "name": tool.name,
-                "description": tool.description or "",
-                "inputSchema": tool.inputSchema,
-            }
-
-    async def call_tool(
-        self, tool_name: str, arguments: dict[str, Any]
-    ) -> tuple[str, bool]:
-        """Execute a tool on the MCP server"""
-        if not self.session:
-            return "Client not connected", False
-
-        # Strip the mcp__servername__ prefix to get the actual tool name
-        actual_tool_name = tool_name.split("__")[-1]
-
-        try:
-            result = await self.session.call_tool(actual_tool_name, arguments)
-
-            # Extract text from content
-            text_parts = []
-            for content in result.content:
-                if isinstance(content, types.TextContent):
-                    text_parts.append(content.text)
-                elif isinstance(content, types.ImageContent):
-                    text_parts.append(f"[Image: {content.mimeType}]")
-                elif isinstance(content, types.EmbeddedResource):
-                    text_parts.append(f"[Resource: {content.resource}]")
-
-            output = "\n".join(text_parts) if text_parts else str(result.content)
-            success = not result.isError
-
-            return output, success
-        except Exception as e:
-            return f"Tool call failed: {str(e)}", False
-
-    def get_tools(self) -> dict[str, dict[str, Any]]:
-        """Get all available tools from this server"""
-        return self.tools.copy()
-
-    async def shutdown(self) -> None:
-        """Shutdown the MCP server connection"""
-        await self.exit_stack.aclose()
-
-
-class McpConnectionManager:
-    """
-    Manages multiple MCP server connections.
-    Based on codex-rs/core/src/mcp_connection_manager.rs
-    """
-
-    def __init__(self):
-        self.clients: dict[str, McpClient] = {}
-
-    async def add_server(
-        self,
-        server_name: str,
-        command: str,
-        args: list[str] | None = None,
-        env: dict[str, str] | None = None,
-    ) -> bool:
-        """Add and start an MCP server"""
-        try:
-            client = McpClient(server_name, command, args, env)
-            await client.start()
-            self.clients[server_name] = client
-            print(
-                f"✅ MCP server '{server_name}' connected with {len(client.tools)} tools"
-            )
-            return True
-        except Exception as e:
-            print(f"❌ Failed to start MCP server '{server_name}': {e}")
-            return False
-
-    def list_all_tools(self) -> dict[str, dict[str, Any]]:
-        """Aggregate tools from all connected servers"""
-        all_tools = {}
-        for client in self.clients.values():
-            all_tools.update(client.get_tools())
-        return all_tools
-
-    async def call_tool(
-        self, tool_name: str, arguments: dict[str, Any]
-    ) -> tuple[str, bool]:
-        """Route tool call to the appropriate MCP server"""
-        # Extract server name from qualified tool name: mcp__servername__toolname
-        if tool_name.startswith("mcp__"):
-            parts = tool_name.split("__")
-            if len(parts) >= 3:
-                server_name = parts[1]
-                if server_name in self.clients:
-                    return await self.clients[server_name].call_tool(
-                        tool_name, arguments
-                    )
-
-        return "Unknown MCP tool", False
-
-    async def shutdown_all(self) -> None:
-        """Shutdown all MCP servers"""
-        for client in self.clients.values():
-            await client.shutdown()

agent/core/tools.py

@@ -4,10 +4,67 @@ Provides ToolSpec and ToolRouter for managing both built-in and MCP tools
 """
 
 import subprocess
+import warnings
 from dataclasses import dataclass
 from typing import Any, Awaitable, Callable, Optional
 
-from agent.core.mcp_client import McpConnectionManager
+from fastmcp import Client
+from mcp.types import EmbeddedResource, ImageContent, TextContent
+
+from agent.config import MCPServerConfig
+
+# Suppress aiohttp deprecation warning
+warnings.filterwarnings(
+    "ignore", category=DeprecationWarning, module="aiohttp.connector"
+)
+
+
+def convert_mcp_content_to_string(content: list) -> str:
+    """
+    Convert MCP content blocks to a string format compatible with LLM messages.
+
+    Based on FastMCP documentation, content can be:
+    - TextContent: has .text field
+    - ImageContent: has .data and .mimeType fields
+    - EmbeddedResource: has .resource field with .text or .blob
+
+    Args:
+        content: List of MCP content blocks
+
+    Returns:
+        String representation of the content suitable for LLM consumption
+    """
+    if not content:
+        return ""
+
+    parts = []
+    for item in content:
+        if isinstance(item, TextContent):
+            # Extract text from TextContent blocks
+            parts.append(item.text)
+        elif isinstance(item, ImageContent):
+            # TODO: Handle images
+            # For images, include a description with MIME type
+            parts.append(f"[Image: {item.mimeType}]")
+        elif isinstance(item, EmbeddedResource):
+            # TODO: Handle embedded resources
+            # For embedded resources, try to extract text
+            resource = item.resource
+            if hasattr(resource, "text") and resource.text:
+                parts.append(resource.text)
+            elif hasattr(resource, "blob") and resource.blob:
+                parts.append(
+                    f"[Binary data: {resource.mimeType if hasattr(resource, 'mimeType') else 'unknown'}]"
+                )
+            else:
+                parts.append(
+                    f"[Resource: {resource.uri if hasattr(resource, 'uri') else 'unknown'}]"
+                )
+        else:
+            # Fallback: try to convert to string
+            parts.append(str(item))
+
+    return "\n".join(parts)
 
 
 @dataclass
@@ -26,30 +83,34 @@ class ToolRouter:
     Based on codex-rs/core/src/tools/router.rs
     """
 
-    def __init__(self, mcp_manager: Optional[McpConnectionManager] = None):
+    def __init__(self, mcp_servers: dict[str, MCPServerConfig]):
         self.tools: dict[str, ToolSpec] = {}
-        self.mcp_manager = mcp_manager
-
-    def register_tool(self, spec: ToolSpec) -> None:
-        """Register a tool with its handler"""
-        self.tools[spec.name] = spec
-
-    def register_mcp_tools(self) -> None:
-        """Register all MCP tools from the connection manager"""
-        if not self.mcp_manager:
-            return
-
-        mcp_tools = self.mcp_manager.list_all_tools()
-        for tool_name, tool_def in mcp_tools.items():
-            spec = ToolSpec(
-                name=tool_name,
-                description=tool_def.get("description", ""),
-                parameters=tool_def.get(
-                    "inputSchema", {"type": "object", "properties": {}}
-                ),
-                handler=None,  # MCP tools use the manager
+        self.mcp_servers: dict[str, dict[str, Any]] = {}
+
+        for tool in create_builtin_tools():
+            self.register_tool(tool)
+
+        if mcp_servers:
+            mcp_servers_payload = {}
+            for name, server in mcp_servers.items():
+                mcp_servers_payload[name] = server.model_dump()
+            self.mcp_client = Client({"mcpServers": mcp_servers_payload})
+        self._mcp_initialized = False
+
+    def register_tool(self, tool: ToolSpec) -> None:
+        self.tools[tool.name] = tool
+
+    async def register_mcp_tools(self) -> None:
+        tools = await self.mcp_client.list_tools()
+        for tool in tools:
+            self.register_tool(
+                ToolSpec(
+                    name=tool.name,
+                    description=tool.description,
+                    parameters=tool.inputSchema,
+                    handler=None,
+                )
             )
-            self.tools[tool_name] = spec
 
     def get_tool_specs_for_llm(self) -> list[dict[str, Any]]:
         """Get tool specifications in OpenAI format"""
@@ -67,24 +128,42 @@ class ToolRouter:
             )
         return specs
 
-    async def execute_tool(
-        self, tool_name: str, arguments: dict[str, Any]
-    ) -> tuple[str, bool]:
-        """Execute a tool by name"""
-        if tool_name not in self.tools:
-            return f"Unknown tool: {tool_name}", False
+    async def __aenter__(self) -> "ToolRouter":
+        if self.mcp_client is not None:
+            await self.mcp_client.__aenter__()
+            await self.mcp_client.initialize()
+            await self.register_mcp_tools()
+            self._mcp_initialized = True
+        print(f"MCP initialized: {self._mcp_initialized}")
+        return self
 
-        tool = self.tools[tool_name]
+    async def __aexit__(self, exc_type, exc, tb) -> None:
+        if self.mcp_client is not None:
+            await self.mcp_client.__aexit__(exc_type, exc, tb)
+            self._mcp_initialized = False
 
-        # MCP tool
-        if tool_name.startswith("mcp__") and self.mcp_manager:
-            return await self.mcp_manager.call_tool(tool_name, arguments)
+    async def call_tool(
+        self, tool_name: str, arguments: dict[str, Any]
+    ) -> tuple[str, bool]:
+        """
+        Call a tool and return (output_string, success_bool).
 
-        # Built-in tool with handler
-        if tool.handler:
+        For MCP tools, converts the CallToolResult content blocks to a string.
+        For built-in tools, calls their handler directly.
+        """
+        # Check if this is a built-in tool with a handler
+        tool = self.tools.get(tool_name)
+        if tool and tool.handler:
             return await tool.handler(arguments)
 
-        return "Tool has no handler", False
+        # Otherwise, use MCP client
+        if self._mcp_initialized:
+            result = await self.mcp_client.call_tool(tool_name, arguments)
+            # Convert MCP content blocks to string
+            output = convert_mcp_content_to_string(result.content)
+            return output, not result.is_error
+
+        return "MCP client not initialized", False
 
 
 # ============================================================================

agent/main.py

@@ -4,11 +4,13 @@ Interactive CLI chat with the agent
 
 import asyncio
 from dataclasses import dataclass
+from pathlib import Path
 from typing import Any, Optional
 
-from agent.config import Config
+from agent.config import load_config
 from agent.core.agent_loop import submission_loop
 from agent.core.session import OpType
+from agent.core.tools import ToolRouter
 
 
 @dataclass
@@ -28,7 +30,9 @@ class Submission:
 
 
 async def event_listener(
-    event_queue: asyncio.Queue, turn_complete_event: asyncio.Event
+    event_queue: asyncio.Queue,
+    turn_complete_event: asyncio.Event,
+    ready_event: asyncio.Event,
 ) -> None:
     """Background task that listens for events and displays them"""
     while True:
@@ -36,7 +40,10 @@ async def event_listener(
             event = await event_queue.get()
 
             # Display event
-            if event.event_type == "assistant_message":
+            if event.event_type == "ready":
+                print("✅ Agent ready")
+                ready_event.set()
+            elif event.event_type == "assistant_message":
                 content = event.data.get("content", "") if event.data else ""
                 if content:
                     print(f"\n🤖 Assistant: {content}")
@@ -54,7 +61,11 @@ async def event_listener(
                 print("✅ Turn complete\n")
                 turn_complete_event.set()
             elif event.event_type == "error":
-                error = event.data.get("error", "Unknown error") if event.data else "Unknown error"
+                error = (
+                    event.data.get("error", "Unknown error")
+                    if event.data
+                    else "Unknown error"
+                )
                 print(f"❌ Error: {error}")
                 turn_complete_event.set()
             elif event.event_type == "shutdown":
@@ -78,7 +89,6 @@ async def get_user_input() -> str:
 
 async def main():
     """Interactive chat with the agent"""
-
     print("=" * 60)
     print("🤖 Interactive Agent Chat")
     print("=" * 60)
@@ -88,30 +98,36 @@ async def main():
     submission_queue = asyncio.Queue()
     event_queue = asyncio.Queue()
 
-    # Event to signal turn completion
+    # Events to signal agent state
     turn_complete_event = asyncio.Event()
-    turn_complete_event.set()  # Ready for first input
+    turn_complete_event.set()
+    ready_event = asyncio.Event()
 
     # Start agent loop in background
+    config_path = Path(__file__).parent / "config_mcp_example.json"
+    config = load_config(config_path)
+
+    # Create tool router
+    print(f"Config: {config.mcpServers}")
+    tool_router = ToolRouter(config.mcpServers)
+
     agent_task = asyncio.create_task(
         submission_loop(
             submission_queue,
             event_queue,
-            config=Config(
-                model_name="anthropic/claude-sonnet-4-5-20250929",
-                tools=[],
-                system_prompt_path="",
-            ),
+            config=config,
+            tool_router=tool_router,
         )
     )
 
     # Start event listener in background
     listener_task = asyncio.create_task(
-        event_listener(event_queue, turn_complete_event)
+        event_listener(event_queue, turn_complete_event, ready_event)
     )
 
     # Wait for agent to initialize
-    await asyncio.sleep(0.2)
+    print("⏳ Initializing agent...")
+    await ready_event.wait()
 
     submission_id = 0
 

agent/tools/__init__.py

@@ -1,8 +0,0 @@
-"""
-Agent tools and actions
-Tools are the actions the agent can take to interact with the environment
-"""
-
-from agent.tools.base import BaseTool, ToolRegistry
-
-__all__ = ["BaseTool", "ToolRegistry"]

pyproject.toml

@@ -14,5 +14,5 @@ dependencies = [
     "python-dotenv>=1.2.1",
     "datasets>=4.3.0",
     "huggingface-hub>=1.0.1",
-    "mcp>=1.21.0",
+    "fastmcp>=2.4.0",
 ]