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 3a9ba1388b7..eb7878337a1 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 @@ -1852,6 +1852,193 @@ class TestMCPServerManager: ) mock_inject.assert_awaited_once() + def test_resolve_mcp_server_for_tool_call_via_prefixed_name(self): + """Resolution succeeds when the prefixed tool name is in the mapping.""" + manager = MCPServerManager() + server = MCPServer( + server_id="jira", + name="jira", + transport=MCPTransport.http, + ) + manager.registry = {"jira": server} + manager.tool_name_to_mcp_server_name_mapping["jira-search_issues"] = "jira" + manager.tool_name_to_mcp_server_name_mapping["search_issues"] = "jira" + + resolved = manager._resolve_mcp_server_for_tool_call("jira", "search_issues") + assert resolved is server + + def test_resolve_mcp_server_for_tool_call_via_alias(self): + """Resolution falls back to alias/server_name match in the registry.""" + manager = MCPServerManager() + server = MCPServer( + server_id="srv-uuid-123", + name="zapier", + alias="zapier-alias", + transport=MCPTransport.http, + ) + manager.registry = {"srv-uuid-123": server} + + resolved = manager._resolve_mcp_server_for_tool_call( + "zapier-alias", "create_zap" + ) + assert resolved is server + + def test_resolve_mcp_server_for_tool_call_fallback_to_unprefixed_lookup(self): + """Fallback to unprefixed _get_mcp_server_from_tool_name when other paths fail.""" + manager = MCPServerManager() + server = MCPServer( + server_id="linear", + name="linear", + transport=MCPTransport.http, + ) + manager.registry = {"linear": server} + manager.tool_name_to_mcp_server_name_mapping["create_issue"] = "linear" + + # server_name is empty so the fallback unprefixed lookup runs and matches. + resolved = manager._resolve_mcp_server_for_tool_call("", "create_issue") + assert resolved is server + + def test_resolve_mcp_server_for_tool_call_raises_when_not_found(self): + """ValueError is raised when no resolution path finds the tool.""" + manager = MCPServerManager() + with pytest.raises(ValueError, match="Tool .* not found"): + manager._resolve_mcp_server_for_tool_call("nonexistent", "ghost_tool") + + def test_resolve_mcp_server_for_tool_call_unknown_tool_with_known_server(self): + """Server-name match alone must not let unknown tools slip through. + + If the registry has tools for this server but neither the prefixed nor + unprefixed tool name is in the mapping, raise rather than returning the + server (would otherwise allow tool enumeration via name spoofing). + """ + manager = MCPServerManager() + server = MCPServer( + server_id="github", + name="github", + transport=MCPTransport.http, + ) + manager.registry = {"github": server} + # Mapping has *some* tools for github but not "missing_tool". + manager.tool_name_to_mcp_server_name_mapping["github-list_repos"] = "github" + manager.tool_name_to_mcp_server_name_mapping["list_repos"] = "github" + + with pytest.raises(ValueError, match="Tool missing_tool not found"): + manager._resolve_mcp_server_for_tool_call("github", "missing_tool") + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_skipped_when_not_user_oauth(self): + """Returns input headers unchanged when server does not need user OAuth.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="plain", + name="plain", + transport=MCPTransport.http, + ) + # needs_user_oauth_token defaults to False. + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="bob") + + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + assert result is None + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_returns_client_supplied_token(self): + """Returns the client's oauth2_headers as-is when already set.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + assert server.needs_user_oauth_token is True + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + supplied = {"Authorization": "Bearer client-supplied"} + + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=supplied, user_api_key_auth=user_auth + ) + assert result is supplied + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_looks_up_stored_token(self): + """Falls back to stored per-user OAuth headers when no token is supplied.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + stored = {"Authorization": "Bearer stored-user-token"} + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new=AsyncMock(return_value=stored), + ) as mock_lookup: + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + + assert result == stored + mock_lookup.assert_awaited_once() + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_swallows_lookup_exception(self): + """Returns supplied headers (None) when the stored-token lookup raises.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new=AsyncMock(side_effect=RuntimeError("redis down")), + ): + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + assert result is None + + @pytest.mark.asyncio + async def test_resolve_oauth2_headers_no_user_id(self): + """Skip lookup entirely when user_api_key_auth has no user_id.""" + from litellm.proxy._types import UserAPIKeyAuth + + manager = MCPServerManager() + server = MCPServer( + server_id="oauth-srv", + name="oauth-srv", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + ) + # user_id is None -> lookup must not happen + user_auth = UserAPIKeyAuth(api_key="sk-test") + + with patch( + "litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db", + new=AsyncMock(return_value={"Authorization": "Bearer x"}), + ) as mock_lookup: + result = await manager._resolve_oauth2_headers_for_tool_call( + server, oauth2_headers=None, user_api_key_auth=user_auth + ) + assert result is None + mock_lookup.assert_not_called() + def test_create_prefixed_tools_updates_mapping_for_both_forms(self): """_create_prefixed_tools should populate mapping for prefixed and original names even when not adding prefix in output.""" manager = MCPServerManager() diff --git a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py index c30abab860f..cb2276ab39d 100644 --- a/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py +++ b/tests/test_litellm/proxy/guardrails/test_mcp_jwt_signer.py @@ -1156,3 +1156,116 @@ async def test_hook_raises_401_when_jwt_verification_fails(): ) assert exc_info.value.status_code == 401 + + +# --- _build_scope branches: call_mcp_tool with empty tool name, list_mcp_tools --- + + +def test_build_scope_call_type_call_mcp_tool_without_tool_name(): + """call_mcp_tool with empty tool name emits a generic mcp:tools/call only.""" + signer = _make_signer() + scope = signer._build_scope("", call_type="call_mcp_tool") + scopes = set(scope.split()) + assert scopes == {"mcp:tools/call"} + + +def test_build_scope_call_type_list_mcp_tools_only_list(): + """list_mcp_tools (no tool) emits only mcp:tools/list, never tools/call.""" + signer = _make_signer() + scope = signer._build_scope("", call_type="list_mcp_tools") + scopes = set(scope.split()) + assert scopes == {"mcp:tools/list"} + + +def test_build_scope_default_is_list_only_when_no_call_type(): + """No call_type and no tool falls through to tools/list (least-privilege default).""" + signer = _make_signer() + scope = signer._build_scope("") + scopes = set(scope.split()) + assert "mcp:tools/list" in scopes + assert "mcp:tools/call" not in scopes + + +# --- inject_mcp_jwt_headers_for_upstream --- + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_returns_unchanged_when_signer_not_configured(): + """No signer configured -> return a fresh copy of extra_headers untouched.""" + import litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer as mod + from litellm.proxy._types import UserAPIKeyAuth + + mod._mcp_jwt_signer_instance = None + headers = {"X-Trace-Id": "abc"} + user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + result = await mod.inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=user_dict, + extra_headers=headers, + ) + assert result == headers + assert result is not headers # must be a copy + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_returns_unchanged_when_user_dict_none(): + """No user_api_key_dict -> short-circuit without invoking the signer.""" + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + inject_mcp_jwt_headers_for_upstream, + ) + + _make_signer() # ensure instance is created + result = await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=None, + extra_headers={"X-Trace-Id": "abc"}, + ) + assert result == {"X-Trace-Id": "abc"} + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_signs_for_list_tools_path(): + """When for_list_tools=True, signer is invoked with list_mcp_tools call_type.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + inject_mcp_jwt_headers_for_upstream, + ) + + _make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300) + user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + result = await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=user_dict, + extra_headers={"X-Trace": "1"}, + raw_headers={"Authorization": "Bearer incoming.opaque.token"}, + for_list_tools=True, + ) + assert result["X-Trace"] == "1" + assert result["Authorization"].startswith("Bearer ") + token = result["Authorization"].removeprefix("Bearer ") + decoded = _decode_unverified(token) + scopes = set(decoded["scope"].split()) + assert scopes == {"mcp:tools/list"} + + +@pytest.mark.asyncio +async def test_inject_mcp_jwt_signs_for_tool_call_path(): + """for_list_tools=False with a tool name signs a call_mcp_tool JWT.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import ( + inject_mcp_jwt_headers_for_upstream, + ) + + _make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300) + user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + result = await inject_mcp_jwt_headers_for_upstream( + user_api_key_dict=user_dict, + for_list_tools=False, + mcp_tool_name="search_web", + ) + assert result["Authorization"].startswith("Bearer ") + token = result["Authorization"].removeprefix("Bearer ") + decoded = _decode_unverified(token) + scopes = set(decoded["scope"].split()) + assert "mcp:tools/call" in scopes + assert "mcp:tools/search_web:call" in scopes