chore: separate toolset creation from init and use typed error (#487)

SamMorrowDrums · web-flow · commit c17ebfe50b8e · 2025-06-06T14:45:03.000+02:00
diff --git a/internal/ghmcp/server.go b/internal/ghmcp/server.go
@@ -113,26 +113,22 @@ func NewMCPServer(cfg MCPServerConfig) (*server.MCPServer, error) {
 	}
 
 	// Create default toolsets
-	toolsets, err := github.InitToolsets(
-		enabledToolsets,
-		cfg.ReadOnly,
-		getClient,
-		getGQLClient,
-		cfg.Translator,
-	)
+	tsg := github.DefaultToolsetGroup(cfg.ReadOnly, getClient, getGQLClient, cfg.Translator)
+	err = tsg.EnableToolsets(enabledToolsets)
+
 	if err != nil {
-		return nil, fmt.Errorf("failed to initialize toolsets: %w", err)
+		return nil, fmt.Errorf("failed to enable toolsets: %w", err)
 	}
 
 	context := github.InitContextToolset(getClient, cfg.Translator)
 	github.RegisterResources(ghServer, getClient, cfg.Translator)
 
 	// Register the tools with the server
-	toolsets.RegisterTools(ghServer)
+	tsg.RegisterTools(ghServer)
 	context.RegisterTools(ghServer)
 
 	if cfg.DynamicToolsets {
-		dynamic := github.InitDynamicToolset(ghServer, toolsets, cfg.Translator)
+		dynamic := github.InitDynamicToolset(ghServer, tsg, cfg.Translator)
 		dynamic.RegisterTools(ghServer)
 	}
 
diff --git a/pkg/github/tools.go b/pkg/github/tools.go
@@ -15,8 +15,7 @@ type GetGQLClientFn func(context.Context) (*githubv4.Client, error)
 
 var DefaultTools = []string{"all"}
 
-func InitToolsets(passedToolsets []string, readOnly bool, getClient GetClientFn, getGQLClient GetGQLClientFn, t translations.TranslationHelperFunc) (*toolsets.ToolsetGroup, error) {
-	// Create a new toolset group
+func DefaultToolsetGroup(readOnly bool, getClient GetClientFn, getGQLClient GetGQLClientFn, t translations.TranslationHelperFunc) *toolsets.ToolsetGroup {
 	tsg := toolsets.NewToolsetGroup(readOnly)
 
 	// Define all available features with their default state (disabled)
@@ -116,13 +115,8 @@ func InitToolsets(passedToolsets []string, readOnly bool, getClient GetClientFn,
 	tsg.AddToolset(secretProtection)
 	tsg.AddToolset(notifications)
 	tsg.AddToolset(experiments)
-	// Enable the requested features
 
-	if err := tsg.EnableToolsets(passedToolsets); err != nil {
-		return nil, err
-	}
-
-	return tsg, nil
+	return tsg
 }
 
 func InitContextToolset(getClient GetClientFn, t translations.TranslationHelperFunc) *toolsets.Toolset {
diff --git a/pkg/toolsets/toolsets.go b/pkg/toolsets/toolsets.go
@@ -7,6 +7,28 @@ import (
 	"github.com/mark3labs/mcp-go/server"
 )
 
+type ToolsetDoesNotExistError struct {
+	Name string
+}
+
+func (e *ToolsetDoesNotExistError) Error() string {
+	return fmt.Sprintf("toolset %s does not exist", e.Name)
+}
+
+func (e *ToolsetDoesNotExistError) Is(target error) bool {
+	if target == nil {
+		return false
+	}
+	if _, ok := target.(*ToolsetDoesNotExistError); ok {
+		return true
+	}
+	return false
+}
+
+func NewToolsetDoesNotExistError(name string) *ToolsetDoesNotExistError {
+	return &ToolsetDoesNotExistError{Name: name}
+}
+
 func NewServerTool(tool mcp.Tool, handler server.ToolHandlerFunc) server.ServerTool {
 	return server.ServerTool{Tool: tool, Handler: handler}
 }
@@ -150,7 +172,7 @@ func (tg *ToolsetGroup) EnableToolsets(names []string) error {
 func (tg *ToolsetGroup) EnableToolset(name string) error {
 	toolset, exists := tg.Toolsets[name]
 	if !exists {
-		return fmt.Errorf("toolset %s does not exist", name)
+		return NewToolsetDoesNotExistError(name)
 	}
 	toolset.Enabled = true
 	tg.Toolsets[name] = toolset
@@ -162,3 +184,11 @@ func (tg *ToolsetGroup) RegisterTools(s *server.MCPServer) {
 		toolset.RegisterTools(s)
 	}
 }
+
+func (tg *ToolsetGroup) GetToolset(name string) (*Toolset, error) {
+	toolset, exists := tg.Toolsets[name]
+	if !exists {
+		return nil, NewToolsetDoesNotExistError(name)
+	}
+	return toolset, nil
+}
diff --git a/pkg/toolsets/toolsets_test.go b/pkg/toolsets/toolsets_test.go
@@ -1,6 +1,7 @@
 package toolsets
 
 import (
+	"errors"
 	"testing"
 )
 
@@ -151,6 +152,9 @@ func TestEnableToolsets(t *testing.T) {
 	if err == nil {
 		t.Error("Expected error when enabling list with non-existent toolset")
 	}
+	if !errors.Is(err, NewToolsetDoesNotExistError("non-existent")) {
+		t.Errorf("Expected ToolsetDoesNotExistError when enabling non-existent toolset, got: %v", err)
+	}
 
 	// Test with empty list
 	err = tsg.EnableToolsets([]string{})
@@ -207,7 +211,7 @@ func TestEnableEverything(t *testing.T) {
 func TestIsEnabledWithEverythingOn(t *testing.T) {
 	tsg := NewToolsetGroup(false)
 
-	// Enable "everything"
+	// Enable "all"
 	err := tsg.EnableToolsets([]string{"all"})
 	if err != nil {
 		t.Errorf("Expected no error when enabling 'all', got: %v", err)
@@ -222,3 +226,27 @@ func TestIsEnabledWithEverythingOn(t *testing.T) {
 		t.Error("Expected IsEnabled to return true for any toolset when everythingOn is true")
 	}
 }
+
+func TestToolsetGroup_GetToolset(t *testing.T) {
+	tsg := NewToolsetGroup(false)
+	toolset := NewToolset("my-toolset", "desc")
+	tsg.AddToolset(toolset)
+
+	// Should find the toolset
+	got, err := tsg.GetToolset("my-toolset")
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+	if got != toolset {
+		t.Errorf("expected to get the same toolset instance")
+	}
+
+	// Should not find a non-existent toolset
+	_, err = tsg.GetToolset("does-not-exist")
+	if err == nil {
+		t.Error("expected error for missing toolset, got nil")
+	}
+	if !errors.Is(err, NewToolsetDoesNotExistError("does-not-exist")) {
+		t.Errorf("expected error to be ToolsetDoesNotExistError, got %v", err)
+	}
+}
