This is page 5 of 5. Use http://codebase.md/grafana/mcp-grafana?lines=true&page={x} to view the full context.
# Directory Structure
```
├── .dockerignore
├── .github
│ ├── dependabot.yml
│ └── workflows
│ ├── docker.yml
│ ├── e2e.yml
│ ├── integration.yml
│ ├── release.yml
│ └── unit.yml
├── .gitignore
├── .golangci.yaml
├── .goreleaser.yaml
├── cmd
│ ├── linters
│ │ └── jsonschema
│ │ └── main.go
│ └── mcp-grafana
│ └── main.go
├── CODEOWNERS
├── docker-compose.yaml
├── Dockerfile
├── examples
│ └── tls_example.go
├── gemini-extension.json
├── go.mod
├── go.sum
├── image-tag
├── internal
│ └── linter
│ └── jsonschema
│ ├── jsonschema_lint_test.go
│ ├── jsonschema_lint.go
│ └── README.md
├── LICENSE
├── Makefile
├── mcpgrafana_test.go
├── mcpgrafana.go
├── proxied_client.go
├── proxied_handler.go
├── proxied_tools_test.go
├── proxied_tools.go
├── README.md
├── renovate.json
├── server.json
├── session_test.go
├── session.go
├── testdata
│ ├── dashboards
│ │ └── demo.json
│ ├── loki-config.yml
│ ├── prometheus-entrypoint.sh
│ ├── prometheus-seed.yml
│ ├── prometheus.yml
│ ├── promtail-config.yml
│ ├── provisioning
│ │ ├── alerting
│ │ │ ├── alert_rules.yaml
│ │ │ └── contact_points.yaml
│ │ ├── dashboards
│ │ │ └── dashboards.yaml
│ │ └── datasources
│ │ └── datasources.yaml
│ ├── tempo-config-2.yaml
│ └── tempo-config.yaml
├── tests
│ ├── .gitignore
│ ├── .python-version
│ ├── admin_test.py
│ ├── conftest.py
│ ├── dashboards_test.py
│ ├── disable_write_test.py
│ ├── health_test.py
│ ├── loki_test.py
│ ├── navigation_test.py
│ ├── pyproject.toml
│ ├── README.md
│ ├── tempo_test.py
│ ├── utils.py
│ └── uv.lock
├── tls_test.go
├── tools
│ ├── admin_test.go
│ ├── admin.go
│ ├── alerting_client_test.go
│ ├── alerting_client.go
│ ├── alerting_test.go
│ ├── alerting_unit_test.go
│ ├── alerting.go
│ ├── annotations_integration_test.go
│ ├── annotations_unit_test.go
│ ├── annotations.go
│ ├── asserts_cloud_test.go
│ ├── asserts_test.go
│ ├── asserts.go
│ ├── cloud_testing_utils.go
│ ├── dashboard_test.go
│ ├── dashboard.go
│ ├── datasources_test.go
│ ├── datasources.go
│ ├── folder.go
│ ├── incident_integration_test.go
│ ├── incident_test.go
│ ├── incident.go
│ ├── loki_test.go
│ ├── loki.go
│ ├── navigation_test.go
│ ├── navigation.go
│ ├── oncall_cloud_test.go
│ ├── oncall.go
│ ├── prometheus_test.go
│ ├── prometheus_unit_test.go
│ ├── prometheus.go
│ ├── pyroscope_test.go
│ ├── pyroscope.go
│ ├── search_test.go
│ ├── search.go
│ ├── sift_cloud_test.go
│ ├── sift.go
│ └── testcontext_test.go
├── tools_test.go
└── tools.go
```
# Files
--------------------------------------------------------------------------------
/mcpgrafana_test.go:
--------------------------------------------------------------------------------
```go
1 | //go:build unit
2 | // +build unit
3 |
4 | package mcpgrafana
5 |
6 | import (
7 | "context"
8 | "net/http"
9 | "testing"
10 |
11 | "github.com/go-openapi/runtime/client"
12 | grafana_client "github.com/grafana/grafana-openapi-client-go/client"
13 | "github.com/mark3labs/mcp-go/mcp"
14 | "github.com/stretchr/testify/assert"
15 | "github.com/stretchr/testify/require"
16 | "go.opentelemetry.io/otel"
17 | "go.opentelemetry.io/otel/attribute"
18 | "go.opentelemetry.io/otel/codes"
19 | sdktrace "go.opentelemetry.io/otel/sdk/trace"
20 | "go.opentelemetry.io/otel/sdk/trace/tracetest"
21 | )
22 |
23 | func TestExtractIncidentClientFromEnv(t *testing.T) {
24 | t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com/")
25 | ctx := ExtractIncidentClientFromEnv(context.Background())
26 |
27 | client := IncidentClientFromContext(ctx)
28 | require.NotNil(t, client)
29 | assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
30 | }
31 |
32 | func TestExtractIncidentClientFromHeaders(t *testing.T) {
33 | t.Run("no headers, no env", func(t *testing.T) {
34 | req, err := http.NewRequest("GET", "http://example.com", nil)
35 | require.NoError(t, err)
36 | ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
37 |
38 | client := IncidentClientFromContext(ctx)
39 | require.NotNil(t, client)
40 | assert.Equal(t, "http://localhost:3000/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
41 | })
42 |
43 | t.Run("no headers, with env", func(t *testing.T) {
44 | t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com/")
45 | req, err := http.NewRequest("GET", "http://example.com", nil)
46 | require.NoError(t, err)
47 | ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
48 |
49 | client := IncidentClientFromContext(ctx)
50 | require.NotNil(t, client)
51 | assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
52 | })
53 |
54 | t.Run("with headers, no env", func(t *testing.T) {
55 | req, err := http.NewRequest("GET", "http://example.com", nil)
56 | req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
57 | require.NoError(t, err)
58 | ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
59 |
60 | client := IncidentClientFromContext(ctx)
61 | require.NotNil(t, client)
62 | assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
63 | })
64 |
65 | t.Run("with headers, with env", func(t *testing.T) {
66 | t.Setenv("GRAFANA_URL", "will-not-be-used")
67 | req, err := http.NewRequest("GET", "http://example.com", nil)
68 | req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
69 | require.NoError(t, err)
70 | ctx := ExtractIncidentClientFromHeaders(context.Background(), req)
71 |
72 | client := IncidentClientFromContext(ctx)
73 | require.NotNil(t, client)
74 | assert.Equal(t, "http://my-test-url.grafana.com/api/plugins/grafana-irm-app/resources/api/v1/", client.RemoteHost)
75 | })
76 | }
77 |
78 | func TestExtractGrafanaInfoFromHeaders(t *testing.T) {
79 | t.Run("no headers, no env", func(t *testing.T) {
80 | // Explicitly clear environment variables to ensure test isolation
81 | t.Setenv("GRAFANA_URL", "")
82 | t.Setenv("GRAFANA_API_KEY", "")
83 | t.Setenv("GRAFANA_SERVICE_ACCOUNT_TOKEN", "")
84 |
85 | req, err := http.NewRequest("GET", "http://example.com", nil)
86 | require.NoError(t, err)
87 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
88 | config := GrafanaConfigFromContext(ctx)
89 | assert.Equal(t, defaultGrafanaURL, config.URL)
90 | assert.Equal(t, "", config.APIKey)
91 | assert.Nil(t, config.BasicAuth)
92 | })
93 |
94 | t.Run("no headers, with env", func(t *testing.T) {
95 | t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com")
96 | t.Setenv("GRAFANA_API_KEY", "my-test-api-key")
97 |
98 | req, err := http.NewRequest("GET", "http://example.com", nil)
99 | require.NoError(t, err)
100 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
101 | config := GrafanaConfigFromContext(ctx)
102 | assert.Equal(t, "http://my-test-url.grafana.com", config.URL)
103 | assert.Equal(t, "my-test-api-key", config.APIKey)
104 | })
105 |
106 | t.Run("no headers, with service account token", func(t *testing.T) {
107 | t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com")
108 | t.Setenv("GRAFANA_SERVICE_ACCOUNT_TOKEN", "my-service-account-token")
109 |
110 | req, err := http.NewRequest("GET", "http://example.com", nil)
111 | require.NoError(t, err)
112 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
113 | config := GrafanaConfigFromContext(ctx)
114 | assert.Equal(t, "http://my-test-url.grafana.com", config.URL)
115 | assert.Equal(t, "my-service-account-token", config.APIKey)
116 | })
117 |
118 | t.Run("no headers, service account token takes precedence over api key", func(t *testing.T) {
119 | t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com")
120 | t.Setenv("GRAFANA_API_KEY", "my-deprecated-api-key")
121 | t.Setenv("GRAFANA_SERVICE_ACCOUNT_TOKEN", "my-service-account-token")
122 |
123 | req, err := http.NewRequest("GET", "http://example.com", nil)
124 | require.NoError(t, err)
125 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
126 | config := GrafanaConfigFromContext(ctx)
127 | assert.Equal(t, "http://my-test-url.grafana.com", config.URL)
128 | assert.Equal(t, "my-service-account-token", config.APIKey)
129 | })
130 |
131 | t.Run("with headers, no env", func(t *testing.T) {
132 | req, err := http.NewRequest("GET", "http://example.com", nil)
133 | require.NoError(t, err)
134 | req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
135 | req.Header.Set(grafanaAPIKeyHeader, "my-test-api-key")
136 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
137 | config := GrafanaConfigFromContext(ctx)
138 | assert.Equal(t, "http://my-test-url.grafana.com", config.URL)
139 | assert.Equal(t, "my-test-api-key", config.APIKey)
140 | })
141 |
142 | t.Run("with headers, with env", func(t *testing.T) {
143 | // Env vars should be ignored if headers are present.
144 | t.Setenv("GRAFANA_URL", "will-not-be-used")
145 | t.Setenv("GRAFANA_API_KEY", "will-not-be-used")
146 | t.Setenv("GRAFANA_SERVICE_ACCOUNT_TOKEN", "will-not-be-used")
147 |
148 | req, err := http.NewRequest("GET", "http://example.com", nil)
149 | require.NoError(t, err)
150 | req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
151 | req.Header.Set(grafanaAPIKeyHeader, "my-test-api-key")
152 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
153 | config := GrafanaConfigFromContext(ctx)
154 | assert.Equal(t, "http://my-test-url.grafana.com", config.URL)
155 | assert.Equal(t, "my-test-api-key", config.APIKey)
156 | })
157 |
158 | t.Run("no headers, with env", func(t *testing.T) {
159 | t.Setenv("GRAFANA_USERNAME", "foo")
160 | t.Setenv("GRAFANA_PASSWORD", "bar")
161 |
162 | req, err := http.NewRequest("GET", "http://example.com", nil)
163 | require.NoError(t, err)
164 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
165 | config := GrafanaConfigFromContext(ctx)
166 | assert.Equal(t, "foo", config.BasicAuth.Username())
167 | password, _ := config.BasicAuth.Password()
168 | assert.Equal(t, "bar", password)
169 | })
170 |
171 | t.Run("user auth with headers, no env", func(t *testing.T) {
172 | req, err := http.NewRequest("GET", "http://example.com", nil)
173 | req.SetBasicAuth("foo", "bar")
174 | require.NoError(t, err)
175 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
176 | config := GrafanaConfigFromContext(ctx)
177 | assert.Equal(t, "foo", config.BasicAuth.Username())
178 | password, _ := config.BasicAuth.Password()
179 | assert.Equal(t, "bar", password)
180 | })
181 |
182 | t.Run("user auth with headers, with env", func(t *testing.T) {
183 | t.Setenv("GRAFANA_USERNAME", "will-not-be-used")
184 | t.Setenv("GRAFANA_PASSWORD", "will-not-be-used")
185 |
186 | req, err := http.NewRequest("GET", "http://example.com", nil)
187 | req.SetBasicAuth("foo", "bar")
188 | require.NoError(t, err)
189 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
190 | config := GrafanaConfigFromContext(ctx)
191 | assert.Equal(t, "foo", config.BasicAuth.Username())
192 | password, _ := config.BasicAuth.Password()
193 | assert.Equal(t, "bar", password)
194 | })
195 |
196 | t.Run("orgID from env", func(t *testing.T) {
197 | t.Setenv("GRAFANA_ORG_ID", "123")
198 |
199 | req, err := http.NewRequest("GET", "http://example.com", nil)
200 | require.NoError(t, err)
201 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
202 | config := GrafanaConfigFromContext(ctx)
203 | assert.Equal(t, int64(123), config.OrgID)
204 | })
205 |
206 | t.Run("orgID from header", func(t *testing.T) {
207 | req, err := http.NewRequest("GET", "http://example.com", nil)
208 | require.NoError(t, err)
209 | req.Header.Set("X-Grafana-Org-Id", "456")
210 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
211 | config := GrafanaConfigFromContext(ctx)
212 | assert.Equal(t, int64(456), config.OrgID)
213 | })
214 |
215 | t.Run("orgID header takes precedence over env", func(t *testing.T) {
216 | t.Setenv("GRAFANA_ORG_ID", "123")
217 |
218 | req, err := http.NewRequest("GET", "http://example.com", nil)
219 | require.NoError(t, err)
220 | req.Header.Set("X-Grafana-Org-Id", "456")
221 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
222 | config := GrafanaConfigFromContext(ctx)
223 | assert.Equal(t, int64(456), config.OrgID)
224 | })
225 |
226 | t.Run("invalid orgID from env ignored", func(t *testing.T) {
227 | t.Setenv("GRAFANA_ORG_ID", "not-a-number")
228 |
229 | req, err := http.NewRequest("GET", "http://example.com", nil)
230 | require.NoError(t, err)
231 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
232 | config := GrafanaConfigFromContext(ctx)
233 | assert.Equal(t, int64(0), config.OrgID)
234 | })
235 |
236 | t.Run("invalid orgID from header ignored", func(t *testing.T) {
237 | req, err := http.NewRequest("GET", "http://example.com", nil)
238 | require.NoError(t, err)
239 | req.Header.Set("X-Grafana-Org-Id", "invalid")
240 | ctx := ExtractGrafanaInfoFromHeaders(context.Background(), req)
241 | config := GrafanaConfigFromContext(ctx)
242 | assert.Equal(t, int64(0), config.OrgID)
243 | })
244 | }
245 |
246 | func TestExtractGrafanaClientPath(t *testing.T) {
247 | t.Run("no custom path", func(t *testing.T) {
248 | t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com/")
249 | ctx := ExtractGrafanaClientFromEnv(context.Background())
250 |
251 | c := GrafanaClientFromContext(ctx)
252 | require.NotNil(t, c)
253 | rt := c.Transport.(*client.Runtime)
254 | assert.Equal(t, "/api", rt.BasePath)
255 | })
256 |
257 | t.Run("custom path", func(t *testing.T) {
258 | t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com/grafana")
259 | ctx := ExtractGrafanaClientFromEnv(context.Background())
260 |
261 | c := GrafanaClientFromContext(ctx)
262 | require.NotNil(t, c)
263 | rt := c.Transport.(*client.Runtime)
264 | assert.Equal(t, "/grafana/api", rt.BasePath)
265 | })
266 |
267 | t.Run("custom path, trailing slash", func(t *testing.T) {
268 | t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com/grafana/")
269 | ctx := ExtractGrafanaClientFromEnv(context.Background())
270 |
271 | c := GrafanaClientFromContext(ctx)
272 | require.NotNil(t, c)
273 | rt := c.Transport.(*client.Runtime)
274 | assert.Equal(t, "/grafana/api", rt.BasePath)
275 | })
276 | }
277 |
278 | // minURL is a helper struct representing what we can extract from a constructed
279 | // Grafana client.
280 | type minURL struct {
281 | host, basePath string
282 | }
283 |
284 | // minURLFromClient extracts some minimal amount of URL info from a Grafana client.
285 | func minURLFromClient(c *grafana_client.GrafanaHTTPAPI) minURL {
286 | rt := c.Transport.(*client.Runtime)
287 | return minURL{rt.Host, rt.BasePath}
288 | }
289 |
290 | func TestExtractGrafanaClientFromHeaders(t *testing.T) {
291 | t.Run("no headers, no env", func(t *testing.T) {
292 | req, err := http.NewRequest("GET", "http://example.com", nil)
293 | require.NoError(t, err)
294 | ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
295 | c := GrafanaClientFromContext(ctx)
296 | url := minURLFromClient(c)
297 | assert.Equal(t, "localhost:3000", url.host)
298 | assert.Equal(t, "/api", url.basePath)
299 | })
300 |
301 | t.Run("no headers, with env", func(t *testing.T) {
302 | t.Setenv("GRAFANA_URL", "http://my-test-url.grafana.com")
303 |
304 | req, err := http.NewRequest("GET", "http://example.com", nil)
305 | require.NoError(t, err)
306 | ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
307 | c := GrafanaClientFromContext(ctx)
308 | url := minURLFromClient(c)
309 | assert.Equal(t, "my-test-url.grafana.com", url.host)
310 | assert.Equal(t, "/api", url.basePath)
311 | })
312 |
313 | t.Run("with headers, no env", func(t *testing.T) {
314 | req, err := http.NewRequest("GET", "http://example.com", nil)
315 | require.NoError(t, err)
316 | req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
317 | ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
318 | c := GrafanaClientFromContext(ctx)
319 | url := minURLFromClient(c)
320 | assert.Equal(t, "my-test-url.grafana.com", url.host)
321 | assert.Equal(t, "/api", url.basePath)
322 | })
323 |
324 | t.Run("with headers, with env", func(t *testing.T) {
325 | // Env vars should be ignored if headers are present.
326 | t.Setenv("GRAFANA_URL", "will-not-be-used")
327 |
328 | req, err := http.NewRequest("GET", "http://example.com", nil)
329 | require.NoError(t, err)
330 | req.Header.Set(grafanaURLHeader, "http://my-test-url.grafana.com")
331 | ctx := ExtractGrafanaClientFromHeaders(context.Background(), req)
332 | c := GrafanaClientFromContext(ctx)
333 | url := minURLFromClient(c)
334 | assert.Equal(t, "my-test-url.grafana.com", url.host)
335 | assert.Equal(t, "/api", url.basePath)
336 | })
337 | }
338 |
339 | func TestToolTracingInstrumentation(t *testing.T) {
340 | // Set up in-memory span recorder
341 | spanRecorder := tracetest.NewSpanRecorder()
342 | tracerProvider := sdktrace.NewTracerProvider(
343 | sdktrace.WithSpanProcessor(spanRecorder),
344 | )
345 | originalProvider := otel.GetTracerProvider()
346 | otel.SetTracerProvider(tracerProvider)
347 | defer otel.SetTracerProvider(originalProvider) // Restore original provider
348 |
349 | t.Run("successful tool execution creates span with correct attributes", func(t *testing.T) {
350 | // Clear any previous spans
351 | spanRecorder.Reset()
352 |
353 | // Define a simple test tool
354 | type TestParams struct {
355 | Message string `json:"message" jsonschema:"description=Test message"`
356 | }
357 |
358 | testHandler := func(ctx context.Context, args TestParams) (string, error) {
359 | return "Hello " + args.Message, nil
360 | }
361 |
362 | // Create tool using MustTool (this applies our instrumentation)
363 | tool := MustTool("test_tool", "A test tool for tracing", testHandler)
364 |
365 | // Create context with argument logging enabled
366 | config := GrafanaConfig{
367 | IncludeArgumentsInSpans: true,
368 | }
369 | ctx := WithGrafanaConfig(context.Background(), config)
370 |
371 | // Create a mock MCP request
372 | request := mcp.CallToolRequest{
373 | Params: struct {
374 | Name string `json:"name"`
375 | Arguments any `json:"arguments,omitempty"`
376 | Meta *mcp.Meta `json:"_meta,omitempty"`
377 | }{
378 | Name: "test_tool",
379 | Arguments: map[string]interface{}{
380 | "message": "world",
381 | },
382 | },
383 | }
384 |
385 | // Execute the tool
386 | result, err := tool.Handler(ctx, request)
387 | require.NoError(t, err)
388 | require.NotNil(t, result)
389 |
390 | // Verify span was created
391 | spans := spanRecorder.Ended()
392 | require.Len(t, spans, 1)
393 |
394 | span := spans[0]
395 | assert.Equal(t, "mcp.tool.test_tool", span.Name())
396 | assert.Equal(t, codes.Ok, span.Status().Code)
397 |
398 | // Check attributes
399 | attributes := span.Attributes()
400 | assertHasAttribute(t, attributes, "mcp.tool.name", "test_tool")
401 | assertHasAttribute(t, attributes, "mcp.tool.description", "A test tool for tracing")
402 | assertHasAttribute(t, attributes, "mcp.tool.arguments", `{"message":"world"}`)
403 | })
404 |
405 | t.Run("tool execution error records error on span", func(t *testing.T) {
406 | // Clear any previous spans
407 | spanRecorder.Reset()
408 |
409 | // Define a test tool that returns an error
410 | type TestParams struct {
411 | ShouldFail bool `json:"shouldFail" jsonschema:"description=Whether to fail"`
412 | }
413 |
414 | testHandler := func(ctx context.Context, args TestParams) (string, error) {
415 | if args.ShouldFail {
416 | return "", assert.AnError
417 | }
418 | return "success", nil
419 | }
420 |
421 | // Create tool
422 | tool := MustTool("failing_tool", "A tool that can fail", testHandler)
423 |
424 | // Create context (spans always created)
425 | config := GrafanaConfig{}
426 | ctx := WithGrafanaConfig(context.Background(), config)
427 |
428 | // Create a mock MCP request that will cause failure
429 | request := mcp.CallToolRequest{
430 | Params: struct {
431 | Name string `json:"name"`
432 | Arguments any `json:"arguments,omitempty"`
433 | Meta *mcp.Meta `json:"_meta,omitempty"`
434 | }{
435 | Name: "failing_tool",
436 | Arguments: map[string]interface{}{
437 | "shouldFail": true,
438 | },
439 | },
440 | }
441 |
442 | // Execute the tool (should fail)
443 | result, err := tool.Handler(ctx, request)
444 | assert.Error(t, err)
445 | assert.Nil(t, result)
446 |
447 | // Verify span was created and marked as error
448 | spans := spanRecorder.Ended()
449 | require.Len(t, spans, 1)
450 |
451 | span := spans[0]
452 | assert.Equal(t, "mcp.tool.failing_tool", span.Name())
453 | assert.Equal(t, codes.Error, span.Status().Code)
454 | assert.Equal(t, assert.AnError.Error(), span.Status().Description)
455 |
456 | // Verify error was recorded (check events for error record)
457 | events := span.Events()
458 | hasErrorEvent := false
459 | for _, event := range events {
460 | if event.Name == "exception" {
461 | hasErrorEvent = true
462 | break
463 | }
464 | }
465 | assert.True(t, hasErrorEvent, "Expected error event to be recorded on span")
466 | })
467 |
468 | t.Run("spans always created for context propagation", func(t *testing.T) {
469 | // Clear any previous spans
470 | spanRecorder.Reset()
471 |
472 | // Define a simple test tool
473 | type TestParams struct {
474 | Message string `json:"message" jsonschema:"description=Test message"`
475 | }
476 |
477 | testHandler := func(ctx context.Context, args TestParams) (string, error) {
478 | return "processed", nil
479 | }
480 |
481 | // Create tool
482 | tool := MustTool("context_prop_tool", "A tool for context propagation", testHandler)
483 |
484 | // Create context with default config (no special flags)
485 | config := GrafanaConfig{}
486 | ctx := WithGrafanaConfig(context.Background(), config)
487 |
488 | // Create a mock MCP request
489 | request := mcp.CallToolRequest{
490 | Params: struct {
491 | Name string `json:"name"`
492 | Arguments any `json:"arguments,omitempty"`
493 | Meta *mcp.Meta `json:"_meta,omitempty"`
494 | }{
495 | Name: "context_prop_tool",
496 | Arguments: map[string]interface{}{
497 | "message": "test",
498 | },
499 | },
500 | }
501 |
502 | // Execute the tool (should always create spans for context propagation)
503 | result, err := tool.Handler(ctx, request)
504 | require.NoError(t, err)
505 | require.NotNil(t, result)
506 |
507 | // Verify spans ARE always created
508 | spans := spanRecorder.Ended()
509 | require.Len(t, spans, 1)
510 |
511 | span := spans[0]
512 | assert.Equal(t, "mcp.tool.context_prop_tool", span.Name())
513 | assert.Equal(t, codes.Ok, span.Status().Code)
514 | })
515 |
516 | t.Run("arguments not logged by default (PII safety)", func(t *testing.T) {
517 | // Clear any previous spans
518 | spanRecorder.Reset()
519 |
520 | // Define a simple test tool
521 | type TestParams struct {
522 | SensitiveData string `json:"sensitiveData" jsonschema:"description=Potentially sensitive data"`
523 | }
524 |
525 | testHandler := func(ctx context.Context, args TestParams) (string, error) {
526 | return "processed", nil
527 | }
528 |
529 | // Create tool
530 | tool := MustTool("sensitive_tool", "A tool with sensitive data", testHandler)
531 |
532 | // Create context with argument logging disabled (default)
533 | config := GrafanaConfig{
534 | IncludeArgumentsInSpans: false, // Default: safe
535 | }
536 | ctx := WithGrafanaConfig(context.Background(), config)
537 |
538 | // Create a mock MCP request with potentially sensitive data
539 | request := mcp.CallToolRequest{
540 | Params: struct {
541 | Name string `json:"name"`
542 | Arguments any `json:"arguments,omitempty"`
543 | Meta *mcp.Meta `json:"_meta,omitempty"`
544 | }{
545 | Name: "sensitive_tool",
546 | Arguments: map[string]interface{}{
547 | "sensitiveData": "[email protected]",
548 | },
549 | },
550 | }
551 |
552 | // Execute the tool (arguments should NOT be logged by default)
553 | result, err := tool.Handler(ctx, request)
554 | require.NoError(t, err)
555 | require.NotNil(t, result)
556 |
557 | // Verify span was created
558 | spans := spanRecorder.Ended()
559 | require.Len(t, spans, 1)
560 |
561 | span := spans[0]
562 | assert.Equal(t, "mcp.tool.sensitive_tool", span.Name())
563 | assert.Equal(t, codes.Ok, span.Status().Code)
564 |
565 | // Check that arguments are NOT logged (PII safety)
566 | attributes := span.Attributes()
567 | assertHasAttribute(t, attributes, "mcp.tool.name", "sensitive_tool")
568 | assertHasAttribute(t, attributes, "mcp.tool.description", "A tool with sensitive data")
569 |
570 | // Verify arguments are NOT present
571 | for _, attr := range attributes {
572 | assert.NotEqual(t, "mcp.tool.arguments", string(attr.Key), "Arguments should not be logged by default for PII safety")
573 | }
574 | })
575 |
576 | t.Run("arguments logged when argument logging enabled", func(t *testing.T) {
577 | // Clear any previous spans
578 | spanRecorder.Reset()
579 |
580 | // Define a simple test tool
581 | type TestParams struct {
582 | SafeData string `json:"safeData" jsonschema:"description=Non-sensitive data"`
583 | }
584 |
585 | testHandler := func(ctx context.Context, args TestParams) (string, error) {
586 | return "processed", nil
587 | }
588 |
589 | // Create tool
590 | tool := MustTool("debug_tool", "A tool for debugging", testHandler)
591 |
592 | // Create context with argument logging enabled
593 | config := GrafanaConfig{
594 | IncludeArgumentsInSpans: true,
595 | }
596 | ctx := WithGrafanaConfig(context.Background(), config)
597 |
598 | // Create a mock MCP request
599 | request := mcp.CallToolRequest{
600 | Params: struct {
601 | Name string `json:"name"`
602 | Arguments any `json:"arguments,omitempty"`
603 | Meta *mcp.Meta `json:"_meta,omitempty"`
604 | }{
605 | Name: "debug_tool",
606 | Arguments: map[string]interface{}{
607 | "safeData": "debug-value",
608 | },
609 | },
610 | }
611 |
612 | // Execute the tool (arguments SHOULD be logged when flag enabled)
613 | result, err := tool.Handler(ctx, request)
614 | require.NoError(t, err)
615 | require.NotNil(t, result)
616 |
617 | // Verify span was created
618 | spans := spanRecorder.Ended()
619 | require.Len(t, spans, 1)
620 |
621 | span := spans[0]
622 | assert.Equal(t, "mcp.tool.debug_tool", span.Name())
623 | assert.Equal(t, codes.Ok, span.Status().Code)
624 |
625 | // Check that arguments ARE logged when flag enabled
626 | attributes := span.Attributes()
627 | assertHasAttribute(t, attributes, "mcp.tool.name", "debug_tool")
628 | assertHasAttribute(t, attributes, "mcp.tool.description", "A tool for debugging")
629 | assertHasAttribute(t, attributes, "mcp.tool.arguments", `{"safeData":"debug-value"}`)
630 | })
631 | }
632 |
633 | func TestHTTPTracingConfiguration(t *testing.T) {
634 | t.Run("HTTP tracing always enabled for context propagation", func(t *testing.T) {
635 | // Create context (HTTP tracing always enabled)
636 | config := GrafanaConfig{}
637 | ctx := WithGrafanaConfig(context.Background(), config)
638 |
639 | // Create Grafana client
640 | client := NewGrafanaClient(ctx, "http://localhost:3000", "test-api-key", nil, 0)
641 | require.NotNil(t, client)
642 |
643 | // Verify the client was created successfully (should not panic)
644 | assert.NotNil(t, client.Transport)
645 | })
646 |
647 | t.Run("tracing works gracefully without OpenTelemetry configured", func(t *testing.T) {
648 | // No OpenTelemetry tracer provider configured
649 |
650 | // Create context (tracing always enabled for context propagation)
651 | config := GrafanaConfig{}
652 | ctx := WithGrafanaConfig(context.Background(), config)
653 |
654 | // Create Grafana client (should not panic even without OTEL configured)
655 | client := NewGrafanaClient(ctx, "http://localhost:3000", "test-api-key", nil, 0)
656 | require.NotNil(t, client)
657 |
658 | // Verify the client was created successfully
659 | assert.NotNil(t, client.Transport)
660 | })
661 | }
662 |
663 | // Helper function to check if an attribute exists with expected value
664 | func assertHasAttribute(t *testing.T, attributes []attribute.KeyValue, key string, expectedValue string) {
665 | for _, attr := range attributes {
666 | if string(attr.Key) == key {
667 | assert.Equal(t, expectedValue, attr.Value.AsString())
668 | return
669 | }
670 | }
671 | t.Errorf("Expected attribute %s with value %s not found", key, expectedValue)
672 | }
673 |
```
--------------------------------------------------------------------------------
/mcpgrafana.go:
--------------------------------------------------------------------------------
```go
1 | package mcpgrafana
2 |
3 | import (
4 | "context"
5 | "crypto/tls"
6 | "crypto/x509"
7 | "fmt"
8 | "log/slog"
9 | "net/http"
10 | "net/url"
11 | "os"
12 | "reflect"
13 | "runtime/debug"
14 | "strconv"
15 | "strings"
16 | "sync"
17 |
18 | "github.com/go-openapi/strfmt"
19 | "github.com/grafana/grafana-openapi-client-go/client"
20 | "github.com/grafana/incident-go"
21 | "github.com/mark3labs/mcp-go/server"
22 | "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
23 | )
24 |
25 | const (
26 | defaultGrafanaHost = "localhost:3000"
27 | defaultGrafanaURL = "http://" + defaultGrafanaHost
28 |
29 | grafanaURLEnvVar = "GRAFANA_URL"
30 | grafanaServiceAccountTokenEnvVar = "GRAFANA_SERVICE_ACCOUNT_TOKEN"
31 | grafanaAPIEnvVar = "GRAFANA_API_KEY" // Deprecated: use GRAFANA_SERVICE_ACCOUNT_TOKEN instead
32 | grafanaOrgIDEnvVar = "GRAFANA_ORG_ID"
33 |
34 | grafanaUsernameEnvVar = "GRAFANA_USERNAME"
35 | grafanaPasswordEnvVar = "GRAFANA_PASSWORD"
36 |
37 | grafanaURLHeader = "X-Grafana-URL"
38 | grafanaAPIKeyHeader = "X-Grafana-API-Key"
39 | )
40 |
41 | func urlAndAPIKeyFromEnv() (string, string) {
42 | u := strings.TrimRight(os.Getenv(grafanaURLEnvVar), "/")
43 |
44 | // Check for the new service account token environment variable first
45 | apiKey := os.Getenv(grafanaServiceAccountTokenEnvVar)
46 | if apiKey != "" {
47 | return u, apiKey
48 | }
49 |
50 | // Fall back to the deprecated API key environment variable
51 | apiKey = os.Getenv(grafanaAPIEnvVar)
52 | if apiKey != "" {
53 | slog.Warn("GRAFANA_API_KEY is deprecated, please use GRAFANA_SERVICE_ACCOUNT_TOKEN instead. See https://grafana.com/docs/grafana/latest/administration/service-accounts/#add-a-token-to-a-service-account-in-grafana for details on creating service account tokens.")
54 | }
55 |
56 | return u, apiKey
57 | }
58 |
59 | func userAndPassFromEnv() *url.Userinfo {
60 | username := os.Getenv(grafanaUsernameEnvVar)
61 | password, exists := os.LookupEnv(grafanaPasswordEnvVar)
62 | if username == "" && password == "" {
63 | return nil
64 | }
65 | if !exists {
66 | return url.User(username)
67 | }
68 | return url.UserPassword(username, password)
69 | }
70 |
71 | func orgIdFromEnv() int64 {
72 | orgIDStr := os.Getenv(grafanaOrgIDEnvVar)
73 | if orgIDStr == "" {
74 | return 0
75 | }
76 | orgID, err := strconv.ParseInt(orgIDStr, 10, 64)
77 | if err != nil {
78 | slog.Warn("Invalid GRAFANA_ORG_ID value, ignoring", "value", orgIDStr, "error", err)
79 | return 0
80 | }
81 | return orgID
82 | }
83 |
84 | func orgIdFromHeaders(req *http.Request) int64 {
85 | orgIDStr := req.Header.Get(client.OrgIDHeader)
86 | if orgIDStr == "" {
87 | return 0
88 | }
89 | orgID, err := strconv.ParseInt(orgIDStr, 10, 64)
90 | if err != nil {
91 | slog.Warn("Invalid X-Grafana-Org-Id header value, ignoring", "value", orgIDStr, "error", err)
92 | return 0
93 | }
94 | return orgID
95 | }
96 |
97 | func urlAndAPIKeyFromHeaders(req *http.Request) (string, string) {
98 | u := strings.TrimRight(req.Header.Get(grafanaURLHeader), "/")
99 | apiKey := req.Header.Get(grafanaAPIKeyHeader)
100 | return u, apiKey
101 | }
102 |
103 | // grafanaConfigKey is the context key for Grafana configuration.
104 | type grafanaConfigKey struct{}
105 |
106 | // TLSConfig holds TLS configuration for Grafana clients.
107 | // It supports mutual TLS authentication with client certificates, custom CA certificates for server verification, and development options like skipping certificate verification.
108 | type TLSConfig struct {
109 | CertFile string
110 | KeyFile string
111 | CAFile string
112 | SkipVerify bool
113 | }
114 |
115 | // GrafanaConfig represents the full configuration for Grafana clients.
116 | // It includes connection details, authentication credentials, debug settings, and TLS options used throughout the MCP server's lifecycle.
117 | type GrafanaConfig struct {
118 | // Debug enables debug mode for the Grafana client.
119 | Debug bool
120 |
121 | // IncludeArgumentsInSpans enables logging of tool arguments in OpenTelemetry spans.
122 | // This should only be enabled in non-production environments or when you're certain
123 | // the arguments don't contain PII. Defaults to false for safety.
124 | // Note: OpenTelemetry spans are always created for context propagation, but arguments
125 | // are only included when this flag is enabled.
126 | IncludeArgumentsInSpans bool
127 |
128 | // URL is the URL of the Grafana instance.
129 | URL string
130 |
131 | // APIKey is the API key or service account token for the Grafana instance.
132 | // It may be empty if we are using on-behalf-of auth.
133 | APIKey string
134 |
135 | // Credentials if user is using basic auth
136 | BasicAuth *url.Userinfo
137 |
138 | // OrgID is the organization ID to use for multi-org support.
139 | // When set, it will be sent as X-Scope-OrgId header regardless of authentication method.
140 | // Works with service account tokens, API keys, and basic authentication.
141 | OrgID int64
142 |
143 | // AccessToken is the Grafana Cloud access policy token used for on-behalf-of auth in Grafana Cloud.
144 | AccessToken string
145 | // IDToken is an ID token identifying the user for the current request.
146 | // It comes from the `X-Grafana-Id` header sent from Grafana to plugin backends.
147 | // It is used for on-behalf-of auth in Grafana Cloud.
148 | IDToken string
149 |
150 | // TLSConfig holds TLS configuration for all Grafana clients.
151 | TLSConfig *TLSConfig
152 | }
153 |
154 | // WithGrafanaConfig adds Grafana configuration to the context.
155 | // This configuration includes API credentials, debug settings, and TLS options that will be used by all Grafana clients created from this context.
156 | func WithGrafanaConfig(ctx context.Context, config GrafanaConfig) context.Context {
157 | return context.WithValue(ctx, grafanaConfigKey{}, config)
158 | }
159 |
160 | // GrafanaConfigFromContext extracts Grafana configuration from the context.
161 | // If no config is found, returns a zero-value GrafanaConfig. This function is typically used by internal components to access configuration set earlier in the request lifecycle.
162 | func GrafanaConfigFromContext(ctx context.Context) GrafanaConfig {
163 | if config, ok := ctx.Value(grafanaConfigKey{}).(GrafanaConfig); ok {
164 | return config
165 | }
166 | return GrafanaConfig{}
167 | }
168 |
169 | // CreateTLSConfig creates a *tls.Config from TLSConfig.
170 | // It supports client certificates, custom CA certificates, and the option to skip TLS verification for development environments.
171 | func (tc *TLSConfig) CreateTLSConfig() (*tls.Config, error) {
172 | if tc == nil {
173 | return nil, nil
174 | }
175 |
176 | tlsConfig := &tls.Config{
177 | InsecureSkipVerify: tc.SkipVerify,
178 | }
179 |
180 | // Load client certificate if both cert and key files are provided
181 | if tc.CertFile != "" && tc.KeyFile != "" {
182 | cert, err := tls.LoadX509KeyPair(tc.CertFile, tc.KeyFile)
183 | if err != nil {
184 | return nil, fmt.Errorf("failed to load client certificate: %w", err)
185 | }
186 | tlsConfig.Certificates = []tls.Certificate{cert}
187 | }
188 |
189 | // Load CA certificate if provided
190 | if tc.CAFile != "" {
191 | caCert, err := os.ReadFile(tc.CAFile)
192 | if err != nil {
193 | return nil, fmt.Errorf("failed to read CA certificate: %w", err)
194 | }
195 | caCertPool := x509.NewCertPool()
196 | if !caCertPool.AppendCertsFromPEM(caCert) {
197 | return nil, fmt.Errorf("failed to parse CA certificate")
198 | }
199 | tlsConfig.RootCAs = caCertPool
200 | }
201 |
202 | return tlsConfig, nil
203 | }
204 |
205 | // HTTPTransport creates an HTTP transport with custom TLS configuration.
206 | // It clones the provided transport and applies the TLS settings, preserving other transport configurations like timeouts and connection pools.
207 | func (tc *TLSConfig) HTTPTransport(defaultTransport *http.Transport) (http.RoundTripper, error) {
208 | transport := defaultTransport.Clone()
209 |
210 | if tc != nil {
211 | tlsCfg, err := tc.CreateTLSConfig()
212 | if err != nil {
213 | return nil, err
214 | }
215 | transport.TLSClientConfig = tlsCfg
216 | }
217 |
218 | return transport, nil
219 | }
220 |
221 | // UserAgentTransport wraps an http.RoundTripper to add a custom User-Agent header.
222 | // This ensures all HTTP requests from the MCP server are properly identified with version information for debugging and analytics.
223 | type UserAgentTransport struct {
224 | rt http.RoundTripper
225 | UserAgent string
226 | }
227 |
228 | func (t *UserAgentTransport) RoundTrip(req *http.Request) (*http.Response, error) {
229 | // Clone the request to avoid modifying the original
230 | clonedReq := req.Clone(req.Context())
231 |
232 | // Add or update the User-Agent header
233 | if clonedReq.Header.Get("User-Agent") == "" {
234 | clonedReq.Header.Set("User-Agent", t.UserAgent)
235 | }
236 |
237 | return t.rt.RoundTrip(clonedReq)
238 | }
239 |
240 | // Version returns the version of the mcp-grafana binary.
241 | // It uses runtime/debug to fetch version information from the build, returning "(devel)" for local development builds.
242 | // The version is computed once and cached for performance.
243 | var Version = sync.OnceValue(func() string {
244 | // Default version string returned by `runtime/debug` if built
245 | // from the source repository rather than with `go install`.
246 | v := "(devel)"
247 | if bi, ok := debug.ReadBuildInfo(); ok && bi.Main.Version != "" {
248 | v = bi.Main.Version
249 | }
250 | return v
251 | })
252 |
253 | // UserAgent returns the user agent string for HTTP requests.
254 | // It includes the mcp-grafana identifier and version number for proper request attribution and debugging.
255 | func UserAgent() string {
256 | return fmt.Sprintf("mcp-grafana/%s", Version())
257 | }
258 |
259 | // NewUserAgentTransport creates a new UserAgentTransport with the specified user agent.
260 | // If no user agent is provided, it uses the default UserAgent() with version information.
261 | // The transport wraps the provided RoundTripper, defaulting to http.DefaultTransport if nil.
262 | func NewUserAgentTransport(rt http.RoundTripper, userAgent ...string) *UserAgentTransport {
263 | if rt == nil {
264 | rt = http.DefaultTransport
265 | }
266 |
267 | ua := UserAgent() // default
268 | if len(userAgent) > 0 {
269 | ua = userAgent[0]
270 | }
271 |
272 | return &UserAgentTransport{
273 | rt: rt,
274 | UserAgent: ua,
275 | }
276 | }
277 |
278 | // wrapWithUserAgent wraps an http.RoundTripper with user agent tracking
279 | func wrapWithUserAgent(rt http.RoundTripper) http.RoundTripper {
280 | return NewUserAgentTransport(rt)
281 | }
282 |
283 | // OrgIDRoundTripper wraps an http.RoundTripper to add the X-Grafana-Org-Id header.
284 | type OrgIDRoundTripper struct {
285 | underlying http.RoundTripper
286 | orgID int64
287 | }
288 |
289 | func (t *OrgIDRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
290 | // clone the request to avoid modifying the original
291 | clonedReq := req.Clone(req.Context())
292 |
293 | if t.orgID > 0 {
294 | clonedReq.Header.Set(client.OrgIDHeader, strconv.FormatInt(t.orgID, 10))
295 | }
296 |
297 | return t.underlying.RoundTrip(clonedReq)
298 | }
299 |
300 | func NewOrgIDRoundTripper(rt http.RoundTripper, orgID int64) *OrgIDRoundTripper {
301 | if rt == nil {
302 | rt = http.DefaultTransport
303 | }
304 |
305 | return &OrgIDRoundTripper{
306 | underlying: rt,
307 | orgID: orgID,
308 | }
309 | }
310 |
311 | // Gets info from environment
312 | func extractKeyGrafanaInfoFromEnv() (url, apiKey string, auth *url.Userinfo, orgId int64) {
313 | url, apiKey = urlAndAPIKeyFromEnv()
314 | if url == "" {
315 | url = defaultGrafanaURL
316 | }
317 | auth = userAndPassFromEnv()
318 | orgId = orgIdFromEnv()
319 | return
320 | }
321 |
322 | // Tries to get grafana info from a request.
323 | // Gets info from environment if it can't get it from request
324 | func extractKeyGrafanaInfoFromReq(req *http.Request) (grafanaUrl, apiKey string, auth *url.Userinfo, orgId int64) {
325 | eUrl, eApiKey, eAuth, eOrgId := extractKeyGrafanaInfoFromEnv()
326 | username, password, _ := req.BasicAuth()
327 |
328 | grafanaUrl, apiKey = urlAndAPIKeyFromHeaders(req)
329 | // If anything is missing, check if we can get it from the environment
330 | if grafanaUrl == "" {
331 | grafanaUrl = eUrl
332 | }
333 |
334 | if apiKey == "" {
335 | apiKey = eApiKey
336 | }
337 |
338 | // Use environment configured auth if nothing was passed in request
339 | if username == "" && password == "" {
340 | auth = eAuth
341 | } else {
342 | auth = url.UserPassword(username, password)
343 | }
344 |
345 | // extract org ID from header, fall back to environment
346 | orgId = orgIdFromHeaders(req)
347 | if orgId == 0 {
348 | orgId = eOrgId
349 | }
350 |
351 | return
352 | }
353 |
354 | // ExtractGrafanaInfoFromEnv is a StdioContextFunc that extracts Grafana configuration from environment variables.
355 | // It reads GRAFANA_URL and GRAFANA_SERVICE_ACCOUNT_TOKEN (or deprecated GRAFANA_API_KEY) environment variables and adds the configuration to the context for use by Grafana clients.
356 | var ExtractGrafanaInfoFromEnv server.StdioContextFunc = func(ctx context.Context) context.Context {
357 | u, apiKey, basicAuth, orgID := extractKeyGrafanaInfoFromEnv()
358 | parsedURL, err := url.Parse(u)
359 | if err != nil {
360 | panic(fmt.Errorf("invalid Grafana URL %s: %w", u, err))
361 | }
362 |
363 | slog.Info("Using Grafana configuration", "url", parsedURL.Redacted(), "api_key_set", apiKey != "", "basic_auth_set", basicAuth != nil, "org_id", orgID)
364 |
365 | // Get existing config or create a new one.
366 | // This will respect the existing debug flag, if set.
367 | config := GrafanaConfigFromContext(ctx)
368 | config.URL = u
369 | config.APIKey = apiKey
370 | config.BasicAuth = basicAuth
371 | config.OrgID = orgID
372 | return WithGrafanaConfig(ctx, config)
373 | }
374 |
375 | // httpContextFunc is a function that can be used as a `server.HTTPContextFunc` or a
376 | // `server.SSEContextFunc`. It is necessary because, while the two types are functionally
377 | // identical, they have distinct types and cannot be passed around interchangeably.
378 | type httpContextFunc func(ctx context.Context, req *http.Request) context.Context
379 |
380 | // ExtractGrafanaInfoFromHeaders is a HTTPContextFunc that extracts Grafana configuration from HTTP request headers.
381 | // It reads X-Grafana-URL and X-Grafana-API-Key headers, falling back to environment variables if headers are not present.
382 | var ExtractGrafanaInfoFromHeaders httpContextFunc = func(ctx context.Context, req *http.Request) context.Context {
383 | u, apiKey, basicAuth, orgID := extractKeyGrafanaInfoFromReq(req)
384 |
385 | // Get existing config or create a new one.
386 | // This will respect the existing debug flag, if set.
387 | config := GrafanaConfigFromContext(ctx)
388 | config.URL = u
389 | config.APIKey = apiKey
390 | config.BasicAuth = basicAuth
391 | config.OrgID = orgID
392 | return WithGrafanaConfig(ctx, config)
393 | }
394 |
395 | // WithOnBehalfOfAuth adds the Grafana access token and user token to the Grafana config.
396 | // These tokens enable on-behalf-of authentication in Grafana Cloud, allowing the MCP server to act on behalf of a specific user with their permissions.
397 | func WithOnBehalfOfAuth(ctx context.Context, accessToken, userToken string) (context.Context, error) {
398 | if accessToken == "" || userToken == "" {
399 | return nil, fmt.Errorf("neither accessToken nor userToken can be empty")
400 | }
401 | cfg := GrafanaConfigFromContext(ctx)
402 | cfg.AccessToken = accessToken
403 | cfg.IDToken = userToken
404 | return WithGrafanaConfig(ctx, cfg), nil
405 | }
406 |
407 | // MustWithOnBehalfOfAuth adds the access and user tokens to the context, panicking if either are empty.
408 | // This is a convenience wrapper around WithOnBehalfOfAuth for cases where token validation has already occurred.
409 | func MustWithOnBehalfOfAuth(ctx context.Context, accessToken, userToken string) context.Context {
410 | ctx, err := WithOnBehalfOfAuth(ctx, accessToken, userToken)
411 | if err != nil {
412 | panic(err)
413 | }
414 | return ctx
415 | }
416 |
417 | type grafanaClientKey struct{}
418 |
419 | func makeBasePath(path string) string {
420 | return strings.Join([]string{strings.TrimRight(path, "/"), "api"}, "/")
421 | }
422 |
423 | // NewGrafanaClient creates a Grafana client with the provided URL and API key.
424 | // The client is automatically configured with the correct HTTP scheme, debug settings from context, custom TLS configuration if present, and OpenTelemetry instrumentation for distributed tracing.
425 | func NewGrafanaClient(ctx context.Context, grafanaURL, apiKey string, auth *url.Userinfo, orgId int64) *client.GrafanaHTTPAPI {
426 | cfg := client.DefaultTransportConfig()
427 |
428 | var parsedURL *url.URL
429 | var err error
430 |
431 | if grafanaURL == "" {
432 | grafanaURL = defaultGrafanaURL
433 | }
434 |
435 | parsedURL, err = url.Parse(grafanaURL)
436 | if err != nil {
437 | panic(fmt.Errorf("invalid Grafana URL: %w", err))
438 | }
439 | cfg.Host = parsedURL.Host
440 | cfg.BasePath = makeBasePath(parsedURL.Path)
441 |
442 | // The Grafana client will always prefer HTTPS even if the URL is HTTP,
443 | // so we need to limit the schemes to HTTP if the URL is HTTP.
444 | if parsedURL.Scheme == "http" {
445 | cfg.Schemes = []string{"http"}
446 | }
447 |
448 | if apiKey != "" {
449 | cfg.APIKey = apiKey
450 | }
451 |
452 | if auth != nil {
453 | cfg.BasicAuth = auth
454 | }
455 |
456 | config := GrafanaConfigFromContext(ctx)
457 | cfg.Debug = config.Debug
458 |
459 | if config.OrgID > 0 {
460 | cfg.OrgID = config.OrgID
461 | }
462 |
463 | // Configure TLS if custom TLS configuration is provided
464 | if tlsConfig := config.TLSConfig; tlsConfig != nil {
465 | tlsCfg, err := tlsConfig.CreateTLSConfig()
466 | if err != nil {
467 | panic(fmt.Errorf("failed to create TLS config: %w", err))
468 | }
469 | cfg.TLSConfig = tlsCfg
470 | slog.Debug("Using custom TLS configuration",
471 | "cert_file", tlsConfig.CertFile,
472 | "ca_file", tlsConfig.CAFile,
473 | "skip_verify", tlsConfig.SkipVerify)
474 | }
475 |
476 | slog.Debug("Creating Grafana client", "url", parsedURL.Redacted(), "api_key_set", apiKey != "", "basic_auth_set", config.BasicAuth != nil, "org_id", cfg.OrgID)
477 | grafanaClient := client.NewHTTPClientWithConfig(strfmt.Default, cfg)
478 |
479 | // Always enable HTTP tracing for context propagation (no-op when no exporter configured)
480 | // Use reflection to wrap the transport without importing the runtime client package
481 | v := reflect.ValueOf(grafanaClient.Transport)
482 | if v.Kind() == reflect.Ptr && !v.IsNil() {
483 | v = v.Elem()
484 | if v.Kind() == reflect.Struct {
485 | transportField := v.FieldByName("Transport")
486 | if transportField.IsValid() && transportField.CanSet() {
487 | if rt, ok := transportField.Interface().(http.RoundTripper); ok {
488 | // Wrap with user agent first, then otel
489 | userAgentWrapped := wrapWithUserAgent(rt)
490 | wrapped := otelhttp.NewTransport(userAgentWrapped)
491 | transportField.Set(reflect.ValueOf(wrapped))
492 | slog.Debug("HTTP tracing and user agent tracking enabled for Grafana client")
493 | }
494 | }
495 | }
496 | }
497 |
498 | return grafanaClient
499 | }
500 |
501 | // ExtractGrafanaClientFromEnv is a StdioContextFunc that creates and injects a Grafana client into the context.
502 | // It uses configuration from GRAFANA_URL, GRAFANA_SERVICE_ACCOUNT_TOKEN (or deprecated GRAFANA_API_KEY), GRAFANA_USERNAME/PASSWORD environment variables to initialize
503 | // the client with proper authentication.
504 | var ExtractGrafanaClientFromEnv server.StdioContextFunc = func(ctx context.Context) context.Context {
505 | // Extract transport config from env vars
506 | grafanaURL, apiKey := urlAndAPIKeyFromEnv()
507 | if grafanaURL == "" {
508 | grafanaURL = defaultGrafanaURL
509 | }
510 | auth := userAndPassFromEnv()
511 | orgId := orgIdFromEnv()
512 | grafanaClient := NewGrafanaClient(ctx, grafanaURL, apiKey, auth, orgId)
513 | return WithGrafanaClient(ctx, grafanaClient)
514 | }
515 |
516 | // ExtractGrafanaClientFromHeaders is a HTTPContextFunc that creates and injects a Grafana client into the context.
517 | // It prioritizes configuration from HTTP headers (X-Grafana-URL, X-Grafana-API-Key) over environment variables for multi-tenant scenarios.
518 | var ExtractGrafanaClientFromHeaders httpContextFunc = func(ctx context.Context, req *http.Request) context.Context {
519 | // Extract transport config from request headers, and set it on the context.
520 | u, apiKey, basicAuth, orgId := extractKeyGrafanaInfoFromReq(req)
521 | slog.Debug("Creating Grafana client", "url", u, "api_key_set", apiKey != "", "basic_auth_set", basicAuth != nil)
522 |
523 | grafanaClient := NewGrafanaClient(ctx, u, apiKey, basicAuth, orgId)
524 | return WithGrafanaClient(ctx, grafanaClient)
525 | }
526 |
527 | // WithGrafanaClient sets the Grafana client in the context.
528 | // The client can be retrieved using GrafanaClientFromContext and will be used by all Grafana-related tools in the MCP server.
529 | func WithGrafanaClient(ctx context.Context, client *client.GrafanaHTTPAPI) context.Context {
530 | return context.WithValue(ctx, grafanaClientKey{}, client)
531 | }
532 |
533 | // GrafanaClientFromContext retrieves the Grafana client from the context.
534 | // Returns nil if no client has been set, which tools should handle gracefully with appropriate error messages.
535 | func GrafanaClientFromContext(ctx context.Context) *client.GrafanaHTTPAPI {
536 | c, ok := ctx.Value(grafanaClientKey{}).(*client.GrafanaHTTPAPI)
537 | if !ok {
538 | return nil
539 | }
540 | return c
541 | }
542 |
543 | type incidentClientKey struct{}
544 |
545 | // ExtractIncidentClientFromEnv is a StdioContextFunc that creates and injects a Grafana Incident client into the context.
546 | // It configures the client using environment variables and applies any custom TLS settings from the context.
547 | var ExtractIncidentClientFromEnv server.StdioContextFunc = func(ctx context.Context) context.Context {
548 | grafanaURL, apiKey := urlAndAPIKeyFromEnv()
549 | if grafanaURL == "" {
550 | grafanaURL = defaultGrafanaURL
551 | }
552 | incidentURL := fmt.Sprintf("%s/api/plugins/grafana-irm-app/resources/api/v1/", grafanaURL)
553 | parsedURL, err := url.Parse(incidentURL)
554 | if err != nil {
555 | panic(fmt.Errorf("invalid incident URL %s: %w", incidentURL, err))
556 | }
557 | slog.Debug("Creating Incident client", "url", parsedURL.Redacted(), "api_key_set", apiKey != "")
558 | client := incident.NewClient(incidentURL, apiKey)
559 |
560 | config := GrafanaConfigFromContext(ctx)
561 | // Configure custom TLS if available
562 | if tlsConfig := config.TLSConfig; tlsConfig != nil {
563 | transport, err := tlsConfig.HTTPTransport(http.DefaultTransport.(*http.Transport))
564 | if err != nil {
565 | slog.Error("Failed to create custom transport for incident client, using default", "error", err)
566 | } else {
567 | orgIDWrapped := NewOrgIDRoundTripper(transport, config.OrgID)
568 | client.HTTPClient.Transport = wrapWithUserAgent(orgIDWrapped)
569 | slog.Debug("Using custom TLS configuration, user agent, and org ID support for incident client",
570 | "cert_file", tlsConfig.CertFile,
571 | "ca_file", tlsConfig.CAFile,
572 | "skip_verify", tlsConfig.SkipVerify)
573 | }
574 | } else {
575 | // No custom TLS, but still add org ID and user agent
576 | orgIDWrapped := NewOrgIDRoundTripper(http.DefaultTransport, config.OrgID)
577 | client.HTTPClient.Transport = wrapWithUserAgent(orgIDWrapped)
578 | }
579 |
580 | return context.WithValue(ctx, incidentClientKey{}, client)
581 | }
582 |
583 | // ExtractIncidentClientFromHeaders is a HTTPContextFunc that creates and injects a Grafana Incident client into the context.
584 | // It uses HTTP headers for configuration with environment variable fallbacks, enabling per-request incident management configuration.
585 | var ExtractIncidentClientFromHeaders httpContextFunc = func(ctx context.Context, req *http.Request) context.Context {
586 | grafanaURL, apiKey, _, orgID := extractKeyGrafanaInfoFromReq(req)
587 | incidentURL := fmt.Sprintf("%s/api/plugins/grafana-irm-app/resources/api/v1/", grafanaURL)
588 | client := incident.NewClient(incidentURL, apiKey)
589 |
590 | config := GrafanaConfigFromContext(ctx)
591 | // Configure custom TLS if available
592 | if tlsConfig := config.TLSConfig; tlsConfig != nil {
593 | transport, err := tlsConfig.HTTPTransport(http.DefaultTransport.(*http.Transport))
594 | if err != nil {
595 | slog.Error("Failed to create custom transport for incident client, using default", "error", err)
596 | } else {
597 | orgIDWrapped := NewOrgIDRoundTripper(transport, orgID)
598 | client.HTTPClient.Transport = wrapWithUserAgent(orgIDWrapped)
599 | slog.Debug("Using custom TLS configuration, user agent, and org ID support for incident client",
600 | "cert_file", tlsConfig.CertFile,
601 | "ca_file", tlsConfig.CAFile,
602 | "skip_verify", tlsConfig.SkipVerify)
603 | }
604 | } else {
605 | // No custom TLS, but still add org ID and user agent
606 | orgIDWrapped := NewOrgIDRoundTripper(http.DefaultTransport, orgID)
607 | client.HTTPClient.Transport = wrapWithUserAgent(orgIDWrapped)
608 | }
609 |
610 | return context.WithValue(ctx, incidentClientKey{}, client)
611 | }
612 |
613 | // WithIncidentClient sets the Grafana Incident client in the context.
614 | // This client is used for managing incidents, activities, and other IRM (Incident Response Management) operations.
615 | func WithIncidentClient(ctx context.Context, client *incident.Client) context.Context {
616 | return context.WithValue(ctx, incidentClientKey{}, client)
617 | }
618 |
619 | // IncidentClientFromContext retrieves the Grafana Incident client from the context.
620 | // Returns nil if no client has been set, indicating that incident management features are not available.
621 | func IncidentClientFromContext(ctx context.Context) *incident.Client {
622 | c, ok := ctx.Value(incidentClientKey{}).(*incident.Client)
623 | if !ok {
624 | return nil
625 | }
626 | return c
627 | }
628 |
629 | // ComposeStdioContextFuncs composes multiple StdioContextFuncs into a single one.
630 | // Functions are applied in order, allowing each to modify the context before passing it to the next.
631 | func ComposeStdioContextFuncs(funcs ...server.StdioContextFunc) server.StdioContextFunc {
632 | return func(ctx context.Context) context.Context {
633 | for _, f := range funcs {
634 | ctx = f(ctx)
635 | }
636 | return ctx
637 | }
638 | }
639 |
640 | // ComposeSSEContextFuncs composes multiple SSEContextFuncs into a single one.
641 | // This enables chaining of context modifications for Server-Sent Events transport, such as extracting headers and setting up clients.
642 | func ComposeSSEContextFuncs(funcs ...httpContextFunc) server.SSEContextFunc {
643 | return func(ctx context.Context, req *http.Request) context.Context {
644 | for _, f := range funcs {
645 | ctx = f(ctx, req)
646 | }
647 | return ctx
648 | }
649 | }
650 |
651 | // ComposeHTTPContextFuncs composes multiple HTTPContextFuncs into a single one.
652 | // This enables chaining of context modifications for HTTP transport, allowing modular setup of authentication, clients, and configuration.
653 | func ComposeHTTPContextFuncs(funcs ...httpContextFunc) server.HTTPContextFunc {
654 | return func(ctx context.Context, req *http.Request) context.Context {
655 | for _, f := range funcs {
656 | ctx = f(ctx, req)
657 | }
658 | return ctx
659 | }
660 | }
661 |
662 | // ComposedStdioContextFunc returns a StdioContextFunc that comprises all predefined StdioContextFuncs.
663 | // It sets up the complete context for stdio transport including Grafana configuration, client initialization from environment variables, and incident management support.
664 | func ComposedStdioContextFunc(config GrafanaConfig) server.StdioContextFunc {
665 | return ComposeStdioContextFuncs(
666 | func(ctx context.Context) context.Context {
667 | return WithGrafanaConfig(ctx, config)
668 | },
669 | ExtractGrafanaInfoFromEnv,
670 | ExtractGrafanaClientFromEnv,
671 | ExtractIncidentClientFromEnv,
672 | )
673 | }
674 |
675 | // ComposedSSEContextFunc returns a SSEContextFunc that comprises all predefined SSEContextFuncs.
676 | // It sets up the complete context for SSE transport, extracting configuration from HTTP headers with environment variable fallbacks.
677 | func ComposedSSEContextFunc(config GrafanaConfig) server.SSEContextFunc {
678 | return ComposeSSEContextFuncs(
679 | func(ctx context.Context, req *http.Request) context.Context {
680 | return WithGrafanaConfig(ctx, config)
681 | },
682 | ExtractGrafanaInfoFromHeaders,
683 | ExtractGrafanaClientFromHeaders,
684 | ExtractIncidentClientFromHeaders,
685 | )
686 | }
687 |
688 | // ComposedHTTPContextFunc returns a HTTPContextFunc that comprises all predefined HTTPContextFuncs.
689 | // It provides the complete context setup for HTTP transport, including header-based authentication and client configuration.
690 | func ComposedHTTPContextFunc(config GrafanaConfig) server.HTTPContextFunc {
691 | return ComposeHTTPContextFuncs(
692 | func(ctx context.Context, req *http.Request) context.Context {
693 | return WithGrafanaConfig(ctx, config)
694 | },
695 | ExtractGrafanaInfoFromHeaders,
696 | ExtractGrafanaClientFromHeaders,
697 | ExtractIncidentClientFromHeaders,
698 | )
699 | }
700 |
```