diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index 28d2535dbd..0574bee1ee 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -18,7 +18,7 @@ LiteLLM Proxy provides an MCP Gateway that allows you to use a fixed endpoint fo | Feature | Description | |---------|-------------| | MCP Operations | • List Tools
• Call Tools | -| Supported MCP Transports | • Streamable HTTP
• SSE | +| Supported MCP Transports | • Streamable HTTP
• SSE
• Standard Input/Output (stdio) | | LiteLLM Permission Management | ✨ Enterprise Only
• By Key
• By Team
• By Organization | ## Adding your MCP @@ -33,12 +33,22 @@ On this form, you should enter your MCP Server URL and the transport you want to LiteLLM supports the following MCP transports: - Streamable HTTP - SSE (Server-Sent Events) +- Standard Input/Output (stdio) +### Adding a stdio MCP Server + +For stdio MCP servers, select "Standard Input/Output (stdio)" as the transport type and provide the stdio configuration in JSON format: + + + @@ -60,6 +70,15 @@ mcp_servers: zapier_mcp: url: "https://actions.zapier.com/mcp/sk-akxxxxx/sse" + # Standard Input/Output (stdio) Server - CircleCI Example + circleci_mcp: + transport: "stdio" + command: "npx" + args: ["-y", "@circleci/mcp-server-circleci"] + env: + CIRCLECI_TOKEN: "your-circleci-token" + CIRCLECI_BASE_URL: "https://circleci.com" + # Full configuration with all optional fields my_http_server: url: "https://my-mcp-server.com/mcp" @@ -70,11 +89,15 @@ mcp_servers: ``` **Configuration Options:** -- **Server Name**: Use any descriptive name for your MCP server (e.g., `zapier_mcp`, `deepwiki_mcp`) -- **URL**: The endpoint URL for your MCP server (required) +- **Server Name**: Use any descriptive name for your MCP server (e.g., `zapier_mcp`, `deepwiki_mcp`, `circleci_mcp`) +- **URL**: The endpoint URL for your MCP server (required for HTTP/SSE transports) - **Transport**: Optional transport type (defaults to `sse`) - `sse` - SSE (Server-Sent Events) transport - `http` - Streamable HTTP transport + - `stdio` - Standard Input/Output transport +- **Command**: The command to execute for stdio transport (required for stdio) +- **Args**: Array of arguments to pass to the command (optional for stdio) +- **Env**: Environment variables to set for the stdio process (optional for stdio) - **Description**: Optional description for the server - **Auth Type**: Optional authentication type - **Spec Version**: Optional MCP specification version (defaults to `2025-03-26`) diff --git a/docs/my-website/img/add_stdio_mcp.png b/docs/my-website/img/add_stdio_mcp.png new file mode 100644 index 0000000000..d82ec72102 Binary files /dev/null and b/docs/my-website/img/add_stdio_mcp.png differ diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index e742cfcd18..843a9fb0e2 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -1,19 +1,26 @@ """ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. """ +import asyncio import base64 from datetime import timedelta from typing import List, Optional -import asyncio -from mcp import ClientSession +from mcp import ClientSession, StdioServerParameters from mcp.client.sse import sse_client +from mcp.client.stdio import stdio_client from mcp.client.streamable_http import streamablehttp_client from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import CallToolResult as MCPCallToolResult from mcp.types import Tool as MCPTool -from litellm.types.mcp import MCPAuth, MCPAuthType, MCPTransport, MCPTransportType +from litellm.types.mcp import ( + MCPAuth, + MCPAuthType, + MCPStdioConfig, + MCPTransport, + MCPTransportType, +) def to_basic_auth(auth_value: str) -> str: @@ -31,11 +38,12 @@ class MCPClient: def __init__( self, - server_url: str, + server_url: str = "", transport_type: MCPTransportType = MCPTransport.http, auth_type: MCPAuthType = None, auth_value: Optional[str] = None, timeout: float = 60.0, + stdio_config: Optional[MCPStdioConfig] = None, ): self.server_url: str = server_url self.transport_type: MCPTransport = transport_type @@ -48,6 +56,7 @@ class MCPClient: self._transport = None self._session_ctx = None self._task: Optional[asyncio.Task] = None + self.stdio_config: Optional[MCPStdioConfig] = stdio_config # handle the basic auth value if provided if auth_value: @@ -70,10 +79,25 @@ class MCPClient: if self._session: return # Already connected - headers = self._get_auth_headers() - try: - if self.transport_type == MCPTransport.sse: + if self.transport_type == MCPTransport.stdio: + # For stdio transport, use stdio_client with command-line parameters + if not self.stdio_config: + raise ValueError("stdio_config is required for stdio transport") + + server_params = StdioServerParameters( + command=self.stdio_config.get("command", ""), + args=self.stdio_config.get("args", []), + env=self.stdio_config.get("env", {}) + ) + + self._transport_ctx = stdio_client(server_params) + self._transport = await self._transport_ctx.__aenter__() + self._session_ctx = ClientSession(self._transport[0], self._transport[1]) + self._session = await self._session_ctx.__aenter__() + await self._session.initialize() + elif self.transport_type == MCPTransport.sse: + headers = self._get_auth_headers() self._transport_ctx = sse_client( url=self.server_url, timeout=self.timeout, @@ -83,7 +107,8 @@ class MCPClient: self._session_ctx = ClientSession(self._transport[0], self._transport[1]) self._session = await self._session_ctx.__aenter__() await self._session.initialize() - else: + else: # http + headers = self._get_auth_headers() self._transport_ctx = streamablehttp_client( url=self.server_url, timeout=timedelta(seconds=self.timeout), diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index e9a033c3ee..b7c2d9756c 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -1,5 +1,5 @@ import uuid -from typing import Iterable, List, Optional, Set +from typing import Any, Dict, Iterable, List, Optional, Set, Union from litellm.proxy._types import ( LiteLLM_MCPServerTable, @@ -13,6 +13,38 @@ from litellm.proxy._types import ( from litellm.proxy.utils import PrismaClient +def _prepare_mcp_server_data( + data: Union[NewMCPServerRequest, UpdateMCPServerRequest] +) -> Dict[str, Any]: + """ + Helper function to prepare MCP server data for database operations. + Handles JSON field serialization for mcp_info and env fields. + + Args: + data: NewMCPServerRequest or UpdateMCPServerRequest object + + Returns: + Dict with properly serialized JSON fields + """ + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + + # Convert model to dict + data_dict = data.model_dump() + + # Handle mcp_info serialization + if data.mcp_info is not None: + data_dict["mcp_info"] = safe_dumps(data.mcp_info) + + # Handle env serialization + if data.env is not None: + data_dict["env"] = safe_dumps(data.env) + + # mcp_access_groups is already List[str], no serialization needed + + + return data_dict + + async def get_all_mcp_servers( prisma_client: PrismaClient, ) -> List[LiteLLM_MCPServerTable]: @@ -215,35 +247,21 @@ async def create_mcp_server( """ Create a new mcp server record in the db """ - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps if data.server_id is None: data.server_id = str(uuid.uuid4()) - # Convert model to dict and handle JSON fields - data_dict = data.model_dump() + # Use helper to prepare data with proper JSON serialization + data_dict = _prepare_mcp_server_data(data) - # Handle mcp_info serialization - mcp_info: Optional[str] = None - if data.mcp_info is not None: - mcp_info = safe_dumps(data.mcp_info) - del data_dict["mcp_info"] - - # Handle mcp_access_groups - it's already a List[str], no need to serialize - mcp_access_groups: Optional[list] = None - if data.mcp_access_groups is not None: - mcp_access_groups = data.mcp_access_groups - del data_dict["mcp_access_groups"] + # Add audit fields + data_dict["created_by"] = touched_by + data_dict["updated_by"] = touched_by - mcp_server_record = await prisma_client.db.litellm_mcpservertable.create( - data={ - **data_dict, - "created_by": touched_by, - "updated_by": touched_by, - "mcp_info": mcp_info, - "mcp_access_groups": mcp_access_groups, - } + new_mcp_server = await prisma_client.db.litellm_mcpservertable.create( + data=data_dict # type: ignore ) - return mcp_server_record + + return new_mcp_server async def update_mcp_server( @@ -252,33 +270,14 @@ async def update_mcp_server( """ Update a new mcp server record in the db """ - from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + # Use helper to prepare data with proper JSON serialization + data_dict = _prepare_mcp_server_data(data) - # Convert model to dict and handle JSON fields - data_dict = data.model_dump() - - # Handle mcp_info serialization - mcp_info: Optional[str] = None - if data.mcp_info is not None: - mcp_info = safe_dumps(data.mcp_info) - del data_dict["mcp_info"] - - # Handle mcp_access_groups - it's already a List[str], no need to serialize - mcp_access_groups: Optional[list] = None - if data.mcp_access_groups is not None: - mcp_access_groups = data.mcp_access_groups - del data_dict["mcp_access_groups"] + # Add audit fields + data_dict["updated_by"] = touched_by - mcp_server_record = await prisma_client.db.litellm_mcpservertable.update( - where={ - "server_id": data.server_id, - }, - data={ - **data_dict, - "created_by": touched_by, - "updated_by": touched_by, - "mcp_info": mcp_info, - "mcp_access_groups": mcp_access_groups, - }, + updated_mcp_server = await prisma_client.db.litellm_mcpservertable.update( + where={"server_id": data.server_id}, data=data_dict # type: ignore ) - return mcp_server_record + + return updated_mcp_server diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 1f9b09c0e8..3b4d6abf57 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -36,9 +36,35 @@ from litellm.proxy._types import ( MCPTransportType, UserAPIKeyAuth, ) +from litellm.types.mcp import MCPStdioConfig from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer +def _deserialize_env_dict(env_data: Any) -> Optional[Dict[str, str]]: + """ + Helper function to deserialize environment dictionary from database storage. + Handles both JSON string and dictionary formats. + + Args: + env_data: The environment data from database (could be JSON string or dict) + + Returns: + Dict[str, str] or None: Deserialized environment dictionary + """ + if not env_data: + return None + + if isinstance(env_data, str): + try: + return json.loads(env_data) + except (json.JSONDecodeError, TypeError): + # If it's not valid JSON, return as-is (shouldn't happen but safety) + return None + else: + # Already a dictionary + return env_data + + class MCPServerManager: def __init__(self): self.registry: Dict[str, MCPServer] = {} @@ -88,7 +114,7 @@ class MCPServerManager: # Generate stable server ID based on parameters server_id = self._generate_stable_server_id( server_name=server_name, - url=server_config["url"], + url=server_config.get("url", None) or "", transport=server_config.get("transport", MCPTransport.http), spec_version=server_config.get("spec_version", MCPSpecVersion.mar_2025), auth_type=server_config.get("auth_type", None), @@ -97,7 +123,10 @@ class MCPServerManager: new_server = MCPServer( server_id=server_id, name=server_name, - url=server_config["url"], + url=server_config.get("url", None) or "", + command=server_config.get("command", None) or "", + args=server_config.get("args", None) or [], + env=server_config.get("env", None) or {}, # TODO: utility fn the default values transport=server_config.get("transport", MCPTransport.http), spec_version=server_config.get("spec_version", MCPSpecVersion.mar_2025), @@ -129,6 +158,10 @@ class MCPServerManager: def add_update_server(self, mcp_server: LiteLLM_MCPServerTable): if mcp_server.server_id not in self.get_registry(): _mcp_info: MCPInfo = mcp_server.mcp_info or {} + + # Use helper to deserialize environment dictionary + env_dict = _deserialize_env_dict(mcp_server.env) + new_server = MCPServer( server_id=mcp_server.server_id, name=mcp_server.alias or mcp_server.server_id, @@ -141,6 +174,10 @@ class MCPServerManager: description=mcp_server.description, mcp_server_cost_info=_mcp_info.get("mcp_server_cost_info", None), ), + # Stdio-specific fields + command=mcp_server.command, + args=mcp_server.args, + env=env_dict, ) self.registry[mcp_server.server_id] = new_server verbose_logger.debug( @@ -227,13 +264,36 @@ class MCPServerManager: MCPClient: Configured MCP client instance """ transport = server.transport or MCPTransport.sse - return MCPClient( - server_url=server.url, - transport_type=transport, - auth_type=server.auth_type, - auth_value=mcp_auth_header or server.authentication_token, - timeout=60.0, - ) + + # Handle stdio transport + if transport == MCPTransport.stdio: + # For stdio, we need to get the stdio config from the server + stdio_config: Optional[MCPStdioConfig] = None + if server.command and server.args is not None: + stdio_config = MCPStdioConfig( + command=server.command, + args=server.args, + env=server.env or {} + ) + + return MCPClient( + server_url="", # Not used for stdio + transport_type=transport, + auth_type=server.auth_type, + auth_value=mcp_auth_header or server.authentication_token, + timeout=60.0, + stdio_config=stdio_config, + ) + else: + # For HTTP/SSE transports + server_url = server.url or "" + return MCPClient( + server_url=server_url, + transport_type=transport, + auth_type=server.auth_type, + auth_value=mcp_auth_header or server.authentication_token, + timeout=60.0, + ) async def _get_tools_from_server(self, server: MCPServer, mcp_auth_header: Optional[str] = None) -> List[MCPTool]: """ diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 811a7b35f0..68de8d8b63 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -847,9 +847,28 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): transport: MCPTransportType = MCPTransport.sse spec_version: MCPSpecVersionType = MCPSpecVersion.mar_2025 auth_type: Optional[MCPAuthType] = None - url: str + url: Optional[str] = None mcp_info: Optional[MCPInfo] = None mcp_access_groups: List[str] = Field(default_factory=list) + # Stdio-specific fields + command: Optional[str] = None + args: Optional[List[str]] = None + env: Optional[Dict[str, str]] = None + + @model_validator(mode="before") + @classmethod + def validate_transport_fields(cls, values): + if isinstance(values, dict): + transport = values.get("transport") + if transport == MCPTransport.stdio: + if not values.get("command"): + raise ValueError("command is required for stdio transport") + if not values.get("args"): + raise ValueError("args is required for stdio transport") + elif transport in [MCPTransport.http, MCPTransport.sse]: + if not values.get("url"): + raise ValueError("url is required for HTTP/SSE transport") + return values @@ -860,9 +879,28 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): transport: MCPTransportType = MCPTransport.sse spec_version: MCPSpecVersionType = MCPSpecVersion.mar_2025 auth_type: Optional[MCPAuthType] = None - url: str + url: Optional[str] = None mcp_info: Optional[MCPInfo] = None mcp_access_groups: List[str] = Field(default_factory=list) + # Stdio-specific fields + command: Optional[str] = None + args: Optional[List[str]] = None + env: Optional[Dict[str, str]] = None + + @model_validator(mode="before") + @classmethod + def validate_transport_fields(cls, values): + if isinstance(values, dict): + transport = values.get("transport") + if transport == MCPTransport.stdio: + if not values.get("command"): + raise ValueError("command is required for stdio transport") + if not values.get("args"): + raise ValueError("args is required for stdio transport") + elif transport in [MCPTransport.http, MCPTransport.sse]: + if not values.get("url"): + raise ValueError("url is required for HTTP/SSE transport") + return values @@ -872,7 +910,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): server_id: str alias: Optional[str] = None description: Optional[str] = None - url: str + url: Optional[str] = None transport: MCPTransportType spec_version: MCPSpecVersionType auth_type: Optional[MCPAuthType] = None @@ -883,6 +921,10 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): teams: List[Dict[str, Optional[str]]] = Field(default_factory=list) mcp_access_groups: List[str] = Field(default_factory=list) mcp_info: Optional[MCPInfo] = None + # Stdio-specific fields + command: Optional[str] = None + args: Optional[List[str]] = None + env: Optional[Dict[str, str]] = None class NewUserRequestTeam(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 351fe75291..b182ce5de1 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -23,8 +23,8 @@ from fastapi.responses import JSONResponse import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.constants import LITELLM_PROXY_ADMIN_NAME -from litellm.proxy.auth.model_checks import get_mcp_server_ids from litellm.proxy._experimental.mcp_server.utils import validate_mcp_server_name +from litellm.proxy.auth.model_checks import get_mcp_server_ids router = APIRouter(prefix="/v1/mcp", tags=["mcp"]) MCP_AVAILABLE: bool = True @@ -244,6 +244,10 @@ if MCP_AVAILABLE: created_at=datetime.now(), updated_at=datetime.now(), mcp_info=_server_config.mcp_info, + # Stdio-specific fields + command=_server_config.command, + args=_server_config.args, + env=_server_config.env, ) ) @@ -263,7 +267,11 @@ if MCP_AVAILABLE: updated_by=server.updated_by, mcp_access_groups=server.mcp_access_groups if server.mcp_access_groups is not None else [], mcp_info=server.mcp_info, - teams=cast(List[Dict[str, str | None]], server_to_teams_map.get(server.server_id, [])) + teams=cast(List[Dict[str, str | None]], server_to_teams_map.get(server.server_id, [])), + # Stdio-specific fields + command=server.command, + args=server.args, + env=server.env, ) for server in LIST_MCP_SERVERS ] diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 3b2cf0af16..b92f4af384 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -168,7 +168,7 @@ model LiteLLM_MCPServerTable { server_id String @id @default(uuid()) alias String? description String? - url String + url String? transport String @default("sse") spec_version String @default("2025-03-26") auth_type String? @@ -178,7 +178,10 @@ model LiteLLM_MCPServerTable { updated_by String? mcp_info Json? @default("{}") mcp_access_groups String[] - + // Stdio-specific fields + command String? + args String[] @default([]) + env Json? @default("{}") } // Generate Tokens for Proxy diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 081bce10b6..dc54694410 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -18,6 +18,7 @@ else: class MCPTransport(str, enum.Enum): sse = "sse" http = "http" + stdio = "stdio" class MCPSpecVersion(str, enum.Enum): @@ -32,7 +33,7 @@ class MCPAuth(str, enum.Enum): # MCP Literals -MCPTransportType = Literal[MCPTransport.sse, MCPTransport.http] +MCPTransportType = Literal[MCPTransport.sse, MCPTransport.http, MCPTransport.stdio] MCPSpecVersionType = Literal[MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025] MCPAuthType = Optional[ Literal[MCPAuth.none, MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic] @@ -52,6 +53,23 @@ class MCPServerCostInfo(TypedDict, total=False): """ +class MCPStdioConfig(TypedDict, total=False): + command: str + """ + Command to run the MCP server (e.g., 'npx', 'python', 'node') + """ + + args: List[str] + """ + Arguments to pass to the command + """ + + env: Optional[Dict[str, str]] + """ + Environment variables to set when running the command + """ + + class MCPPostCallResponseObject(BaseModel): """ Pydantic object used for MCP post_call_hook response diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 1eeab925e2..8ed95b8394 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Dict, List, Optional from pydantic import BaseModel, ConfigDict from typing_extensions import TypedDict @@ -17,10 +17,14 @@ class MCPInfo(TypedDict, total=False): class MCPServer(BaseModel): server_id: str name: str - url: str + url: Optional[str] = None transport: MCPTransportType spec_version: MCPSpecVersionType auth_type: Optional[MCPAuthType] = None authentication_token: Optional[str] = None mcp_info: Optional[MCPInfo] = None + # Stdio-specific fields + command: Optional[str] = None + args: Optional[List[str]] = None + env: Optional[Dict[str, str]] = None model_config = ConfigDict(arbitrary_types_allowed=True) diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py new file mode 100644 index 0000000000..cf1af288d5 --- /dev/null +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -0,0 +1,79 @@ +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +# Add the parent directory to the path so we can import litellm +sys.path.insert(0, '../../../') + +from litellm.experimental_mcp_client.client import MCPClient +from litellm.types.mcp import MCPStdioConfig, MCPTransport + + +class TestMCPClient: + """Test MCP Client stdio functionality""" + + def test_mcp_client_stdio_init(self): + """Test MCPClient initialization with stdio config""" + stdio_config = MCPStdioConfig( + command="python", + args=["-m", "my_mcp_server"], + env={"DEBUG": "1"} + ) + + client = MCPClient( + transport_type=MCPTransport.stdio, + stdio_config=stdio_config + ) + + assert client.transport_type == MCPTransport.stdio + assert client.stdio_config == stdio_config + assert client.stdio_config["command"] == "python" + assert client.stdio_config["args"] == ["-m", "my_mcp_server"] + + @pytest.mark.asyncio + async def test_mcp_client_stdio_connect_error(self): + """Test MCP client stdio connection error handling""" + # Test missing stdio_config + client = MCPClient(transport_type=MCPTransport.stdio) + + with pytest.raises(ValueError, match="stdio_config is required for stdio transport"): + await client.connect() + + @pytest.mark.asyncio + @patch('litellm.experimental_mcp_client.client.stdio_client') + @patch('litellm.experimental_mcp_client.client.ClientSession') + async def test_mcp_client_stdio_connect_success(self, mock_session, mock_stdio_client): + """Test successful stdio connection""" + # Setup mocks + mock_transport = (MagicMock(), MagicMock()) + mock_stdio_client.return_value.__aenter__ = AsyncMock(return_value=mock_transport) + + mock_session_instance = MagicMock() + mock_session_instance.__aenter__ = AsyncMock(return_value=mock_session_instance) + mock_session_instance.initialize = AsyncMock() + mock_session.return_value = mock_session_instance + + stdio_config = MCPStdioConfig( + command="python", + args=["-m", "my_mcp_server"], + env={"DEBUG": "1"} + ) + + client = MCPClient( + transport_type=MCPTransport.stdio, + stdio_config=stdio_config + ) + + await client.connect() + + # Verify stdio_client was called with correct parameters + mock_stdio_client.assert_called_once() + call_args = mock_stdio_client.call_args[0][0] + assert call_args.command == "python" + assert call_args.args == ["-m", "my_mcp_server"] + assert call_args.env == {"DEBUG": "1"} + + +if __name__ == "__main__": + pytest.main([__file__]) \ No newline at end of file diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py new file mode 100644 index 0000000000..486ad8fe33 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -0,0 +1,94 @@ +import sys +from datetime import datetime +from unittest.mock import MagicMock + +import pytest + +# Add the parent directory to the path so we can import litellm +sys.path.insert(0, '../../../../../') + +from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + _deserialize_env_dict, +) +from litellm.proxy._types import LiteLLM_MCPServerTable, MCPSpecVersion, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + + +class TestMCPServerManager: + """Test MCP Server Manager stdio functionality""" + + def test_deserialize_env_dict(self): + """Test environment dictionary deserialization""" + # Test JSON string + env_json = '{"PATH": "/usr/bin", "DEBUG": "1"}' + result = _deserialize_env_dict(env_json) + assert result == {"PATH": "/usr/bin", "DEBUG": "1"} + + # Test already dict + env_dict = {"PATH": "/usr/bin", "DEBUG": "1"} + result = _deserialize_env_dict(env_dict) + assert result == {"PATH": "/usr/bin", "DEBUG": "1"} + + # Test invalid JSON + invalid_json = '{"PATH": "/usr/bin", "DEBUG": 1' + result = _deserialize_env_dict(invalid_json) + assert result is None + + def test_add_update_server_stdio(self): + """Test adding stdio MCP server""" + manager = MCPServerManager() + + stdio_server = LiteLLM_MCPServerTable( + server_id="stdio-server-1", + alias="test_stdio_server", + description="Test stdio server", + url=None, + transport=MCPTransport.stdio, + spec_version=MCPSpecVersion.mar_2025, + command="python", + args=["-m", "server"], + env={"DEBUG": "1", "TEST": "1"}, + created_at=datetime.now(), + updated_at=datetime.now() + ) + + manager.add_update_server(stdio_server) + + # Verify server was added + assert "stdio-server-1" in manager.registry + added_server = manager.registry["stdio-server-1"] + + assert added_server.server_id == "stdio-server-1" + assert added_server.name == "test_stdio_server" + assert added_server.transport == MCPTransport.stdio + assert added_server.command == "python" + assert added_server.args == ["-m", "server"] + assert added_server.env == {"DEBUG": "1", "TEST": "1"} + + def test_create_mcp_client_stdio(self): + """Test creating MCP client for stdio transport""" + manager = MCPServerManager() + + stdio_server = MCPServer( + server_id="stdio-server-2", + name="test_stdio_server", + url=None, + transport=MCPTransport.stdio, + spec_version=MCPSpecVersion.mar_2025, + command="node", + args=["server.js"], + env={"NODE_ENV": "test"} + ) + + client = manager._create_mcp_client(stdio_server) + + assert client.transport_type == MCPTransport.stdio + assert client.stdio_config is not None + assert client.stdio_config["command"] == "node" + assert client.stdio_config["args"] == ["server.js"] + assert client.stdio_config["env"] == {"NODE_ENV": "test"} + + +if __name__ == "__main__": + pytest.main([__file__]) \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/mcp_tools/StdioConfiguration.tsx b/ui/litellm-dashboard/src/components/mcp_tools/StdioConfiguration.tsx new file mode 100644 index 0000000000..2c3c4e515f --- /dev/null +++ b/ui/litellm-dashboard/src/components/mcp_tools/StdioConfiguration.tsx @@ -0,0 +1,58 @@ +import React from "react"; +import { Form, Input, Tooltip } from "antd"; +import { InfoCircleOutlined } from "@ant-design/icons"; + +interface StdioConfigurationProps { + isVisible: boolean; +} + +const StdioConfiguration: React.FC = ({ isVisible }) => { + if (!isVisible) return null; + + return ( + + Stdio Configuration (JSON) + + + + + } + name="stdio_config" + rules={[ + { required: true, message: "Please enter stdio configuration" }, + { + validator: (_, value) => { + if (!value) return Promise.resolve(); + try { + JSON.parse(value); + return Promise.resolve(); + } catch { + return Promise.reject("Please enter valid JSON"); + } + }, + }, + ]} + > + + + ); +}; + +export default StdioConfiguration; \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx b/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx index 4718dcd5d2..6f2845a63d 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/ToolTestPanel.tsx @@ -4,24 +4,6 @@ import { MCPTool, InputSchema } from "./types"; import { Form, Tooltip, message } from "antd"; import { InfoCircleOutlined, ClockCircleOutlined } from "@ant-design/icons"; -const AuthBanner = ({ needsAuth, authValue }: { needsAuth: boolean; authValue?: string | null }) => { - if (!needsAuth || (needsAuth && authValue)) { - return ( - - This tool does not require authentication or has authentication added. - - ); - } - - if (needsAuth && !authValue) { - return ( - - Please provide authentication details if this tool call requires auth. - - ); - } - return null; -}; export function ToolTestPanel({ tool, @@ -65,11 +47,59 @@ export function ToolTestPanel({ return tool.inputSchema as InputSchema; }, [tool.inputSchema]); + // Check if this is a nested params structure and extract the actual parameters + const actualSchema: InputSchema = React.useMemo(() => { + if (schema.properties && schema.properties.params && + schema.properties.params.type === "object" && + schema.properties.params.properties) { + // This is a nested params structure, extract the actual parameters + return { + type: "object", + properties: schema.properties.params.properties, + required: schema.properties.params.required || [], + }; + } + return schema; + }, [schema]); + const handleSubmit = (values: Record) => { const start = Date.now(); setStartTime(start); setDuration(null); - onSubmit(values); + + // Convert form values to proper types based on schema + const convertedValues: Record = {}; + const schemaToUse = actualSchema; + + Object.entries(values).forEach(([key, value]) => { + const prop = schemaToUse.properties?.[key]; + if (prop && value !== null && value !== undefined && value !== "") { + switch (prop.type) { + case "boolean": + convertedValues[key] = value === "true" || value === true; + break; + case "number": + convertedValues[key] = Number(value); + break; + case "string": + convertedValues[key] = String(value); + break; + default: + convertedValues[key] = value; + } + } else if (value !== null && value !== undefined && value !== "") { + convertedValues[key] = value; + } + }); + + // If this was a nested params structure, wrap the values back in params + const submitValues = (schema.properties && schema.properties.params && + schema.properties.params.type === "object" && + schema.properties.params.properties) + ? { params: convertedValues } + : convertedValues; + + onSubmit(submitValues); }; // Track when result changes to calculate duration @@ -183,9 +213,7 @@ export function ToolTestPanel({ - {/* Auth Banner */} - - + {/* Two Column Layout - Always Side by Side */}
{/* Left Column - Input Parameters */} @@ -219,7 +247,7 @@ export function ToolTestPanel({ />
- ) : schema.properties === undefined ? ( + ) : actualSchema.properties === undefined ? (

No Parameters Required

@@ -228,13 +256,13 @@ export function ToolTestPanel({
) : (
- {Object.entries(schema.properties).map(([key, prop]) => ( + {Object.entries(actualSchema.properties).map(([key, prop]) => ( {key}{" "} - {schema.required?.includes(key) && *} + {actualSchema.required?.includes(key) && *} {prop.description && ( @@ -245,13 +273,29 @@ export function ToolTestPanel({ name={key} rules={[ { - required: schema.required?.includes(key), + required: actualSchema.required?.includes(key), message: `Please enter ${key}`, }, ]} className="mb-3" > - {prop.type === "string" && ( + {prop.type === "string" && prop.enum && ( + + )} + + {prop.type === "string" && !prop.enum && ( - - Enable this option -
+ )} ))} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 22654fa67d..65d493a495 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -6,15 +6,15 @@ import { Select, message, Button as AntdButton, - Space, Input, } from "antd"; -import { InfoCircleOutlined, MinusCircleOutlined, PlusOutlined } from "@ant-design/icons"; +import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TextInput } from "@tremor/react"; import { createMCPServer } from "../networking"; import { MCPServer, MCPServerCostInfo } from "./types"; import MCPServerCostConfig from "./mcp_server_cost_config"; import MCPConnectionStatus from "./mcp_connection_status"; +import StdioConfiguration from "./StdioConfiguration"; import { isAdminRole } from "@/utils/roles"; @@ -38,6 +38,7 @@ const CreateMCPServer: React.FC = ({ const [mcpAccessGroups, setMcpAccessGroups] = useState([]); const [formValues, setFormValues] = useState>({}); const [tools, setTools] = useState([]); + const [transportType, setTransportType] = useState('sse'); const handleCreate = async (formValues: Record) => { setIsLoading(true); @@ -46,9 +47,51 @@ const CreateMCPServer: React.FC = ({ const accessGroups = formValues.mcp_access_groups + // Process stdio configuration if present + let stdioFields = {}; + if (formValues.stdio_config && transportType === 'stdio') { + try { + const stdioConfig = JSON.parse(formValues.stdio_config); + + // Handle both formats: + // 1. Full mcpServers structure: {"mcpServers": {"server-name": {...}}} + // 2. Direct config: {"command": "...", "args": [...], "env": {...}} + + let actualConfig = stdioConfig; + + // If it's the full mcpServers structure, extract the first server config + if (stdioConfig.mcpServers && typeof stdioConfig.mcpServers === 'object') { + const serverNames = Object.keys(stdioConfig.mcpServers); + if (serverNames.length > 0) { + const firstServerName = serverNames[0]; + actualConfig = stdioConfig.mcpServers[firstServerName]; + + // If no alias is provided, use the server name from the JSON + if (!formValues.alias) { + formValues.alias = firstServerName.replace(/-/g, '_'); // Replace hyphens with underscores + } + } + } + + stdioFields = { + command: actualConfig.command, + args: actualConfig.args, + env: actualConfig.env + }; + + console.log('Parsed stdio config:', stdioFields); + } catch (error) { + message.error("Invalid JSON in stdio configuration"); + return; + } + } + // Prepare the payload with cost configuration const payload = { ...formValues, + ...stdioFields, + // Remove the raw stdio_config field as we've extracted its components + stdio_config: undefined, mcp_info: { server_name: formValues.alias || formValues.url, description: formValues.description, @@ -89,7 +132,15 @@ const CreateMCPServer: React.FC = ({ setModalVisible(false); }; - + const handleTransportChange = (value: string) => { + setTransportType(value); + // Clear fields that are not relevant for the selected transport + if (value === 'stdio') { + form.setFieldsValue({ url: undefined, auth_type: undefined }); + } else { + form.setFieldsValue({ command: undefined, args: undefined, env: undefined }); + } + }; // rendering if (!isAdminRole(userRole)) { @@ -186,41 +237,48 @@ const CreateMCPServer: React.FC = ({ - MCP Server URL + Transport Type } - name="url" - rules={[ - { required: true, message: "Please enter a server URL" }, - { type: 'url', message: "Please enter a valid URL" } - ]} + name="transport" + rules={[{ required: true, message: "Please select a transport type" }]} > - + -
+ {/* URL field - only show for HTTP and SSE */} + {transportType !== 'stdio' && ( - Transport Type + MCP Server URL } - name="transport" - rules={[{ required: true, message: "Please select a transport type" }]} + name="url" + rules={[ + { required: true, message: "Please enter a server URL" }, + { type: 'url', message: "Please enter a valid URL" } + ]} > - + + )} + {/* Authentication - only show for HTTP and SSE */} + {transportType !== 'stdio' && ( @@ -241,7 +299,10 @@ const CreateMCPServer: React.FC = ({ Basic Auth -
+ )} + + {/* Stdio Configuration - only show for stdio transport */} + { export interface InputSchemaProperty { type: string; description?: string; + properties?: Record; // For nested object properties + required?: string[]; // For required fields in nested objects + enum?: string[]; // For enum values + default?: any; // For default values } // Define the structure for the input schema of a tool