This is page 2 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
--------------------------------------------------------------------------------
/tools/datasources.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "context"
5 | "fmt"
6 | "strings"
7 |
8 | "github.com/mark3labs/mcp-go/mcp"
9 | "github.com/mark3labs/mcp-go/server"
10 |
11 | "github.com/grafana/grafana-openapi-client-go/models"
12 | mcpgrafana "github.com/grafana/mcp-grafana"
13 | )
14 |
15 | type ListDatasourcesParams struct {
16 | Type string `json:"type,omitempty" jsonschema:"description=The type of datasources to search for. For example\\, 'prometheus'\\, 'loki'\\, 'tempo'\\, etc..."`
17 | }
18 |
19 | type dataSourceSummary struct {
20 | ID int64 `json:"id"`
21 | UID string `json:"uid"`
22 | Name string `json:"name"`
23 | Type string `json:"type"`
24 | IsDefault bool `json:"isDefault"`
25 | }
26 |
27 | func listDatasources(ctx context.Context, args ListDatasourcesParams) ([]dataSourceSummary, error) {
28 | c := mcpgrafana.GrafanaClientFromContext(ctx)
29 | resp, err := c.Datasources.GetDataSources()
30 | if err != nil {
31 | return nil, fmt.Errorf("list datasources: %w", err)
32 | }
33 | datasources := filterDatasources(resp.Payload, args.Type)
34 | return summarizeDatasources(datasources), nil
35 | }
36 |
37 | // filterDatasources returns only datasources of the specified type `t`. If `t`
38 | // is an empty string no filtering is done.
39 | func filterDatasources(datasources models.DataSourceList, t string) models.DataSourceList {
40 | if t == "" {
41 | return datasources
42 | }
43 | filtered := models.DataSourceList{}
44 | t = strings.ToLower(t)
45 | for _, ds := range datasources {
46 | if strings.Contains(strings.ToLower(ds.Type), t) {
47 | filtered = append(filtered, ds)
48 | }
49 | }
50 | return filtered
51 | }
52 |
53 | func summarizeDatasources(dataSources models.DataSourceList) []dataSourceSummary {
54 | result := make([]dataSourceSummary, 0, len(dataSources))
55 | for _, ds := range dataSources {
56 | result = append(result, dataSourceSummary{
57 | ID: ds.ID,
58 | UID: ds.UID,
59 | Name: ds.Name,
60 | Type: ds.Type,
61 | IsDefault: ds.IsDefault,
62 | })
63 | }
64 | return result
65 | }
66 |
67 | var ListDatasources = mcpgrafana.MustTool(
68 | "list_datasources",
69 | "List available Grafana datasources. Optionally filter by datasource type (e.g., 'prometheus', 'loki'). Returns a summary list including ID, UID, name, type, and default status.",
70 | listDatasources,
71 | mcp.WithTitleAnnotation("List datasources"),
72 | mcp.WithIdempotentHintAnnotation(true),
73 | mcp.WithReadOnlyHintAnnotation(true),
74 | )
75 |
76 | type GetDatasourceByUIDParams struct {
77 | UID string `json:"uid" jsonschema:"required,description=The uid of the datasource"`
78 | }
79 |
80 | func getDatasourceByUID(ctx context.Context, args GetDatasourceByUIDParams) (*models.DataSource, error) {
81 | c := mcpgrafana.GrafanaClientFromContext(ctx)
82 | datasource, err := c.Datasources.GetDataSourceByUID(args.UID)
83 | if err != nil {
84 | // Check if it's a 404 Not Found Error
85 | if strings.Contains(err.Error(), "404") {
86 | return nil, fmt.Errorf("datasource with UID '%s' not found. Please check if the datasource exists and is accessible", args.UID)
87 | }
88 | return nil, fmt.Errorf("get datasource by uid %s: %w", args.UID, err)
89 | }
90 | return datasource.Payload, nil
91 | }
92 |
93 | var GetDatasourceByUID = mcpgrafana.MustTool(
94 | "get_datasource_by_uid",
95 | "Retrieves detailed information about a specific datasource using its UID. Returns the full datasource model, including name, type, URL, access settings, JSON data, and secure JSON field status.",
96 | getDatasourceByUID,
97 | mcp.WithTitleAnnotation("Get datasource by UID"),
98 | mcp.WithIdempotentHintAnnotation(true),
99 | mcp.WithReadOnlyHintAnnotation(true),
100 | )
101 |
102 | type GetDatasourceByNameParams struct {
103 | Name string `json:"name" jsonschema:"required,description=The name of the datasource"`
104 | }
105 |
106 | func getDatasourceByName(ctx context.Context, args GetDatasourceByNameParams) (*models.DataSource, error) {
107 | c := mcpgrafana.GrafanaClientFromContext(ctx)
108 | datasource, err := c.Datasources.GetDataSourceByName(args.Name)
109 | if err != nil {
110 | return nil, fmt.Errorf("get datasource by name %s: %w", args.Name, err)
111 | }
112 | return datasource.Payload, nil
113 | }
114 |
115 | var GetDatasourceByName = mcpgrafana.MustTool(
116 | "get_datasource_by_name",
117 | "Retrieves detailed information about a specific datasource using its name. Returns the full datasource model, including UID, type, URL, access settings, JSON data, and secure JSON field status.",
118 | getDatasourceByName,
119 | mcp.WithTitleAnnotation("Get datasource by name"),
120 | mcp.WithIdempotentHintAnnotation(true),
121 | mcp.WithReadOnlyHintAnnotation(true),
122 | )
123 |
124 | func AddDatasourceTools(mcp *server.MCPServer) {
125 | ListDatasources.Register(mcp)
126 | GetDatasourceByUID.Register(mcp)
127 | GetDatasourceByName.Register(mcp)
128 | }
129 |
```
--------------------------------------------------------------------------------
/tests/conftest.py:
--------------------------------------------------------------------------------
```python
1 | import pytest
2 | import os
3 | import asyncio
4 | import gc
5 | import base64
6 | from dotenv import load_dotenv
7 | from mcp.client.sse import sse_client
8 | from mcp.client.stdio import stdio_client
9 | from mcp.client.streamable_http import streamablehttp_client
10 | from mcp import ClientSession, StdioServerParameters
11 |
12 | load_dotenv()
13 |
14 | DEFAULT_GRAFANA_URL = "http://localhost:3000"
15 | DEFAULT_MCP_URL = "http://localhost:8000"
16 | DEFAULT_MCP_TRANSPORT = "sse"
17 |
18 | models = ["gpt-4o", "claude-3-5-sonnet-20240620"]
19 |
20 | pytestmark = pytest.mark.anyio
21 |
22 |
23 | @pytest.fixture
24 | def anyio_backend():
25 | return "asyncio"
26 |
27 |
28 | @pytest.fixture(autouse=True)
29 | async def cleanup_sessions():
30 | """Clean up any lingering HTTP sessions after each test."""
31 | yield
32 | # Force garbage collection to clean up any unclosed sessions
33 | gc.collect()
34 | # Give a brief moment for cleanup
35 | await asyncio.sleep(0.01)
36 |
37 |
38 | @pytest.fixture
39 | def mcp_transport():
40 | return os.environ.get("MCP_TRANSPORT", DEFAULT_MCP_TRANSPORT)
41 |
42 |
43 | @pytest.fixture
44 | def mcp_url():
45 | return os.environ.get("MCP_GRAFANA_URL", DEFAULT_MCP_URL)
46 |
47 |
48 | @pytest.fixture
49 | def grafana_env():
50 | env = {"GRAFANA_URL": os.environ.get("GRAFANA_URL", DEFAULT_GRAFANA_URL)}
51 | # Check for the new service account token environment variable first
52 | if key := os.environ.get("GRAFANA_SERVICE_ACCOUNT_TOKEN"):
53 | env["GRAFANA_SERVICE_ACCOUNT_TOKEN"] = key
54 | elif key := os.environ.get("GRAFANA_API_KEY"):
55 | env["GRAFANA_API_KEY"] = key
56 | import warnings
57 |
58 | warnings.warn(
59 | "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.",
60 | DeprecationWarning,
61 | )
62 | elif (username := os.environ.get("GRAFANA_USERNAME")) and (
63 | password := os.environ.get("GRAFANA_PASSWORD")
64 | ):
65 | env["GRAFANA_USERNAME"] = username
66 | env["GRAFANA_PASSWORD"] = password
67 | return env
68 |
69 |
70 | @pytest.fixture
71 | def grafana_headers():
72 | headers = {
73 | "X-Grafana-URL": os.environ.get("GRAFANA_URL", DEFAULT_GRAFANA_URL),
74 | }
75 | # Check for the new service account token environment variable first
76 | if key := os.environ.get("GRAFANA_SERVICE_ACCOUNT_TOKEN"):
77 | headers["X-Grafana-API-Key"] = key
78 | elif key := os.environ.get("GRAFANA_API_KEY"):
79 | headers["X-Grafana-API-Key"] = key
80 | import warnings
81 |
82 | warnings.warn(
83 | "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.",
84 | DeprecationWarning,
85 | )
86 | elif (username := os.environ.get("GRAFANA_USERNAME")) and (
87 | password := os.environ.get("GRAFANA_PASSWORD")
88 | ):
89 | credentials = f"{username}:{password}"
90 | headers["Authorization"] = (
91 | "Basic " + base64.b64encode(credentials.encode("utf-8")).decode()
92 | )
93 | return headers
94 |
95 |
96 | @pytest.fixture
97 | async def mcp_client(mcp_transport, mcp_url, grafana_env, grafana_headers):
98 | if mcp_transport == "stdio":
99 | params = StdioServerParameters(
100 | command=os.environ.get("MCP_GRAFANA_PATH", "../dist/mcp-grafana"),
101 | args=["--debug", "--log-level", "debug"],
102 | env=grafana_env,
103 | )
104 | async with stdio_client(params) as (read, write):
105 | async with ClientSession(read, write) as session:
106 | await session.initialize()
107 | yield session
108 | elif mcp_transport == "sse":
109 | url = f"{mcp_url}/sse"
110 | async with sse_client(url, headers=grafana_headers) as (read, write):
111 | async with ClientSession(read, write) as session:
112 | await session.initialize()
113 | yield session
114 | elif mcp_transport == "streamable-http":
115 | # Use HTTP client for streamable-http transport
116 | url = f"{mcp_url}/mcp"
117 | async with streamablehttp_client(url, headers=grafana_headers) as (
118 | read,
119 | write,
120 | _,
121 | ):
122 | async with ClientSession(read, write) as session:
123 | await session.initialize()
124 | yield session
125 | else:
126 | raise ValueError(f"Unsupported transport: {mcp_transport}")
127 |
```
--------------------------------------------------------------------------------
/tests/disable_write_test.py:
--------------------------------------------------------------------------------
```python
1 | import pytest
2 | import os
3 | from mcp.client.stdio import stdio_client
4 | from mcp import ClientSession, StdioServerParameters
5 |
6 | pytestmark = pytest.mark.anyio
7 |
8 |
9 | @pytest.fixture
10 | def grafana_env():
11 | env = {"GRAFANA_URL": os.environ.get("GRAFANA_URL", "http://localhost:3000")}
12 | # Check for the new service account token environment variable first
13 | if key := os.environ.get("GRAFANA_SERVICE_ACCOUNT_TOKEN"):
14 | env["GRAFANA_SERVICE_ACCOUNT_TOKEN"] = key
15 | elif key := os.environ.get("GRAFANA_API_KEY"):
16 | env["GRAFANA_API_KEY"] = key
17 | return env
18 |
19 |
20 | async def test_disable_write_flag_disables_write_tools(grafana_env):
21 | """Test that --disable-write flag disables write tools."""
22 | params = StdioServerParameters(
23 | command=os.environ.get("MCP_GRAFANA_PATH", "../dist/mcp-grafana"),
24 | args=["--disable-write"],
25 | env=grafana_env,
26 | )
27 | async with stdio_client(params) as (read, write):
28 | async with ClientSession(read, write) as session:
29 | await session.initialize()
30 |
31 | # List all available tools
32 | tools_result = await session.list_tools()
33 | tool_names = [tool.name for tool in tools_result.tools]
34 |
35 | # Verify write tools are NOT present
36 | write_tools = [
37 | "update_dashboard",
38 | "create_folder",
39 | "create_incident",
40 | "add_activity_to_incident",
41 | "create_alert_rule",
42 | "update_alert_rule",
43 | "delete_alert_rule",
44 | "create_annotation",
45 | "create_graphite_annotation",
46 | "update_annotation",
47 | "patch_annotation",
48 | "find_error_pattern_logs",
49 | "find_slow_requests",
50 | ]
51 |
52 | for tool in write_tools:
53 | assert tool not in tool_names, f"Write tool '{tool}' should not be available with --disable-write flag"
54 |
55 | # Verify read tools ARE still present
56 | read_tools = [
57 | "get_dashboard_by_uid",
58 | "list_alert_rules",
59 | "get_alert_rule_by_uid",
60 | "list_contact_points",
61 | "list_incidents",
62 | "get_incident",
63 | "get_sift_investigation",
64 | "get_annotations",
65 | "get_annotation_tags",
66 | ]
67 |
68 | for tool in read_tools:
69 | assert tool in tool_names, f"Read tool '{tool}' should still be available with --disable-write flag"
70 |
71 |
72 | async def test_without_disable_write_flag_enables_write_tools(grafana_env):
73 | """Test that without --disable-write flag, write tools are enabled."""
74 | params = StdioServerParameters(
75 | command=os.environ.get("MCP_GRAFANA_PATH", "../dist/mcp-grafana"),
76 | args=[], # No --disable-write flag
77 | env=grafana_env,
78 | )
79 | async with stdio_client(params) as (read, write):
80 | async with ClientSession(read, write) as session:
81 | await session.initialize()
82 |
83 | # List all available tools
84 | tools_result = await session.list_tools()
85 | tool_names = [tool.name for tool in tools_result.tools]
86 |
87 | # Verify write tools ARE present
88 | write_tools = [
89 | "update_dashboard",
90 | "create_folder",
91 | "create_incident",
92 | "add_activity_to_incident",
93 | "create_alert_rule",
94 | "update_alert_rule",
95 | "delete_alert_rule",
96 | "create_annotation",
97 | "create_graphite_annotation",
98 | "update_annotation",
99 | "patch_annotation",
100 | "find_error_pattern_logs",
101 | "find_slow_requests",
102 | ]
103 |
104 | for tool in write_tools:
105 | assert tool in tool_names, f"Write tool '{tool}' should be available without --disable-write flag"
106 |
107 | # Verify read tools are also present
108 | read_tools = [
109 | "get_dashboard_by_uid",
110 | "list_alert_rules",
111 | "get_alert_rule_by_uid",
112 | "list_contact_points",
113 | "list_incidents",
114 | "get_incident",
115 | "get_sift_investigation",
116 | "get_annotations",
117 | "get_annotation_tags",
118 | ]
119 |
120 | for tool in read_tools:
121 | assert tool in tool_names, f"Read tool '{tool}' should be available without --disable-write flag"
122 |
123 |
```
--------------------------------------------------------------------------------
/testdata/provisioning/alerting/alert_rules.yaml:
--------------------------------------------------------------------------------
```yaml
1 | apiVersion: 1
2 |
3 | groups:
4 | - orgId: 1
5 | name: Test Alert Rules
6 | folder: Tests
7 | interval: 1m
8 | rules:
9 | - uid: test_alert_rule_1
10 | title: Test Alert Rule 1
11 | condition: B
12 | data:
13 | - refId: A
14 | relativeTimeRange:
15 | from: 600
16 | to: 0
17 | datasourceUid: prometheus
18 | model:
19 | datasource:
20 | type: prometheus
21 | uid: prometheus
22 | editorMode: code
23 | expr: vector(1)
24 | hide: false
25 | instant: true
26 | legendFormat: __auto
27 | range: false
28 | refId: A
29 | - refId: B
30 | datasourceUid: __expr__
31 | model:
32 | conditions:
33 | - evaluator:
34 | params:
35 | - 0
36 | - 0
37 | type: gt
38 | operator:
39 | type: and
40 | query:
41 | params: []
42 | reducer:
43 | params: []
44 | type: avg
45 | type: query
46 | datasource:
47 | name: Expression
48 | type: __expr__
49 | uid: __expr__
50 | expression: A
51 | hide: false
52 | refId: B
53 | type: threshold
54 | noDataState: NoData
55 | execErrState: Error
56 | for: 1m
57 | keepFiringFor: 0s
58 | annotations:
59 | description: This is a test alert rule that is always firing
60 | labels:
61 | severity: info
62 | type: test
63 | rule: first
64 | isPaused: false
65 |
66 | - uid: test_alert_rule_2
67 | title: Test Alert Rule 2
68 | condition: B
69 | data:
70 | - refId: A
71 | relativeTimeRange:
72 | from: 600
73 | to: 0
74 | datasourceUid: prometheus
75 | model:
76 | datasource:
77 | type: prometheus
78 | uid: prometheus
79 | editorMode: code
80 | expr: vector(0)
81 | hide: false
82 | instant: true
83 | legendFormat: __auto
84 | range: false
85 | refId: A
86 | - refId: B
87 | datasourceUid: __expr__
88 | model:
89 | conditions:
90 | - evaluator:
91 | params:
92 | - 0
93 | - 0
94 | type: gt
95 | operator:
96 | type: and
97 | query:
98 | params: []
99 | reducer:
100 | params: []
101 | type: avg
102 | type: query
103 | datasource:
104 | name: Expression
105 | type: __expr__
106 | uid: __expr__
107 | expression: A
108 | hide: false
109 | refId: B
110 | type: threshold
111 | noDataState: NoData
112 | execErrState: Error
113 | for: 1m
114 | keepFiringFor: 0s
115 | annotations:
116 | description: This is a test alert rule that is always normal
117 | labels:
118 | severity: info
119 | type: test
120 | rule: second
121 | isPaused: false
122 |
123 | - uid: test_alert_rule_paused
124 | title: Test Alert Rule (Paused)
125 | condition: B
126 | data:
127 | - refId: A
128 | relativeTimeRange:
129 | from: 600
130 | to: 0
131 | datasourceUid: prometheus
132 | model:
133 | datasource:
134 | type: prometheus
135 | uid: prometheus
136 | editorMode: code
137 | expr: vector(1)
138 | hide: false
139 | instant: true
140 | legendFormat: __auto
141 | range: false
142 | refId: A
143 | - refId: B
144 | datasourceUid: __expr__
145 | model:
146 | conditions:
147 | - evaluator:
148 | params:
149 | - 0
150 | - 0
151 | type: gt
152 | operator:
153 | type: and
154 | query:
155 | params: []
156 | reducer:
157 | params: []
158 | type: avg
159 | type: query
160 | datasource:
161 | name: Expression
162 | type: __expr__
163 | uid: __expr__
164 | expression: A
165 | hide: false
166 | refId: B
167 | type: threshold
168 | noDataState: NoData
169 | execErrState: Error
170 | for: 1m
171 | keepFiringFor: 0s
172 | annotations:
173 | description: This is a paused alert rule
174 | labels:
175 | severity: info
176 | type: test
177 | rule: third
178 | isPaused: true
179 |
```
--------------------------------------------------------------------------------
/tools/asserts.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "bytes"
5 | "context"
6 | "encoding/json"
7 | "fmt"
8 | "io"
9 | "net/http"
10 | "strings"
11 | "time"
12 |
13 | mcpgrafana "github.com/grafana/mcp-grafana"
14 | "github.com/mark3labs/mcp-go/mcp"
15 | "github.com/mark3labs/mcp-go/server"
16 | )
17 |
18 | func newAssertsClient(ctx context.Context) (*Client, error) {
19 | cfg := mcpgrafana.GrafanaConfigFromContext(ctx)
20 | url := fmt.Sprintf("%s/api/plugins/grafana-asserts-app/resources/asserts/api-server", strings.TrimRight(cfg.URL, "/"))
21 |
22 | // Create custom transport with TLS configuration if available
23 | var transport = http.DefaultTransport
24 | if tlsConfig := cfg.TLSConfig; tlsConfig != nil {
25 | var err error
26 | transport, err = tlsConfig.HTTPTransport(transport.(*http.Transport))
27 | if err != nil {
28 | return nil, fmt.Errorf("failed to create custom transport: %w", err)
29 | }
30 | }
31 |
32 | transport = NewAuthRoundTripper(transport, cfg.AccessToken, cfg.IDToken, cfg.APIKey, cfg.BasicAuth)
33 | transport = mcpgrafana.NewOrgIDRoundTripper(transport, cfg.OrgID)
34 |
35 | client := &http.Client{
36 | Transport: mcpgrafana.NewUserAgentTransport(
37 | transport,
38 | ),
39 | }
40 |
41 | return &Client{
42 | httpClient: client,
43 | baseURL: url,
44 | }, nil
45 | }
46 |
47 | type GetAssertionsParams struct {
48 | StartTime time.Time `json:"startTime" jsonschema:"required,description=The start time in RFC3339 format"`
49 | EndTime time.Time `json:"endTime" jsonschema:"required,description=The end time in RFC3339 format"`
50 | EntityType string `json:"entityType" jsonschema:"description=The type of the entity to list (e.g. Service\\, Node\\, Pod\\, etc.)"`
51 | EntityName string `json:"entityName" jsonschema:"description=The name of the entity to list"`
52 | Env string `json:"env,omitempty" jsonschema:"description=The env of the entity to list"`
53 | Site string `json:"site,omitempty" jsonschema:"description=The site of the entity to list"`
54 | Namespace string `json:"namespace,omitempty" jsonschema:"description=The namespace of the entity to list"`
55 | }
56 |
57 | type scope struct {
58 | Env string `json:"env,omitempty"`
59 | Site string `json:"site,omitempty"`
60 | Namespace string `json:"namespace,omitempty"`
61 | }
62 |
63 | type entity struct {
64 | Name string `json:"name"`
65 | Type string `json:"type"`
66 | Scope scope `json:"scope"`
67 | }
68 |
69 | type requestBody struct {
70 | StartTime int64 `json:"startTime"`
71 | EndTime int64 `json:"endTime"`
72 | EntityKeys []entity `json:"entityKeys"`
73 | SuggestionSrcEntities []entity `json:"suggestionSrcEntities"`
74 | AlertCategories []string `json:"alertCategories"`
75 | }
76 |
77 | func (c *Client) fetchAssertsData(ctx context.Context, urlPath string, method string, reqBody any) (string, error) {
78 | jsonData, err := json.Marshal(reqBody)
79 | if err != nil {
80 | return "", fmt.Errorf("failed to marshal request body: %w", err)
81 | }
82 |
83 | req, err := http.NewRequestWithContext(ctx, method, c.baseURL+urlPath, bytes.NewBuffer(jsonData))
84 | if err != nil {
85 | return "", fmt.Errorf("failed to create request: %w", err)
86 | }
87 | req.Header.Set("Content-Type", "application/json")
88 |
89 | resp, err := c.httpClient.Do(req)
90 | if err != nil {
91 | return "", fmt.Errorf("failed to execute request: %w", err)
92 | }
93 | defer func() {
94 | _ = resp.Body.Close() //nolint:errcheck
95 | }()
96 |
97 | body, err := io.ReadAll(resp.Body)
98 | if err != nil {
99 | return "", fmt.Errorf("failed to read response body: %w", err)
100 | }
101 |
102 | if resp.StatusCode != http.StatusOK {
103 | return "", fmt.Errorf("request failed with status %d: %s", resp.StatusCode, string(body))
104 | }
105 |
106 | return string(body), nil
107 | }
108 |
109 | func getAssertions(ctx context.Context, args GetAssertionsParams) (string, error) {
110 | client, err := newAssertsClient(ctx)
111 | if err != nil {
112 | return "", fmt.Errorf("failed to create Asserts client: %w", err)
113 | }
114 |
115 | // Create request body
116 | reqBody := requestBody{
117 | StartTime: args.StartTime.UnixMilli(),
118 | EndTime: args.EndTime.UnixMilli(),
119 | EntityKeys: []entity{
120 | {
121 | Name: args.EntityName,
122 | Type: args.EntityType,
123 | Scope: scope{},
124 | },
125 | },
126 | SuggestionSrcEntities: []entity{},
127 | AlertCategories: []string{"saturation", "amend", "anomaly", "failure", "error"},
128 | }
129 |
130 | if args.Env != "" {
131 | reqBody.EntityKeys[0].Scope.Env = args.Env
132 | }
133 | if args.Site != "" {
134 | reqBody.EntityKeys[0].Scope.Site = args.Site
135 | }
136 | if args.Namespace != "" {
137 | reqBody.EntityKeys[0].Scope.Namespace = args.Namespace
138 | }
139 |
140 | data, err := client.fetchAssertsData(ctx, "/v1/assertions/llm-summary", "POST", reqBody)
141 | if err != nil {
142 | return "", fmt.Errorf("failed to fetch data: %w", err)
143 | }
144 |
145 | return data, nil
146 | }
147 |
148 | var GetAssertions = mcpgrafana.MustTool(
149 | "get_assertions",
150 | "Get assertion summary for a given entity with its type, name, env, site, namespace, and a time range",
151 | getAssertions,
152 | mcp.WithTitleAnnotation("Get assertions summary"),
153 | mcp.WithIdempotentHintAnnotation(true),
154 | mcp.WithReadOnlyHintAnnotation(true),
155 | )
156 |
157 | func AddAssertsTools(mcp *server.MCPServer) {
158 | GetAssertions.Register(mcp)
159 | }
160 |
```
--------------------------------------------------------------------------------
/tests/admin_test.py:
--------------------------------------------------------------------------------
```python
1 | from typing import Dict
2 | import pytest
3 | from langevals import expect
4 | from langevals_langevals.llm_boolean import (
5 | CustomLLMBooleanEvaluator,
6 | CustomLLMBooleanSettings,
7 | )
8 | from litellm import Message, acompletion
9 | from mcp import ClientSession
10 | import aiohttp
11 | import uuid
12 | import os
13 | from conftest import DEFAULT_GRAFANA_URL
14 |
15 | from conftest import models
16 | from utils import (
17 | get_converted_tools,
18 | llm_tool_call_sequence,
19 | )
20 |
21 | pytestmark = pytest.mark.anyio
22 |
23 |
24 | @pytest.fixture
25 | async def grafana_team():
26 | """Create a temporary test team and clean it up after the test is done."""
27 | # Generate a unique team name to avoid conflicts
28 | team_name = f"test-team-{uuid.uuid4().hex[:8]}"
29 |
30 | # Get Grafana URL and service account token from environment
31 | grafana_url = os.environ.get("GRAFANA_URL", DEFAULT_GRAFANA_URL)
32 |
33 | auth_header = None
34 | # Check for the new service account token environment variable first
35 | if api_key := os.environ.get("GRAFANA_SERVICE_ACCOUNT_TOKEN"):
36 | auth_header = {"Authorization": f"Bearer {api_key}"}
37 | elif api_key := os.environ.get("GRAFANA_API_KEY"):
38 | auth_header = {"Authorization": f"Bearer {api_key}"}
39 | import warnings
40 |
41 | warnings.warn(
42 | "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.",
43 | DeprecationWarning,
44 | )
45 |
46 | if not auth_header:
47 | pytest.skip("No authentication credentials available to create team")
48 |
49 | # Create the team using Grafana API
50 | team_id = None
51 | async with aiohttp.ClientSession() as session:
52 | create_url = f"{grafana_url}/api/teams"
53 | async with session.post(
54 | create_url,
55 | headers=auth_header,
56 | json={"name": team_name, "email": f"{team_name}@example.com"},
57 | ) as response:
58 | if response.status != 200:
59 | resp_text = await response.text()
60 | pytest.skip(f"Failed to create team: {resp_text}")
61 | resp_data = await response.json()
62 | team_id = resp_data.get("teamId")
63 |
64 | # Yield the team info for the test to use
65 | yield {"id": team_id, "name": team_name}
66 |
67 | # Clean up after the test
68 | if team_id:
69 | async with aiohttp.ClientSession() as session:
70 | delete_url = f"{grafana_url}/api/teams/{team_id}"
71 | async with session.delete(delete_url, headers=auth_header) as response:
72 | if response.status != 200:
73 | resp_text = await response.text()
74 | print(f"Warning: Failed to delete team: {resp_text}")
75 |
76 |
77 | @pytest.mark.parametrize("model", models)
78 | @pytest.mark.flaky(max_runs=3)
79 | async def test_list_teams_tool(
80 | model: str, mcp_client: ClientSession, grafana_team: Dict[str, str]
81 | ):
82 | tools = await get_converted_tools(mcp_client)
83 | team_name = grafana_team["name"]
84 | prompt = "Can you list the teams in Grafana?"
85 |
86 | messages = [
87 | Message(role="system", content="You are a helpful assistant."),
88 | Message(role="user", content=prompt),
89 | ]
90 |
91 | # 1. Call the list teams tool
92 | messages = await llm_tool_call_sequence(
93 | model,
94 | messages,
95 | tools,
96 | mcp_client,
97 | "list_teams",
98 | )
99 |
100 | # 2. Final LLM response
101 | response = await acompletion(model=model, messages=messages, tools=tools)
102 | content = response.choices[0].message.content
103 | panel_queries_checker = CustomLLMBooleanEvaluator(
104 | settings=CustomLLMBooleanSettings(
105 | prompt=(
106 | "Does the response contain specific information about "
107 | "the teams in Grafana?"
108 | f"There should be a team named {team_name}. "
109 | ),
110 | )
111 | )
112 | expect(input=prompt, output=content).to_pass(panel_queries_checker)
113 |
114 |
115 | @pytest.mark.parametrize("model", models)
116 | @pytest.mark.flaky(max_runs=3)
117 | async def test_list_users_by_org_tool(model: str, mcp_client: ClientSession):
118 | tools = await get_converted_tools(mcp_client)
119 | prompt = "Can you list the users in Grafana?"
120 |
121 | messages = [
122 | Message(role="system", content="You are a helpful assistant."),
123 | Message(role="user", content=prompt),
124 | ]
125 |
126 | # 1. Call the list_users_by_org tool
127 | messages = await llm_tool_call_sequence(
128 | model, messages, tools, mcp_client, "list_users_by_org"
129 | )
130 |
131 | # 2. Final LLM response
132 | response = await acompletion(model=model, messages=messages, tools=tools)
133 | content = response.choices[0].message.content
134 | user_checker = CustomLLMBooleanEvaluator(
135 | settings=CustomLLMBooleanSettings(
136 | prompt="Does the response contain specific information about users in Grafana, such as usernames, emails, or roles?",
137 | )
138 | )
139 | expect(input=prompt, output=content).to_pass(user_checker)
140 |
```
--------------------------------------------------------------------------------
/tools/alerting_client.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "context"
5 | "encoding/json"
6 | "fmt"
7 | "io"
8 | "net/http"
9 | "net/url"
10 | "strconv"
11 | "strings"
12 | "time"
13 |
14 | "github.com/prometheus/prometheus/model/labels"
15 |
16 | mcpgrafana "github.com/grafana/mcp-grafana"
17 | )
18 |
19 | const (
20 | defaultTimeout = 30 * time.Second
21 | rulesEndpointPath = "/api/prometheus/grafana/api/v1/rules"
22 | )
23 |
24 | type alertingClient struct {
25 | baseURL *url.URL
26 | accessToken string
27 | idToken string
28 | apiKey string
29 | basicAuth *url.Userinfo
30 | orgID int64
31 | httpClient *http.Client
32 | }
33 |
34 | func newAlertingClientFromContext(ctx context.Context) (*alertingClient, error) {
35 | cfg := mcpgrafana.GrafanaConfigFromContext(ctx)
36 | baseURL := strings.TrimRight(cfg.URL, "/")
37 | parsedBaseURL, err := url.Parse(baseURL)
38 | if err != nil {
39 | return nil, fmt.Errorf("invalid Grafana base URL %q: %w", baseURL, err)
40 | }
41 |
42 | client := &alertingClient{
43 | baseURL: parsedBaseURL,
44 | accessToken: cfg.AccessToken,
45 | idToken: cfg.IDToken,
46 | apiKey: cfg.APIKey,
47 | basicAuth: cfg.BasicAuth,
48 | orgID: cfg.OrgID,
49 | httpClient: &http.Client{
50 | Timeout: defaultTimeout,
51 | },
52 | }
53 |
54 | // Create custom transport with TLS configuration if available
55 | if tlsConfig := mcpgrafana.GrafanaConfigFromContext(ctx).TLSConfig; tlsConfig != nil {
56 | client.httpClient.Transport, err = tlsConfig.HTTPTransport(http.DefaultTransport.(*http.Transport))
57 | if err != nil {
58 | return nil, fmt.Errorf("failed to create custom transport: %w", err)
59 | }
60 | // Wrap with user agent
61 | client.httpClient.Transport = mcpgrafana.NewUserAgentTransport(
62 | client.httpClient.Transport,
63 | )
64 | } else {
65 | // No custom TLS, but still add user agent
66 | client.httpClient.Transport = mcpgrafana.NewUserAgentTransport(
67 | http.DefaultTransport,
68 | )
69 | }
70 |
71 | return client, nil
72 | }
73 |
74 | func (c *alertingClient) makeRequest(ctx context.Context, path string) (*http.Response, error) {
75 | p := c.baseURL.JoinPath(path).String()
76 |
77 | req, err := http.NewRequestWithContext(ctx, http.MethodGet, p, nil)
78 | if err != nil {
79 | return nil, fmt.Errorf("failed to create request to %s: %w", p, err)
80 | }
81 |
82 | req.Header.Set("Accept", "application/json")
83 | req.Header.Set("Content-Type", "application/json")
84 |
85 | // If accessToken is set we use that first and fall back to normal Authorization.
86 | if c.accessToken != "" && c.idToken != "" {
87 | req.Header.Set("X-Access-Token", c.accessToken)
88 | req.Header.Set("X-Grafana-Id", c.idToken)
89 | } else if c.apiKey != "" {
90 | req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.apiKey))
91 | } else if c.basicAuth != nil {
92 | password, _ := c.basicAuth.Password()
93 | req.SetBasicAuth(c.basicAuth.Username(), password)
94 | }
95 |
96 | // Add org ID header for multi-org support
97 | if c.orgID > 0 {
98 | req.Header.Set("X-Scope-OrgId", strconv.FormatInt(c.orgID, 10))
99 | }
100 |
101 | resp, err := c.httpClient.Do(req)
102 | if err != nil {
103 | return nil, fmt.Errorf("failed to execute request to %s: %w", p, err)
104 | }
105 | if resp.StatusCode != http.StatusOK {
106 | bodyBytes, _ := io.ReadAll(resp.Body)
107 | _ = resp.Body.Close() //nolint:errcheck
108 | return nil, fmt.Errorf("grafana API returned status code %d: %s", resp.StatusCode, string(bodyBytes))
109 | }
110 |
111 | return resp, nil
112 | }
113 |
114 | func (c *alertingClient) GetRules(ctx context.Context) (*rulesResponse, error) {
115 | resp, err := c.makeRequest(ctx, rulesEndpointPath)
116 | if err != nil {
117 | return nil, fmt.Errorf("failed to get alert rules from Grafana API: %w", err)
118 | }
119 | defer func() {
120 | _ = resp.Body.Close() //nolint:errcheck
121 | }()
122 |
123 | var rulesResponse rulesResponse
124 | decoder := json.NewDecoder(resp.Body)
125 | if err := decoder.Decode(&rulesResponse); err != nil {
126 | return nil, fmt.Errorf("failed to decode rules response from %s: %w", rulesEndpointPath, err)
127 | }
128 |
129 | return &rulesResponse, nil
130 | }
131 |
132 | type rulesResponse struct {
133 | Data struct {
134 | RuleGroups []ruleGroup `json:"groups"`
135 | NextToken string `json:"groupNextToken,omitempty"`
136 | Totals map[string]int64 `json:"totals,omitempty"`
137 | } `json:"data"`
138 | }
139 |
140 | type ruleGroup struct {
141 | Name string `json:"name"`
142 | FolderUID string `json:"folderUid"`
143 | Rules []alertingRule `json:"rules"`
144 | Interval float64 `json:"interval"`
145 | LastEvaluation time.Time `json:"lastEvaluation"`
146 | EvaluationTime float64 `json:"evaluationTime"`
147 | }
148 |
149 | type alertingRule struct {
150 | State string `json:"state,omitempty"`
151 | Name string `json:"name,omitempty"`
152 | Query string `json:"query,omitempty"`
153 | Duration float64 `json:"duration,omitempty"`
154 | KeepFiringFor float64 `json:"keepFiringFor,omitempty"`
155 | Annotations labels.Labels `json:"annotations,omitempty"`
156 | ActiveAt *time.Time `json:"activeAt,omitempty"`
157 | Alerts []alert `json:"alerts,omitempty"`
158 | Totals map[string]int64 `json:"totals,omitempty"`
159 | TotalsFiltered map[string]int64 `json:"totalsFiltered,omitempty"`
160 | UID string `json:"uid"`
161 | FolderUID string `json:"folderUid"`
162 | Labels labels.Labels `json:"labels,omitempty"`
163 | Health string `json:"health"`
164 | LastError string `json:"lastError,omitempty"`
165 | Type string `json:"type"`
166 | LastEvaluation time.Time `json:"lastEvaluation"`
167 | EvaluationTime float64 `json:"evaluationTime"`
168 | }
169 |
170 | type alert struct {
171 | Labels labels.Labels `json:"labels"`
172 | Annotations labels.Labels `json:"annotations"`
173 | State string `json:"state"`
174 | ActiveAt *time.Time `json:"activeAt"`
175 | Value string `json:"value"`
176 | }
177 |
```
--------------------------------------------------------------------------------
/internal/linter/jsonschema/jsonschema_lint_test.go:
--------------------------------------------------------------------------------
```go
1 | package linter
2 |
3 | import (
4 | "os"
5 | "path/filepath"
6 | "testing"
7 | )
8 |
9 | func TestFindUnescapedCommas(t *testing.T) {
10 | // Create a temporary directory for test files
11 | tmpDir, err := os.MkdirTemp("", "jsonschema-linter-test")
12 | if err != nil {
13 | t.Fatalf("Failed to create temp dir: %v", err)
14 | }
15 | defer func() {
16 | if err := os.RemoveAll(tmpDir); err != nil {
17 | t.Logf("Failed to remove temp dir: %v", err)
18 | }
19 | }()
20 |
21 | // Create test files
22 | testFiles := map[string]string{
23 | "valid.go": `package test
24 |
25 | // Valid has properly escaped commas
26 | type Valid struct {
27 | Name string ` + "`json:\"name\" jsonschema:\"description=A valid field\\, with escaped comma\"`" + `
28 | Age int ` + "`json:\"age\" jsonschema:\"description=Another valid field\"`" + `
29 | }
30 | `,
31 | "invalid.go": `package test
32 |
33 | // Invalid has unescaped commas
34 | type Invalid struct {
35 | Name string ` + "`json:\"name\" jsonschema:\"description=An invalid field, with unescaped comma\"`" + `
36 | Age int ` + "`json:\"age\" jsonschema:\"description=Another valid field\"`" + `
37 | }
38 | `,
39 | "mixed.go": `package test
40 |
41 | // Mixed has both valid and invalid fields
42 | type Mixed struct {
43 | Valid string ` + "`json:\"valid\" jsonschema:\"description=A valid field\\, with escaped comma\"`" + `
44 | Invalid string ` + "`json:\"invalid\" jsonschema:\"description=An invalid field, with unescaped comma\"`" + `
45 | }
46 | `,
47 | }
48 |
49 | for filename, content := range testFiles {
50 | filePath := filepath.Join(tmpDir, filename)
51 | if err := os.WriteFile(filePath, []byte(content), 0644); err != nil {
52 | t.Fatalf("Failed to write test file %s: %v", filename, err)
53 | }
54 | }
55 |
56 | // Run the linter
57 | linter := &JSONSchemaLinter{}
58 | err = linter.FindUnescapedCommas(tmpDir)
59 | if err != nil {
60 | t.Fatalf("Linter failed: %v", err)
61 | }
62 |
63 | // Check if we found the expected errors
64 | if len(linter.Errors) != 2 {
65 | t.Errorf("Expected 2 errors, got %d", len(linter.Errors))
66 | }
67 |
68 | // Check if the errors are in the expected files
69 | fileErrors := make(map[string]int)
70 | for _, e := range linter.Errors {
71 | fileName := filepath.Base(e.FilePath)
72 | fileErrors[fileName]++
73 | }
74 |
75 | if fileErrors["invalid.go"] != 1 {
76 | t.Errorf("Expected 1 error in invalid.go, got %d", fileErrors["invalid.go"])
77 | }
78 |
79 | if fileErrors["mixed.go"] != 1 {
80 | t.Errorf("Expected 1 error in mixed.go, got %d", fileErrors["mixed.go"])
81 | }
82 |
83 | if fileErrors["valid.go"] != 0 {
84 | t.Errorf("Expected 0 errors in valid.go, got %d", fileErrors["valid.go"])
85 | }
86 | }
87 |
88 | // TestEscapedQuotesWithComma tests if the regex correctly identifies unescaped commas
89 | // in jsonschema tags that contain escaped quotes
90 | func TestEscapedQuotesWithComma(t *testing.T) {
91 | testCases := []struct {
92 | tag string
93 | shouldMatch bool
94 | description string
95 | }{
96 | {`jsonschema:"description=This has an unescaped, comma"`, true, "Simple unescaped comma"},
97 | {`jsonschema:"description=This has escaped quote \"followed by, comma"`, true, "Escaped quote then unescaped comma"},
98 | {`jsonschema:"description=This has escaped quote \", comma"`, true, "Escaped quote, comma with space"},
99 | {`jsonschema:"description=This has escaped quote \\\"and escaped\\, comma"`, false, "Properly escaped quote and comma"},
100 | {`jsonschema:"description=No comma here"`, false, "No comma at all"},
101 | }
102 |
103 | for _, tc := range testCases {
104 | t.Run(tc.description, func(t *testing.T) {
105 | matches := tagPattern.FindStringSubmatch(tc.tag)
106 | hasMatch := len(matches) > 0
107 | if hasMatch != tc.shouldMatch {
108 | t.Fatalf("Test failed for %s: expected match=%v, got=%v\n", tc.description, tc.shouldMatch, hasMatch)
109 | }
110 | })
111 | }
112 | }
113 |
114 | func TestFixUnescapedCommas(t *testing.T) {
115 | // Create a temporary directory for test files
116 | tmpDir, err := os.MkdirTemp("", "jsonschema-linter-test")
117 | if err != nil {
118 | t.Fatalf("Failed to create temp dir: %v", err)
119 | }
120 | defer func() {
121 | if err := os.RemoveAll(tmpDir); err != nil {
122 | t.Logf("Failed to remove temp dir: %v", err)
123 | }
124 | }()
125 |
126 | // Create a test file with unescaped commas
127 | invalidContent := `package test
128 |
129 | // Invalid has unescaped commas
130 | type Invalid struct {
131 | Name string ` + "`json:\"name\" jsonschema:\"description=An invalid field, with unescaped comma\"`" + `
132 | Age int ` + "`json:\"age\" jsonschema:\"description=Another field, also with unescaped comma\"`" + `
133 | }
134 | `
135 |
136 | // Expected content after fixing
137 | // Note: We need double backslashes in the actual file, so we use double escaped backslashes here
138 | expectedContent := `package test
139 |
140 | // Invalid has unescaped commas
141 | type Invalid struct {
142 | Name string ` + "`json:\"name\" jsonschema:\"description=An invalid field\\\\, with unescaped comma\"`" + `
143 | Age int ` + "`json:\"age\" jsonschema:\"description=Another field\\\\, also with unescaped comma\"`" + `
144 | }
145 | `
146 |
147 | filePath := filepath.Join(tmpDir, "invalid.go")
148 | if err := os.WriteFile(filePath, []byte(invalidContent), 0644); err != nil {
149 | t.Fatalf("Failed to write test file: %v", err)
150 | }
151 |
152 | // Run the linter with fix mode enabled
153 | linter := &JSONSchemaLinter{FixMode: true}
154 | err = linter.FindUnescapedCommas(tmpDir)
155 | if err != nil {
156 | t.Fatalf("Linter failed: %v", err)
157 | }
158 |
159 | // Check if we found the expected errors
160 | if len(linter.Errors) != 2 {
161 | t.Errorf("Expected 2 errors, got %d", len(linter.Errors))
162 | }
163 |
164 | // Verify the file was fixed
165 | fixedContent, err := os.ReadFile(filePath)
166 | if err != nil {
167 | t.Fatalf("Failed to read fixed file: %v", err)
168 | }
169 |
170 | if string(fixedContent) != expectedContent {
171 | t.Errorf("File not fixed correctly.\nExpected:\n%s\n\nGot:\n%s", expectedContent, string(fixedContent))
172 | }
173 |
174 | // Verify the fixed field was correctly tracked
175 | if !linter.Fixed[filePath] {
176 | t.Errorf("Fixed file not tracked in linter.Fixed")
177 | }
178 | }
179 |
```
--------------------------------------------------------------------------------
/tools/sift_cloud_test.go:
--------------------------------------------------------------------------------
```go
1 | //go:build cloud
2 | // +build cloud
3 |
4 | // This file contains cloud integration tests that run against a dedicated test instance
5 | // at mcptests.grafana-dev.net. This instance is configured with a minimal setup on the Sift side:
6 | // - 2 test investigations
7 | // These tests expect this configuration to exist and will skip if the required
8 | // environment variables (GRAFANA_URL, GRAFANA_SERVICE_ACCOUNT_TOKEN or GRAFANA_API_KEY) are not set.
9 | // The GRAFANA_API_KEY variable is deprecated.
10 |
11 | package tools
12 |
13 | import (
14 | "testing"
15 | "time"
16 |
17 | "github.com/stretchr/testify/assert"
18 | "github.com/stretchr/testify/require"
19 | )
20 |
21 | func TestCloudSiftInvestigations(t *testing.T) {
22 | ctx := createCloudTestContext(t, "Sift", "GRAFANA_URL", "GRAFANA_API_KEY")
23 |
24 | // Test listing all investigations
25 | t.Run("list all investigations", func(t *testing.T) {
26 | result, err := listSiftInvestigations(ctx, ListSiftInvestigationsParams{})
27 | require.NoError(t, err, "Should not error when listing investigations")
28 | assert.NotNil(t, result, "Result should not be nil")
29 | assert.GreaterOrEqual(t, len(result), 1, "Should have at least one investigation")
30 | })
31 |
32 | // Test listing investigations with a limit
33 | t.Run("list investigations with limit", func(t *testing.T) {
34 | // Get the client
35 | client, err := siftClientFromContext(ctx)
36 | require.NoError(t, err, "Should not error when getting Sift client")
37 |
38 | // List investigations with a limit of 1
39 | investigations, err := client.listSiftInvestigations(ctx, 1)
40 | require.NoError(t, err, "Should not error when listing investigations with limit")
41 | assert.NotNil(t, investigations, "Investigations should not be nil")
42 | assert.LessOrEqual(t, len(investigations), 1, "Should have at most one investigation")
43 |
44 | // If there are investigations, verify their structure
45 | if len(investigations) > 0 {
46 | investigation := investigations[0]
47 | assert.NotEmpty(t, investigation.ID, "Investigation should have an ID")
48 | assert.NotEmpty(t, investigation.Name, "Investigation should have a name")
49 | assert.NotEmpty(t, investigation.TenantID, "Investigation should have a tenant ID")
50 | }
51 | })
52 |
53 | // Get an investigation ID from the list to test getting a specific investigation
54 | investigations, err := listSiftInvestigations(ctx, ListSiftInvestigationsParams{Limit: 10})
55 | require.NoError(t, err, "Should not error when listing investigations")
56 | require.NotEmpty(t, investigations, "Should have at least one investigation to test with")
57 |
58 | // Find an investigation with at least one analysis.
59 | var investigationID string
60 | for _, investigation := range investigations {
61 | if len(investigation.Analyses.Items) > 0 {
62 | investigationID = investigation.ID.String()
63 | break
64 | }
65 | }
66 | require.NotEmpty(t, investigationID, "Should have at least one investigation with at least one analysis")
67 |
68 | // Test getting a specific investigation
69 | t.Run("get specific investigation", func(t *testing.T) {
70 | result, err := getSiftInvestigation(ctx, GetSiftInvestigationParams{
71 | ID: investigationID,
72 | })
73 | require.NoError(t, err, "Should not error when getting specific investigation")
74 | assert.NotNil(t, result, "Result should not be nil")
75 | assert.Equal(t, investigationID, result.ID.String(), "Should return the correct investigation")
76 |
77 | // Verify all required fields are present
78 | assert.NotEmpty(t, result.Name, "Investigation should have a name")
79 | assert.NotEmpty(t, result.TenantID, "Investigation should have a tenant ID")
80 | assert.NotNil(t, result.GrafanaURL, "Investigation should have a Grafana URL")
81 | assert.NotNil(t, result.Status, "Investigation should have a status")
82 | assert.NotNil(t, result.FailureReason, "Investigation should have a failure reason")
83 | })
84 |
85 | // Test getting a non-existent investigation
86 | t.Run("get non-existent investigation", func(t *testing.T) {
87 | _, err := getSiftInvestigation(ctx, GetSiftInvestigationParams{
88 | ID: "00000000-0000-0000-0000-000000000000",
89 | })
90 | assert.NoError(t, err, "Should not error when getting non-existent investigation")
91 | })
92 |
93 | // Test getting analyses for an investigation
94 | t.Run("get analyses for investigation", func(t *testing.T) {
95 | // Get the investigation
96 | result, err := getSiftInvestigation(ctx, GetSiftInvestigationParams{
97 | ID: investigationID,
98 | })
99 | require.NoError(t, err, "Should not error when getting specific investigation")
100 | assert.NotNil(t, result, "Result should not be nil")
101 |
102 | // Get an analysis ID
103 | analysisID := result.Analyses.Items[0].ID
104 |
105 | // Get the analysis
106 | analysis, err := getSiftAnalysis(ctx, GetSiftAnalysisParams{
107 | InvestigationID: investigationID,
108 | AnalysisID: analysisID.String(),
109 | })
110 | require.NoError(t, err, "Should not error when getting specific analysis")
111 | assert.NotNil(t, analysis, "Analysis should not be nil")
112 |
113 | // Verify all required fields are present
114 | assert.NotEmpty(t, analysis.Name, "Analysis should have a name")
115 | assert.NotEmpty(t, analysis.InvestigationID, "Analysis should have an investigation ID")
116 | assert.NotNil(t, analysis.Result, "Analysis should have a result")
117 | })
118 |
119 | t.Run("find error patterns", func(t *testing.T) {
120 | // Find error patterns
121 | analysis, err := findErrorPatternLogs(ctx, FindErrorPatternLogsParams{
122 | Name: "Test Sift",
123 | Labels: map[string]string{
124 | "namespace": "hosted-grafana",
125 | "cluster": "dev-eu-west-2",
126 | "slug": "mcptests",
127 | },
128 | Start: time.Now().Add(-5 * time.Minute),
129 | End: time.Now(),
130 | })
131 | require.NoError(t, err, "Should not error when finding error patterns")
132 | assert.NotNil(t, analysis, "Result should not be nil")
133 |
134 | // Verify all required fields are present
135 | assert.NotEmpty(t, analysis.Name, "Analysis should have a name")
136 | assert.NotEmpty(t, analysis.InvestigationID, "Analysis should have an investigation ID")
137 | assert.NotEmpty(t, analysis.Result.Message, "Analysis should have a message")
138 | })
139 | }
140 |
```
--------------------------------------------------------------------------------
/tls_test.go:
--------------------------------------------------------------------------------
```go
1 | package mcpgrafana
2 |
3 | import (
4 | "fmt"
5 | "net/http"
6 | "testing"
7 |
8 | "github.com/stretchr/testify/assert"
9 | "github.com/stretchr/testify/require"
10 | )
11 |
12 | func TestTLSConfig_CreateTLSConfig(t *testing.T) {
13 | t.Run("nil config returns nil", func(t *testing.T) {
14 | var config *TLSConfig
15 | tlsCfg, err := config.CreateTLSConfig()
16 | assert.NoError(t, err)
17 | assert.Nil(t, tlsCfg)
18 | })
19 |
20 | t.Run("skip verify only", func(t *testing.T) {
21 | config := &TLSConfig{SkipVerify: true}
22 | tlsCfg, err := config.CreateTLSConfig()
23 | assert.NoError(t, err)
24 | require.NotNil(t, tlsCfg)
25 | assert.True(t, tlsCfg.InsecureSkipVerify)
26 | assert.Empty(t, tlsCfg.Certificates)
27 | assert.Nil(t, tlsCfg.RootCAs)
28 | })
29 |
30 | t.Run("invalid cert file", func(t *testing.T) {
31 | config := &TLSConfig{
32 | CertFile: "nonexistent.pem",
33 | KeyFile: "nonexistent.key",
34 | }
35 | _, err := config.CreateTLSConfig()
36 | assert.Error(t, err)
37 | assert.Contains(t, err.Error(), "failed to load client certificate")
38 | })
39 |
40 | t.Run("invalid CA file", func(t *testing.T) {
41 | config := &TLSConfig{
42 | CAFile: "nonexistent-ca.pem",
43 | }
44 | _, err := config.CreateTLSConfig()
45 | assert.Error(t, err)
46 | assert.Contains(t, err.Error(), "failed to read CA certificate")
47 | })
48 | }
49 |
50 | func TestHTTPTransport(t *testing.T) {
51 | t.Run("nil TLS config", func(t *testing.T) {
52 | var tlsConfig *TLSConfig
53 | transport, err := tlsConfig.HTTPTransport(http.DefaultTransport.(*http.Transport))
54 | assert.NoError(t, err)
55 | assert.NotNil(t, transport)
56 |
57 | // Should be default transport clone
58 | httpTransport, ok := transport.(*http.Transport)
59 | require.True(t, ok)
60 | assert.NotNil(t, httpTransport)
61 | })
62 |
63 | t.Run("skip verify config", func(t *testing.T) {
64 | tlsConfig := &TLSConfig{SkipVerify: true}
65 | transport, err := tlsConfig.HTTPTransport(http.DefaultTransport.(*http.Transport))
66 | assert.NoError(t, err)
67 | require.NotNil(t, transport)
68 |
69 | httpTransport, ok := transport.(*http.Transport)
70 | require.True(t, ok)
71 | require.NotNil(t, httpTransport.TLSClientConfig)
72 | assert.True(t, httpTransport.TLSClientConfig.InsecureSkipVerify)
73 | })
74 |
75 | t.Run("invalid TLS config", func(t *testing.T) {
76 | tlsConfig := &TLSConfig{
77 | CertFile: "nonexistent.pem",
78 | KeyFile: "nonexistent.key",
79 | }
80 | _, err := tlsConfig.HTTPTransport(http.DefaultTransport.(*http.Transport))
81 | assert.Error(t, err)
82 | })
83 | }
84 |
85 | // mockRoundTripper is a mock implementation of http.RoundTripper for testing
86 | type mockRoundTripper struct {
87 | capturedRequest *http.Request
88 | response *http.Response
89 | err error
90 | }
91 |
92 | func (m *mockRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
93 | m.capturedRequest = req
94 | if m.response != nil {
95 | return m.response, m.err
96 | }
97 | // Return a default successful response
98 | return &http.Response{
99 | StatusCode: 200,
100 | Header: make(http.Header),
101 | Body: http.NoBody,
102 | }, m.err
103 | }
104 |
105 | func TestUserAgentTransport(t *testing.T) {
106 | tests := []struct {
107 | name string
108 | userAgent string
109 | existingUserAgent string
110 | expectedUserAgent string
111 | }{
112 | {
113 | name: "sets user agent when not present",
114 | userAgent: "mcp-grafana/1.0.0",
115 | existingUserAgent: "",
116 | expectedUserAgent: "mcp-grafana/1.0.0",
117 | },
118 | {
119 | name: "does not override existing user agent",
120 | userAgent: "mcp-grafana/1.0.0",
121 | existingUserAgent: "existing-client/2.0.0",
122 | expectedUserAgent: "existing-client/2.0.0",
123 | },
124 | }
125 |
126 | for _, tt := range tests {
127 | t.Run(tt.name, func(t *testing.T) {
128 | // Create mock round tripper
129 | mockRT := &mockRoundTripper{}
130 |
131 | // Create user agent transport
132 | transport := &UserAgentTransport{
133 | rt: mockRT,
134 | UserAgent: tt.userAgent,
135 | }
136 |
137 | // Create request
138 | req, err := http.NewRequest("GET", "http://example.com", nil)
139 | require.NoError(t, err)
140 |
141 | // Set existing user agent if specified
142 | if tt.existingUserAgent != "" {
143 | req.Header.Set("User-Agent", tt.existingUserAgent)
144 | }
145 |
146 | // Make request through transport
147 | _, err = transport.RoundTrip(req)
148 | require.NoError(t, err)
149 |
150 | // Verify user agent header
151 | assert.Equal(t, tt.expectedUserAgent, mockRT.capturedRequest.Header.Get("User-Agent"))
152 | })
153 | }
154 | }
155 |
156 | func TestVersion(t *testing.T) {
157 | version := Version()
158 | assert.NotEmpty(t, version)
159 | // Version should be either "(devel)" for development builds or a proper version
160 | assert.True(t, version == "(devel)" || len(version) > 0)
161 | }
162 |
163 | func TestUserAgent(t *testing.T) {
164 | userAgent := UserAgent()
165 | assert.Contains(t, userAgent, "mcp-grafana/")
166 | assert.NotEqual(t, "mcp-grafana/", userAgent) // Should have version appended
167 |
168 | // Should match the pattern mcp-grafana/{version}
169 | version := Version()
170 | expected := fmt.Sprintf("mcp-grafana/%s", version)
171 | assert.Equal(t, expected, userAgent)
172 | }
173 |
174 | func TestNewUserAgentTransport(t *testing.T) {
175 | t.Run("with explicit user agent", func(t *testing.T) {
176 | mockRT := &mockRoundTripper{}
177 | userAgent := "test-agent/1.0.0"
178 |
179 | transport := NewUserAgentTransport(mockRT, userAgent)
180 |
181 | assert.Equal(t, mockRT, transport.rt)
182 | assert.Equal(t, userAgent, transport.UserAgent)
183 | })
184 |
185 | t.Run("with default user agent", func(t *testing.T) {
186 | mockRT := &mockRoundTripper{}
187 |
188 | transport := NewUserAgentTransport(mockRT)
189 |
190 | assert.Equal(t, mockRT, transport.rt)
191 | assert.Equal(t, UserAgent(), transport.UserAgent)
192 | assert.Contains(t, transport.UserAgent, "mcp-grafana/")
193 | })
194 | }
195 |
196 | func TestNewUserAgentTransportWithNilRoundTripper(t *testing.T) {
197 | t.Run("with explicit user agent", func(t *testing.T) {
198 | userAgent := "test-agent/1.0.0"
199 |
200 | transport := NewUserAgentTransport(nil, userAgent)
201 |
202 | assert.Equal(t, http.DefaultTransport, transport.rt)
203 | assert.Equal(t, userAgent, transport.UserAgent)
204 | })
205 |
206 | t.Run("with default user agent", func(t *testing.T) {
207 | transport := NewUserAgentTransport(nil)
208 |
209 | assert.Equal(t, http.DefaultTransport, transport.rt)
210 | assert.Equal(t, UserAgent(), transport.UserAgent)
211 | assert.Contains(t, transport.UserAgent, "mcp-grafana/")
212 | })
213 | }
214 |
```
--------------------------------------------------------------------------------
/tools/incident.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "context"
5 | "fmt"
6 |
7 | "github.com/grafana/incident-go"
8 | mcpgrafana "github.com/grafana/mcp-grafana"
9 | "github.com/mark3labs/mcp-go/mcp"
10 | "github.com/mark3labs/mcp-go/server"
11 | )
12 |
13 | type ListIncidentsParams struct {
14 | Limit int `json:"limit" jsonschema:"description=The maximum number of incidents to return"`
15 | Drill bool `json:"drill" jsonschema:"description=Whether to include drill incidents"`
16 | Status string `json:"status" jsonschema:"description=The status of the incidents to include. Valid values: 'active'\\, 'resolved'"`
17 | }
18 |
19 | func listIncidents(ctx context.Context, args ListIncidentsParams) (*incident.QueryIncidentPreviewsResponse, error) {
20 | c := mcpgrafana.IncidentClientFromContext(ctx)
21 | is := incident.NewIncidentsService(c)
22 |
23 | // Set default limit to 10 if not specified
24 | limit := args.Limit
25 | if limit <= 0 {
26 | limit = 10
27 | }
28 |
29 | query := ""
30 | if !args.Drill {
31 | query = "isdrill:false"
32 | }
33 | if args.Status != "" {
34 | query += fmt.Sprintf(" status:%s", args.Status)
35 | }
36 | incidents, err := is.QueryIncidentPreviews(ctx, incident.QueryIncidentPreviewsRequest{
37 | Query: incident.IncidentPreviewsQuery{
38 | QueryString: query,
39 | OrderDirection: "DESC",
40 | Limit: limit,
41 | },
42 | })
43 | if err != nil {
44 | return nil, fmt.Errorf("list incidents: %w", err)
45 | }
46 | return incidents, nil
47 | }
48 |
49 | var ListIncidents = mcpgrafana.MustTool(
50 | "list_incidents",
51 | "List Grafana incidents. Allows filtering by status ('active', 'resolved') and optionally including drill incidents. Returns a preview list with basic details.",
52 | listIncidents,
53 | mcp.WithTitleAnnotation("List incidents"),
54 | mcp.WithIdempotentHintAnnotation(true),
55 | mcp.WithReadOnlyHintAnnotation(true),
56 | )
57 |
58 | type CreateIncidentParams struct {
59 | Title string `json:"title" jsonschema:"description=The title of the incident"`
60 | Severity string `json:"severity" jsonschema:"description=The severity of the incident"`
61 | RoomPrefix string `json:"roomPrefix" jsonschema:"description=The prefix of the room to create the incident in"`
62 | IsDrill bool `json:"isDrill" jsonschema:"description=Whether the incident is a drill incident"`
63 | Status string `json:"status" jsonschema:"description=The status of the incident"`
64 | AttachCaption string `json:"attachCaption" jsonschema:"description=The caption of the attachment"`
65 | AttachURL string `json:"attachUrl" jsonschema:"description=The URL of the attachment"`
66 | Labels []incident.IncidentLabel `json:"labels" jsonschema:"description=The labels to add to the incident"`
67 | }
68 |
69 | func createIncident(ctx context.Context, args CreateIncidentParams) (*incident.Incident, error) {
70 | c := mcpgrafana.IncidentClientFromContext(ctx)
71 | is := incident.NewIncidentsService(c)
72 | incident, err := is.CreateIncident(ctx, incident.CreateIncidentRequest{
73 | Title: args.Title,
74 | Severity: args.Severity,
75 | RoomPrefix: args.RoomPrefix,
76 | IsDrill: args.IsDrill,
77 | Status: args.Status,
78 | AttachCaption: args.AttachCaption,
79 | AttachURL: args.AttachURL,
80 | Labels: args.Labels,
81 | })
82 | if err != nil {
83 | return nil, fmt.Errorf("create incident: %w", err)
84 | }
85 | return &incident.Incident, nil
86 | }
87 |
88 | var CreateIncident = mcpgrafana.MustTool(
89 | "create_incident",
90 | "Create a new Grafana incident. Requires title, severity, and room prefix. Allows setting status and labels. This tool should be used judiciously and sparingly, and only after confirmation from the user, as it may notify or alarm lots of people.",
91 | createIncident,
92 | mcp.WithTitleAnnotation("Create incident"),
93 | )
94 |
95 | type AddActivityToIncidentParams struct {
96 | IncidentID string `json:"incidentId" jsonschema:"description=The ID of the incident to add the activity to"`
97 | Body string `json:"body" jsonschema:"description=The body of the activity. URLs will be parsed and attached as context"`
98 | EventTime string `json:"eventTime" jsonschema:"description=The time that the activity occurred. If not provided\\, the current time will be used"`
99 | }
100 |
101 | func addActivityToIncident(ctx context.Context, args AddActivityToIncidentParams) (*incident.ActivityItem, error) {
102 | c := mcpgrafana.IncidentClientFromContext(ctx)
103 | as := incident.NewActivityService(c)
104 | activity, err := as.AddActivity(ctx, incident.AddActivityRequest{
105 | IncidentID: args.IncidentID,
106 | ActivityKind: "userNote",
107 | Body: args.Body,
108 | EventTime: args.EventTime,
109 | })
110 | if err != nil {
111 | return nil, fmt.Errorf("add activity to incident: %w", err)
112 | }
113 | return &activity.ActivityItem, nil
114 | }
115 |
116 | var AddActivityToIncident = mcpgrafana.MustTool(
117 | "add_activity_to_incident",
118 | "Add a note (userNote activity) to an existing incident's timeline using its ID. The note body can include URLs which will be attached as context. Use this to add context to an incident.",
119 | addActivityToIncident,
120 | mcp.WithTitleAnnotation("Add activity to incident"),
121 | )
122 |
123 | func AddIncidentTools(mcp *server.MCPServer, enableWriteTools bool) {
124 | ListIncidents.Register(mcp)
125 | if enableWriteTools {
126 | CreateIncident.Register(mcp)
127 | AddActivityToIncident.Register(mcp)
128 | }
129 | GetIncident.Register(mcp)
130 | }
131 |
132 | type GetIncidentParams struct {
133 | ID string `json:"id" jsonschema:"description=The ID of the incident to retrieve"`
134 | }
135 |
136 | func getIncident(ctx context.Context, args GetIncidentParams) (*incident.Incident, error) {
137 | c := mcpgrafana.IncidentClientFromContext(ctx)
138 | is := incident.NewIncidentsService(c)
139 |
140 | incidentResp, err := is.GetIncident(ctx, incident.GetIncidentRequest{
141 | IncidentID: args.ID,
142 | })
143 | if err != nil {
144 | return nil, fmt.Errorf("get incident by ID: %w", err)
145 | }
146 |
147 | return &incidentResp.Incident, nil
148 | }
149 |
150 | var GetIncident = mcpgrafana.MustTool(
151 | "get_incident",
152 | "Get a single incident by ID. Returns the full incident details including title, status, severity, labels, timestamps, and other metadata.",
153 | getIncident,
154 | mcp.WithTitleAnnotation("Get incident details"),
155 | mcp.WithIdempotentHintAnnotation(true),
156 | mcp.WithReadOnlyHintAnnotation(true),
157 | )
158 |
```
--------------------------------------------------------------------------------
/examples/tls_example.go:
--------------------------------------------------------------------------------
```go
1 | package main
2 |
3 | import (
4 | "context"
5 | "fmt"
6 | "log"
7 | "net/http"
8 | "os"
9 |
10 | mcpgrafana "github.com/grafana/mcp-grafana"
11 | "github.com/grafana/mcp-grafana/tools"
12 | "github.com/mark3labs/mcp-go/server"
13 | )
14 |
15 | func main() {
16 | // Example 1: Basic TLS configuration with skip verify (for testing)
17 | fmt.Println("Example 1: Basic TLS configuration with skip verify")
18 | basicTLSExample()
19 |
20 | // Example 2: Full mTLS configuration with client certificates
21 | fmt.Println("\nExample 2: Full mTLS configuration")
22 | fullTLSExample()
23 |
24 | // Example 3: Running an MCP server with TLS support
25 | fmt.Println("\nExample 3: MCP server with TLS support")
26 | if len(os.Args) > 1 && os.Args[1] == "run-server" {
27 | runServerWithTLS()
28 | } else {
29 | fmt.Println("Use 'go run tls_example.go run-server' to actually start the server")
30 | showServerExample()
31 | }
32 | }
33 |
34 | func basicTLSExample() {
35 | // Create a TLS config that skips certificate verification
36 | // This is useful for testing against self-signed certificates
37 | tlsConfig := &mcpgrafana.TLSConfig{SkipVerify: true}
38 |
39 | // Create a Grafana config with TLS support
40 | grafanaConfig := mcpgrafana.GrafanaConfig{
41 | Debug: true,
42 | TLSConfig: tlsConfig,
43 | }
44 |
45 | // Create a context function that includes TLS configuration
46 | contextFunc := mcpgrafana.ComposedStdioContextFunc(grafanaConfig)
47 |
48 | // Test the context function
49 | ctx := contextFunc(context.Background())
50 |
51 | // Verify the configuration is applied
52 | retrievedConfig := mcpgrafana.GrafanaConfigFromContext(ctx)
53 | if retrievedConfig.TLSConfig != nil {
54 | fmt.Printf("✓ TLS configuration applied: SkipVerify=%v\n", retrievedConfig.TLSConfig.SkipVerify)
55 | }
56 |
57 | fmt.Printf("✓ Debug mode enabled: %v\n", retrievedConfig.Debug)
58 | }
59 |
60 | func fullTLSExample() {
61 | // Example paths for certificate files
62 | // In a real scenario, these would point to actual certificate files
63 | certFile := "/path/to/client.crt"
64 | keyFile := "/path/to/client.key"
65 | caFile := "/path/to/ca.crt"
66 |
67 | // Create TLS config with client certificates and CA verification
68 | tlsConfig := &mcpgrafana.TLSConfig{
69 | CertFile: certFile,
70 | KeyFile: keyFile,
71 | CAFile: caFile,
72 | }
73 |
74 | // Create Grafana config with TLS support
75 | grafanaConfig := mcpgrafana.GrafanaConfig{
76 | Debug: false,
77 | TLSConfig: tlsConfig,
78 | }
79 |
80 | fmt.Printf("✓ TLS configuration created:\n")
81 | fmt.Printf(" - Client cert: %s\n", tlsConfig.CertFile)
82 | fmt.Printf(" - Client key: %s\n", tlsConfig.KeyFile)
83 | fmt.Printf(" - CA file: %s\n", tlsConfig.CAFile)
84 | fmt.Printf(" - Skip verify: %v\n", tlsConfig.SkipVerify)
85 | fmt.Printf(" - Debug mode: %v\n", grafanaConfig.Debug)
86 |
87 | // Create context functions for different transport types
88 | stdioFunc := mcpgrafana.ComposedStdioContextFunc(grafanaConfig)
89 | sseFunc := mcpgrafana.ComposedSSEContextFunc(grafanaConfig)
90 | httpFunc := mcpgrafana.ComposedHTTPContextFunc(grafanaConfig)
91 |
92 | fmt.Printf("✓ Context functions created for all transport types\n")
93 |
94 | _ = stdioFunc
95 | _ = sseFunc
96 | _ = httpFunc
97 | }
98 |
99 | func showServerExample() {
100 | fmt.Println("Example MCP server configuration with TLS:")
101 | fmt.Println(`// Create TLS configuration
102 | tlsConfig := &mcpgrafana.TLSConfig{
103 | CertFile: "/path/to/client.crt",
104 | KeyFile: "/path/to/client.key",
105 | CAFile: "/path/to/ca.crt",
106 | }
107 |
108 | // Create Grafana configuration
109 | grafanaConfig := mcpgrafana.GrafanaConfig{
110 | Debug: true,
111 | TLSConfig: tlsConfig,
112 | }
113 |
114 | // Create MCP server
115 | s := server.NewMCPServer("mcp-grafana", "1.0.0")
116 |
117 | // Add tools
118 | tools.AddSearchTools(s)
119 | tools.AddDatasourceTools(s)
120 | // ... add other tools as needed
121 |
122 | // Create stdio server with TLS support
123 | srv := server.NewStdioServer(s)
124 | srv.SetContextFunc(mcpgrafana.ComposedStdioContextFunc(grafanaConfig))
125 |
126 | // Start server
127 | srv.Listen(ctx, os.Stdin, os.Stdout)`)
128 | }
129 |
130 | func runServerWithTLS() {
131 | // Set up environment variables (in practice, these would be set externally)
132 | if os.Getenv("GRAFANA_URL") == "" {
133 | if err := os.Setenv("GRAFANA_URL", "https://localhost:3000"); err != nil {
134 | log.Printf("Failed to set GRAFANA_URL: %v", err)
135 | }
136 | }
137 | // Check for service account token first, then fall back to deprecated API key
138 | if os.Getenv("GRAFANA_SERVICE_ACCOUNT_TOKEN") == "" {
139 | if os.Getenv("GRAFANA_API_KEY") == "" {
140 | fmt.Println("Warning: Neither GRAFANA_SERVICE_ACCOUNT_TOKEN nor GRAFANA_API_KEY is set")
141 | } else {
142 | fmt.Println("Warning: GRAFANA_API_KEY is deprecated, please use GRAFANA_SERVICE_ACCOUNT_TOKEN instead")
143 | }
144 | }
145 |
146 | // Create TLS configuration that skips verification for demo purposes
147 | // In production, you would use real certificates
148 | tlsConfig := &mcpgrafana.TLSConfig{SkipVerify: true}
149 | grafanaConfig := mcpgrafana.GrafanaConfig{
150 | Debug: true,
151 | TLSConfig: tlsConfig,
152 | }
153 |
154 | // Create MCP server
155 | s := server.NewMCPServer("mcp-grafana-tls-example", "1.0.0")
156 |
157 | // Add some basic tools
158 | tools.AddSearchTools(s)
159 | tools.AddDatasourceTools(s)
160 | tools.AddDashboardTools(s, false) // Read-only mode (no write tools)
161 |
162 | // Create stdio server with TLS-enabled context function
163 | srv := server.NewStdioServer(s)
164 | srv.SetContextFunc(mcpgrafana.ComposedStdioContextFunc(grafanaConfig))
165 |
166 | fmt.Printf("Starting MCP Grafana server with TLS support...\n")
167 | fmt.Printf("Grafana URL: %s\n", os.Getenv("GRAFANA_URL"))
168 | fmt.Printf("TLS Skip Verify: %v\n", tlsConfig.SkipVerify)
169 |
170 | // Start the server
171 | ctx := context.Background()
172 | if err := srv.Listen(ctx, os.Stdin, os.Stdout); err != nil {
173 | log.Fatalf("Server error: %v", err)
174 | }
175 | }
176 |
177 | // Example of creating custom HTTP clients with TLS configuration
178 | func customClientExample() { //nolint:unused // Example function for documentation
179 | ctx := context.Background()
180 |
181 | // Add Grafana configuration to context
182 | tlsConfig := &mcpgrafana.TLSConfig{
183 | CertFile: "/path/to/cert.pem",
184 | KeyFile: "/path/to/key.pem",
185 | CAFile: "/path/to/ca.pem",
186 | }
187 | config := mcpgrafana.GrafanaConfig{
188 | TLSConfig: tlsConfig,
189 | }
190 | ctx = mcpgrafana.WithGrafanaConfig(ctx, config)
191 | _ = ctx // Use ctx to avoid ineffectual assignment warning
192 |
193 | // Create custom HTTP transport with TLS
194 | transport, err := tlsConfig.HTTPTransport(http.DefaultTransport.(*http.Transport))
195 | if err != nil {
196 | log.Fatalf("Failed to create transport: %v", err)
197 | }
198 |
199 | // Use the transport in your HTTP client
200 | _ = transport
201 | fmt.Println("✓ Custom HTTP transport created with TLS configuration")
202 | }
203 |
```
--------------------------------------------------------------------------------
/internal/linter/jsonschema/jsonschema_lint.go:
--------------------------------------------------------------------------------
```go
1 | package linter
2 |
3 | import (
4 | "fmt"
5 | "go/ast"
6 | "go/parser"
7 | "go/token"
8 | "os"
9 | "path/filepath"
10 | "regexp"
11 | "sort"
12 | "strings"
13 | )
14 |
15 | // JSONSchemaLinter checks for unescaped commas in jsonschema struct tags
16 | type JSONSchemaLinter struct {
17 | FilePaths []string
18 | Errors []JSONSchemaError
19 | FixMode bool
20 | Fixed map[string]bool
21 | }
22 |
23 | // JSONSchemaError represents a linting error with file position details
24 | type JSONSchemaError struct {
25 | FilePath string
26 | Line int
27 | Column int
28 | Offset int // Byte offset in the file
29 | Struct string
30 | Field string
31 | Tag string
32 | FixedTag string
33 | }
34 |
35 | // tagPattern matches jsonschema tags with description containing unescaped commas
36 | // It captures:
37 | // 1. The jsonschema tag
38 | // 2. Parts of the description containing unescaped commas
39 | // The pattern correctly handles:
40 | // - Simple unescaped comma: "description=Something, with comma"
41 | // - Escaped quote followed by unescaped comma: "description=With \"quote, and comma"
42 | // - But not match escaped comma: "description=With escaped\, comma"
43 | var tagPattern = regexp.MustCompile(`jsonschema:"([^"]*)description=(.*?[^\\],)([^"]*)"`)
44 |
45 | // FindUnescapedCommas scans Go files for jsonschema struct tags with unescaped commas in descriptions
46 | func (l *JSONSchemaLinter) FindUnescapedCommas(baseDir string) error {
47 | // Reset errors
48 | l.Errors = nil
49 | if l.FixMode {
50 | l.Fixed = make(map[string]bool)
51 | }
52 |
53 | // Walk through the directory and find Go files
54 | err := filepath.Walk(baseDir, func(path string, info os.FileInfo, err error) error {
55 | if err != nil {
56 | return err
57 | }
58 |
59 | // Skip non-Go files
60 | if !info.IsDir() && strings.HasSuffix(path, ".go") {
61 | l.FilePaths = append(l.FilePaths, path)
62 | }
63 |
64 | return nil
65 | })
66 |
67 | if err != nil {
68 | return fmt.Errorf("error walking directory: %v", err)
69 | }
70 |
71 | // Parse all Go files and check for the unescaped commas
72 | for _, path := range l.FilePaths {
73 | fset := token.NewFileSet()
74 | f, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
75 | if err != nil {
76 | return fmt.Errorf("error parsing file %s: %v", path, err)
77 | }
78 |
79 | fileErrors := []JSONSchemaError{}
80 |
81 | // Visit all struct types
82 | ast.Inspect(f, func(n ast.Node) bool {
83 | ts, ok := n.(*ast.TypeSpec)
84 | if !ok || ts.Type == nil {
85 | return true
86 | }
87 |
88 | st, ok := ts.Type.(*ast.StructType)
89 | if !ok {
90 | return true
91 | }
92 |
93 | structName := ts.Name.Name
94 |
95 | // Check each field of the struct
96 | for _, field := range st.Fields.List {
97 | if field.Tag == nil {
98 | continue
99 | }
100 |
101 | tag := field.Tag.Value
102 |
103 | // Check if the tag has a jsonschema description with unescaped comma
104 | matches := tagPattern.FindStringSubmatch(tag)
105 | if len(matches) > 0 {
106 | fieldName := ""
107 | if len(field.Names) > 0 {
108 | fieldName = field.Names[0].Name
109 | }
110 |
111 | // Generate the fixed tag by escaping the commas in the description
112 | fixedTag := tag
113 | if len(matches) > 2 {
114 | descWithUnescapedCommas := matches[2]
115 | // Escape all unescaped commas
116 | fixedDesc := escapeUnescapedCommas(descWithUnescapedCommas)
117 | // Replace the original description with the fixed one
118 | fixedTag = strings.Replace(tag, descWithUnescapedCommas, fixedDesc, 1)
119 | }
120 |
121 | pos := fset.Position(field.Tag.Pos())
122 | errorInfo := JSONSchemaError{
123 | FilePath: path,
124 | Line: pos.Line,
125 | Column: pos.Column,
126 | Offset: pos.Offset,
127 | Struct: structName,
128 | Field: fieldName,
129 | Tag: tag,
130 | FixedTag: fixedTag,
131 | }
132 | fileErrors = append(fileErrors, errorInfo)
133 | }
134 | }
135 |
136 | return true
137 | })
138 |
139 | // Add all errors for this file
140 | l.Errors = append(l.Errors, fileErrors...)
141 |
142 | // If in fix mode and we found errors, fix the file
143 | if l.FixMode && len(fileErrors) > 0 {
144 | err := l.fixFile(path, fileErrors)
145 | if err != nil {
146 | return fmt.Errorf("error fixing file %s: %v", path, err)
147 | }
148 | l.Fixed[path] = true
149 | }
150 | }
151 |
152 | return nil
153 | }
154 |
155 | // escapeUnescapedCommas escapes any unescaped commas in the description
156 | func escapeUnescapedCommas(desc string) string {
157 | // Use regex to find all commas that are not preceded by a backslash
158 | r := regexp.MustCompile(`([^\\]),`)
159 | // Replace them with the same text but with an escaped comma
160 | return r.ReplaceAllString(desc, `$1\\,`)
161 | }
162 |
163 | // fixFile applies the fixes to a file
164 | func (l *JSONSchemaLinter) fixFile(path string, errors []JSONSchemaError) error {
165 | // Read the file content
166 | content, err := os.ReadFile(path)
167 | if err != nil {
168 | return fmt.Errorf("error reading file %s: %v", path, err)
169 | }
170 |
171 | // Convert to string for easier manipulation
172 | fileContent := string(content)
173 |
174 | // Sort errors by offset in reverse order to avoid offset changes
175 | sort.Slice(errors, func(i, j int) bool {
176 | return errors[i].Offset > errors[j].Offset
177 | })
178 |
179 | // Apply fixes
180 | for _, e := range errors {
181 | // Find the tag in the file content
182 | tagStart := strings.Index(fileContent[e.Offset:], e.Tag)
183 | if tagStart == -1 {
184 | continue
185 | }
186 | absOffset := e.Offset + tagStart
187 |
188 | // Replace the tag with the fixed version
189 | fixedContent := fileContent[:absOffset] + e.FixedTag + fileContent[absOffset+len(e.Tag):]
190 | fileContent = fixedContent
191 | }
192 |
193 | // Write back to the file
194 | err = os.WriteFile(path, []byte(fileContent), 0644)
195 | if err != nil {
196 | return fmt.Errorf("error writing file %s: %v", path, err)
197 | }
198 |
199 | return nil
200 | }
201 |
202 | // PrintErrors outputs all the found errors
203 | func (l *JSONSchemaLinter) PrintErrors() {
204 | if len(l.Errors) == 0 {
205 | fmt.Println("No unescaped commas found in jsonschema descriptions.")
206 | return
207 | }
208 |
209 | if l.FixMode {
210 | fmt.Printf("Found and fixed %d unescaped commas in jsonschema descriptions:\n\n", len(l.Errors))
211 | } else {
212 | fmt.Printf("Found %d unescaped commas in jsonschema descriptions:\n\n", len(l.Errors))
213 | }
214 |
215 | for i, err := range l.Errors {
216 | relPath, _ := filepath.Rel(".", err.FilePath)
217 | fmt.Printf("%d. %s:%d:%d - Struct: %s, Field: %s\n",
218 | i+1, relPath, err.Line, err.Column, err.Struct, err.Field)
219 | fmt.Printf(" - %s\n", err.Tag)
220 | if l.FixMode {
221 | fmt.Printf(" - Fixed to: %s\n\n", err.FixedTag)
222 | } else {
223 | fmt.Printf(" - Commas in description must be escaped with \\\\,\n\n")
224 | }
225 | }
226 |
227 | if !l.FixMode {
228 | fmt.Println("Please escape all commas in jsonschema descriptions with \\\\, to prevent truncation.")
229 | fmt.Println("You can run with --fix to automatically fix these issues.")
230 | } else {
231 | fixedFileCount := len(l.Fixed)
232 | fmt.Printf("Fixed %d file(s).\n", fixedFileCount)
233 | }
234 | }
235 |
```
--------------------------------------------------------------------------------
/tests/navigation_test.py:
--------------------------------------------------------------------------------
```python
1 | import json
2 | import pytest
3 | from langevals import expect
4 | from langevals_langevals.llm_boolean import (
5 | CustomLLMBooleanEvaluator,
6 | CustomLLMBooleanSettings,
7 | )
8 | from litellm import Message, acompletion
9 | from mcp import ClientSession
10 | from mcp.types import TextContent
11 |
12 | from conftest import models
13 | from utils import (
14 | get_converted_tools,
15 | llm_tool_call_sequence,
16 | flexible_tool_call,
17 | )
18 |
19 | pytestmark = pytest.mark.anyio
20 |
21 |
22 | @pytest.mark.parametrize("model", models)
23 | @pytest.mark.flaky(max_runs=3)
24 | async def test_generate_dashboard_deeplink(model: str, mcp_client: ClientSession):
25 | tools = await get_converted_tools(mcp_client)
26 |
27 | prompt = """Please create a dashboard deeplink for dashboard with UID 'test-uid'."""
28 |
29 | messages = [
30 | Message(role="system", content="You are a helpful assistant."),
31 | Message(role="user", content=prompt),
32 | ]
33 |
34 | messages = await llm_tool_call_sequence(
35 | model, messages, tools, mcp_client, "generate_deeplink",
36 | {"resourceType": "dashboard", "dashboardUid": "test-uid"}
37 | )
38 |
39 | response = await acompletion(model=model, messages=messages, tools=tools)
40 | content = response.choices[0].message.content
41 |
42 | assert "/d/test-uid" in content, f"Expected dashboard URL with /d/test-uid, got: {content}"
43 |
44 | dashboard_link_checker = CustomLLMBooleanEvaluator(
45 | settings=CustomLLMBooleanSettings(
46 | prompt="Does the response contain a URL with /d/ path and the dashboard UID?",
47 | )
48 | )
49 | print("Dashboard deeplink content:", content)
50 | expect(input=prompt, output=content).to_pass(dashboard_link_checker)
51 |
52 |
53 | @pytest.mark.parametrize("model", models)
54 | @pytest.mark.flaky(max_runs=3)
55 | async def test_generate_panel_deeplink(model: str, mcp_client: ClientSession):
56 | tools = await get_converted_tools(mcp_client)
57 | prompt = "Generate a deeplink for panel 5 in dashboard with UID 'test-uid'"
58 |
59 | messages = [
60 | Message(role="system", content="You are a helpful assistant."),
61 | Message(role="user", content=prompt),
62 | ]
63 |
64 | messages = await llm_tool_call_sequence(
65 | model, messages, tools, mcp_client, "generate_deeplink",
66 | {
67 | "resourceType": "panel",
68 | "dashboardUid": "test-uid",
69 | "panelId": 5
70 | }
71 | )
72 |
73 | response = await acompletion(model=model, messages=messages, tools=tools)
74 | content = response.choices[0].message.content
75 |
76 | assert "viewPanel=5" in content, f"Expected panel URL with viewPanel=5, got: {content}"
77 |
78 | panel_link_checker = CustomLLMBooleanEvaluator(
79 | settings=CustomLLMBooleanSettings(
80 | prompt="Does the response contain a URL with viewPanel parameter?",
81 | )
82 | )
83 | print("Panel deeplink content:", content)
84 | expect(input=prompt, output=content).to_pass(panel_link_checker)
85 |
86 |
87 | @pytest.mark.parametrize("model", models)
88 | @pytest.mark.flaky(max_runs=3)
89 | async def test_generate_explore_deeplink(model: str, mcp_client: ClientSession):
90 | tools = await get_converted_tools(mcp_client)
91 | prompt = "Generate a deeplink for Grafana Explore with datasource 'test-uid'"
92 |
93 | messages = [
94 | Message(role="system", content="You are a helpful assistant."),
95 | Message(role="user", content=prompt),
96 | ]
97 |
98 | messages = await llm_tool_call_sequence(
99 | model, messages, tools, mcp_client, "generate_deeplink",
100 | {"resourceType": "explore", "datasourceUid": "test-uid"}
101 | )
102 |
103 | response = await acompletion(model=model, messages=messages, tools=tools)
104 | content = response.choices[0].message.content
105 |
106 | assert "/explore" in content, f"Expected explore URL with /explore path, got: {content}"
107 |
108 | explore_link_checker = CustomLLMBooleanEvaluator(
109 | settings=CustomLLMBooleanSettings(
110 | prompt="Does the response contain a URL with /explore path?",
111 | )
112 | )
113 | print("Explore deeplink content:", content)
114 | expect(input=prompt, output=content).to_pass(explore_link_checker)
115 |
116 |
117 | @pytest.mark.parametrize("model", models)
118 | @pytest.mark.flaky(max_runs=3)
119 | async def test_generate_deeplink_with_time_range(model: str, mcp_client: ClientSession):
120 | tools = await get_converted_tools(mcp_client)
121 | prompt = "Generate a dashboard deeplink for 'test-uid' showing the last 6 hours"
122 |
123 | messages = [
124 | Message(role="system", content="You are a helpful assistant."),
125 | Message(role="user", content=prompt),
126 | ]
127 |
128 | messages = await llm_tool_call_sequence(
129 | model, messages, tools, mcp_client, "generate_deeplink",
130 | {
131 | "resourceType": "dashboard",
132 | "dashboardUid": "test-uid",
133 | "timeRange": {
134 | "from": "now-6h",
135 | "to": "now"
136 | }
137 | }
138 | )
139 |
140 | response = await acompletion(model=model, messages=messages, tools=tools)
141 | content = response.choices[0].message.content
142 |
143 | assert "from=now-6h" in content and "to=now" in content, f"Expected time range parameters, got: {content}"
144 |
145 | time_range_checker = CustomLLMBooleanEvaluator(
146 | settings=CustomLLMBooleanSettings(
147 | prompt="Does the response contain a URL with time range parameters?",
148 | )
149 | )
150 | print("Time range deeplink content:", content)
151 | expect(input=prompt, output=content).to_pass(time_range_checker)
152 |
153 |
154 | @pytest.mark.parametrize("model", models)
155 | @pytest.mark.flaky(max_runs=3)
156 | async def test_generate_deeplink_with_query_params(model: str, mcp_client: ClientSession):
157 | tools = await get_converted_tools(mcp_client)
158 | prompt = "Use the generate_deeplink tool to create a dashboard link for UID 'test-uid' with var-datasource=prometheus and refresh=30s as query parameters"
159 |
160 | messages = [
161 | Message(role="system", content="You are a helpful assistant."),
162 | Message(role="user", content=prompt),
163 | ]
164 |
165 | # Use flexible tool call with required parameters
166 | messages = await flexible_tool_call(
167 | model, messages, tools, mcp_client, "generate_deeplink",
168 | required_params={"resourceType": "dashboard", "dashboardUid": "test-uid"}
169 | )
170 |
171 | response = await acompletion(model=model, messages=messages, tools=tools)
172 | content = response.choices[0].message.content
173 |
174 | # Verify both specific query parameters are in the final URL
175 | assert "var-datasource=prometheus" in content, f"Expected var-datasource=prometheus in URL, got: {content}"
176 | assert "refresh=30s" in content, f"Expected refresh=30s in URL, got: {content}"
177 |
178 | custom_params_checker = CustomLLMBooleanEvaluator(
179 | settings=CustomLLMBooleanSettings(
180 | prompt="Does the response contain a URL with custom query parameters?",
181 | )
182 | )
183 | print("Custom params deeplink content:", content)
184 | expect(input=prompt, output=content).to_pass(custom_params_checker)
185 |
186 |
187 |
```
--------------------------------------------------------------------------------
/tests/utils.py:
--------------------------------------------------------------------------------
```python
1 | import json
2 | from litellm.types.utils import ModelResponse
3 | from litellm import acompletion, Choices, Message
4 | from mcp.types import TextContent, Tool
5 |
6 |
7 | async def assert_and_handle_tool_call(
8 | response: ModelResponse,
9 | mcp_client,
10 | expected_tool: str,
11 | expected_args: dict = None,
12 | ) -> list:
13 | messages = []
14 | tool_calls = []
15 | for c in response.choices:
16 | assert isinstance(c, Choices)
17 | tool_calls.extend(c.message.tool_calls or [])
18 | messages.append(c.message)
19 |
20 | # Better error message if wrong number of tool calls
21 | if len(tool_calls) != 1:
22 | actual_calls = [tc.function.name for tc in tool_calls] if tool_calls else []
23 | assert len(tool_calls) == 1, (
24 | f"\n❌ Expected exactly 1 tool call, got {len(tool_calls)}\n"
25 | f"Expected tool: {expected_tool}\n"
26 | f"Actual tools called: {actual_calls}\n"
27 | f"LLM response: {response.choices[0].message.content if response.choices else 'N/A'}"
28 | )
29 |
30 | for tool_call in tool_calls:
31 | actual_tool = tool_call.function.name
32 | if actual_tool != expected_tool:
33 | # Parse arguments to understand what LLM was trying to do
34 | try:
35 | actual_args = (
36 | json.loads(tool_call.function.arguments)
37 | if tool_call.function.arguments
38 | else {}
39 | )
40 | except:
41 | actual_args = tool_call.function.arguments
42 |
43 | assert False, (
44 | f"\n❌ LLM called wrong tool!\n"
45 | f"Expected: {expected_tool}\n"
46 | f"Got: {actual_tool}\n"
47 | f"With args: {json.dumps(actual_args, indent=2)}\n"
48 | f"\n💡 Debugging tips:\n"
49 | f" - Check if the prompt clearly indicates which tool to use\n"
50 | f" - Verify the expected tool exists in the available tools\n"
51 | f" - Consider if the tool description is clear enough\n"
52 | )
53 | arguments = (
54 | {}
55 | if len(tool_call.function.arguments) == 0
56 | else json.loads(tool_call.function.arguments)
57 | )
58 | if expected_args:
59 | for key, value in expected_args.items():
60 | if key not in arguments:
61 | assert False, (
62 | f"\n❌ Missing expected parameter '{key}'\n"
63 | f"Expected args: {json.dumps(expected_args, indent=2)}\n"
64 | f"Actual args: {json.dumps(arguments, indent=2)}\n"
65 | )
66 | if arguments[key] != value:
67 | assert False, (
68 | f"\n❌ Wrong value for parameter '{key}'\n"
69 | f"Expected: {value}\n"
70 | f"Got: {arguments[key]}\n"
71 | f"Full args: {json.dumps(arguments, indent=2)}\n"
72 | )
73 | result = await mcp_client.call_tool(tool_call.function.name, arguments)
74 | assert len(result.content) == 1, (
75 | f"Expected one result for tool {tool_call.function.name}, got {len(result.content)}"
76 | )
77 | assert isinstance(result.content[0], TextContent), (
78 | f"Expected TextContent for tool {tool_call.function.name}, got {type(result.content[0])}"
79 | )
80 | messages.append(
81 | Message(
82 | role="tool", tool_call_id=tool_call.id, content=result.content[0].text
83 | )
84 | )
85 | return messages
86 |
87 |
88 | def convert_tool(tool: Tool) -> dict:
89 | return {
90 | "type": "function",
91 | "function": {
92 | "name": tool.name,
93 | "description": tool.description,
94 | "parameters": {
95 | **tool.inputSchema,
96 | "properties": tool.inputSchema.get("properties", {}),
97 | },
98 | },
99 | }
100 |
101 |
102 | async def llm_tool_call_sequence(
103 | model, messages, tools, mcp_client, tool_name, tool_args=None
104 | ):
105 | print(f"\n🤖 Calling LLM ({model}) and expecting tool: {tool_name}")
106 | print(f"📝 Last message: {messages[-1].get('content', messages[-1])[:200]}...")
107 |
108 | response = await acompletion(
109 | model=model,
110 | messages=messages,
111 | tools=tools,
112 | )
113 | assert isinstance(response, ModelResponse)
114 |
115 | # Print what tool was actually called for debugging
116 | if response.choices and response.choices[0].message.tool_calls:
117 | actual_tool = response.choices[0].message.tool_calls[0].function.name
118 | print(f"✅ LLM called: {actual_tool}")
119 | if actual_tool != tool_name:
120 | print(f"⚠️ WARNING: Expected {tool_name} but got {actual_tool}")
121 |
122 | messages.extend(
123 | await assert_and_handle_tool_call(
124 | response, mcp_client, tool_name, tool_args or {}
125 | )
126 | )
127 | return messages
128 |
129 |
130 | async def get_converted_tools(mcp_client):
131 | tools = await mcp_client.list_tools()
132 | return [convert_tool(t) for t in tools.tools]
133 |
134 |
135 | async def flexible_tool_call(
136 | model, messages, tools, mcp_client, expected_tool_name, required_params=None
137 | ):
138 | """
139 | Make a flexible tool call that only checks essential parameters.
140 | Returns updated messages list.
141 |
142 | Args:
143 | model: The LLM model to use
144 | messages: Current conversation messages
145 | tools: Available tools
146 | mcp_client: MCP client session
147 | expected_tool_name: Name of the tool we expect to be called
148 | required_params: Dict of essential parameters to check (optional)
149 |
150 | Returns:
151 | Updated messages list including tool call and result
152 | """
153 | response = await acompletion(model=model, messages=messages, tools=tools)
154 |
155 | # Check that a tool call was made
156 | assert response.choices[0].message.tool_calls is not None, (
157 | f"Expected tool call for {expected_tool_name}"
158 | )
159 | assert len(response.choices[0].message.tool_calls) >= 1, (
160 | f"Expected at least one tool call for {expected_tool_name}"
161 | )
162 |
163 | tool_call = response.choices[0].message.tool_calls[0]
164 | assert tool_call.function.name == expected_tool_name, (
165 | f"Expected {expected_tool_name} tool, got {tool_call.function.name}"
166 | )
167 |
168 | arguments = json.loads(tool_call.function.arguments)
169 |
170 | # Check required parameters if specified
171 | if required_params:
172 | for key, expected_value in required_params.items():
173 | assert key in arguments, f"Expected parameter '{key}' in tool arguments"
174 | if expected_value is not None:
175 | assert arguments[key] == expected_value, (
176 | f"Expected {key}='{expected_value}', got {key}='{arguments.get(key)}'"
177 | )
178 |
179 | # Call the tool to verify it works
180 | result = await mcp_client.call_tool(tool_call.function.name, arguments)
181 | assert len(result.content) == 1
182 | assert isinstance(result.content[0], TextContent)
183 |
184 | # Add both the tool call and result to message history
185 | messages.append(response.choices[0].message)
186 | messages.append(
187 | Message(role="tool", tool_call_id=tool_call.id, content=result.content[0].text)
188 | )
189 |
190 | return messages
191 |
```
--------------------------------------------------------------------------------
/tools/annotations_unit_test.go:
--------------------------------------------------------------------------------
```go
1 | //go:build unit
2 |
3 | package tools
4 |
5 | import (
6 | "context"
7 | "encoding/json"
8 | "net/http"
9 | "net/http/httptest"
10 | "net/url"
11 | "strconv"
12 | "testing"
13 |
14 | "github.com/grafana/grafana-openapi-client-go/client"
15 | "github.com/grafana/grafana-openapi-client-go/models"
16 | mcpgrafana "github.com/grafana/mcp-grafana"
17 | "github.com/stretchr/testify/assert"
18 | "github.com/stretchr/testify/require"
19 | )
20 |
21 | func mockCtxWithClient(server *httptest.Server) context.Context {
22 | u, _ := url.Parse(server.URL)
23 | cfg := client.DefaultTransportConfig()
24 | cfg.Host = u.Host
25 | cfg.Schemes = []string{"http"}
26 | cfg.APIKey = "test"
27 |
28 | c := client.NewHTTPClientWithConfig(nil, cfg)
29 | return mcpgrafana.WithGrafanaClient(context.Background(), c)
30 | }
31 |
32 | func TestGetAnnotations_UsesCorrectQueryParams(t *testing.T) {
33 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
34 | assert.Equal(t, "/api/annotations", r.URL.Path)
35 |
36 | q := r.URL.Query()
37 | assert.Equal(t, "50", q.Get("limit"))
38 | assert.Equal(t, "dash-1", q.Get("dashboardUID"))
39 | assert.Equal(t, "true", q.Get("matchAny"))
40 | assert.Equal(t, "tagA", q["tags"][0])
41 | assert.Equal(t, "tagB", q["tags"][1])
42 |
43 | w.Header().Set("Content-Type", "application/json")
44 | w.WriteHeader(http.StatusOK)
45 | _ = json.NewEncoder(w).Encode([]interface{}{})
46 | }))
47 | defer server.Close()
48 |
49 | ctx := mockCtxWithClient(server)
50 | limit := int64(50)
51 | uid := "dash-1"
52 | matchAny := true
53 |
54 | _, err := getAnnotations(ctx, GetAnnotationsInput{
55 | Limit: &limit,
56 | DashboardUID: &uid,
57 | MatchAny: &matchAny,
58 | Tags: []string{"tagA", "tagB"},
59 | })
60 | require.NoError(t, err)
61 | }
62 |
63 | func TestGetAnnotations_PropagatesError(t *testing.T) {
64 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
65 | w.WriteHeader(http.StatusInternalServerError)
66 | _, _ = w.Write([]byte(`oops`))
67 | }))
68 | defer server.Close()
69 |
70 | ctx := mockCtxWithClient(server)
71 |
72 | _, err := getAnnotations(ctx, GetAnnotationsInput{})
73 | require.Error(t, err)
74 | assert.Contains(t, err.Error(), "get annotations:")
75 | }
76 |
77 | func TestCreateAnnotationGraphiteFormat_SendsCorrectBody_Minimal(t *testing.T) {
78 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
79 | assert.Equal(t, "/api/annotations/graphite", r.URL.Path)
80 | assert.Equal(t, "POST", r.Method)
81 | assert.Equal(t, "Bearer test", r.Header.Get("Authorization"))
82 |
83 | var body models.PostGraphiteAnnotationsCmd
84 | require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
85 | assert.Equal(t, "deploy", body.What)
86 | assert.Equal(t, int64(1710000000000), body.When)
87 | assert.Nil(t, body.Tags)
88 | assert.Empty(t, body.Data)
89 |
90 | w.Header().Set("Content-Type", "application/json")
91 | w.WriteHeader(http.StatusOK)
92 | _, _ = w.Write([]byte(`{"message":"annotation created"}`))
93 | }))
94 | defer server.Close()
95 |
96 | ctx := mockCtxWithClient(server)
97 |
98 | _, err := createAnnotationGraphiteFormat(ctx, CreateGraphiteAnnotationInput{
99 | What: "deploy",
100 | When: 1710000000000,
101 | })
102 | require.NoError(t, err)
103 | }
104 |
105 | func TestCreateAnnotationGraphiteFormat_SendsCorrectBody_WithTagsAndData(t *testing.T) {
106 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
107 | assert.Equal(t, "/api/annotations/graphite", r.URL.Path)
108 | assert.Equal(t, "POST", r.Method)
109 |
110 | var body map[string]interface{}
111 | require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
112 |
113 | assert.Equal(t, "incident", body["what"])
114 | assert.Equal(t, float64(1720000000000), body["when"])
115 | assert.ElementsMatch(t, []interface{}{"sev1", "network"}, body["tags"].([]interface{}))
116 | assert.Equal(t, "context", body["data"])
117 |
118 | w.Header().Set("Content-Type", "application/json")
119 | w.WriteHeader(http.StatusOK)
120 | _, _ = w.Write([]byte(`{"message":"ok"}`))
121 | }))
122 | defer server.Close()
123 |
124 | ctx := mockCtxWithClient(server)
125 |
126 | _, err := createAnnotationGraphiteFormat(ctx, CreateGraphiteAnnotationInput{
127 | What: "incident",
128 | When: 1720000000000,
129 | Tags: []string{"sev1", "network"},
130 | Data: "context",
131 | })
132 | require.NoError(t, err)
133 | }
134 |
135 | func TestCreateAnnotation_SendsCorrectBody(t *testing.T) {
136 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
137 | assert.Equal(t, "/api/annotations", r.URL.Path)
138 | assert.Equal(t, "POST", r.Method)
139 |
140 | var body models.PostAnnotationsCmd
141 | require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
142 |
143 | assert.Equal(t, int64(7), body.PanelID)
144 | assert.Equal(t, "hello", *body.Text)
145 |
146 | w.Header().Set("Content-Type", "application/json")
147 | w.WriteHeader(http.StatusOK)
148 | _, _ = w.Write([]byte(`{"id": 1}`))
149 | }))
150 | defer server.Close()
151 |
152 | ctx := mockCtxWithClient(server)
153 |
154 | _, err := createAnnotation(ctx, CreateAnnotationInput{
155 | PanelID: 7,
156 | Text: "hello",
157 | })
158 | require.NoError(t, err)
159 | }
160 |
161 | func TestCreateAnnotation_ErrorWrapped(t *testing.T) {
162 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
163 | w.WriteHeader(http.StatusInternalServerError)
164 | }))
165 | defer server.Close()
166 |
167 | ctx := mockCtxWithClient(server)
168 |
169 | _, err := createAnnotation(ctx, CreateAnnotationInput{Text: "t"})
170 | require.Error(t, err)
171 | assert.Contains(t, err.Error(), "create annotation:")
172 | }
173 |
174 | func TestCreateAnnotationGraphiteFormat_HTTPError(t *testing.T) {
175 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
176 | w.WriteHeader(http.StatusInternalServerError)
177 | _, _ = w.Write([]byte(`internal error`))
178 | }))
179 | defer server.Close()
180 |
181 | ctx := mockCtxWithClient(server)
182 |
183 | _, err := createAnnotationGraphiteFormat(ctx, CreateGraphiteAnnotationInput{
184 | What: "bad",
185 | When: 1700000000000,
186 | })
187 | require.Error(t, err)
188 | assert.Contains(t, err.Error(), "create graphite annotation")
189 | }
190 |
191 | func TestUpdateAnnotation_SendsCorrectBodyAndPath(t *testing.T) {
192 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
193 | assert.Equal(t, "/api/annotations/"+strconv.Itoa(55), r.URL.Path)
194 | assert.Equal(t, "PUT", r.Method)
195 |
196 | var body models.UpdateAnnotationsCmd
197 | _ = json.NewDecoder(r.Body).Decode(&body)
198 |
199 | assert.Equal(t, int64(111), body.Time)
200 | assert.Equal(t, int64(222), body.TimeEnd)
201 | assert.Equal(t, "hello", body.Text)
202 | assert.Equal(t, []string{"a", "b"}, body.Tags)
203 |
204 | w.Header().Set("Content-Type", "application/json")
205 | w.WriteHeader(http.StatusOK)
206 | _, _ = w.Write([]byte(`{}`))
207 | }))
208 | defer server.Close()
209 |
210 | ctx := mockCtxWithClient(server)
211 |
212 | _, err := updateAnnotation(ctx, UpdateAnnotationInput{
213 | ID: 55,
214 | Time: 111,
215 | TimeEnd: 222,
216 | Text: "hello",
217 | Tags: []string{"a", "b"},
218 | })
219 | require.NoError(t, err)
220 | }
221 |
222 | func TestPatchAnnotation_SendsOnlyProvidedFields(t *testing.T) {
223 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
224 | assert.Equal(t, "/api/annotations/"+strconv.Itoa(9), r.URL.Path)
225 | assert.Equal(t, "PATCH", r.Method)
226 |
227 | var body map[string]interface{}
228 | _ = json.NewDecoder(r.Body).Decode(&body)
229 |
230 | assert.Equal(t, "patched", body["text"])
231 | assert.ElementsMatch(t, []interface{}{"x"}, body["tags"].([]interface{}))
232 | assert.Nil(t, body["time"])
233 | assert.Nil(t, body["timeEnd"])
234 |
235 | w.Header().Set("Content-Type", "application/json")
236 | w.WriteHeader(http.StatusOK)
237 | _, _ = w.Write([]byte(`{}`))
238 | }))
239 | defer server.Close()
240 |
241 | ctx := mockCtxWithClient(server)
242 | text := "patched"
243 |
244 | _, err := patchAnnotation(ctx, PatchAnnotationInput{
245 | ID: 9,
246 | Text: &text,
247 | Tags: []string{"x"},
248 | })
249 | require.NoError(t, err)
250 | }
251 |
252 | func TestGetAnnotationTags_UsesCorrectQueryParams(t *testing.T) {
253 | server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
254 | assert.Equal(t, "/api/annotations/tags", r.URL.Path)
255 |
256 | q := r.URL.Query()
257 | assert.Equal(t, "error", q.Get("tag"))
258 | assert.Equal(t, "50", q.Get("limit"))
259 |
260 | w.Header().Set("Content-Type", "application/json")
261 | w.WriteHeader(http.StatusOK)
262 | _, _ = w.Write([]byte(`{"result":{"tags":[]}}`))
263 | }))
264 | defer server.Close()
265 |
266 | ctx := mockCtxWithClient(server)
267 | tag := "error"
268 | limit := "50"
269 |
270 | _, err := getAnnotationTags(ctx, GetAnnotationTagsInput{
271 | Tag: &tag,
272 | Limit: &limit,
273 | })
274 | require.NoError(t, err)
275 | }
276 |
```
--------------------------------------------------------------------------------
/tools/alerting_unit_test.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "testing"
5 |
6 | "github.com/stretchr/testify/require"
7 | )
8 |
9 | // Unit tests for parameter validation (no integration tag needed)
10 | func TestCreateAlertRuleParams_Validate(t *testing.T) {
11 | t.Run("valid parameters", func(t *testing.T) {
12 | params := CreateAlertRuleParams{
13 | Title: "Test Rule",
14 | RuleGroup: "test-group",
15 | FolderUID: "test-folder",
16 | Condition: "A",
17 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
18 | NoDataState: "OK",
19 | ExecErrState: "OK",
20 | For: "5m",
21 | OrgID: 1,
22 | }
23 | err := params.validate()
24 | require.NoError(t, err)
25 | })
26 |
27 | t.Run("missing title", func(t *testing.T) {
28 | params := CreateAlertRuleParams{
29 | RuleGroup: "test-group",
30 | FolderUID: "test-folder",
31 | Condition: "A",
32 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
33 | NoDataState: "OK",
34 | ExecErrState: "OK",
35 | For: "5m",
36 | OrgID: 1,
37 | }
38 | err := params.validate()
39 | require.Error(t, err)
40 | require.Contains(t, err.Error(), "title is required")
41 | })
42 |
43 | t.Run("missing rule group", func(t *testing.T) {
44 | params := CreateAlertRuleParams{
45 | Title: "Test Rule",
46 | FolderUID: "test-folder",
47 | Condition: "A",
48 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
49 | NoDataState: "OK",
50 | ExecErrState: "OK",
51 | For: "5m",
52 | OrgID: 1,
53 | }
54 | err := params.validate()
55 | require.Error(t, err)
56 | require.Contains(t, err.Error(), "ruleGroup is required")
57 | })
58 |
59 | t.Run("missing folder UID", func(t *testing.T) {
60 | params := CreateAlertRuleParams{
61 | Title: "Test Rule",
62 | RuleGroup: "test-group",
63 | Condition: "A",
64 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
65 | NoDataState: "OK",
66 | ExecErrState: "OK",
67 | For: "5m",
68 | OrgID: 1,
69 | }
70 | err := params.validate()
71 | require.Error(t, err)
72 | require.Contains(t, err.Error(), "folderUID is required")
73 | })
74 |
75 | t.Run("missing condition", func(t *testing.T) {
76 | params := CreateAlertRuleParams{
77 | Title: "Test Rule",
78 | RuleGroup: "test-group",
79 | FolderUID: "test-folder",
80 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
81 | NoDataState: "OK",
82 | ExecErrState: "OK",
83 | For: "5m",
84 | OrgID: 1,
85 | }
86 | err := params.validate()
87 | require.Error(t, err)
88 | require.Contains(t, err.Error(), "condition is required")
89 | })
90 |
91 | t.Run("missing data", func(t *testing.T) {
92 | params := CreateAlertRuleParams{
93 | Title: "Test Rule",
94 | RuleGroup: "test-group",
95 | FolderUID: "test-folder",
96 | Condition: "A",
97 | NoDataState: "OK",
98 | ExecErrState: "OK",
99 | For: "5m",
100 | OrgID: 1,
101 | }
102 | err := params.validate()
103 | require.Error(t, err)
104 | require.Contains(t, err.Error(), "data is required")
105 | })
106 |
107 | t.Run("missing no data state", func(t *testing.T) {
108 | params := CreateAlertRuleParams{
109 | Title: "Test Rule",
110 | RuleGroup: "test-group",
111 | FolderUID: "test-folder",
112 | Condition: "A",
113 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
114 | ExecErrState: "OK",
115 | For: "5m",
116 | OrgID: 1,
117 | }
118 | err := params.validate()
119 | require.Error(t, err)
120 | require.Contains(t, err.Error(), "noDataState is required")
121 | })
122 |
123 | t.Run("missing exec error state", func(t *testing.T) {
124 | params := CreateAlertRuleParams{
125 | Title: "Test Rule",
126 | RuleGroup: "test-group",
127 | FolderUID: "test-folder",
128 | Condition: "A",
129 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
130 | NoDataState: "OK",
131 | For: "5m",
132 | OrgID: 1,
133 | }
134 | err := params.validate()
135 | require.Error(t, err)
136 | require.Contains(t, err.Error(), "execErrState is required")
137 | })
138 |
139 | t.Run("missing for duration", func(t *testing.T) {
140 | params := CreateAlertRuleParams{
141 | Title: "Test Rule",
142 | RuleGroup: "test-group",
143 | FolderUID: "test-folder",
144 | Condition: "A",
145 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
146 | NoDataState: "OK",
147 | ExecErrState: "OK",
148 | OrgID: 1,
149 | }
150 | err := params.validate()
151 | require.Error(t, err)
152 | require.Contains(t, err.Error(), "for duration is required")
153 | })
154 |
155 | t.Run("invalid org ID", func(t *testing.T) {
156 | params := CreateAlertRuleParams{
157 | Title: "Test Rule",
158 | RuleGroup: "test-group",
159 | FolderUID: "test-folder",
160 | Condition: "A",
161 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
162 | NoDataState: "OK",
163 | ExecErrState: "OK",
164 | For: "5m",
165 | OrgID: 0,
166 | }
167 | err := params.validate()
168 | require.Error(t, err)
169 | require.Contains(t, err.Error(), "orgID is required and must be greater than 0")
170 | })
171 | }
172 |
173 | func TestUpdateAlertRuleParams_Validate(t *testing.T) {
174 | t.Run("valid parameters", func(t *testing.T) {
175 | params := UpdateAlertRuleParams{
176 | UID: "test-uid",
177 | Title: "Test Rule",
178 | RuleGroup: "test-group",
179 | FolderUID: "test-folder",
180 | Condition: "A",
181 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
182 | NoDataState: "OK",
183 | ExecErrState: "OK",
184 | For: "5m",
185 | OrgID: 1,
186 | }
187 | err := params.validate()
188 | require.NoError(t, err)
189 | })
190 |
191 | t.Run("missing UID", func(t *testing.T) {
192 | params := UpdateAlertRuleParams{
193 | Title: "Test Rule",
194 | RuleGroup: "test-group",
195 | FolderUID: "test-folder",
196 | Condition: "A",
197 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
198 | NoDataState: "OK",
199 | ExecErrState: "OK",
200 | For: "5m",
201 | OrgID: 1,
202 | }
203 | err := params.validate()
204 | require.Error(t, err)
205 | require.Contains(t, err.Error(), "uid is required")
206 | })
207 |
208 | t.Run("invalid org ID", func(t *testing.T) {
209 | params := UpdateAlertRuleParams{
210 | UID: "test-uid",
211 | Title: "Test Rule",
212 | RuleGroup: "test-group",
213 | FolderUID: "test-folder",
214 | Condition: "A",
215 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
216 | NoDataState: "OK",
217 | ExecErrState: "OK",
218 | For: "5m",
219 | OrgID: -1,
220 | }
221 | err := params.validate()
222 | require.Error(t, err)
223 | require.Contains(t, err.Error(), "orgID is required and must be greater than 0")
224 | })
225 | }
226 |
227 | func TestDeleteAlertRuleParams_Validate(t *testing.T) {
228 | t.Run("valid parameters", func(t *testing.T) {
229 | params := DeleteAlertRuleParams{
230 | UID: "test-uid",
231 | }
232 | err := params.validate()
233 | require.NoError(t, err)
234 | })
235 |
236 | t.Run("missing UID", func(t *testing.T) {
237 | params := DeleteAlertRuleParams{
238 | UID: "",
239 | }
240 | err := params.validate()
241 | require.Error(t, err)
242 | require.Contains(t, err.Error(), "uid is required")
243 | })
244 | }
245 |
246 | func TestBuiltInValidationCatchesInvalidData(t *testing.T) {
247 | t.Run("invalid NoDataState enum value", func(t *testing.T) {
248 | params := CreateAlertRuleParams{
249 | Title: "Test Rule",
250 | RuleGroup: "test-group",
251 | FolderUID: "test-folder",
252 | Condition: "A",
253 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
254 | NoDataState: "InvalidValue", // Invalid enum
255 | ExecErrState: "OK",
256 | For: "5m",
257 | OrgID: 1,
258 | }
259 |
260 | // Our simple validation won't catch this, but it would fail at API call
261 | err := params.validate()
262 | require.NoError(t, err, "Simple validation doesn't check enum values")
263 | })
264 |
265 | t.Run("invalid ExecErrState enum value", func(t *testing.T) {
266 | params := CreateAlertRuleParams{
267 | Title: "Test Rule",
268 | RuleGroup: "test-group",
269 | FolderUID: "test-folder",
270 | Condition: "A",
271 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
272 | NoDataState: "OK",
273 | ExecErrState: "BadValue", // Invalid enum
274 | For: "5m",
275 | OrgID: 1,
276 | }
277 |
278 | // Our simple validation won't catch this
279 | err := params.validate()
280 | require.NoError(t, err, "Simple validation doesn't check enum values")
281 | })
282 |
283 | t.Run("title too long", func(t *testing.T) {
284 | longTitle := make([]byte, 200) // Max is 190
285 | for i := range longTitle {
286 | longTitle[i] = 'A'
287 | }
288 |
289 | params := CreateAlertRuleParams{
290 | Title: string(longTitle),
291 | RuleGroup: "test-group",
292 | FolderUID: "test-folder",
293 | Condition: "A",
294 | Data: []interface{}{map[string]interface{}{"refId": "A"}},
295 | NoDataState: "OK",
296 | ExecErrState: "OK",
297 | For: "5m",
298 | OrgID: 1,
299 | }
300 |
301 | // Simple validation only checks if title is empty, not length
302 | err := params.validate()
303 | require.NoError(t, err, "Simple validation doesn't check length constraints")
304 | })
305 | }
306 |
```
--------------------------------------------------------------------------------
/tools.go:
--------------------------------------------------------------------------------
```go
1 | package mcpgrafana
2 |
3 | import (
4 | "context"
5 | "encoding/json"
6 | "errors"
7 | "fmt"
8 | "reflect"
9 |
10 | "github.com/invopop/jsonschema"
11 | "github.com/mark3labs/mcp-go/mcp"
12 | "github.com/mark3labs/mcp-go/server"
13 | "go.opentelemetry.io/otel"
14 | "go.opentelemetry.io/otel/attribute"
15 | "go.opentelemetry.io/otel/codes"
16 | )
17 |
18 | // Tool represents a tool definition and its handler function for the MCP server.
19 | // It encapsulates both the tool metadata (name, description, schema) and the function that executes when the tool is called.
20 | // The simplest way to create a Tool is to use MustTool for compile-time tool creation,
21 | // or ConvertTool if you need runtime tool creation with proper error handling.
22 | type Tool struct {
23 | Tool mcp.Tool
24 | Handler server.ToolHandlerFunc
25 | }
26 |
27 | // Register adds the Tool to the given MCPServer.
28 | // It is a convenience method that calls server.MCPServer.AddTool with the Tool's metadata and handler,
29 | // allowing fluent tool registration in a single statement:
30 | //
31 | // mcpgrafana.MustTool(name, description, toolHandler).Register(server)
32 | func (t *Tool) Register(mcp *server.MCPServer) {
33 | mcp.AddTool(t.Tool, t.Handler)
34 | }
35 |
36 | // MustTool creates a new Tool from the given name, description, and toolHandler.
37 | // It panics if the tool cannot be created, making it suitable for compile-time tool definitions where creation errors indicate programming mistakes.
38 | func MustTool[T any, R any](
39 | name, description string,
40 | toolHandler ToolHandlerFunc[T, R],
41 | options ...mcp.ToolOption,
42 | ) Tool {
43 | tool, handler, err := ConvertTool(name, description, toolHandler, options...)
44 | if err != nil {
45 | panic(err)
46 | }
47 | return Tool{Tool: tool, Handler: handler}
48 | }
49 |
50 | // ToolHandlerFunc is the type of a handler function for a tool.
51 | // T is the request parameter type (must be a struct with jsonschema tags), and R is the response type which can be a string, struct, or *mcp.CallToolResult.
52 | type ToolHandlerFunc[T any, R any] = func(ctx context.Context, request T) (R, error)
53 |
54 | // ConvertTool converts a toolHandler function to an MCP Tool and ToolHandlerFunc.
55 | // The toolHandler must accept a context.Context and a struct with jsonschema tags for parameter documentation.
56 | // The struct fields define the tool's input schema, while the return value can be a string, struct, or *mcp.CallToolResult.
57 | // This function automatically generates JSON schema from the struct type and wraps the handler with OpenTelemetry instrumentation.
58 | func ConvertTool[T any, R any](name, description string, toolHandler ToolHandlerFunc[T, R], options ...mcp.ToolOption) (mcp.Tool, server.ToolHandlerFunc, error) {
59 | zero := mcp.Tool{}
60 | handlerValue := reflect.ValueOf(toolHandler)
61 | handlerType := handlerValue.Type()
62 | if handlerType.Kind() != reflect.Func {
63 | return zero, nil, errors.New("tool handler must be a function")
64 | }
65 | if handlerType.NumIn() != 2 {
66 | return zero, nil, errors.New("tool handler must have 2 arguments")
67 | }
68 | if handlerType.NumOut() != 2 {
69 | return zero, nil, errors.New("tool handler must return 2 values")
70 | }
71 | if handlerType.In(0) != reflect.TypeOf((*context.Context)(nil)).Elem() {
72 | return zero, nil, errors.New("tool handler first argument must be context.Context")
73 | }
74 | // We no longer check the type of the first return value
75 | if handlerType.Out(1).Kind() != reflect.Interface {
76 | return zero, nil, errors.New("tool handler second return value must be error")
77 | }
78 |
79 | argType := handlerType.In(1)
80 | if argType.Kind() != reflect.Struct {
81 | return zero, nil, errors.New("tool handler second argument must be a struct")
82 | }
83 |
84 | handler := func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
85 | // Create OpenTelemetry span for tool execution (no-op when no exporter configured)
86 | config := GrafanaConfigFromContext(ctx)
87 | ctx, span := otel.Tracer("mcp-grafana").Start(ctx, fmt.Sprintf("mcp.tool.%s", name))
88 | defer span.End()
89 |
90 | // Add tool metadata as span attributes
91 | span.SetAttributes(
92 | attribute.String("mcp.tool.name", name),
93 | attribute.String("mcp.tool.description", description),
94 | )
95 |
96 | argBytes, err := json.Marshal(request.Params.Arguments)
97 | if err != nil {
98 | span.RecordError(err)
99 | span.SetStatus(codes.Error, "failed to marshal arguments")
100 | return nil, fmt.Errorf("marshal args: %w", err)
101 | }
102 |
103 | // Add arguments as span attribute only if adding args to trace attributes is enabled
104 | if config.IncludeArgumentsInSpans {
105 | span.SetAttributes(attribute.String("mcp.tool.arguments", string(argBytes)))
106 | }
107 |
108 | unmarshaledArgs := reflect.New(argType).Interface()
109 | if err := json.Unmarshal(argBytes, unmarshaledArgs); err != nil {
110 | span.RecordError(err)
111 | span.SetStatus(codes.Error, "failed to unmarshal arguments")
112 | return nil, fmt.Errorf("unmarshal args: %s", err)
113 | }
114 |
115 | // Need to dereference the unmarshaled arguments
116 | of := reflect.ValueOf(unmarshaledArgs)
117 | if of.Kind() != reflect.Ptr || !of.Elem().CanInterface() {
118 | err := errors.New("arguments must be a struct")
119 | span.RecordError(err)
120 | span.SetStatus(codes.Error, "invalid arguments structure")
121 | return nil, err
122 | }
123 |
124 | // Pass the instrumented context to the tool handler
125 | args := []reflect.Value{reflect.ValueOf(ctx), of.Elem()}
126 |
127 | output := handlerValue.Call(args)
128 | if len(output) != 2 {
129 | err := errors.New("tool handler must return 2 values")
130 | span.RecordError(err)
131 | span.SetStatus(codes.Error, "invalid tool handler return")
132 | return nil, err
133 | }
134 | if !output[0].CanInterface() {
135 | err := errors.New("tool handler first return value must be interfaceable")
136 | span.RecordError(err)
137 | span.SetStatus(codes.Error, "tool handler return value not interfaceable")
138 | return nil, err
139 | }
140 |
141 | // Handle the error return value first
142 | var handlerErr error
143 | var ok bool
144 | if output[1].Kind() == reflect.Interface && !output[1].IsNil() {
145 | handlerErr, ok = output[1].Interface().(error)
146 | if !ok {
147 | err := errors.New("tool handler second return value must be error")
148 | span.RecordError(err)
149 | span.SetStatus(codes.Error, "invalid error return type")
150 | return nil, err
151 | }
152 | }
153 |
154 | // If there's an error, record it and return
155 | if handlerErr != nil {
156 | span.RecordError(handlerErr)
157 | span.SetStatus(codes.Error, handlerErr.Error())
158 | return nil, handlerErr
159 | }
160 |
161 | // Tool execution completed successfully
162 | span.SetStatus(codes.Ok, "tool execution completed")
163 |
164 | // Check if the first return value is nil (only for pointer, interface, map, etc.)
165 | isNilable := output[0].Kind() == reflect.Ptr ||
166 | output[0].Kind() == reflect.Interface ||
167 | output[0].Kind() == reflect.Map ||
168 | output[0].Kind() == reflect.Slice ||
169 | output[0].Kind() == reflect.Chan ||
170 | output[0].Kind() == reflect.Func
171 |
172 | if isNilable && output[0].IsNil() {
173 | return nil, nil
174 | }
175 |
176 | returnVal := output[0].Interface()
177 | returnType := output[0].Type()
178 |
179 | // Case 1: Already a *mcp.CallToolResult
180 | if callResult, ok := returnVal.(*mcp.CallToolResult); ok {
181 | return callResult, nil
182 | }
183 |
184 | // Case 2: An mcp.CallToolResult (not a pointer)
185 | if returnType.ConvertibleTo(reflect.TypeOf(mcp.CallToolResult{})) {
186 | callResult := returnVal.(mcp.CallToolResult)
187 | return &callResult, nil
188 | }
189 |
190 | // Case 3: String or *string
191 | if str, ok := returnVal.(string); ok {
192 | if str == "" {
193 | return nil, nil
194 | }
195 | return mcp.NewToolResultText(str), nil
196 | }
197 |
198 | if strPtr, ok := returnVal.(*string); ok {
199 | if strPtr == nil || *strPtr == "" {
200 | return nil, nil
201 | }
202 | return mcp.NewToolResultText(*strPtr), nil
203 | }
204 |
205 | // Case 4: Any other type - marshal to JSON
206 | returnBytes, err := json.Marshal(returnVal)
207 | if err != nil {
208 | return nil, fmt.Errorf("failed to marshal return value: %s", err)
209 | }
210 |
211 | return mcp.NewToolResultText(string(returnBytes)), nil
212 | }
213 |
214 | jsonSchema := createJSONSchemaFromHandler(toolHandler)
215 | properties := make(map[string]any, jsonSchema.Properties.Len())
216 | for pair := jsonSchema.Properties.Oldest(); pair != nil; pair = pair.Next() {
217 | properties[pair.Key] = pair.Value
218 | }
219 | inputSchema := mcp.ToolInputSchema{
220 | Type: jsonSchema.Type,
221 | Properties: properties,
222 | Required: jsonSchema.Required,
223 | }
224 |
225 | t := mcp.Tool{
226 | Name: name,
227 | Description: description,
228 | InputSchema: inputSchema,
229 | }
230 | for _, option := range options {
231 | option(&t)
232 | }
233 | return t, handler, nil
234 | }
235 |
236 | // Creates a full JSON schema from a user provided handler by introspecting the arguments
237 | func createJSONSchemaFromHandler(handler any) *jsonschema.Schema {
238 | handlerValue := reflect.ValueOf(handler)
239 | handlerType := handlerValue.Type()
240 | argumentType := handlerType.In(1)
241 | inputSchema := jsonSchemaReflector.ReflectFromType(argumentType)
242 | return inputSchema
243 | }
244 |
245 | var (
246 | jsonSchemaReflector = jsonschema.Reflector{
247 | BaseSchemaID: "",
248 | Anonymous: true,
249 | AssignAnchor: false,
250 | AllowAdditionalProperties: true,
251 | RequiredFromJSONSchemaTags: true,
252 | DoNotReference: true,
253 | ExpandedStruct: true,
254 | FieldNameTag: "",
255 | IgnoredTypes: nil,
256 | Lookup: nil,
257 | Mapper: nil,
258 | Namer: nil,
259 | KeyNamer: nil,
260 | AdditionalFields: nil,
261 | CommentMap: nil,
262 | }
263 | )
264 |
```
--------------------------------------------------------------------------------
/tools/prometheus_test.go:
--------------------------------------------------------------------------------
```go
1 | // Requires a Grafana instance running on localhost:3000,
2 | // with a Prometheus datasource provisioned.
3 | // Run with `go test -tags integration`.
4 | //go:build integration
5 |
6 | package tools
7 |
8 | import (
9 | "fmt"
10 | "testing"
11 | "time"
12 |
13 | "github.com/prometheus/common/model"
14 | "github.com/prometheus/prometheus/model/labels"
15 | "github.com/stretchr/testify/assert"
16 | "github.com/stretchr/testify/require"
17 | )
18 |
19 | func TestPrometheusTools(t *testing.T) {
20 | t.Run("list prometheus metric metadata", func(t *testing.T) {
21 | ctx := newTestContext()
22 | result, err := listPrometheusMetricMetadata(ctx, ListPrometheusMetricMetadataParams{
23 | DatasourceUID: "prometheus",
24 | })
25 | require.NoError(t, err)
26 | assert.Len(t, result, 10)
27 | })
28 |
29 | t.Run("list prometheus metric names", func(t *testing.T) {
30 | ctx := newTestContext()
31 | result, err := listPrometheusMetricNames(ctx, ListPrometheusMetricNamesParams{
32 | DatasourceUID: "prometheus",
33 | Regex: ".*",
34 | Limit: 10,
35 | })
36 | require.NoError(t, err)
37 | assert.Len(t, result, 10)
38 | })
39 |
40 | t.Run("list prometheus label names", func(t *testing.T) {
41 | ctx := newTestContext()
42 | result, err := listPrometheusLabelNames(ctx, ListPrometheusLabelNamesParams{
43 | DatasourceUID: "prometheus",
44 | Matches: []Selector{
45 | {
46 | Filters: []LabelMatcher{
47 | {Name: "job", Value: "prometheus"},
48 | },
49 | },
50 | },
51 | Limit: 10,
52 | })
53 | require.NoError(t, err)
54 | assert.Len(t, result, 10)
55 | })
56 |
57 | t.Run("list prometheus label values", func(t *testing.T) {
58 | ctx := newTestContext()
59 | result, err := listPrometheusLabelValues(ctx, ListPrometheusLabelValuesParams{
60 | DatasourceUID: "prometheus",
61 | LabelName: "job",
62 | Matches: []Selector{
63 | {
64 | Filters: []LabelMatcher{
65 | {Name: "job", Value: "prometheus"},
66 | },
67 | },
68 | },
69 | })
70 | require.NoError(t, err)
71 | assert.Len(t, result, 1)
72 | })
73 | }
74 |
75 | func TestSelectorMatches(t *testing.T) {
76 | testCases := []struct {
77 | name string
78 | selector Selector
79 | labels map[string]string
80 | expected bool
81 | expectErr bool
82 | }{
83 | {
84 | name: "Equal match",
85 | selector: Selector{
86 | Filters: []LabelMatcher{
87 | {Name: "job", Type: "=", Value: "prometheus"},
88 | },
89 | },
90 | labels: map[string]string{"job": "prometheus"},
91 | expected: true,
92 | },
93 | {
94 | name: "Equal no match",
95 | selector: Selector{
96 | Filters: []LabelMatcher{
97 | {Name: "job", Type: "=", Value: "prometheus"},
98 | },
99 | },
100 | labels: map[string]string{"job": "node-exporter"},
101 | expected: false,
102 | },
103 | {
104 | name: "Not equal match",
105 | selector: Selector{
106 | Filters: []LabelMatcher{
107 | {Name: "job", Type: "!=", Value: "prometheus"},
108 | },
109 | },
110 | labels: map[string]string{"job": "node-exporter"},
111 | expected: true,
112 | },
113 | {
114 | name: "Not equal no match",
115 | selector: Selector{
116 | Filters: []LabelMatcher{
117 | {Name: "job", Type: "!=", Value: "prometheus"},
118 | },
119 | },
120 | labels: map[string]string{"job": "prometheus"},
121 | expected: false,
122 | },
123 | {
124 | name: "Regex match",
125 | selector: Selector{
126 | Filters: []LabelMatcher{
127 | {Name: "job", Type: "=~", Value: "prom.*"},
128 | },
129 | },
130 | labels: map[string]string{"job": "prometheus"},
131 | expected: true,
132 | },
133 | {
134 | name: "Regex no match",
135 | selector: Selector{
136 | Filters: []LabelMatcher{
137 | {Name: "job", Type: "=~", Value: "node.*"},
138 | },
139 | },
140 | labels: map[string]string{"job": "prometheus"},
141 | expected: false,
142 | },
143 | {
144 | name: "Not regex match",
145 | selector: Selector{
146 | Filters: []LabelMatcher{
147 | {Name: "job", Type: "!~", Value: "node.*"},
148 | },
149 | },
150 | labels: map[string]string{"job": "prometheus"},
151 | expected: true,
152 | },
153 | {
154 | name: "Not regex no match",
155 | selector: Selector{
156 | Filters: []LabelMatcher{
157 | {Name: "job", Type: "!~", Value: "prom.*"},
158 | },
159 | },
160 | labels: map[string]string{"job": "prometheus"},
161 | expected: false,
162 | },
163 | {
164 | name: "Multiple filters all match",
165 | selector: Selector{
166 | Filters: []LabelMatcher{
167 | {Name: "job", Type: "=", Value: "prometheus"},
168 | {Name: "instance", Type: "=~", Value: "localhost.*"},
169 | },
170 | },
171 | labels: map[string]string{"job": "prometheus", "instance": "localhost:9090"},
172 | expected: true,
173 | },
174 | {
175 | name: "Multiple filters one doesn't match",
176 | selector: Selector{
177 | Filters: []LabelMatcher{
178 | {Name: "job", Type: "=", Value: "prometheus"},
179 | {Name: "instance", Type: "=~", Value: "remote.*"},
180 | },
181 | },
182 | labels: map[string]string{"job": "prometheus", "instance": "localhost:9090"},
183 | expected: false,
184 | },
185 | {
186 | name: "Label doesn't exist with = operator",
187 | selector: Selector{
188 | Filters: []LabelMatcher{
189 | {Name: "missing", Type: "=", Value: "value"},
190 | },
191 | },
192 | labels: map[string]string{"job": "prometheus"},
193 | expected: false,
194 | },
195 | {
196 | name: "Label doesn't exist with != operator",
197 | selector: Selector{
198 | Filters: []LabelMatcher{
199 | {Name: "missing", Type: "!=", Value: "value"},
200 | },
201 | },
202 | labels: map[string]string{"job": "prometheus"},
203 | expected: true,
204 | },
205 | {
206 | name: "Invalid matcher type",
207 | selector: Selector{
208 | Filters: []LabelMatcher{
209 | {Name: "job", Type: "<>", Value: "prometheus"},
210 | },
211 | },
212 | labels: map[string]string{"job": "prometheus"},
213 | expected: false,
214 | expectErr: true,
215 | },
216 | }
217 |
218 | for _, tc := range testCases {
219 | t.Run(tc.name, func(t *testing.T) {
220 | lbls := labels.FromMap(tc.labels)
221 | result, err := tc.selector.Matches(lbls)
222 |
223 | if tc.expectErr {
224 | assert.Error(t, err)
225 | return
226 | }
227 |
228 | assert.NoError(t, err)
229 | assert.Equal(t, tc.expected, result)
230 | })
231 | }
232 | }
233 |
234 | func TestPrometheusQueries(t *testing.T) {
235 | t.Run("query prometheus range", func(t *testing.T) {
236 | end := time.Now()
237 | start := end.Add(-10 * time.Minute)
238 | for _, step := range []int{15, 60, 300} {
239 | t.Run(fmt.Sprintf("step=%d", step), func(t *testing.T) {
240 | ctx := newTestContext()
241 | result, err := queryPrometheus(ctx, QueryPrometheusParams{
242 | DatasourceUID: "prometheus",
243 | Expr: "test",
244 | StartTime: start.Format(time.RFC3339),
245 | EndTime: end.Format(time.RFC3339),
246 | StepSeconds: step,
247 | QueryType: "range",
248 | })
249 | require.NoError(t, err)
250 | matrix := result.(model.Matrix)
251 | require.Len(t, matrix, 1)
252 | expectedLen := int(end.Sub(start).Seconds()/float64(step)) + 1
253 | assert.Len(t, matrix[0].Values, expectedLen)
254 | assert.Less(t, matrix[0].Values[0].Timestamp.Sub(model.TimeFromUnix(start.Unix())), time.Duration(step)*time.Second)
255 | assert.Equal(t, matrix[0].Metric["__name__"], model.LabelValue("test"))
256 | })
257 | }
258 | })
259 |
260 | t.Run("query prometheus instant", func(t *testing.T) {
261 | ctx := newTestContext()
262 | result, err := queryPrometheus(ctx, QueryPrometheusParams{
263 | DatasourceUID: "prometheus",
264 | Expr: "up",
265 | StartTime: time.Now().Format(time.RFC3339),
266 | QueryType: "instant",
267 | })
268 | require.NoError(t, err)
269 | scalar := result.(model.Vector)
270 | assert.Equal(t, scalar[0].Value, model.SampleValue(1))
271 | assert.Equal(t, scalar[0].Timestamp, model.TimeFromUnix(time.Now().Unix()))
272 | assert.Equal(t, scalar[0].Metric["__name__"], model.LabelValue("up"))
273 | })
274 |
275 | t.Run("query prometheus instant with relative timestamps", func(t *testing.T) {
276 | ctx := newTestContext()
277 | beforeQuery := model.TimeFromUnix(time.Now().Unix())
278 | result, err := queryPrometheus(ctx, QueryPrometheusParams{
279 | DatasourceUID: "prometheus",
280 | Expr: "up",
281 | StartTime: "now",
282 | QueryType: "instant",
283 | })
284 | afterQuery := model.TimeFromUnix(time.Now().Unix())
285 | require.NoError(t, err)
286 | scalar := result.(model.Vector)
287 | assert.Equal(t, scalar[0].Value, model.SampleValue(1))
288 |
289 | // Check that the timestamp is within the expected range
290 | buffer := 5 * time.Second
291 | assert.True(t, scalar[0].Timestamp >= beforeQuery,
292 | "Result timestamp should be after or equal to the time before the query")
293 | assert.True(t, scalar[0].Timestamp <= afterQuery.Add(buffer),
294 | "Result timestamp should be before or equal to the time after the query (with 5s buffer)")
295 |
296 | assert.Equal(t, scalar[0].Metric["__name__"], model.LabelValue("up"))
297 | })
298 |
299 | t.Run("query prometheus range with relative timestamps", func(t *testing.T) {
300 | ctx := newTestContext()
301 | beforeQuery := model.TimeFromUnix(time.Now().Unix())
302 | result, err := queryPrometheus(ctx, QueryPrometheusParams{
303 | DatasourceUID: "prometheus",
304 | Expr: "test",
305 | StartTime: "now-1h",
306 | EndTime: "now",
307 | StepSeconds: 60,
308 | QueryType: "range",
309 | })
310 | afterQuery := model.TimeFromUnix(time.Now().Unix())
311 | require.NoError(t, err)
312 | matrix := result.(model.Matrix)
313 | require.Len(t, matrix, 1)
314 |
315 | // Should have approximately 60 samples (one per minute for an hour)
316 | assert.InDelta(t, 60, len(matrix[0].Values), 2)
317 |
318 | buffer := 5 * time.Second
319 | oneHour := time.Hour
320 |
321 | firstSampleTime := matrix[0].Values[0].Timestamp
322 | // Check that the start timestamp is within the expected range
323 | assert.True(t, firstSampleTime >= beforeQuery.Add(-oneHour),
324 | "First timestamp should be after or equal to the time before the query minus one hour")
325 | assert.True(t, firstSampleTime <= afterQuery.Add(buffer).Add(-oneHour),
326 | "First timestamp should be before or equal to the time after the query minus one hour (with 5s buffer)")
327 |
328 | // Check that the end timestamp is is within the expected range
329 | lastSampleTime := matrix[0].Values[len(matrix[0].Values)-1].Timestamp
330 | assert.True(t, lastSampleTime >= beforeQuery,
331 | "Last timestamp should be after or equal to the time before the query")
332 | assert.True(t, lastSampleTime <= afterQuery.Add(buffer),
333 | "Last timestamp should be before or equal to the time after the query (with 5s buffer)")
334 |
335 | assert.Equal(t, matrix[0].Metric["__name__"], model.LabelValue("test"))
336 | })
337 | }
338 |
```
--------------------------------------------------------------------------------
/tests/tempo_test.py:
--------------------------------------------------------------------------------
```python
1 | from mcp import ClientSession
2 | import pytest
3 | from langevals import expect
4 | from langevals_langevals.llm_boolean import (
5 | CustomLLMBooleanEvaluator,
6 | CustomLLMBooleanSettings,
7 | )
8 | from litellm import Message, acompletion
9 | from mcp import ClientSession
10 |
11 | from conftest import models
12 | from utils import (
13 | get_converted_tools,
14 | llm_tool_call_sequence,
15 | )
16 |
17 | pytestmark = pytest.mark.anyio
18 |
19 |
20 | class TestTempoProxiedToolsBasic:
21 | """Test Tempo proxied MCP tools functionality.
22 |
23 | These tests verify that Tempo datasources with MCP support are discovered
24 | per-session and their tools are registered with a datasourceUid parameter
25 | for multi-datasource support.
26 |
27 | Requires:
28 | - Docker compose services running (includes 2 Tempo instances)
29 | - GRAFANA_USERNAME and GRAFANA_PASSWORD environment variables
30 | - MCP server running
31 | """
32 |
33 | @pytest.mark.anyio
34 | async def test_tempo_tools_discovered_and_registered(
35 | self, mcp_client: ClientSession
36 | ):
37 | """Test that Tempo tools are discovered and registered with datasourceUid parameter."""
38 |
39 | # List all tools
40 | list_response = await mcp_client.list_tools()
41 | all_tool_names = [tool.name for tool in list_response.tools]
42 |
43 | # Find tempo-prefixed tools (should preserve hyphens from original tool names)
44 | tempo_tools = [name for name in all_tool_names if name.startswith("tempo_")]
45 |
46 | # Expected tools from Tempo MCP server
47 | expected_tempo_tools = [
48 | "tempo_traceql-search",
49 | "tempo_traceql-metrics-instant",
50 | "tempo_traceql-metrics-range",
51 | "tempo_get-trace",
52 | "tempo_get-attribute-names",
53 | "tempo_get-attribute-values",
54 | "tempo_docs-traceql",
55 | ]
56 |
57 | assert len(tempo_tools) == len(expected_tempo_tools), (
58 | f"Expected {len(expected_tempo_tools)} unique tempo tools, found {len(tempo_tools)}: {tempo_tools}"
59 | )
60 |
61 | for expected_tool in expected_tempo_tools:
62 | assert expected_tool in tempo_tools, (
63 | f"Tool {expected_tool} should be available"
64 | )
65 |
66 | @pytest.mark.anyio
67 | async def test_tempo_tools_have_datasourceUid_parameter(self, mcp_client):
68 | """Test that all tempo tools have a required datasourceUid parameter."""
69 |
70 | list_response = await mcp_client.list_tools()
71 | tempo_tools = [
72 | tool for tool in list_response.tools if tool.name.startswith("tempo_")
73 | ]
74 |
75 | assert len(tempo_tools) > 0, "Should have at least one tempo tool"
76 |
77 | for tool in tempo_tools:
78 | # Verify the tool has input schema
79 | assert hasattr(tool, "inputSchema"), (
80 | f"Tool {tool.name} should have inputSchema"
81 | )
82 | assert isinstance(tool.inputSchema, dict), (
83 | f"Tool {tool.name} inputSchema should be a dict"
84 | )
85 |
86 | # Verify datasourceUid parameter exists (camelCase)
87 | properties = tool.inputSchema.get("properties", {})
88 | assert "datasourceUid" in properties, (
89 | f"Tool {tool.name} should have datasourceUid parameter (camelCase)"
90 | )
91 |
92 | # Verify it's required
93 | required = tool.inputSchema.get("required", [])
94 | assert "datasourceUid" in required, (
95 | f"Tool {tool.name} should require datasourceUid parameter"
96 | )
97 |
98 | # Verify parameter has proper description
99 | datasource_uid_prop = properties["datasourceUid"]
100 | assert "type" in datasource_uid_prop, (
101 | f"datasourceUid should have type defined"
102 | )
103 | assert datasource_uid_prop["type"] == "string", (
104 | f"datasourceUid should be type string"
105 | )
106 |
107 | @pytest.mark.anyio
108 | async def test_tempo_tool_call_with_valid_datasource(self, mcp_client):
109 | """Test calling a tempo tool with a valid datasourceUid."""
110 |
111 | # Call docs-traceql which should return documentation (doesn't require data)
112 | try:
113 | call_response = await mcp_client.call_tool(
114 | "tempo_docs-traceql",
115 | arguments={"datasourceUid": "tempo", "name": "basic"},
116 | )
117 |
118 | # Verify we got a response
119 | assert call_response.content, "Tool should return content"
120 |
121 | # Should have text content (documentation)
122 | response_text = call_response.content[0].text
123 | assert len(response_text) > 0, "Response should have content"
124 | assert "traceql" in response_text.lower(), (
125 | "Response should contain TraceQL documentation"
126 | )
127 | print(response_text)
128 |
129 | except Exception as e:
130 | # If this fails, it might be because Tempo doesn't have data yet
131 | # but at least verify the error isn't about missing datasourceUid
132 | error_msg = str(e).lower()
133 | assert "datasourceuid" not in error_msg, (
134 | f"Should not fail due to datasourceUid parameter: {e}"
135 | )
136 | print(error_msg)
137 |
138 | @pytest.mark.anyio
139 | async def test_tempo_tool_call_missing_datasourceUid(self, mcp_client):
140 | """Test that calling a tempo tool without datasourceUid fails appropriately."""
141 |
142 | with pytest.raises(Exception) as exc_info:
143 | await mcp_client.call_tool(
144 | "tempo_docs-traceql",
145 | arguments={"name": "basic"}, # Missing datasourceUid
146 | )
147 |
148 | error_msg = str(exc_info.value).lower()
149 | assert "datasourceuid" in error_msg and "required" in error_msg, (
150 | f"Should require datasourceUid parameter: {exc_info.value}"
151 | )
152 |
153 | @pytest.mark.anyio
154 | async def test_tempo_tool_call_invalid_datasourceUid(self, mcp_client):
155 | """Test that calling a tempo tool with invalid datasourceUid returns helpful error."""
156 |
157 | with pytest.raises(Exception) as exc_info:
158 | await mcp_client.call_tool(
159 | "tempo_docs-traceql",
160 | arguments={"datasourceUid": "nonexistent-tempo", "name": "basic"},
161 | )
162 |
163 | error_msg = str(exc_info.value).lower()
164 | # Should mention that datasource wasn't found
165 | assert "not found" in error_msg or "not accessible" in error_msg, (
166 | f"Should indicate datasource not found: {exc_info.value}"
167 | )
168 |
169 | # Should mention available datasources to help user
170 | assert "tempo" in error_msg or "available" in error_msg, (
171 | f"Error should be helpful and mention available datasources: {exc_info.value}"
172 | )
173 |
174 | @pytest.mark.anyio
175 | async def test_tempo_tool_works_with_multiple_datasources(self, mcp_client):
176 | """Test that the same tool works with different datasources via datasourceUid."""
177 |
178 | # Both tempo and tempo-secondary should be available in our test environment
179 | datasources = ["tempo", "tempo-secondary"]
180 |
181 | for datasource_uid in datasources:
182 | try:
183 | # Call the same tool with different datasources
184 | call_response = await mcp_client.call_tool(
185 | "tempo_get-attribute-names",
186 | arguments={"datasourceUid": datasource_uid},
187 | )
188 |
189 | # Verify we got a response
190 | assert call_response.content, (
191 | f"Tool should return content for datasource {datasource_uid}"
192 | )
193 |
194 | # Response should be valid JSON or text
195 | response_text = call_response.content[0].text
196 | assert len(response_text) > 0, (
197 | f"Response should have content for datasource {datasource_uid}"
198 | )
199 |
200 | except Exception as e:
201 | # If this fails, it's acceptable if Tempo doesn't have trace data yet
202 | # But verify it's not a routing/config error
203 | error_msg = str(e).lower()
204 | assert (
205 | "not found" not in error_msg or datasource_uid not in error_msg
206 | ), f"Datasource {datasource_uid} should be accessible: {e}"
207 |
208 |
209 | class TestTempoProxiedToolsWithLLM:
210 | """LLM integration tests for Tempo proxied tools."""
211 |
212 | @pytest.mark.parametrize("model", models)
213 | @pytest.mark.flaky(max_runs=3)
214 | async def test_llm_can_list_trace_attributes(
215 | self, model: str, mcp_client: ClientSession
216 | ):
217 | """Test that an LLM can list available trace attributes from Tempo."""
218 | tools = await get_converted_tools(mcp_client)
219 | prompt = (
220 | "Use the tempo tools to get a list of all available trace attribute names "
221 | "from the datasource with UID 'tempo'. I want to know what attributes "
222 | "I can use in my TraceQL queries."
223 | )
224 |
225 | messages = [
226 | Message(role="system", content="You are a helpful assistant."),
227 | Message(role="user", content=prompt),
228 | ]
229 |
230 | # LLM should call tempo_get-attribute-names with datasourceUid
231 | messages = await llm_tool_call_sequence(
232 | model,
233 | messages,
234 | tools,
235 | mcp_client,
236 | "tempo_get-attribute-names",
237 | {"datasourceUid": "tempo"},
238 | )
239 |
240 | # Final LLM response should mention attributes
241 | response = await acompletion(model=model, messages=messages, tools=tools)
242 | content = response.choices[0].message.content
243 |
244 | attributes_checker = CustomLLMBooleanEvaluator(
245 | settings=CustomLLMBooleanSettings(
246 | prompt="Does the response list or describe trace attributes that are available for querying?",
247 | )
248 | )
249 | expect(input=prompt, output=content).to_pass(attributes_checker)
250 |
```
--------------------------------------------------------------------------------
/tools/annotations.go:
--------------------------------------------------------------------------------
```go
1 | package tools
2 |
3 | import (
4 | "context"
5 | "fmt"
6 | "github.com/mark3labs/mcp-go/mcp"
7 | "github.com/mark3labs/mcp-go/server"
8 | "strconv"
9 |
10 | mcpgrafana "github.com/grafana/mcp-grafana"
11 |
12 | "github.com/grafana/grafana-openapi-client-go/client/annotations"
13 | "github.com/grafana/grafana-openapi-client-go/models"
14 | )
15 |
16 | // GetAnnotationsInput filters annotation search.
17 | type GetAnnotationsInput struct {
18 | From *int64 `jsonschema:"description=Epoch ms start time"`
19 | To *int64 `jsonschema:"description=Epoch ms end time"`
20 | Limit *int64 `jsonschema:"description=Max results default 100"`
21 | AlertID *int64 `jsonschema:"description=Deprecated. Use AlertUID"`
22 | AlertUID *string `jsonschema:"description=Filter by alert UID"`
23 | DashboardID *int64 `jsonschema:"description=Deprecated. Use DashboardUID"`
24 | DashboardUID *string `jsonschema:"description=Filter by dashboard UID"`
25 | PanelID *int64 `jsonschema:"description=Filter by panel ID"`
26 | UserID *int64 `jsonschema:"description=Filter by creator user ID"`
27 | Type *string `jsonschema:"description=annotation or alert"`
28 | Tags []string `jsonschema:"description=Multiple tags allowed tags=tag1&tags=tag2"`
29 | MatchAny *bool `jsonschema:"description=true OR tag match false AND"`
30 | }
31 |
32 | // getAnnotations retrieves Grafana annotations using filters.
33 | func getAnnotations(ctx context.Context, args GetAnnotationsInput) (*annotations.GetAnnotationsOK, error) {
34 | c := mcpgrafana.GrafanaClientFromContext(ctx)
35 |
36 | req := annotations.GetAnnotationsParams{
37 | From: args.From,
38 | To: args.To,
39 | Limit: args.Limit,
40 | AlertID: args.AlertID,
41 | AlertUID: args.AlertUID,
42 | DashboardID: args.DashboardID,
43 | DashboardUID: args.DashboardUID,
44 | PanelID: args.PanelID,
45 | UserID: args.UserID,
46 | Type: args.Type,
47 | Tags: args.Tags,
48 | MatchAny: args.MatchAny,
49 | Context: ctx,
50 | }
51 |
52 | resp, err := c.Annotations.GetAnnotations(&req)
53 | if err != nil {
54 | return nil, fmt.Errorf("get annotations: %w", err)
55 | }
56 |
57 | return resp, nil
58 | }
59 |
60 | var GetAnnotationsTool = mcpgrafana.MustTool(
61 | "get_annotations",
62 | "Fetch Grafana annotations using filters such as dashboard UID, time range and tags.",
63 | getAnnotations,
64 | mcp.WithTitleAnnotation("Get Annotations"),
65 | mcp.WithIdempotentHintAnnotation(true),
66 | mcp.WithReadOnlyHintAnnotation(true),
67 | )
68 |
69 | // CreateAnnotationInput creates a new annotation.
70 | type CreateAnnotationInput struct {
71 | DashboardID int64 `json:"dashboardId,omitempty" jsonschema:"description=Deprecated. Use dashboardUID"`
72 | DashboardUID string `json:"dashboardUID,omitempty" jsonschema:"description=Preferred dashboard UID"`
73 | PanelID int64 `json:"panelId,omitempty" jsonschema:"description=Panel ID"`
74 | Time int64 `json:"time,omitempty" jsonschema:"description=Start time epoch ms"`
75 | TimeEnd int64 `json:"timeEnd,omitempty" jsonschema:"description=End time epoch ms"`
76 | Tags []string `json:"tags,omitempty" jsonschema:"description=Optional list of tags"`
77 | Text string `json:"text" jsonschema:"description=Annotation text required"`
78 | Data map[string]any `json:"data,omitempty" jsonschema:"description=Optional JSON payload"`
79 | }
80 |
81 | // createAnnotation sends a POST request to create a Grafana annotation.
82 | func createAnnotation(ctx context.Context, args CreateAnnotationInput) (*annotations.PostAnnotationOK, error) {
83 | c := mcpgrafana.GrafanaClientFromContext(ctx)
84 |
85 | req := models.PostAnnotationsCmd{
86 | DashboardID: args.DashboardID,
87 | DashboardUID: args.DashboardUID,
88 | PanelID: args.PanelID,
89 | Time: args.Time,
90 | TimeEnd: args.TimeEnd,
91 | Tags: args.Tags,
92 | Text: &args.Text,
93 | Data: args.Data,
94 | }
95 |
96 | resp, err := c.Annotations.PostAnnotation(&req)
97 | if err != nil {
98 | return nil, fmt.Errorf("create annotation: %w", err)
99 | }
100 |
101 | return resp, nil
102 | }
103 |
104 | var CreateAnnotationTool = mcpgrafana.MustTool(
105 | "create_annotation",
106 | "Create a new annotation on a dashboard or panel.",
107 | createAnnotation,
108 | mcp.WithTitleAnnotation("Create Annotation"),
109 | mcp.WithIdempotentHintAnnotation(false),
110 | )
111 |
112 | // CreateGraphiteAnnotationInput represents the payload format for creating a Graphite-style annotation.
113 | type CreateGraphiteAnnotationInput struct {
114 | What string `json:"what" jsonschema:"description=Annotation text"`
115 | When int64 `json:"when" jsonschema:"description=Epoch ms timestamp"`
116 | Tags []string `json:"tags,omitempty" jsonschema:"description=Optional list of tags"`
117 | Data string `json:"data,omitempty" jsonschema:"description=Optional payload"`
118 | }
119 |
120 | // createAnnotationGraphiteFormat creates an annotation using the Graphite annotation format.
121 | func createAnnotationGraphiteFormat(ctx context.Context, args CreateGraphiteAnnotationInput) (*annotations.PostGraphiteAnnotationOK, error) {
122 | c := mcpgrafana.GrafanaClientFromContext(ctx)
123 |
124 | req := &models.PostGraphiteAnnotationsCmd{
125 | What: args.What,
126 | When: args.When,
127 | Tags: args.Tags,
128 | Data: args.Data,
129 | }
130 |
131 | resp, err := c.Annotations.PostGraphiteAnnotation(req)
132 | if err != nil {
133 | return nil, fmt.Errorf("create graphite annotation: %w", err)
134 | }
135 |
136 | return resp, nil
137 | }
138 |
139 | var CreateGraphiteAnnotationTool = mcpgrafana.MustTool(
140 | "create_graphite_annotation",
141 | "Create an annotation using Graphite annotation format.",
142 | createAnnotationGraphiteFormat,
143 | mcp.WithTitleAnnotation("Create Graphite Annotation"),
144 | mcp.WithIdempotentHintAnnotation(false),
145 | )
146 |
147 | // UpdateAnnotationInput represents the payload used to update an existing annotation by ID.
148 | type UpdateAnnotationInput struct {
149 | ID int64 `json:"id" jsonschema:"description=Annotation ID to update"`
150 | Time int64 `json:"time,omitempty" jsonschema:"description=Start time epoch ms"`
151 | TimeEnd int64 `json:"timeEnd,omitempty" jsonschema:"description=End time epoch ms"`
152 | Text string `json:"text,omitempty" jsonschema:"description=Annotation text"`
153 | Tags []string `json:"tags,omitempty" jsonschema:"description=Tags to replace existing tags"`
154 | Data map[string]any `json:"data,omitempty" jsonschema:"description=Optional JSON payload"`
155 | }
156 |
157 | // updateAnnotation updates an annotation using its ID.
158 | func updateAnnotation(ctx context.Context, args UpdateAnnotationInput) (*annotations.UpdateAnnotationOK, error) {
159 | c := mcpgrafana.GrafanaClientFromContext(ctx)
160 | annotationID := strconv.FormatInt(args.ID, 10)
161 | req := &models.UpdateAnnotationsCmd{
162 | Time: args.Time,
163 | TimeEnd: args.TimeEnd,
164 | Text: args.Text,
165 | Tags: args.Tags,
166 | Data: args.Data,
167 | }
168 |
169 | resp, err := c.Annotations.UpdateAnnotation(annotationID, req)
170 | if err != nil {
171 | return nil, fmt.Errorf("update annotation: %w", err)
172 | }
173 |
174 | return resp, nil
175 | }
176 |
177 | var UpdateAnnotationTool = mcpgrafana.MustTool(
178 | "update_annotation",
179 | "Updates all properties of an annotation that matches the specified ID. Sends a full update (PUT). For partial updates, use patch_annotation instead.",
180 | updateAnnotation,
181 | mcp.WithTitleAnnotation("Update Annotation"),
182 | mcp.WithIdempotentHintAnnotation(false),
183 | )
184 |
185 | // PatchAnnotationInput updates only the provided fields.
186 | type PatchAnnotationInput struct {
187 | ID int64 `json:"id" jsonschema:"description=Annotation ID"`
188 | Text *string `json:"text,omitempty" jsonschema:"description=Optional new text"`
189 | Tags []string `json:"tags,omitempty" jsonschema:"description=Optional replace tags"`
190 | Time *int64 `json:"time,omitempty" jsonschema:"description=Optional new start epoch ms"`
191 | TimeEnd *int64 `json:"timeEnd,omitempty" jsonschema:"description=Optional new end epoch ms"`
192 | Data map[string]any `json:"data,omitempty" jsonschema:"description=Optional metadata"`
193 | }
194 |
195 | // patchAnnotation patches only the provided annotation fields.
196 | func patchAnnotation(ctx context.Context, args PatchAnnotationInput) (*annotations.PatchAnnotationOK, error) {
197 | c := mcpgrafana.GrafanaClientFromContext(ctx)
198 | id := strconv.FormatInt(args.ID, 10)
199 |
200 | body := &models.PatchAnnotationsCmd{}
201 |
202 | if args.Text != nil {
203 | body.Text = *args.Text
204 | }
205 | if args.Time != nil {
206 | body.Time = *args.Time
207 | }
208 | if args.TimeEnd != nil {
209 | body.TimeEnd = *args.TimeEnd
210 | }
211 | if args.Tags != nil {
212 | body.Tags = args.Tags
213 | }
214 | if args.Data != nil {
215 | body.Data = args.Data
216 | }
217 |
218 | resp, err := c.Annotations.PatchAnnotation(id, body)
219 | if err != nil {
220 | return nil, fmt.Errorf("patch annotation: %w", err)
221 | }
222 | return resp, nil
223 | }
224 |
225 | var PatchAnnotationTool = mcpgrafana.MustTool(
226 | "patch_annotation",
227 | "Updates only the provided properties of an annotation. Fields omitted are not modified. Use update_annotation for full replacement.",
228 | patchAnnotation,
229 | mcp.WithTitleAnnotation("Patch Annotation"),
230 | mcp.WithIdempotentHintAnnotation(false),
231 | )
232 |
233 | // GetAnnotationTagsInput defines filters for retrieving annotation tags.
234 | type GetAnnotationTagsInput struct {
235 | Tag *string `json:"tag,omitempty" jsonschema:"description=Optional filter by tag name"`
236 | Limit *string `json:"limit,omitempty" jsonschema:"description=Max results\\, default 100"`
237 | }
238 |
239 | func getAnnotationTags(ctx context.Context, args GetAnnotationTagsInput) (*annotations.GetAnnotationTagsOK, error) {
240 | c := mcpgrafana.GrafanaClientFromContext(ctx)
241 |
242 | req := annotations.GetAnnotationTagsParams{
243 | Tag: args.Tag,
244 | Limit: args.Limit,
245 | Context: ctx,
246 | }
247 |
248 | resp, err := c.Annotations.GetAnnotationTags(&req)
249 | if err != nil {
250 | return nil, fmt.Errorf("get annotation tags: %w", err)
251 | }
252 |
253 | return resp, nil
254 | }
255 |
256 | var GetAnnotationTagsTool = mcpgrafana.MustTool(
257 | "get_annotation_tags",
258 | "Returns annotation tags with optional filtering by tag name. Only the provided filters are applied.",
259 | getAnnotationTags,
260 | mcp.WithTitleAnnotation("Get Annotation Tags"),
261 | mcp.WithIdempotentHintAnnotation(true),
262 | mcp.WithReadOnlyHintAnnotation(true),
263 | )
264 |
265 | func AddAnnotationTools(mcp *server.MCPServer, enableWriteTools bool) {
266 | GetAnnotationsTool.Register(mcp)
267 | if enableWriteTools {
268 | CreateAnnotationTool.Register(mcp)
269 | CreateGraphiteAnnotationTool.Register(mcp)
270 | UpdateAnnotationTool.Register(mcp)
271 | PatchAnnotationTool.Register(mcp)
272 | }
273 | GetAnnotationTagsTool.Register(mcp)
274 | }
275 |
```