refactor: make tool mutator generic between cluster/context naming

Signed-off-by: Calum Murray <[email protected]>

Author: Cali0707

pkg/mcp/tool_mutator.go

@@ -10,9 +10,9 @@ import (
 
 type ToolMutator func(tool api.ServerTool) api.ServerTool
 
-const maxClustersInEnum = 15 // TODO: test and validate that this is a reasonable cutoff
+const maxTargetsInEnum = 15 // TODO: test and validate that this is a reasonable cutoff
 
-func WithClusterParameter(defaultCluster string, clusters, skipToolNames []string) ToolMutator {
+func WithTargetParameter(defaultCluster, targetParameterName string, targets, skipToolNames []string) ToolMutator {
 	skipNames := make(map[string]struct{}, len(skipToolNames))
 	for _, n := range skipToolNames {
 		skipNames[n] = struct{}{}
@@ -31,28 +31,34 @@ func WithClusterParameter(defaultCluster string, clusters, skipToolNames []strin
 			tool.Tool.InputSchema.Properties = make(map[string]*jsonschema.Schema)
 		}
 
-		if len(clusters) > 1 {
-			tool.Tool.InputSchema.Properties["cluster"] = createClusterProperty(defaultCluster, clusters)
+		if len(targets) > 1 {
+			tool.Tool.InputSchema.Properties[targetParameterName] = createTargetProperty(
+				defaultCluster,
+				targetParameterName,
+				targets,
+			)
 		}
 
 		return tool
 	}
 }
 
-func createClusterProperty(defaultCluster string, clusters []string) *jsonschema.Schema {
+func createTargetProperty(defaultCluster, targetName string, targets []string) *jsonschema.Schema {
 	baseSchema := &jsonschema.Schema{
 		Type: "string",
 		Description: fmt.Sprintf(
-			"Optional parameter selecting which cluster to run the tool in. Defaults to %s if not set", defaultCluster,
+			"Optional parameter selecting which %s to run the tool in. Defaults to %s if not set",
+			targetName,
+			defaultCluster,
 		),
 	}
 
-	if len(clusters) <= maxClustersInEnum {
+	if len(targets) <= maxTargetsInEnum {
 		// Sort clusters to ensure consistent enum ordering
-		sort.Strings(clusters)
+		sort.Strings(targets)
 
-		enumValues := make([]any, 0, len(clusters))
-		for _, c := range clusters {
+		enumValues := make([]any, 0, len(targets))
+		for _, c := range targets {
 			enumValues = append(enumValues, c)
 		}
 		baseSchema.Enum = enumValues

pkg/mcp/tool_mutator_test.go

@@ -66,15 +66,16 @@ func createTestToolWithExistingProperties(name string) api.ServerTool {
 
 func TestWithClusterParameter(t *testing.T) {
 	tests := []struct {
-		name           string
-		defaultCluster string
-		clusters       []string
-		skipToolNames  []string
-		toolName       string
-		toolFactory    func(string) api.ServerTool
-		expectCluster  bool
-		expectEnum     bool
-		enumCount      int
+		name                string
+		defaultCluster      string
+		targetParameterName string
+		clusters            []string
+		skipToolNames       []string
+		toolName            string
+		toolFactory         func(string) api.ServerTool
+		expectCluster       bool
+		expectEnum          bool
+		enumCount           int
 	}{
 		{
 			name:           "adds cluster parameter when multiple clusters provided",
@@ -146,7 +147,10 @@ func TestWithClusterParameter(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			mutator := WithClusterParameter(tt.defaultCluster, tt.clusters, tt.skipToolNames)
+			if tt.targetParameterName == "" {
+				tt.targetParameterName = "cluster"
+			}
+			mutator := WithTargetParameter(tt.defaultCluster, tt.targetParameterName, tt.clusters, tt.skipToolNames)
 			tool := tt.toolFactory(tt.toolName)
 			originalTool := tool // Keep reference to check if tool was unchanged
 
@@ -192,41 +196,47 @@ func TestCreateClusterProperty(t *testing.T) {
 	tests := []struct {
 		name           string
 		defaultCluster string
+		targetName     string
 		clusters       []string
 		expectEnum     bool
 		expectedCount  int
 	}{
 		{
 			name:           "creates property with enum when clusters <= maxClustersInEnum",
 			defaultCluster: "default",
+			targetName:     "cluster",
 			clusters:       []string{"cluster1", "cluster2", "cluster3"},
 			expectEnum:     true,
 			expectedCount:  3,
 		},
 		{
 			name:           "creates property without enum when clusters > maxClustersInEnum",
 			defaultCluster: "default",
-			clusters:       make([]string, maxClustersInEnum+5), // 20 clusters
+			targetName:     "cluster",
+			clusters:       make([]string, maxTargetsInEnum+5), // 20 clusters
 			expectEnum:     false,
 			expectedCount:  0,
 		},
 		{
 			name:           "creates property with exact maxClustersInEnum clusters",
 			defaultCluster: "default",
-			clusters:       make([]string, maxClustersInEnum),
+			targetName:     "cluster",
+			clusters:       make([]string, maxTargetsInEnum),
 			expectEnum:     true,
-			expectedCount:  maxClustersInEnum,
+			expectedCount:  maxTargetsInEnum,
 		},
 		{
 			name:           "handles single cluster",
 			defaultCluster: "default",
+			targetName:     "cluster",
 			clusters:       []string{"single-cluster"},
 			expectEnum:     true,
 			expectedCount:  1,
 		},
 		{
 			name:           "handles empty clusters list",
 			defaultCluster: "default",
+			targetName:     "cluster",
 			clusters:       []string{},
 			expectEnum:     true,
 			expectedCount:  0,
@@ -242,7 +252,7 @@ func TestCreateClusterProperty(t *testing.T) {
 				}
 			}
 
-			property := createClusterProperty(tt.defaultCluster, tt.clusters)
+			property := createTargetProperty(tt.defaultCluster, tt.targetName, tt.clusters)
 
 			assert.Equal(t, "string", property.Type)
 			assert.Contains(t, property.Description, tt.defaultCluster)
@@ -279,6 +289,6 @@ func TestToolMutatorType(t *testing.T) {
 
 func TestMaxClustersInEnumConstant(t *testing.T) {
 	t.Run("maxClustersInEnum has expected value", func(t *testing.T) {
-		assert.Equal(t, 15, maxClustersInEnum, "maxClustersInEnum should be 15")
+		assert.Equal(t, 15, maxTargetsInEnum, "maxClustersInEnum should be 15")
 	})
 }