test: add unit tests for MCP initialize instructions feature

Extend existing test modules with coverage for the instructions merge logic, upstream cache, ContextVar-based injection, and client-side capture — following each file's established patterns. Made-with: Cursor
2026-07-18 12:17:48 +00:00 · 2026-04-14 15:54:27 +03:00
parent 8c2ebee4de
commit 7e656f4329
4 changed files with 334 additions and 14 deletions
@@ -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
@@ -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.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__])