From 4099c5257421fd26f3e33b75abd30f4262589bd4 Mon Sep 17 00:00:00 2001 From: Sam Morrow Date: Fri, 6 Jun 2025 13:07:04 +0200 Subject: [PATCH] chore: separate toolset creation from init and use typed error --- internal/ghmcp/server.go | 16 ++++++---------- pkg/github/tools.go | 10 ++-------- pkg/toolsets/toolsets.go | 32 +++++++++++++++++++++++++++++++- pkg/toolsets/toolsets_test.go | 30 +++++++++++++++++++++++++++++- 4 files changed, 68 insertions(+), 20 deletions(-) diff --git a/internal/ghmcp/server.go b/internal/ghmcp/server.go index 8f5e16bc0..593411ae3 100644 --- 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 index ab0528174..f8e05fc85 100644 --- 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 index 7400119c8..fcb5e93b3 100644 --- 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 index 6d634fc4d..d74c94bbb 100644 --- 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) + } +}