scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-10 08:28:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix(mcp): keep cached initialize instructions accurate across catalog reloads

Browse source 
Every full catalog reload (the 30 second reload_mcp_servers_job, the reload after
an admin write, add_deployment) wiped every server's cached upstream initialize
instructions. Servers that need per-user auth cannot be prefetched, so gateway
initialize returned no instructions for them until a user listed tools again.

The reload already drops the entry of each server whose definition changed or was
removed, so the global wipe only ever cleared unchanged servers. Remove it and the
method that only it used.

That wipe was also the only thing that expired two kinds of stale text, so fix
both at the source. Each upstream session now records whether initialize finished
and what it said, and an answer without instructions drops the cached entry. A
per-server generation, bumped whenever a server's definition changes, is read
before every upstream session, and a session that outlived an edit of its server
no longer writes the old upstream's text back.

Resolves LIT-9851
This commit is contained in:
shrey kharbanda 2026-10-26 18:52:07 -07:00
parent c905cd065d
commit b8cd3a7142
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):