From 5d19cb6ee28d79572bd089d58582ffcd0009d3b7 Mon Sep 17 00:00:00 2001 From: Hashim1999164 Date: Thu, 16 Jul 2026 01:28:59 +0500 Subject: [PATCH 1/6] fix(mcp): make mcp_tool_search default top_k configurable Allow operators to set a global litellm_settings default or per-key object_permission override instead of always defaulting to 5. Fixes #33440 --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/models/object_permission.py | 1 + .../mcp_server/rest_endpoints.py | 11 +++- .../proxy/_experimental/mcp_server/server.py | 12 +++- .../_experimental/mcp_server/tool_search.py | 48 +++++++++++++-- litellm/proxy/_types.py | 1 + litellm/proxy/schema.prisma | 1 + litellm/types/object_permission.py | 1 + schema.prisma | 1 + .../mcp_server/test_mcp_tool_search.py | 60 +++++++++++++++++++ 11 files changed, 129 insertions(+), 10 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_mcp_tool_search_top_k/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_mcp_tool_search_top_k/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_mcp_tool_search_top_k/migration.sql new file mode 100644 index 00000000000..94240b245ea --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716000000_add_mcp_tool_search_top_k/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_ObjectPermissionTable" ADD COLUMN "mcp_tool_search_top_k" INTEGER; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index a23cecc3911..faa66fae557 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -280,6 +280,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + mcp_tool_search_top_k Int? teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] diff --git a/litellm/models/object_permission.py b/litellm/models/object_permission.py index 3052a2af459..ec9b6671e8b 100644 --- a/litellm/models/object_permission.py +++ b/litellm/models/object_permission.py @@ -25,3 +25,4 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): blocked_tools: Optional[List[str]] = [] search_tools: Optional[List[str]] = [] mcp_tool_search_enabled: Optional[bool] = None + mcp_tool_search_top_k: Optional[int] = None diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 111fde86ea0..0511c574801 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -139,9 +139,11 @@ if MCP_AVAILABLE: ) from litellm.proxy._experimental.mcp_server.tool_search import ( MCP_TOOL_SEARCH_TOOL_NAME, - coerce_top_k, + get_mcp_tool_search_default_top_k, + get_virtual_tool_definitions, handle_mcp_tool_call, handle_mcp_tool_search, + resolve_mcp_tool_search_top_k, ) from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.proxy_server import general_settings, proxy_config, proxy_logging_obj @@ -162,7 +164,7 @@ if MCP_AVAILABLE: if tool_name == MCP_TOOL_SEARCH_TOOL_NAME: return await handle_mcp_tool_search( query=tool_arguments.get("query", ""), - top_k=coerce_top_k(tool_arguments.get("top_k", 5)), + top_k=resolve_mcp_tool_search_top_k(tool_arguments.get("top_k"), user_api_key_dict), user_api_key_dict=user_api_key_dict, client_ip=rest_client_ip, mcp_auth_header=virtual_mcp_auth_header, @@ -733,11 +735,14 @@ if MCP_AVAILABLE: ) ): from litellm.proxy._experimental.mcp_server.tool_search import ( + get_mcp_tool_search_default_top_k, get_virtual_tool_definitions, ) return { - "tools": get_virtual_tool_definitions(), + "tools": get_virtual_tool_definitions( + default_top_k=get_mcp_tool_search_default_top_k(user_api_key_dict) + ), "error": None, "message": "Successfully retrieved tools", } diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 68a61b85175..0432af9ff97 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -702,10 +702,16 @@ if MCP_AVAILABLE: from mcp.types import Tool from litellm.proxy._experimental.mcp_server.tool_search import ( + get_mcp_tool_search_default_top_k, get_virtual_tool_definitions, ) - return [Tool(**d) for d in get_virtual_tool_definitions()] + return [ + Tool(**d) + for d in get_virtual_tool_definitions( + default_top_k=get_mcp_tool_search_default_top_k(user_api_key_auth) + ) + ] # Get mcp_servers from context variable verbose_logger.debug("MCP list_tools - Calling _list_mcp_tools") @@ -821,9 +827,9 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.tool_search import ( MCP_TOOL_CALL_TOOL_NAME, MCP_TOOL_SEARCH_TOOL_NAME, - coerce_top_k, handle_mcp_tool_call, handle_mcp_tool_search, + resolve_mcp_tool_search_top_k, ) if name not in (MCP_TOOL_SEARCH_TOOL_NAME, MCP_TOOL_CALL_TOOL_NAME): @@ -848,7 +854,7 @@ if MCP_AVAILABLE: if name == MCP_TOOL_SEARCH_TOOL_NAME: return await handle_mcp_tool_search( query=args.get("query", ""), - top_k=coerce_top_k(args.get("top_k", 5)), + top_k=resolve_mcp_tool_search_top_k(args.get("top_k"), user_api_key_auth), user_api_key_dict=user_api_key_auth, client_ip=client_ip, mcp_servers=mcp_servers, diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index fa57a2b3eb2..dd59f9f6860 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -12,16 +12,54 @@ if TYPE_CHECKING: MCP_TOOL_SEARCH_TOOL_NAME: str = "mcp_tool_search" MCP_TOOL_CALL_TOOL_NAME: str = "mcp_tool_call" +DEFAULT_MCP_TOOL_SEARCH_TOP_K: int = 5 -def coerce_top_k(value: Any, default: int = 5) -> int: +def _get_litellm_settings() -> dict[str, Any]: + try: + from litellm.proxy.proxy_server import proxy_config + + return proxy_config.get_config_state().get("litellm_settings") or {} + except Exception: + return {} + + +def get_mcp_tool_search_default_top_k( + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, +) -> int: + """Resolve the default top_k for mcp_tool_search (per-key, then global, then 5).""" + if user_api_key_dict is not None: + object_permission = getattr(user_api_key_dict, "object_permission", None) + if object_permission is not None: + key_top_k = getattr(object_permission, "mcp_tool_search_top_k", None) + if key_top_k is not None: + return coerce_top_k(key_top_k, default=DEFAULT_MCP_TOOL_SEARCH_TOP_K) + + global_top_k = _get_litellm_settings().get("mcp_tool_search_default_top_k") + if global_top_k is not None: + return coerce_top_k(global_top_k, default=DEFAULT_MCP_TOOL_SEARCH_TOP_K) + + return DEFAULT_MCP_TOOL_SEARCH_TOP_K + + +def resolve_mcp_tool_search_top_k( + explicit_top_k: Any, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, +) -> int: + default_top_k = get_mcp_tool_search_default_top_k(user_api_key_dict) + if explicit_top_k is None: + return default_top_k + return coerce_top_k(explicit_top_k, default=default_top_k) + + +def coerce_top_k(value: Any, default: int = DEFAULT_MCP_TOOL_SEARCH_TOP_K) -> int: try: return int(value) except (TypeError, ValueError): return default -def search_tools(query: str, tools: list[dict[str, Any]], top_k: int = 5) -> list[dict[str, Any]]: +def search_tools(query: str, tools: list[dict[str, Any]], top_k: int = DEFAULT_MCP_TOOL_SEARCH_TOP_K) -> list[dict[str, Any]]: if not query: return [] tokens = query.lower().split() @@ -34,7 +72,9 @@ def search_tools(query: str, tools: list[dict[str, Any]], top_k: int = 5) -> lis return [tool for _, tool in sorted(scored, key=lambda x: x[0], reverse=True)[:top_k]] -def get_virtual_tool_definitions() -> list[dict[str, Any]]: +def get_virtual_tool_definitions( + default_top_k: int = DEFAULT_MCP_TOOL_SEARCH_TOP_K, +) -> list[dict[str, Any]]: return [ { "name": MCP_TOOL_SEARCH_TOOL_NAME, @@ -49,7 +89,7 @@ def get_virtual_tool_definitions() -> list[dict[str, Any]]: "top_k": { "type": "integer", "description": "Maximum number of results to return.", - "default": 5, + "default": default_top_k, }, }, "required": ["query"], diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5e3ea4b7dcb..dd8150cba35 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1009,6 +1009,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): models: Optional[List[str]] = None search_tools: Optional[List[str]] = None mcp_tool_search_enabled: Optional[bool] = None + mcp_tool_search_top_k: Optional[int] = None from litellm.types.object_permission import ( # noqa: E402 diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index a23cecc3911..faa66fae557 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -280,6 +280,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + mcp_tool_search_top_k Int? teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] diff --git a/litellm/types/object_permission.py b/litellm/types/object_permission.py index d0458173fbf..024e141781d 100644 --- a/litellm/types/object_permission.py +++ b/litellm/types/object_permission.py @@ -25,3 +25,4 @@ class ObjectPermissionDict(TypedDict, total=False): models: Optional[list[str]] search_tools: Optional[list[str]] mcp_tool_search_enabled: Optional[bool] + mcp_tool_search_top_k: Optional[int] diff --git a/schema.prisma b/schema.prisma index a23cecc3911..faa66fae557 100644 --- a/schema.prisma +++ b/schema.prisma @@ -280,6 +280,7 @@ model LiteLLM_ObjectPermissionTable { mcp_toolsets String[] @default([]) // Toolset IDs granted to this key/team/user search_tools String[] @default([]) // search_tool_name values this key/team/user may call mcp_tool_search_enabled Boolean? + mcp_tool_search_top_k Int? teams LiteLLM_TeamTable[] projects LiteLLM_ProjectTable[] verification_tokens LiteLLM_VerificationToken[] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index b8f0b205831..e860365954e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -19,8 +19,11 @@ from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.proxy._experimental.mcp_server.tool_search import ( MCP_TOOL_CALL_TOOL_NAME, MCP_TOOL_SEARCH_TOOL_NAME, + DEFAULT_MCP_TOOL_SEARCH_TOP_K, coerce_top_k, + get_mcp_tool_search_default_top_k, get_virtual_tool_definitions, + resolve_mcp_tool_search_top_k, search_tools, ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -72,6 +75,58 @@ class TestCoerceTopK: assert coerce_top_k("nope", default=10) == 10 +class TestMcpToolSearchDefaultTopK: + def test_builtin_default(self) -> None: + assert get_mcp_tool_search_default_top_k() == DEFAULT_MCP_TOOL_SEARCH_TOP_K + + def test_per_key_override(self) -> None: + uak = UserAPIKeyAuth( + api_key="k", + object_permission=_make_perm(mcp_tool_search_top_k=10), + ) + assert get_mcp_tool_search_default_top_k(uak) == 10 + + def test_global_litellm_settings_override(self, monkeypatch: pytest.MonkeyPatch) -> None: + mock_config = MagicMock() + mock_config.get_config_state.return_value = { + "litellm_settings": {"mcp_tool_search_default_top_k": 12} + } + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_config", + mock_config, + ) + assert get_mcp_tool_search_default_top_k() == 12 + + def test_per_key_beats_global(self, monkeypatch: pytest.MonkeyPatch) -> None: + mock_config = MagicMock() + mock_config.get_config_state.return_value = { + "litellm_settings": {"mcp_tool_search_default_top_k": 12} + } + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_config", + mock_config, + ) + uak = UserAPIKeyAuth( + api_key="k", + object_permission=_make_perm(mcp_tool_search_top_k=8), + ) + assert get_mcp_tool_search_default_top_k(uak) == 8 + + def test_resolve_uses_explicit_top_k(self) -> None: + uak = UserAPIKeyAuth( + api_key="k", + object_permission=_make_perm(mcp_tool_search_top_k=10), + ) + assert resolve_mcp_tool_search_top_k(3, uak) == 3 + + def test_resolve_uses_default_when_omitted(self) -> None: + uak = UserAPIKeyAuth( + api_key="k", + object_permission=_make_perm(mcp_tool_search_top_k=10), + ) + assert resolve_mcp_tool_search_top_k(None, uak) == 10 + + class TestSearchTools: def test_returns_matching_tools(self) -> None: results = search_tools("github issue", SAMPLE_TOOLS) @@ -139,6 +194,11 @@ class TestGetVirtualToolDefinitions: assert "arguments" in props assert "tool_name" in call_tool["inputSchema"]["required"] + def test_mcp_tool_search_schema_top_k_default(self) -> None: + tools = get_virtual_tool_definitions(default_top_k=10) + search_tool = next(t for t in tools if t["name"] == MCP_TOOL_SEARCH_TOOL_NAME) + assert search_tool["inputSchema"]["properties"]["top_k"]["default"] == 10 + def test_all_tools_have_description(self) -> None: for tool in get_virtual_tool_definitions(): assert tool.get("description"), f"{tool['name']} missing description" From d8868be936739d9110f14272c789c1fc7deb4b07 Mon Sep 17 00:00:00 2001 From: Hashim Khan Date: Wed, 29 Jul 2026 19:06:48 +0500 Subject: [PATCH 2/6] fix(mcp): pass search settings without a cyclic import --- .../mcp_server/rest_endpoints.py | 13 ++++++-- .../proxy/_experimental/mcp_server/server.py | 15 +++++++-- .../_experimental/mcp_server/tool_search.py | 17 ++++------ .../mcp_server/test_mcp_tool_search.py | 32 ++++++++----------- 4 files changed, 43 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index d2c267c0b47..5dba5ce435e 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -168,7 +168,11 @@ if MCP_AVAILABLE: if tool_name == MCP_TOOL_SEARCH_TOOL_NAME: return await handle_mcp_tool_search( query=tool_arguments.get("query", ""), - top_k=resolve_mcp_tool_search_top_k(tool_arguments.get("top_k"), user_api_key_dict), + top_k=resolve_mcp_tool_search_top_k( + tool_arguments.get("top_k"), + user_api_key_dict, + proxy_config.get_config_state().get("litellm_settings") or {}, + ), user_api_key_dict=user_api_key_dict, client_ip=rest_client_ip, mcp_auth_header=virtual_mcp_auth_header, @@ -736,10 +740,15 @@ if MCP_AVAILABLE: get_mcp_tool_search_default_top_k, get_virtual_tool_definitions, ) + from litellm.proxy.proxy_server import proxy_config return { "tools": get_virtual_tool_definitions( - default_top_k=get_mcp_tool_search_default_top_k(user_api_key_dict) + default_top_k=get_mcp_tool_search_default_top_k( + user_api_key_dict, + proxy_config.get_config_state().get("litellm_settings") + or {}, + ) ), "error": None, "message": "Successfully retrieved tools", diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 231dc828e85..7df2c4a73f3 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -778,11 +778,16 @@ if MCP_AVAILABLE: get_mcp_tool_search_default_top_k, get_virtual_tool_definitions, ) + from litellm.proxy.proxy_server import proxy_config return [ Tool(**d) for d in get_virtual_tool_definitions( - default_top_k=get_mcp_tool_search_default_top_k(user_api_key_auth) + default_top_k=get_mcp_tool_search_default_top_k( + user_api_key_auth, + proxy_config.get_config_state().get("litellm_settings") + or {}, + ) ) ] @@ -933,9 +938,15 @@ if MCP_AVAILABLE: args = arguments or {} if name == MCP_TOOL_SEARCH_TOOL_NAME: + from litellm.proxy.proxy_server import proxy_config + return await handle_mcp_tool_search( query=args.get("query", ""), - top_k=resolve_mcp_tool_search_top_k(args.get("top_k"), user_api_key_auth), + top_k=resolve_mcp_tool_search_top_k( + args.get("top_k"), + user_api_key_auth, + proxy_config.get_config_state().get("litellm_settings") or {}, + ), user_api_key_dict=user_api_key_auth, client_ip=client_ip, mcp_servers=mcp_servers, diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index d4059deb706..8d205bfdfd1 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -15,17 +15,9 @@ MCP_TOOL_CALL_TOOL_NAME: str = "mcp_tool_call" DEFAULT_MCP_TOOL_SEARCH_TOP_K: int = 5 -def _get_litellm_settings() -> dict[str, Any]: - try: - from litellm.proxy.proxy_server import proxy_config - - return proxy_config.get_config_state().get("litellm_settings") or {} - except Exception: - return {} - - def get_mcp_tool_search_default_top_k( user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + litellm_settings: Optional[dict[str, Any]] = None, ) -> int: """Resolve the default top_k for mcp_tool_search (per-key, then global, then 5).""" if user_api_key_dict is not None: @@ -35,7 +27,7 @@ def get_mcp_tool_search_default_top_k( if key_top_k is not None: return coerce_top_k(key_top_k, default=DEFAULT_MCP_TOOL_SEARCH_TOP_K) - global_top_k = _get_litellm_settings().get("mcp_tool_search_default_top_k") + global_top_k = (litellm_settings or {}).get("mcp_tool_search_default_top_k") if global_top_k is not None: return coerce_top_k(global_top_k, default=DEFAULT_MCP_TOOL_SEARCH_TOP_K) @@ -45,8 +37,11 @@ def get_mcp_tool_search_default_top_k( def resolve_mcp_tool_search_top_k( explicit_top_k: Any, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + litellm_settings: Optional[dict[str, Any]] = None, ) -> int: - default_top_k = get_mcp_tool_search_default_top_k(user_api_key_dict) + default_top_k = get_mcp_tool_search_default_top_k( + user_api_key_dict, litellm_settings + ) if explicit_top_k is None: return default_top_k return coerce_top_k(explicit_top_k, default=default_top_k) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index 5695932b54e..cae95c2b6cd 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -91,31 +91,25 @@ class TestMcpToolSearchDefaultTopK: ) assert get_mcp_tool_search_default_top_k(uak) == 10 - def test_global_litellm_settings_override(self, monkeypatch: pytest.MonkeyPatch) -> None: - mock_config = MagicMock() - mock_config.get_config_state.return_value = { - "litellm_settings": {"mcp_tool_search_default_top_k": 12} - } - monkeypatch.setattr( - "litellm.proxy.proxy_server.proxy_config", - mock_config, + def test_global_litellm_settings_override(self) -> None: + assert ( + get_mcp_tool_search_default_top_k( + litellm_settings={"mcp_tool_search_default_top_k": 12} + ) + == 12 ) - assert get_mcp_tool_search_default_top_k() == 12 - def test_per_key_beats_global(self, monkeypatch: pytest.MonkeyPatch) -> None: - mock_config = MagicMock() - mock_config.get_config_state.return_value = { - "litellm_settings": {"mcp_tool_search_default_top_k": 12} - } - monkeypatch.setattr( - "litellm.proxy.proxy_server.proxy_config", - mock_config, - ) + def test_per_key_beats_global(self) -> None: uak = UserAPIKeyAuth( api_key="k", object_permission=_make_perm(mcp_tool_search_top_k=8), ) - assert get_mcp_tool_search_default_top_k(uak) == 8 + assert ( + get_mcp_tool_search_default_top_k( + uak, {"mcp_tool_search_default_top_k": 12} + ) + == 8 + ) def test_resolve_uses_explicit_top_k(self) -> None: uak = UserAPIKeyAuth( From f97c1441a3676a695db69881a404a4bcb6d15a39 Mon Sep 17 00:00:00 2001 From: Hashim Khan Date: Wed, 29 Jul 2026 19:23:38 +0500 Subject: [PATCH 3/6] Fix MCP search CI validation --- .../_experimental/mcp_server/rest_endpoints.py | 5 +---- litellm/proxy/_experimental/mcp_server/server.py | 3 +-- .../proxy/_experimental/mcp_server/tool_search.py | 8 ++++---- .../mcp_server/test_mcp_tool_search.py | 14 ++------------ .../test_customer_endpoints.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 ++++ 6 files changed, 13 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 5dba5ce435e..03a25b30a3a 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -143,8 +143,6 @@ if MCP_AVAILABLE: ) from litellm.proxy._experimental.mcp_server.tool_search import ( MCP_TOOL_SEARCH_TOOL_NAME, - get_mcp_tool_search_default_top_k, - get_virtual_tool_definitions, handle_mcp_tool_call, handle_mcp_tool_search, resolve_mcp_tool_search_top_k, @@ -746,8 +744,7 @@ if MCP_AVAILABLE: "tools": get_virtual_tool_definitions( default_top_k=get_mcp_tool_search_default_top_k( user_api_key_dict, - proxy_config.get_config_state().get("litellm_settings") - or {}, + proxy_config.get_config_state().get("litellm_settings") or {}, ) ), "error": None, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7df2c4a73f3..c0a7974011e 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -785,8 +785,7 @@ if MCP_AVAILABLE: for d in get_virtual_tool_definitions( default_top_k=get_mcp_tool_search_default_top_k( user_api_key_auth, - proxy_config.get_config_state().get("litellm_settings") - or {}, + proxy_config.get_config_state().get("litellm_settings") or {}, ) ) ] diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 8d205bfdfd1..21d0db6ca55 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -39,9 +39,7 @@ def resolve_mcp_tool_search_top_k( user_api_key_dict: Optional["UserAPIKeyAuth"] = None, litellm_settings: Optional[dict[str, Any]] = None, ) -> int: - default_top_k = get_mcp_tool_search_default_top_k( - user_api_key_dict, litellm_settings - ) + default_top_k = get_mcp_tool_search_default_top_k(user_api_key_dict, litellm_settings) if explicit_top_k is None: return default_top_k return coerce_top_k(explicit_top_k, default=default_top_k) @@ -55,7 +53,9 @@ def coerce_top_k(value: Any, default: int = DEFAULT_MCP_TOOL_SEARCH_TOP_K) -> in return result if result > 0 else default -def search_tools(query: str, tools: list[dict[str, Any]], top_k: int = DEFAULT_MCP_TOOL_SEARCH_TOP_K) -> list[dict[str, Any]]: +def search_tools( + query: str, tools: list[dict[str, Any]], top_k: int = DEFAULT_MCP_TOOL_SEARCH_TOP_K +) -> list[dict[str, Any]]: if not query: return [] tokens = query.lower().split() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py index cae95c2b6cd..6ef9ad4d475 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_tool_search.py @@ -92,24 +92,14 @@ class TestMcpToolSearchDefaultTopK: assert get_mcp_tool_search_default_top_k(uak) == 10 def test_global_litellm_settings_override(self) -> None: - assert ( - get_mcp_tool_search_default_top_k( - litellm_settings={"mcp_tool_search_default_top_k": 12} - ) - == 12 - ) + assert get_mcp_tool_search_default_top_k(litellm_settings={"mcp_tool_search_default_top_k": 12}) == 12 def test_per_key_beats_global(self) -> None: uak = UserAPIKeyAuth( api_key="k", object_permission=_make_perm(mcp_tool_search_top_k=8), ) - assert ( - get_mcp_tool_search_default_top_k( - uak, {"mcp_tool_search_default_top_k": 12} - ) - == 8 - ) + assert get_mcp_tool_search_default_top_k(uak, {"mcp_tool_search_default_top_k": 12}) == 8 def test_resolve_uses_explicit_top_k(self) -> None: uak = UserAPIKeyAuth( diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 5fbc3c4869b..5e497853d32 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -719,6 +719,7 @@ _EXPECTED_CUSTOMER = { "blocked_tools": [], "search_tools": [], "mcp_tool_search_enabled": None, + "mcp_tool_search_top_k": None, }, } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 541d7a17ae1..d3390c1ed0b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25611,6 +25611,8 @@ export interface components { } | null; /** Mcp Tool Search Enabled */ mcp_tool_search_enabled?: boolean | null; + /** Mcp Tool Search Top K */ + mcp_tool_search_top_k?: number | null; /** Mcp Toolsets */ mcp_toolsets?: string[] | null; /** Models */ @@ -25656,6 +25658,8 @@ export interface components { } | null; /** Mcp Tool Search Enabled */ mcp_tool_search_enabled?: boolean | null; + /** Mcp Tool Search Top K */ + mcp_tool_search_top_k?: number | null; /** Mcp Toolsets */ mcp_toolsets?: string[] | null; /** From a4f724316cb27e33c25e486c0c4a34a8bd374efd Mon Sep 17 00:00:00 2001 From: Hashim Khan Date: Wed, 29 Jul 2026 19:35:09 +0500 Subject: [PATCH 4/6] Resolve strict MCP search lint violations --- litellm/models/object_permission.py | 2 +- .../proxy/_experimental/mcp_server/tool_search.py | 14 ++++++++------ litellm/proxy/_types.py | 2 +- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/litellm/models/object_permission.py b/litellm/models/object_permission.py index ec9b6671e8b..218957037cd 100644 --- a/litellm/models/object_permission.py +++ b/litellm/models/object_permission.py @@ -25,4 +25,4 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase): blocked_tools: Optional[List[str]] = [] search_tools: Optional[List[str]] = [] mcp_tool_search_enabled: Optional[bool] = None - mcp_tool_search_top_k: Optional[int] = None + mcp_tool_search_top_k: int | None = None diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 21d0db6ca55..ca553ea8013 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -16,8 +16,8 @@ DEFAULT_MCP_TOOL_SEARCH_TOP_K: int = 5 def get_mcp_tool_search_default_top_k( - user_api_key_dict: Optional["UserAPIKeyAuth"] = None, - litellm_settings: Optional[dict[str, Any]] = None, + user_api_key_dict: UserAPIKeyAuth | None = None, + litellm_settings: dict[str, object] | None = None, ) -> int: """Resolve the default top_k for mcp_tool_search (per-key, then global, then 5).""" if user_api_key_dict is not None: @@ -35,9 +35,9 @@ def get_mcp_tool_search_default_top_k( def resolve_mcp_tool_search_top_k( - explicit_top_k: Any, - user_api_key_dict: Optional["UserAPIKeyAuth"] = None, - litellm_settings: Optional[dict[str, Any]] = None, + explicit_top_k: object, + user_api_key_dict: UserAPIKeyAuth | None = None, + litellm_settings: dict[str, object] | None = None, ) -> int: default_top_k = get_mcp_tool_search_default_top_k(user_api_key_dict, litellm_settings) if explicit_top_k is None: @@ -45,7 +45,9 @@ def resolve_mcp_tool_search_top_k( return coerce_top_k(explicit_top_k, default=default_top_k) -def coerce_top_k(value: Any, default: int = DEFAULT_MCP_TOOL_SEARCH_TOP_K) -> int: +def coerce_top_k(value: object, default: int = DEFAULT_MCP_TOOL_SEARCH_TOP_K) -> int: + if not isinstance(value, (int, float, str, bytes, bytearray)): + return default try: result = int(value) except (TypeError, ValueError): diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2e4278f2ad2..a414d12b2f9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1018,7 +1018,7 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): models: Optional[List[str]] = None search_tools: Optional[List[str]] = None mcp_tool_search_enabled: Optional[bool] = None - mcp_tool_search_top_k: Optional[int] = None + mcp_tool_search_top_k: int | None = None from litellm.models.team import BudgetLimitEntry as BudgetLimitEntry # noqa: E402 From 5777efa2847ae24cbaa3c726d81c431e5b6369f8 Mon Sep 17 00:00:00 2001 From: Hashim Khan Date: Wed, 29 Jul 2026 19:51:37 +0500 Subject: [PATCH 5/6] Use immutable MCP settings type --- litellm/proxy/_experimental/mcp_server/tool_search.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index ca553ea8013..dc13b81129f 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +from collections.abc import Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Optional @@ -17,7 +18,7 @@ DEFAULT_MCP_TOOL_SEARCH_TOP_K: int = 5 def get_mcp_tool_search_default_top_k( user_api_key_dict: UserAPIKeyAuth | None = None, - litellm_settings: dict[str, object] | None = None, + litellm_settings: Mapping[str, object] | None = None, ) -> int: """Resolve the default top_k for mcp_tool_search (per-key, then global, then 5).""" if user_api_key_dict is not None: @@ -37,7 +38,7 @@ def get_mcp_tool_search_default_top_k( def resolve_mcp_tool_search_top_k( explicit_top_k: object, user_api_key_dict: UserAPIKeyAuth | None = None, - litellm_settings: dict[str, object] | None = None, + litellm_settings: Mapping[str, object] | None = None, ) -> int: default_top_k = get_mcp_tool_search_default_top_k(user_api_key_dict, litellm_settings) if explicit_top_k is None: From dd4d2f44e33118f1d0b1ea3884a3c5808cc6b195 Mon Sep 17 00:00:00 2001 From: Hashim1999164 <64767361+Hashim1999164@users.noreply.github.com> Date: Thu, 13 Aug 2026 04:13:51 +0500 Subject: [PATCH 6/6] Avoid empty dict fallbacks in mcp tool search settings --- litellm/proxy/_experimental/mcp_server/rest_endpoints.py | 8 ++++++-- litellm/proxy/_experimental/mcp_server/server.py | 8 ++++++-- litellm/proxy/_experimental/mcp_server/tool_search.py | 5 ++++- litellm/types/object_permission.py | 2 +- 4 files changed, 17 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 1f49214fcd7..d6da662fd09 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -189,12 +189,14 @@ if MCP_AVAILABLE: ) = _extract_mcp_headers_from_request(request, MCPRequestHandler) virtual_oauth2_headers: Final = MCPRequestHandler._get_oauth2_headers_from_headers(request.headers) if tool_name == MCP_TOOL_SEARCH_TOOL_NAME: + raw_settings = proxy_config.get_config_state().get("litellm_settings") + settings = raw_settings if isinstance(raw_settings, Mapping) else None return await handle_mcp_tool_search( query=tool_arguments.get("query", ""), top_k=resolve_mcp_tool_search_top_k( tool_arguments.get("top_k"), user_api_key_dict, - proxy_config.get_config_state().get("litellm_settings") or {}, + settings, ), user_api_key_dict=user_api_key_dict, client_ip=rest_client_ip, @@ -767,11 +769,13 @@ if MCP_AVAILABLE: ) from litellm.proxy.proxy_server import proxy_config + raw_settings = proxy_config.get_config_state().get("litellm_settings") + settings = raw_settings if isinstance(raw_settings, Mapping) else None return { "tools": get_virtual_tool_definitions( default_top_k=get_mcp_tool_search_default_top_k( user_api_key_dict, - proxy_config.get_config_state().get("litellm_settings") or {}, + settings, ) ), "error": None, diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 350335577bd..157351c429b 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -791,12 +791,14 @@ if MCP_AVAILABLE: ) from litellm.proxy.proxy_server import proxy_config + raw_settings = proxy_config.get_config_state().get("litellm_settings") + settings = raw_settings if isinstance(raw_settings, Mapping) else None return [ Tool.model_validate(d) for d in get_virtual_tool_definitions( default_top_k=get_mcp_tool_search_default_top_k( user_api_key_auth, - proxy_config.get_config_state().get("litellm_settings") or {}, + settings, ) ) ] @@ -950,12 +952,14 @@ if MCP_AVAILABLE: if name == MCP_TOOL_SEARCH_TOOL_NAME: from litellm.proxy.proxy_server import proxy_config + raw_settings = proxy_config.get_config_state().get("litellm_settings") + settings = raw_settings if isinstance(raw_settings, Mapping) else None return await handle_mcp_tool_search( query=args.get("query", ""), top_k=resolve_mcp_tool_search_top_k( args.get("top_k"), user_api_key_auth, - proxy_config.get_config_state().get("litellm_settings") or {}, + settings, ), user_api_key_dict=user_api_key_auth, client_ip=client_ip, diff --git a/litellm/proxy/_experimental/mcp_server/tool_search.py b/litellm/proxy/_experimental/mcp_server/tool_search.py index 6f2d94320e4..226a54bbf5c 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_search.py +++ b/litellm/proxy/_experimental/mcp_server/tool_search.py @@ -28,7 +28,10 @@ def get_mcp_tool_search_default_top_k( if key_top_k is not None: return coerce_top_k(key_top_k, default=DEFAULT_MCP_TOOL_SEARCH_TOP_K) - global_top_k = (litellm_settings or {}).get("mcp_tool_search_default_top_k") + if litellm_settings is None: + global_top_k = None + else: + global_top_k = litellm_settings.get("mcp_tool_search_default_top_k") if global_top_k is not None: return coerce_top_k(global_top_k, default=DEFAULT_MCP_TOOL_SEARCH_TOP_K) diff --git a/litellm/types/object_permission.py b/litellm/types/object_permission.py index a84355ebb11..52cf9f06422 100644 --- a/litellm/types/object_permission.py +++ b/litellm/types/object_permission.py @@ -23,4 +23,4 @@ class ObjectPermissionDict(TypedDict, total=False): models: list[str] | None search_tools: list[str] | None mcp_tool_search_enabled: bool | None - mcp_tool_search_top_k: int | None + mcp_tool_search_top_k: int | None # writable-ok: mutated on in-memory dict payloads before persistence