scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-04 07:27:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 195

Merge b8cd3a7142 into 0657ed46cf

Browse source
This commit is contained in:
shrey 2026-10-27 10:38:40 -07:00 • committed by GitHub
commit b9dd48d902
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 136 additions and 47 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm/experimental_mcp_client/client.py
View file

@ -792,7 +792,7 @@ class MCPClient:
init_result: Final = await self._prepare_session(session)
instructions: Final = getattr(init_result, "instructions", None)
self._last_initialize_instructions = (
instructions.strip() or None if isinstance(instructions, str) else None
instructions.strip() if isinstance(instructions, str) else ""
)
result: Final = await operation(session)
except BaseException as operation_error:

2
litellm/proxy/_experimental/mcp_server/catalog.py
View file

@ -525,8 +525,6 @@ class CatalogSnapshots:
self.manager.invalidate_oauth_discovery_state(server_id)
self._database_identity = database_identity
self.manager.registry = registered_registry
if not reuse_unchanged:
self.manager.clear_initialize_instructions()
warn_on_shared_identifier_prefixes(registered_registry.values())
# A discovery task may have published into ``previous_registry`` while
# this replacement was being staged. Reconcile every published entry

37
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -2076,6 +2076,7 @@ class MCPServerManager:
self._listed_tools_by_server_id: dict[str, _ListedToolsByCaller] = {} # mutable-ok: refreshed per tools/list
self._listed_tools_generations: dict[str, int] = {} # mutable-ok: bumped per server save
self._upstream_initialize_instructions_by_server_id: dict[str, str] = {}
self._initialize_instructions_generations: Mapping[str, int] = MappingProxyType({})
# Per-server monotonic timestamp of last upstream prefetch attempt (success,
# empty result, or failure). Used to throttle re-probes for servers that do
# not return instructions, and to apply a short cooldown after failures.
@ -2448,10 +2449,18 @@ class MCPServerManager:
return current
raise HTTPException(status_code=503, detail="OAuth metadata discovery changed repeatedly; retry shortly")
def _remember_upstream_initialize_instructions(self, server: MCPServer, client: MCPClient) -> None:
def initialize_instructions_generation(self, server_id: str) -> int:
return self._initialize_instructions_generations.get(server_id, 0)
def _remember_upstream_initialize_instructions(self, server: MCPServer, client: MCPClient, generation: int) -> None:
raw: Final[str | None] = 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()
if raw is None or generation != self.initialize_instructions_generation(server.server_id):
return
instructions: Final = raw.strip()
if instructions:
self._upstream_initialize_instructions_by_server_id[server.server_id] = instructions
else:
self._upstream_initialize_instructions_by_server_id.pop(server.server_id, None)
async def ensure_upstream_initialize_instructions_cached(self, server: MCPServer) -> None:
"""
@ -2495,6 +2504,7 @@ class MCPServerManager:
# Record the attempt up-front so that a failure / empty response does not
# cause every subsequent initialize request to re-open the upstream session.
self._upstream_initialize_instructions_probed_at[server.server_id] = time.monotonic()
generation: Final = self.initialize_instructions_generation(server.server_id)
try:
resolved_static_headers: Final = await self._resolve_static_headers_with_env_vars(
@ -2514,7 +2524,7 @@ class MCPServerManager:
return "ok"
await asyncio.wait_for(client.run_with_session(_noop), timeout=MCP_HEALTH_CHECK_TIMEOUT)
self._remember_upstream_initialize_instructions(server, client)
self._remember_upstream_initialize_instructions(server, client, generation)
except Exception as e:
verbose_logger.debug(
"Upstream initialize instructions prefetch failed for %s: %s",
@ -4475,6 +4485,7 @@ class MCPServerManager:
oauth2_headers=oauth2_headers,
)
listed_generation: Final = self._listed_tools_generations.get(server.server_id, 0)
instructions_generation: Final = self.initialize_instructions_generation(server.server_id)
try:
# Tool *listing* must not be blocked by missing per-user env vars —
@ -4592,7 +4603,7 @@ class MCPServerManager:
else ListToolsResult(tools=await self._fetch_tools_with_timeout(client, server.name))
)
tools = page.tools
self._remember_upstream_initialize_instructions(server, client)
self._remember_upstream_initialize_instructions(server, client, instructions_generation)
guarded_tools: Final = await self._guard_tool_catalog(
server=server,
@ -4650,10 +4661,6 @@ class MCPServerManager:
)
raise_classified_list_failure(e, server.name, suppress_challenge=server.is_dcr_bridge)
def clear_initialize_instructions(self) -> None:
self._upstream_initialize_instructions_by_server_id.clear()
self._upstream_initialize_instructions_probed_at.clear()
def invalidate_discovery_lists(self, server_id: str) -> None:
self._upstream_initialize_instructions_by_server_id.pop(server_id, None)
self._upstream_initialize_instructions_probed_at.pop(server_id, None)
@ -4667,6 +4674,12 @@ class MCPServerManager:
)
self.invalidate_discovery_lists(server_id)
self._initialize_instructions_generations = MappingProxyType(
{
**self._initialize_instructions_generations,
server_id: self.initialize_instructions_generation(server_id) + 1,
}
)
self._drop_listed_tools(server_id)
invalidate_oauth_metadata_cache(server_id)
@ -6532,6 +6545,7 @@ class MCPServerManager:
GuardrailRaisedException: If guardrails block the call
HTTPException: If an HTTP error occurs
"""
instructions_generation: Final = self.initialize_instructions_generation(mcp_server.server_id)
server_auth_header: Final = _server_auth_header_for(mcp_server, mcp_server_auth_headers, mcp_auth_header)
# Extract subject token for OAuth2 Token Exchange (OBO) and ID-JAG flows
@ -6761,7 +6775,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: Final = 1 if proxy_logging_obj else 0
result: Final = mcp_responses[result_index]
self._remember_upstream_initialize_instructions(mcp_server, client)
self._remember_upstream_initialize_instructions(mcp_server, client, instructions_generation)
return (
bind_target(result, mcp_server)
@ -7513,6 +7527,7 @@ class MCPServerManager:
"""
from datetime import datetime
instructions_generation: Final = self.initialize_instructions_generation(server_id)
server: Final = self.get_mcp_server_by_id(server_id)
if not server:
verbose_logger.warning("MCP Server %s not found", server_id)
@ -7574,7 +7589,7 @@ class MCPServerManager:
client.run_with_session(_noop, quiet_on_error=True),
timeout=MCP_HEALTH_CHECK_TIMEOUT,
)
self._remember_upstream_initialize_instructions(server, client)
self._remember_upstream_initialize_instructions(server, client, instructions_generation)
status = "healthy"
except asyncio.TimeoutError:
health_check_error = f"Health check timed out after {MCP_HEALTH_CHECK_TIMEOUT} seconds"

5
tests/unit/experimental_mcp_client/test_mcp_client.py
View file

@ -699,8 +699,7 @@ class TestMCPClientInstructionsCapture:
@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."""
async def test_missing_instructions_are_recorded_as_empty(self, mock_session_cls: MagicMock) -> None:
client = MCPClient(
server_url="http://example.com/mcp",
transport_type="http",
@ -723,7 +722,7 @@ class TestMCPClientInstructionsCapture:
return "done"
await client._execute_session_operation(transport_ctx, _op)
assert client._last_initialize_instructions is None
assert client._last_initialize_instructions == ""
# ---------------------------------------------------------------------------

135
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -10461,29 +10461,82 @@ class TestMCPServerManagerUpstreamInstructionsCache:
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"
@pytest.mark.parametrize(
("answered", "cached"),
[(" hello \n", "hello"), ("", None), (" ", None), (None, "previous text")],
ids=["replaces_with_stripped_text", "drops_when_upstream_sent_none", "drops_blank", "keeps_without_initialize"],
)
def test_remember_applies_the_last_completed_initialize(self, answered: str | None, cached: str | None) -> None:
manager: Final = MCPServerManager()
manager._upstream_initialize_instructions_by_server_id["srv"] = "previous text"
manager._remember_upstream_initialize_instructions(
MagicMock(server_id="srv"),
MagicMock(_last_initialize_instructions=answered),
manager.initialize_instructions_generation("srv"),
)
assert manager._upstream_initialize_instructions_by_server_id.get("srv") == cached
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
@pytest.mark.asyncio
@pytest.mark.parametrize("writer", ["prefetch", "health_check", "tools_list", "tool_call"])
@pytest.mark.parametrize(
("mid_flight", "cached"),
[("nothing", "upstream text"), ("discovery_refresh", "upstream text"), ("server_edit", None)],
ids=["unchanged", "discovery_refresh_mid_flight", "edited_mid_flight"],
)
async def test_writer_skips_instructions_from_a_session_that_outlived_a_server_edit(
self, writer: str, mid_flight: str, cached: str | None
) -> None:
manager: Final = MCPServerManager()
server: Final = MCPServer(
server_id="srv",
name="srv",
url="https://srv.example/mcp",
transport=MCPTransport.http,
auth_type=MCPAuth.none,
)
manager.registry = {server.server_id: server}
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
def upstream(result: object) -> AsyncMock:
async def respond(*args: object, **kwargs: object) -> object:
match mid_flight:
case "server_edit":
manager.invalidate_server_definition_caches(server.server_id)
case "discovery_refresh":
manager.invalidate_discovery_lists(server.server_id)
return result
return AsyncMock(side_effect=respond)
manager.create_mcp_client = AsyncMock(
return_value=SimpleNamespace(
_last_initialize_instructions="upstream text",
run_with_session=upstream("ok"),
list_tools_page=upstream(ListToolsResult(tools=[])),
call_tool=upstream(CallToolResult(content=[], isError=False)),
)
)
match writer:
case "prefetch":
await manager.ensure_upstream_initialize_instructions_cached(server)
case "health_check":
await manager.health_check_server(server.server_id)
case "tools_list":
await manager.get_tools_page(server, params=PaginatedRequestParams())
case _:
await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="tool",
arguments={},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers=None,
raw_headers=None,
proxy_logging_obj=None,
)
assert manager._upstream_initialize_instructions_by_server_id.get(server.server_id) == cached
@pytest.mark.asyncio
async def test_load_servers_from_config_clears_cache(self, config_only_mcp_manager_factory):
@ -18730,16 +18783,40 @@ async def test_catalog_failed_openapi_row_does_not_publish_partial_handlers(monk
@pytest.mark.asyncio
async def test_catalog_unchanged_read_preserves_derived_initialize_instructions(monkeypatch):
read_rows = AsyncMock(return_value=[_catalog_row()])
@pytest.mark.parametrize("full_reload", [False, True], ids=["request_path_refresh", "scheduled_reload"])
async def test_catalog_refresh_drops_initialize_instructions_only_for_changed_servers(
monkeypatch: pytest.MonkeyPatch, full_reload: bool
) -> None:
def row(server_id: str, updated_day: int) -> LiteLLM_MCPServerTable:
return LiteLLM_MCPServerTable(
server_id=server_id,
server_name=server_id,
transport=MCPTransport.http,
url=f"https://{server_id}.example/mcp",
created_at=datetime(2026, 1, 1),
updated_at=datetime(2026, 1, updated_day),
)
read_rows: Final = AsyncMock(return_value=[row("unchanged", 1), row("edited", 1), row("deleted", 1)])
_catalog_database(monkeypatch, read_rows)
manager = MCPServerManager()
manager: Final = MCPServerManager()
await manager.catalog.list()
manager._upstream_initialize_instructions_by_server_id = {"catalog-server": "cached instructions"}
manager._upstream_initialize_instructions_probed_at = {"catalog-server": 123.0}
await manager.catalog.list()
assert manager._upstream_initialize_instructions_by_server_id == {"catalog-server": "cached instructions"}
assert manager._upstream_initialize_instructions_probed_at == {"catalog-server": 123.0}
manager._upstream_initialize_instructions_by_server_id = {
"unchanged": "unchanged text",
"edited": "stale edited text",
"deleted": "stale deleted text",
"not-in-db": "config server text",
}
manager._upstream_initialize_instructions_probed_at = {"unchanged": 1.0, "edited": 2.0, "deleted": 3.0}
read_rows.return_value = [row("unchanged", 1), row("edited", 2)]
await (manager.reload_servers_from_database() if full_reload else manager.catalog.list())
assert manager._upstream_initialize_instructions_by_server_id == {
"unchanged": "unchanged text",
"not-in-db": "config server text",
}
assert manager._upstream_initialize_instructions_probed_at == {"unchanged": 1.0}
@pytest.mark.asyncio

2
tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py
View file

@ -6572,7 +6572,7 @@ class TestEnsureUpstreamInitializeInstructionsCached:
server = _make_instruction_server(server_id="empty-server", instructions=None)
fake_client = MagicMock()
fake_client.run_with_session = AsyncMock(return_value="ok")
fake_client._last_initialize_instructions = None # upstream sent nothing
fake_client._last_initialize_instructions = "" # upstream sent nothing
create = AsyncMock(return_value=fake_client)
with patch.object(global_mcp_server_manager, "create_mcp_client", create):