From a28712f0510401ca4ccb97fca51c84ac48b28d47 Mon Sep 17 00:00:00 2001 From: Paul Nechifor Date: Wed, 18 Feb 2026 03:55:23 +0200 Subject: [PATCH] feat(mcp): add mcp client and server --- .gitignore | 2 + dimos/agents/agent.py | 5 +- dimos/agents/conftest.py | 5 - dimos/agents/mcp/README.md | 55 ++++ dimos/{protocol => agents}/mcp/__init__.py | 0 dimos/agents/mcp/conftest.py | 103 ++++++++ .../test_can_call_again_on_error[False].json | 34 +++ .../test_can_call_again_on_error[True].json | 34 +++ .../fixtures/test_can_call_tool[False].json | 22 ++ .../fixtures/test_can_call_tool[True].json | 22 ++ dimos/agents/mcp/fixtures/test_image.json | 23 ++ ...ple_tool_calls_with_multiple_messages.json | 116 ++++++++ dimos/agents/mcp/fixtures/test_prompt.json | 8 + dimos/agents/mcp/mcp_client.py | 250 ++++++++++++++++++ dimos/agents/mcp/mcp_server.py | 197 ++++++++++++++ dimos/agents/mcp/test_mcp_client.py | 210 +++++++++++++++ dimos/agents/mcp/test_mcp_client_unit.py | 145 ++++++++++ .../mcp/test_mcp_server.py} | 73 ++--- dimos/protocol/mcp/README.md | 35 --- dimos/protocol/mcp/__main__.py | 36 --- dimos/protocol/mcp/bridge.py | 53 ---- dimos/protocol/mcp/mcp.py | 139 ---------- dimos/robot/all_blueprints.py | 1 + .../agentic/unitree_go2_agentic_mcp.py | 12 +- pyproject.toml | 3 - uv.lock | 52 ---- 26 files changed, 1260 insertions(+), 375 deletions(-) create mode 100644 dimos/agents/mcp/README.md rename dimos/{protocol => agents}/mcp/__init__.py (100%) create mode 100644 dimos/agents/mcp/conftest.py create mode 100644 dimos/agents/mcp/fixtures/test_can_call_again_on_error[False].json create mode 100644 dimos/agents/mcp/fixtures/test_can_call_again_on_error[True].json create mode 100644 dimos/agents/mcp/fixtures/test_can_call_tool[False].json create mode 100644 dimos/agents/mcp/fixtures/test_can_call_tool[True].json create mode 100644 dimos/agents/mcp/fixtures/test_image.json create mode 100644 dimos/agents/mcp/fixtures/test_multiple_tool_calls_with_multiple_messages.json create mode 100644 dimos/agents/mcp/fixtures/test_prompt.json create mode 100644 dimos/agents/mcp/mcp_client.py create mode 100644 dimos/agents/mcp/mcp_server.py create mode 100644 dimos/agents/mcp/test_mcp_client.py create mode 100644 dimos/agents/mcp/test_mcp_client_unit.py rename dimos/{protocol/mcp/test_mcp_module.py => agents/mcp/test_mcp_server.py} (62%) delete mode 100644 dimos/protocol/mcp/README.md delete mode 100644 dimos/protocol/mcp/__main__.py delete mode 100644 dimos/protocol/mcp/bridge.py delete mode 100644 dimos/protocol/mcp/mcp.py diff --git a/.gitignore b/.gitignore index 24a3dd8919..f97d9f906a 100644 --- a/.gitignore +++ b/.gitignore @@ -69,3 +69,5 @@ yolo11n.pt CLAUDE.MD /assets/teleop_certs/ + +/.mcp.json diff --git a/dimos/agents/agent.py b/dimos/agents/agent.py index 76195ccea0..98f23d7e8d 100644 --- a/dimos/agents/agent.py +++ b/dimos/agents/agent.py @@ -46,9 +46,8 @@ class AgentConfig(ModuleConfig): model_fixture: str | None = None -class Agent(Module): - default_config: type[AgentConfig] = AgentConfig - config: AgentConfig +class Agent(Module[AgentConfig]): + default_config = AgentConfig agent: Out[BaseMessage] human_input: In[str] agent_idle: Out[bool] diff --git a/dimos/agents/conftest.py b/dimos/agents/conftest.py index 23d888b0fe..1be2aadc0c 100644 --- a/dimos/agents/conftest.py +++ b/dimos/agents/conftest.py @@ -31,11 +31,6 @@ FIXTURE_DIR = Path(__file__).parent / "fixtures" -@pytest.fixture -def fixture_dir() -> Path: - return FIXTURE_DIR - - @pytest.fixture def agent_setup(request): coordinator = None diff --git a/dimos/agents/mcp/README.md b/dimos/agents/mcp/README.md new file mode 100644 index 0000000000..f9e887beb1 --- /dev/null +++ b/dimos/agents/mcp/README.md @@ -0,0 +1,55 @@ +# DimOS MCP Server + +Expose DimOS robot skills to Claude Code via Model Context Protocol. + +## Setup + +```bash +uv sync --extra base --extra unitree +``` + +Add to Claude Code (one command) + +```bash +claude mcp add --transport http --scope project dimos http://localhost:9990/mcp +``` + +Verify that it was added: + +```bash +claude mcp list +``` + +## MCP Inspector + +If you want to inspect the server manually, you can use MCP Inspector. + +Install it: + +```bash +npx -y @modelcontextprotocol/inspector +``` + +It will open a browser window. + +Change **Transport Type** to "Streamable HTTP", change **URL** to `http://localhost:9990/mcp`, and **Connection Type** to "Direct". Then click on "Connect". + +## Usage + +**Terminal 1** - Start DimOS: +```bash +uv run dimos run unitree-go2-agentic-mcp +``` + +**Claude Code** - Use robot skills: +``` +> move forward 1 meter +> go to the kitchen +> tag this location as "desk" +``` + +## How It Works + +1. `McpServer` in the blueprint starts a FastAPI server on port 9990 +2. Claude Code connects directly to `http://localhost:9990/mcp` +3. Skills are exposed as MCP tools (e.g., `relative_move`, `navigate_with_text`) diff --git a/dimos/protocol/mcp/__init__.py b/dimos/agents/mcp/__init__.py similarity index 100% rename from dimos/protocol/mcp/__init__.py rename to dimos/agents/mcp/__init__.py diff --git a/dimos/agents/mcp/conftest.py b/dimos/agents/mcp/conftest.py new file mode 100644 index 0000000000..532ef16592 --- /dev/null +++ b/dimos/agents/mcp/conftest.py @@ -0,0 +1,103 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from pathlib import Path +from threading import Event + +from dotenv import load_dotenv +from langchain_core.messages.base import BaseMessage +import pytest + +from dimos.agents.agent_test_runner import AgentTestRunner +from dimos.agents.mcp.mcp_client import McpClient +from dimos.agents.mcp.mcp_server import McpServer +from dimos.core.blueprints import autoconnect +from dimos.core.global_config import global_config +from dimos.core.transport import pLCMTransport + +load_dotenv() + +FIXTURE_DIR = Path(__file__).parent / "fixtures" + + +@pytest.fixture +def agent_setup(request): + coordinator = None + transports: list[pLCMTransport] = [] + unsubs: list = [] + recording = bool(os.getenv("RECORD")) + + def fn( + *, + blueprints, + messages: list[BaseMessage], + dask: bool = False, + system_prompt: str | None = None, + fixture: str | None = None, + ) -> list[BaseMessage]: + history: list[BaseMessage] = [] + finished_event = Event() + + agent_transport: pLCMTransport = pLCMTransport("/agent") + finished_transport: pLCMTransport = pLCMTransport("/finished") + transports.extend([agent_transport, finished_transport]) + + def on_message(msg: BaseMessage) -> None: + history.append(msg) + + unsubs.append(agent_transport.subscribe(on_message)) + unsubs.append(finished_transport.subscribe(lambda _: finished_event.set())) + + # Derive fixture path from test name if not explicitly provided. + if fixture is not None: + fixture_path = FIXTURE_DIR / fixture + else: + fixture_path = FIXTURE_DIR / f"{request.node.name}.json" + + client_kwargs: dict = {"system_prompt": system_prompt} + + if recording or fixture_path.exists(): + client_kwargs["model_fixture"] = str(fixture_path) + + blueprint = autoconnect( + *blueprints, + McpServer.blueprint(), + McpClient.blueprint(**client_kwargs), + AgentTestRunner.blueprint(messages=messages), + ) + + global_config.update( + viewer_backend="none", + dask=dask, + ) + + nonlocal coordinator + coordinator = blueprint.build() + + if not finished_event.wait(60): + raise TimeoutError("Timed out waiting for agent to finish processing messages.") + + return history + + yield fn + + if coordinator is not None: + coordinator.stop() + + for transport in transports: + transport.stop() + + for unsub in unsubs: + unsub() diff --git a/dimos/agents/mcp/fixtures/test_can_call_again_on_error[False].json b/dimos/agents/mcp/fixtures/test_can_call_again_on_error[False].json new file mode 100644 index 0000000000..8cfe2f69c7 --- /dev/null +++ b/dimos/agents/mcp/fixtures/test_can_call_again_on_error[False].json @@ -0,0 +1,34 @@ +{ + "responses": [ + { + "content": "", + "tool_calls": [ + { + "name": "register_user", + "args": { + "name": "Paul" + }, + "id": "call_NrrizXSIFaeCLuG9i05IwDy3", + "type": "tool_call" + } + ] + }, + { + "content": "", + "tool_calls": [ + { + "name": "register_user", + "args": { + "name": "paul" + }, + "id": "call_2QPx4GsL61Xjrggbq7afXTjn", + "type": "tool_call" + } + ] + }, + { + "content": "The user named 'paul' has been registered successfully.", + "tool_calls": [] + } + ] +} diff --git a/dimos/agents/mcp/fixtures/test_can_call_again_on_error[True].json b/dimos/agents/mcp/fixtures/test_can_call_again_on_error[True].json new file mode 100644 index 0000000000..3d3765f43a --- /dev/null +++ b/dimos/agents/mcp/fixtures/test_can_call_again_on_error[True].json @@ -0,0 +1,34 @@ +{ + "responses": [ + { + "content": "", + "tool_calls": [ + { + "name": "register_user", + "args": { + "name": "Paul" + }, + "id": "call_XSy1Dx1dGtQv5zPaEJtb2hd7", + "type": "tool_call" + } + ] + }, + { + "content": "", + "tool_calls": [ + { + "name": "register_user", + "args": { + "name": "paul" + }, + "id": "call_aYFug1g3TATnaYus9HUVxoQS", + "type": "tool_call" + } + ] + }, + { + "content": "The user named \"paul\" has been registered successfully.", + "tool_calls": [] + } + ] +} diff --git a/dimos/agents/mcp/fixtures/test_can_call_tool[False].json b/dimos/agents/mcp/fixtures/test_can_call_tool[False].json new file mode 100644 index 0000000000..7d1ac3075b --- /dev/null +++ b/dimos/agents/mcp/fixtures/test_can_call_tool[False].json @@ -0,0 +1,22 @@ +{ + "responses": [ + { + "content": "", + "tool_calls": [ + { + "name": "add", + "args": { + "x": 33333, + "y": 100 + }, + "id": "call_RssRDDd9apDjNoVLz4jRLVk0", + "type": "tool_call" + } + ] + }, + { + "content": "The result of 33333 + 100 is 33433.", + "tool_calls": [] + } + ] +} diff --git a/dimos/agents/mcp/fixtures/test_can_call_tool[True].json b/dimos/agents/mcp/fixtures/test_can_call_tool[True].json new file mode 100644 index 0000000000..d375c82235 --- /dev/null +++ b/dimos/agents/mcp/fixtures/test_can_call_tool[True].json @@ -0,0 +1,22 @@ +{ + "responses": [ + { + "content": "", + "tool_calls": [ + { + "name": "add", + "args": { + "x": 33333, + "y": 100 + }, + "id": "call_pzzddF9mBynGYZVdCmGHOB5V", + "type": "tool_call" + } + ] + }, + { + "content": "The result of 33333 + 100 is 33433.", + "tool_calls": [] + } + ] +} diff --git a/dimos/agents/mcp/fixtures/test_image.json b/dimos/agents/mcp/fixtures/test_image.json new file mode 100644 index 0000000000..0e4816b8ee --- /dev/null +++ b/dimos/agents/mcp/fixtures/test_image.json @@ -0,0 +1,23 @@ +{ + "responses": [ + { + "content": "", + "tool_calls": [ + { + "name": "take_a_picture", + "args": {}, + "id": "call_7Qwsr8QMLWhKRMektcGiKYf7", + "type": "tool_call" + } + ] + }, + { + "content": "I've taken a picture. Let me analyze and describe it for you.\nThe image features an expansive outdoor stadium. From the camera's perspective, the word 'stadium' best matches the image. Is there anything else you'd like to know or do?", + "tool_calls": [] + }, + { + "content": "The image shows a group of people sitting and enjoying their time at an outdoor cafe. Therefore, the word 'cafe' best matches the image.", + "tool_calls": [] + } + ] +} diff --git a/dimos/agents/mcp/fixtures/test_multiple_tool_calls_with_multiple_messages.json b/dimos/agents/mcp/fixtures/test_multiple_tool_calls_with_multiple_messages.json new file mode 100644 index 0000000000..5c0d551e13 --- /dev/null +++ b/dimos/agents/mcp/fixtures/test_multiple_tool_calls_with_multiple_messages.json @@ -0,0 +1,116 @@ +{ + "responses": [ + { + "content": "", + "tool_calls": [ + { + "name": "locate_person", + "args": { + "name": "John" + }, + "id": "call_eOoKTtyvvXBk171ro4bXzW5C", + "type": "tool_call" + } + ] + }, + { + "content": "", + "tool_calls": [ + { + "name": "register_person", + "args": { + "name": "John" + }, + "id": "call_tTB5A3q60teaBrdonRvCwcwM", + "type": "tool_call" + } + ] + }, + { + "content": "", + "tool_calls": [ + { + "name": "locate_person", + "args": { + "name": "John" + }, + "id": "call_uEhafkL3f7BLQKhRuZlEAany", + "type": "tool_call" + } + ] + }, + { + "content": "", + "tool_calls": [ + { + "name": "go_to_location", + "args": { + "description": "kitchen" + }, + "id": "call_oxnH4gCGi6aSeVLPrhnp31yP", + "type": "tool_call" + } + ] + }, + { + "content": "I have moved to the kitchen where John is located.", + "tool_calls": [] + }, + { + "content": "", + "tool_calls": [ + { + "name": "locate_person", + "args": { + "name": "Jane" + }, + "id": "call_2HinxBmffnafloaP4b7DkBZW", + "type": "tool_call" + } + ] + }, + { + "content": "", + "tool_calls": [ + { + "name": "register_person", + "args": { + "name": "Jane" + }, + "id": "call_XtHavMmgpzrhmVi3XB6RUFrW", + "type": "tool_call" + } + ] + }, + { + "content": "", + "tool_calls": [ + { + "name": "locate_person", + "args": { + "name": "Jane" + }, + "id": "call_fRHHO4cPWDXi4IvQ4qQqidwT", + "type": "tool_call" + } + ] + }, + { + "content": "", + "tool_calls": [ + { + "name": "go_to_location", + "args": { + "description": "living room" + }, + "id": "call_Hcc7C0FMWS8rfKwMP0sUL7XN", + "type": "tool_call" + } + ] + }, + { + "content": "I have moved to the living room where Jane is located.", + "tool_calls": [] + } + ] +} diff --git a/dimos/agents/mcp/fixtures/test_prompt.json b/dimos/agents/mcp/fixtures/test_prompt.json new file mode 100644 index 0000000000..acb77fe350 --- /dev/null +++ b/dimos/agents/mcp/fixtures/test_prompt.json @@ -0,0 +1,8 @@ +{ + "responses": [ + { + "content": "Hello! My name is Johnny. How can I assist you today?", + "tool_calls": [] + } + ] +} diff --git a/dimos/agents/mcp/mcp_client.py b/dimos/agents/mcp/mcp_client.py new file mode 100644 index 0000000000..7c5eda5302 --- /dev/null +++ b/dimos/agents/mcp/mcp_client.py @@ -0,0 +1,250 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dataclasses import dataclass +from queue import Empty, Queue +from threading import Event, RLock, Thread +import time +from typing import Any +import uuid + +import httpx +from langchain.agents import create_agent +from langchain_core.messages import HumanMessage +from langchain_core.messages.base import BaseMessage +from langchain_core.tools import StructuredTool +from langgraph.graph.state import CompiledStateGraph +from reactivex.disposable import Disposable + +from dimos.agents.system_prompt import SYSTEM_PROMPT +from dimos.agents.utils import pretty_print_langchain_message +from dimos.core.core import rpc +from dimos.core.module import Module, ModuleConfig +from dimos.core.rpc_client import RPCClient +from dimos.core.stream import In, Out +from dimos.utils.logging_config import setup_logger +from dimos.utils.sequential_ids import SequentialIds + +logger = setup_logger() + + +@dataclass +class McpClientConfig(ModuleConfig): + system_prompt: str | None = SYSTEM_PROMPT + model: str = "gpt-4o" + model_fixture: str | None = None + mcp_server_url: str = "http://localhost:9990/mcp" + + +class McpClient(Module[McpClientConfig]): + default_config = McpClientConfig + agent: Out[BaseMessage] + human_input: In[str] + agent_idle: Out[bool] + + _lock: RLock + _state_graph: CompiledStateGraph[Any, Any, Any, Any] | None + _message_queue: Queue[BaseMessage] + _history: list[BaseMessage] + _thread: Thread + _stop_event: Event + _http_client: httpx.Client + _seq_ids: SequentialIds + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._lock = RLock() + self._state_graph = None + self._message_queue = Queue() + self._history = [] + self._thread = Thread( + target=self._thread_loop, + name=f"{self.__class__.__name__}-thread", + daemon=True, + ) + self._stop_event = Event() + self._http_client = httpx.Client(timeout=120.0) + self._seq_ids = SequentialIds() + + def __reduce__(self) -> Any: + return (self.__class__, (), {}) + + def _mcp_request(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]: + body: dict[str, Any] = { + "jsonrpc": "2.0", + "id": self._seq_ids.next(), + "method": method, + } + if params is not None: + body["params"] = params + + resp = self._http_client.post(self.config.mcp_server_url, json=body) + resp.raise_for_status() + data = resp.json() + + if "error" in data: + raise RuntimeError(f"MCP error {data['error']['code']}: {data['error']['message']}") + + result: dict[str, Any] = data.get("result") + return result + + def _fetch_tools(self, timeout: float = 60.0, interval: float = 1.0) -> list[StructuredTool]: + result = self._try_fetch_tools(timeout=timeout, interval=interval) + if result is None: + raise RuntimeError( + f"Failed to fetch tools from MCP server {self.config.mcp_server_url}" + ) + + tools = [self._mcp_tool_to_langchain(t) for t in result.get("tools", [])] + + if not tools: + logger.warning("No tools found from MCP server.") + else: + tool_names = [t.name for t in tools] + logger.info("Discovered tools from MCP server.", tools=tool_names, n_tools=len(tools)) + + return tools + + def _try_fetch_tools(self, timeout: float, interval: float) -> dict[str, Any] | None: + deadline = time.monotonic() + timeout + + while True: + try: + self._mcp_request("initialize") + break + except (httpx.ConnectError, httpx.RemoteProtocolError): + if time.monotonic() >= deadline: + return None + time.sleep(interval) + + return self._mcp_request("tools/list") + + def _mcp_tool_to_langchain(self, mcp_tool: dict[str, Any]) -> StructuredTool: + name = mcp_tool["name"] + description = mcp_tool.get("description", "") + input_schema = mcp_tool.get("inputSchema", {"type": "object", "properties": {}}) + + def call_tool(**kwargs: Any) -> str: + result = self._mcp_request("tools/call", {"name": name, "arguments": kwargs}) + content = result.get("content", []) + parts = [c.get("text", "") for c in content if c.get("type") == "text"] + text = "\n".join(parts) + + # Images need to be added to the history separately because they + # cannot be included in the tool response for OpenAI models and + # probably others. + for item in content: + if item.get("type") != "text": + uuid_ = str(uuid.uuid4()) + text += f"Tool call started with UUID: {uuid_}. You will be updated with the result soon." + _append_image_to_history(self, name, uuid_, item) + + return text + + return StructuredTool( + name=name, + description=description, + func=call_tool, + args_schema=input_schema, + ) + + @rpc + def start(self) -> None: + super().start() + + def _on_human_input(string: str) -> None: + self._message_queue.put(HumanMessage(content=string)) + + self._disposables.add(Disposable(self.human_input.subscribe(_on_human_input))) + + @rpc + def on_system_modules(self, _modules: list[RPCClient]) -> None: + tools = self._fetch_tools() + + model: str | Any = self.config.model + if self.config.model_fixture is not None: + from dimos.agents.testing import MockModel + + model = MockModel(json_path=self.config.model_fixture) + + with self._lock: + self._state_graph = create_agent( + model=model, + tools=tools, + system_prompt=self.config.system_prompt, + ) + self._thread.start() + + @rpc + def stop(self) -> None: + self._stop_event.set() + if self._thread.is_alive(): + self._thread.join(timeout=2.0) + self._http_client.close() + super().stop() + + @rpc + def add_message(self, message: BaseMessage) -> None: + self._message_queue.put(message) + + def _thread_loop(self) -> None: + while not self._stop_event.is_set(): + try: + message = self._message_queue.get(timeout=0.5) + except Empty: + continue + + with self._lock: + if not self._state_graph: + raise ValueError("No state graph initialized") + self._process_message(self._state_graph, message) + + def _process_message( + self, state_graph: CompiledStateGraph[Any, Any, Any, Any], message: BaseMessage + ) -> None: + self.agent_idle.publish(False) + self._history.append(message) + pretty_print_langchain_message(message) + self.agent.publish(message) + + for update in state_graph.stream({"messages": self._history}, stream_mode="updates"): + for node_output in update.values(): + for msg in node_output.get("messages", []): + self._history.append(msg) + pretty_print_langchain_message(msg) + self.agent.publish(msg) + + if self._message_queue.empty(): + self.agent_idle.publish(True) + + +def _append_image_to_history( + mcp_client: McpClient, func_name: str, uuid_: str, result: Any +) -> None: + mcp_client.add_message( + HumanMessage( + content=[ + { + "type": "text", + "text": f"This is the artefact for the '{func_name}' tool with UUID:={uuid_}.", + }, + result, + ] + ) + ) + + +mcp_client = McpClient.blueprint + +__all__ = ["McpClient", "McpClientConfig", "mcp_client"] diff --git a/dimos/agents/mcp/mcp_server.py b/dimos/agents/mcp/mcp_server.py new file mode 100644 index 0000000000..1f8ce92888 --- /dev/null +++ b/dimos/agents/mcp/mcp_server.py @@ -0,0 +1,197 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import asyncio +import json +from typing import TYPE_CHECKING, Any + +from fastapi import FastAPI +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import JSONResponse +from starlette.responses import Response +import uvicorn + +from dimos.utils.logging_config import setup_logger + +logger = setup_logger() + + +from dimos.core import Module, rpc # noqa: I001 +from dimos.core.rpc_client import RpcCall, RPCClient + +from starlette.requests import Request # noqa: TC002 + +if TYPE_CHECKING: + import concurrent.futures + + from dimos.core.module import SkillInfo + + +app = FastAPI() +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_methods=["POST"], + allow_headers=["*"], +) +app.state.skills = [] +app.state.rpc_calls = {} + + +def _jsonrpc_result(req_id: Any, result: Any) -> dict[str, Any]: + return {"jsonrpc": "2.0", "id": req_id, "result": result} + + +def _jsonrpc_result_text(req_id: Any, text: str) -> dict[str, Any]: + return _jsonrpc_result(req_id, {"content": [{"type": "text", "text": text}]}) + + +def _jsonrpc_error(req_id: Any, code: int, message: str) -> dict[str, Any]: + return {"jsonrpc": "2.0", "id": req_id, "error": {"code": code, "message": message}} + + +def _handle_initialize(req_id: Any) -> dict[str, Any]: + return _jsonrpc_result( + req_id, + { + "protocolVersion": "2025-11-25", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "dimensional", "version": "1.0.0"}, + }, + ) + + +def _handle_tools_list(req_id: Any, skills: list[SkillInfo]) -> dict[str, Any]: + tools = [] + + for skill in skills: + schema = json.loads(skill.args_schema) + description = schema.pop("description", None) + schema.pop("title", None) + tool = {"name": skill.func_name, "inputSchema": schema} + if description: + tool["description"] = description + tools.append(tool) + + return _jsonrpc_result(req_id, {"tools": tools}) + + +async def _handle_tools_call( + req_id: Any, params: dict[str, Any], rpc_calls: dict[str, Any] +) -> dict[str, Any]: + name = params.get("name", "") + args: dict[str, Any] = params.get("arguments") or {} + + rpc_call = rpc_calls.get(name) + if rpc_call is None: + return _jsonrpc_result_text(req_id, f"Tool not found: {name}") + + try: + result = await asyncio.get_event_loop().run_in_executor(None, lambda: rpc_call(**args)) + except Exception as e: + logger.exception("Error running tool", tool_name=name, exc_info=True) + return _jsonrpc_result_text(req_id, f"Error running tool '{name}': {e}") + + if result is None: + return _jsonrpc_result_text(req_id, "It has started. You will be updated later.") + + if hasattr(result, "agent_encode"): + return _jsonrpc_result(req_id, {"content": result.agent_encode()}) + + return _jsonrpc_result_text(req_id, str(result)) + + +async def handle_request( + request: dict[str, Any], + skills: list[SkillInfo], + rpc_calls: dict[str, Any], +) -> dict[str, Any] | None: + """Handle a single MCP JSON-RPC request. + + Returns None for JSON-RPC notifications (no ``id``), which must not + receive a response. + """ + method = request.get("method", "") + params = request.get("params", {}) or {} + req_id = request.get("id") + + # JSON-RPC notifications have no "id" – the server must not reply. + if "id" not in request: + return None + + if method == "initialize": + return _handle_initialize(req_id) + if method == "tools/list": + return _handle_tools_list(req_id, skills) + if method == "tools/call": + return await _handle_tools_call(req_id, params, rpc_calls) + return _jsonrpc_error(req_id, -32601, f"Unknown: {method}") + + +@app.post("/mcp") +async def mcp_endpoint(request: Request) -> Response: + raw = await request.body() + try: + body = json.loads(raw) + except Exception: + logger.exception("POST /mcp JSON parse failed") + return JSONResponse( + {"jsonrpc": "2.0", "id": None, "error": {"code": -32700, "message": "Parse error"}}, + status_code=400, + ) + result = await handle_request(body, request.app.state.skills, request.app.state.rpc_calls) + if result is None: + return Response(status_code=204) + return JSONResponse(result) + + +class McpServer(Module): + def __init__(self) -> None: + super().__init__() + self._uvicorn_server: uvicorn.Server | None = None + self._serve_future: concurrent.futures.Future[None] | None = None + + @rpc + def start(self) -> None: + super().start() + self._start_server() + + @rpc + def stop(self) -> None: + if self._uvicorn_server: + self._uvicorn_server.should_exit = True + loop = self._loop + if loop is not None and self._serve_future is not None: + self._serve_future.result(timeout=5.0) + self._uvicorn_server = None + self._serve_future = None + super().stop() + + @rpc + def on_system_modules(self, modules: list[RPCClient]) -> None: + assert self.rpc is not None + app.state.skills = [skill for module in modules for skill in (module.get_skills() or [])] + app.state.rpc_calls = { + skill.func_name: RpcCall(None, self.rpc, skill.func_name, skill.class_name, []) + for skill in app.state.skills + } + + def _start_server(self, port: int = 9990) -> None: + config = uvicorn.Config(app, host="0.0.0.0", port=port, log_level="info") + server = uvicorn.Server(config) + self._uvicorn_server = server + loop = self._loop + assert loop is not None + self._serve_future = asyncio.run_coroutine_threadsafe(server.serve(), loop) diff --git a/dimos/agents/mcp/test_mcp_client.py b/dimos/agents/mcp/test_mcp_client.py new file mode 100644 index 0000000000..be4a09d5b9 --- /dev/null +++ b/dimos/agents/mcp/test_mcp_client.py @@ -0,0 +1,210 @@ +# Copyright 2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from langchain_core.messages import HumanMessage +import pytest + +from dimos.agents.annotation import skill +from dimos.core.module import Module +from dimos.msgs.sensor_msgs import Image +from dimos.utils.data import get_data + + +class Adder(Module): + @skill + def add(self, x: int, y: int) -> str: + """adds x and y.""" + return str(x + y) + + +@pytest.mark.integration +@pytest.mark.parametrize("dask", [False, True]) +def test_can_call_tool(dask, agent_setup): + history = agent_setup( + blueprints=[Adder.blueprint()], + messages=[HumanMessage("What is 33333 + 100? Use the tool.")], + dask=dask, + ) + + assert "33433" in history[-1].content + + +class UserRegistration(Module): + def __init__(self): + super().__init__() + self._first_call = True + self._use_upper = False + + @skill + def register_user(self, name: str) -> str: + """registers a user by name.""" + + # If the agent calls with "paul" or "Paul", always say it's the wrong way + # to force it to try again. + + if self._first_call: + self._first_call = False + self._use_upper = not name[0].isupper() + + if self._use_upper and not name[0].isupper(): + raise ValueError("Names must start with an uppercase letter.") + if not self._use_upper and name[0].isupper(): + raise ValueError("The names must only use lowercase letters.") + + return "User name registered successfully." + + +@pytest.mark.integration +@pytest.mark.parametrize("dask", [False, True]) +def test_can_call_again_on_error(dask, agent_setup): + history = agent_setup( + blueprints=[UserRegistration.blueprint()], + messages=[ + HumanMessage( + "Register a user named 'Paul'. If there are errors, just try again until you succeed." + ) + ], + dask=dask, + ) + + assert any(message.content == "User name registered successfully." for message in history) + + +class MultipleTools(Module): + def __init__(self): + super().__init__() + self._people = {"Ben": "office", "Bob": "garage"} + + @skill + def register_person(self, name: str) -> str: + """Registers a person by name.""" + if name.lower() == "john": + self._people[name] = "kitchen" + elif name.lower() == "jane": + self._people[name] = "living room" + return f"'{name}' has been registered." + + @skill + def locate_person(self, name: str) -> str: + """Locates a person by name.""" + if name not in self._people: + known_people = list(self._people.keys()) + return ( + f"Error: '{name}' is not registered. People cannot be located until they've " + f"been registered in the system. People known so far: {', '.join(known_people)}. " + "Use register_person to register a person." + ) + return f"'{name}' is located at '{self._people[name]}'." + + +class NavigationSkill(Module): + @skill + def go_to_location(self, description: str) -> str: + """Go to a location by a description.""" + if description.strip().lower() not in ["kitchen", "living room"]: + return f"Error: Unknown location description: '{description}'." + return f"Going to the {description}." + + +@pytest.mark.integration +def test_multiple_tool_calls_with_multiple_messages(agent_setup): + history = agent_setup( + blueprints=[MultipleTools.blueprint(), NavigationSkill.blueprint()], + messages=[ + HumanMessage( + "You are a robot assistant. Move to the location where John is. Don't ask me for feedback, just go there." + ), + HumanMessage("Nice job. You did it. Now go to the location where Jane is."), + ], + ) + + # Collect all go_to_location calls from the history + go_to_location_calls = [] + for message in history: + if hasattr(message, "tool_calls"): + for tool_call in message.tool_calls: + if tool_call["name"] == "go_to_location": + go_to_location_calls.append(tool_call) + + # Find the index of the second HumanMessage to split first/second prompt + second_human_idx = None + human_count = 0 + for i, message in enumerate(history): + if isinstance(message, HumanMessage): + human_count += 1 + if human_count == 2: + second_human_idx = i + break + + # Collect go_to_location calls before and after the second prompt + calls_after_first_prompt = [] + calls_after_second_prompt = [] + for i, message in enumerate(history): + if hasattr(message, "tool_calls"): + for tool_call in message.tool_calls: + if tool_call["name"] == "go_to_location": + if i < second_human_idx: + calls_after_first_prompt.append(tool_call) + else: + calls_after_second_prompt.append(tool_call) + + # After the first prompt, go_to_location should be called with "kitchen" + assert len(calls_after_first_prompt) == 1 + assert "kitchen" in calls_after_first_prompt[0]["args"]["description"].lower() + + # After the second prompt, go_to_location should be called with "living room" + assert len(calls_after_second_prompt) == 1 + assert "living room" in calls_after_second_prompt[0]["args"]["description"].lower() + + # There should be exactly two go_to_location calls total + assert len(go_to_location_calls) == 2 + + +@pytest.mark.integration +def test_prompt(agent_setup): + history = agent_setup( + blueprints=[], + messages=[HumanMessage("What is your name?")], + system_prompt="You are a helpful assistant named Johnny.", + ) + + assert "Johnny" in history[-1].content + + +class Visualizer(Module): + @skill + def take_a_picture(self) -> Image: + """Takes a picture.""" + return Image.from_file(get_data("cafe-smol.jpg")).to_rgb() + + +@pytest.mark.integration +def test_image(agent_setup): + history = agent_setup( + blueprints=[Visualizer.blueprint()], + messages=[ + HumanMessage( + "What do you see? Take a picture using your camera and describe it. " + "Please mention one of the words which best match the image: " + "'stadium', 'cafe', 'battleship'." + ) + ], + system_prompt="You are a helpful assistant that can use a camera to take pictures.", + ) + + response = history[-1].content.lower() + assert "cafe" in response + assert "stadium" not in response + assert "battleship" not in response diff --git a/dimos/agents/mcp/test_mcp_client_unit.py b/dimos/agents/mcp/test_mcp_client_unit.py new file mode 100644 index 0000000000..8cd888f851 --- /dev/null +++ b/dimos/agents/mcp/test_mcp_client_unit.py @@ -0,0 +1,145 @@ +# Copyright 2025-2026 Dimensional Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import json +from unittest.mock import MagicMock, patch + +import pytest + +from dimos.agents.mcp.mcp_client import McpClient +from dimos.utils.sequential_ids import SequentialIds + + +def _mock_post(url: str, **kwargs: object) -> MagicMock: + """Return a fake httpx response based on the JSON-RPC method.""" + body = kwargs.get("json") or (kwargs.get("content") and json.loads(kwargs["content"])) + assert isinstance(body, dict) + method = body["method"] + req_id = body["id"] + + result: object + if method == "initialize": + result = { + "protocolVersion": "2024-11-05", + "capabilities": {"tools": {}}, + "serverInfo": {"name": "dimensional", "version": "1.0.0"}, + } + elif method == "tools/list": + result = { + "tools": [ + { + "name": "add", + "description": "Add two numbers", + "inputSchema": { + "type": "object", + "properties": { + "x": {"type": "integer"}, + "y": {"type": "integer"}, + }, + "required": ["x", "y"], + }, + }, + { + "name": "greet", + "description": "Say hello", + "inputSchema": { + "type": "object", + "properties": { + "name": {"type": "string"}, + }, + }, + }, + ] + } + elif method == "tools/call": + name = body["params"]["name"] + args = body["params"].get("arguments", {}) + if name == "add": + text = str(args.get("x", 0) + args.get("y", 0)) + elif name == "greet": + text = f"Hello, {args.get('name', 'world')}!" + else: + text = "Skill not found" + result = {"content": [{"type": "text", "text": text}]} + else: + resp = MagicMock() + resp.status_code = 200 + resp.raise_for_status = MagicMock() + resp.json.return_value = { + "jsonrpc": "2.0", + "id": req_id, + "error": {"code": -32601, "message": f"Unknown: {method}"}, + } + return resp + + resp = MagicMock() + resp.status_code = 200 + resp.raise_for_status = MagicMock() + resp.json.return_value = {"jsonrpc": "2.0", "id": req_id, "result": result} + return resp + + +@pytest.fixture +def mcp_client() -> McpClient: + """Build an McpClient wired to the mock MCP post handler.""" + mock_http = MagicMock() + mock_http.post.side_effect = _mock_post + + with patch("dimos.agents.mcp.mcp_client.httpx.Client", return_value=mock_http): + client = McpClient.__new__(McpClient) + + client._http_client = mock_http + client._seq_ids = SequentialIds() + client.config = MagicMock() + client.config.mcp_server_url = "http://localhost:9990/mcp" + return client + + +def test_fetch_tools_from_mcp_server(mcp_client: McpClient) -> None: + tools = mcp_client._fetch_tools() + + assert len(tools) == 2 + assert tools[0].name == "add" + assert tools[1].name == "greet" + + +def test_tool_invocation_via_mcp(mcp_client: McpClient) -> None: + tools = mcp_client._fetch_tools() + add_tool = next(t for t in tools if t.name == "add") + greet_tool = next(t for t in tools if t.name == "greet") + + assert add_tool.func(x=2, y=3) == "5" + assert greet_tool.func(name="Alice") == "Hello, Alice!" + + +def test_mcp_request_error_propagation(mcp_client: McpClient) -> None: + def error_post(url: str, **kwargs: object) -> MagicMock: + resp = MagicMock() + resp.status_code = 200 + resp.raise_for_status = MagicMock() + resp.json.return_value = { + "jsonrpc": "2.0", + "id": 1, + "error": {"code": -32601, "message": "Unknown: bad/method"}, + } + return resp + + mcp_client._http_client.post.side_effect = error_post + + try: + mcp_client._mcp_request("bad/method") + raise AssertionError("Expected RuntimeError") + except RuntimeError as e: + assert "Unknown: bad/method" in str(e) diff --git a/dimos/protocol/mcp/test_mcp_module.py b/dimos/agents/mcp/test_mcp_server.py similarity index 62% rename from dimos/protocol/mcp/test_mcp_module.py rename to dimos/agents/mcp/test_mcp_server.py index 050e24f13b..1cbca9e3e4 100644 --- a/dimos/protocol/mcp/test_mcp_module.py +++ b/dimos/agents/mcp/test_mcp_server.py @@ -16,34 +16,25 @@ import asyncio import json -from pathlib import Path from unittest.mock import MagicMock +from dimos.agents.mcp.mcp_server import handle_request from dimos.core.module import SkillInfo -from dimos.protocol.mcp.mcp import MCPModule -def _make_mcp(skills: list[SkillInfo], call_results: dict[str, object]) -> MCPModule: - """Create an MCPModule with pre-populated skills and mock RPC calls.""" - mcp = MCPModule.__new__(MCPModule) - mcp._skills = skills - mcp._rpc_calls = {} +def _make_rpc_calls( + skills: list[SkillInfo], call_results: dict[str, object] +) -> dict[str, MagicMock]: + """Create mock RPC calls for the given skills.""" + rpc_calls: dict[str, MagicMock] = {} for skill in skills: mock_call = MagicMock() if skill.func_name in call_results: mock_call.return_value = call_results[skill.func_name] else: mock_call.return_value = None - mcp._rpc_calls[skill.func_name] = mock_call - return mcp - - -def test_unitree_blueprint_has_mcp() -> None: - contents = Path( - "dimos/robot/unitree/go2/blueprints/agentic/unitree_go2_agentic_mcp.py" - ).read_text() - assert "agentic_mcp" in contents - assert "MCPModule.blueprint()" in contents + rpc_calls[skill.func_name] = mock_call + return rpc_calls def test_mcp_module_request_flow() -> None: @@ -56,20 +47,21 @@ def test_mcp_module_request_flow() -> None: } ) skills = [SkillInfo(class_name="TestSkills", func_name="add", args_schema=schema)] + rpc_calls = _make_rpc_calls(skills, {"add": 5}) - mcp = _make_mcp(skills, {"add": 5}) - - response = asyncio.run(mcp._handle_request({"method": "tools/list", "id": 1})) + response = asyncio.run(handle_request({"method": "tools/list", "id": 1}, skills, rpc_calls)) assert response["result"]["tools"][0]["name"] == "add" assert response["result"]["tools"][0]["description"] == "Add two numbers" response = asyncio.run( - mcp._handle_request( + handle_request( { "method": "tools/call", "id": 2, "params": {"name": "add", "arguments": {"x": 2, "y": 3}}, - } + }, + skills, + rpc_calls, ) ) assert response["result"]["content"][0]["text"] == "5" @@ -82,49 +74,40 @@ def test_mcp_module_handles_errors() -> None: SkillInfo(class_name="TestSkills", func_name="fail_skill", args_schema=schema), ] - mcp = _make_mcp(skills, {"ok_skill": "done"}) - mcp._rpc_calls["fail_skill"] = MagicMock(side_effect=RuntimeError("boom")) + rpc_calls = _make_rpc_calls(skills, {"ok_skill": "done"}) + rpc_calls["fail_skill"] = MagicMock(side_effect=RuntimeError("boom")) # All skills listed - response = asyncio.run(mcp._handle_request({"method": "tools/list", "id": 1})) + response = asyncio.run(handle_request({"method": "tools/list", "id": 1}, skills, rpc_calls)) tool_names = {tool["name"] for tool in response["result"]["tools"]} assert "ok_skill" in tool_names assert "fail_skill" in tool_names # Error skill returns error text response = asyncio.run( - mcp._handle_request( - {"method": "tools/call", "id": 2, "params": {"name": "fail_skill", "arguments": {}}} + handle_request( + {"method": "tools/call", "id": 2, "params": {"name": "fail_skill", "arguments": {}}}, + skills, + rpc_calls, ) ) - assert "Error:" in response["result"]["content"][0]["text"] + assert "Error running tool" in response["result"]["content"][0]["text"] assert "boom" in response["result"]["content"][0]["text"] # Unknown skill returns not found response = asyncio.run( - mcp._handle_request( - {"method": "tools/call", "id": 3, "params": {"name": "no_such", "arguments": {}}} + handle_request( + {"method": "tools/call", "id": 3, "params": {"name": "no_such", "arguments": {}}}, + skills, + rpc_calls, ) ) assert "not found" in response["result"]["content"][0]["text"].lower() def test_mcp_module_initialize_and_unknown() -> None: - mcp = _make_mcp([], {}) - - response = asyncio.run(mcp._handle_request({"method": "initialize", "id": 1})) + response = asyncio.run(handle_request({"method": "initialize", "id": 1}, [], {})) assert response["result"]["serverInfo"]["name"] == "dimensional" - response = asyncio.run(mcp._handle_request({"method": "unknown/method", "id": 2})) + response = asyncio.run(handle_request({"method": "unknown/method", "id": 2}, [], {})) assert response["error"]["code"] == -32601 - - -def test_mcp_module_invalid_tool_name() -> None: - mcp = _make_mcp([], {}) - - response = asyncio.run( - mcp._handle_request( - {"method": "tools/call", "id": 1, "params": {"name": 123, "arguments": {}}} - ) - ) - assert response["error"]["code"] == -32602 diff --git a/dimos/protocol/mcp/README.md b/dimos/protocol/mcp/README.md deleted file mode 100644 index 233e852669..0000000000 --- a/dimos/protocol/mcp/README.md +++ /dev/null @@ -1,35 +0,0 @@ -# DimOS MCP Server - -Expose DimOS robot skills to Claude Code via Model Context Protocol. - -## Setup - -```bash -uv sync --extra base --extra unitree -``` - -Add to Claude Code (one command): -```bash -claude mcp add --transport stdio dimos --scope project -- python -m dimos.protocol.mcp -``` - - -## Usage - -**Terminal 1** - Start DimOS: -```bash -uv run dimos run unitree-go2-agentic-mcp -``` - -**Claude Code** - Use robot skills: -``` -> move forward 1 meter -> go to the kitchen -> tag this location as "desk" -``` - -## How It Works - -1. `MCPModule` in the blueprint starts a TCP server on port 9990 -2. Claude Code spawns the bridge (`--bridge`) which connects to `localhost:9990` -3. Skills are exposed as MCP tools (e.g., `relative_move`, `navigate_with_text`) diff --git a/dimos/protocol/mcp/__main__.py b/dimos/protocol/mcp/__main__.py deleted file mode 100644 index a58e59d367..0000000000 --- a/dimos/protocol/mcp/__main__.py +++ /dev/null @@ -1,36 +0,0 @@ -# Copyright 2025-2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""CLI entry point for Dimensional MCP Bridge. - -Connects Claude Code (or other MCP clients) to a running DimOS agent. - -Usage: - python -m dimos.protocol.mcp # Bridge to running DimOS on default port -""" - -from __future__ import annotations - -import asyncio - -from dimos.protocol.mcp.bridge import main as bridge_main - - -def main() -> None: - """Main entry point - connects to running DimOS via bridge.""" - asyncio.run(bridge_main()) - - -if __name__ == "__main__": - main() diff --git a/dimos/protocol/mcp/bridge.py b/dimos/protocol/mcp/bridge.py deleted file mode 100644 index 0b09997798..0000000000 --- a/dimos/protocol/mcp/bridge.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright 2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -"""MCP Bridge - Connects stdio (Claude Code) to TCP (DimOS Agent).""" - -import asyncio -import os -import sys - -DEFAULT_PORT = 9990 - - -async def main() -> None: - port = int(os.environ.get("MCP_PORT", DEFAULT_PORT)) - host = os.environ.get("MCP_HOST", "localhost") - - reader, writer = await asyncio.open_connection(host, port) - sys.stderr.write(f"MCP Bridge connected to {host}:{port}\n") - - async def stdin_to_tcp() -> None: - loop = asyncio.get_event_loop() - while True: - line = await loop.run_in_executor(None, sys.stdin.readline) - if not line: - break - writer.write(line.encode()) - await writer.drain() - - async def tcp_to_stdout() -> None: - while True: - data = await reader.readline() - if not data: - break - sys.stdout.write(data.decode()) - sys.stdout.flush() - - await asyncio.gather(stdin_to_tcp(), tcp_to_stdout()) - - -if __name__ == "__main__": - asyncio.run(main()) diff --git a/dimos/protocol/mcp/mcp.py b/dimos/protocol/mcp/mcp.py deleted file mode 100644 index 78d19c64db..0000000000 --- a/dimos/protocol/mcp/mcp.py +++ /dev/null @@ -1,139 +0,0 @@ -# Copyright 2025-2026 Dimensional Inc. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -from __future__ import annotations - -import asyncio -import json -from typing import TYPE_CHECKING, Any - -from dimos.core import Module, rpc -from dimos.core.rpc_client import RpcCall, RPCClient - -if TYPE_CHECKING: - from dimos.core.module import SkillInfo - - -class MCPModule(Module): - _skills: list[SkillInfo] - _rpc_calls: dict[str, RpcCall] - - def __init__(self, *args: Any, **kwargs: Any) -> None: - super().__init__(*args, **kwargs) - self._skills = [] - self._rpc_calls = {} - self._server: asyncio.AbstractServer | None = None - self._server_future: object | None = None - - @rpc - def start(self) -> None: - super().start() - self._start_server() - - @rpc - def stop(self) -> None: - if self._server: - self._server.close() - loop = self._loop - assert loop is not None - asyncio.run_coroutine_threadsafe(self._server.wait_closed(), loop).result() - self._server = None - if self._server_future and hasattr(self._server_future, "cancel"): - self._server_future.cancel() - super().stop() - - @rpc - def on_system_modules(self, modules: list[RPCClient]) -> None: - assert self.rpc is not None - self._skills = [skill for module in modules for skill in (module.get_skills() or [])] - self._rpc_calls = { - skill.func_name: RpcCall(None, self.rpc, skill.func_name, skill.class_name, []) - for skill in self._skills - } - - def _start_server(self, port: int = 9990) -> None: - async def handle_client(reader, writer) -> None: # type: ignore[no-untyped-def] - while True: - if not (data := await reader.readline()): - break - response = await self._handle_request(json.loads(data.decode())) - writer.write(json.dumps(response).encode() + b"\n") - await writer.drain() - writer.close() - - async def start_server() -> None: - self._server = await asyncio.start_server(handle_client, "0.0.0.0", port) - await self._server.serve_forever() - - loop = self._loop - assert loop is not None - self._server_future = asyncio.run_coroutine_threadsafe(start_server(), loop) - - async def _handle_request(self, request: dict[str, Any]) -> dict[str, Any]: - method = request.get("method", "") - params = request.get("params", {}) or {} - req_id = request.get("id") - if method == "initialize": - init_result = { - "protocolVersion": "2024-11-05", - "capabilities": {"tools": {}}, - "serverInfo": {"name": "dimensional", "version": "1.0.0"}, - } - return {"jsonrpc": "2.0", "id": req_id, "result": init_result} - if method == "tools/list": - tools = [] - for skill in self._skills: - schema = json.loads(skill.args_schema) - tools.append( - { - "name": skill.func_name, - "description": schema.get("description", ""), - "inputSchema": schema, - } - ) - return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": tools}} - if method == "tools/call": - name = params.get("name") - args = params.get("arguments") or {} - if not isinstance(name, str): - return { - "jsonrpc": "2.0", - "id": req_id, - "error": {"code": -32602, "message": "Missing or invalid tool name"}, - } - if not isinstance(args, dict): - args = {} - rpc_call = self._rpc_calls.get(name) - if rpc_call is None: - return { - "jsonrpc": "2.0", - "id": req_id, - "result": {"content": [{"type": "text", "text": "Skill not found"}]}, - } - try: - result = await asyncio.get_event_loop().run_in_executor( - None, lambda: rpc_call(**args) - ) - text = str(result) if result is not None else "Completed" - except Exception as e: - text = f"Error: {e}" - return { - "jsonrpc": "2.0", - "id": req_id, - "result": {"content": [{"type": "text", "text": text}]}, - } - return { - "jsonrpc": "2.0", - "id": req_id, - "error": {"code": -32601, "message": f"Unknown: {method}"}, - } diff --git a/dimos/robot/all_blueprints.py b/dimos/robot/all_blueprints.py index 19d7e7db29..750c6f00e1 100644 --- a/dimos/robot/all_blueprints.py +++ b/dimos/robot/all_blueprints.py @@ -112,6 +112,7 @@ "keyboard_teleop_module": "dimos.teleop.keyboard.keyboard_teleop_module", "manipulation_module": "dimos.manipulation.manipulation_module", "mapper": "dimos.robot.unitree.type.map", + "mcp_client": "dimos.agents.mcp.mcp_client", "mid360_module": "dimos.hardware.sensors.lidar.livox.module", "navigation_skill": "dimos.agents.skills.navigation", "object_scene_registration_module": "dimos.perception.object_scene_registration", diff --git a/dimos/robot/unitree/go2/blueprints/agentic/unitree_go2_agentic_mcp.py b/dimos/robot/unitree/go2/blueprints/agentic/unitree_go2_agentic_mcp.py index bbc3e4c216..e75b31e511 100644 --- a/dimos/robot/unitree/go2/blueprints/agentic/unitree_go2_agentic_mcp.py +++ b/dimos/robot/unitree/go2/blueprints/agentic/unitree_go2_agentic_mcp.py @@ -13,13 +13,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +from dimos.agents.mcp.mcp_client import mcp_client +from dimos.agents.mcp.mcp_server import McpServer from dimos.core.blueprints import autoconnect -from dimos.protocol.mcp.mcp import MCPModule -from dimos.robot.unitree.go2.blueprints.agentic.unitree_go2_agentic import unitree_go2_agentic +from dimos.robot.unitree.go2.blueprints.agentic._common_agentic import _common_agentic +from dimos.robot.unitree.go2.blueprints.smart.unitree_go2_spatial import unitree_go2_spatial unitree_go2_agentic_mcp = autoconnect( - unitree_go2_agentic, - MCPModule.blueprint(), + unitree_go2_spatial, + McpServer.blueprint(), + mcp_client(), + _common_agentic, ) __all__ = ["unitree_go2_agentic_mcp"] diff --git a/pyproject.toml b/pyproject.toml index 68c15fe221..8f0d3be488 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -141,9 +141,6 @@ agents = [ "openai", "openai-whisper", "sounddevice", - - # MCP Server - "mcp>=1.0.0", ] web = [ diff --git a/uv.lock b/uv.lock index 9db7e7332e..4de833a774 100644 --- a/uv.lock +++ b/uv.lock @@ -1794,7 +1794,6 @@ agents = [ { name = "langchain-ollama" }, { name = "langchain-openai" }, { name = "langchain-text-splitters" }, - { name = "mcp" }, { name = "ollama" }, { name = "openai" }, { name = "openai-whisper" }, @@ -1815,7 +1814,6 @@ base = [ { name = "langchain-openai" }, { name = "langchain-text-splitters" }, { name = "lap" }, - { name = "mcp" }, { name = "moondream" }, { name = "mujoco" }, { name = "ollama" }, @@ -2013,7 +2011,6 @@ unitree = [ { name = "langchain-openai" }, { name = "langchain-text-splitters" }, { name = "lap" }, - { name = "mcp" }, { name = "moondream" }, { name = "mujoco" }, { name = "ollama" }, @@ -2091,7 +2088,6 @@ requires-dist = [ { name = "llvmlite", specifier = ">=0.42.0" }, { name = "lxml-stubs", marker = "extra == 'dev'", specifier = ">=0.5.1,<1" }, { name = "matplotlib", marker = "extra == 'manipulation'", specifier = ">=3.7.1" }, - { name = "mcp", marker = "extra == 'agents'", specifier = ">=1.0.0" }, { name = "md-babel-py", marker = "extra == 'dev'", specifier = "==1.1.1" }, { name = "moondream", marker = "extra == 'perception'" }, { name = "mujoco", marker = "extra == 'sim'", specifier = ">=3.3.4" }, @@ -3211,15 +3207,6 @@ http2 = [ { name = "h2" }, ] -[[package]] -name = "httpx-sse" -version = "0.4.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/0f/4c/751061ffa58615a32c31b2d82e8482be8dd4a89154f003147acee90f2be9/httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d", size = 15943, upload-time = "2025-10-10T21:48:22.271Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" }, -] - [[package]] name = "huggingface-hub" version = "0.36.2" @@ -4890,31 +4877,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/27/1a/1f68f9ba0c207934b35b86a8ca3aad8395a3d6dd7921c0686e23853ff5a9/mccabe-0.7.0-py2.py3-none-any.whl", hash = "sha256:6c2d30ab6be0e4a46919781807b4f0d834ebdd6c6e3dca0bda5a15f863427b6e", size = 7350, upload-time = "2022-01-24T01:14:49.62Z" }, ] -[[package]] -name = "mcp" -version = "1.26.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "anyio" }, - { name = "httpx" }, - { name = "httpx-sse" }, - { name = "jsonschema" }, - { name = "pydantic" }, - { name = "pydantic-settings" }, - { name = "pyjwt", extra = ["crypto"] }, - { name = "python-multipart" }, - { name = "pywin32", marker = "sys_platform == 'win32'" }, - { name = "sse-starlette" }, - { name = "starlette" }, - { name = "typing-extensions" }, - { name = "typing-inspection" }, - { name = "uvicorn", marker = "sys_platform != 'emscripten'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/fc/6d/62e76bbb8144d6ed86e202b5edd8a4cb631e7c8130f3f4893c3f90262b10/mcp-1.26.0.tar.gz", hash = "sha256:db6e2ef491eecc1a0d93711a76f28dec2e05999f93afd48795da1c1137142c66", size = 608005, upload-time = "2026-01-24T19:40:32.468Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fd/d9/eaa1f80170d2b7c5ba23f3b59f766f3a0bb41155fbc32a69adfa1adaaef9/mcp-1.26.0-py3-none-any.whl", hash = "sha256:904a21c33c25aa98ddbeb47273033c435e595bbacfdb177f4bd87f6dceebe1ca", size = 233615, upload-time = "2026-01-24T19:40:30.652Z" }, -] - [[package]] name = "md-babel-py" version = "1.1.1" @@ -7696,20 +7658,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" }, ] -[[package]] -name = "pyjwt" -version = "2.11.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/5c/5a/b46fa56bf322901eee5b0454a34343cdbdae202cd421775a8ee4e42fd519/pyjwt-2.11.0.tar.gz", hash = "sha256:35f95c1f0fbe5d5ba6e43f00271c275f7a1a4db1dab27bf708073b75318ea623", size = 98019, upload-time = "2026-01-30T19:59:55.694Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6f/01/c26ce75ba460d5cd503da9e13b21a33804d38c2165dec7b716d06b13010c/pyjwt-2.11.0-py3-none-any.whl", hash = "sha256:94a6bde30eb5c8e04fee991062b534071fd1439ef58d2adc9ccb823e7bcd0469", size = 28224, upload-time = "2026-01-30T19:59:54.539Z" }, -] - -[package.optional-dependencies] -crypto = [ - { name = "cryptography" }, -] - [[package]] name = "pylibsrtp" version = "1.0.0"