diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 864b17afe3..b7dbdeed5f 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -6,7 +6,6 @@ LiteLLM MCP Server Routes import asyncio import contextlib -import contextvars import time import types import traceback @@ -1161,18 +1160,23 @@ if MCP_AVAILABLE: return texts[0][1] return "\n\n---\n\n".join(f"[{lbl}]\n{txt}" for lbl, txt in texts) - async def _set_mcp_gateway_initialize_instructions_token( + @contextlib.asynccontextmanager + async def _gateway_initialize_instructions_request_scope( user_api_key_auth: Optional[UserAPIKeyAuth], mcp_servers: Optional[List[str]], client_ip: Optional[str], - ) -> contextvars.Token[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) - return _mcp_gateway_initialize_instructions.set(merged) + 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], @@ -2741,15 +2745,12 @@ if MCP_AVAILABLE: # Request was fully handled (e.g., DELETE on non-existent session) return - _instr_tok = await _set_mcp_gateway_initialize_instructions_token( + async with _gateway_initialize_instructions_request_scope( user_api_key_auth, mcp_servers, _client_ip, - ) - try: + ): await session_manager.handle_request(scope, receive, send) - finally: - _mcp_gateway_initialize_instructions.reset(_instr_tok) except HTTPException: # Re-raise HTTP exceptions to preserve status codes and details raise @@ -2808,15 +2809,12 @@ if MCP_AVAILABLE: await initialize_session_managers() await asyncio.sleep(0.1) - _sse_instr_tok = await _set_mcp_gateway_initialize_instructions_token( + async with _gateway_initialize_instructions_request_scope( user_api_key_auth, mcp_servers, _sse_client_ip, - ) - try: + ): await sse_session_manager.handle_request(scope, receive, send) - finally: - _mcp_gateway_initialize_instructions.reset(_sse_instr_tok) except Exception as e: verbose_logger.exception(f"Error handling MCP request: {e}") # Instead of re-raising, try to send a graceful error response diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 13a09f54e6..46d483c248 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -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__]) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 384d428888..acba06afee 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 656a9c616e..503ef71173 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -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.get_upstream_initialize_instructions("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.get_upstream_initialize_instructions("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.get_upstream_initialize_instructions("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.get_upstream_initialize_instructions("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.get_upstream_initialize_instructions("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__])