feat: allow customizing tools for select backends, add thanos support

This is a feature I've talked about with several people already, and now
that the groundwork is laid for customizing toolsets loaded to the
server, this adds a further abstraction to allow customizing the tools
loaded by the MCP server for a specific prometheus compatible backend.

This is done by maintaining a mapping of tool name -> {description +
tool implementation}, so it's possible to maintain a mapping of
available tools and their implementations per-backend that we choose to
support.

This commit also adds initial customized support for thanos. Other
backends and tool customizations can be added in future work.

Signed-off-by: TJ Hoplock <[email protected]>

Author: tjhop

README.md

@@ -109,6 +109,33 @@ Would result in the following tools being loaded:
 - `runtime_info`
 - `series`
 
+#### Prometheus Compatible Backends
+
+There are many Prometheus compatible backends that can be used to extend prometheus in a variety of ways, often with the goals of offering long term storage or query aggregation from multiple prometheus instances.
+Some examples can be found in the [Remote Storage](https://prometheus.io/docs/operating/integrations/#remote-endpoints-and-storage) of prometheus' docs.
+
+Many of those services also offer a "prometheus compatible" API that can be used to query/interact with the data using native promQL.
+In general, this MCP server should at a minimum work for other prometheus API compatible services to execute queries and interact with the series/labels/metadata endpoints for metric and label discovery.
+Beyond that, there may be API differences as the different systems implement different parts/extensions of the API for their needs.
+
+Examples:
+- Thanos does not use a centralized config, so the config endpoint is not implemented and thus the config tool fails.
+- Mimir and Cortex implement extra endpoints to manage/add/remove rules
+
+To workaround this and provide a better experience on some of the commonly used Prometheus compatible systems, this project may add direct support for select systems to provide different/more tools.
+Choosing a specific prometheus backend implementation can be done with the [`--prometheus.backend` flag](#command-line-flags).
+The list of available backend implementations on a given release of the MCP server can be found in the output of the [`--help` flag](#command-line-flags).
+Qualifications and support criteria are still under consideration, please open an issue to request support/features for a specific backend for further discussion.
+
+##### Prometheus Backend Implementation Differences
+
+| Backend | Tool | Add/Remove/Change | Notes |
+| --- | --- | --- | --- |
+| `prometheus` | n/a | none | Standard prometheus tools. Functionally equivalent to `--mcp.tools="all"`. The default MCP server toolset. |
+| [`thanos`](https://thanos.io/) | `alertmanagers` | remove | Thanos does not implement the endpoint and the tool returns a `404`. |
+| [`thanos`](https://thanos.io/) | `config` | remove | Thanos does not use a centralized config, so it doesn't implement the endpoint and the tool returns a `404`. |
+| [`thanos`](https://thanos.io/) | `wal_replay_status` | remove | Thanos does not implement the endpoint and the tool returns a `404`. |
+
 ### Resources
 
 | Resource Name | Resource URI | Description | 

cmd/prometheus-mcp-server/main.go

@@ -12,6 +12,7 @@ import (
 	"os"
 	"os/signal"
 	"runtime"
+	"strings"
 	"syscall"
 	"time"
 
@@ -51,6 +52,12 @@ var (
 			" Please see project README for more information and the full list of tools.",
 	).Default("all").Strings()
 
+	flagPrometheusBackend = kingpin.Flag(
+		"prometheus.backend",
+		"Customize the toolset for a specific Prometheus API compatible backend."+
+			" Supported backends include: "+strings.Join(mcp.PrometheusBackends, ","),
+	).String()
+
 	flagPrometheusUrl = kingpin.Flag(
 		"prometheus.url",
 		"URL of the Prometheus instance to connect to",
@@ -134,7 +141,7 @@ func main() {
 		docsFs = docs
 	}
 
-	mcpServer := mcp.NewServer(ctx, logger, *flagPrometheusUrl, rt, *flagEnableTsdbAdminTools, *flagMcpTools, docsFs)
+	mcpServer := mcp.NewServer(ctx, logger, *flagPrometheusUrl, *flagPrometheusBackend, rt, *flagEnableTsdbAdminTools, *flagMcpTools, docsFs)
 	srv := setupServer(logger)
 
 	var g run.Group

pkg/mcp/server.go

@@ -213,7 +213,14 @@ func (m *telemetryMiddleware) ResourceMiddleware(next server.ResourceHandlerFunc
 	}
 }
 
-func NewServer(ctx context.Context, logger *slog.Logger, promUrl string, promRt http.RoundTripper, enableTsdbAdminTools bool, enabledTools []string, docs fs.FS) *server.MCPServer {
+func NewServer(ctx context.Context, logger *slog.Logger,
+	promUrl string,
+	prometheusBackend string,
+	promRt http.RoundTripper,
+	enableTsdbAdminTools bool,
+	enabledTools []string,
+	docs fs.FS,
+) *server.MCPServer {
 	hooks := &server.Hooks{}
 
 	hooks.AddAfterInitialize(func(ctx context.Context, id any, message *mcp.InitializeRequest, result *mcp.InitializeResult) {
@@ -319,6 +326,31 @@ func NewServer(ctx context.Context, logger *slog.Logger, promUrl string, promRt
 		}
 	}
 
+	// If a specific prometheus compatible backend was provided, that
+	// overrides defined tools, since dynamic tool registration is specific
+	// to prometheus tools and they may not be compatible with all
+	// prometheus backends.
+	backend := strings.ToLower(prometheusBackend)
+	switch backend {
+	case "": // If no backend entered, keep loaded toolset.
+	case "prometheus":
+		logger.Info("Setting tools based on provided prometheus backend", "backend", backend)
+		var backendToolset []server.ServerTool
+		for _, tool := range prometheusToolset {
+			backendToolset = append(backendToolset, tool)
+		}
+		toolset = backendToolset
+	case "thanos":
+		logger.Info("Setting tools based on provided prometheus backend", "backend", backend)
+		var backendToolset []server.ServerTool
+		for _, tool := range thanosToolset {
+			backendToolset = append(backendToolset, tool)
+		}
+		toolset = backendToolset
+	default:
+		logger.Warn("Prometheus backend does not have custom tool support, keeping the existing loaded toolset", "backend", backend, "toolset", enabledTools)
+	}
+
 	// Add tools.
 	mcpServer.AddTools(toolset...)
 

pkg/mcp/tools.go

@@ -36,6 +36,7 @@ var (
 		"snapshot",
 	}
 
+	// prometheusToolset contains all the tools to interact with standard prometheus through the HTTP API.
 	prometheusToolset = map[string]server.ServerTool{
 		"alertmanagers":     {Tool: prometheusAlertmanagersTool, Handler: prometheusAlertmanagersToolHandler},
 		"build_info":        {Tool: prometheusBuildinfoTool, Handler: prometheusBuildinfoToolHandler},
@@ -63,6 +64,46 @@ var (
 		"wal_replay_status": {Tool: prometheusWalReplayTool, Handler: prometheusWalReplayToolHandler},
 	}
 
+	// thanosToolset contains all the tools to interact with thanos as a
+	// prometheus HTTP API compatible backend.
+	//
+	// Currently, the only difference between thanosToolset and
+	// prometheusToolset is that thanosToolset has the following tools
+	// removed because they are not implemented in Thanos and return 404s:
+	// [alertmanagers, config, wal_replay_status].
+	thanosToolset = map[string]server.ServerTool{
+		"build_info":       {Tool: prometheusBuildinfoTool, Handler: prometheusBuildinfoToolHandler},
+		"clean_tombstones": {Tool: prometheusCleanTombstonesTool, Handler: prometheusCleanTombstonesToolHandler},
+		"delete_series":    {Tool: prometheusDeleteSeriesTool, Handler: prometheusDeleteSeriesToolHandler},
+		"docs_list":        {Tool: prometheusDocsListTool, Handler: prometheusDocsListToolHandler},
+		"docs_read":        {Tool: prometheusDocsReadTool, Handler: prometheusDocsReadToolHandler},
+		"docs_search":      {Tool: prometheusDocsSearchTool, Handler: prometheusDocsSearchToolHandler},
+		"exemplar_query":   {Tool: prometheusExemplarQueryTool, Handler: prometheusExemplarQueryToolHandler},
+		"flags":            {Tool: prometheusFlagsTool, Handler: prometheusFlagsToolHandler},
+		"label_names":      {Tool: prometheusLabelNamesTool, Handler: prometheusLabelNamesToolHandler},
+		"label_values":     {Tool: prometheusLabelValuesTool, Handler: prometheusLabelValuesToolHandler},
+		"list_alerts":      {Tool: prometheusListAlertsTool, Handler: prometheusListAlertsToolHandler},
+		"list_rules":       {Tool: prometheusRulesTool, Handler: prometheusRulesToolHandler},
+		"metric_metadata":  {Tool: prometheusMetricMetadataTool, Handler: prometheusMetricMetadataToolHandler},
+		"query":            {Tool: prometheusQueryTool, Handler: prometheusQueryToolHandler},
+		"range_query":      {Tool: prometheusRangeQueryTool, Handler: prometheusRangeQueryToolHandler},
+		"runtime_info":     {Tool: prometheusRuntimeinfoTool, Handler: prometheusRuntimeinfoToolHandler},
+		"series":           {Tool: prometheusSeriesTool, Handler: prometheusSeriesToolHandler},
+		"snapshot":         {Tool: prometheusSnapshotTool, Handler: prometheusSnapshotToolHandler},
+		"targets_metadata": {Tool: prometheusTargetsMetadataTool, Handler: prometheusTargetsMetadataToolHandler},
+		"list_targets":     {Tool: prometheusTargetsTool, Handler: prometheusTargetsToolHandler},
+		"tsdb_stats":       {Tool: prometheusTsdbStatsTool, Handler: prometheusTsdbStatsToolHandler},
+	}
+
+	// PrometheusBackends is a list of directly supported Prometheus API
+	// compatible backends. Backends other than prometheus itself may
+	// expose a different set of tools more tailored to the backend and/or
+	// change functionality of existing tools.
+	PrometheusBackends = []string{
+		"prometheus",
+		"thanos",
+	}
+
 	// Tools Definitions.
 
 	// Tools for Prometheus API.