feat: MCP server enhancements + dimos mcp CLI + agent context (DIM-686, DIM-687)

dimos/agents/mcp/mcp_adapter.py

@@ -0,0 +1,190 @@
+# 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.
+
+"""Lightweight MCP JSON-RPC client adapter.
+
+``McpAdapter`` provides a typed Python interface to a running MCP server.
+It is used by:
+
+* The ``dimos mcp`` CLI commands
+* Integration / e2e tests
+* Any code that needs to talk to a local MCP server
+
+Usage::
+
+    adapter = McpAdapter("http://localhost:9990/mcp")
+    adapter.wait_for_ready(timeout=10)
+    tools = adapter.list_tools()
+    result = adapter.call_tool("echo", {"message": "hi"})
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+import uuid
+
+import requests
+
+from dimos.utils.logging_config import setup_logger
+
+logger = setup_logger()
+
+DEFAULT_TIMEOUT = 30
+
+
+class McpError(Exception):
+    """Raised when the MCP server returns a JSON-RPC error."""
+
+    def __init__(self, message: str, code: int | None = None) -> None:
+        self.code = code
+        super().__init__(message)
+
+
+class McpAdapter:
+    """Thin JSON-RPC client for a running MCP server."""
+
+    def __init__(self, url: str | None = None, timeout: int = DEFAULT_TIMEOUT) -> None:
+        if url is None:
+            from dimos.core.global_config import global_config
+
+            url = f"http://localhost:{global_config.mcp_port}/mcp"
+        self.url = url
+        self.timeout = timeout
+
+    # ------------------------------------------------------------------
+    # Low-level JSON-RPC
+    # ------------------------------------------------------------------
+
+    def call(self, method: str, params: dict[str, Any] | None = None) -> dict[str, Any]:
+        """Send a JSON-RPC request and return the parsed response.
+
+        Raises ``requests.ConnectionError`` if the server is unreachable.
+        """
+        payload: dict[str, Any] = {
+            "jsonrpc": "2.0",
+            "id": str(uuid.uuid4()),
+            "method": method,
+        }
+        if params:
+            payload["params"] = params
+
+        resp = requests.post(self.url, json=payload, timeout=self.timeout)
+        try:
+            resp.raise_for_status()
+        except requests.HTTPError as e:
+            raise McpError(f"HTTP {resp.status_code}: {e}") from e
+        return resp.json()  # type: ignore[no-any-return]
+
+    # ------------------------------------------------------------------
+    # MCP standard methods
+    # ------------------------------------------------------------------
+
+    def initialize(self) -> dict[str, Any]:
+        """Send ``initialize`` and return server info."""
+        return self.call("initialize")
+
+    def list_tools(self) -> list[dict[str, Any]]:
+        """Return the list of available tools."""
+        result = self._unwrap(self.call("tools/list"))
+        return result.get("tools", [])  # type: ignore[no-any-return]
+
+    def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> dict[str, Any]:
+        """Call a tool by name and return the result dict."""
+        return self._unwrap(self.call("tools/call", {"name": name, "arguments": arguments or {}}))
+
+    def call_tool_text(self, name: str, arguments: dict[str, Any] | None = None) -> str:
+        """Call a tool and return just the first text content item."""
+        result = self.call_tool(name, arguments)
+        content = result.get("content", [])
+        if not content:
+            return ""
+        return content[0].get("text", str(content[0]))  # type: ignore[no-any-return]
+
+    # ------------------------------------------------------------------
+    # Readiness probes
+    # ------------------------------------------------------------------
+
+    def wait_for_ready(self, timeout: float = 10.0, interval: float = 0.5) -> bool:
+        """Poll until the MCP server responds, or return False on timeout."""
+        deadline = time.monotonic() + timeout
+        while time.monotonic() < deadline:
+            try:
+                resp = requests.post(
+                    self.url,
+                    json={"jsonrpc": "2.0", "id": "probe", "method": "initialize"},
+                    timeout=2,
+                )
+                if resp.status_code == 200:
+                    return True
+            except requests.ConnectionError:
+                pass
+            time.sleep(interval)
+        return False
+
+    def wait_for_down(self, timeout: float = 10.0, interval: float = 0.5) -> bool:
+        """Poll until the MCP server stops responding."""
+        deadline = time.monotonic() + timeout
+        while time.monotonic() < deadline:
+            try:
+                requests.post(
+                    self.url,
+                    json={"jsonrpc": "2.0", "id": "probe", "method": "initialize"},
+                    timeout=1,
+                )
+            except (requests.ConnectionError, requests.ReadTimeout):
+                return True
+            time.sleep(interval)
+        return False
+
+    # ------------------------------------------------------------------
+    # Class methods for discovery
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def from_run_entry(cls, entry: Any | None = None, timeout: int = DEFAULT_TIMEOUT) -> McpAdapter:
+        """Create an adapter from a RunEntry, or discover the latest one.
+
+        Falls back to the default URL if no entry is found.
+        """
+        if entry is None:
+            from dimos.core.run_registry import list_runs
+
+            runs = list_runs(alive_only=True)
+            entry = runs[0] if runs else None
+
+        if entry is not None and hasattr(entry, "mcp_url") and entry.mcp_url:
+            return cls(url=entry.mcp_url, timeout=timeout)
+
+        # Fall back to default URL using GlobalConfig port
+        from dimos.core.global_config import global_config
+
+        url = f"http://localhost:{global_config.mcp_port}/mcp"
+        return cls(url=url, timeout=timeout)
+
+    # ------------------------------------------------------------------
+    # Internals
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _unwrap(response: dict[str, Any]) -> dict[str, Any]:
+        """Extract the ``result`` from a JSON-RPC response, raising on error."""
+        if "error" in response:
+            err = response["error"]
+            msg = err.get("message", str(err)) if isinstance(err, dict) else str(err)
+            raise McpError(msg, code=err.get("code") if isinstance(err, dict) else None)
+        return response.get("result", {})  # type: ignore[no-any-return]
+
+    def __repr__(self) -> str:
+        return f"McpAdapter(url={self.url!r})"

dimos/agents/mcp/mcp_server.py

@@ -15,25 +15,25 @@
 
 import asyncio
 import json
+import os
 from typing import TYPE_CHECKING, Any
 
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
+from starlette.requests import Request  # noqa: TC002
 from starlette.responses import Response
 import uvicorn
 
+from dimos.agents.annotation import skill
+from dimos.core.core import rpc
+from dimos.core.module import Module
+from dimos.core.rpc_client import RpcCall, RPCClient
 from dimos.utils.logging_config import setup_logger
 
 logger = setup_logger()
 
 
-from starlette.requests import Request  # noqa: TC002
-
-from dimos.core.core import rpc
-from dimos.core.module import Module
-from dimos.core.rpc_client import RpcCall, RPCClient
-
 if TYPE_CHECKING:
     import concurrent.futures
 
@@ -51,6 +51,11 @@
 app.state.rpc_calls = {}
 
 
+# ---------------------------------------------------------------------------
+# JSON-RPC helpers
+# ---------------------------------------------------------------------------
+
+
 def _jsonrpc_result(req_id: Any, result: Any) -> dict[str, Any]:
     return {"jsonrpc": "2.0", "id": req_id, "result": result}
 
@@ -63,6 +68,11 @@ 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}}
 
 
+# ---------------------------------------------------------------------------
+# JSON-RPC handlers (standard MCP protocol only)
+# ---------------------------------------------------------------------------
+
+
 def _handle_initialize(req_id: Any) -> dict[str, Any]:
     return _jsonrpc_result(
         req_id,
@@ -77,11 +87,11 @@ def _handle_initialize(req_id: Any) -> dict[str, Any]:
 def _handle_tools_list(req_id: Any, skills: list[SkillInfo]) -> dict[str, Any]:
     tools = []
 
-    for skill in skills:
-        schema = json.loads(skill.args_schema)
+    for s in skills:
+        schema = json.loads(s.args_schema)
         description = schema.pop("description", None)
         schema.pop("title", None)
-        tool = {"name": skill.func_name, "inputSchema": schema}
+        tool: dict[str, Any] = {"name": s.func_name, "inputSchema": schema}
         if description:
             tool["description"] = description
         tools.append(tool)
@@ -128,7 +138,7 @@ async def handle_request(
     params = request.get("params", {}) or {}
     req_id = request.get("id")
 
-    # JSON-RPC notifications have no "id" – the server must not reply.
+    # JSON-RPC notifications have no "id" -- the server must not reply.
     if "id" not in request:
         return None
 
@@ -158,6 +168,11 @@ async def mcp_endpoint(request: Request) -> Response:
     return JSONResponse(result)
 
 
+# ---------------------------------------------------------------------------
+# McpServer Module
+# ---------------------------------------------------------------------------
+
+
 class McpServer(Module):
     def __init__(self) -> None:
         super().__init__()
@@ -183,12 +198,62 @@ def stop(self) -> None:
     @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.skills = [
+            skill_info for module in modules for skill_info 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
+            skill_info.func_name: RpcCall(
+                None, self.rpc, skill_info.func_name, skill_info.class_name, []
+            )
+            for skill_info in app.state.skills
         }
 
+    # ------------------------------------------------------------------
+    # Introspection skills (exposed as MCP tools via tools/list)
+    # ------------------------------------------------------------------
+
+    @skill
+    def server_status(self) -> str:
+        """Get MCP server status: main process PID, deployed modules, and skill count."""
+        from dimos.core.run_registry import get_most_recent
+
+        skills: list[SkillInfo] = app.state.skills
+        modules = list(dict.fromkeys(s.class_name for s in skills))
+        entry = get_most_recent()
+        pid = entry.pid if entry else os.getpid()
+        return json.dumps(
+            {
+                "pid": pid,
+                "modules": modules,
+                "skills": [s.func_name for s in skills],
+            }
+        )
+
+    @skill
+    def list_modules(self) -> str:
+        """List deployed modules and their skills."""
+        skills: list[SkillInfo] = app.state.skills
+        modules: dict[str, list[str]] = {}
+        for s in skills:
+            modules.setdefault(s.class_name, []).append(s.func_name)
+        return json.dumps({"modules": modules})
+
+    @skill
+    def agent_send(self, message: str) -> str:
+        """Send a message to the running DimOS agent via LCM."""
+        if not message:
+            raise ValueError("Message cannot be empty")
+
+        from dimos.core.transport import pLCMTransport
+
+        transport: pLCMTransport[str] = pLCMTransport("/human_input")
+        try:
+            transport.start()
+            transport.publish(message)
+            return f"Message sent to agent: {message[:100]}"
+        finally:
+            transport.stop()
+
     def _start_server(self, port: int | None = None) -> None:
         from dimos.core.global_config import global_config