Make Tool APIs more type safe and require fewer temporaries

sd2k · sd2k · commit b8a111c9904e · 2025-03-11T10:40:43.000Z
diff --git a/tools.go b/tools.go
@@ -11,15 +11,48 @@ import (
 	"github.com/mark3labs/mcp-go/server"
 )
 
-func MustTool(name, description string, toolHandler any) (mcp.Tool, server.ToolHandlerFunc) {
+// Tool is a struct that represents a tool definition and the function used
+// to handle tool calls.
+//
+// The simplest way to create a Tool is to use `MustTool`, or `ConvertTool`
+// if you wish to create tools at runtime and need to handle errors without
+// panicking.
+type Tool struct {
+	Tool    mcp.Tool
+	Handler server.ToolHandlerFunc
+}
+
+// Register adds the Tool to the given MCPServer.
+//
+// It is a convenience method that calls `server.MCPServer.Register` with the
+// Tool's Tool and Handler fields, allowing you to add the tool in a single
+// statement:
+//
+//	mcpgrafana.MustTool(name, description, toolHandler).Register(server)
+func (t *Tool) Register(mcp *server.MCPServer) {
+	mcp.AddTool(t.Tool, t.Handler)
+}
+
+// MustTool creates a new Tool from the given name, description, and toolHandler.
+// It panics if the tool cannot be created.
+func MustTool[T any](name, description string, toolHandler ToolHandlerFunc[T]) Tool {
 	tool, handler, err := ConvertTool(name, description, toolHandler)
 	if err != nil {
 		panic(err)
 	}
-	return tool, handler
+	return Tool{Tool: tool, Handler: handler}
 }
 
-func ConvertTool(name, description string, toolHandler any) (mcp.Tool, server.ToolHandlerFunc, error) {
+// ToolHandlerFunc is the type of a handler function for a tool.
+type ToolHandlerFunc[T any] = func(ctx context.Context, request T) (*mcp.CallToolResult, error)
+
+// ConvertTool converts a toolHandler function to a Tool and ToolHandlerFunc.
+//
+// The toolHandler function must have two arguments: a context.Context and a struct
+// to be used as the parameters for the tool. The second argument must not be a pointer,
+// should be marshalable to JSON, and the fields should have a `jsonschema` tag with the
+// description of the parameter.
+func ConvertTool[T any](name, description string, toolHandler ToolHandlerFunc[T]) (mcp.Tool, server.ToolHandlerFunc, error) {
 	zero := mcp.Tool{}
 	handlerValue := reflect.ValueOf(toolHandler)
 	handlerType := handlerValue.Type()
diff --git a/tools/datasources.go b/tools/datasources.go
@@ -26,7 +26,7 @@ func listDatasources(ctx context.Context, args ListDatasourcesParams) (*mcp.Call
 	return mcp.NewToolResultText(string(b)), nil
 }
 
-var ListDatasourcesTool, ListDatasourcesHandler = mcpgrafana.MustTool(
+var ListDatasources = mcpgrafana.MustTool(
 	"list_datasources",
 	"List datasources",
 	listDatasources,
@@ -49,7 +49,7 @@ func getDatasourceByUID(ctx context.Context, args GetDatasourceByUIDParams) (*mc
 	return mcp.NewToolResultText(string(b)), nil
 }
 
-var GetDatasourceByUIDTool, GetDatasourceByUIDHandler = mcpgrafana.MustTool(
+var GetDatasourceByUID = mcpgrafana.MustTool(
 	"get_datasource_by_uid",
 	"Get datasource by uid",
 	getDatasourceByUID,
@@ -72,14 +72,14 @@ func getDatasourceByName(ctx context.Context, args GetDatasourceByNameParams) (*
 	return mcp.NewToolResultText(string(b)), nil
 }
 
-var GetDatasourceByNameTool, GetDatasourceByNameHandler = mcpgrafana.MustTool(
+var GetDatasourceByName = mcpgrafana.MustTool(
 	"get_datasource_by_name",
 	"Get datasource by name",
 	getDatasourceByName,
 )
 
 func AddDatasourceTools(mcp *server.MCPServer) {
-	mcp.AddTool(ListDatasourcesTool, ListDatasourcesHandler)
-	mcp.AddTool(GetDatasourceByUIDTool, GetDatasourceByUIDHandler)
-	mcp.AddTool(GetDatasourceByNameTool, GetDatasourceByNameHandler)
+	ListDatasources.Register(mcp)
+	GetDatasourceByUID.Register(mcp)
+	GetDatasourceByName.Register(mcp)
 }
diff --git a/tools/incident.go b/tools/incident.go
@@ -17,7 +17,7 @@ type ListIncidentsParams struct {
 	Status string `json:"status" jsonschema:"description=The status of the incidents to include"`
 }
 
-func ListIncidents(ctx context.Context, args ListIncidentsParams) (*mcp.CallToolResult, error) {
+func listIncidents(ctx context.Context, args ListIncidentsParams) (*mcp.CallToolResult, error) {
 	c := mcpgrafana.IncidentClientFromContext(ctx)
 	is := incident.NewIncidentsService(c)
 	query := ""
@@ -44,10 +44,10 @@ func ListIncidents(ctx context.Context, args ListIncidentsParams) (*mcp.CallTool
 	return mcp.NewToolResultText(string(b)), nil
 }
 
-var ListIncidentsTool, ListIncidentsHandler = mcpgrafana.MustTool(
+var ListIncidents = mcpgrafana.MustTool(
 	"list_incidents",
 	"List incidents",
-	ListIncidents,
+	listIncidents,
 )
 
 type CreateIncidentParams struct {
@@ -61,7 +61,7 @@ type CreateIncidentParams struct {
 	Labels        []incident.IncidentLabel `json:"labels" jsonschema:"description=The labels to add to the incident"`
 }
 
-func CreateIncident(ctx context.Context, args CreateIncidentParams) (*mcp.CallToolResult, error) {
+func createIncident(ctx context.Context, args CreateIncidentParams) (*mcp.CallToolResult, error) {
 	c := mcpgrafana.IncidentClientFromContext(ctx)
 	is := incident.NewIncidentsService(c)
 	incident, err := is.CreateIncident(ctx, incident.CreateIncidentRequest{
@@ -84,10 +84,10 @@ func CreateIncident(ctx context.Context, args CreateIncidentParams) (*mcp.CallTo
 	return mcp.NewToolResultText(string(b)), nil
 }
 
-var CreateIncidentTool, CreateIncidentHandler = mcpgrafana.MustTool(
+var CreateIncident = mcpgrafana.MustTool(
 	"create_incident",
 	"Create an incident",
-	CreateIncident,
+	createIncident,
 )
 
 type AddActivityToIncidentParams struct {
@@ -96,7 +96,7 @@ type AddActivityToIncidentParams struct {
 	EventTime  string `json:"eventTime" jsonschema:"description=The time that the activity occurred. If not provided, the current time will be used"`
 }
 
-func AddActivityToIncident(ctx context.Context, args AddActivityToIncidentParams) (*mcp.CallToolResult, error) {
+func addActivityToIncident(ctx context.Context, args AddActivityToIncidentParams) (*mcp.CallToolResult, error) {
 	c := mcpgrafana.IncidentClientFromContext(ctx)
 	as := incident.NewActivityService(c)
 	activity, err := as.AddActivity(ctx, incident.AddActivityRequest{
@@ -115,14 +115,14 @@ func AddActivityToIncident(ctx context.Context, args AddActivityToIncidentParams
 	return mcp.NewToolResultText(string(b)), nil
 }
 
-var AddActivityToIncidentTool, AddActivityToIncidentHandler = mcpgrafana.MustTool(
+var AddActivityToIncident = mcpgrafana.MustTool(
 	"add_activity_to_incident",
 	"Add an activity to an incident",
-	AddActivityToIncident,
+	addActivityToIncident,
 )
 
 func AddIncidentTools(mcp *server.MCPServer) {
-	mcp.AddTool(ListIncidentsTool, ListIncidentsHandler)
-	mcp.AddTool(CreateIncidentTool, CreateIncidentHandler)
-	mcp.AddTool(AddActivityToIncidentTool, AddActivityToIncidentHandler)
+	ListIncidents.Register(mcp)
+	CreateIncident.Register(mcp)
+	AddActivityToIncident.Register(mcp)
 }
diff --git a/tools/prometheus.go b/tools/prometheus.go
@@ -42,7 +42,7 @@ type ListPrometheusMetricMetadataParams struct {
 	Metric         string `json:"metric" jsonschema:"description=The metric to query"`
 }
 
-func ListPrometheusMetricMetadata(ctx context.Context, args ListPrometheusMetricMetadataParams) (*mcp.CallToolResult, error) {
+func listPrometheusMetricMetadata(ctx context.Context, args ListPrometheusMetricMetadataParams) (*mcp.CallToolResult, error) {
 	promClient, err := promClientFromContext(ctx, args.DatasourceUID)
 	if err != nil {
 		return nil, fmt.Errorf("getting Prometheus client: %w", err)
@@ -64,10 +64,10 @@ func ListPrometheusMetricMetadata(ctx context.Context, args ListPrometheusMetric
 	return mcp.NewToolResultText(string(b)), nil
 }
 
-var ListPrometheusMetricMetadataTool, ListPrometheusMetricMetadataHandler = mcpgrafana.MustTool(
+var ListPrometheusMetricMetadata = mcpgrafana.MustTool(
 	"list_prometheus_metric_metadata",
 	"List Prometheus metric metadata",
-	ListPrometheusMetricMetadata,
+	listPrometheusMetricMetadata,
 )
 
 type QueryPrometheusParams struct {
@@ -79,7 +79,7 @@ type QueryPrometheusParams struct {
 	QueryType     string `json:"queryType,omitempty" jsonschema:"description=The type of query to use. Either 'range' or 'instant'"`
 }
 
-func QueryPrometheus(ctx context.Context, args QueryPrometheusParams) (*mcp.CallToolResult, error) {
+func queryPrometheus(ctx context.Context, args QueryPrometheusParams) (*mcp.CallToolResult, error) {
 	promClient, err := promClientFromContext(ctx, args.DatasourceUID)
 	if err != nil {
 		return nil, fmt.Errorf("getting Prometheus client: %w", err)
@@ -136,10 +136,10 @@ func QueryPrometheus(ctx context.Context, args QueryPrometheusParams) (*mcp.Call
 	return nil, fmt.Errorf("invalid query type: %s", queryType)
 }
 
-var QueryPrometheusTool, QueryPrometheusHandler = mcpgrafana.MustTool(
+var QueryPrometheus = mcpgrafana.MustTool(
 	"query_prometheus",
 	"Query Prometheus using a range or instant request",
-	QueryPrometheus,
+	queryPrometheus,
 )
 
 type ListPrometheusMetricNamesParams struct {
@@ -149,7 +149,7 @@ type ListPrometheusMetricNamesParams struct {
 	Page          int    `json:"page,omitempty" jsonschema:"description=The page number to return"`
 }
 
-func ListPrometheusMetricNames(ctx context.Context, args ListPrometheusMetricNamesParams) (*mcp.CallToolResult, error) {
+func listPrometheusMetricNames(ctx context.Context, args ListPrometheusMetricNamesParams) (*mcp.CallToolResult, error) {
 	promClient, err := promClientFromContext(ctx, args.DatasourceUID)
 	if err != nil {
 		return nil, fmt.Errorf("getting Prometheus client: %w", err)
@@ -207,10 +207,10 @@ func ListPrometheusMetricNames(ctx context.Context, args ListPrometheusMetricNam
 	return mcp.NewToolResultText(string(b)), nil
 }
 
-var ListPrometheusMetricNamesTool, ListPrometheusMetricNamesHandler = mcpgrafana.MustTool(
+var ListPrometheusMetricNames = mcpgrafana.MustTool(
 	"list_prometheus_metric_names",
 	"List metric names in a Prometheus datasource that match the given regex",
-	ListPrometheusMetricNames,
+	listPrometheusMetricNames,
 )
 
 type LabelMatcher struct {
@@ -244,7 +244,7 @@ type ListPrometheusLabelNamesParams struct {
 	Limit         int        `json:"limit,omitempty" jsonschema:"description=Optionally, the maximum number of results to return"`
 }
 
-func ListPrometheusLabelNames(ctx context.Context, args ListPrometheusLabelNamesParams) (*mcp.CallToolResult, error) {
+func listPrometheusLabelNames(ctx context.Context, args ListPrometheusLabelNamesParams) (*mcp.CallToolResult, error) {
 	promClient, err := promClientFromContext(ctx, args.DatasourceUID)
 	if err != nil {
 		return nil, fmt.Errorf("getting Prometheus client: %w", err)
@@ -289,10 +289,10 @@ func ListPrometheusLabelNames(ctx context.Context, args ListPrometheusLabelNames
 	return mcp.NewToolResultText(string(b)), nil
 }
 
-var ListPrometheusLabelNamesTool, ListPrometheusLabelNamesHandler = mcpgrafana.MustTool(
+var ListPrometheusLabelNames = mcpgrafana.MustTool(
 	"list_prometheus_label_names",
 	"List the label names in a Prometheus datasource",
-	ListPrometheusLabelNames,
+	listPrometheusLabelNames,
 )
 
 type ListPrometheusLabelValuesParams struct {
@@ -304,7 +304,7 @@ type ListPrometheusLabelValuesParams struct {
 	Limit         int        `json:"limit,omitempty" jsonschema:"description=Optionally, the maximum number of results to return"`
 }
 
-func ListPrometheusLabelValues(ctx context.Context, args ListPrometheusLabelValuesParams) (*mcp.CallToolResult, error) {
+func listPrometheusLabelValues(ctx context.Context, args ListPrometheusLabelValuesParams) (*mcp.CallToolResult, error) {
 	promClient, err := promClientFromContext(ctx, args.DatasourceUID)
 	if err != nil {
 		return nil, fmt.Errorf("getting Prometheus client: %w", err)
@@ -349,16 +349,16 @@ func ListPrometheusLabelValues(ctx context.Context, args ListPrometheusLabelValu
 	return mcp.NewToolResultText(string(b)), nil
 }
 
-var ListPrometheusLabelValuesTool, ListPrometheusLabelValuesHandler = mcpgrafana.MustTool(
+var ListPrometheusLabelValues = mcpgrafana.MustTool(
 	"list_prometheus_label_values",
 	"Get the values of a label in Prometheus",
-	ListPrometheusLabelValues,
+	listPrometheusLabelValues,
 )
 
 func AddPrometheusTools(mcp *server.MCPServer) {
-	mcp.AddTool(ListPrometheusMetricMetadataTool, ListPrometheusMetricMetadataHandler)
-	mcp.AddTool(QueryPrometheusTool, QueryPrometheusHandler)
-	mcp.AddTool(ListPrometheusMetricNamesTool, ListPrometheusMetricNamesHandler)
-	mcp.AddTool(ListPrometheusLabelNamesTool, ListPrometheusLabelNamesHandler)
-	mcp.AddTool(ListPrometheusLabelValuesTool, ListPrometheusLabelValuesHandler)
+	ListPrometheusMetricMetadata.Register(mcp)
+	QueryPrometheus.Register(mcp)
+	ListPrometheusMetricNames.Register(mcp)
+	ListPrometheusLabelNames.Register(mcp)
+	ListPrometheusLabelValues.Register(mcp)
 }
diff --git a/tools/search.go b/tools/search.go
@@ -33,12 +33,12 @@ func searchDashboards(ctx context.Context, args SearchDashboardsParams) (*mcp.Ca
 	return mcp.NewToolResultText(string(b)), nil
 }
 
-var SearchDashboardsTool, SearchDashboardsHandler = mcpgrafana.MustTool(
+var SearchDashboards = mcpgrafana.MustTool(
 	"search_dashboards",
 	"Search for dashboards",
 	searchDashboards,
 )
 
 func AddSearchTools(mcp *server.MCPServer) {
-	mcp.AddTool(SearchDashboardsTool, SearchDashboardsHandler)
+	SearchDashboards.Register(mcp)
 }
diff --git a/tools_test.go b/tools_test.go
@@ -131,48 +131,11 @@ func TestConvertTool(t *testing.T) {
 	})
 
 	t.Run("invalid handler types", func(t *testing.T) {
-		// Test non-function handler
-		_, _, err := ConvertTool("invalid", "description", "not a function")
-		assert.Error(t, err)
-		assert.Contains(t, err.Error(), "must be a function")
-
-		// Test wrong number of arguments
-		wrongArgsFunc := func(ctx context.Context) (*mcp.CallToolResult, error) {
-			return nil, nil
-		}
-		_, _, err = ConvertTool("invalid", "description", wrongArgsFunc)
-		assert.Error(t, err)
-		assert.Contains(t, err.Error(), "must have 2 arguments")
-
-		// Test wrong number of return values
-		wrongReturnFunc := func(ctx context.Context, params testToolParams) *mcp.CallToolResult {
-			return nil
-		}
-		_, _, err = ConvertTool("invalid", "description", wrongReturnFunc)
-		assert.Error(t, err)
-		assert.Contains(t, err.Error(), "must return 2 values")
-
-		// Test wrong first argument type
-		wrongFirstArgFunc := func(s string, params testToolParams) (*mcp.CallToolResult, error) {
-			return nil, nil
-		}
-		_, _, err = ConvertTool("invalid", "description", wrongFirstArgFunc)
-		assert.Error(t, err)
-		assert.Contains(t, err.Error(), "first argument must be context.Context")
-
-		// Test wrong first return value type
-		wrongFirstReturnFunc := func(ctx context.Context, params testToolParams) (string, error) {
-			return "", nil
-		}
-		_, _, err = ConvertTool("invalid", "description", wrongFirstReturnFunc)
-		assert.Error(t, err)
-		assert.Contains(t, err.Error(), "first return value must be mcp.CallToolResult")
-
 		// Test wrong second argument type (not a struct)
 		wrongSecondArgFunc := func(ctx context.Context, s string) (*mcp.CallToolResult, error) {
 			return nil, nil
 		}
-		_, _, err = ConvertTool("invalid", "description", wrongSecondArgFunc)
+		_, _, err := ConvertTool("invalid", "description", wrongSecondArgFunc)
 		assert.Error(t, err)
 		assert.Contains(t, err.Error(), "second argument must be a struct")
 	})

Original file line number	Diff line number	Diff line change
`@@ -33,12 +33,12 @@ func searchDashboards(ctx context.Context, args SearchDashboardsParams) (*mcp.Ca`
`33`	`33`	`return mcp.NewToolResultText(string(b)), nil`
`34`	`34`	`}`
`35`	`35`
`36`		`-var SearchDashboardsTool, SearchDashboardsHandler = mcpgrafana.MustTool(`
	`36`	`+var SearchDashboards = mcpgrafana.MustTool(`
`37`	`37`	`"search_dashboards",`
`38`	`38`	`"Search for dashboards",`
`39`	`39`	`searchDashboards,`
`40`	`40`	`)`
`41`	`41`
`42`	`42`	`func AddSearchTools(mcp *server.MCPServer) {`
`43`		`- mcp.AddTool(SearchDashboardsTool, SearchDashboardsHandler)`
	`43`	`+ SearchDashboards.Register(mcp)`
`44`	`44`	`}`