mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(mcp): record OpenAPI tool definitions for pre-call guardrail metadata
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
c23ae322c1
commit
eba05f4333
2 changed files with 39 additions and 3 deletions
|
|
@ -4442,11 +4442,9 @@ class MCPServerManager:
|
|||
# applied (e.g. "test_petstore-getinventory"). Do NOT pass them
|
||||
# through _create_prefixed_tools — that would add the prefix a second
|
||||
# time producing "test_petstore-test_petstore-getinventory".
|
||||
if add_prefix:
|
||||
return tools
|
||||
prefix: Final = get_server_prefix(server)
|
||||
sep: Final = MCP_TOOL_PREFIX_SEPARATOR
|
||||
return [ # mutable-ok: returned through the list[MCPTool] listing contract
|
||||
unprefixed_tools: Final = [ # mutable-ok: returned through the list[MCPTool] listing contract
|
||||
(
|
||||
t.model_copy(update={"name": t.name[len(prefix) + len(sep) :]})
|
||||
if t.name.startswith(f"{prefix}{sep}")
|
||||
|
|
@ -4454,6 +4452,10 @@ class MCPServerManager:
|
|||
)
|
||||
for t in tools
|
||||
]
|
||||
self._listed_tools_by_server_id[server.server_id] = MappingProxyType(
|
||||
{t.name: t for t in unprefixed_tools}
|
||||
)
|
||||
return tools if add_prefix else unprefixed_tools
|
||||
else:
|
||||
tools = await self._fetch_tools_with_timeout(client, server.name)
|
||||
self._remember_upstream_initialize_instructions(server, client)
|
||||
|
|
|
|||
|
|
@ -6632,6 +6632,40 @@ class TestMCPServerManager:
|
|||
listed = manager.get_listed_tool(server, "echo")
|
||||
assert listed is not None and listed.description == "shared"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("add_prefix", [True, False])
|
||||
async def test_openapi_listing_records_listed_tools(self, add_prefix):
|
||||
from litellm.proxy._experimental.mcp_server.tool_registry import global_mcp_tool_registry
|
||||
|
||||
server = MCPServer(
|
||||
server_id="petstore-id",
|
||||
name="petstore",
|
||||
alias="petstore",
|
||||
transport=MCPTransport.http,
|
||||
url=None,
|
||||
spec_path="/spec.yaml",
|
||||
)
|
||||
manager = MCPServerManager()
|
||||
manager._create_mcp_client = AsyncMock(return_value=AsyncMock())
|
||||
|
||||
async def _handler(**kwargs):
|
||||
return None
|
||||
|
||||
with patch.dict(global_mcp_tool_registry.tools, {}, clear=True):
|
||||
global_mcp_tool_registry.register_tool(
|
||||
name="petstore-list_pets",
|
||||
description="List pets",
|
||||
input_schema={"type": "object", "properties": {"limit": {"type": "integer"}}},
|
||||
handler=_handler,
|
||||
)
|
||||
listed = await manager._get_tools_from_server(server=server, add_prefix=add_prefix)
|
||||
|
||||
assert [t.name for t in listed] == ["petstore-list_pets" if add_prefix else "list_pets"]
|
||||
for name in ("list_pets", "petstore-list_pets"):
|
||||
tool = manager.get_listed_tool(server, name)
|
||||
assert tool is not None and tool.description == "List pets"
|
||||
assert tool.inputSchema["properties"] == {"limit": {"type": "integer"}}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_allowed_mcp_servers_with_user_api_key_auth(self):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue