fix(mcp): coerce skill_search top_k instead of raising 500 on malformed input

The MCP-REST skill_search dispatch validated raw tool arguments through a
pydantic model directly, so a non-numeric top_k raised a ValidationError
that the endpoint's catch-all turned into an HTTP 500. Mirrors the
agent_search branch's tolerant coerce_top_k handling instead.
This commit is contained in:
Yassin Kortam 2026-09-02 15:45:48 -07:00
parent cce83d649a
commit 6f3447bb67
2 changed files with 32 additions and 11 deletions

View file

@ -8,7 +8,6 @@ from typing import TYPE_CHECKING, Any, Final, Literal
import anyio
import httpx
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from pydantic import BaseModel, ConfigDict
from litellm._logging import verbose_logger
from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT
@ -163,12 +162,6 @@ if MCP_AVAILABLE:
request_path=request.scope.get("_original_path") or request.url.path,
)
class _SkillSearchToolArguments(BaseModel):
model_config = ConfigDict(extra="ignore")
query: str = ""
top_k: int = DEFAULT_SKILL_SEARCH_TOP_K
async def _handle_virtual_mcp_tool(
request: Request,
data: dict[str, Any],
@ -211,10 +204,11 @@ if MCP_AVAILABLE:
user_api_key_dict=user_api_key_dict,
)
if tool_name == SKILL_SEARCH_TOOL_NAME:
skill_search_args: Final = _SkillSearchToolArguments.model_validate(tool_arguments)
return await handle_skill_search(
query=skill_search_args.query,
top_k=skill_search_args.top_k,
query=str(tool_arguments.get("query", "")),
top_k=coerce_top_k(
tool_arguments.get("top_k", DEFAULT_SKILL_SEARCH_TOP_K), default=DEFAULT_SKILL_SEARCH_TOP_K
),
user_api_key_dict=user_api_key_dict,
)
rest_client_ip: Final = IPAddressUtils.get_mcp_client_ip(request)

View file

@ -744,6 +744,31 @@ class TestCallToolRestApiVirtualTools:
assert mock_search.await_args.kwargs["top_k"] == 1
assert mock_search.await_args.kwargs["agents"] == (translator,)
@pytest.mark.asyncio
async def test_skill_search_call_tolerates_malformed_top_k(self) -> None:
"""Regression: a caller-supplied non-numeric top_k must be coerced to the default,
the same as agent_search, instead of raising a pydantic ValidationError that the
endpoint's catch-all turns into an HTTP 500."""
from mcp.types import CallToolResult, TextContent
from litellm.llms.litellm_proxy.skills.skill_search import DEFAULT_SKILL_SEARCH_TOP_K
user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
request = self._make_request(
{"name": SKILL_SEARCH_TOOL_NAME, "arguments": {"query": "translate a document", "top_k": "not-a-number"}}
)
fake_result = CallToolResult(content=[TextContent(type="text", text="[]")], isError=False)
with patch( # test-quality-ok: the embedding router only resolves via proxy_server globals, no injection seam
"litellm.proxy._experimental.mcp_server.tool_search.handle_skill_search",
new_callable=AsyncMock,
return_value=fake_result,
) as mock_search:
result = await self._get_call_fn()(request=request, user_api_key_dict=user_api_key_dict)
assert result.isError is False
assert mock_search.await_args.kwargs["top_k"] == DEFAULT_SKILL_SEARCH_TOP_K
assert mock_search.await_args.kwargs["query"] == "translate a document"
@pytest.mark.asyncio
async def test_agent_search_call_reports_missing_embedding_model_as_tool_error(self) -> None:
from litellm.proxy.agent_endpoints.agent_search import AgentSearchNotConfigured
@ -817,7 +842,9 @@ class TestCallToolRestApiVirtualTools:
assert "mcp_tool_search.embedding_model" in result.content[0].text
@pytest.mark.asyncio
async def test_mcp_tool_search_reports_invalid_settings_as_tool_error(self, monkeypatch: pytest.MonkeyPatch) -> None:
async def test_mcp_tool_search_reports_invalid_settings_as_tool_error(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "mcp_tool_search", {"top_k": 0})
user_api_key_dict = UserAPIKeyAuth(api_key="k", object_permission=_make_perm(mcp_tool_search_enabled=True))
result = await self._get_call_fn()(request=self._semantic_request(), user_api_key_dict=user_api_key_dict)