From f9427ab04e7dbfd3c7509d111a877a8d0aeb4900 Mon Sep 17 00:00:00 2001
From: William Martin <williammartin@github.com>
Date: Fri, 25 Apr 2025 13:58:00 +0200
Subject: [PATCH] Ensure toolsets are configurable via env var

---
 cmd/github-mcp-server/main.go | 10 +++++++++-
 e2e/e2e_test.go               | 34 +++++++++++++++++++++++++++++++++-
 2 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/cmd/github-mcp-server/main.go b/cmd/github-mcp-server/main.go
index 5ca0e21c..cf459f47 100644
--- a/cmd/github-mcp-server/main.go
+++ b/cmd/github-mcp-server/main.go
@@ -45,7 +45,15 @@ var (
 				stdlog.Fatal("Failed to initialize logger:", err)
 			}
 
-			enabledToolsets := viper.GetStringSlice("toolsets")
+			// If you're wondering why we're not using viper.GetStringSlice("toolsets"),
+			// it's because viper doesn't handle comma-separated values correctly for env
+			// vars when using GetStringSlice.
+			// https://github.com/spf13/viper/issues/380
+			var enabledToolsets []string
+			err = viper.UnmarshalKey("toolsets", &enabledToolsets)
+			if err != nil {
+				stdlog.Fatal("Failed to unmarshal toolsets:", err)
+			}
 
 			logCommands := viper.GetBool("enable-command-logging")
 			cfg := runConfig{
diff --git a/e2e/e2e_test.go b/e2e/e2e_test.go
index 5d5c2032..a3f5df6f 100644
--- a/e2e/e2e_test.go
+++ b/e2e/e2e_test.go
@@ -7,6 +7,7 @@ import (
 	"encoding/json"
 	"os"
 	"os/exec"
+	"slices"
 	"sync"
 	"testing"
 	"time"
@@ -115,6 +116,9 @@ func setupMCPClient(t *testing.T, options ...ClientOption) *mcpClient.Client {
 	t.Log("Starting Stdio MCP client...")
 	client, err := mcpClient.NewStdioMCPClient(args[0], []string{}, args[1:]...)
 	require.NoError(t, err, "expected to create client successfully")
+	t.Cleanup(func() {
+		require.NoError(t, client.Close(), "expected to close client successfully")
+	})
 
 	// Initialize the client
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
@@ -166,5 +170,33 @@ func TestGetMe(t *testing.T) {
 	require.NoError(t, err, "expected to get user successfully")
 	require.Equal(t, trimmedContent.Login, *user.Login, "expected login to match")
 
-	require.NoError(t, mcpClient.Close(), "expected to close client successfully")
+}
+
+func TestToolsets(t *testing.T) {
+	mcpClient := setupMCPClient(
+		t,
+		WithEnvVars(map[string]string{
+			"GITHUB_TOOLSETS": "repos,issues",
+		}),
+	)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	request := mcp.ListToolsRequest{}
+	response, err := mcpClient.ListTools(ctx, request)
+	require.NoError(t, err, "expected to list tools successfully")
+
+	// We could enumerate the tools here, but we'll need to expose that information
+	// declaratively in the MCP server, so for the moment let's just check the existence
+	// of an issue and repo tool, and the non-existence of a pull_request tool.
+	var toolsContains = func(expectedName string) bool {
+		return slices.ContainsFunc(response.Tools, func(tool mcp.Tool) bool {
+			return tool.Name == expectedName
+		})
+	}
+
+	require.True(t, toolsContains("get_issue"), "expected to find 'get_issue' tool")
+	require.True(t, toolsContains("list_branches"), "expected to find 'list_branches' tool")
+	require.False(t, toolsContains("get_pull_request"), "expected not to find 'get_pull_request' tool")
 }