mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Merge pull request #25694 from milan-berri/feat/mcp-initialize-instructions
feat(mcp): expose per-server InitializeResult.instructions from gateway
This commit is contained in:
ishaan-berri
GitHub
commit
d479234f0e
14 changed files with 452 additions and 4 deletions
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "instructions" TEXT;
|
||||
|
|
@ -289,6 +289,7 @@ model LiteLLM_MCPServerTable {
|
|||
server_name String?
|
||||
alias String?
|
||||
description String?
|
||||
instructions String?
|
||||
url String?
|
||||
spec_path String?
|
||||
transport String @default("sse")
|
||||
|
|
|
|||
|
|
@ -221,6 +221,7 @@ class MCPClient:
|
|||
self.extra_headers: Optional[Dict[str, str]] = extra_headers
|
||||
self.ssl_verify: Optional[VerifyTypes] = ssl_verify
|
||||
self._aws_auth: Optional[httpx.Auth] = aws_auth
|
||||
self._last_initialize_instructions: Optional[str] = None
|
||||
# handle the basic auth value if provided
|
||||
if auth_value:
|
||||
self.update_auth_value(auth_value)
|
||||
|
|
@ -296,7 +297,12 @@ class MCPClient:
|
|||
session_ctx = ClientSession(read_stream, write_stream)
|
||||
session = await session_ctx.__aenter__()
|
||||
try:
|
||||
await session.initialize()
|
||||
init_result = await session.initialize()
|
||||
self._last_initialize_instructions = None
|
||||
if init_result is not None:
|
||||
ins = getattr(init_result, "instructions", None)
|
||||
if isinstance(ins, str) and ins.strip():
|
||||
self._last_initialize_instructions = ins.strip()
|
||||
return await operation(session)
|
||||
finally:
|
||||
try:
|
||||
|
|
@ -315,6 +321,7 @@ class MCPClient:
|
|||
"""Open a session, run the provided coroutine, and clean up."""
|
||||
http_client: Optional[httpx.AsyncClient] = None
|
||||
try:
|
||||
self._last_initialize_instructions = None
|
||||
transport_ctx, http_client = self._create_transport_context()
|
||||
return await self._execute_session_operation(transport_ctx, operation)
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -14,3 +14,8 @@ from typing import Optional
|
|||
_mcp_active_toolset_id: ContextVar[Optional[str]] = ContextVar(
|
||||
"_mcp_active_toolset_id", default=None
|
||||
)
|
||||
|
||||
# Per-request merged InitializeResult.instructions; set in MCP HTTP/SSE handlers.
|
||||
_mcp_gateway_initialize_instructions: ContextVar[Optional[str]] = ContextVar(
|
||||
"_mcp_gateway_initialize_instructions", default=None
|
||||
)
|
||||
|
|
|
|||
|
|
@ -184,6 +184,16 @@ class MCPServerManager:
|
|||
"gmail_send_email": "zapier_mcp_server",
|
||||
}
|
||||
"""
|
||||
self._upstream_initialize_instructions_by_server_id: Dict[str, str] = {}
|
||||
|
||||
def _remember_upstream_initialize_instructions(
|
||||
self, server: MCPServer, client: MCPClient
|
||||
) -> None:
|
||||
raw = getattr(client, "_last_initialize_instructions", None)
|
||||
if raw and str(raw).strip():
|
||||
self._upstream_initialize_instructions_by_server_id[server.server_id] = (
|
||||
str(raw).strip()
|
||||
)
|
||||
|
||||
def get_registry(self) -> Dict[str, MCPServer]:
|
||||
"""
|
||||
|
|
@ -204,6 +214,7 @@ class MCPServerManager:
|
|||
mcp_aliases: Optional dictionary mapping aliases to server names from litellm_settings
|
||||
"""
|
||||
verbose_logger.debug("Loading MCP Servers from config-----")
|
||||
self._upstream_initialize_instructions_by_server_id.clear()
|
||||
|
||||
# Track which aliases have been used to ensure only first occurrence is used
|
||||
used_aliases = set()
|
||||
|
|
@ -351,6 +362,7 @@ class MCPServerManager:
|
|||
aws_service_name=server_config.get("aws_service_name", None),
|
||||
aws_role_name=server_config.get("aws_role_name", None),
|
||||
aws_session_name=server_config.get("aws_session_name", None),
|
||||
instructions=server_config.get("instructions", None),
|
||||
)
|
||||
self.config_mcp_servers[server_id] = new_server
|
||||
|
||||
|
|
@ -693,6 +705,7 @@ class MCPServerManager:
|
|||
aws_service_name=aws_creds.get("aws_service_name"),
|
||||
aws_role_name=aws_creds.get("aws_role_name"),
|
||||
aws_session_name=aws_creds.get("aws_session_name"),
|
||||
instructions=mcp_server.instructions,
|
||||
)
|
||||
return new_server
|
||||
|
||||
|
|
@ -1247,6 +1260,7 @@ class MCPServerManager:
|
|||
return tools
|
||||
else:
|
||||
tools = await self._fetch_tools_with_timeout(client, server.name)
|
||||
self._remember_upstream_initialize_instructions(server, client)
|
||||
|
||||
prefixed_or_original_tools = self._create_prefixed_tools(
|
||||
tools, server, add_prefix=add_prefix
|
||||
|
|
@ -2383,6 +2397,7 @@ class MCPServerManager:
|
|||
# If proxy_logging_obj is not None, the tool call result is at index 1 (after the during hook task)
|
||||
result_index = 1 if proxy_logging_obj else 0
|
||||
result = mcp_responses[result_index]
|
||||
self._remember_upstream_initialize_instructions(mcp_server, client)
|
||||
|
||||
return cast(CallToolResult, result)
|
||||
|
||||
|
|
@ -2622,6 +2637,7 @@ class MCPServerManager:
|
|||
)
|
||||
|
||||
verbose_logger.debug("Loading MCP servers from database into registry...")
|
||||
self._upstream_initialize_instructions_by_server_id.clear()
|
||||
|
||||
# perform authz check to filter the mcp servers user has access to
|
||||
prisma_client = get_prisma_client_or_throw(
|
||||
|
|
@ -2905,6 +2921,7 @@ class MCPServerManager:
|
|||
await asyncio.wait_for(
|
||||
client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT
|
||||
)
|
||||
self._remember_upstream_initialize_instructions(server, client)
|
||||
status = "healthy"
|
||||
except asyncio.TimeoutError:
|
||||
health_check_error = (
|
||||
|
|
@ -2946,6 +2963,7 @@ class MCPServerManager:
|
|||
token_url=server.token_url,
|
||||
registration_url=server.registration_url,
|
||||
allow_all_keys=server.allow_all_keys,
|
||||
instructions=server.instructions,
|
||||
)
|
||||
|
||||
async def get_all_mcp_servers_with_health_and_teams(
|
||||
|
|
@ -3041,6 +3059,7 @@ class MCPServerManager:
|
|||
is_byok=server.is_byok,
|
||||
byok_description=server.byok_description,
|
||||
byok_api_key_help_url=server.byok_api_key_help_url,
|
||||
instructions=server.instructions,
|
||||
)
|
||||
|
||||
async def get_all_mcp_servers_unfiltered(self) -> List[LiteLLM_MCPServerTable]:
|
||||
|
|
|
|||
|
|
@ -933,6 +933,7 @@ if MCP_AVAILABLE:
|
|||
authorization_url=request.authorization_url,
|
||||
registration_url=request.registration_url,
|
||||
oauth2_flow=_oauth2_flow,
|
||||
instructions=request.instructions,
|
||||
)
|
||||
|
||||
stdio_env = global_mcp_server_manager._build_stdio_env(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ LiteLLM MCP Server Routes
|
|||
import asyncio
|
||||
import contextlib
|
||||
import time
|
||||
import types
|
||||
import traceback
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
|
@ -37,7 +38,10 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
|||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
get_request_base_url,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_context import _mcp_active_toolset_id
|
||||
from litellm.proxy._experimental.mcp_server.mcp_context import (
|
||||
_mcp_active_toolset_id,
|
||||
_mcp_gateway_initialize_instructions,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_debug import MCPDebug
|
||||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
LITELLM_MCP_SERVER_DESCRIPTION,
|
||||
|
|
@ -122,6 +126,8 @@ _INITIALIZATION_LOCK = asyncio.Lock()
|
|||
|
||||
if MCP_AVAILABLE:
|
||||
from mcp.server import Server
|
||||
from mcp.server.lowlevel.server import NotificationOptions
|
||||
from mcp.server.models import InitializationOptions
|
||||
|
||||
# Import auth context variables and middleware
|
||||
from mcp.server.auth.middleware.auth_context import (
|
||||
|
|
@ -200,6 +206,21 @@ if MCP_AVAILABLE:
|
|||
)
|
||||
return normalized
|
||||
|
||||
def _gateway_create_initialization_options(
|
||||
self,
|
||||
notification_options: Optional[NotificationOptions] = None,
|
||||
experimental_capabilities: Optional[Dict[str, Dict[str, Any]]] = None,
|
||||
) -> InitializationOptions:
|
||||
opts = Server.create_initialization_options(
|
||||
self,
|
||||
notification_options=notification_options,
|
||||
experimental_capabilities=experimental_capabilities or {},
|
||||
)
|
||||
merged = _mcp_gateway_initialize_instructions.get()
|
||||
if merged is not None:
|
||||
return opts.model_copy(update={"instructions": merged})
|
||||
return opts
|
||||
|
||||
########################################################
|
||||
############ Initialize the MCP Server #################
|
||||
########################################################
|
||||
|
|
@ -207,6 +228,9 @@ if MCP_AVAILABLE:
|
|||
name=LITELLM_MCP_SERVER_NAME,
|
||||
version=LITELLM_MCP_SERVER_VERSION,
|
||||
)
|
||||
server.create_initialization_options = types.MethodType( # type: ignore[method-assign]
|
||||
_gateway_create_initialization_options, server
|
||||
)
|
||||
sse: SseServerTransport = SseServerTransport("/mcp/sse/messages")
|
||||
|
||||
# Create session managers
|
||||
|
|
@ -1103,6 +1127,57 @@ if MCP_AVAILABLE:
|
|||
|
||||
return server_auth_header, extra_headers
|
||||
|
||||
def _merge_gateway_initialize_instructions(
|
||||
allowed_mcp_servers: List[MCPServer],
|
||||
) -> Optional[str]:
|
||||
"""YAML/DB override, else in-memory upstream text from list_tools / health_check / call_tool."""
|
||||
if not allowed_mcp_servers:
|
||||
return None
|
||||
|
||||
texts: List[Tuple[str, str]] = []
|
||||
for server in allowed_mcp_servers:
|
||||
label = (
|
||||
server.alias
|
||||
or server.server_name
|
||||
or server.name
|
||||
or server.server_id
|
||||
or "mcp"
|
||||
)
|
||||
if server.instructions and server.instructions.strip():
|
||||
texts.append((label, server.instructions.strip()))
|
||||
continue
|
||||
if server.spec_path:
|
||||
continue
|
||||
cached = global_mcp_server_manager._upstream_initialize_instructions_by_server_id.get(
|
||||
server.server_id
|
||||
)
|
||||
if cached and cached.strip():
|
||||
texts.append((label, cached.strip()))
|
||||
|
||||
if not texts:
|
||||
return None
|
||||
if len(texts) == 1:
|
||||
return texts[0][1]
|
||||
return "\n\n---\n\n".join(f"[{lbl}]\n{txt}" for lbl, txt in texts)
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def _gateway_initialize_instructions_request_scope(
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth],
|
||||
mcp_servers: Optional[List[str]],
|
||||
client_ip: Optional[str],
|
||||
) -> AsyncIterator[None]:
|
||||
allowed = await _get_allowed_mcp_servers(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
mcp_servers=mcp_servers,
|
||||
client_ip=client_ip,
|
||||
)
|
||||
merged = _merge_gateway_initialize_instructions(allowed_mcp_servers=allowed)
|
||||
tok = _mcp_gateway_initialize_instructions.set(merged)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_mcp_gateway_initialize_instructions.reset(tok)
|
||||
|
||||
async def _get_tools_from_mcp_servers( # noqa: PLR0915
|
||||
user_api_key_auth: Optional[UserAPIKeyAuth],
|
||||
mcp_auth_header: Optional[str],
|
||||
|
|
@ -2670,7 +2745,12 @@ if MCP_AVAILABLE:
|
|||
# Request was fully handled (e.g., DELETE on non-existent session)
|
||||
return
|
||||
|
||||
await session_manager.handle_request(scope, receive, send)
|
||||
async with _gateway_initialize_instructions_request_scope(
|
||||
user_api_key_auth,
|
||||
mcp_servers,
|
||||
_client_ip,
|
||||
):
|
||||
await session_manager.handle_request(scope, receive, send)
|
||||
except HTTPException:
|
||||
# Re-raise HTTP exceptions to preserve status codes and details
|
||||
raise
|
||||
|
|
@ -2729,7 +2809,12 @@ if MCP_AVAILABLE:
|
|||
await initialize_session_managers()
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
await sse_session_manager.handle_request(scope, receive, send)
|
||||
async with _gateway_initialize_instructions_request_scope(
|
||||
user_api_key_auth,
|
||||
mcp_servers,
|
||||
_sse_client_ip,
|
||||
):
|
||||
await sse_session_manager.handle_request(scope, receive, send)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error handling MCP request: {e}")
|
||||
# Instead of re-raising, try to send a graceful error response
|
||||
|
|
|
|||
|
|
@ -1137,6 +1137,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
|
|||
tool_name_to_description: Optional[Dict[str, str]] = None
|
||||
extra_headers: Optional[List[str]] = None
|
||||
static_headers: Optional[Dict[str, str]] = None
|
||||
instructions: Optional[str] = None
|
||||
# Stdio-specific fields
|
||||
command: Optional[str] = None
|
||||
args: List[str] = Field(default_factory=list)
|
||||
|
|
@ -1219,6 +1220,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
|
|||
tool_name_to_description: Optional[Dict[str, str]] = None
|
||||
extra_headers: Optional[List[str]] = None
|
||||
static_headers: Optional[Dict[str, str]] = None
|
||||
instructions: Optional[str] = None
|
||||
# Stdio-specific fields
|
||||
command: Optional[str] = None
|
||||
args: List[str] = Field(default_factory=list)
|
||||
|
|
@ -1270,6 +1272,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
|
|||
transport: MCPTransportType
|
||||
auth_type: Optional[MCPAuthType] = None
|
||||
credentials: Optional[MCPCredentials] = None
|
||||
instructions: Optional[str] = None
|
||||
created_at: Optional[datetime] = None
|
||||
created_by: Optional[str] = None
|
||||
updated_at: Optional[datetime] = None
|
||||
|
|
|
|||
|
|
@ -289,6 +289,7 @@ model LiteLLM_MCPServerTable {
|
|||
server_name String?
|
||||
alias String?
|
||||
description String?
|
||||
instructions String?
|
||||
url String?
|
||||
spec_path String?
|
||||
transport String @default("sse")
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ class MCPServer(BaseModel):
|
|||
spec_path: Optional[str] = None
|
||||
auth_type: Optional[MCPAuthType] = None
|
||||
authentication_token: Optional[str] = None
|
||||
instructions: Optional[str] = None
|
||||
mcp_info: Optional[MCPInfo] = None
|
||||
extra_headers: Optional[
|
||||
List[str]
|
||||
|
|
|
|||
|
|
@ -289,6 +289,7 @@ model LiteLLM_MCPServerTable {
|
|||
server_name String?
|
||||
alias String?
|
||||
description String?
|
||||
instructions String?
|
||||
url String?
|
||||
spec_path String?
|
||||
transport String @default("sse")
|
||||
|
|
|
|||
|
|
@ -312,5 +312,80 @@ class TestMCPClient:
|
|||
assert MCPAuth.token.value == "token"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _last_initialize_instructions capture
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMCPClientInstructionsCapture:
|
||||
"""Tests for _last_initialize_instructions capture during session init."""
|
||||
|
||||
def test_initial_value_is_none(self):
|
||||
"""Fresh client has no cached instructions."""
|
||||
client = MCPClient(
|
||||
server_url="http://example.com/mcp",
|
||||
transport_type="http",
|
||||
)
|
||||
assert client._last_initialize_instructions is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("litellm.experimental_mcp_client.client.ClientSession")
|
||||
async def test_captures_instructions_from_initialize(self, mock_session_cls):
|
||||
"""Instructions from upstream initialize() are captured and stripped."""
|
||||
client = MCPClient(
|
||||
server_url="http://example.com/mcp",
|
||||
transport_type="http",
|
||||
)
|
||||
|
||||
mock_session = AsyncMock()
|
||||
init_result = MagicMock()
|
||||
init_result.instructions = " upstream says hello "
|
||||
mock_session.initialize = AsyncMock(return_value=init_result)
|
||||
|
||||
session_ctx = MagicMock()
|
||||
session_ctx.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
session_ctx.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_session_cls.return_value = session_ctx
|
||||
|
||||
transport_ctx = MagicMock()
|
||||
transport_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock()))
|
||||
transport_ctx.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
async def _op(session):
|
||||
return "done"
|
||||
|
||||
await client._execute_session_operation(transport_ctx, _op)
|
||||
assert client._last_initialize_instructions == "upstream says hello"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("litellm.experimental_mcp_client.client.ClientSession")
|
||||
async def test_none_instructions_stays_none(self, mock_session_cls):
|
||||
"""When upstream returns no instructions the field stays None."""
|
||||
client = MCPClient(
|
||||
server_url="http://example.com/mcp",
|
||||
transport_type="http",
|
||||
)
|
||||
|
||||
mock_session = AsyncMock()
|
||||
init_result = MagicMock()
|
||||
init_result.instructions = None
|
||||
mock_session.initialize = AsyncMock(return_value=init_result)
|
||||
|
||||
session_ctx = MagicMock()
|
||||
session_ctx.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
session_ctx.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_session_cls.return_value = session_ctx
|
||||
|
||||
transport_ctx = MagicMock()
|
||||
transport_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock()))
|
||||
transport_ctx.__aexit__ = AsyncMock(return_value=False)
|
||||
|
||||
async def _op(session):
|
||||
return "done"
|
||||
|
||||
await client._execute_session_operation(transport_ctx, _op)
|
||||
assert client._last_initialize_instructions is None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
|
|
|
|||
|
|
@ -2421,3 +2421,178 @@ async def test_get_tools_from_mcp_servers_injects_stored_oauth2_token():
|
|||
assert call_kwargs["extra_headers"] == {"Authorization": f"Bearer {STORED_TOKEN}"}
|
||||
|
||||
assert tools == [tool_1]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _merge_gateway_initialize_instructions + ContextVar / InitializationOptions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_instruction_server(
|
||||
server_id="s1",
|
||||
name="s1",
|
||||
*,
|
||||
alias=None,
|
||||
server_name=None,
|
||||
instructions=None,
|
||||
spec_path=None,
|
||||
url="https://example.com",
|
||||
):
|
||||
return MCPServer(
|
||||
server_id=server_id,
|
||||
name=name,
|
||||
alias=alias,
|
||||
server_name=server_name,
|
||||
url=url,
|
||||
transport=MCPTransport.http,
|
||||
instructions=instructions,
|
||||
spec_path=spec_path,
|
||||
)
|
||||
|
||||
|
||||
class TestMergeGatewayInitializeInstructions:
|
||||
"""Tests for _merge_gateway_initialize_instructions."""
|
||||
|
||||
def _merge(self, servers):
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
_merge_gateway_initialize_instructions,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
return _merge_gateway_initialize_instructions(servers)
|
||||
|
||||
def test_empty_server_list_returns_none(self):
|
||||
"""No servers yields no instructions."""
|
||||
assert self._merge([]) is None
|
||||
|
||||
def test_single_server_yaml_instructions(self):
|
||||
"""A single server with YAML instructions returns them verbatim."""
|
||||
s = _make_instruction_server(instructions="Use add() for sums.")
|
||||
assert self._merge([s]) == "Use add() for sums."
|
||||
|
||||
def test_yaml_instructions_strips_whitespace(self):
|
||||
"""Leading/trailing whitespace is stripped."""
|
||||
s = _make_instruction_server(instructions=" padded \n")
|
||||
assert self._merge([s]) == "padded"
|
||||
|
||||
def test_yaml_override_beats_upstream_cache(self):
|
||||
"""YAML/DB instructions take precedence over upstream cache."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
global_mcp_server_manager._upstream_initialize_instructions_by_server_id["s1"] = "upstream"
|
||||
try:
|
||||
s = _make_instruction_server(instructions="yaml wins")
|
||||
assert self._merge([s]) == "yaml wins"
|
||||
finally:
|
||||
global_mcp_server_manager._upstream_initialize_instructions_by_server_id.pop("s1", None)
|
||||
|
||||
def test_upstream_cache_used_when_no_yaml(self):
|
||||
"""Upstream cached instructions are used when no YAML override is set."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
global_mcp_server_manager._upstream_initialize_instructions_by_server_id["s1"] = "from upstream"
|
||||
try:
|
||||
s = _make_instruction_server(instructions=None)
|
||||
assert self._merge([s]) == "from upstream"
|
||||
finally:
|
||||
global_mcp_server_manager._upstream_initialize_instructions_by_server_id.pop("s1", None)
|
||||
|
||||
def test_spec_path_servers_skipped(self):
|
||||
"""OpenAPI (spec_path) servers do not contribute instructions."""
|
||||
s = _make_instruction_server(spec_path="/openapi.json", url=None)
|
||||
assert self._merge([s]) is None
|
||||
|
||||
def test_no_instructions_no_cache_returns_none(self):
|
||||
"""Server with no instructions and no cache yields None."""
|
||||
s = _make_instruction_server()
|
||||
assert self._merge([s]) is None
|
||||
|
||||
def test_multiple_servers_merged_with_labels(self):
|
||||
"""Multiple servers get label-prefixed and separator-joined."""
|
||||
s1 = _make_instruction_server(server_id="a", name="a", alias="Alpha", instructions="instr A")
|
||||
s2 = _make_instruction_server(server_id="b", name="b", alias="Beta", instructions="instr B")
|
||||
result = self._merge([s1, s2])
|
||||
assert result is not None
|
||||
assert "[Alpha]" in result and "[Beta]" in result
|
||||
assert "instr A" in result and "instr B" in result
|
||||
assert "---" in result
|
||||
|
||||
def test_single_server_no_label_wrapping(self):
|
||||
"""A single server's instructions are not wrapped with a label."""
|
||||
s = _make_instruction_server(alias="MyServer", instructions="single")
|
||||
result = self._merge([s])
|
||||
assert result == "single"
|
||||
assert "[MyServer]" not in result
|
||||
|
||||
def test_mixed_yaml_cache_specpath(self):
|
||||
"""YAML, upstream-cache, and spec_path servers are handled correctly together."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
global_mcp_server_manager._upstream_initialize_instructions_by_server_id["c"] = "cached C"
|
||||
try:
|
||||
s_yaml = _make_instruction_server(server_id="a", name="a", alias="A", instructions="yaml A")
|
||||
s_spec = _make_instruction_server(server_id="b", name="b", alias="B", spec_path="/spec.json", url=None)
|
||||
s_cached = _make_instruction_server(server_id="c", name="c", alias="C")
|
||||
result = self._merge([s_yaml, s_spec, s_cached])
|
||||
assert "yaml A" in result
|
||||
assert "cached C" in result
|
||||
assert "[B]" not in result
|
||||
finally:
|
||||
global_mcp_server_manager._upstream_initialize_instructions_by_server_id.pop("c", None)
|
||||
|
||||
|
||||
class TestGatewayCreateInitializationOptions:
|
||||
"""Tests for the patched server.create_initialization_options via ContextVar."""
|
||||
|
||||
def test_no_contextvar_returns_default_options(self):
|
||||
"""When ContextVar is None, instructions are absent."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.mcp_context import (
|
||||
_mcp_gateway_initialize_instructions,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.server import server
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
tok = _mcp_gateway_initialize_instructions.set(None)
|
||||
try:
|
||||
opts = server.create_initialization_options()
|
||||
assert getattr(opts, "instructions", None) is None
|
||||
finally:
|
||||
_mcp_gateway_initialize_instructions.reset(tok)
|
||||
|
||||
def test_contextvar_set_injects_instructions(self):
|
||||
"""When ContextVar has a value, it appears in InitializationOptions."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.mcp_context import (
|
||||
_mcp_gateway_initialize_instructions,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.server import server
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
tok = _mcp_gateway_initialize_instructions.set("hello from merge")
|
||||
try:
|
||||
opts = server.create_initialization_options()
|
||||
assert opts.instructions == "hello from merge"
|
||||
finally:
|
||||
_mcp_gateway_initialize_instructions.reset(tok)
|
||||
|
||||
def test_contextvar_reset_removes_instructions(self):
|
||||
"""After resetting the ContextVar, instructions disappear."""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.mcp_context import (
|
||||
_mcp_gateway_initialize_instructions,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.server import server
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
tok = _mcp_gateway_initialize_instructions.set("temporary")
|
||||
_mcp_gateway_initialize_instructions.reset(tok)
|
||||
opts = server.create_initialization_options()
|
||||
assert getattr(opts, "instructions", None) is None
|
||||
|
|
|
|||
|
|
@ -2483,5 +2483,77 @@ class TestHasClientCredentialsOAuth2Flow:
|
|||
assert server.needs_user_oauth_token is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Upstream initialize-instructions cache
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMCPServerManagerUpstreamInstructionsCache:
|
||||
"""Tests for the upstream initialize-instructions cache."""
|
||||
|
||||
def test_get_returns_none_when_empty(self):
|
||||
"""Empty cache returns None for any key."""
|
||||
manager = MCPServerManager()
|
||||
assert manager._upstream_initialize_instructions_by_server_id.get("nonexistent") is None
|
||||
|
||||
def test_remember_stores_stripped_value(self):
|
||||
"""_remember_upstream_initialize_instructions stores a stripped string."""
|
||||
manager = MCPServerManager()
|
||||
fake_server = MagicMock(server_id="srv")
|
||||
fake_client = MagicMock(_last_initialize_instructions=" hello \n")
|
||||
manager._remember_upstream_initialize_instructions(fake_server, fake_client)
|
||||
assert manager._upstream_initialize_instructions_by_server_id.get("srv") == "hello"
|
||||
|
||||
def test_remember_ignores_empty_string(self):
|
||||
"""Whitespace-only instructions are not stored."""
|
||||
manager = MCPServerManager()
|
||||
fake_server = MagicMock(server_id="srv")
|
||||
fake_client = MagicMock(_last_initialize_instructions=" ")
|
||||
manager._remember_upstream_initialize_instructions(fake_server, fake_client)
|
||||
assert manager._upstream_initialize_instructions_by_server_id.get("srv") is None
|
||||
|
||||
def test_remember_ignores_none(self):
|
||||
"""None instructions are not stored."""
|
||||
manager = MCPServerManager()
|
||||
fake_server = MagicMock(server_id="srv")
|
||||
fake_client = MagicMock(_last_initialize_instructions=None)
|
||||
manager._remember_upstream_initialize_instructions(fake_server, fake_client)
|
||||
assert manager._upstream_initialize_instructions_by_server_id.get("srv") is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_servers_from_config_clears_cache(self):
|
||||
"""Reloading config clears any previously cached upstream instructions."""
|
||||
manager = MCPServerManager()
|
||||
manager._upstream_initialize_instructions_by_server_id["old"] = "stale"
|
||||
await manager.load_servers_from_config(
|
||||
mcp_servers_config={
|
||||
"fresh_srv": {
|
||||
"url": "https://example.com",
|
||||
"instructions": "from yaml",
|
||||
}
|
||||
}
|
||||
)
|
||||
assert manager._upstream_initialize_instructions_by_server_id.get("old") is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_servers_reads_instructions_from_config(self):
|
||||
"""instructions field from YAML config is persisted on the MCPServer."""
|
||||
manager = MCPServerManager()
|
||||
await manager.load_servers_from_config(
|
||||
mcp_servers_config={
|
||||
"srv_a": {
|
||||
"url": "https://a.example.com",
|
||||
"instructions": "A instructions",
|
||||
},
|
||||
"srv_b": {
|
||||
"url": "https://b.example.com",
|
||||
},
|
||||
}
|
||||
)
|
||||
by_name = {s.server_name: s for s in manager.config_mcp_servers.values()}
|
||||
assert "srv_a" in by_name and by_name["srv_a"].instructions == "A instructions"
|
||||
assert "srv_b" in by_name and by_name["srv_b"].instructions is None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue