Skip to content

Commit 6cda0e7

Browse files
sd2ksd2k
committed
Add Prometheus tools
1 parent 317a563 commit 6cda0e7

File tree

5 files changed

+536
-37
lines changed

5 files changed

+536
-37
lines changed

‎client.go

+80-9
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,6 @@ import (
1313
"github.com/mark3labs/mcp-go/server"
1414
)
1515

16-
type clientKey struct{}
17-
1816
const (
1917
grafanaURLEnvVar = "GRAFANA_URL"
2018
grafanaAPIEnvVar = "GRAFANA_API_KEY"
@@ -35,9 +33,48 @@ func urlAndAPIKeyFromHeaders(req *http.Request) (string, string) {
3533
return u, apiKey
3634
}
3735

38-
// ExtractClientFromEnv is a StdioContextFunc that extracts Grafana configuration
36+
type grafanaURLKey struct{}
37+
type grafanaAPIKeyKey struct{}
38+
39+
// ExtractGrafanaInfoFromEnv is a StdioContextFunc that extracts Grafana configuration
40+
// from environment variables and injects a configured client into the context.
41+
var ExtractGrafanaInfoFromEnv server.StdioContextFunc = func(ctx context.Context) context.Context {
42+
u, apiKey := urlAndAPIKeyFromEnv()
43+
return WithGrafanaURL(WithGrafanaAPIKey(ctx, apiKey), u)
44+
}
45+
46+
// ExtractGrafanaInfoFromHeaders is a SSEContextFunc that extracts Grafana configuration
47+
// from request headers and injects a configured client into the context.
48+
var ExtractGrafanaInfoFromHeaders server.SSEContextFunc = func(ctx context.Context, req *http.Request) context.Context {
49+
u, apiKey := urlAndAPIKeyFromHeaders(req)
50+
return WithGrafanaURL(WithGrafanaAPIKey(ctx, apiKey), u)
51+
}
52+
53+
// WithGrafanaURL adds the Grafana URL to the context.
54+
func WithGrafanaURL(ctx context.Context, url string) context.Context {
55+
return context.WithValue(ctx, grafanaURLKey{}, url)
56+
}
57+
58+
// WithGrafanaAPIKey adds the Grafana API key to the context.
59+
func WithGrafanaAPIKey(ctx context.Context, apiKey string) context.Context {
60+
return context.WithValue(ctx, grafanaAPIKeyKey{}, apiKey)
61+
}
62+
63+
// GrafanaURLFromContext extracts the Grafana URL from the context.
64+
func GrafanaURLFromContext(ctx context.Context) string {
65+
return ctx.Value(grafanaURLKey{}).(string)
66+
}
67+
68+
// GrafanaAPIKeyFromContext extracts the Grafana API key from the context.
69+
func GrafanaAPIKeyFromContext(ctx context.Context) string {
70+
return ctx.Value(grafanaAPIKeyKey{}).(string)
71+
}
72+
73+
type grafanaClientKey struct{}
74+
75+
// ExtractGrafanaClientFromEnv is a StdioContextFunc that extracts Grafana configuration
3976
// from environment variables and injects a configured client into the context.
40-
var ExtractClientFromEnv server.StdioContextFunc = func(ctx context.Context) context.Context {
77+
var ExtractGrafanaClientFromEnv server.StdioContextFunc = func(ctx context.Context) context.Context {
4178
cfg := client.DefaultTransportConfig()
4279
// Extract transport config from env vars, and set it on the context.
4380
if u, ok := os.LookupEnv(grafanaURLEnvVar); ok {
@@ -57,12 +94,12 @@ var ExtractClientFromEnv server.StdioContextFunc = func(ctx context.Context) con
5794
}
5895

5996
client := client.NewHTTPClientWithConfig(strfmt.Default, cfg)
60-
return context.WithValue(ctx, clientKey{}, client)
97+
return context.WithValue(ctx, grafanaClientKey{}, client)
6198
}
6299

63-
// ExtractClientFromHeaders is a SSEContextFunc that extracts Grafana configuration
100+
// ExtractGrafanaClientFromHeaders is a SSEContextFunc that extracts Grafana configuration
64101
// from request headers and injects a configured client into the context.
65-
var ExtractClientFromHeaders server.SSEContextFunc = func(ctx context.Context, req *http.Request) context.Context {
102+
var ExtractGrafanaClientFromHeaders server.SSEContextFunc = func(ctx context.Context, req *http.Request) context.Context {
66103
cfg := client.DefaultTransportConfig()
67104
// Extract transport config from request headers, and set it on the context.
68105
u, apiKey := urlAndAPIKeyFromHeaders(req)
@@ -85,12 +122,12 @@ var ExtractClientFromHeaders server.SSEContextFunc = func(ctx context.Context, r
85122
//
86123
// It can be retrieved using GrafanaClientFromContext.
87124
func WithGrafanaClient(ctx context.Context, client *client.GrafanaHTTPAPI) context.Context {
88-
return context.WithValue(ctx, clientKey{}, client)
125+
return context.WithValue(ctx, grafanaClientKey{}, client)
89126
}
90127

91128
// GrafanaClientFromContext retrieves the Grafana client from the context.
92129
func GrafanaClientFromContext(ctx context.Context) *client.GrafanaHTTPAPI {
93-
c, ok := ctx.Value(clientKey{}).(*client.GrafanaHTTPAPI)
130+
c, ok := ctx.Value(grafanaClientKey{}).(*client.GrafanaHTTPAPI)
94131
if !ok {
95132
return nil
96133
}
@@ -124,3 +161,37 @@ func IncidentClientFromContext(ctx context.Context) *incident.Client {
124161
}
125162
return c
126163
}
164+
165+
// ComposeStdioContextFuncs composes multiple StdioContextFuncs into a single one.
166+
func ComposeStdioContextFuncs(funcs ...server.StdioContextFunc) server.StdioContextFunc {
167+
return func(ctx context.Context) context.Context {
168+
for _, f := range funcs {
169+
ctx = f(ctx)
170+
}
171+
return ctx
172+
}
173+
}
174+
175+
// ComposeSSEContextFuncs composes multiple SSEContextFuncs into a single one.
176+
func ComposeSSEContextFuncs(funcs ...server.SSEContextFunc) server.SSEContextFunc {
177+
return func(ctx context.Context, req *http.Request) context.Context {
178+
for _, f := range funcs {
179+
ctx = f(ctx, req)
180+
}
181+
return ctx
182+
}
183+
}
184+
185+
// ComposedStdioContextFunc is a StdioContextFunc that comprises all predefined StdioContextFuncs.
186+
var ComposedStdioContextFunc = ComposeStdioContextFuncs(
187+
ExtractGrafanaInfoFromEnv,
188+
ExtractGrafanaClientFromEnv,
189+
ExtractIncidentClientFromEnv,
190+
)
191+
192+
// ComposedSSEContextFunc is a SSEContextFunc that comprises all predefined SSEContextFuncs.
193+
var ComposedSSEContextFunc = ComposeSSEContextFuncs(
194+
ExtractGrafanaInfoFromHeaders,
195+
ExtractGrafanaClientFromHeaders,
196+
ExtractIncidentClientFromHeaders,
197+
)

‎cmd/main.go

+3-2
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ func newServer() *server.MCPServer {
2121
tools.AddSearchTools(s)
2222
tools.AddDatasourceTools(s)
2323
tools.AddIncidentTools(s)
24+
tools.AddPrometheusTools(s)
2425
return s
2526
}
2627

@@ -30,12 +31,12 @@ func run(transport string) error {
3031
switch transport {
3132
case "stdio":
3233
srv := server.NewStdioServer(s)
33-
srv.SetContextFunc(mcpgrafana.ExtractClientFromEnv)
34+
srv.SetContextFunc(mcpgrafana.ComposedStdioContextFunc)
3435
return srv.Listen(context.Background(), os.Stdin, os.Stdout)
3536
case "sse":
3637
addr := "http://localhost:8080"
3738
srv := server.NewSSEServer(s, addr)
38-
srv.SetContextFunc(mcpgrafana.ExtractClientFromHeaders)
39+
srv.SetContextFunc(mcpgrafana.ComposedSSEContextFunc)
3940
log.Printf("SSE server listening on %s", addr)
4041
if err := srv.Start("localhost:8080"); err != nil {
4142
return fmt.Errorf("Server error: %v", err)

‎go.mod

+29-8
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,23 @@ go 1.24.0
55
require (
66
github.com/go-openapi/strfmt v0.23.0
77
github.com/grafana/grafana-openapi-client-go v0.0.0-20250108132429-8d7e1f158f65
8+
github.com/grafana/incident-go v0.0.0-20250211094540-dc6a98fdae43
89
github.com/invopop/jsonschema v0.13.0
910
github.com/mark3labs/mcp-go v0.8.5
11+
github.com/prometheus/client_golang v1.21.0
12+
github.com/prometheus/common v0.62.0
13+
github.com/prometheus/prometheus v0.302.1
1014
github.com/stretchr/testify v1.10.0
1115
)
1216

1317
require (
1418
github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 // indirect
1519
github.com/bahlo/generic-list-go v0.2.0 // indirect
20+
github.com/beorn7/perks v1.0.1 // indirect
1621
github.com/buger/jsonparser v1.1.1 // indirect
17-
github.com/davecgh/go-spew v1.1.1 // indirect
18-
github.com/go-logr/logr v1.4.1 // indirect
22+
github.com/cespare/xxhash/v2 v2.3.0 // indirect
23+
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
24+
github.com/go-logr/logr v1.4.2 // indirect
1925
github.com/go-logr/stdr v1.2.2 // indirect
2026
github.com/go-openapi/analysis v0.23.0 // indirect
2127
github.com/go-openapi/errors v0.22.0 // indirect
@@ -27,19 +33,34 @@ require (
2733
github.com/go-openapi/swag v0.23.0 // indirect
2834
github.com/go-openapi/validate v0.24.0 // indirect
2935
github.com/google/uuid v1.6.0 // indirect
30-
github.com/grafana/incident-go v0.0.0-20250211094540-dc6a98fdae43 // indirect
36+
github.com/grafana/regexp v0.0.0-20240518133315-a468a5bfb3bc // indirect
3137
github.com/josharian/intern v1.0.0 // indirect
38+
github.com/jpillora/backoff v1.0.0 // indirect
39+
github.com/json-iterator/go v1.1.12 // indirect
3240
github.com/mailru/easyjson v0.7.7 // indirect
3341
github.com/mitchellh/mapstructure v1.5.0 // indirect
42+
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
43+
github.com/modern-go/reflect2 v1.0.2 // indirect
44+
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
45+
github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f // indirect
3446
github.com/oklog/ulid v1.3.1 // indirect
3547
github.com/opentracing/opentracing-go v1.2.0 // indirect
36-
github.com/pmezard/go-difflib v1.0.0 // indirect
48+
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
49+
github.com/prometheus/client_model v0.6.1 // indirect
50+
github.com/prometheus/procfs v0.15.1 // indirect
3751
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
3852
go.mongodb.org/mongo-driver v1.14.0 // indirect
39-
go.opentelemetry.io/otel v1.24.0 // indirect
40-
go.opentelemetry.io/otel/metric v1.24.0 // indirect
41-
go.opentelemetry.io/otel/trace v1.24.0 // indirect
42-
golang.org/x/sync v0.6.0 // indirect
53+
go.opentelemetry.io/auto/sdk v1.1.0 // indirect
54+
go.opentelemetry.io/otel v1.34.0 // indirect
55+
go.opentelemetry.io/otel/metric v1.34.0 // indirect
56+
go.opentelemetry.io/otel/trace v1.34.0 // indirect
57+
golang.org/x/net v0.34.0 // indirect
58+
golang.org/x/oauth2 v0.25.0 // indirect
59+
golang.org/x/sync v0.10.0 // indirect
60+
golang.org/x/sys v0.29.0 // indirect
61+
golang.org/x/text v0.21.0 // indirect
62+
google.golang.org/protobuf v1.36.4 // indirect
63+
gopkg.in/yaml.v2 v2.4.0 // indirect
4364
gopkg.in/yaml.v3 v3.0.1 // indirect
4465
)
4566

0 commit comments

Comments
 (0)