Skip to content

Commit b65e541

Browse files
tommaso-morotommaso-moro
committed
add ability to enable specific tools
1 parent 100e682 commit b65e541

File tree

3 files changed

+116
-16
lines changed

3 files changed

+116
-16
lines changed

‎cmd/github-mcp-server/main.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,22 @@ var (
4545
return fmt.Errorf("failed to unmarshal toolsets: %w", err)
4646
}
4747

48+
var enabledTools []string
49+
if err := viper.UnmarshalKey("tools", &enabledTools); err != nil {
50+
return fmt.Errorf("failed to unmarshal tools: %w", err)
51+
}
52+
53+
// If specific tools are requested, ignore toolsets completely
54+
if len(enabledTools) > 0 {
55+
enabledToolsets = nil
56+
}
57+
4858
stdioServerConfig := ghmcp.StdioServerConfig{
4959
Version: version,
5060
Host: viper.GetString("host"),
5161
Token: token,
5262
EnabledToolsets: enabledToolsets,
63+
EnabledTools: enabledTools,
5364
DynamicToolsets: viper.GetBool("dynamic_toolsets"),
5465
ReadOnly: viper.GetBool("read-only"),
5566
ExportTranslations: viper.GetBool("export-translations"),
@@ -70,6 +81,7 @@ func init() {
7081

7182
// Add global flags that will be shared by all commands
7283
rootCmd.PersistentFlags().StringSlice("toolsets", github.DefaultTools, "An optional comma separated list of groups of tools to allow, defaults to enabling all")
84+
rootCmd.PersistentFlags().StringSlice("tools", []string{}, "An optional comma separated list of specific tools to enable (requires their toolsets to be enabled)")
7385
rootCmd.PersistentFlags().Bool("dynamic-toolsets", false, "Enable dynamic toolsets")
7486
rootCmd.PersistentFlags().Bool("read-only", false, "Restrict the server to read-only operations")
7587
rootCmd.PersistentFlags().String("log-file", "", "Path to log file")
@@ -80,6 +92,7 @@ func init() {
8092

8193
// Bind flag to viper
8294
_ = viper.BindPFlag("toolsets", rootCmd.PersistentFlags().Lookup("toolsets"))
95+
_ = viper.BindPFlag("tools", rootCmd.PersistentFlags().Lookup("tools"))
8396
_ = viper.BindPFlag("dynamic_toolsets", rootCmd.PersistentFlags().Lookup("dynamic-toolsets"))
8497
_ = viper.BindPFlag("read-only", rootCmd.PersistentFlags().Lookup("read-only"))
8598
_ = viper.BindPFlag("log-file", rootCmd.PersistentFlags().Lookup("log-file"))
@@ -90,6 +103,7 @@ func init() {
90103

91104
// Add subcommands
92105
rootCmd.AddCommand(stdioCmd)
106+
rootCmd.AddCommand(wizardCmd)
93107
}
94108

95109
func initConfig() {

‎internal/ghmcp/server.go‎

Lines changed: 60 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ type MCPServerConfig struct {
3838
// See: https://github.com/github/github-mcp-server?tab=readme-ov-file#tool-configuration
3939
EnabledToolsets []string
4040

41+
EnabledTools []string
42+
4143
// Whether to enable dynamic toolsets
4244
// See: https://github.com/github/github-mcp-server?tab=readme-ov-file#dynamic-tool-discovery
4345
DynamicToolsets bool
@@ -141,11 +143,61 @@ func NewMCPServer(cfg MCPServerConfig) (*server.MCPServer, error) {
141143
}
142144

143145
// Create default toolsets
144-
tsg := github.DefaultToolsetGroup(cfg.ReadOnly, getClient, getGQLClient, getRawClient, cfg.Translator, cfg.ContentWindowSize)
145-
err = tsg.EnableToolsets(enabledToolsets)
146-
147-
if err != nil {
148-
return nil, fmt.Errorf("failed to enable toolsets: %w", err)
146+
tsg := github.DefaultToolsetGroup(cfg.ReadOnly, getClient, getGQLClient, getRawClient, cfg.Translator, cfg.ContentWindowSize)
147+
148+
// When tools are specified, we enable only those tools
149+
if len(cfg.EnabledTools) > 0 {
150+
fmt.Fprintf(os.Stderr, "DEBUG: Specific tools mode activated with tools: %v\n", cfg.EnabledTools)
151+
152+
// Build a map of tool name -> toolset name
153+
toolToToolsetMap := make(map[string]string)
154+
for toolsetName, toolset := range tsg.Toolsets {
155+
for _, tool := range toolset.GetAvailableTools() {
156+
toolToToolsetMap[tool.Tool.Name] = toolsetName
157+
}
158+
}
159+
160+
fmt.Fprintf(os.Stderr, "DEBUG: Built tool map with %d tools\n", len(toolToToolsetMap))
161+
162+
// Determine which toolsets need to be enabled based on requested tools
163+
toolsetToToolsMap := make(map[string][]string)
164+
for _, toolName := range cfg.EnabledTools {
165+
if toolsetName, ok := toolToToolsetMap[toolName]; ok {
166+
toolsetToToolsMap[toolsetName] = append(toolsetToToolsMap[toolsetName], toolName)
167+
fmt.Fprintf(os.Stderr, "DEBUG: Tool '%s' belongs to toolset '%s'\n", toolName, toolsetName)
168+
} else {
169+
fmt.Fprintf(os.Stderr, "DEBUG: WARNING - Tool '%s' not found in any toolset!\n", toolName)
170+
}
171+
}
172+
173+
// Enable the required toolsets
174+
requiredToolsets := make([]string, 0, len(toolsetToToolsMap))
175+
for toolsetName := range toolsetToToolsMap {
176+
requiredToolsets = append(requiredToolsets, toolsetName)
177+
}
178+
179+
fmt.Fprintf(os.Stderr, "DEBUG: Enabling %d toolsets: %v\n", len(requiredToolsets), requiredToolsets)
180+
181+
err = tsg.EnableToolsets(requiredToolsets)
182+
if err != nil {
183+
return nil, fmt.Errorf("failed to enable toolsets: %w", err)
184+
}
185+
186+
// Enable only specific tools in each toolset
187+
for toolsetName, toolNames := range toolsetToToolsMap {
188+
toolset, err := tsg.GetToolset(toolsetName)
189+
if err != nil {
190+
return nil, fmt.Errorf("failed to get toolset %s: %w", toolsetName, err)
191+
}
192+
fmt.Fprintf(os.Stderr, "DEBUG: Enabling specific tools %v for toolset '%s'\n", toolNames, toolsetName)
193+
toolset.EnableSpecificTools(toolNames)
194+
}
195+
} else {
196+
// Normal toolset mode - enable requested toolsets (or all by default)
197+
err = tsg.EnableToolsets(enabledToolsets)
198+
if err != nil {
199+
return nil, fmt.Errorf("failed to enable toolsets: %w", err)
200+
}
149201
}
150202

151203
// Register all mcp functionality with the server
@@ -173,6 +225,8 @@ type StdioServerConfig struct {
173225
// See: https://github.com/github/github-mcp-server?tab=readme-ov-file#tool-configuration
174226
EnabledToolsets []string
175227

228+
EnabledTools []string
229+
176230
// Whether to enable dynamic toolsets
177231
// See: https://github.com/github/github-mcp-server?tab=readme-ov-file#dynamic-tool-discovery
178232
DynamicToolsets bool
@@ -207,6 +261,7 @@ func RunStdioServer(cfg StdioServerConfig) error {
207261
Host: cfg.Host,
208262
Token: cfg.Token,
209263
EnabledToolsets: cfg.EnabledToolsets,
264+
EnabledTools: cfg.EnabledTools,
210265
DynamicToolsets: cfg.DynamicToolsets,
211266
ReadOnly: cfg.ReadOnly,
212267
Translator: t,

‎pkg/toolsets/toolsets.go‎

Lines changed: 42 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package toolsets
22

33
import (
44
"fmt"
5+
"os"
56

67
"github.com/mark3labs/mcp-go/mcp"
78
"github.com/mark3labs/mcp-go/server"
@@ -60,6 +61,16 @@ type Toolset struct {
6061
resourceTemplates []server.ServerResourceTemplate
6162
// prompts are also not tools but are namespaced similarly
6263
prompts []server.ServerPrompt
64+
enabledTools map[string]bool
65+
}
66+
67+
func (t *Toolset) EnableSpecificTools(toolNames []string) {
68+
if t.enabledTools == nil {
69+
t.enabledTools = make(map[string]bool)
70+
}
71+
for _, name := range toolNames {
72+
t.enabledTools[name] = true
73+
}
6374
}
6475

6576
func (t *Toolset) GetActiveTools() []server.ServerTool {
@@ -80,17 +91,37 @@ func (t *Toolset) GetAvailableTools() []server.ServerTool {
8091
}
8192

8293
func (t *Toolset) RegisterTools(s *server.MCPServer) {
83-
if !t.Enabled {
84-
return
85-
}
86-
for _, tool := range t.readTools {
87-
s.AddTool(tool.Tool, tool.Handler)
88-
}
89-
if !t.readOnly {
90-
for _, tool := range t.writeTools {
91-
s.AddTool(tool.Tool, tool.Handler)
92-
}
93-
}
94+
if !t.Enabled {
95+
fmt.Fprintf(os.Stderr, "DEBUG RegisterTools: Toolset '%s' is not enabled, skipping\n", t.Name)
96+
return
97+
}
98+
fmt.Fprintf(os.Stderr, "DEBUG RegisterTools: Toolset '%s' is enabled, enabledTools=%v\n", t.Name, t.enabledTools)
99+
for _, tool := range t.readTools {
100+
// Check if granular control is active
101+
if t.enabledTools != nil {
102+
// In selective mode, only register if tool is explicitly enabled
103+
if enabled, exists := t.enabledTools[tool.Tool.Name]; !exists || !enabled {
104+
fmt.Fprintf(os.Stderr, "DEBUG RegisterTools: Skipping read tool '%s' (exists=%v, enabled=%v)\n", tool.Tool.Name, exists, enabled)
105+
continue // Skip this tool
106+
}
107+
}
108+
fmt.Fprintf(os.Stderr, "DEBUG RegisterTools: Registering read tool '%s'\n", tool.Tool.Name)
109+
s.AddTool(tool.Tool, tool.Handler)
110+
}
111+
if !t.readOnly {
112+
for _, tool := range t.writeTools {
113+
// Same check for write tools
114+
if t.enabledTools != nil {
115+
// In selective mode, only register if tool is explicitly enabled
116+
if enabled, exists := t.enabledTools[tool.Tool.Name]; !exists || !enabled {
117+
fmt.Fprintf(os.Stderr, "DEBUG RegisterTools: Skipping write tool '%s' (exists=%v, enabled=%v)\n", tool.Tool.Name, exists, enabled)
118+
continue
119+
}
120+
}
121+
fmt.Fprintf(os.Stderr, "DEBUG RegisterTools: Registering write tool '%s'\n", tool.Tool.Name)
122+
s.AddTool(tool.Tool, tool.Handler)
123+
}
124+
}
94125
}
95126

96127
func (t *Toolset) AddResourceTemplates(templates ...server.ServerResourceTemplate) *Toolset {

0 commit comments

Comments
 (0)