diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 402e12d9356..39268565f8f 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -185,12 +185,31 @@ class MCPServerManager: } """ + # Async-safe equivalent of @functools.lru_cache for stdio tool fetching. + # Keyed by server_id → raw (unprefixed) List[MCPTool]. + # Populated once at startup; use _cache_stdio_tools / _invalidate_stdio_cache + # to manage entries. Prevents re-spawning the subprocess (and re-downloading + # packages via npx / uvx) on every list_tools call. + self._stdio_tools_cache: Dict[str, List[MCPTool]] = {} + def get_registry(self) -> Dict[str, MCPServer]: """ Get the registered MCP Servers from the registry and union with the config MCP Servers """ return self.config_mcp_servers | self.registry + # ------------------------------------------------------------------ + # stdio tools cache helpers (async-safe @functools.lru_cache analog) + # ------------------------------------------------------------------ + + def _cache_stdio_tools(self, server_id: str, tools: List[MCPTool]) -> None: + """Store raw (unprefixed) tools for a stdio server. Call after first fetch.""" + self._stdio_tools_cache[server_id] = tools + + def _invalidate_stdio_cache(self, server_id: str) -> None: + """Evict a stdio server's cached tools (e.g. after update or delete).""" + self._stdio_tools_cache.pop(server_id, None) + async def load_servers_from_config( self, mcp_servers_config: Dict[str, Any], @@ -532,6 +551,7 @@ class MCPServerManager: """ Remove a server from the registry """ + self._invalidate_stdio_cache(mcp_server.server_id) if mcp_server.server_name in self.get_registry(): del self.registry[mcp_server.server_name] verbose_logger.debug(f"Removed MCP Server: {mcp_server.server_name}") @@ -726,6 +746,9 @@ class MCPServerManager: if mcp_server.server_id in self.registry: new_server = await self.build_mcp_server_from_table(mcp_server) self.registry[mcp_server.server_id] = new_server + # Invalidate the stdio cache so the updated command/args are picked + # up on the next fetch (re-downloads package if version changed). + self._invalidate_stdio_cache(mcp_server.server_id) await self._maybe_register_openapi_tools(new_server) verbose_logger.debug(f"Updated MCP Server: {new_server.name}") @@ -1057,6 +1080,104 @@ class MCPServerManager: ######################################################### # Methods that call the upstream MCP servers ######################################################### + + # ------------------------------------------------------------------ + # stdio package-security helpers + # ------------------------------------------------------------------ + # Pinned semver: digits only, no range operators (^, ~, >, <, *). + # Allows pre-release suffixes like 1.2.3-rc.1. + _PINNED_SEMVER_RE = re.compile(r"^\d+\.\d+(\.\d+)?([-.][a-zA-Z0-9.]+)?$") + + def _parse_package_and_version( + self, command: str, args: List[str] + ) -> Optional[Tuple[str, str]]: + """ + Extract (package_name, version) from stdio args for npx or uvx commands. + + Returns None if the command is not npx/uvx, or the args don't contain a + package spec with a pinned version. + + npx examples: + ["-y", "@scope/pkg@1.2.3"] → ("@scope/pkg", "1.2.3") + ["pkg@1.0.0"] → ("pkg", "1.0.0") + ["pkg"] → None (no version) + ["-y", "pkg@latest"] → None (not pinned) + + uvx / pip examples: + ["pkg@1.2.3"] → ("pkg", "1.2.3") + ["pkg==1.2.3"] → ("pkg", "1.2.3") + ["pkg>=1.0"] → None (range, not pinned) + """ + base_command = os.path.basename(command) + if base_command not in ("npx", "uvx"): + return None + + for arg in args: + if arg.startswith("-"): + continue + + # uvx / pip style: pkg==1.2.3 + if "==" in arg: + pkg, _, ver = arg.partition("==") + if pkg and self._PINNED_SEMVER_RE.match(ver): + return pkg.strip(), ver.strip() + return None + + # npx / uvx style: pkg@version (rfind handles @scope/pkg@1.0) + at_idx = arg.rfind("@") + if at_idx > 0: + pkg = arg[:at_idx] + ver = arg[at_idx + 1:] + if pkg and self._PINNED_SEMVER_RE.match(ver): + return pkg, ver + # First non-flag arg found but no valid pinned version + return None + + return None + + def _validate_pinned_version(self, server: MCPServer) -> None: + """ + Raise HTTPException(403) if a stdio server uses npx/uvx without a + pinned semver package version (e.g. @latest or bare package names). + + Pinning prevents silent upgrades that could introduce malicious code. + """ + if server.transport != MCPTransport.stdio or not server.command: + return + base_command = os.path.basename(server.command) + if base_command not in ("npx", "uvx"): + return + if not server.args: + raise HTTPException( + status_code=403, + detail=( + f"MCP stdio server '{server.name}' uses {base_command} but " + "has no args. Specify a package with a pinned version " + "(e.g. @modelcontextprotocol/server-everything@1.2.3)." + ), + ) + if self._parse_package_and_version(server.command, server.args) is None: + pkg_arg = next((a for a in server.args if not a.startswith("-")), "") + raise HTTPException( + status_code=403, + detail=( + f"MCP stdio server '{server.name}': package '{pkg_arg}' must " + "specify an exact pinned version (e.g. pkg@1.2.3 or pkg==1.2.3). " + "Tags like @latest or bare package names are not allowed." + ), + ) + + def _inject_ignore_scripts(self, args: List[str]) -> List[str]: + """ + Return a copy of args with --ignore-scripts prepended (if absent). + + Prevents npm postinstall/preinstall lifecycle scripts — a common + supply-chain attack vector — from executing when npx installs a package. + """ + if "--ignore-scripts" in args: + return args + return ["--ignore-scripts"] + list(args) + def _build_stdio_env( self, server: MCPServer, @@ -1135,11 +1256,23 @@ class MCPServerManager: f"Add it to LITELLM_MCP_STDIO_EXTRA_COMMANDS to allow this command.", ) + # Require an exact pinned version for npx/uvx packages so that + # silent upgrades (e.g. @latest) cannot introduce malicious code. + self._validate_pinned_version(server) + stdio_config: Optional[MCPStdioConfig] = None if server.command and server.args is not None: + base_command = os.path.basename(server.command) + # Prevent npm lifecycle scripts (postinstall etc.) from running — + # a common supply-chain attack vector. + resolved_args = ( + self._inject_ignore_scripts(server.args) + if base_command == "npx" + else list(server.args) + ) stdio_config = MCPStdioConfig( command=server.command, - args=server.args, + args=resolved_args, env=resolved_env, ) @@ -1204,6 +1337,17 @@ class MCPServerManager: verbose_logger.debug(f"Connecting to url: {server.url}") verbose_logger.info(f"_get_tools_from_server for {server.name}...") + # For stdio servers, serve tools from the in-memory cache that was populated at + # startup. This avoids re-spawning the subprocess (and re-downloading packages + # via npx/uvx) on every list_tools / call_tool request. + if server.transport == MCPTransport.stdio and server.server_id in self._stdio_tools_cache: + verbose_logger.debug( + f"Returning cached tools for stdio server '{server.name}' (server_id={server.server_id})" + ) + return self._create_prefixed_tools( + self._stdio_tools_cache[server.server_id], server, add_prefix=add_prefix + ) + client = None try: @@ -1248,6 +1392,14 @@ class MCPServerManager: else: tools = await self._fetch_tools_with_timeout(client, server.name) + # Populate the stdio cache on first successful fetch so that + # subsequent calls skip subprocess spawning entirely. + if server.transport == MCPTransport.stdio and tools: + self._cache_stdio_tools(server.server_id, tools) + verbose_logger.debug( + f"Cached {len(tools)} tools for stdio server '{server.name}' (server_id={server.server_id})" + ) + prefixed_or_original_tools = self._create_prefixed_tools( tools, server, add_prefix=add_prefix ) 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 index 656a9c616e8..186b0abf31f 100644 --- 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 @@ -141,7 +141,7 @@ class TestMCPServerManager: url=None, transport=MCPTransport.stdio, command="npx", - args=["-y", "@modelcontextprotocol/server-everything"], + args=["-y", "@modelcontextprotocol/server-everything@1.0.0"], env={}, ) client = await manager._create_mcp_client(server_no_cache) @@ -154,12 +154,90 @@ class TestMCPServerManager: url=None, transport=MCPTransport.stdio, command="npx", - args=["-y", "@modelcontextprotocol/server-everything"], + args=["-y", "@modelcontextprotocol/server-everything@1.0.0"], env={"NPM_CONFIG_CACHE": "/custom/cache"}, ) client2 = await manager._create_mcp_client(server_with_cache) assert client2.stdio_config["env"]["NPM_CONFIG_CACHE"] == "/custom/cache" + def test_validate_pinned_version_rejects_latest(self): + """npx pkg@latest must be rejected — not a pinned semver.""" + manager = MCPServerManager() + server = MCPServer( + server_id="s1", + name="s1", + transport=MCPTransport.stdio, + command="npx", + args=["-y", "@scope/pkg@latest"], + ) + with pytest.raises(Exception) as exc_info: + manager._validate_pinned_version(server) + assert "pinned version" in str(exc_info.value).lower() or "403" in str(exc_info.value) + + def test_validate_pinned_version_rejects_bare_package(self): + """npx @scope/pkg with no version must be rejected.""" + manager = MCPServerManager() + server = MCPServer( + server_id="s2", + name="s2", + transport=MCPTransport.stdio, + command="npx", + args=["-y", "@scope/pkg"], + ) + with pytest.raises(Exception): + manager._validate_pinned_version(server) + + def test_validate_pinned_version_accepts_exact_semver(self): + """npx pkg@1.2.3 must pass validation without raising.""" + manager = MCPServerManager() + server = MCPServer( + server_id="s3", + name="s3", + transport=MCPTransport.stdio, + command="npx", + args=["-y", "@modelcontextprotocol/server-everything@1.2.3"], + ) + manager._validate_pinned_version(server) # must not raise + + def test_validate_pinned_version_skips_non_npx_uvx(self): + """python / node commands are exempt from the version-pinning check.""" + manager = MCPServerManager() + for cmd in ("python", "python3", "node"): + server = MCPServer( + server_id="s4", + name="s4", + transport=MCPTransport.stdio, + command=cmd, + args=["server.py"], + ) + manager._validate_pinned_version(server) # must not raise + + def test_inject_ignore_scripts_prepends_flag(self): + """--ignore-scripts must be prepended when absent.""" + manager = MCPServerManager() + result = manager._inject_ignore_scripts(["-y", "pkg@1.0.0"]) + assert result == ["--ignore-scripts", "-y", "pkg@1.0.0"] + + def test_inject_ignore_scripts_idempotent(self): + """--ignore-scripts must not be duplicated if already present.""" + manager = MCPServerManager() + args = ["--ignore-scripts", "-y", "pkg@1.0.0"] + assert manager._inject_ignore_scripts(args) == args + + @pytest.mark.asyncio + async def test_create_mcp_client_npx_injects_ignore_scripts(self): + """_create_mcp_client must inject --ignore-scripts for npx commands.""" + manager = MCPServerManager() + server = MCPServer( + server_id="npx-scripts", + name="npx_scripts", + transport=MCPTransport.stdio, + command="npx", + args=["-y", "@modelcontextprotocol/server-everything@1.0.0"], + ) + client = await manager._create_mcp_client(server) + assert "--ignore-scripts" in client.stdio_config["args"] + def test_build_stdio_env_only_accepts_x_prefixed_placeholders(self): """Ensure only ${X-*} placeholders are substituted from headers.""" manager = MCPServerManager() @@ -207,6 +285,90 @@ class TestMCPServerManager: # When the header isn't provided, the key is omitted entirely assert env == {} + @pytest.mark.asyncio + async def test_stdio_tools_cache_prevents_repeated_subprocess_spawn(self): + """ + After the first successful fetch the cached tools must be returned on + subsequent calls — the subprocess must NOT be re-spawned. + + Proof: on the second call we swap the mock to return a *different* tool + list. If the response still matches the first result, the cache was + used and the subprocess was not re-invoked. + """ + manager = MCPServerManager() + server = MCPServer( + server_id="stdio-cache-server", + name="cache_test_server", + transport=MCPTransport.stdio, + command="npx", + args=["--ignore-scripts", "@modelcontextprotocol/server-test@1.0.0"], + env={}, + ) + + first_tools = [ + MCPTool(name="tool_a", description="Tool A", inputSchema={"type": "object"}), + MCPTool(name="tool_b", description="Tool B", inputSchema={"type": "object"}), + ] + second_tools = [ + MCPTool(name="different_tool", description="Different", inputSchema={"type": "object"}), + ] + + fetch_call_count = 0 + + async def fake_fetch(client, server_name): + nonlocal fetch_call_count + fetch_call_count += 1 + return first_tools if fetch_call_count == 1 else second_tools + + manager._create_mcp_client = AsyncMock(return_value=MagicMock()) + manager._fetch_tools_with_timeout = fake_fetch + + # First call — subprocess spawned, tools cached + result1 = await manager._get_tools_from_server(server, add_prefix=False) + assert fetch_call_count == 1 + assert {t.name for t in result1} == {"tool_a", "tool_b"} + + # Second call — if cache is used, we still get tool_a/tool_b, not different_tool + result2 = await manager._get_tools_from_server(server, add_prefix=False) + assert fetch_call_count == 1, "subprocess was re-spawned; cache was not used" + assert {t.name for t in result2} == {"tool_a", "tool_b"}, ( + "got different tools on second call — cache was bypassed" + ) + + @pytest.mark.asyncio + async def test_update_server_invalidates_stdio_tools_cache(self): + """ + Updating a stdio server must evict its cache entry so the next + _get_tools_from_server re-spawns the subprocess with the new config. + """ + manager = MCPServerManager() + + stdio_server = LiteLLM_MCPServerTable( + server_id="stdio-update-cache", + alias="update_cache_server", + description="", + url=None, + transport=MCPTransport.stdio, + command="npx", + args=["--ignore-scripts", "@modelcontextprotocol/server-test@1.0.0"], + env={}, + created_at=datetime.now(), + updated_at=datetime.now(), + ) + + # Seed the cache as if startup already fetched tools for this server + manager._cache_stdio_tools( + "stdio-update-cache", + [MCPTool(name="old_tool", description="", inputSchema={"type": "object"})], + ) + assert "stdio-update-cache" in manager._stdio_tools_cache + + await manager.add_server(stdio_server) + assert "stdio-update-cache" in manager._stdio_tools_cache # add doesn't evict + + await manager.update_server(stdio_server) + assert "stdio-update-cache" not in manager._stdio_tools_cache + @pytest.mark.asyncio async def test_load_servers_from_config_warns_on_invalid_alias(self, caplog): """Invalid aliases from config should emit warnings during load."""